Ускорение моделирования и поиска ближайших соседей. Обсудить

Текст публикуется в переводе, автор оригинальной статьи Max Miller.

До начала пандемии я опубликовал статью, обсуждавшую создание эпидемиологических моделей на Python с целью создания анимированной визуализации – упрощенных версий анимированных изображений, которые Washington Post и 3Blue1Brown использовали в своих популярных постах на тему пандемии.

Игрушечная модель эпидемии, в которой красные точки представляют зараженных людей

В то время моей целью в основном была работа с анимационной библиотекой Matplotlib, чтобы создать визуализации, и в результате код, который я собрал в спешке, чтобы создать эти модели, не был особенно хорошим. Я имею в виду не только то, что модели были не особенно интересными – Грант Сандерсон из 3Blue1Brown выполнил реально сложную интеллектуальную работу по экспериментированию с различными параметрами и настройками моделей, чтобы из них можно было почерпнуть какие-то догадки, а я всего лишь поддерживал занавес в некоторых технических аспектах – но также и то, что их код был неэффективным, что стало бы серьезной проблемой, если бы я попытался масштабировать эти модели.

Как вы можете представить, вероятно, самый важный шаг в эпидемиологической модели – это моделирование передачи болезни. В каждом моделируемом периоде времени наша модель должна определить, подхватит ли уязвимый и в данный момент не зараженный человек болезнь от того, кто в данный момент заражен. В простой модели, которую мы сделали в прошлый раз, это включает определение, есть ли зараженные индивиды в заданном радиусе заражения, и именно в этом заключается проблема. Компьютеру не так-то легко понять, находятся ли индивиды рядом друг с другом, особенно учитывая то, что мы моделируем их передвижение, и единственное, что компьютер может сделать – это взять двух индивидов и рассчитать расстояние между ними. Чтобы гарантировать, что компьютер не пропустил ни одной потенциальной пары индивидов, ему придется проверить множество пар, которые явно не находятся на достаточно близком для заражения расстоянии.

Простейший способ гарантировать, что вы не упустили ни одной возможной пары – это устроить цикл по всем индивидам, который сработает, но в конце концов потребует огромного количества вычислений, поскольку количество всех возможных пар пропорционально квадрату количества индивидов. Несмотря на то, что вычисление расстояния для одной пары требует немного времени, при масштабировании модели до десятков тысяч индивидов количество вычислений станет настолько большим, что процесс расчета займет значительное время.

Вопрос, на который мы должны ответить – это как сократить количество вычислений, не рискуя упустить какие-либо возможные заражения. Я собираюсь исследовать одну из возможных стратегий добиться этого, включающую разделение пространства. Как оказалось, это важно не только для того вида моделирования, которое мы здесь обсуждали, но также для широкого спектра других вычислительных областей, включая компьютерную графику и некоторые алгоритмы data science, такие, как модели «ближайших соседей», включающие похожую задачу поиска среди точек тренировочного набора данных точки, ближайшей к новой точке. Мы начнем с рассмотрения простого, но эффективного метода разделения пространства, обсудим некоторые его ограничения, а затем исследуем немного другой подход, деревья разделения пространства и методы их использования для нашей цели и других приложений data science.

Простые разделения пространства

Чтобы продемонстрировать, что мы здесь делаем, давайте визуализируем нашу задачу. Вспомните, что мы начали с попыток моделирования распространения гипотетической болезни от зараженных индивидов к уязвимым индивидам. Мы можем визуализировать начальное состояние нашей модели примерно так (красные точки представляют зараженных индивидов, а синие – не зараженных):

Начальное состояние простой модели разделения пространства

Зараженные индивиды имеют некоторую вероятность заражения любых в данное время уязвимых индивидов в определенном радиусе. Мы можем добавить этот радиус возможного заражения к нашей визуализации:

Каждый зараженный индивид представляет риск для уязвимых индивидов в определенном радиусе

На начальной стадии модели некоторые из уязвимых индивидов находятся на безопасном расстоянии от всех зараженных, а другие – не находятся. Мы видим это совершенно ясно, но как нам заставить компьютер, не имеющий нашей визуальной интуиции, определить, находится ли данный индивид в опасности заражения или нет? Единственное, что компьютер может сделать – это рассчитать расстояние между двумя индивидами и проверить, достаточно ли оно мало, чтобы возникла опасность заражения. Придется ли нашему компьютеру рассчитывать это расстояние между каждыми двумя индивидами в популяции?

Вы могли подумать, что, хотя компьютеру, очевидно, придется проверять множество пар индивидов, некоторые их этих пар мы могли бы отбраковать без всяких вычислений, поскольку очевидно, что они достаточно далеки друг от друга. Возможно, мы могли бы разделить нашу схему на квадранты примерно так:

Разделение индивидов на квадранты

Мы еще не решили, кто, скорее всего, рядом с кем располагается, но мы, по крайней мере, можем утверждать, что некто из верхнего левого квадранта едва ли будут рядом с кем-либо из нижнего правого. Таким образом, нам нужно протестировать каждого индивида только на близость с другими индивидами из его квадранта, а не со всеми индивидами модели. Вы, наверное, уже заметили одну проблему этой стратегии: некоторые индивиды расположены близко к границам квадранта и могут быть рядом с индивидами из других квадрантов, также близкими к границе квадранта. Мы можем справиться с этой проблемой, приписав индивидов, находящихся близко к границе между квадрантами, к обоим квадрантам сразу: к тому, в котором они реально находятся, и к тому, к которому они находятся достаточно близко, чтобы иметь потенциальную возможность взаимодействия с его индивидами. Это можно представить, как небольшое расширение границ каждого квадранта, чтобы они занимали чуть больше четверти схемы и немного перекрывались:

Разделение на «квадранты», немного большие четверти изображения и слегка перекрывающиеся

Это простое равномерное сеточное разбиение моделируемого пространства, и мы можем использовать его для существенного ускорения нашей модели. Теперь на каждом шагу времени, который мы моделируем, мы сначала проходим по всем индивидам и назначаем каждому из них квадрант (или несколько квадрантов, если они близки к границе квадранта). Затем, когда мы проверяем возможность заражения, компьютеру приходится считать только расстояния между элементами того же квадранта. Мы добавили в наш процесс моделирования дополнительный шаг – определение, в какой квадрант попадает каждый элемент – но мы сократили количество вычислений на шаге проверки заражения намного больше.

Даже это простейшее сеточное разбиение может обеспечить огромное ускорение вычислений. Когда я добавил простое разбиение вроде этого в свою модель, использованную для генерации визуализаций, я сократил требуемое время более чем вдвое. Повышение количества секторов, на которые мы разбиваем пространство, может дать еще больший прирост скорости, но до определенного предела.

Конечно, у этой стратегии есть недостатки. Первый, неизбежный для всех разделений пространства такого типа, заключается в повышенных требованиях к памяти, поскольку, кроме всего прочего, что наша модель должна держать в локальной памяти, нам нужна какая-то таблица или список, хранящая информацию о том, в каких секторах находится каждый индивид. Фактически, мы «купили» скорость модели, заплатив за нее памятью. Для наших маленьких игрушечных эпидемиологических моделей это всего лишь мелочь, но для большой модели или для приложения компьютерной графики, где памяти может не хватать, это придется принимать во внимание.

Вторая проблема, относящаяся только к данному методу разделения пространства, заключается в том, что такие равномерные сетки далеко не всегда эффективно делят пространство. Смотрите, что произойдет с нашим примером, если увеличить количество разбиений по обеим координатам до сетки 6*6:

Разбиение пространства на 36 секторов

С одной стороны, это выглядит даже лучше, чем квадранты: количество индивидов в каждом секторе довольно мало, так что вам придется проверить совсем немного комбинаций. Но такое разбиение неэффективно для других целей. Из-за случайного распределения индивидов некоторые секторы будут содержать намного больше индивидов, чем прочие, а некоторые вообще окажутся пустыми. Эта неэффективность вырастает в большую проблему при большем количестве измерений – если наша модель в трехмерном пространстве или мы разбиваем набор данных с большим количеством переменных. При множестве переменных пространство будет более разреженным, и равномерная сетка вроде этой может стать менее целесообразной.

K-мерные деревья

Один из способов справиться с этой неэффективностью, а также получить вдобавок пару других преимуществ – использовать более умный метод разделения пространства, который не даст в результате равномерную сетку. Популярный подход – использовать алгоритм K-мерных деревьев. Этот алгоритм начинает с выбора оси в наборе данных, находит медианное значение координат по этой оси для всех точек и делит пространство по этой оси в точке медианы. В нашем примере давайте начнем с оси x. Найдем значение медианы по x и нарисуем в этом месте разделяющую линию:

Первое разделение пространства в нашем K-мерном дереве

Как видим, эта линия чуть правее центра схемы, поскольку случайно сгенерированных точек справа оказалось немного больше, чем слева, так что медиана сдвинулась вправо. Теперь, для каждого региона по обе стороны от этой линии мы переключаемся на ось y и делим пространство тем же способом, вычисляя медиану по y только для точек этого региона:

Второе разделение пространства в K-мерном дереве, на этот раз по y

Мы можем продолжать так и дальше, постоянно меняя ось и разбивая каждое пространство. Вот как будет выглядеть следующий уровень разбиения:

Еще одно разделение по оси x

А вот так – еще один уровень:

Это результирующее разделение пространства может выглядеть не так аккуратно, как наша простая равномерная сетка, но у этого метода есть два важных преимущества. Во-первых, итоговое разбиение «сбалансировано», то есть каждая созданная область пространства имеет примерно одинаковое количество элементов – это будет чрезвычайно полезным при повышении количества измерений. Во-вторых, алгоритм естественным образом строит древовидную структуру данных. Алгоритм K-мерных деревьев создает не сетку, а последовательности вложенных порогов. На первом шаге нашего разбиения точки с x-координатами больше порогового значения пошли в правую область пространства, а точки с меньшими – в левую. В каждой области есть новый порог разбиения по оси y, и так далее. В некоторых задачах компьютерам проще передвигаться по древовидным структурам данных, чем по данным в виде сетки. Эта древовидная структура – причина того, почему алгоритм K-мерных деревьев так полезен в алгоритмах поиска вроде моделей KNN (K ближайших соседей). Давайте посмотрим, как они работают.

K-мерные деревья для KNN

Давайте рассмотрим немного другую задачу. Предположим, у нас есть некоторый набор точек, который мы разбили с помощью алгоритма K-мерных деревьев:

K-мерное древовидное разбиение нового набора данных

А теперь мы вводим новую целевую точку, для которой попробуем оценить какое-либо значение. На следующей схеме она показана красным цветом:

Какая точка ближайшая к красной?

Вопрос, который мы ставим перед компьютером: какие точки из набора находятся ближе всего к новой целевой точке? Как и в нашем прошлом упражнении в моделировании, мы хотим избежать расчета расстояний между целевой точкой и каждой точкой из набора. Древовидная структура K-мерного дерева позволит нам эффективно сузить область поиска. Мы можем начать, просто следуя по структуре дерева и смотря, в какую область попадает наша новая точка. Сначала мы разделили пространство по оси x, и целевая точка попадает в правую область этого разделения:

Мы следуем по дереву, выбирая правую область на первом разделении пространства

Затем идет разделение по оси y, и новая точка попадает в нижнюю область:

На втором разделении наша новая точка попала в нижнюю область

Потом следует еще одно разделение по оси x:

И так далее, пока область не будет содержать совсем немного точек:

Нахождение расположения целевой точки в дереве не потребовало расчета расстояния между ней и множеством других точек – всего лишь сравнения ее координат x и y с четырьмя или пятью пороговыми значениями. Это и делает дерево крайне полезным инструментом для сокращения количества потенциальных кандидатов на роль точек, ближайших к целевой.

Конечно, иногда возникает осложнение, если координата целевой точки оказывается близкой к линии разделения, поскольку при этом ее ближайшие соседи могут оказаться за порогом и совсем на другой ветви дерева. Это значит, что в процессе поиска по дереву вам иногда придется возвращаться назад и проверять области по другую сторону порога. Но даже в этом случае дерево позволит вам быстро исключать целые регионы пространства/поддеревья. Ни одна точка из соседней области пространства не может быть ближе к целевой, чем пороговое значение, так что мы можем исключать области, если расстояние до их линии разделения больше, чем расстояние до ближайшего соседа из своей области.

Примерно так работает реализация KNN в широко используемой библиотеке sklearn, когда вы обучаете модель. Эта реализация разбивает тренировочный набор данных, используя что-то вроде K-мерного дерева. Когда приходит время использовать эту модель, чтобы предсказать значения для новых точек, модель должна как можно быстрее найти несколько ближайших соседей новой точки, что она и делает с помощью этого дерева.

Деревья шаров (Ball trees)

Есть и другой популярный вид деревьев разделения пространства, используемый алгоритмами поиска, достигающий похожих целей немного другим способом. Это алгоритм «дерева шаров», группирующий точки во вложенные «шары» или сферы. Он работает примерно так: сначала определяем центроид для всех точек. Затем находим точку, наиболее удаленную от этого центроида, она будет первым «дочерним» узлом нашего графа. Расстояние от центроида до этой точки – «радиус» первого шара, причем ни одна точка из набора не выйдет за пределы «шара» этого радиуса с центром в центроиде. Самая дальняя точка от первой дочерней станет вторым дочерним узлом. Остальные точки группируются вокруг первого или второго дочернего узла в зависимости от того, какой из них ближе. Используя наши данные из предыдущего примера, мы можем визуализировать первое разбиение пространства таким образом:

Первое разбиение алгоритма дерева шаров. Два дочерних узла обозначены звездами.

Мы можем последовательно продолжать этот процесс, разделяя каждый шар на два с помощью нахождения его центроида и дочерних узлов на следующие группы:

Дальнейшее разделение

И еще раз:

Разумеется, мы могли бы продолжать этот процесс и дальше. Вложенные группы можно было изобразить в виде шаров, но я обозначил их рамками, чтобы сделать визуализации более разборчивыми. Как и алгоритм K-мерных деревьев, алгоритм дерева шаров создает древовидную структуру, хотя здесь она не сбалансирована. Точно так же, как расстояние до разделительных линий в K-мерном дереве позволяло нам проверить, должны ли мы искать возможных соседей и в соседней области, или ее можно полностью дисквалифицировать, для той же цели можно использовать радиусы и расстояния до центроидов в деревьях шаров.