🪆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://t.me/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.
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://t.me/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.
Источник: gonzo-обзоры ML статей
2023-11-03 20:10:56