Google Research представила Retrieve-for-Train (R4T): диффузионный ретривер, скомпилированный с помощью RL, для ускорения разветвления запросов в 12–20 раз

Системам поиска и рекомендаций всё чаще требуется возвращать набор результатов, а не одно наилучшее совпадение. По запросу вроде «снаряжение для кемпинга» следует возвращать палатку, спальный мешок, горелку и налобный фонарь. Не следует возвращать 10 почти одинаковых палаток.

Google Research представила Retrieve-for-Train (R4T) для решения этой проблемы. Фреймворк один раз использует обучение с подкреплением (RL) в офлайн-режиме, чтобы научиться правильно разветвлять запросы. Затем он дистиллирует это поведение в небольшую диффузионную модель, которая за 1 проход создаёт все направления поиска.

Почему стандартного разветвления недостаточно

Разветвление запросов разделяет один широкий запрос на несколько подзапросов. Исследовательская команда выделяет 2 проблемы, возникающие, когда обычная LLM делает это во время инференса.

Первая — парафразный коллапс. Для запроса «богемный стиль для фестиваля» модель Qwen3-4B без обучения написала «богемная фестивальная мода» и «богемная фестивальная одежда». Эти почти синонимичные варианты возвращают однородный набор результатов.

Вторая — задержка. Авторегрессионная генерация в сочетании с повторными вызовами поиска работает медленно. Сэмплирование Best-of-N повышает качество, но многократно увеличивает стоимость инференса.

Конвейер R4T из 3 этапов

Обучение LM для разветвления: языковая модель разветвления (FOLM) генерирует k подзапросов. Замороженный плотный поисковик выполняет их. Награда на уровне набора оценивает весь полученный набор, а не каждый элемент по отдельности.
Синтез обучающих данных: обученная FOLM сэмплирует 128 вариантов разветвления для каждого запроса при температуре 0,9. Они становятся обучающими парами (запрос, целевой набор) без использования человеческой разметки. Для открытых задач целями являются эмбеддинги найденного контента. Для композиционных задач целями являются эмбеддинги подзапросов.
Обучение диффузионного поисковика: диффузионный трансформер с 53,9 млн параметров учится сопоставлять эмбеддинг запроса с полным набором целевых эмбеддингов. Он использует формулировку с взрывным ростом дисперсии в рамках фреймворка EDM. Во время инференса он генерирует все эмбеддинги за один недautoreгрессионный проход. Затем поиск ближайших соседей сопоставляет каждый эмбеддинг с элементами базы данных.

Проектирование награды и взлом награды

Для открытого абстрактного поиска (OAR) награда объединяет 3 взвешенных компонента:

Обоснованность (λg = 0,6): штрафует за расстояние между эмбеддингом каждого подзапроса и ближайшим элементом базы данных.
Разнообразие (λd = 0,2): оценка Vendi Score для репрезентативных найденных элементов, например элемента top-1 для каждого подзапроса.
Согласованность (λa = 0,2): среднее косинусное сходство между каждым подзапросом и исходным запросом.

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

Абляционное исследование объясняет, почему важны все 3 компонента OAR. При использовании только обоснованности Gemma3-4B сходилась к строкам вроде «конец строки конец строки конец строки». Добавление согласованности ещё быстрее приводило к коллапсу, поскольку политика повторяла парафразы запроса. Добавление разнообразия устраняло оба пути к такому результату.

В обучении используется GRPO с мягкой регуляризацией PPO, которая добавляет прямые и обратные штрафы KL. Ключевые настройки включают размер группы 8, скорость обучения 1×10⁻⁷ и глобальный размер пакета 512.

Интерактивное объяснение

Результаты

В экспериментах использовался датасет модных образов Polyvore с энкодером matryoshka на основе CLIP с размерностью 128. Также использовался проприетарный датасет музыки с экспертными плейлистами и эмбеддингами MuLan. Каждый метод разветвления создавал k = 10 подзапросов, а Best-of-N использовал N = 5.

Качество OAR оценивалось судьёй-LLM по 5-балльной шкале Лайкерта. На Polyvore средний результат Gemma3-4B R4T-FOLM составил 49,1 против 40,9 у Best-of-N и 38,5 в режиме без обучения. Разнообразие выросло с 56,0 без обучения до 76,8. R4T-Diffusion сохранила большую часть этого показателя — 74,3. На Music средний результат Gemma3-4B R4T-FOLM составил 58,1 против 49,2 у Best-of-N. Обоснованность для R4T-Diffusion не приводится, поскольку модель не создаёт текстовые подзапросы.

Результаты WSCR на Polyvore показывают компромисс между полнотой и разнообразием. R4T-FOLM (Qwen) достигла Recall@5K 20,9 и Hit@5K 64,6 против 15,7 и 52,1 у Gemini-2.5-Flash. Однако её оценка Vendi Score снизилась до 27,5. Авторы связывают это со снижением энтропии выходных данных при интенсивной оптимизации с помощью RL. R4T-Diffusion (Qwen) сохранила более высокую оценку Vendi Score — 34,7 — при Recall@5K 16,5.

Эффективность

При размере пакета 8 авторегрессионное разветвление занимало около 1,46 секунды. Диффузионной модели требовалось 0,07 секунды. При размере пакета 1024 авторегрессионное разветвление занимало почти 50 секунд против 4,21 секунды для диффузионной модели. Авторы сообщают об устойчивом ускорении в 12–20 раз.

Основные выводы

R4T однократно выполняет RL в офлайн-режиме, а затем разворачивает диффузионный поисковик с 53,9 млн параметров.
3 вида награды (обоснованность, разнообразие и согласованность) препятствуют взлому награды и парафразному коллапсу.
Gemma3-4B R4T-FOLM повысила средний показатель OAR на Polyvore с 40,9 (Best-of-N) до 49,1.
Диффузионное разветвление работает в 12–20 раз быстрее авторегрессионного.

Ознакомьтесь с исследованием и техническими подробностями. Вся заслуга принадлежит исследователям этого проекта. Также подписывайтесь на нас в Twitter и не забудьте присоединиться к нашему сабреддиту о машинном обучении с более чем 150 тыс. участников и подписаться на нашу рассылку. Постойте! Вы пользуетесь Telegram? теперь вы также можете присоединиться к нам в Telegram.

Хотите сотрудничать с нами в продвижении вашего репозитория GitHub, страницы Hugging Face, выпуска продукта, вебинара и т. д.? Свяжитесь с нами

Переведено автоматически с английского. Оригинал статьи — по ссылке ниже.