2D Matryoshka Sentence Embeddings (2402.14776v3)

Abstract: Common approaches rely on fixed-length embedding vectors from LLMs as sentence embeddings for downstream tasks such as semantic textual similarity (STS). Such methods are limited in their flexibility due to unknown computational constraints and budgets across various applications. Matryoshka Representation Learning (MRL) \cite{aditya2022matryoshka} encodes information at finer granularities, i.e., with lower embedding dimensions, to adaptively accommodate \emph{ad hoc} tasks. Similar accuracy can be achieved with a smaller embedding size, leading to speedups in downstream tasks. Despite its improved efficiency, MRL still requires traversing all Transformer layers before obtaining the embedding, which remains the dominant factor in time and memory consumption. This prompts consideration of whether the fixed number of Transformer layers affects representation quality and whether using intermediate layers for sentence representation is feasible. In this paper, we introduce a novel sentence embedding model called \textit{Two-dimensional Matryoshka Sentence Embedding} (2DMSE)\footnote{Our code is available at \url{https://github.com/SeanLee97/AnglE/blob/main/README_2DMSE.md}.}. It supports elastic settings for both embedding sizes and Transformer layers, offering greater flexibility and efficiency than MRL. We conduct extensive experiments on STS tasks and downstream applications. The experimental results demonstrate the effectiveness of our proposed model in dynamically supporting different embedding sizes and Transformer layers, allowing it to be highly adaptable to various scenarios.

Enhancing Sentence Embedding Flexibility with 2D Matryoshka Sentence Embeddings

Introduction to 2D Matryoshka Sentence Embeddings (2DMSE)

The landscape of sentence embedding has been significantly advanced by the introduction of Two-dimensional Matryoshka Sentence Embeddings (2DMSE). This novel framework inherits the principles of Matryoshka Representation Learning (MRL), extending its capabilities to harness both the depth of Transformer layers and the granularity of embedding sizes for sentence embedding tasks. In this context, 2DMSE distinguishes itself by offering a unique approach to generate sentence embeddings that are not only efficient but also maintain high semantic accuracy across a variety of benchmarks.

Technical Overview

2DMSE is conceptualized around the flexibility of generating sentence embeddings through the integration of elastic settings for both model depth and embedding dimensions. This two-dimensional scalability allows for significant improvements in both computational efficiency and flexibility in embedding adaptation. The framework employs a process where, at each training step, a Transformer layer is randomly sampled, and embeddings from this layer, alongside those from the last layer, are fine-tuned using a customized loss function. This is complemented by an alignment method that minimizes the Kullback-Leibler divergence between embeddings from the chosen layer and the last layer, facilitating a coherence in semantic representation across the model's depth.

Experimental Insights

The efficacy of 2DMSE has been comprehensively established through extensive experiments on standard Semantic Textual Similarity (STS) benchmarks. Key findings include:

• The capability of 2DMSE to produce embedding vectors from intermediate Transformer layers that exhibit significant qualitative improvements over those generated by traditional sentence embedding methods and the MRL framework.
• Demonstrated superiority in embedding performance across various layers and dimensions, with 2DMSE achieving remarkable scores on STS benchmarks that exceed those of powerful baselines including SBERT, USE, and the state-of-the-art AnglE framework.

Theoretical and Practical Implications

From a theoretical standpoint, 2DMSE introduces a paradigm shift in how sentence embeddings are generated, leveraging the inherently hierarchical structure of language representation within Transformer models to a fuller extent. Practically, it equips users with an unmatched level of adaptability, enabling the tailoring of embedding generation to specific computational budgets without sacrificing performance. This adaptability is particularly crucial in resource-constrained environments, where efficiency in processing and memory usage is paramount.

Looking Forward

The introduction of 2DMSE opens up new avenues for research and application in the field of NLP. Future explorations could delve into the optimization of the layer sampling strategy or the exploration of different alignment methodologies to further boost the quality of sentence embeddings. Additionally, the framework's inherent flexibility suggests potential for widespread adoption in diverse NLP tasks ranging from information retrieval to real-time language understanding systems.

In conclusion, the Two-dimensional Matryoshka Sentence Embeddings framework sets a new benchmark in the development of scalable, efficient, and high-performance sentence embedding models. By addressing the dual aspects of model depth and embedding size adaptability, 2DMSE paves the way for more nuanced and versatile approaches to capturing semantic nuances in textual data.