Можно сделать самую лучшую на свете модель, но от нее будет мало проку, если не обеспечить ее интеграцию в реальные бизнес-процессы. Всем привет, я Илья Бадекин — Data Scientist в команде товарных рекомендаций Wildberries и в данной статье расскажу о том, зачем текстовый энкодер в команде «Товарных рекомендаций» Wildberries, на что он способен и как мы сжимали его эмбеддинги для онлайн-доранжирование рекламных баннеров по запросам пользователей.
Зачем текстовый энкодер в «Похожих товарах»
Математика, статистика и данные (в нашем случае история взаимодействий) — основа работы любого прикладного инструмента. Рекомендательные системы не исключение. Если условные Петя и Вася купили себе мороженное, коктейль и стейк и у нас есть Андрей, который тоже купил мороженное и коктейль, то вполне релевантно ему тоже предложить стейк (коллаборативная фильтрация для самых маленьких).

А если истории нет? Ну на нет и суда нет?) Увы, суд всё-таки есть и в реальных системах это базовый минимум, а не роскошный максимум.
Мы работаем с большим количеством холодных товаров (товары без богатой истории взаимодействий) и более того, этот каталог регулярно пополняется. Кроме того, такие товары могут быть рекламными, а значит — их нужно уметь хорошо продвигать.
Окей, если представить товар через историю взаимодействий невозможно, то давайте спустимся на уровень ниже и представим его через контент (текст, звук, видео и другие фичи). В нашем случае — через текст.
В чём плюсы контентной модели? Во-первых, независимость от наличия истории взаимодействий с конкретным товаром. Во-вторых, относительная интерпретируемость по сочетанию цвета, бренда, материала и других характеристик (хотя это тоже все еще черепикинг).
Недостатки тоже есть. Нам всё ещё нужен контент и контент качественный (модель не волшебник). А ещё в модель необходимо прокидывать информацию о домене (дообучать на данных из нужного корпуса), чтобы не получать тривиальные и посредственные рекомендации.
Итоговый профит: мы учимся на подмножестве сильно меньшего размера, но получаем рекомендации на всём наборе.
Наш текстовый энкодер
За основу мы взяли Multilingual-E5-large — мультиязычный текстовый трансформер. Модель все еще показывает хорошие результаты на большинстве бенчмарков и лучший (на тот момент) результат на данных Wildberries без дообучения. 24 слоя, размерность эмбеддинга 1024, контекст 512.

E5 — тюн BERT для задачи поиска в две стадии. На первом этапе модели обучались примерно на миллиарде пар текстов на различных языках, используя контрастивный метод с функцией потерь InfoNCE-loss и батчем 32 тыс. Суть InfoNCE-loss — сделать соответствующие пары ближе, чем любые другие.
На втором этапе модели дообучались на комбинации высококачественных размеченных датасетов (MSMARCO, NQ, TriviaQA, SQuAD, NLI и другие) с общим объёмом около 1,6 млн примеров. Данных меньше, но их качество выше.

Обучение для «Похожих» строилось по тому же принципу: контрастивное дообучение с in-batch негативами на нашем датасете «похожих товаров».
Какой текст идёт на вход?
расцветка товара,
заголовок товара,
пометка 18+ или нет,
название бренда,
все характеристики товара в формате по строкам — «название характеристики: значение».
Пример текстового описания товара:
passage: кремовый, бежевый, темно-синий Блузка "AAA" <sep>бренда: brandname<sep>
Страна производства: Россия
Коллекция: Базовая коллекция
Комплектация: Блузка — 1 шт
Фактура материала: гладкая
Вырез горловины: округлый
Мы пробовали докидывать тексты: в офлайне показатели растут, в проде — эффект не стат. значимый.
Как устроен инференс модели?
Перегоняем предобученный чекпоинт во что-то быстрое на ваш вкус — например, BetterTransformer, TensorRT, Тритон и так далее. Векторизуем каталог во float16 — он кратно быстрее благодаря оптимизации NVIDIA и занимает в два раза меньше места без потери качества. Закидываем в Faiss-GPU (или любой другой ANN индекс) — рекомендации готовы.
Пример выдачи

Первый товар — это якорный товар, к которому мы подбираем рекомендации. В верхнем ряду выдача промежуточной E5, которую частично доучили на наших данных. Результат средний: сборники ОГЭ и ЕГЭ похожи на исходный сборник ВПР, но всё-таки девятый и одиннадцатый класс не имеют отношения к тетради для восьмого класса. В нижнем ряду финальная продовая E5, и здесь рекомендации гораздо точнее.
Как мы докатились до жизни такой…
До этого эмбеддинги нашей модели использовались по сути оффлайн и все в нашей жизни было хорошо, но коллеги подкинули нам интересную задачу.
Когда вы вводите запрос в строке поиска Wildberries, то получаете релевантную выдачу и несколько рекламных баннеров «в подарок».
За баннеры уплачены деньги и хорошо бы показывать их контекстно уместно: чтоб всех радовать и никого не огорчать. Вот, например, результат по запросу «картина по номерам космос» старого варианта подбора:

Не претендую на истину в последней инстанции, но связь между картиной по номерам, кухонной техникой, моющим средство и постельным бельем кажется не тривиальной. Баннер же в сути своей некоторая картинка, ее текст и товары связанные с этим баннером. Поэтому платформа продвижения попросила нас о помощи.
Баннеры сущность изменчивая: у них плавающее время жизни, они могут содержать в себе довольно разные товары, в том числе холодные, и сам состав баннеров может меняться (добавили новый товар, убрали старый). Следовательно необходимо уметь оперативно поддерживать актуальность представления и эмбеддинги в размерности 1024 — слишком дорогое удовольствие, если учесть требуемый уровень нагрузки.
Как принято уменьшать размерность в ML?
На ум сразу может прийти набор вполне классических методов, таких как: t-SNE, PCA, TruncatedSVD, UMAP и так далее.
Но размерность эмбеддингов и биг дата накладывает органичения с которыми нужно как-то мириться. Всеми любимый t-SNE не поддерживает размерность выходных векторов больше четырёх, поэтому его мы не рассматривали. MDS и Isomap аппроксимируют матрицу расстояния в пространстве меньшей размерности, но вычислительно это дорогое удовольствие при нашем объёме данных.
С учетом ограничений из классического ML были взяты на апробацию следующие методы: PCA, TruncatedSVD, GaussianRandomProjection и UMAP.
PCA. Очень известный метод. Faiss PCA — то же самое, только с упором на скорость и масштабируемость для обработки высокоразмерных и объёмных матриц на GPU и CPU. Наша задача — найти подпространство, которое максимально сохранит дисперсию исходных данных. Для этого мы строим матрицу ковариации, считаем собственные векторы и числа и строим оператор отображения. Итог: сохраняется 64% от исходных метрик. Хороший baseline!

TruncatedSVD. Если взять набор векторов и объединить их в матрицу, получится линейный оператор, который отображает векторы из одного пространства в векторы другого пространства. Линейная алгебра говорит, что этот оператор можно разложить (декомпозировать) на три отдельных оператора, где первый поворачивает, второй сжимает, третий снова поворачивает. Собственно такое разложение и называется Singular Value Decomposition или SVD.

В свою очередь TruncatedSVD считает усечённое/приближённое SVD вычисляя только первые (наиболее значимые) сингулярные векторы матрицы. Тестировали реализацию от cuML с GPU-ускорением. Итог: 68% исходных метрик.
GaussianRandomProjection и UMAP показали только 53% и 24% от исходных метрик соответственно, так что не вижу особого смысла заострять на них внимание.
А теперь добавим чуть-чуть deep и вспомним про автоэнкодеры. Автоэнкодер — полносвязная нейронная сеть, которая состоит из энкодера и декодера. Идея следующая: сжимаем исходный эмбеддинг → разжимаем его → сравниваем разжатый эмбеддинг с исходным → считаем loss → отрезаем декодер → получаем модель для сжатия эмбеддингов.
Автоэнкодер штука крайне кастомизируемая: можно варьировать функции активации, слои, loss и даже, если очень хочется, добавить квантизацию (Rate-Adaptive Quantized Variational Autoencoder). Кроме того, модель довольно легко интегрировать в пайплайн. Из грустного, модель крайне склонна к переобучению, но это болезнь всех полносвязных нейронок. Итог: 75% от исходных метрик.

По результатам первой итерации победили автоэнкодеры.
Можно ли лучше?
К началу второй итерации дела обстояли вот так:

У нас есть текстовый энкодер, мы даем ему текст, она выдает нам вектор размерности 1024, где фичи как-то раскиданы по этим 1024 измерениям.
А дальше классика последних лет: «не учи много моделей — обучи один хороший трансфор».
Вот мы тоже так подумали и решили попробовать прикрутить «Матрешку» к нашей E5.

MRL (Matryoshka Representation Learning) — это метод обучения представлений, который позволяет создавать гибкие эмбеддинги разной размерности в рамках одной модели, подобно вложенным матрёшкам.

Основной эмбеддинг — это самая большая и детализированная матрёшка. Мы обрезаем её на меньшую размерность (512, 256 и т. д.) и получаем меньше деталей, сохраняя суть. Еще можно сравнить данный подход с рядами Тейлора, где мы регулируем точность разложения с помощью количества членов.
Идея классная, как реализовать? Идейно довольно просто.
При подсчете обычного loss-а, ошибка по векторам распределяется условно случайно (как-то как нашей моделе удобно в рамках оптимизационной задачи), так как мы не накладываем никаких ограничений на порядок. Но если посчитать loss не только от оригинальной размерности, но и от редуцированных, то получим зоны с кратным штрафом, что автоматически заставит модель переносить в начало наиболее важные признаки для минимизации ошибки.
Думаю будет чуть понятнее, если посмотреть на картинку: в серой зоне X2, в жёлтой зоне будет копиться loss с ошибкой X2, в синей — X3, и далее подобным образом.

Окей, но получиться дотюнить модель или нужно учить с нуля? Хватит ли адаптера или LoRA? Не просядет ли качество на основной размерности?
Ожидание и реальность
Если верить авторам статьи, то с нуля все должно быть все хорошо.
Вот график ResNet50 на ImageNet-1K:

FF — модель с изначально фиксированной размерностью, MRL — матрёшка. Видно, что MRL лучше во всех размерностях, а к основной размерности 2048 FF и MRL почти выравниваются.
Теперь то же самое, только для MRL и с распределением по классам:

А что у нас? Попробовали ванильно прикрутить MRL на четыре размерности (1024, 768, 512, 256) и по ошибке на обучении кажется все хорошо — модель сходится.

Но увы, метрики на редуцированной размерности оказались сильно хуже. Основная размерность обучилась, нижняя же — не сошлась.
В целом на этом когда-то и завершилась наша первая попытка с матрешкой и метод был отложен в долгий ящик.
Спустя время мы вернулись к методу, докинули больше логирования, поигрались с гиперпараметрами и спустя K попыток моделька сошлась на размерностях 1024, 128, 64.

Что может помочь в интеграции матрёшки, по нашему опыту:
выбор меньшей гранулярности,
настройка весов для размерностей;
голова под каждую размерность, чтобы оптимизировать сходимость.
Чего мы добились
Во-первых, мы сохранили метрики на основной размерности.
Во-вторых, мы сохранили 95–96% от исходных метрик на размерности 128. Потеря составляет всего 4–5% при сжатии в восемь раз.
В-третьих, рекламные баннеры теперь доранжируются онлайн по запросам пользователей.
А если по-человечески, то теперь при запросе «вертикальный пылесос» пользователь видит баннеры только с вертикальными пылесосами — и никаких кофемашин ;)
Поделитесь впечатлениями от моего рассказа в комментвриях. Как вы действовали бы на нашем месте?

