Matryoshka Representation Learning

Fuente: arXiv
Saved in:
Bibliographic Details
Main Authors: Kusupati, Aditya, Bhatt, Gantavya, Rege, Aniket, Wallingford, Matthew, Sinha, Aditya, Ramanujan, Vivek, Howard-Snyder, William, Chen, Kaifeng, Kakade, Sham, Jain, Prateek, Farhadi, Ali
Format: Preprint
Published: 2022
Subjects:
Online Access:
Tags: Add Tag
No Tags, Be the first to tag this record!
_version_ 1866909097895395328
author Kusupati, Aditya
Bhatt, Gantavya
Rege, Aniket
Wallingford, Matthew
Sinha, Aditya
Ramanujan, Vivek
Howard-Snyder, William
Chen, Kaifeng
Kakade, Sham
Jain, Prateek
Farhadi, Ali
author_facet Kusupati, Aditya
Bhatt, Gantavya
Rege, Aniket
Wallingford, Matthew
Sinha, Aditya
Ramanujan, Vivek
Howard-Snyder, William
Chen, Kaifeng
Kakade, Sham
Jain, Prateek
Farhadi, Ali
contents Learned representations are a central component in modern ML systems, serving a multitude of downstream tasks. When training such representations, it is often the case that computational and statistical constraints for each downstream task are unknown. In this context rigid, fixed capacity representations can be either over or under-accommodating to the task at hand. This leads us to ask: can we design a flexible representation that can adapt to multiple downstream tasks with varying computational resources? Our main contribution is Matryoshka Representation Learning (MRL) which encodes information at different granularities and allows a single embedding to adapt to the computational constraints of downstream tasks. MRL minimally modifies existing representation learning pipelines and imposes no additional cost during inference and deployment. MRL learns coarse-to-fine representations that are at least as accurate and rich as independently trained low-dimensional representations. The flexibility within the learned Matryoshka Representations offer: (a) up to 14x smaller embedding size for ImageNet-1K classification at the same level of accuracy; (b) up to 14x real-world speed-ups for large-scale retrieval on ImageNet-1K and 4K; and (c) up to 2% accuracy improvements for long-tail few-shot classification, all while being as robust as the original representations. Finally, we show that MRL extends seamlessly to web-scale datasets (ImageNet, JFT) across various modalities -- vision (ViT, ResNet), vision + language (ALIGN) and language (BERT). MRL code and pretrained models are open-sourced at https://github.com/RAIVNLab/MRL.
format Preprint
id arxiv_https___arxiv_org_abs_2205_13147
institution arXiv
publishDate 2022
record_format arxiv
spellingShingle Matryoshka Representation Learning
Kusupati, Aditya
Bhatt, Gantavya
Rege, Aniket
Wallingford, Matthew
Sinha, Aditya
Ramanujan, Vivek
Howard-Snyder, William
Chen, Kaifeng
Kakade, Sham
Jain, Prateek
Farhadi, Ali
Machine Learning
Computer Vision and Pattern Recognition
Learned representations are a central component in modern ML systems, serving a multitude of downstream tasks. When training such representations, it is often the case that computational and statistical constraints for each downstream task are unknown. In this context rigid, fixed capacity representations can be either over or under-accommodating to the task at hand. This leads us to ask: can we design a flexible representation that can adapt to multiple downstream tasks with varying computational resources? Our main contribution is Matryoshka Representation Learning (MRL) which encodes information at different granularities and allows a single embedding to adapt to the computational constraints of downstream tasks. MRL minimally modifies existing representation learning pipelines and imposes no additional cost during inference and deployment. MRL learns coarse-to-fine representations that are at least as accurate and rich as independently trained low-dimensional representations. The flexibility within the learned Matryoshka Representations offer: (a) up to 14x smaller embedding size for ImageNet-1K classification at the same level of accuracy; (b) up to 14x real-world speed-ups for large-scale retrieval on ImageNet-1K and 4K; and (c) up to 2% accuracy improvements for long-tail few-shot classification, all while being as robust as the original representations. Finally, we show that MRL extends seamlessly to web-scale datasets (ImageNet, JFT) across various modalities -- vision (ViT, ResNet), vision + language (ALIGN) and language (BERT). MRL code and pretrained models are open-sourced at https://github.com/RAIVNLab/MRL.
title Matryoshka Representation Learning
topic Machine Learning
Computer Vision and Pattern Recognition
url https://arxiv.org/abs/2205.13147