Length-Induced Embedding Collapse in Transformer-based Models

W3

The improvement of TempScale on each model is inconsistent, and the main improvement does not come from the long text part (as shown in Figure 6), so more explanation and exploration are necessary.

Thank you for your feedback.

• For Consistency improvement： We recognize that TempScale's improvements vary across models due to differences in architecture. For example, GTR is based on T5, ANCE on RoBERTa, and both GIST and BGE on BERT, which leads to similar improvement patterns in GIST and BGE. Additionally, for simplicity, we used the same temperature setting across models in our main table, which may have impacted performance on some tasks. Setting optimal temperatures per task could yield more consistent improvements across models.

• For Short-text Performance: Furthermore, as for Figure 6 showing greater improvement on short texts than long texts, this is due to the distributional properties of embeddings. For example, in MTEB’s KNN classifier, short texts are more dispersed in the embedding space, while long texts, due to length collapse, cluster toward the center. During KNN optimization, long texts, being closer to the embedding center, influence voting results more strongly. By applying TempScale (as shown in Figure 1c), short and long texts become more aligned in distribution, enabling short texts to have a greater influence and leading to a larger improvement in performance on short texts.

• For Long-Text Performance: Finally, Figures 7, 8, and 9 demonstrate that TempScale helps mitigate length collapse, providing more contextual representations for long texts and enhancing model performance on these examples.

W4

The relevant work [3] and the important long-context embedder [4] have been overlooked.

Thank you for providing your work and model. For the relevant work, our study differs as follows:

• Context Differences: While the relevant work focuses on high similarity in repeated text segments, our study identifies this phenomenon across all natural language sequences.
• Mechanism Explanation: [3] attributes the similarity issue in repeated texts to reduced attention to the [CLS] token. In contrast, we rigorously analyze this from a low-pass filtering perspective, showing that long text embeddings cluster in a narrower space, increasing similarity between long texts compared to shorter ones. Additionally, we reveal that the distributional inconsistency between long and short text embeddings contributes to the reduced performance of long texts in downstream tasks.

Limitations with the important long-context embedder Model (GTE): Regarding the GTE model, we initially considered it as a base model for our main experiment. However, as GTE relies on a custom-defined model, we were unable to access the code for its attention module, making it incompatible with our TempScale method. To address the absence of a long-text base model, we used the e5-4k-base model instead, which can handle texts up to 4k tokens and has demonstrated strong performance on the LongEmbed Benchmark. Furthermore, we still observed the Length Collapse phenomenon in the GTE model (Alibaba-NLP/gte-base-en-v1.5). Specifically, We conducted experiments using three commonly used long-text datasets, including LongRAG and LongAlpaca-12k from HuggingFace, which were selected based on relevance to the keyword "long" and chosen for their evenly distributed text lengths. We analyzed distribution shifts in embedding space across different text length intervals, calculating the average Euclidean distance of these embeddings from the central embedding and the cosine similarity between each pair of embeddings.

Euclidean distance

Pair Cosine Similarity

The experimental results show that as text length increases, embeddings from the GTE model display a convergence trend, with decreasing average Euclidean distances between embeddings and increasing pairwise cosine similarities. This suggests that even mainstream long-context embedding models experience embedding convergence (Length Collapse).

Furthermore, we also explored whether such Length Collapse affects the model's performance. To investigate this, we tested the GTE model on the IMDB dataset across different text length intervals, and the results are as follows:

Acc on ImbdClassification

The results indicate that, despite the GTE model having a longer context window, it still experiences a performance decline on longer texts.

W2

• I am wondering whether long-context llms still suffer this length collapse, and perhaps furthers discussions on this direction would be more exciting and expected.

Thank you for your suggestion. To explore whether length collapse also occurs in long-context LLMs, we selected three widely used LLM-based embedding models from the MTEB benchmark—bge-multilingual-gemma2[2], NV-Embed-v2[3], and e5-mistral-7b-instruct[4]—all of which rank within the top 20 on the MTEB leaderboard. We conducted experiments using three commonly used long-text datasets, including LongRAG[5] and LongAlpaca-12k[6] from HuggingFace, which were selected based on relevance to the keyword "long" and chosen for their evenly distributed text lengths. Specifically, we analyzed shifts in embedding space across different text length intervals by calculating the average Euclidean distance of each embedding from the central embedding (the mean of all embeddings) and computing cosine similarity between each pair of embeddings as follows.

Euclidean distance

Pair Cosine Similarity

The experimental results show that as text length increases, the embeddings from different LLM-based embedding models exhibit a gradual convergence trend, with the average Euclidean distance between embeddings decreasing and pairwise cosine similarity increasing. This indicates that even mainstream long-context LLM embedding models tend to experience embedding convergence (Length Collapse) and reduced distinctiveness in long-text processing due to the low-pass filtering effect.

Thank you very much for your detailed review and valuable comments on our paper. Your suggestions have provided important perspectives that helped us improve the content of our manuscript.

[1] Improving Text Embeddings with Large Language Models. ACL2024.

[2] BGE M3-Embedding: Multi-Lingual, Multi-Functionality, Multi-Granularity Text Embeddings Through Self-Knowledge Distillation. ACL2024.

[3] NV-Embed: Improved Techniques for Training LLMs as Generalist Embedding Models. arXiv2024.

[4] Text Embeddings by Weakly-Supervised Contrastive Pre-training. arXiv2022.

[5] LongRAG: Enhancing Retrieval-Augmented Generation with Long-context LLMs. arXiv2024.

[6] LongLoRA: Efficient Fine-tuning of Long-Context Large Language Models. ICLR2024.

2. Theoretical Comparison:

Method1 offers a solution using a specially designed Transformer example, but its motivation is limited by the specificity of this example. In contrast, Method2 introduces an empirically derived scaling formula, though the YaRN paper doesn’t explain the reasoning behind it; we attempt to provide this derivation. TempScale, however, is based on a rigorous low-pass filtering analysis, giving it a strong theoretical foundation and an intuitive explanation.

• Method1: This method involves building a Transformer to identify if the first string in a sequence of binary strings is "1". The authors observed that, as the sequence length $n$ increased, model performance dropped for sequences longer than the training length. They found that applying Method1 to the attention matrix mitigated performance declines on unseen lengths. However, this solution does not provide a theoretical basis, serving instead as an empirical fix within a specific model and task framework.

• Method2: While no specific motivation is mentioned for Method2, its scaling factor was empirically determined using the LLaMA2 model. After reviewing the original YaRN paper and OpenReview sources, the reasoning behind the formula choice remains unclear. Here, I attempt to derive a mathematical rationale for Method2. Assuming the attention matrix has the form $\text{softmax}\left(\frac{\tau}{\sqrt{d}} QK^T \right) V$, each element $a_{i,j}$ of the matrix can be interpreted as a conditional distribution, where the entropy of the attention at position $i$ can be defined as $\mathcal{H}_i=-\sum _ {j=1}^{n} a _ {i,j} \log a _ {i,j}$ . Method2’s extension principle combines "high-frequency extrapolation" with "low-frequency interpolation." For interpolation, we aim for consistent entropy $\mathcal{H} _ i$ at position $i$ before and after RoPE interpolation, requiring careful design of $\tau$ in the attention matrix. We can calculate $\mathcal{H} _ i$ after incorporating $\tau$ as follows:

  $$\mathcal{H} _ {i} = \log\sum _ {j=1}^{n} e^{\tau q _ {i}\cdot k _ {j}} - \frac{\sum _ {j=1}^{n} e^{\tau q _ {i}\cdot k _ {j}}(\tau q _ {i}\cdot k _ {j})}{\sum _ {j=1}^{n} e^{\tau q _ {i}\cdot k _ {j}}},$$

  Assuming $q$ and $k$ are normally distributed with mean 0 and variance 1, we approximate:

  $$\sum _ {j=1}^{n} e^{\tau q _ i \cdot k _ j} \approx n \mathbb{E} _ j[e^{\tau q _ i \cdot k _ j}],$$

  leading to

  $$\mathcal{H} _ {i} \approx \log n + \log \mathbb{E} _ {j}\left[e^{\tau q _ {i} \cdot k _ {j}}\right] - \frac{\tau \mathbb{E} _ {j}\left[e^{\tau q _ {i} \cdot k _ {j}} (q _ {i} \cdot k _ {j})\right]}{\mathbb{E} _ {j}\left[e^{\tau q _ {i} \cdot k _ {j}}\right]}.$$

  Substituting $q _ i \cdot k _ j$ with $d \cos \theta$, we obtain

  $$\mathcal{H} _ {i} \approx \log n + \log \mathbb{E} _ {\theta}\left[e^{\tau d \cos \theta}\right] - \frac{\tau d \mathbb{E} _ {\theta}\left[e^{\tau d \cos \theta} \cos \theta\right]}{\mathbb{E} _ {\theta}\left[e^{\tau d \cos \theta}\right]}.$$

  Applying Laplace approximation yields

  $$\mathcal{H} _ i \approx \log n - 0.24\tau d + \mathcal{O}(1).$$

  Since we aim for consistent entropy $\mathcal{H} _ i$ before and after RoPE interpolation, we set

  $$\log n - 0.24 d + \mathcal{O}(1) = \log sn -0.24\tau d + \mathcal{O}(1),$$

  leading to $\tau \approx \frac{\log s}{0.24} + 1$.

• TempScale: The main motivation of this method is to introduce a temperature that reduces the low-pass filtering effect in the attention matrix, allowing the final embedding to retain more high-frequency information, thus providing a more "contextual" representation. This method starts from the attention mechanism itself and does not address length generalization. Instead, it focuses on refining attention within the model’s training window.

Method2’s proof does not consider low-pass filtering because it is tailored to the specific scenario Method2 addresses—primarily, how to enable the original position embeddings from training to function effectively on longer texts. The solution includes interpolation, leading to the intuitive idea that the “information aggregation” capability of the attention matrix should remain consistent before and after interpolation. In the proof, “information aggregation” is measured by the entropy of the attention distribution. Moreover, we can also view Method2 from a low-pass filtering perspective: position embeddings originally help the attention matrix to better select information, which preserves high-frequency details in the sequence. When using the original $n$ position embeddings to model a sequence of length $sn$, scaling is needed to focus the model more on specific tokens, thereby retaining more high-frequency information.

Dear authors,

Thanks for you explanation. I realized that there's a typo in my previous response, where the baseline I intended to express should be $softmax(\frac{s}{\sqrt{d}}QK^T)V$, where $s=max(1,log_{L_o}^{n})$. But, anyway, since you are focusing on the issue of long texts within the model's training context window, it seems that $s$ will always be $1$. Looking forward to the version with comparison of Method1 & 2 included.

W2

The paper attributes its poor performance in modeling long texts to the low-pass filtering effect of the attention mechanism, though I question whether this is the primary cause. My concerns include:

• In Figure 1b, the t-SNE visualization shows long text embeddings clustering near the origin, but this proximity is in the reduced-dimensional space and doesn't confirm their position in the original embedding space.

Thank you for your attention to our t-SNE visualization. Our main purpose in using t-SNE was to better illustrate the overall distribution trends of long and short texts on a 2D plane through dimensionality reduction, rather than to obtain precise embedding positions. We are aware that t-SNE introduces a certain degree of spatial distortion during dimensionality reduction, so the clustering positions should not be directly equated with distances in the original embedding space. Additionally, the trend shown in Figure 1c—where the average pairwise cosine similarity between embeddings of different length intervals and the central embedding (the mean of all embeddings) increases with text length—partially indicates that embeddings for longer texts tend to converge toward the center of all embeddings.

Furthermore, to provide a more intuitive explanation of the differences among embeddings of varying lengths, we calculated the average Euclidean distance from the embeddings in different length intervals to the central embedding on the NFCorpus and SciFact datasets.

NFCorpus

SciFact

The results in the table indicate that as text length increases, embeddings are positioned closer to the central embedding, meaning that their distribution in space becomes increasingly centered around the overall embedding distribution. This leads to the Length Collapse.

Dear Reviewer j2YY,

We sincerely appreciate your valuable feedback on our manuscript. In response to your suggestions, we have enhanced the revised manuscript with additional experiments and analyses.

As the rebuttal period concludes in less than two days, we hope our efforts align with your expectations. If you find our response satisfactory, we would be grateful if you could consider revising the score.

Thank you once again for your insightful guidance.

Warm regards,

Authors