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 підвищує якість, але збільшує вартість інференсу.

3-етапний конвеєр R4T

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

Проєктування нагороди та злам нагороди

Для відкритого абстрактного пошуку (OAR) нагорода поєднує 3 зважені компоненти:

Обґрунтованість (λg = 0,6): штрафує за відстань між ембеддингом кожного підзапиту та найближчим елементом бази даних.
Різноманітність (λd = 0,2): оцінка Vendi для репрезентативних отриманих елементів, наприклад елемента 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 знизилася до 27,5. Автори пов’язують це зі зменшенням ентропії виходу за умов сильної оптимізації RL. R4T-Diffusion (Qwen) зберегла вищу оцінку Vendi — 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, релізу продукту, вебінару тощо? Зв’яжіться з нами

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