2D Matryoshka Training for Information Retrieval

TL;DR

This study reproduces and evaluates the 2D Matryoshka Training method for embedding models across various tasks, confirming its effectiveness and exploring optimal configurations for retrieval applications.

Contribution

It provides a comprehensive reproducibility analysis of 2D Matryoshka Training, compares two implementation versions, and extends evaluation to retrieval tasks with insights on loss functions and training strategies.

Findings

Both versions outperform traditional training on sub-dimensions but not on specific sub-layer setups.

Effectiveness generalizes well to retrieval tasks in supervised and zero-shot settings.

Certain loss configurations, like full-dimension loss, improve retrieval performance.

Abstract

2D Matryoshka Training is an advanced embedding representation training approach designed to train an encoder model simultaneously across various layer-dimension setups. This method has demonstrated higher effectiveness in Semantic Text Similarity (STS) tasks over traditional training approaches when using sub-layers for embeddings. Despite its success, discrepancies exist between two published implementations, leading to varied comparative results with baseline models. In this reproducibility study, we implement and evaluate both versions of 2D Matryoshka Training on STS tasks and extend our analysis to retrieval tasks. Our findings indicate that while both versions achieve higher effectiveness than traditional Matryoshka training on sub-dimensions, and traditional full-sized model training approaches, they do not outperform models trained separately on specific sub-layer and sub-dimension setups. Moreover, these results generalize well to retrieval tasks, both in supervised (MSMARCO) and zero-shot (BEIR) settings. Further explorations of different loss computations reveals more suitable implementations for retrieval tasks, such as incorporating full-dimension loss and training on a broader range of target dimensions. Conversely, some intuitive approaches, such as fixing document encoders to full model outputs, do not yield improvements. Our reproduction code is available at https://github.com/ielab/2DMSE-Reproduce.