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

Обучение на LM за разклоняване: Езиков модел за разклоняване (FOLM) генерира k подзаявки. Замразен плътен ретривър ги изпълнява. Награда на ниво набор оценява целия извлечен набор, а не всеки елемент поотделно.
Синтезиране на надзорни данни: Обученият FOLM семплира 128 разклонения за всяка заявка при температура 0.9. Те се превръщат в обучаващи двойки (заявка, целеви набор) без човешки етикети. За отворени задачи целите са ембедингите на извлеченото съдържание. За композиционни задачи целите са ембединги на подзаявките.
Обучение на дифузионен ретривър: Дифузионен трансформър с 53,9 млн. параметъра се научава да преобразува ембединг на заявка в пълен набор от целеви ембединги. Той използва формулировка с експлодираща дисперсия в рамките на рамката EDM. При инференция генерира всички ембединги в един-единствен неавторегресивен проход. След това търсенето на най-близки съседи съпоставя всеки ембединг с елементи от базата данни.

Проектиране на наградата и хакване на наградата

За отворено абстрактно извличане (OAR) наградата комбинира 3 претеглени компонента:

Обоснованост (λg = 0.6): наказва разстоянието между всеки ембединг на подзаявка и най-близкия елемент в базата данни.
Разнообразие (λd = 0.2): оценката Vendi върху представителни извлечени елементи, например елемента с ранг 1 за всяка подзаявка.
Съответствие (λa = 0.2): средната косинусова близост между всяка подзаявка и оригиналната заявка.

За композиционно извличане със слаб надзор (WSCR) наградата е делът на елементите от референтния набор, които разклонението извлича.

Аблационният анализ обяснява защо и 3-те компонента на OAR са важни. Само с обоснованост Gemma3-4B се насочи към низове като „край на ред край на ред край на ред“. Добавянето на съответствие ускори колапса още повече, тъй като политиката повтаряше парафрази на заявката. Добавянето на разнообразие премахна и двата лесни пътя.

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

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

Резултати

Експериментите използваха набора от данни за модни тоалети Polyvore с базиран на 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 повиши средната стойност на Polyvore OAR от 40.9 (Best-of-N) до 49.1.
Дифузионното разклоняване работи 12× до 20× по-бързо от авторегресивното разклоняване.

Разгледайте статията и техническите подробности. Цялата заслуга е на изследователите на този проект. Също така, не се колебайте да ни последвате в Twitter и не забравяйте да се присъедините към нашия 150k+ ML SubReddit и да се абонирате за нашия бюлетин. Чакайте! В Telegram ли сте? Вече можете да се присъедините към нас и в Telegram.

Имате нужда от партньорство с нас за популяризиране на вашето GitHub хранилище ИЛИ страница в Hugging Face ИЛИ продуктово представяне ИЛИ уебинар и т.н.? Свържете се с нас

Преведено автоматично от английски. Оригиналната статия е на връзката по-долу.