R4T от Google: диффузионный ретривер ускоряет поиск в 12–20 раз

Современные поисковые и рекомендательные системы должны возвращать набор разнообразных результатов, а не один лучший. Например, на запрос «походное снаряжение» нужны палатка, спальник, горелка и фонарь, а не десять почти одинаковых палаток. Google Research предложила Retrieve-for-Train (R4T) — фреймворк, который с помощью обучения с подкреплением (RL) один раз, офлайн, находит оптимальное разбиение запроса на подзапросы (fan-out), а затем дистиллирует это поведение в маленькую диффузионную модель (53,9 млн параметров). На инференсе модель генерирует все направления поиска за один неавторегрессивный проход, что даёт ускорение в 12–20 раз.

Стандартный fan-out страдает от двух проблем. Первая — парафрастический коллапс: LLM генерирует почти одинаковые подзапросы, и результаты дублируются. Вторая — задержка: авторегрессивная генерация плюс многократные вызовы поиска медленны. R4T решает это через трёхшаговый пайплайн. Сначала обучается языковая модель fan-out (FOLM), которая генерирует k подзапросов, а замороженный плотный ретривер выполняет их. Награда считается для всего набора, а не для каждого элемента. Затем FOLM семплирует 128 вариантов на запрос, создавая обучающие пары (запрос, целевой набор) без человеческой разметки. Наконец, диффузионный трансформер учится отображать эмбеддинг запроса в полный набор целевых эмбеддингов.

Особое внимание уделено дизайну функции награды, чтобы избежать reward hacking. Для абстрактного поиска (OAR) награда комбинирует три взвешенных компонента: обоснованность (λ=0,6), разнообразие (λ=0,2) и согласованность с исходным запросом (λ=0,2). Абляции показали, что без компонента разнообразия модель схлопывается в повторяющиеся строки. Для композиционного поиска (WSCR) награда — доля элементов эталонного набора, которые удалось найти. Обучение использует GRPO с мягкой PPO-регуляризацией, группа из 8, learning rate 1×10⁻⁷, глобальный batch 512.

R4T — это практичный способ ускорить и улучшить качество поиска в продакшене, не требуя дорогого инференса LLM на каждом запросе. Диффузионная модель размером 53 млн параметров заменяет генерацию десятков подзапросов большой моделью, а ближайшие соседи по эмбеддингам быстро находят нужные объекты в базе. Результат — более разнообразные и релевантные наборы результатов при радикально меньшей задержке.