HNSW-индекс ускорил декодирование Gemma 3 270M на величину до 82%

Исследователи предложили способ ускорить автогенерацию текста языковыми моделями (LLM), заменив стандартную плотную проекцию выходного слоя на индекс приближённого поиска ближайших соседей по алгоритму HNSW (Hierarchical Navigable Small World).

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

Авторы переформулировали связку "проекция + отбор топ-k токенов" как задачу поиска максимального внутреннего произведения (maximum inner product search) между вектором скрытого состояния и эмбеддингами токенов словаря. Вместо перемножения со всей матрицей метод строит HNSW-индекс по эмбеддингам токенов и извлекает лишь небольшой набор кандидатов с наибольшим скорингом; извлечённые логиты затем "разбрасываются" в разреженный тензор полного размера словаря, так что метод встраивается в существующие пайплайны декодирования без переделки остальной части модели.

Метод проверили на CPU-инференсе трёх моделей, Gemma 3, Llama 3.2 и Qwen 3. Числовой результат приведён только для одной конфигурации: на модели Gemma 3 270M при батче размера 1 сквозная (end-to-end) пропускная способность декодирования выросла на величину до 82% по сравнению с базовым методом; для Llama 3.2 и Qwen 3 количественные цифры в тексте не даны. Качество генерации проверили на бенчмарке AlpacaEval, по заявлению авторов, оно сохраняется, но конкретный балл в тексте не приведён.

Авторы заключают, что приближённый поиск, практичная альтернатива плотной проекции выходного слоя в сценариях декодирования с малым батчем, где критична задержка.

Ключевые факты

  • Метод заменяет плотную проекцию выходного слоя LLM на приближённый поиск максимального внутреннего произведения (MIPS) через HNSW-индекс по эмбеддингам токенов словаря.
  • Извлекается лишь небольшой набор кандидатных токенов с высоким скорингом; логиты "разбрасываются" в разреженный тензор полного словаря, что позволяет встроить метод в существующие пайплайны декодирования.
  • На CPU-инференсе модели Gemma 3 270M при батче размера 1 сквозная пропускная способность декодирования выросла на величину до 82%.
  • Метод также проверен на Llama 3.2 и Qwen 3, но числовые показатели для них в статье не приведены.
  • Качество генерации проверено на бенчмарке AlpacaEval и, по словам авторов, сохраняется, конкретный балл не указан.

Почему это важно

На каждом шаге автогенерации языковая модель должна перемножить скрытое состояние на матрицу эмбеддингов всего словаря, чтобы выбрать следующий токен. Для компактных моделей с большими многоязычными словарями эта матрица не помещается в быстрый кэш процессора и становится узким местом по пропускной способности памяти, именно она, а не арифметика, ограничивает скорость генерации при батче размера 1. Метод сводит эту операцию к приближённому поиску ближайших соседей через HNSW-индекс и извлекает лишь небольшой набор кандидатных токенов вместо перемножения со всей матрицей, что и даёт прирост пропускной способности до 82% на модели Gemma 3 270M.

Кому это важно

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

Как это применить

Метод встраивается в существующий пайплайн декодирования: вместо плотного перемножения со всей матрицей эмбеддингов строится HNSW-индекс по эмбеддингам токенов словаря, из него извлекается небольшой набор кандидатов с наибольшим скорингом, а полученные логиты "разбрасываются" в разреженный тензор полного размера словаря, остальная часть модели и логика отбора топ-k токенов не меняются. В статье метод проверен на CPU-инференсе для Gemma 3, Llama 3.2 и Qwen 3; конкретный числовой прирост (до 82% пропускной способности) приведён только для Gemma 3 270M при батче размера 1.

Можно ли доверять

Имена авторов и институты в тексте не названы. Числовой результат (до 82%) приведён только для одной модели и одной конфигурации, Gemma 3 270M, CPU, батч размера 1; для Llama 3.2 и Qwen 3 количественных цифр нет, как и абсолютных значений задержки, пропускной способности или потребления памяти, только относительный процент. Сохранение качества генерации проверено на бенчмарке AlpacaEval, но конкретный балл в тексте не приведён, только качественное утверждение авторов. Название конференции или издания, а также дата публикации в тексте не указаны.

Риски и подводные камни

Приближённый поиск ближайших соседей по своей природе может пропустить истинный топ-k токен, если индекс HNSW его не найдёт, цена скорости в точности количественно не оценена, кроме общего утверждения о сохранении качества по AlpacaEval. Выигрыш показан только для наименьшей из трёх моделей (Gemma 3 270M) и только при батче размера 1 на CPU, неясно, сохраняется ли прирост при большем батче, на GPU или для более крупных моделей. Не приведены цифры по времени построения индекса, его памяти или размеру по сравнению с плотной матрицей эмбеддингов, что мешает оценить реальную экономию ресурсов на практике.