"Applied ML", an interesting case :)

https://twitter.com/kchonyc/status/1719781112703054245

🪆Matryoshka Representation Learning
Aditya Kusupati, Gantavya Bhatt, Aniket Rege, Matthew Wallingford, Aditya Sinha, Vivek Ramanujan, William Howard-Snyder, Kaifeng Chen, Sham Kakade, Prateek Jain, Ali Farhadi
Статья: https://arxiv.org/abs/2205.13147
Код: https://github.com/RAIVNLab/MRL

Работа с NeurIPS 2022, которую я давно хотел разобрать. Лежала у меня среди тысячи открытых вкладок (где ещё уйма интересного, ждущего своего часа), пока я не нарвался на свежую “Matryoshka Diffusion Models” (https://arxiv.org/abs/2310.15111). Но про diffusion models не так интересно, как про репрезентации.

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

Собственно вопрос, можно ли придумать какую-то гибкую репрезентацию, которая сможет адаптироваться под разные задачи с разными требованиями по вычислительным ресурсам? Ну вы уже поняли, что можно!

Авторы предлагают подход 🪆Matryoshka Representation Learning (MRL), позволяющий получать такие репрезентации. Идея проста: в эмбеддинге фиксированного размера (d) выделим вложенные друг в друга части (подпространства) размеров d/2, d/4, d/8, … каждая из которых будет хорошим эмбеддингом для своего размера.

Задача в том, чтобы модифицировать процедуру обучения так, чтобы в векторе размерности d первые m измерений содержали общие и переносимые на другие задачи репрезентации.

Работу проводят на supervised мультиклассовой классификации. Берётся ResNet50, который эмбеддит в вектор размерности d = 2048, датасет ImageNet-1K с тысячью классов, и набор вложенных размерностей M = {8, 16, . . . , 1024, 2048}.

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

Линейные классификаторы тоже могут быть сделаны с общими весами, так что веса классификатора от маленького эмбеддинга являются подмножеством весов от классификатора для большего эмбеддинга. Это помогает сэкономить память в случае очень больших пространств эмбеддингов. Такой вариант подхода называется Efficient Matryoshka Representation Learning (MRL–E).

Подход проверили на разных задачах.

Для representation learning взяты задачи:
(a) Supervised learning для изображений: ResNet50 на ImageNet-1K и ViT-B/16 на JFT-300M.
(b) Contrastive learning для изображений и языка: ALIGN с ViT-B/16 vision encoder и энкодер BERT на данных ALIGN (https://news.1rj.ru/str/gonzo_ML/679).
(c) Masked language modelling: BERT на English Wikipedia и BooksCorpus.

Лучшие гиперпараметры не подбирали, взяли те, что были у независимо обученных бейзлайнов.

У резнета размер эмбеддинга 2048, у ViT и BERT 768. Лестницы размерностей соответственно следующие: M = {8, 16, 32, 64, 128, 256, 512, 1024, 2048} и M = {12, 24, 48, 96, 192, 384, 768}.

Для сравнения взяты низкоразмерные репрезентации фиксированного размера (FF), уменьшенные через SVD, из slimmable networks и рандомно выбранные фичи из FF максимального размера.

Оценивают качество классификации на ImageNet через linear classification/probe (LP) и 1-nearest neighbour (1-NN). MRL рулит, при малом размере эмбеддинга качество даже чуть выше, чем для фиксированного эмбеддинга такого же размера. И намного лучше и рандомно выбранных фич, и SVD.

Для ViT и JFT-300M эксперименты дорогие, это web-scale датасет, здесь обучали только самую большую FF модель. И здесь MRL показывает свою работоспособность. Кроме того, полученные эмбеддинги можно интерполировать, если нужен какой-то промежуточный размер -- точность на таким образом полученных репрезентациях лежит на кривой ровно там, где и можно было бы ожидать.

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

Отдельная тема -- retrieval. Цель -- найти картинки из того же класса, что запрос, используя эмбеддинги. Оценивают по mean Average Precision@10. Эмбеддинги нормализованы и достаются через L2 близость. MRL бьёт бейзлайны, при этом ещё и лучше отдельных FF эмбеддингов. MRL-E чуть похуже.

В сочетании с adaptive retrieval (AR) можно сэкономить кучу места, не храня полные репрезентации. В AR сначала получают шорт-лист (K=200) кандидатов через низкоразмерную репрезентацию (Ds = 16), затем список переранжируют через более высокоразмерную (Dr = 2048) -- это сильно дешевле по вычислениям, чем сразу делать поиск соседей по полному эмбеддингу. AR с указанными выше параметрами настолько же точна, что и выборка по полному эмбеддингу (d = 2048), но в 128 раз более эффективна в теории, и в 14 раз быстрее на практике.

Определить оптимальные значения Ds и Dr может быть непросто, поэтому авторы предлагают подход Funnel Retrieval, где используется каскад последовательно увеличивающихся эмбеддингов. Начинаем с самых малых, затем полученный на каждом этапе список переранжируем более широким эмбеддингом, в два раза сокращая список и в два раза увеличивая размерность каждый раз.

Интересно, какие-нибудь из векторных баз это уже поддерживают? Ash? ;)

По крайней мере есть следующая работа про adaptive approximate nearest neighbour search с использованием матрёшек -- работа побывавшая одновременно и на ICLR 2023 (https://iclr.cc/virtual/2023/13999) и на NeurIPS 2023 (AdANNS, https://arxiv.org/abs/2305.19435) — что, так можно?

На out-of-domain датасетах робастность MRL репрезентаций не хуже обычных. Для retrieval она даже немного повыше. На few-shot learning через nearest class mean тоже хорошо. На новых классах из длинного хвоста даже улучшение.

В целом с увеличением эмбеддинга качество классификации растёт. Но есть некоторое число случаев, когда на малом размере эмбеддинга классификация лучше. Если оценивать качество на супер-классе (когда в разметке есть иерархия), то там не такая сильная просадка на малом размере репрезентации. То есть какую-то иерархию эти эмбеддинги тоже ловят.

Matryoshka Representations можно получить частичным файнтюном уже обученных моделей.

Итого, почти халявный способ улучшения репрезентаций. Было бы интересно обучить эмбеддинговую модель типа text-embedding-ada или тех, что делает Cohere, в такой постановке. Ну и вообще большой простор для улучшений -- отдельные лоссы, поиск весов, специальные структуры для поиска.