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