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

В промышленных системах рекомендаций переход к генеративному поиску (Generative Retrieval, GR) заменяет традиционный поиск по вложенностям ближайших соседей на большие языковые модели (LLMs). Эти модели представляют элементы в виде семантических идентификаторов (SIDs) — дискретных последовательностей токенов — и рассматривают поиск как задачу авторегрессионного декодирования. Однако промышленные приложения часто требуют строгого соблюдения бизнес-логики, например, обеспечения актуальности контента или наличия товаров в наличии. Стандартное авторегрессионное декодирование не может нативно обеспечить эти ограничения, что часто приводит к тому, что модель «выдумывает» недопустимые идентификаторы элементов или идентификаторы товаров, которых нет в наличии.

Узкое место ускорителя: попытки против TPU/GPU

Чтобы обеспечить допустимый результат, разработчики обычно используют префиксное дерево (trie) для маскировки недопустимых токенов на каждом этапе декодирования. Хотя концептуально это просто, традиционные реализации trie неэффективны на аппаратных ускорителях, таких как TPU и GPU.

Причины низкой эффективности:
* Задержка памяти: структуры, основанные на указателях, приводят к несмежному, случайному доступу к памяти. Это препятствует объединению памяти и не позволяет использовать возможности памяти с высокой пропускной способностью (HBM) современных ускорителей.
* Несовместимость компиляции: ускорители полагаются на статические вычислительные графы для компиляции машинного обучения (например, Google XLA). Стандартные trie используют зависимый от данных поток управления и рекурсивное ветвление, которые несовместимы с этой парадигмой и часто приводят к дорогостоящим циклам хост-устройство.

STATIC: ускоренный индекс trie на основе разреженной матрицы перехода

Исследователи Google DeepMind и YouTube представили STATIC (Sparse Transition Matrix-Accelerated Trie Index for Constrained Decoding) для решения этих проблем. Вместо того чтобы рассматривать trie как граф для обхода, STATIC преобразует его в статическую матрицу Compressed Sparse Row (CSR). Это преобразование позволяет выполнять нерегулярные обходы дерева в виде полностью векторизованных операций с разреженными матрицами.

Гибридная архитектура декодирования

STATIC использует двухэтапную стратегию поиска, чтобы сбалансировать использование памяти и скорость:
* Плотное маскирование (t-1 < d): для первых d=2 слоёв, где коэффициент ветвления наибольший, STATIC использует битовый плотный логический тензор. Это позволяет выполнять поиск за O(1) на самых вычислительно затратных начальных этапах.
* Векторизованное ядро перехода узлов (VNTK): для более глубоких слоёв (l ≥ 3) STATIC использует ядро без ветвлений. Это ядро выполняет «спекулятивный срез» фиксированного количества записей (Bt), соответствующего максимальному коэффициенту ветвления на этом уровне. Используя срез фиксированного размера независимо от фактического количества дочерних элементов, весь процесс декодирования остаётся единым статическим вычислительным графом.

Такой подход обеспечивает сложность ввода-вывода O(1) относительно размера набора ограничений, тогда как предыдущие методы двоичного поиска с аппаратным ускорением масштабировались логарифмически (O(log|C|)).

Производительность и масштабируемость

Оценено на ускорителях Google TPU v6e с использованием модели на 3 миллиарда параметров с размером пакета 2 и размером луча (M) 70, STATIC продемонстрировал значительное повышение производительности по сравнению с существующими методами.

| Метод | Задержка на шаг (мс) | % от общего времени вывода |
| — | — | — |
| STATIC (Ours) | +0,033 | 0,25% |
| PPV Approximate | +1,56 | 11,9% |
| Hash Bitmap | +12,39 | 4,0% |
| CPU Trie | +31,32 | 39% |
| PPV Exact | +34,12 | 60% |

STATIC достиг ускорения в 948 раз по сравнению с CPU-выгруженными trie и превзошёл точный базовый уровень двоичного поиска (PPV) в 1033 раза. Его задержка остаётся практически постоянной даже при увеличении размера словаря семантических идентификаторов (|V|).

Результаты развёртывания

STATIC был развёрнут на YouTube для обеспечения соблюдения ограничения актуальности «за последние 7 дней» для видеорекомендаций. Система обслуживала словарь из 20 миллионов свежих элементов со 100% соблюдением требований.

Онлайн-A/B-тестирование показало:
* Увеличение просмотров свежих видео за 7 дней на 5,1%.
* Увеличение просмотров свежих видео за 3 дня на 2,9%.
* Увеличение коэффициента кликабельности (CTR) на 0,15%.

Холодный старт

Фреймворк также решает проблему «холодного старта» генеративного поиска — рекомендации элементов, не увиденных во время обучения. Ограничивая модель набором элементов «холодного старта» на наборах данных Amazon Reviews, STATIC значительно улучшил производительность по сравнению с неограниченными базовыми показателями, которые зафиксировали 0,00% Recall@1. Для этих тестов использовалась архитектура Gemma с 1 миллиардом параметров, L = 4 токена и размером словаря |V|=256.

Ключевые выводы

* Векторизованная эффективность: STATIC преобразует декодирование с ограничениями из задачи обхода графа в аппаратно-ориентированные векторизованные операции с разреженными матрицами, преобразуя префиксные деревья в статические матрицы Compressed Sparse Row (CSR).
* Массивное ускорение: система обеспечивает задержку в 0,033 мс на шаг, что представляет собой ускорение в 948 раз по сравнению с CPU-выгруженными trie и в 47–1033 раза по сравнению с базовыми показателями с аппаратным ускорением двоичного поиска.
* Масштабируемая сложность O(1): достигая сложности ввода-вывода O(1) относительно размера набора ограничений, STATIC поддерживает высокую производительность при низком использовании памяти — примерно 90 МБ на 1 миллион элементов.
* Результаты, проверенные в производстве: развёртывание на YouTube показало 100% соблюдение бизнес-логики, что привело к увеличению просмотров свежих видео на 5,1% и повышению коэффициента кликабельности на 0,15%.
* Решение для холодного старта: фреймворк позволяет моделям генеративного поиска успешно рекомендовать элементы «холодного старта», повышая производительность Recall@1 с 0,00% до нетривиального уровня на тестах Amazon Reviews.

1. Какие проблемы решает система STATIC в контексте генеративного поиска?

В статье указано, что система STATIC решает проблему обеспечения допустимых результатов при декодировании с ограничениями в промышленных приложениях. Это достигается за счёт использования разреженных матриц для ускорения процесса поиска и обхода префиксного дерева.

2. Какие преимущества предоставляет STATIC по сравнению с традиционными методами поиска?

STATIC демонстрирует значительное ускорение по сравнению с существующими методами, такими как CPU Trie, PPV Approximate, Hash Bitmap и PPV Exact. Это достигается за счёт векторизации операций с разреженными матрицами и преобразования префиксных деревьев в статические матрицы Compressed Sparse Row (CSR).

3. Какие результаты были получены при развёртывании STATIC на YouTube?

При развёртывании на YouTube система STATIC обеспечила соблюдение ограничения актуальности «за последние 7 дней» для видеорекомендаций. Также было отмечено увеличение просмотров свежих видео за 7 дней на 5,1% и за 3 дня на 2,9%, а также повышение коэффициента кликабельности (CTR) на 0,15%.

4. Как STATIC решает проблему «холодного старта» в генеративном поиске?

Фреймворк STATIC решает проблему «холодного старта» за счёт ограничения модели набором элементов «холодного старта» на наборах данных Amazon Reviews. Это позволило значительно улучшить производительность по сравнению с неограниченными базовыми показателями, которые зафиксировали 0,00% Recall@1.

5. Какие ключевые выводы можно сделать из статьи о системе STATIC?

Ключевые выводы включают:
* Векторизованная эффективность STATIC преобразует декодирование с ограничениями из задачи обхода графа в аппаратно-ориентированные векторизованные операции с разреженными матрицами.
* Массивное ускорение система обеспечивает задержку в 0,033 мс на шаг, что представляет собой значительное ускорение по сравнению с традиционными методами.
* Масштабируемая сложность O(1) STATIC поддерживает высокую производительность при низком использовании памяти.
* Результаты, проверенные в производстве, развёртывание на YouTube показало 100% соблюдение бизнес-логики и положительные результаты в виде увеличения просмотров и CTR.
* Решение для холодного старта фреймворк позволяет моделям генеративного поиска успешно рекомендовать элементы «холодного старта».

Источник