提升2D套娃嵌入模型训练,性能逼近独立训练模型。
Starbucks-v2: Improved Training for 2D Matryoshka Embeddings
- 用结构化微调+MAE预训练,按层维固定顺序优化
- 在多个任务上达到独立模型性能,且保持高效可调
- 适合需要多尺寸嵌入的部署场景,支持深度/宽度混合
2D套娃嵌入训练使单一模型生成不同层数和维度的子网络表示,适应多样计算与任务需求。但其性能远低于同等规模的独立训练模型。为此,我们提出Starbucks,一种新的套娃嵌入模型训练策略,结合结构化微调与掩码自编码器(MAE)预训练。微调阶段,对一组固定的层-维度组合(从小到大)计算损失,显著优于随机采样子网络,性能匹配独立训练模型。基于MAE的预训练进一步提升子网络表示质量,为下游任务提供更强基础。在域内(语义相似度、段落检索)和域外(BEIR)基准测试中,Starbucks始终优于传统2D套娃模型,并匹配或超越独立训练模型性能,同时保持高效率与可调性。消融实验验证了损失设计合理性、SMAE预训练效果及跨主干模型的适用性。我们还发现深度与宽度方向的Starbucks变体捕捉互补信息,其融合可带来额外增益,延迟开销极小(因可并行化)。代码见https://github.com/ielab/Starbucks。
原文摘要 · Abstract (English)
2D Matryoshka training enables a single embedding model to generate sub-network representations across different layers and embedding dimensions, offering adaptability to diverse computational and task constraints. However, its effectiveness remains well below that of individually trained models of equivalent 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, from small to large, which significantly improves performance over randomly sampled sub-networks and matches that of separately trained models. Our MAE-based pre-training further enhances the representation quality of sub-networks, providing a stronger 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 adaptability. Ablation studies confirm our loss design choices, the impact of SMAE pre-training and demonstrate the applicability of Starbucks across backbones. We further show that depth- and width-wise Starbucks variants capture complementary information, and that their hybridization yields additional performance gains with minimal latency overhead due to parallelization. Code available at https://github.com/ielab/Starbucks
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。