🌟 Как YouTube ускорил проверку валидных рекомендаций в LLM в 948 раз.

YouTube и Google DeepMind опубликовали статью и код фреймворка STATIC.

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

Очевидное решение - префиксное дерево: на каждом шаге декодирования маска блокирует невалидные токены. Работает в целом нормально, но убивает производительность на TPU и GPU. Причины две:

🟠Обход через цепочки указателей создает случайный, несмежный паттерн доступа к памяти, а память ускорителей рассчитана на потоковое чтение блоками, а не на такой режим.

🟠Гугловский XLA-компилятор требует статических вычислительных графов, а префиксное дерево с управлением потоком, зависящим от данных, в это не вписывается.

В итоге префиксное дерево на CPU удваивает время инференса, что для системы с целевой задержкой ≤10 мс на шаг неприемлемо.

STATIC меняет подход принципиально: дерево разворачивается в статическую разреженную матрицу формата Compressed Sparse Row. Обход превращается в векторизованную операцию, которую акселератор умеет выполнять нативно.

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

Для глубоких уровней работает специализированное ядро Vectorized Node Transition Kernel, оно читает фиксированный блок данных вне зависимости от реального числа дочерних узлов, не создавая условных переходов. В этом и есть вся соль: весь граф остается статическим, XLA не перекомпилирует ничего на ходу.

🟡Тесты и результаты

Замеры проводили на TPU v6e с 3B-моделью и словарем из 20 млн. свежих видео.

STATIC добился задержки 0,033 мс на шаг декодирования (это всего 0,25% от общего времени инференса.

Для сравнения: префиксное дерево на CPU давало +31,3 мс (239% от инференса), лучший из конкурирующих методов непосредственно на акселераторе, PPV Approximate +1,56 мс (11,9%). Итоговый профит:

🟢948x против дерева на CPU;
🟢47x против PPV Approximate;
🟢1033x против точного PPV.

По памяти: ~90 МБ на 1 млн. элементов. Для словаря в 20 млн. верхний предел по HBM примерно 1,5 ГБ, на практике - около 75% от этого значения.

🟡STATIC развернут на YouTube в продакшне.

A/B-тест с условием «только видео за последние 7 дней» показал +5,1% просмотров свежего контента, +2,9% для трехдневного окна, +0,15% по CTR и 100% соответствие бизнес-правилам.

Дополнительный бонус: метод решает проблему рекомендации новых товаров, не представленных в обучающей выборке. На датасете Amazon Reviews Recall@1 вырос с 0% до 1,2–4,4% в зависимости от категории.

В репозитории лежит ноутбук, на котором без утомительной настройки можно сразу посмотреть, как строится индекс из Semantic ID и как запускается декодирование с ограничениями.


📌Лицензирование: Apache 2.0 License.


🟡Arxiv
🖥GitHub


@ai_machinelearning_big_data

#AI #ML #LLM #STATIC #DeepMind