ARGUS: как масштабировать рекомендательные трансформеры

Привет! Меня зовут Кирилл Хрыльченко. Я руковожу командой, которая занимается R&D для рекомендательных технологий в Яндексе. Одна из наших основных задач — развивать трансформерные технологии в контексте рекомендательных систем, и мы активно занимаемся этим уже примерно пять лет. Не так давно у нас произошёл новый виток в развитии рекомендательных технологий, которым мы хотим поделиться с вами в этой статье.

Актуальность рекомендательных систем в мире и для Яндекса обосновать несложно: количество контента растёт очень быстро, всё просматривать самостоятельно невозможно, поэтому для борьбы с информационной перегрузкой нужны рексистемы. Рекомендации музыки, фильмов, книг, товаров, видеороликов, постов, друзей — бо́льшая часть этого есть и у нас в Яндексе. При этом важно не забывать, что эти сервисы помогают не только пользователям, но и создателям контента, которым нужно искать свою аудиторию.

Мы уже внедрили новое поколение рекомендательных трансформеров во множество сервисов — Яндекс Музыку, Алису, Маркет, Лавку — и активно работаем над внедрением в другие. Везде получилось значительно улучшить качество рекомендаций. Если вы ML‑инженер в области рекомендаций — надеюсь, что после этой статьи у вас появятся идеи, как сделать что‑то похожее для вашей рекомендательной системы. А если вы пользователь рекомендаций — то у вас есть возможность побольше узнать о том, как работает та самая рекомендательная система.

Какими бывают рексистемы и как они работают

Сама задача рекомендаций математически формулируется просто: хотим для каждого пользователя подобрать такие айтемы, объекты, документы или товары , которые ему понравятся . Но есть нюансы:

• каталоги, в которых надо искать, очень большие (вплоть до миллиардов объектов);

• пользователей очень много, их интересы постоянно меняются;

• взаимодействия между пользователями и айтемами очень «разреженные»;

• и что такое «пользователю понравился айтем» — тоже не всегда однозначно задаётся.

Поэтому нужно изобретать нетривиальные модели и применять машинное обучение.

Нейросети — очень мощный инструмент машинного обучения. Особенно в сценариях, когда есть много не очень структурированных данных, например текстов или картинок. В то время как традиционное классическое машинное обучение подразумевает экспертные знания в моделируемой области и кучу ручной работы (feature engineering), нейросети позволяют почти автоматически извлекать сложные взаимосвязи и паттерны из сырых данных.

У нас очень много данных, буквально триллионы обезличенных пользовательских взаимодействий, причём существенная их часть как раз неструктурированная — у айтемов есть контент (названия, описания, картинки, аудио), а пользователя можно представить в виде цифровой последовательности событий, каких‑то взаимодействий с сервисом. Нам очень важно, чтобы рексистема хорошо работала на новых айтемах и на «холодных» пользователях, и с этим как раз помогает кодирование пользователей и айтемов через контент.

Время, за которое мы должны формировать для пользователя рекомендации, очень жёстко ограничено. Важна каждая миллисекунда:) И количество ресурсов (железа) у нас тоже не бесконечное, а каталоги, из которых нужно что‑то порекомендовать, очень большие. Поэтому мы делаем рекомендации многостадийно:

• сначала из всего каталога отбираем относительно небольшое множество кандидатов с помощью легковесных моделей (стадия генерации кандидатов, она же retrieval);

• затем обрабатываем этих кандидатов более сложными моделями, в которых используется больше информации и больше тяжёлых вычислений для каждого кандидата (стадия ранжирования, ranking).

Архитектурно модели настолько сильно различаются между стадиями, что сложно что‑либо обсуждать, не ссылаясь при этом на конкретные стадии рексистемы.

Очень популярная нейросетевая архитектура для стадии генерации кандидатов — двухбашенные нейросети. Пользователи и айтемы (на языке информационного поиска это были бы «запросы» и «документы») независимо кодируются в векторные представления, и для подсчёта близости между ними используется скалярное произведение. Ещё говорят, что такие модели «вкладывают» пользователей и айтемы в единое «семантическое пространство». Семантическое оно в том смысле, что близость объектов в этом пространстве является индикатором их реальной похожести.

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

• «башню пользователя» один раз на запрос;

• векторы всех айтемов‑кандидатов, для которых вы хотите оценить интерес пользователя;

• скалярные произведения.

На самом деле даже векторы айтемов‑кандидатов пересчитывать на каждый пользовательский запрос не нужно, так как они совпадают для всех пользователей и очень редко меняются: мы не предполагаем, что, скажем, у фильма или трека часто меняется название. На практике мы регулярно пересчитываем векторы айтемов для всего каталога в офлайне (например, ежедневно) и подгружаем их либо в сам сервис, в котором нужно посчитать скалярное произведение, либо в какой‑то другой сервис, в который можно сходить по сети, чтобы забрать нужные векторы айтемов.

Но это я описал сценарий использования, когда у вас есть какое‑то разумное небольшое количество кандидатов, для которых вы хотите посчитать близости. Такое справедливо для стадии ранжирования. А на стадии генерации кандидатов задача сложнее: нужно посчитать близости для всех айтемов из каталога, отобрать среди них топ‑N (где N обычно в сотнях‑тысячах) с самыми большими значениями близости и затем отправить на следующие стадии.

Здесь двухбашенные модели незаменимы: мы умеем очень быстро формировать примерный топ‑N по скалярному произведению даже среди очень больших каталогов с помощью методов приближенного поиска. Для набора уже посчитанных векторов айтемов строится некий «индекс» (обычно это графовая структура, как, например, в методе HNSW), затем можно держать его в сервисе и «ходить» в него пользовательскими векторами, вытаскивая для них приближённый топ. Процедура построения индекса довольно тяжёлая и медленная (и отдельный челлендж — обновлять и перестраивать индексы быстро), но её можно делать в офлайне, тоже на регулярной основе, и просто подгружать бинарь с индексом в сервис, в который будем ходить в рантайме за кандидатами.

Отдельный интерес представляет следующий вопрос: как закодировать пользователя в вектор? Классические алгоритмы решали его довольно просто: в методах матричной факторизации (типа ALS) вектор пользователя был «обучаемый», являлся параметрами модели и определялся в рамках процедуры оптимизации. В методах user2item коллаборативной фильтрации пользователь задавался вектором размерности каталога, в котором i‑я координата соответствовала какому‑то конкретному айтему и обозначала, как часто пользователь взаимодействовал с этим айтемом (например, как покупал его или какой рейтинг ему ставил).

А вот современная наука говорит, что пользователя можно закодировать в вектор с помощью трансформеров: берём анонимизированную историю пользователя, то есть последовательность каких‑то событий, кодируем эти события в векторы и затем применяем трансформер. В простейшем случае события — это, например, покупки или лайки. А может быть, это вся история взаимодействий с экосистемой какой‑то компании.

Изначально, когда трансформеры только начинали применять в рекомендациях, очень сильно вдохновлялись аналогиями с обработкой естественного языка: пользователь — это как бы предложение, а слова в нём — это покупки, лайки и другие взаимодействия.

Я также должен упомянуть и другой тип нейросетевых рекомендательных моделей — модели с ранним связыванием. В них мы не делим пользователей и айтемы на две башни, а обрабатываем всю информацию вместе. То есть делаем раннее связывание всей информации про пользователя, айтем и их взаимодействие. Про двухбашенные модели как раз говорят, что они делают позднее связывание — через скалярное произведение. Такие модели, с ранним связыванием, более выразительные, чем двухбашенные, и позволяют улавливать более сложный сигнал, выучивать более нетривиальные зависимости. Но при этом их сложно применять вне стадии ранжирования, так как они вычислительно тяжелей и требуют перевычислять целиком всю модель на каждый пользовательский запрос и для каждого кандидата. Нет факторизации вычислений, как в двухбашенных моделях.

У нас в Яндексе используются все возможные типы архитектур: есть и двухбашенные модели с трансформерами, и модели с ранним связыванием. Чаще мы используем всё‑таки двухбашенные архитектуры, потому что они довольно эффективны, подходят сразу для всех стадий и за гораздо более скромные ресурсы всё ещё показывают хорошие приросты качества.

Раньше мы учили двухбашенные модели в две стадии:

1. Предобучение с помощью contrastive learning. Учим модель определять понравившиеся пользователю айтемы с помощью контрастивного обучения: сближаем векторные представления пользователя с понравившимися ему айтемами и отдаляем от других, случайных айтемов.

2. Дообучение под задачу. Дообучение, как и в NLP, зависит от самой задачи. Если модель планируется использовать для ранжирования, мы учим её правильно ранжировать показанные пользователю рекомендации (показали два айтема; один понравился, другой нет; хотим ранжировать в таком же порядке). На генерации кандидатов задача скорее похожа на предобучение, но использует дополнительные трюки, которые растят «полноту» кандидатогенерации.

Но так было раньше. И дальше я расскажу вам про наши новые модели, которые гораздо сильнее, быстрее и умнее.

Масштабирование рексистем

Нас долгое время заботил следующий вопрос: есть ли какой‑то предел размера рекомендательных моделей, после которого от увеличения модели мы уже не видим улучшения качества рекомендаций?

Долгое время наши рекомендательные модели (и не только наши, а во всей индустрии и академии) были очень маленькие, то есть ответ был скорее положительный. А в глубоком обучении есть гипотеза масштабирования, которая гласит, что по мере увеличения моделей и роста количества данных качество моделей должно существенно расти!

Бóльшую часть прогресса в глубоком обучении за последние 10 лет можно обосновать этой гипотезой. Даже самые первые успехи глубокого обучения базировались на масштабировании: появился очень большой датасет для классификации картинок — ImageNet. Затем на нём хорошо себя показали нейросети (AlexNet).

Если посмотреть на языковые модели и обработку естественного языка (NLP), то там гипотеза масштабирования проявляется ещё более ярко: зависимость роста качества от объёма используемых вычислений буквально можно предсказывать, формулировать соответствующие степенны́е законы (scaling law).

Что же я имею в виду, когда говорю, что рекомендательные модели можно сделать больше? Есть целых четыре разных оси для масштабирования:

Эмбеддинги. Мы знаем про пользователей и айтемы очень много различных кусочков информации, а значит, нам доступно очень много различных признаков в модели. И существенная часть из них — категориальные. Самый яркий пример категориального признака — Item ID, но бывают и другие, например ID исполнителя, жанр, язык.

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

При этом эмбеддинги — это «узкое горлышко» между входными данными и моделью, поэтому для хорошего качества их действительно необходимо делать большими. Например, у Меты* размерность матриц эмбеддингов достигает от 675 млрд до 13 трлн параметров, Гугл репортил как минимум про 1 млрд параметров в YoutubeDNN в далёком 2016 году. Даже Пинтерест, который очень долго продвигал индуктивные графовые эмбеддинги из PinSage [1, 2], с недавних пор стал использовать большие матрицы эмбеддингов.

Длина контекста. Десятилетиями инженеры рекомендательных систем занимаются генерацией признаков (feature engineering). В современных ранжирующих системах количество признаков достигает сотен или даже тысяч — и в Яндексе такие сервисы тоже есть.

Другой пример «контекста» в модели — история пользователя в трансформере. Здесь размер контекста определяется длиной истории. И в индустрии, и в академии она, как правило, очень маленькая — в лучшем случае несколько сотен событий.

Размер обучающей выборки. Как уже обсуждалось выше, у нас очень много данных. Рексистемы ежедневно производят сотни датасетов для обучения GPT-3. В индустрии легко можно найти примеры использования очень больших датасетов с миллиардами обучающих примеров: 2 млрд, 2,1 млрд, 3 млрд, 60 млрд, 100 млрд, 146 млрд, 500 млрд.

Размер энкодера. Для моделей с ранним связыванием стандартные размеры будут в миллионах или десятках миллионов параметров. Гугл в статьях использовал для экспериментов «упрощённые» версии своих продовых Wide&Deep‑моделей с количеством параметров от 1 до 68 млн [1, 2]. А если применять над тысячей вещественных признаков двухслойный DCN‑v2 (это такой популярный нейросетевой слой для моделей с ранним связыванием), то получим не больше 10 млн параметров.

В двухбашенных моделях для кодирования пользователя чаще всего используют очень маленькие трансформеры: например, два трансформерных блока с размерностью скрытого слоя, не превышающей пары сотен. В такой конфигурации будет от силы несколько миллионов параметров.

И если размеры матриц эмбеддингов и обучающих датасетов и так уже довольно большие, то масштабирование длины истории пользователя и ёмкости энкодерной части модели — это открытый вопрос. Есть ли какое‑то существенное масштабирование по этим параметрам или нет.

Такой была ситуация на момент февраля 2024 года. А потом вышла статья от инженеров из Меты* «Actions speak louder than words», которая немного нас всех взбодрила. Инженеры представили новую архитектуру энкодера под названием HSTU, а также сформулировали и задачу ранжирования, и задачу генерации кандидатов в виде генеративной модели. В модели была очень большая длина истории (8 тысяч событий!) одновременно с очень большим датасетом для обучения (100 млрд примеров). И энкодер истории пользователя был гораздо больше прежних нескольких миллионов параметров. Но даже здесь самая большая конфигурация энкодера, которую упоминают авторы, — всего лишь 176 млн параметров. И непонятно, внедрили ли они её (судя по следующим статьям, нет).

176 млн параметров в энкодере — это много или мало? Если посмотреть на языковые модели, то кажется, что ответ однозначный: LLM с 176 млн параметров в энкодере будет очень сильно уступать по способностям и качеству решения задач современным SOTA‑моделям, насчитывающим миллиарды или даже триллионы параметров.

Почему же тогда у нас в рекомендательных системах такие маленькие модели? Почему мы не можем достичь похожего скачка в качестве, если заменяем тексты естественного языка обезличенными историями пользователей, в которых словами выступают действия? Неужели рекомендательные модели уже достигли потолка своего базового качества и нам остаётся только делать небольшие инкрементальные улучшения, перебирать признаки и таргеты?

Именно этими экзистенциальными вопросами мы и задались, когда придумывали наш собственный новый подход ARGUS.

RecSys × LLM × RL

После штудирования многочисленной литературы по скейлингу мы пришли к выводу, что есть три главных условия успеха нейросетевого масштабирования:

• у вас должно быть очень много данных;

• достаточно выразительная архитектура с большой ёмкостью модели;

• и максимальная общая, фундаментальная задача обучения.

Например, LLM — это очень выразительные и мощные трансформеры, которые учатся буквально на всех данных из интернета. А ещё и задача предсказания следующего слова — это суперфундаментальная задача, которая в реальности декомпозируется на мириаду очень разных задач, связанных с чем угодно: грамматикой, эрудицией, математикой, физикой, программированием. Все три условия соблюдены!

Если посмотреть на рекомендательные системы, то:

• у нас тоже очень много данных — буквально триллионы взаимодействий между пользователями и айтемами;

• мы точно так же можем использовать трансформеры;

• и нужно всего лишь подобрать правильную задачу обучения, при которой рекомендательные модели будут масштабироваться.

Этим мы и занялись.

В предобучении больших языковых моделей есть следующая интересная интуиция. Если просто взять и спросить о чём‑то предобученную LLM, она выдаст среднестатистический ответ — самый вероятный ответ, который она встречала в обучающих данных. И этот ответ необязательно будет хорошим или правильным. Но если добавить перед вопросом промт вида «Представь, что ты хороший специалист в области X», то она начнёт выдавать гораздо более хорошие и правильные ответы!

Всё потому, что LLM не просто учится имитировать ответы из интернета, но ещё и приобретает более фундаментальные представления о мире в попытках сжать всю информацию из обучения. Учит паттерны, абстракции. И как раз за счёт того, что LLM знает всевозможные ответы и при этом имеет какое‑то фундаментальное представление о мире, мы и можем «выуживать» из неё хорошие ответы.

Мы попробовали переложить эту логику на рекомендательные системы. Для начала нужно сформулировать рекомендации в виде задачи обучения с подкреплением:

• Рексистема — это агент.

• Действия — это рекомендации. В простейшем случае рексистема рекомендует по одному айтему (например, каждый раз рекомендует один трек в Моей волне).

• Окружение — это пользователи, их поведение, паттерны, предпочтения, интересы.

• Политика — это вероятностное распределение по айтемам.

• Награда — положительный фидбек от пользователя в ответ на рекомендацию. А может, и что‑то более разреженное и отложенное, например подписка через месяц или покупка, совершённая через пару недель.

Есть прямая аналогия с примером про LLM. «Ответы из интернета» — это действия прошлых рексистем (логирующей политики), а фундаментальные знания о мире — это понимание пользователей, их паттернов и предпочтений. Мы хотим, чтобы наша новая модель:

• умела имитировать действия прошлых рексистем;

• при этом имела хорошее понимание пользователей;

• смогла скорректировать их действия в более хорошую сторону.

Прежде чем мы перейдём к нашему новому подходу, давайте посмотрим на самую популярную постановку для обучения рекомендательных трансформеров — next item prediction. Яркий представитель — модель SASRec. Для пользователя формируется история его положительных взаимодействий с сервисом (например, покупок), и модель учится предсказывать, какая покупка идёт дальше в последовательности. То есть вместо next token prediction, как было в NLP, делаем next item prediction.

Такой подход (SASRec и обычный next item prediction) не очень соответствует той философии, которую я описал ранее: про коррекцию логирующей политики за счёт фундаментальных знаний о мире. Казалось бы, чтобы предсказать, что́ пользователь дальше купит, модель должна как раз соответствовать этой философии:

• Она должна понимать, что́ потенциально могла показать пользователю та рексистема, которая работала в продакшне на момент времени, для которого нужно сделать предсказание. То есть она должна иметь хорошую модель поведения логирующей политики (иметь модель = уметь имитировать).

• Она должна понимать, что́ могло понравиться пользователю из показанного прошлой рексистемой. То есть понять его предпочтения (это и есть те самые фундаментальные представления о мире).

Но модели типа SASRec ни то ни другое в явном виде не моделируют. У них нет полной информации о логирующей политике в истории (видим только рекомендации с положительным исходом), и имитировать логирующую политику мы тоже не учимся. Выучить, что́ потенциально могла предложить прошлая рексистема, не получится. И при этом модель мира, то есть пользователя, мы тоже полноценно не учим — игнорируем весь отрицательный фидбек, смотрим только на положительный.

ARGUS: AutoRegressive Generative User Sequential modeling

AutoRegressive Generative User Sequential modeling (ARGUS) — наш новый подход к обучению рекомендательных трансформеров.

Во‑первых, мы смотрим целиком на всю анонимизированную историю пользователя — не только на положительные взаимодействия, но и на все остальные. А ещё добавляем сущность контекста взаимодействия: в какое время оно происходило, на каком устройстве был пользователь, на какой страничке продукта, какие были настройки Моей волны и т. д.

История пользователя — это некая последовательность троек вида (context, item, feedback), где context — это контекст взаимодействия, item — объект, с которым взаимодействует пользователь, а feedback — реакция пользователя на взаимодействие (понравился ли пользователю айтем, купил ли он его и т. д.).

Во‑вторых, мы определяем две новых задачи обучения, обе из которых выходят за рамки традиционного next item prediction, широко используемого в индустрии и академии.

Next Item Prediction

Наша первая задача тоже называется Next Item Prediction. Глядя на историю и текущий контекст взаимодействия, предсказываем, с каким айтемом будет взаимодействие: P(item | history, context).

• Если в истории есть только рекомендательный трафик (события, которые порождаются непосредственно рексистемой), то модель учится имитировать логирующую политику (рекомендации прошлой рексистемы).

• Если есть также органический трафик (любой другой трафик, кроме рекомендательного: например, из поиска, или же пользователь сам зашёл в библиотеку и послушал свой любимый трек), то мы также приобретаем и более фундаментальные знания про пользователя, не связанные с логирующей политикой.

Важно: хотя она и называется так же, как у SASRec (next item prediction), это совсем не та же самая задача. Мы предсказываем не только положительные взаимодействия, но и отрицательные. А ещё учитываем текущий контекст. Контекст помогает понять, органическое сейчас действие или нет, а если это рекомендация — то на какой «поверхности» (в каком месте, на какой странице или карусели). Да и в он целом снижает уровень шума при обучении модели.

Для музыкальных рекомендаций контекст очень важен: настроение пользователя, его текущая ситуация сильно влияет на то, какую музыку он хочет послушать.

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

И если в LLM размеры словарей в худшем случае достигают сотен тысяч объектов и подсчёт софтмакса не вызывает особенных трудностей, то в рекомендательных системах это очень большая проблема. Здесь каталоги состоят из миллионов, а где‑то даже миллиардов айтемов, и посчитать полный софтмакс невозможно. Эта тема для отдельной большой статьи, но в конечном счёте приходится использовать хитрую функцию потерь под названием «семплированный софтмакс» с logQ‑коррекцией:

• 

• — смесь in‑batch и равномерных негативов,

• — logQ‑коррекция,

• Температура — обучаемый параметр, клипается до [0.01, 100].

Feedback Prediction

Вторая задача обучения — Feedback Prediction. Глядя на историю, текущий контекст и айтем, предсказываем фидбек пользователя: P(feedback | history, context, item).

Если первая задача, next item prediction, учит нас имитировать логирующую политику (и немного учит нас понимать пользователей, если есть органический трафик), то эта задача, наоборот, направлена исключительно на фундаментальные знания про пользователей, их предпочтения и интересы. Она очень похожа на то, как в статье «Action Speak Louder than Words» учится ранжирующий вариант модели на последовательности из двоек (item, action), но здесь отдельно рассматривается токен контекста, а также присутствуют не только рекомендательные контексты.

На практике фидбек может быть многокомпонентный: был ли лайк, дизлайк, добавление в плейлист, какая доля трека была прослушана. Мы предсказываем все типы фидбека, декомпозируя их на отдельные функции потерь. В качестве конкретной функции потерь можно использовать что угодно, включая кросс‑энтропию, регрессию и т. д. Например, для предсказания, был лайк или нет, достаточно использовать бинарную кросс‑энтропию.

Несмотря на то что какой‑то фидбек встречается гораздо чаще другого (как правило, лайков гораздо меньше, чем долгих прослушиваний), модель хорошо учится предсказывать все сигналы. Чем больше модель, тем проще ей учиться сразу на все задачи без конфликтов. Более того, частый фидбек (прослушивания), наоборот, помогает модели научиться моделировать редкий, разреженный фидбек (лайк). Knowledge transfer!

Если совместить это всё в единую задачу обучения, то получается следующее:

• формируем для пользователя истории из троек (context, item, feedback);

• применяем трансформер;

• из скрытого состояния контекста предсказываем, что за item идёт дальше;

• из скрытого состояния айтема предсказываем, какой будет фидбек от пользователя после взаимодействия с этим айтемом.

Сразу же прокомментирую и отличие от HSTU. В статье «Actions speak louder…» авторы учат две отдельные модели для генерации кандидатов и ранжирования. Кандген‑модель содержит всю историю, но моделирует, как и SASRec, только положительные взаимодействия — там, где дальше идёт отрицательное взаимодействие, лосс в модели не считается. А ранжирующая модель, как я уже писал ранее, учится на задачу, похожую на наш feedback prediction.

То есть в нашем решении более полная задача next item prediction, более полная задача feedback prediction, и к тому же учатся они одновременно!

Simplified ARGUS

Одна большая проблема нашего подхода — мы сильно раздуваем историю пользователя. Так как каждое взаимодействие с айтемом представляется сразу тремя токенами (context, item, feedback), чтобы проанализировать 8192 последних прослушиваний пользователя, нам придётся подавать в трансформер почти 25 тысяч токенов.

Здесь можно возразить, что это всё равно не так уж и много и что в LLM на самом деле длина контекста гораздо больше, — но это неправда. Особенно при предобучении у LLM контексты в среднем гораздо меньше, в сотнях токенов. А, например, в Яндекс Музыке у пользователей очень часто набираются тысячи или даже десятки тысяч событий. Поэтому у нас длины контекстов и так гораздо больше, а раздувание этих длин в три раза ещё сильнее влияет на скорость обучения.

Поэтому мы сделали упрощённую версию модели, в которой каждая тройка (context, item, feedback) сворачивается в единый вектор. С точки зрения формата входных данных это больше похоже на то, как выглядели наши прошлые поколения трансформерных моделей. Но при этом мы оставляем всё те же две задачи обучения — next item prediction и feedback prediction.

Чтобы предсказать следующий айтем, мы берём скрытое состояние из трансформера, соответствующее тройке (c, i, f) на прошлый момент времени, конкатенируем к нему вектор текущего контекста, сжимаем с помощью MLP в меньшую размерность и затем с помощью семплированного софтмакса учимся предсказывать следующий айтем. Для предсказания фидбека мы конкатенируем туда же вектор самого́ текущего айтема и затем с помощью MLP предсказываем все нужные целевые переменные. На языке рекомендательных трансформерных архитектур наша модель становится менее target‑aware и менее context‑aware, но тем не менее это всё ещё хорошо работает и позволяет нам ускориться в три раза.

Внедрение ARGUS

Модель, обученную в таком двухголовом режиме сразу на обе задачи (next item prediction и feedback prediction), можно внедрить as is. NIP‑голова позволяет делать отбор кандидатов, а FP‑голова — делать финальное ранжирование.

Но как минимум для первого внедрения мы так делать не захотели:

• У нас была цель внедрить очень большую модель, и поэтому мы изначально целились в офлайн‑внедрение. При внедрении в офлайне отдельно регулярным процессом ежедневно пересчитываются векторы пользователей, айтемов, а в рантайме нужно посчитать только скалярное произведение. Предобученная версия ARGUS подразумевает доступ к истории пользователя без какой‑либо задержки: мы видим все события в истории пользователя вплоть до текущего момента времени, в который делается предсказание. То есть её нужно применять в рантайме.

• NIP‑голова предсказывает все взаимодействия пользователя, а обычно для генерации кандидатов модель учат предсказывать только будущие положительные взаимодействия. Предсказание положительных взаимодействий — это на самом деле эвристика. Суррогатная задача обучения. Возможно, использовать голову, предсказывающую все взаимодействия, будет даже правильней, потому что она учится на согласованность с ранжированием (если объект пророс до рекомендации, значит, ранжированию он понравился). Но в данной ситуации мы не были готовы с этим экспериментировать и хотели скорее пойти по проторённой дорожке.

• FP‑голова учится на поточечные лоссы (pointwise): будет лайк или нет, какая доля трека будет прослушана и т. д. А для ранжирования мы всё‑таки чаще учим модели на попарное ранжирование — учимся ранжировать объекты, которые были рекомендованы «рядом» и получили разный фидбек. Есть аргументы в пользу того, что pointwise‑лоссов для обучения ранжирующих моделей должно быть достаточно. Но в данной ситуации мы не заменяем целиком ранжирующий стек, а хотим добавить в финальную ранжирующую модель новый мощный признак на основе нейросети. Если финальная ранжирующая модель учится на определённую задачу (попарное ранжирование), то нейросеть, генерирующую признак, эффективней всего учить на эту задачу. Иначе финальная модель будет меньше опираться на наш признак. Соответственно, для применения в ранжировании мы бы хотели дообучить ARGUS на ту же задачу, на которую учится исходная ранжирующая модель.

На самом деле есть и другие сценарии внедрения, не связанные с традиционными стадиями генерации кандидатов и ранжирования, и их мы тоже активно исследуем. Но в качестве первого внедрения мы пошли по пути офлайнового двухбашенного ранжирования:

• Решили дообучить ARGUS так, чтобы его можно было использовать как офлайновую двухбашенную модель. С её помощью ежедневно пересчитываются векторы пользователей и айтемов, а предпочтения пользователя определяются через скалярное произведение пользователя с айтемами.

• При этом мы дообучили его на задачу попарного ранжирования, похожую на ту, на которой учится финальный ранкер. Это значит, что есть каким‑то образом отобранные пары треков, которые пользователь слышал и при этом оставил разный по степени положительности фидбек, и мы хотим научиться их правильно ранжировать.

Такие модели мы делаем довольно часто: их просто внедрять с точки зрения ресурсов и затрат на разработку, а также их просто обучать. Но наши прошлые модели все были гораздо меньше по размеру и учились по‑другому — не с помощью ARGUS‑процедуры, а сначала на обычный contrastive learning между пользователями и позитивами, а затем дообучались под задачу.

Наша прошлая процедура предобучения (contrastive pre‑training) подразумевала, что для пользователя составляется много обучающих примеров — если у него было n покупок, значит, в датасете будет n семплов. При этом мы не учились авторегрессивно — то есть при обучении буквально n раз прогонялся трансформер. Такой подход позволял нам очень гибко составлять пары (пользователь, айтем) для обучения, использовать любой формат истории, кодировать вместе с пользователем контекст, учитывать лаг (при предсказании лайка применять для пользователя историю с лагом в один день). Однако всё работало довольно медленно.

ARGUS‑предобучение — авторегрессивное, то есть за один прогон трансформера мы учимся сразу на всех событиях в жизни пользователя. Это очень сильное ускорение, которое позволило нам за те же ресурсы обучать модели гораздо большего размера.

При дообучении мы тоже раньше много раз прогоняли трансформер для одного пользователя. То, что называется impression‑level‑обучение, как было у Меты* раньше, до HSTU. Если пользователю в определённый момент показали айтем, формируем семпл вида (пользователь, айтем). Таких показов для одного пользователя в датасете может быть очень много, и для каждого мы будем повторно прогонять трансформер. Для попарного ранжирования рассматривались тройки вида (пользователь, айтем1, айтем2). Так мы раньше и делали.

Глядя на то, как ускорилась стадия предобучения, мы решили провернуть похожий трюк и с дообучением. Сделали такую процедуру дообучения двухбашенной модели на ранжирование, при которой трансформер надо применять только один раз.

Пусть у нас есть вся история пользователя за год, а также все показанные пользователю рекомендации за этот же период. Применив трансформер с каузальной маской над всей историей, мы получаем векторные представления пользователя сразу на все моменты времени, которые были в этот год:

• можно отдельно посчитать векторы показанных айтемов,

• глядя на таймстемпы, сопоставить показы рекомендаций с векторами пользователя, соответствующими нужной задержке доставки истории пользователя,

• посчитать все нужные скалярные произведения и все члены функции потерь.

И все это — сразу за весь этот год — за одно применение трансформера.

Раньше на каждую пару показов мы бы применяли трансформер повторно, а теперь за одно применение обрабатываем сразу все показы! Это очень большое ускорение: в десятки, сотни или даже тысячи раз. Чтобы потом применять такую двухбашенную модель, достаточно взять для пользователя векторное представление на последний момент времени (соответствующее последнему событию в истории) в качестве текущего векторного представления, а для айтемов применять тот энкодер, который применялся над показами при обучении. При обучении мы эмулируем задержку истории пользователя в один день и затем применяем модель как офлайновую, пересчитывая векторы пользователей раз в день.

Когда я пишу, что мы за одно применение трансформера обрабатываем всю историю пользователя за год, я немного лукавлю. В реальности есть определённое ограничение на максимальную длину истории, которое мы соблюдаем, и пользователь в датасете может иметь несколько семплов, чанков. Для предобучения эти чанки не пересекаются, а на дообучении есть ограничение не только на максимальную длину истории, но и на минимальную длину, а также максимальное количество показов рекомендаций в одном обучающем примере, на котором учится ранжирование.

Результаты

В качестве первого сервиса для экспериментов мы выбрали Яндекс Музыку. Здесь очень важны рекомендации и есть много активных пользователей. Мы собрали гигантский обучающий датасет, в котором больше 300 млрд прослушиваний от миллионов пользователей. Это в десятки или даже сотни раз больше, чем обучающие датасеты, которые мы использовали раньше.

Что в Музыке является тройкой (context, item, feedback)?

• Контекст — является ли текущее взаимодействие рекомендательным или органическим. Если рекомендательным — то на какой «поверхности» оно происходит, а если это Моя волна — какие у неё настройки.

• Айтем — это музыкальный трек. Самый важный признак для кодирования айтема — это item ID. Мы используем технику unified embeddings для кодирования признаков с высокой кардинальностью, в данном случае для айтема берётся три хэша размерности 512к. Матрица unified‑эмбеддингов в наших экспериментах зафиксирована, и у неё 130 млн параметров.

• Фидбек пользователя — был ли лайк и какая доля трека была прослушана.

Для оценки качества в офлайне мы используем следующую неделю после обучающих данных, то есть применяем global temporal split.

Чтобы оценить качество предобученной модели, мы смотрим на значение функции потерь на самих задачах предобучения — next item prediction и feedback prediction. То есть насколько хорошо модель научилась решать те задачи, которые мы перед ней поставили. Чем значение меньше, тем лучше.

Важно: история пользователя берётся за большой промежуток времени, но лосс считается только по тем событиям, которые попали в тестовый период.

Так как на дообучении мы учимся правильно ранжировать пары айтемов исходя из фидбека, оставленного пользователем, для нас хорошая офлайн‑метрика — доля правильно упорядоченных моделью пар: PairAccuracy. На практике мы ещё чуть‑чуть перевзвешиваем пары исходя из фидбека: например, пары, в которых произошёл лайк и скип, имеют вес выше, чем те, в которых было прослушивание и скип.

Наш сценарий внедрения — это добавить новый мощный признак в финальный ранкер. Поэтому мы замеряем относительный прирост PairAccuracy для финального ранкера с добавлением нового признака в сравнении с финальным ранкером без него. Финальный ранкер в Яндекс Музыке — это градиентный бустинг.

Результаты A/B-тестов и замеры

Наша изначальная цель — масштабировать рекомендательные трансформеры. Для проверки масштабирования мы выбрали четыре разных по размеру конфигурации трансформера: от 3,2 млн до 1,007 млрд параметров.

Также мы решили проверить эффективность архитектуры HSTU. Напомню, что авторы «Actions speak louder…» предложили свою новую архитектуру энкодера, которая довольно сильно отличается от трансформера. Исходя из экспериментов авторов, эта архитектура на рекомендательных задачах показывает себя гораздо лучше, чем трансформеры.

Cкейлинг есть! Получилось, что каждый новый скачок в размере архитектуры приводит к приросту качества. Как на предобучении, так и на дообучении.

HSTU показал себя не лучше трансформера. Мы взяли самую большую конфигурацию, которую авторы упоминают в «Actions speak louder…». У неё в полтора раза больше параметров, чем у нашей Medium‑конфигурации трансформера, а качество примерно такое же.

Если визуализировать метрики из таблички в виде графика, то виден scaling law по крайней мере по нашим четырём точкам — зависимость качества от логарифма количества параметров выглядит примерно линейной

А ещё мы провели небольшое ablation study, чтобы понять, можно ли как‑то упростить нашу модель или что‑то убрать из процедуры обучения.

Если убрать предобучение — качество сильно падает.

Если уменьшить длительность дообучения — качество падает ещё больше.

В начале этой статьи обсуждали, что в «Actions speak louder…» авторы обучили модель с длиной истории 8 тыс. айтемов. Мы тоже решили попробовать: получилось, что от такой глубокой музыкальной истории пользователя есть заметное улучшение рекомендаций. Раньше наши модели использовали максимум полторы‑две тысячи событий, и так было в любом домене, не только в Музыке. Здесь у нас впервые получилось преодолеть этот порог.

Итоги внедрения

Вместе с командой Яндекс Музыки мы занимаемся разработкой трансформеров для музыкальных рекомендаций уже примерно три года. Мы прошли большой путь:

• Первые три трансформера были офлайновыми. Векторы пользователей и айтемов пересчитывались ежедневно, затем векторы пользователей загружались в key‑value‑хранилище, а векторы айтемов помещались в оперативную память сервиса — в рантайме считалось только скалярное произведение. Некоторые из этих моделей использовались не только для ранжирования, но и для генерации кандидатов (умеем делать многоголовые модели, которые делают и то и другое). Тогда в оперативной памяти сервиса также лежит HNSW‑индекс, из которого можно забрать кандидатов.

• В первой модели был только сигнал про лайки, во второй — про прослушивания (включая скипы), а в третьей мы совместили оба типа сигнала (и explicit, и implicit).

• v4-версия модели‑ адаптация v3, которая применяется в рантайме с небольшой задержкой истории пользователя. При этом энкодер у неё в 6 раз меньше, чем был у v3-модели.

• В новой модели ARGUS длина истории пользователя в 8 раз больше, а энкодер — в 10. И учится она по новой схеме, которую я описал в этой статье.

TLT — это суммарное время прослушивания, а like likelihood — вероятность лайка при включении рекомендации. Каждое внедрение дало хороший прирост для метрик Моей волны. А первый ARGUS дал примерно такой же прирост метрик, как все прошлые внедрения суммарно!

А ещё в Моей волне есть особая настройка, для которой работает отдельный стек ранжирования: «Незнакомое». В неё ARGUS внедрялся отдельно. Для этой настройки ARGUS на 12% увеличил TLT и на 10% — вероятность лайка. Настройкой «Незнакомое» пользуются люди, которым важна дискаверийность рекомендаций. Тот факт, что в ней получился большой прирост, подтверждает, что ARGUS лучше обрабатывает нетривиальные сценарии.

Также мы поработали с командой медиарекомендаций, которая занимается музыкальными рекомендациями в умных колонках. Мы внедрили ARGUS в музыкальные сценарии в Станциях — там получилось увеличить на 0,75% суммарное время, которое пользователи проводят с активной колонкой. При этом здесь в качестве финального ранкера выступает не градиентный бустинг, а полноценная ранжирующая нейросеть; её архитектура похожа на ту, про которую я рассказывал в Яндексе на ML Party. Поэтому мы смогли не просто подать один признак‑скаляр из ARGUS, а передать полноценные векторы пользователя и айтема на вход финальному ранкеру. По сравнению с одним признаком‑скаляром это позволило увеличить прирост качества ещё в полтора‑два раза.

ARGUS уже успешно внедрился в Маркет (не только как ранжирующий признак, но и как генератор кандидатов) и Лавку. Команда Яндекс Музыки совместно с рекомендательной службой Портала адаптировала офлайн‑ARGUS в рантайм. Все эти внедрения дали хорошие приросты ключевых метрик, про которые мы расскажем вам в другой раз.

Так или иначе будущее рекомендательных систем — за нейросетями. Предстоит ещё очень большое путешествие: рекомендательные нейросети будут становиться всё больше, умнее, полезней.

Спасибо, что дочитали, буду рад ответить на вопросы в комментариях!

* Компания Meta признана экстремистской организацией, а её продукты Facebook и Instagram запрещены на территории РФ