Новый алгоритм Prof-K ускоряет выборку top-k в 1,5-10 раз для ИИ и баз данных

Операция выбора top-k — извлечение k наибольших или наименьших элементов из массива — лежит в основе множества систем: от баз данных и поиска до обработки сигналов и современных моделей машинного обучения, включая разреженные активации и обрезку механизмов внимания. С ростом объёмов данных классические точные методы требуют всё больше памяти и вычислений, а приближённые часто зависят от хрупких эвристик, которые дают сбой на сложных или неравномерно распределённых данных.

В новой работе, опубликованной в архиве препринтов arXiv, представлен алгоритм Prof-K. Его ключевая особенность — вероятностная гарантия корректности, которая не зависит от распределения входных данных. Метод выполняет однопроходную фильтрацию: сначала по небольшой случайной выборке оценивается адаптивный порог, затем все N элементов за один проход записываются в компактный буфер, после чего точная процедура top-k на этом буфере восстанавливает истинные k элементов с вероятностью не менее 1 - ?, где ? задаётся пользователем.

Авторы вывели теоретические оценки для вероятности успеха и требуемого размера буфера, а также нашли приближённо оптимальный объём выборки, который минимизирует накладные расходы в зависимости от N и k. Это позволяет избежать перерасхода памяти и времени, характерного для точных алгоритмов, и одновременно гарантирует устойчивость к «атакующим» или тяжелохвостовым распределениям, где многие приближённые методы деградируют.

Эксперименты показали, что Prof-K достигает ускорения от 1,5 до 10 раз по сравнению с высокооптимизированной функцией topk из PyTorch и недавней реализацией RadiK. Наибольший выигрыш получен в режиме больших массивов и малого или умеренного k — именно там, где существующие подходы испытывают наибольшие трудности. Дополнительно алгоритм позволяет гибко управлять точностью: если отказаться от полного восстановления top-k и согласиться на 95% совпадение, можно ещё сильнее ускорить вычисления.

Практическая значимость работы подтверждена экспериментом с обучением разреженных автоэнкодеров BatchTopK, где выбор top-k составляет значительную часть вычислительных затрат. По словам авторов, использование Prof-K позволяет заметно сократить время тренировки таких моделей, не жертвуя качеством. В перспективе этот алгоритм может найти применение в recommendation systems, обработке графов и любых средах, где требуется быстрая и надёжная фильтрация экстремальных значений.