2D Matryoshka training enables a single embedding model to produce sub-network representations across varying layers and embedding dimensions, offering flexibility under different computational and task constraints. However, its performance remains below that of individually trained models of comparable sizes. To address this, we propose Starbucks, a new training strategy for Matryoshka-style embedding models that combines structured fine-tuning with masked autoencoder (MAE) pre-training. During fine-tuning, we compute the loss over a fixed set of layer-dimension pairs, ordered from small to large, which significantly improves over random sub-network sampling and matches the performance of separately trained models. Our MAE-based pre-training further strengthens sub-network representations, providing a more robust backbone for downstream tasks. Experiments on both in-domain (semantic similarity and passage retrieval) and out-of-domain (BEIR) benchmarks show that Starbucks consistently outperforms 2D Matryoshka models and matches or exceeds the performance of individually trained models, while maintaining high efficiency and flexibility. Ablation studies validate our loss design, the benefits of SMAE pre-training, and demonstrate Starbucks’ applicability across backbones. We further show that depth- and width-wise Starbucks variants encode complementary information, and that combining them yields further gains with minimal latency overhead via parallelization (Code at https://github.com/ielab/Starbucks .).

错误:搜索内容不能为空,请输入英文关键词
错误:关键词超出字数限制,请精简
高级检索

A simple sketch of a disposable coffee cup with a lid. The cup features a green circle in the center, possibly representing a logo or design element. The drawing is outlined in black with minimal detail. Starbucks: Improved Training for 2D Matryoshka Embeddings

  • Shengyao Zhuang,
  • Shuai Wang,
  • Fabio Zheng,
  • Bevan Koopman,
  • Guido Zuccon

摘要

2D Matryoshka training enables a single embedding model to produce sub-network representations across varying layers and embedding dimensions, offering flexibility under different computational and task constraints. However, its performance remains below that of individually trained models of comparable sizes. To address this, we propose Starbucks, a new training strategy for Matryoshka-style embedding models that combines structured fine-tuning with masked autoencoder (MAE) pre-training. During fine-tuning, we compute the loss over a fixed set of layer-dimension pairs, ordered from small to large, which significantly improves over random sub-network sampling and matches the performance of separately trained models. Our MAE-based pre-training further strengthens sub-network representations, providing a more robust backbone for downstream tasks. Experiments on both in-domain (semantic similarity and passage retrieval) and out-of-domain (BEIR) benchmarks show that Starbucks consistently outperforms 2D Matryoshka models and matches or exceeds the performance of individually trained models, while maintaining high efficiency and flexibility. Ablation studies validate our loss design, the benefits of SMAE pre-training, and demonstrate Starbucks’ applicability across backbones. We further show that depth- and width-wise Starbucks variants encode complementary information, and that combining them yields further gains with minimal latency overhead via parallelization (Code at https://github.com/ielab/Starbucks .).