JTCSE: Joint Tensor-Modulus Constraints and Cross-Attention for Unsupervised Contrastive Learning of Sentence Embeddings

TL;DR

This paper addresses limitations in unsupervised sentence embedding by introducing JTCSE, a twin-tower framework that jointly enforces tensor modulus constraints and cross-attention to improve representation magnitude and CLS pooling. The method integrates a modulus-based loss $L_{TMC}$ with cross-attention between encoders, producing enhanced positive alignment and richer global representations via increased CLS energy. Empirical results on seven STS tasks set new state-of-the-art performance, and extensive zero-shot evaluations across 130 downstream tasks demonstrate robust generalization; the approach also achieves substantially lower inference overhead than traditional large ensembles, with distillation yielding competitive single-encoder performance. The work provides detailed ablations and analyses, validating the complementary roles of modulus constraints and cross-attention and offering open-source code and checkpoints for reproducibility and further research.

Abstract

Unsupervised contrastive learning has become a hot research topic in natural language processing. Existing works usually aim at constraining the orientation distribution of the representations of positive and negative samples in the high-dimensional semantic space in contrastive learning, but the semantic representation tensor possesses both modulus and orientation features, and the existing works ignore the modulus feature of the representations and cause insufficient contrastive learning. % Therefore, we firstly propose a training objective that aims at modulus constraints on the semantic representation tensor, to strengthen the alignment between the positive samples in contrastive learning. Therefore, we first propose a training objective that is designed to impose modulus constraints on the semantic representation tensor, to strengthen the alignment between positive samples in contrastive learning. Then, the BERT-like model suffers from the phenomenon of sinking attention, leading to a lack of attention to CLS tokens that aggregate semantic information. In response, we propose a cross-attention structure among the twin-tower ensemble models to enhance the model's attention to CLS token and optimize the quality of CLS Pooling. Combining the above two motivations, we propose a new \textbf{J}oint \textbf{T}ensor representation modulus constraint and \textbf{C}ross-attention unsupervised contrastive learning \textbf{S}entence \textbf{E}mbedding representation framework JTCSE, which we evaluate in seven semantic text similarity computation tasks, and the experimental results show that JTCSE's twin-tower ensemble model and single-tower distillation model outperform the other baselines and become the current SOTA. In addition, we have conducted an extensive zero-shot downstream task evaluation, which shows that JTCSE outperforms other baselines overall on more than 130 tasks.

Figures (10)

• Figure 1: Subfigure a. represents the traditional ensemble modeling approach (EDFSEzong), which naively trains multiple sub-encoders separately and then sums the outputs. This approach causes a large inference overhead. Subfigure b. represents the optimized ensemble learning framework JTCSE proposed in this work. It incorporates semantic representation tensor modulus constraints and joint modeling of cross-attention between sub-encoders. This framework contains only two sub-encoders. It significantly reduces inference overhead while improving the quality of sentence embeddings relative to a.
• Figure 2: This figure represents the distribution of the positions of a pair of positive sample semantic representation tensors $h$ and $h^{+}$ in three-dimensional space and the vectors $h-h^{+}$ for which they are subtracted. According to the principle of similar triangles, when the angle $\gamma$ is specific, the larger the modulus of $h$ or $h^{+}$, the larger the modulus of $h-h^{+}$ will be, and the greater the value of being constrained.
• Figure 3: This figure illustrates the binary loss function $L_{TMC}$, with respect to the range of values of the two independent variables $t$ and $k$ over part of its domain of definition.
• Figure 4: This figure shows the structure of the proposed unsupervised sentence embedding representation framework, JTCSE, which consists of two main parts: the semantic representation tensor modulus constraints and the joint modeling of subencoder cross-attention. Subfigure b. shows the overall structure of JTCSE, which contains two subencoders, I and II. Each is a fine-tuned BERT-like model that includes an embedding layer, an encoder, and a pooler layer. Before the training, we specify the cross-attention encoder layer's(CAEL) position in the encoder, the position of CAEL in both subencoders is the same. During training, a mini-batch is fed into the embedding layer of two sub-encoders simultaneously, and the hidden state output from each embedding layer goes into its own encoder; if CAEL is encountered, in addition to the normal forward propagation within each sub-encoder, it is also necessary to mutually pass through the attention network in each other's EncoderLayer to achieve the computation of cross-attention. Both the primitive last hidden state (LHS) and the cross-attention's LHS pass through the IC-InfoNCE constraints. The primitive LHS also passes through the pooler layer to get the pooler output, which in turn passes through the tensor modulus-constrained loss function. Subfigure a. represents the details of CAEL, the Query, Key, and Value weights in MHA and MHCA are identical. Subfigure c. represents the details of ICTM loss and IC-InfoNCE loss.
• Figure 5: This figure reports the cosine similarity density distribution plots of different models on the STS-B dataset, where sentence pairs are uniformly divided into five groups, with the vertical coordinates of each subplot denoting the group and the horizontal coordinates denoting the model scoring. Each subplot should have an overall “sub-diagonal” distribution, indicating closer to the labeling distribution. We use the same code to report the performance of all models.