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

Зачем текстовый энкодер в «Похожих товарах»

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

Коллаборативная фильтрация для самых маленьких
Коллаборативная фильтрация для самых маленьких

А если истории нет? Ну на нет и суда нет?) Увы, суд всё-таки есть и в реальных системах это базовый минимум, а не роскошный максимум. 

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

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

В чём плюсы контентной модели? Во-первых, независимость от наличия истории взаимодействий с конкретным товаром. Во-вторых, относительная интерпретируемость по сочетанию цвета, бренда, материала и других характеристик (хотя это тоже все еще черепикинг).

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

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

Наш текстовый энкодер

За основу мы взяли Multilingual-E5-large — мультиязычный текстовый трансформер. Модель все еще показывает хорошие результаты на большинстве бенчмарков и лучший (на тот момент) результат на данных Wildberries без дообучения. 24 слоя, размерность эмбеддинга 1024, контекст 512.

карточка E5 на Hugging Face
карточка E5 на Hugging Face

E5 — тюн BERT для задачи поиска в две стадии. На первом этапе модели обучались примерно на миллиарде пар текстов на различных языках, используя контрастивный метод с функцией потерь InfoNCE-loss и батчем 32 тыс. Суть InfoNCE-loss — сделать соответствующие пары ближе, чем любые другие.

На втором этапе модели дообучались на комбинации высококачественных размеченных датасетов (MSMARCO, NQ, TriviaQA, SQuAD, NLI и другие) с общим объёмом около 1,6 млн примеров. Данных меньше, но их качество выше.

InfoNCE Loss
InfoNCE Loss

Обучение для «Похожих» строилось по тому же принципу: контрастивное дообучение с in-batch негативами на нашем датасете «похожих товаров». 

Какой текст идёт на вход?

  • расцветка товара,

  • заголовок товара,

  • пометка 18+ или нет,

  • название бренда,

  • все характеристики товара в формате по строкам — «название характеристики: значение».

Пример текстового описания товара:

passage: кремовый, бежевый, темно-синий  Блузка "AAA" <sep>бренда: brandname<sep>

Страна производства: Россия

Коллекция: Базовая коллекция

Комплектация: Блузка — 1 шт

Фактура материала: гладкая

Вырез горловины: округлый

Мы пробовали докидывать тексты: в офлайне показатели растут, в проде — эффект не стат. значимый.

Как устроен инференс модели?

Перегоняем предобученный чекпоинт во что-то быстрое на ваш вкус — например, BetterTransformer, TensorRT, Тритон и так далее. Векторизуем каталог во float16 — он кратно быстрее благодаря оптимизации NVIDIA и занимает в два раза меньше места без потери качества. Закидываем в Faiss-GPU (или любой другой ANN индекс) — рекомендации готовы.

Пример выдачи

Пример рекомендаций: якорь - 3 рекомендации
Пример рекомендаций: якорь - 3 рекомендации

Первый товар — это якорный товар, к которому мы подбираем рекомендации. В верхнем ряду выдача промежуточной 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!

PCA
PCA

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

SVD
SVD

В свою очередь TruncatedSVD считает усечённое/приближённое SVD вычисляя только первые (наиболее значимые) сингулярные векторы матрицы. Тестировали реализацию от cuML с GPU-ускорением. Итог: 68% исходных метрик.

GaussianRandomProjection и UMAP показали только 53% и 24% от исходных метрик соответственно, так что не вижу особого смысла заострять на них внимание.

А теперь добавим чуть-чуть deep и вспомним про автоэнкодеры. Автоэнкодер — полносвязная нейронная сеть, которая состоит из энкодера и декодера. Идея следующая: сжимаем исходный эмбеддинг → разжимаем его → сравниваем разжатый эмбеддинг с исходным → считаем loss → отрезаем декодер → получаем модель для сжатия эмбеддингов.

Автоэнкодер штука крайне кастомизируемая: можно варьировать функции активации, слои, loss и даже, если очень хочется, добавить квантизацию (Rate-Adaptive Quantized Variational Autoencoder). Кроме того, модель довольно легко интегрировать в пайплайн. Из грустного, модель крайне склонна к переобучению, но это болезнь всех полносвязных нейронок. Итог: 75% от исходных метрик.

Autoencoder
Autoencoder

По результатам первой итерации победили автоэнкодеры.

Можно ли лучше?

К началу второй итерации дела обстояли вот так:

Векторизация текста
Векторизация текста

У нас есть текстовый энкодер, мы даем ему текст, она выдает нам вектор размерности 1024, где фичи как-то раскиданы по этим 1024 измерениям. 

А дальше классика последних лет:  «не учи много моделей — обучи  один хороший трансфор».

Вот мы тоже  так подумали и решили попробовать прикрутить «Матрешку» к нашей E5. 

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

Как работает MRL
Как работает MRL

Основной эмбеддинг — это самая большая и детализированная матрёшка. Мы обрезаем её на меньшую размерность (512, 256 и т. д.) и получаем меньше деталей, сохраняя суть. Еще можно сравнить данный подход с рядами Тейлора, где мы регулируем точность разложения с помощью количества членов.

Идея классная, как реализовать? Идейно довольно просто. 

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

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

Обучение MRL
Обучение MRL

Окей, но получиться дотюнить модель или нужно учить с нуля? Хватит ли адаптера или LoRA? Не просядет ли качество на основной размерности? 

Ожидание и реальность

Если верить авторам статьи, то с нуля все должно быть все хорошо. 

Вот график  ResNet50 на ImageNet-1K:

Качество модели на разных размерностях из статьи
Качество модели на разных размерностях из статьи

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

Теперь то же самое, только для MRL и с распределением по классам:

Качество модели на разных размерностях по классам  из статьи
Качество модели на разных размерностях по классам из статьи

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

e5_mrl.jpg
Первая попытка

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

В целом на этом когда-то и завершилась наша первая попытка с матрешкой и метод был отложен в долгий ящик. 

Спустя время мы вернулись к методу, докинули больше логирования, поигрались с гиперпараметрами и спустя K попыток моделька сошлась на размерностях 1024, 128, 64. 

newplot.png
Оранжевый loss ведёт нашу модель к светлому будущему, остальные подтягиваются!

Что может помочь в интеграции матрёшки, по нашему опыту:

  • выбор меньшей гранулярности,

  • настройка весов для размерностей;

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

Чего мы добились

Во-первых, мы сохранили метрики на основной размерности.

Во-вторых, мы сохранили 95–96% от исходных метрик на размерности 128.  Потеря составляет всего 4–5% при сжатии в восемь раз.

В-третьих, рекламные баннеры теперь доранжируются онлайн по запросам пользователей.

А если по-человечески, то теперь при запросе «вертикальный пылесос» пользователь видит баннеры только с вертикальными пылесосами — и никаких кофемашин ;)

Поделитесь впечатлениями от моего рассказа в комментвриях. Как вы действовали бы на нашем месте?