Pooling And Attention: What Are Effective Designs For LLM-Based Embedding Models?

Thank you for highlighting the need for deeper analysis of task-specific performance variations. We have conducted additional experiments that provide important insights. Below we address the main concerns:

1.Explanation About STS/Retrieval vs Classification/Clustering

Inspired by reviewers' questions about task-specific performance variations, we conducted a controlled experiment to investigate how different training data compositions affect model performance.

This experiment confirmed our findings in the paper, which suggested that the optimization objectives of STS/Retrieval and Classification/Clustering may not be consistent.

We compared two settings:

• Baseline: 68,039 samples (random sampled from the original dataset)
• Enhanced: 136,078 samples (68,039 retrieval + 68,039 classification/clustering tasks)
• All training data has no overlap with testing data.

This suggests that:

1. Different tasks have potentially conflicting optimization objectives
2. Performance trade-offs are inherent, not architectural limitations
3. Task-specific training data is crucial for optimal performance

2. Additional Datasets

We appreciate your feedback on the robustness of work. We evaluated our original models in the paper on additional datasets to ensure their generalizability. These additional datasets further confirm our findings, demonstrating the robustness and generalizability of our method.

The results are summarized below:

STS Datasets: FINAL (Ju et al., 2023), FinSTS (Liu et al., 2024)

Retrieval Datasets: FiQA2018 (FiQA, 2018), FinanceBench (Islam et al.,2023), HC3 (Guo et al., 2023), Apple10KRetrieval (Tang et al., 2024), FinQA (Chen et al., 2021)

Classification Datasets: FinancialPhrasebank (Malo et al., 2014), FinSent (Yang et al., 2023), FiQA, SemEva2017 Headline (Cortis et al., 2017), FLS (Yang et al., 2023), ESG (Yang et al., 2023), FOMC (Shah et al., 2023)

Clustering Datasets: MInDS-14-zh (Gerz et al., 2021), WikiCompany2Industry (Tang et al., 2024)

3. NV-embed Style Model

We would like to clarify that Model 4 in our paper serves as an NV-Embedded-style trainable last-layer pooling with a slight difference: While NV-Embed transforms the last layer's hidden states to the query matrix Q, our method transforms it into K matrix and V matrix for apple-to-apple comparison with Model 5 (Add trainable layer weights).

For the rigorous evaluation, we also implemented a NV-Embed model using the dataset from question 1. The results are illustrated in the table:

The performance of the NV-Embed model is also consistent with the findings in the paper and did not surpass the model 5 (with layer weights).

Thank you again for your valuable feedback. If you have any other questions or require more clarification, please do not hesitate to let us know!

Thank you for your thorough review. We appreciate your recognition of our work's strengths and your constructive feedback. Let us address each of your concerns:

1. Missing Training Data of Classification/Clustering Tasks

Regarding our training dataset selection, we initially aligned with current standard training pipelines (such as LLM2Vec [1]) to ensure fair comparison. However, your question raised an important point about potential task-specific biases.

Inspired by your suggestion and other reviewers' feedback, we conducted a new controlled experiment to investigate how training data composition affects different downstream tasks.

Key findings show that while balanced training data improved clustering performance, it led to slight decreases in retrieval/STS tasks. This trade-off suggests different tasks may indeed require different optimization objectives, aligning with our paper's findings.

We compared two settings:

• Baseline: 68,039 samples (random sampled from the original dataset)
• Enhanced: 136,078 samples (68,039 retrieval + 68,039 classification/clustering tasks)
• All training data has no overlap with testing data.

2. Task Dependent Trainable Q

Following your suggestion, we implemented a task-dependent Q by concatenating task instruction embeddings with the trainable Q matrix. The experimental setup remained consistent with Model 5, trained on the enhanced dataset. Results below show that this approach led to performance degradation across all metrics:

While this simple instruction-based adaptation did not yield improvements, we believe your suggestion opens up promising future directions! We will explore more approaches in this field.

3. Compare with NV-Embed Style Model

We would like to clarify that Model 4 in our paper serves as an NV-Embedded-style trainable last-layer pooling with a slight difference: While NV-Embed transforms the last layer's hidden states to the query matrix Q, our method transforms it into K matrix and V matrix for apple-to-apple comparison with Model 5 (Add trainable layer weights).

For the rigorous evaluation, we also implemented a complete NV-Embed model using the dataset from question 1. The results are illustrated in the table:

The performance of the NV-Embed model is also consistent with the findings in the paper and did not surpass the model 5 (with layer weights).

We deeply appreciate the time and effort you have dedicated to reviewing our work.

• Regarding the statement:

Instead, layer weights are added to the hidden layer features to incorporate positional information into the input features.

The positional information here refers to the weights at different layer indices, which is consistent with our statement that "we introduce a trainable layer weights matrix that captures the significance of each layer." The purpose is to find the optimal combination of layers in a trainable way, as the last layer alone may not be optimal for semantic tasks.

• Regarding the Classification/Clustering Tasks:

  • Our use of label_text as "documents" follows the experimental setting established in NV-Embed.
  • Since we use gold-standard labels and treat other labels from the same dataset as negative samples, the question of "whether items could be categorized into other classes" is not applicable in this context.

While we remain confident in our research findings and results, we sincerely appreciate your thoughtful suggestions for deeper exploration. We are excited to continue our investigation in this direction.

Thank you again for your valuable feedback. We have learned a lot from your insights. Wishing you all the best!

We appreciate your recognition of the value of our work in the multi-layer trainable pooling and thoughtful feedback! Below, we address each of your concerns:

1. Full MTEB Evaluation

The primary objective of our paper is to empirically establish best practices for LLM-based embedding models, rather than to achieve state-of-the-art performance. This focus explains why we didn't directly compare our models with other models in the MTEB benchmark. Instead, we concentrated on four fundamental tasks (STS, retrieval, classification, and clustering) that are most frequently used in real-world applications.

However, for completeness and to address your question, we present the full MTEB evaluation results in the table below:

We will also include this table in the Appendix.

2. STS/Retrieval vs Classification/Clustering Tasks

Regarding our training dataset selection, we initially aligned with current standard training pipelines (such as LLM2Vec(BehnamGhader et al., 2024)) to ensure fair comparison. However, your question raised an important point about potential task-specific biases.

Inspired by your suggestion and other reviewers' feedback, we conducted a new controlled experiment to investigate how training data composition affects different downstream tasks.

Key Findings: This experiment confirmed our findings in the paper, which suggested that the optimization objectives of STS/Retrieval and Classification/Clustering may not be consistent. This explains why a model optimized for one type of task might not naturally excel at the other.

We compared two settings:

• Baseline: 68,039 samples (random sampled from the original dataset)
• Enhanced: 136,078 samples (68,039 retrieval + 68,039 classification/clustering tasks)
• All training data has no overlap with testing data.

We can find:

• Task Balance Impact: Models trained on the enhanced dataset showed improved performance in clustering tasks, supporting your hypothesis about the importance of diverse training data.

• Performance Trade-offs: The observed performance decrease in STS/Retrieval tasks after adding classification/clustering data provides important insights. This suggests that similarity-based tasks and category-based tasks indeed have different, potentially inconsistent optimization objectives, which is consistent with our findings in paper.

To conclude, your suggestions also provide valuable guidance in selecting training data based on their specific application needs. We plan to incorporate these insights into our final paper and remain available for further discussion if you have any additional questions or require more clarification.

Thank you to the author for conducting additional ablation studies. The *-enhanced model's average scores now align more closely with the ideas of the paper across various embedding benchmark tasks. However, the *-enhanced models perform worse than the baseline retrieval and classification scores. I believe these sub-optimal results can be addressed through improvements in training strategies, dataset blends, hyperparameter tuning, etc. I increase the score of the paper.

Thank you for your thoughtful evaluation and constructive feedback on our manuscript. We appreciate the opportunity to clarify our contributions and address your concerns.

1. Clarification on Contributions&Novelty

• Existing Challenge:

Given that most existing LLM-based embedding models are trained using different datasets with different base models, it is difficult to draw conclusions regarding the contribution of different model architecture design choices (as shown in Table 1).

• Primary Objective:

Thus, the primary objective of our paper is to empirically establish best practices for LLM-based embedding models, rather than to achieve state-of-the-art performance. Then, provide insights and guidance on choosing pooling and attention strategies for LLM-based embedding models.

- Methodology
  - Conduct controlled experiments using identical base models, training data, and hyperparameters.
  - Focus on five widely-used model architecture combinations.
  - Use empirical experiments to ensure that findings are statistically meaningful rather than due to random variation.

• Secondary Objective

To better enhance the pooling strategy, we introduced the implementation of multi-layer pooling in LLM-based embedding via cross-attention. This is based on our observation that the last layer hidden state is not always the most semantically relevant layer.

To our knowledge, this is the first method using multi-layer information for LLM-based embedding models.

2. Regarding "Negligible Improvements"

We respectfully address the concern about "negligible improvements" with two key points:

• Context of Existing Works: As shown in Table, improvements across recent embedding models are typically modest. Even state-of-the-art models achieve only incremental gains, that's why we want to find an optimal model setting using controlled/empirical experiments.

• Statistical Validation: Unlike many previous works, we provide statistical significance tests for all reported improvements, ensuring that our gains, though modest, are reliable and statistically meaningful rather than due to random variation.

3.Explanation About STS/Retrieval vs Classification/Clustering

Inspired by reviewers' questions about task-specific performance variations, we conducted a new controlled experiment to investigate how different training data compositions affect model performance.

This experiment confirmed our findings in the paper, which suggested that the optimization objectives of STS/Retrieval and Classification/Clustering may not be consistent. This explains why a model optimized for one type of task might not naturally excel at the other.

• Methodology

We compared two settings:

1. Baseline: 68,039 samples (random sampled from the original dataset).

2. Enhanced: 136,078 samples (68,039 retrieval + 68,039 classification/clustering tasks)

Our findings reveal that:

1. With retrieval-only training data (Baseline), models with trainable layers (Model 2-5) significantly outperform EOS-token pooling in similarity-based tasks (STS/Retrieval), but show comparable or slightly worse performance in classification/clustering tasks.

2. When adding classification/clustering training data (Enhanced), we observe:

   • Performance drops in STS/Retrieval tasks
   • Improved performance in clustering tasks
   • Relatively stable classification performance

Once again, we sincerely thank you for your valuable feedback and remain available for further discussion if you have any additional questions or require more clarification.