Decoder-Only LLMs can be Masked Auto-Encoders
Authors: Dan Qiao, Yuan Gao, Zheming Yang, Di Yang, Ziheng Wu, Pengcheng Lu, Minghui Qiu, Juntao Li, Min Zhang
Accepted at: acl25

Decoder-Only LLMs can be Masked Auto-Encoders

Dan Qiao1,2, Yuan Gao2 Zheming Yang2, Di Yang2, Ziheng Wu2, Pengcheng Lu2, Minghui Qiu2, Juntao Li1* Min Zhang1
1School of Computer Science and Technology, Soochow University
2ByteDance Inc
qiaodan.23@bytedance.com

Abstract
Modern NLP workflows (e.g., RAG systems) require different models for generation and embedding tasks, where bidirectional pre-trained encoders and decoder-only Large Language Models (LLMs) dominate respective tasks. Structural differences between models result in extra development costs and limit knowledge sharing between tasks. In this work, we present UniMAE, a novel unsupervised training method that transforms a Decoder-Only LLM into a Uni-Directional Masked Auto-Encoder. UniMAE compresses high-quality semantic information into the [EOS] embedding while preserving the generation capabilities of LLMs. Comprehensive evaluations across 56 MTEB datasets demonstrate that UniMAE can achieve state-of-the-art results under unsupervised settings with merely 100 training steps, establishing the first effective approach to unifying generation and representation learning in decoder-only architectures.

1 Introduction
Pre-trained language models have been widely applied in various scenarios (Devlin, 2018; Ouyang et al., 2022; Brown et al., 2020; Wu et al., 2025; ?). However, the workflows of real-world applications like RAG (Fan et al., 2024) or retrieval systems (Zhu et al., 2023) often involve collaboration between generation tasks and embedding tasks. For example, RAG requires retrieval followed by generation (Asai et al., 2023), while retrieval systems often need to rewrite queries before performing the search (Ma et al., 2023; Liu and Mozafari, 2024).

For embedding tasks, encoder models with bidirectional attention mechanisms have been the mainstream choice (Devlin, 2018; Liu, 2019). More recently, decoder-only large language models (LLMs) with unidirectional attention have demonstrated the ability to compress knowledge from trillions of tokens during pre-training (Radford, 2018; Ouyang et al., 2022; Guo et al., 2025), thus have been widely utilized in various generative tasks. This structural separation leads to extra training and deployment costs and impedes knowledge sharing between tasks (Asai et al., 2023). In the long run, enabling large language models (LLMs) to perform both generation and embedding tasks simultaneously holds significant promise.

In this paper, we introduce UniMAE, an unsupervised method that compresses the semantic content of input text into the [EOS] (end-of-sequence) embedding of the last transformer layer, while preserving the capabilities of original LLMs (Figure 1). We scale our method across models with 1B, 3B, and 8B parameters and conduct evaluations on 56 datasets of the Massive Text Embeddings Benchmark (MTEB) (Muenchhoff et al., 2022). UniMAE achieves state-of-the-art performance on MTEB under an unsupervised setting within only 100 training steps. Evaluations on language modeling tasks demonstrate that UniMAE not only enhances representation capabilities but also maintains the generation abilities of LLMs. Domain post-training experiments indicate that UniMAE can simultaneously improve the generation and embedding performance in a single training process.

![Overview of UniMAE training. The left part shows the Masked Auto-Regressive process, while the right part shows the reconstruction process using a tiny decoder with [EOS] embeddings to recover the input.](page_728_370_670_246.png)

Figure 1: Overview of UniMAE training. The left part shows the Masked Auto-Regressive process, while the right part shows the reconstruction process using a tiny decoder with [EOS] embeddings to recover the input.
To our knowledge, UniMAE is the first effective approach to unifying generation and representation learning in decoder-only architectures.

2 Method

2.1 Motivation

Unlike pre-trained encoder models that utilize special tokens like [CLS] for sentence embeddings, LLMs primarily focus on predicting the next token during pre-training and lack a dedicated position for sentence representation. As a result, neither averaging the token embeddings nor directly using the [EOS] embedding can effectively represent the entire sentence. (Springer et al., 2024; Jiang et al., 2023; BehnamGhader et al., 2024).

To bridge this gap, recent studies have focused on enhancing the representational capabilities of large models. Many directly use the model as an encoder, employing average embeddings or token-level embeddings for supervised contrastive fine-tuning (Lee et al., 2024; Li et al., 2024). Springer et al. (2024) and BehnamGhader et al. (2024) argue that unidirectional attention prevents each token from incorporating information from subsequent tokens, thereby limiting the model’s contextual understanding. Springer et al. (2024) addresses this limitation by repeating the sentence twice. Meanwhile, BehnamGhader et al. (2024) applies a Bidirectional Attention Mechanism and adapts the MLM task to LLM, then finally uses average pooling to get sentence embeddings.

However, these methods still fail to address the lack of a unified pooling mechanism in LLMs. Moreover, they simply repurpose LLMs as encoders during fine-tuning, thereby overlooking the potential loss of their generative capabilities.

2.2 Overall framework

Similar to Visual-MAE (He et al., 2022), UniMAE employs the [EOS] embedding as the latent representation of the corrupted input and initializes a small decoder to reconstruct the original input based on this representation. Unlike other approaches, our method maintains a unidirectional attention mechanism characteristic of large language models (LLMs). We illustrate the entire training process in Figure 1.

The overall objective consists of two components: 1. Masked Auto-Regressive (MAR) and 2. Masked Re-Construct (MRC). The complete training process involves jointly optimizing two sets of parameters: the LLM parameters \( \Phi_{llm} \) and small decoder parameters \( \Phi_{dec} \).

2.3 Masked Auto-Regressive (MAR)

The sole distinction between the classic Auto-Regressive (Radford, 2018; Radford et al., 2019; Brown et al., 2020) and MAR is that the input is corrupted with random masks of ratio \( p_{mar} \). For a sentence \( X \), it is corrupted by replacing some tokens with mask tokens, resulting in the masked sentence \( \tilde{X} \). The MAR loss can be expressed as:

\[
L_{MAR} = \sum_t CE(x_t \mid \tilde{x}_{i < t}; \Phi_{llm}),
\]

where \( x_i \in X \) is the i-th token or the original text, \( \tilde{x}_i \in \tilde{X} \) is the i-th token of the corrupted text. \( \Phi_{llm} \) denotes the parameters of the original LLM.

2.4 Masked Re-Construct (MRC)

We believe that the representation at [EOS] position has the potential to encapsulate the semantic meaning of the entire input. Since only the last token can attend to all the input tokens under the unidirectional attention mechanism. So the latent representation \( h_{eos} \) is:

\[
h_{eos} \leftarrow \Phi_{llm} \left( \tilde{X} \right)
\]

The Masked Re-Construct (MRC) process utilizes a newly initialized small decoder, which employs the [EOS] embedding combined with partial input information to reconstruct the original sentence. The input of the tiny decoder can be formulated as:

\[
H_1 \leftarrow [h_{eos} + p_0, \ldots, h_{eos} + p_N],
\]
\[
H_2 \leftarrow [h_{eos}, e_{x_1} + p_1, \ldots, e_{x_N} + p_N],
\]

where \( p_i \) is the trainable positional embedding, \( e_i \) is the input token embedding. The only difference between MRC and traditional cross-attention is that we applied a mask to the attention matrix, forcing the model to rely more on the \( h_{eos} \) for sentence recovery. We set a mask ratio \( p_{mrc} \) and let \( B_{ij} \sim \text{Bernoulli}(p_{mrc}) \in \{0, 1\} \). Then the computation process of attention can be expressed with the following formula:

\[
Q = H_1 W^Q, K = H_2 W^K, V = H_2 W^V
\]
\[
M_{ij} = \begin{cases}
0, & \text{if } i \neq j \land B_{ij} = 1 \\
-\infty, & \text{else}
\end{cases},
\]
\[
A = \text{softmax} \left( \frac{Q^T K}{\sqrt{d}} + M \right) V.
\]
<table>
  <tr>
    <th rowspan="2">Attn.</th>
    <th rowspan="2">Categories → # of datasets →</th>
    <th colspan="1">Retr.</th>
    <th colspan="1">Rerank.</th>
    <th colspan="1">Clust.</th>
    <th colspan="1">PairClass.</th>
    <th colspan="1">Class.</th>
    <th colspan="1">STS</th>
    <th colspan="1">Summ.</th>
    <th colspan="1">Avg</th>
  </tr>
  <tr></tr>
  <tr>
    <td colspan="10">LLaMA3.2-1B</td>
  </tr>
  <tr>
    <td>Uni.</td>
    <td>MEAN</td>
    <td>15.13</td>
    <td>42.96</td>
    <td>34.30</td>
    <td>53.50</td>
    <td>56.20</td>
    <td>56.27</td>
    <td>28.06</td>
    <td>39.32</td>
  </tr>
  <tr>
    <td>Uni.</td>
    <td>ECHO</td>
    <td>24.57</td>
    <td>48.27</td>
    <td>36.68</td>
    <td>67.65</td>
    <td>65.60</td>
    <td>71.86</td>
    <td>29.79</td>
    <td>48.28</td>
  </tr>
  <tr>
    <td>Bi.</td>
    <td>MNTP</td>
    <td>18.33</td>
    <td>43.54</td>
    <td>34.20</td>
    <td>56.30</td>
    <td>60.24</td>
    <td>61.15</td>
    <td>27.21</td>
    <td>42.12</td>
  </tr>
  <tr>
    <td>Uni.</td>
    <td>UniMAE</td>
    <td><b>34.48</b></td>
    <td>50.31</td>
    <td>38.42</td>
    <td>79.91</td>
    <td>68.88</td>
    <td>73.14</td>
    <td>29.93</td>
    <td>52.48</td>
  </tr>
  <tr>
    <td>Bi.</td>
    <td>MNTP + SimCSE</td>
    <td>30.39</td>
    <td>49.67</td>
    <td>37.48</td>
    <td>74.41</td>
    <td>63.70</td>
    <td>72.83</td>
    <td><b>30.71</b></td>
    <td>50.73</td>
  </tr>
  <tr>
    <td>Uni.</td>
    <td>UniMAE + SimCSE</td>
    <td><u>31.54</u></td>
    <td>54.23</td>
    <td><u>43.22</u></td>
    <td><u>78.04</u></td>
    <td><u>69.26</u></td>
    <td><u>76.95</u></td>
    <td>30.14</td>
    <td><b>54.11</b></td>
  </tr>
  <tr>
    <td colspan="10">LLaMA3.2-3B</td>
  </tr>
  <tr>
    <td>Uni.</td>
    <td>MEAN</td>
    <td>14.61</td>
    <td>43.22</td>
    <td>34.73</td>
    <td>55.50</td>
    <td>55.55</td>
    <td>56.00</td>
    <td>22.13</td>
    <td>39.09</td>
  </tr>
  <tr>
    <td>Uni.</td>
    <td>ECHO</td>
    <td>25.44</td>
    <td>48.64</td>
    <td>36.86</td>
    <td>70.68</td>
    <td>67.96</td>
    <td>71.95</td>
    <td>24.31</td>
    <td>49.16</td>
  </tr>
  <tr>
    <td>Bi.</td>
    <td>MNTP</td>
    <td>17.75</td>
    <td>44.69</td>
    <td>36.06</td>
    <td>57.47</td>
    <td>59.79</td>
    <td>61.87</td>
    <td>23.06</td>
    <td>42.38</td>
  </tr>
  <tr>
    <td>Uni.</td>
    <td>UniMAE</td>
    <td>34.25</td>
    <td>51.79</td>
    <td>38.95</td>
    <td>78.31</td>
    <td>70.14</td>
    <td>73.06</td>
    <td>29.25</td>
    <td><u>53.30</u></td>
  </tr>
  <tr>
    <td>Bi.</td>
    <td>MNTP + SimCSE</td>
    <td>34.35</td>
    <td>50.91</td>
    <td>39.83</td>
    <td>76.21</td>
    <td>66.00</td>
    <td>73.77</td>
    <td><u>30.64</u></td>
    <td>52.61</td>
  </tr>
  <tr>
    <td>Uni.</td>
    <td>UniMAE + SimCSE</td>
    <td><b>36.09</b></td>
    <td><u>55.41</u></td>
    <td><u>43.77</u></td>
    <td><u>80.90</u></td>
    <td><u>71.60</u></td>
    <td><u>76.53</u></td>
    <td>30.33</td>
    <td><b>56.11</b></td>
  </tr>
  <tr>
    <td colspan="10">LLaMA3.1-8B</td>
  </tr>
  <tr>
    <td>Uni.</td>
    <td>MEAN</td>
    <td>14.26</td>
    <td>46.89</td>
    <td>37.84</td>
    <td>70.12</td>
    <td>66.73</td>
    <td>62.32</td>
    <td>24.98</td>
    <td>42.95</td>
  </tr>
  <tr>
    <td>Uni.</td>
    <td>ECHO</td>
    <td>25.53</td>
    <td>49.75</td>
    <td>37.77</td>
    <td>70.11</td>
    <td>67.03</td>
    <td>71.83</td>
    <td>27.32</td>
    <td>49.24</td>
  </tr>
  <tr>
    <td>Bi.</td>
    <td>MNTP</td>
    <td>22.83</td>
    <td>48.46</td>
    <td>40.96</td>
    <td>64.26</td>
    <td>60.11</td>
    <td>64.50</td>
    <td>27.40</td>
    <td>45.95</td>
  </tr>
  <tr>
    <td>Uni.</td>
    <td>UniMAE</td>
    <td>35.37</td>
    <td>52.74</td>
    <td>41.25</td>
    <td>78.23</td>
    <td>68.49</td>
    <td>73.45</td>
    <td>30.40</td>
    <td>53.87</td>
  </tr>
  <tr>
    <td>Bi.</td>
    <td>MNTP + SimCSE</td>
    <td><b>39.75</b></td>
    <td>53.59</td>
    <td>42.70</td>
    <td>78.71</td>
    <td>67.38</td>
    <td><u>76.84</u></td>
    <td><u>31.77</u></td>
    <td>55.81</td>
  </tr>
  <tr>
    <td>Uni.</td>
    <td>UniMAE + SimCSE</td>
    <td><u>38.51</u></td>
    <td><b>56.91</b></td>
    <td><u>46.06</u></td>
    <td><u>80.15</u></td>
    <td><u>68.77</u></td>
    <td><u>76.21</u></td>
    <td>29.80</td>
    <td><b>56.60</b></td>
  </tr>
</table>

Table 1: Results on 56 MTEB datasets with best results highlighted in bold, and the second-best results underlined.

The MRC object can be formalized as:

\[
\mathcal{L}_{MRC} = \sum_t CE(x_t \mid A, H_1, H_2, \Phi_{dec}).
\] (5)

The joint optimization of MAR and MRC is expressed as follows:

\[
\mathcal{L}_{UniMAE} = \alpha \mathcal{L}_{MAR} + \mathcal{L}_{MRC},
\] (6)

where \( \alpha \) is the weight to control the MAR objective.

3 Experiment

3.1 Settings

For embedding tasks, we conduct evaluations on the Massive Text Embedding Benchmark (MTEB) (Muennighoff et al., 2022), a collection of 7 diverse embedding task categories covering a total of 56 datasets. All evaluations are done using the official MTEB code repository. For preliminary experiments and ablations, we use MTEB-15 a subset including 15 representative tasks. Details are attached in Appendix A.

For baselines, we compare our method with average pooling, Echo (Springer et al., 2024), which repeats the sentence twice and averages the token embedding of the second sentence, and LLM2Vec (BehnamGhader et al., 2024), which sequentially employs MNTP and SimCSE (Gao et al., 2021b) to train the model. To make a comparison with LLM2Vec with SimCSE training, we also provide UniMAE with further SimCSE training. We provide training details in Appendix C.

3.2 Embedding Performance

Table 1 shows that, after UniMAE + SimCSE training, we observe a considerable improvement in performance for all three models compared with direct mean pooling, which is 28% for the 1B model, 43% for the 3B model and 32% for the 8B model. Using only UniMAE can significantly surpass all single unsupervised learning methods, UniMAE even outperforms MNTP + SimCSE training on 1B and 3B models. Overall, the UniMAE framework achieves the state-of-the-art results across models of all scales. Scores for each dataset are shown in Appendix G.

3.3 Will UniMAE hurt LLMs?

We selected 10 commonly used datasets for evaluating generative language models, including ARC (Clark et al., 2018), HellaSwag (Zellers et al., 2019) and MMLU (Hendrycks et al., 2021). Details can be found in Appendix E. As shown in Figure 3b, after UniMAE training, the generative capabilities of the model are nearly equivalent to those of the original foundation models. We also tested the baseline method, MNTP training, which modifies causal attention to bidirectional attention. We observe that the structural modification significantly
reduces the model’s performance on language modeling tasks. The same situation occurs in supervised / unsupervised contrastive learning applied to the representations after average pooling.

3.4 UniMAE optimizes Vector Space

We apply t-SNE transformations to the top 6 classes’ sentence embeddings from BiorxivClusteringS2S dataset, using PCA initialization, 2000 iterations, and early exaggeration of 20.

Figure 2 shows the distribution of the UniMAE and baseline embeddings. UniMAE improves spatial distribution by bringing samples closer within each group using [EOS] token embeddings. This improved distribution highlights the effectiveness of our approach in boosting LLM performance as an encoder.

![t-SNE visualization of sentence embeddings from top 6 classes on BiorxivClusteringS2S testset.](page_186_573_1077_246.png)

(a) Base Model (b) UniMAE

Figure 2: t-SNE visualization of sentence embeddings from top 6 classes on BiorxivClusteringS2S testset.

3.5 Is MAR necessary?

We designed MAR to prevent the model from learning the shortcut of simply memorizing all the input tokens during training. Therefore, we introduce noise to the input, allowing the model to extract semantic representations even from incomplete inputs, rather than simply memorizing them. To prove that MAR is superior to AR, we trained different models using input masking ranging from 0% to 80% and performed inference on MTEB-15. As shown in Figure 3a, for models of different sizes, the optimal mask ratio falls between 40% and 60%. For convenience, we set this ratio to 50% for all models. In the detailed results, we observe that input masking significantly improved performance on tasks related to semantic similarity, such as Clustering and STS, while the enhancement for retrieval tasks was relatively minor.

3.6 Domain Post-Training Result

We believe that UniMAE can not only enhance the representation ability of LLMs but also serve as a novel pre-training and post-training method to simultaneously improve both generative and representational capabilities. We conduct E-Commerce domain post-training for 100B tokens, we describe more details in the Appendix F. The evaluation contains three types of tasks: domain language modeling, retrieval, and classification. Language model tasks (LM.) are primarily QA, retrieval (Retr.), and classification (Class.) tasks that are similar to those in MTEB. Table 2 demonstrates that UniMAE training slightly outperforms Auto-Regressive in generation tasks, and significantly outperforms Auto-Regressive in embedding tasks.

<table>
  <tr>
    <th>Categories → # of sub-tasks →</th>
    <th>LM.</th>
    <th>Retr.</th>
    <th>Class.</th>
  </tr>
  <tr>
    <td>Qwen2-0.5B</td>
    <td>71.12</td>
    <td>36.91</td>
    <td>81.43</td>
  </tr>
  <tr>
    <td>Qwen2-0.5B + AR</td>
    <td>73.16</td>
    <td>35.89</td>
    <td>83.98</td>
  </tr>
  <tr>
    <td>Qwen2-0.5B + UniMAE</td>
    <td><b>74.55</b></td>
    <td><b>50.84</b></td>
    <td><b>95.38</b></td>
  </tr>
</table>

Table 2: Model performance after domain post-training

![Results under different MAR mask ratio on MTEB-15. (b) Performance on language model tasks.](page_1012_120_410_246.png)

(a) Mask Ratio Search (b) LM Evaluation

Figure 3: (a) Results under different MAR mask ratio on MTEB-15. (b) Performance on language model tasks.

4 Conclusion

We introduce UniMAE, a method capable of training models for both generation and embedding tasks. We evaluate UniMAE on diverse datasets, including open-source and domain-specific corpora, across both types of tasks. UniMAE achieves state-of-the-art performance on the MTEB benchmark without compromising the language modeling capabilities of LLMs. As a domain post-training approach, UniMAE enhances both the representational and generative capacities of models simultaneously. In future work, we plan to apply UniMAE for large-scale pre-training from scratch and instruction alignment, aiming to develop a new generation of LLMs that excel in both generation and embedding tasks.
5 Limitations

There is an inherent limitation: we initialized a small decoder from scratch to reconstruct the original input. However, during the inference process, we directly discarded this small decoder and instead utilized the results generated by the LLM for both generation and representation tasks. This approach leads to a certain degree of parameter waste. If this method is scaled up for large-scale zero pre-training, such waste would be suboptimal. In the future, we will explore new structural designs to address this issue.

Acknowledgments

We want to thank all the anonymous reviewers for their valuable comments. This work was supported by the National Science Foundation of China (NSFC No. 62206194), the Natural Science Foundation of Jiangsu Province, China (Grant No. BK20220488), the Young Elite Scientists Sponsorship Program by CAST (2023QNRC001), and the Priority Academic Program Development of Jiangsu Higher Education Institutions.

References

Akari Asai, Zeqiu Wu, Yizhong Wang, Avirup Sil, and Hannaneh Hajishirzi. 2023. Self-rag: Learning to retrieve, generate, and critique through self-reflection. arXiv preprint arXiv:2310.11511.

Parishad BehnamGhader, Vaibhav Adlakha, Marius Mosbach, Dzmitry Bahdanau, Nicolas Chapados, and Siva Reddy. 2024. Llm2vec: Large language models are secretly powerful text encoders. arXiv preprint arXiv:2404.05961.

Yonatan Bisk, Rowan Zellers, Jianfeng Gao, Yejin Choi, et al. 2020. Piqa: Reasoning about physical commonsense in natural language. In Proceedings of the AAAI conference on artificial intelligence, volume 34, pages 7432–7439.

Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. 2020. Language models are few-shot learners. Advances in neural information processing systems, 33:1877–1901.

Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and Oyvind Tafjord. 2018. Think you have solved question answering? try arc, the ai2 reasoning challenge. arXiv preprint arXiv:1803.05457.

Jacob Devlin. 2018. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.04805.

Wenqi Fan, Yujuan Ding, Liangbo Ning, Shijie Wang, Hengyun Li, Dawei Yin, Tat-Seng Chua, and Qing Li. 2024. A survey on rag meeting llms: Towards retrieval-augmented large language models. In Proceedings of the 30th ACM SIGKDD Conference on Knowledge Discovery and Data Mining, pages 6491–6501.

Leo Gao, Jonathan Tow, Stella Biderman, Sid Black, Anthony DiPofi, Charles Foster, Laurence Golding, Jeffrey Hsu, Kyle McDonell, Niklas Muennighoff, Jason Phang, Laria Reynolds, Eric Tang, Anish Thite, Ben Wang, Kevin Wang, and Andy Zou. 2021a. A framework for few-shot language model evaluation.

Tianyu Gao, Xingcheng Yao, and Danqi Chen. 2021b. Simcse: Simple contrastive learning of sentence embeddings. arXiv preprint arXiv:2104.08821.

Daya Guo, Dejian Yang, Haowei Zhang, Junxiao Song, Ruoyu Zhang, Runxin Xu, Qihao Zhu, Shirong Ma, Peiyi Wang, Xiao Bi, et al. 2025. Deepseek-rl: Incentivizing reasoning capability in llms via reinforcement learning. arXiv preprint arXiv:2501.12948.

Faegheh Hasibi, Fedor Nikolaev, Chenyan Xiong, Krisztian Balog, Svein Erik Bratsberg, Alexander Kotov, and Jamie Callan. 2017. Dbpedia-entity v2: A test collection for entity search. In Proceedings of the 40th International ACM SIGIR Conference on Research and Development in Information Retrieval, SIGIR ’17, pages 1265–1268. ACM.

Kaiming He, Xinlei Chen, Saining Xie, Yanghao Li, Piotr Dollár, and Ross Girshick. 2022. Masked autoencoders are scalable vision learners. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 16000–16009.

Dan Hendrycks, Collin Burns, Steven Basart, Andy Zou, Mantas Mazeika, Dawn Song, and Jacob Steinhardt. 2021. Measuring massive multitask language understanding. In International Conference on Learning Representations.

Edward J Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen. 2021. Lora: Low-rank adaptation of large language models. arXiv preprint arXiv:2106.09685.

Ting Jiang, Shaohan Huang, Zhongzhi Luan, Deqing Wang, and Fuzhen Zhuang. 2023. Scaling sentence embeddings with large language models. arXiv preprint arXiv:2307.16645.

Tom Kwiatkowski, Jennimaria Palomaki, Olivia Redfield, Michael Collins, Ankur Parikh, Chris Alberti, Danielle Epstein, Illia Polosukhin, Jacob Devlin, Kenton Lee, Kristina Toutanova, Llion Jones, Matthew Kelcey, Ming-Wei Chang, Andrew M. Dai, Jakob Uszkoreit, Quoc Le, and Slav Petrov. 2019. Natural questions: A benchmark for question answering research. Transactions of the Association for Computational Linguistics, 7:452–466.
Chankyu Lee, Rajarshi Roy, Mengyao Xu, Jonathan Raiman, Mohammad Shoeybi, Bryan Catanzaro, and Wei Ping. 2024. Nv-embed: Improved techniques for training llms as generalist embedding models. arXiv preprint arXiv:2405.17428.

Chaofan Li, MingHao Qin, Shitao Xiao, Jianlyu Chen, Kun Luo, Yingxia Shao, Defu Lian, and Zheng Liu. 2024. Making text embedders few-shot learners. arXiv preprint arXiv:2409.15700.

Jian Liu, Leyang Cui, Hannmeng Liu, Dandan Huang, Yile Wang, and Yue Zhang. 2020. Logica: A challenge dataset for machine reading comprehension with logical reasoning. In Proceedings of the Twenty-Ninth International Joint Conference on Artificial Intelligence, IJCAI-20, pages 3622–3628.

Jie Liu and Barzan Mozafari. 2024. Query rewriting via large language models. arXiv preprint arXiv:2403.09060.

Yinhan Liu. 2019. Roberta: A robustly optimized bert pretraining approach. arXiv preprint arXiv:1907.11692, 364.

Xinbei Ma, Yeyun Gong, Pengcheng He, Hai Zhao, and Nan Duan. 2023. Query rewriting for retrieval-augmented large language models. arXiv preprint arXiv:2305.14283.

Niklas Muenninghoff, Nouamane Tazi, Loic Magne, and Nils Reimers. 2022. Mteb: Massive text embedding benchmark. arXiv preprint arXiv:2210.07316.

Tri Nguyen, Mir Rosenberg, Xia Song, Jianfeng Gao, Saurabh Tiwary, Rangan Majumder, and Li Deng. 2016. Ms marco: A human-generated machine reading comprehension dataset.

Long Ouyang, Jeffrey Wu, Xu Jiang, Diogo Almeida, Carroll Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, et al. 2022. Training language models to follow instructions with human feedback. Advances in neural information processing systems, 35:27730–27744.

Denis Paperno, Germán Kruszewski, Angeliki Lazari-dou, Ngoc Quan Pham, Raffaella Bernardi, Sandro Pezzelle, Marco Baroni, Gemma Boleda, and Raquel Fernández. 2016. The LAMBADA dataset: Word prediction requiring a broad discourse context. In Proceedings of the 54th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pages 1525–1534.

Alec Radford. 2018. Improving language understanding by generative pre-training.

Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. 2019. Language models are unsupervised multitask learners. OpenAI blog, 1(8):9.

Keisuke Sakaguchi, Ronan Le Bras, Chandra Bhagavatula, and Yejin Choi. 2021. Winogrande: An adversarial winograd schema challenge at scale. Communications of the ACM, 64(9):99–106.

Jacob Mitchell Springer, Suhas Kotha, Daniel Fried, Graham Neubig, and Aditi Raghunathan. 2024. Repetition improves language model embeddings. arXiv preprint arXiv:2402.15449.

Johannes Welbl, Nelson F. Liu, and Matt Gardner. 2017. Crowdsourcing multiple choice science questions. In Proceedings of the 3rd Workshop on Noisy User-generated Text, pages 94–106.

Ziheng Wu, Zhenghao Chen, Ruipu Luo, Can Zhang, Yuan Gao, Zhentao He, Xian Wang, Haoran Lin, and Minghui Qiu. 2025. Valley2: Exploring multimodal models with scalable vision-language design. arXiv preprint arXiv:2501.05901.

Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. 2019. HellaSwag: Can a machine really finish your sentence? In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, pages 4791–4800.

Yutao Zhu, Huaying Yuan, Shuting Wang, Jiongnan Liu, Wenhan Liu, Chenlong Deng, Haonan Chen, Zheng Liu, Zhicheng Dou, and Ji-Rong Wen. 2023. Large language models for information retrieval: A survey. arXiv preprint arXiv:2308.07107.

A Massive Text Embeddings Benchmark (MTEB)

The original MTEB benchmark (Muenninghoff et al., 2022) consists of 56 datasets of various embedding tasks. Some tasks like classification and clustering require additional use of small clustering models or classification heads after extracting embeddings. For a fair comparison, we use the official MTEB code repository 1 to test all the models, simply overriding the encode function to extract representations for the given text. Since the models under evaluation have different architectures and pooling methods, they may have different optimal instructions for their respective tasks. For fair comparisons, we do not use any instruction and only extract embeddings from the test texts.

MTEB consists of diverse small and large embedding tasks. Some datasets like MSMARCO (Nguyen et al., 2016) and DBPedia (Hasibi et al., 2017) have even more than 5,000,000 samples. Considering the test speed, for preliminary experiments and ablation studies, we select 15 representative tasks as a subset in Table 3, which aligns with LLM2VEC (BehnamGhader et al., 2024). For each

1https://github.com/embeddings-benchmark/mteb
task category, we selected the same proportion of datasets to ensure that the results in MTEB-15 are not biased towards MTEB-56.

<table>
  <tr>
    <th>Category</th>
    <th>Dataset</th>
  </tr>
  <tr>
    <td rowspan="3">Retrieval (3)</td>
    <td>SciFact</td>
  </tr>
  <tr>
    <td>ArguAna</td>
  </tr>
  <tr>
    <td>NPCorpus</td>
  </tr>
  <tr>
    <td>Reranking (2)</td>
    <td>StackOverflowDupQuestions<br>SciDocsRR</td>
  </tr>
  <tr>
    <td rowspan="3">Clustering (3)</td>
    <td>BiorxivClusteringS2S</td>
  </tr>
  <tr>
    <td>MedrxivClusteringS2S</td>
  </tr>
  <tr>
    <td>TwentyNewsgroupsClustering</td>
  </tr>
  <tr>
    <td>Pair Classification (1)</td>
    <td>SprintDuplicateQuestions</td>
  </tr>
  <tr>
    <td rowspan="3">Classification (3)</td>
    <td>Banking77Classification</td>
  </tr>
  <tr>
    <td>EmotionClassification</td>
  </tr>
  <tr>
    <td>MassiveIntentClassification</td>
  </tr>
  <tr>
    <td rowspan="3">STS (3)</td>
    <td>STS17</td>
  </tr>
  <tr>
    <td>SICK-R</td>
  </tr>
  <tr>
    <td>STSBenchmark</td>
  </tr>
  <tr>
    <td>Overall</td>
    <td>15 datasets</td>
  </tr>
</table>

Table 3: Datasets of MTEB-15

B Usage Instructions

In the entire framework, there are two important hyper-parameters: the mask ratio \( p_{mar} \) of MAR, and the weight \( \alpha \) of MAR. Based on our experience, a higher \( p_{mar} \) often provides greater benefits for tasks such as STS and clustering, which are sensitive to vector space comparisons, while the improvement for retrieval tasks is relatively smaller. For the weight \( \alpha \), it depends on your expectations for the LLM base. If you aim for the LLM itself to learn from the data (wiki-text), you can set a relatively high \( \alpha \). However, if your goal is to enhance the representation capability while maintaining the original abilities of the LLM, a value of 0.1 would be sufficient.

Regarding the number of training steps, we only trained for 100 steps on general open-source data. This is mainly because LLMs have typically encountered this open-source data during their pre-training. Numerous studies have shown that performing autoregressive training on repetitive data can easily lead to model collapse. Therefore, if there is data that can enhance the LLM itself, it can be trained extensively, and full-parameter training can be employed. For instance, we conducted full-parameter training on domain-specific data with 100 billion tokens, which continuously improved the model’s overall capabilities. However, if training on data does not benefit LLMs or even hurt LLMs, we recommend using LoRA to converge new parameters instead.

C Unsupervised Training Settings

C.1 Baselines

For Echo (Springer et al., 2024), we simply repeat the input text and take the average output embedding of the second sentence as the global embedding. For MNTP and MNTP + SimCSE (BehnamGhader et al., 2024), we train the LLaMA-3.2-1B and LLaMA-3.2-3B using the official code with the setting mentioned for Sheared-LLaMA-1.3B, as they do not release models based on these two foundation models. For LLaMA-3.1-8B, we directly use the released lora weights 2.

C.2 Our method

We keep the same training settings across all the models. For UniMAE training, we train all the models using Peft 3 package with LORA (Hu et al., 2021) adapter to train the models, with \( \alpha_{lora} = 32, r = 16, lr = 1e - 4 \). We only add trainable lora weight to the attention block and mlp block. We train the model on randomly selected wiki-text data 4 for 100 steps with a global batch size of 32 samples. The max length of each sample is 512. The learning rate is \( 1e - 4 \) and we use a constant learning rate scheduler. The MAE weight \( \alpha \) is set to 0.1 and MRC weight \( \beta \) is set to 1. MAR mask ratio \( p_{mar} \) is set to 0.5, and MRC mask ratio \( p_{mrc} \) is set to 0.5.

For SimCSE training, we train all the models using Peft package with LORA (Hu et al., 2021) adapter to train the models, with \( \alpha = 32, r = 16, lr = 1e - 4 \). We only add trainable lora weight to the attention block and mlp block. We train the model on randomly selected wiki-data for 100 steps with a global batch size of 64 samples. The max length of each sample is 512. The learning rate is \( 1e - 4 \) and we use a constant learning rate scheduler.

D Combine UniMAE with Supervised Contrastive Learning

Experiments in table 1 have shown that UniMAE significantly improves the embedding performance

2https://huggingface.co/McGill-NLP/LLM2Vec-MetaLLama-3-8B-Instruct-mntp
3https://github.com/huggingface/peft
4https://huggingface.co/datasets/wikimedia/wikipedia
<table>
  <tr>
    <th>Categories →</th>
    <th colspan="3">Retr.</th>
    <th colspan="2">Rerank.</th>
    <th colspan="2">Clust.</th>
    <th colspan="2">PairClass.</th>
    <th colspan="2">Class.</th>
    <th colspan="2">STS</th>
    <th>Avg</th>
  </tr>
  <tr>
    <th># of datasets →</th>
    <th>3</th>
    <th>2</th>
    <th>3</th>
    <th>1</th>
    <th>3</th>
    <th>3</th>
    <th>3</th>
    <th>3</th>
    <th>15</th>
  </tr>
  <tr>
    <td colspan="16">LLaMA3.2-1B</td>
  </tr>
  <tr>
    <td>Base Model</td>
    <td>51.17</td>
    <td>66.08</td>
    <td>37.23</td>
    <td>93.76</td>
    <td>69.70</td>
    <td>81.85</td>
    <td>63.05</td>
  </tr>
  <tr>
    <td>UniMAE + SimCSE</td>
    <td><b>53.95</b></td>
    <td><b>67.17</b></td>
    <td><b>38.45</b></td>
    <td><b>94.57</b></td>
    <td><b>70.41</b></td>
    <td><b>84.00</b></td>
    <td><b>64.62</b></td>
  </tr>
  <tr>
    <td colspan="16">LLaMA3.2-3B</td>
  </tr>
  <tr>
    <td>Base Model</td>
    <td>50.77</td>
    <td>66.18</td>
    <td><b>40.69</b></td>
    <td>88.89</td>
    <td>67.53</td>
    <td>79.81</td>
    <td>62.51</td>
  </tr>
  <tr>
    <td>UniMAE + SimCSE</td>
    <td><b>56.24</b></td>
    <td><b>69.32</b></td>
    <td>40.23</td>
    <td><b>94.25</b></td>
    <td><b>70.28</b></td>
    <td><b>83.50</b></td>
    <td><b>65.58</b></td>
  </tr>
  <tr>
    <td colspan="16">LLaMA3.1-8B</td>
  </tr>
  <tr>
    <td>Base Model</td>
    <td>53.07</td>
    <td>67.19</td>
    <td>40.21</td>
    <td>95.04</td>
    <td>69.42</td>
    <td>82.44</td>
    <td>64.32</td>
  </tr>
  <tr>
    <td>UniMAE + SimCSE</td>
    <td><b>57.93</b></td>
    <td><b>70.67</b></td>
    <td><b>42.65</b></td>
    <td><b>96.13</b></td>
    <td><b>70.77</b></td>
    <td><b>84.51</b></td>
    <td><b>67.00</b></td>
  </tr>
</table>

Table 4: Results on MTEB-15 after supervised contrastive learning.

of LLMs. Many studies have proven that fine-tuning supervised contrastive learning data can further improve the performance of LLMs on embedding tasks. To see whether our method can further improve the performance under supervised settings, we collect open-sourced supervised contrastive learning data to train the base model and model trained using UniMAE + SimCSE. We follow the setting of BGE-ICL (Li et al., 2024), which also does not modify the model architecture, using [EOS] embedding as the sentence embedding for contrastive learning. We use the official training code of BGE-ICL 5. Data can be found here 6. For base models and models trained using UniMAE + SimCSE, we both use the [EOS] embedding of the last transformer layer as the sentence embedding to conduct supervised contrastive learning. We train each model for 4000 steps with a global batch size of 64 samples. The max length of each text is set to 512. We use the Peft package with LORA (Hu et al., 2021) adapter to train the models. The learning rate is \( 1e - 4 \) and we use a constant learning rate scheduler.

The result can be found in table 4. It demonstrates that our approach not only directly enhances the base model’s capabilities but also shows improved performance after additional contrastive learning compared to the original base model.

E Language Model Evaluation

We use the lm-evaluation-harness package (Gao et al., 2021a) to evaluate the LLMs on 10 downstream tasks: ARC-E (Clark et al., 2018), LAMBADA (Paperno et al., 2016), LogiQA (Liu et al., 2020), PIQA (Bisk et al., 2020), SciQ (Welbl et al., 2017), and WinoGrande (Sakaguchi et al., 2021), HellaSwag (Zellers et al., 2019), ARC-C (Clark et al., 2018)), Natural Questions (Kwiatkowski et al., 2019), MMLU (Hendrycks et al., 2021). We keep a zero-shot setting for each task.

F E-commerce Domain Post-training

In the E-commerce domain, the pre-training dataset consists of approximately 100 billion tokens (Due to the principle of anonymity, we cannot disclose the specific platform). To maintain a certain level of general capability for the LLM, we incorporated a substantial amount of open-source data, with domain-specific data accounting for 25% of the total. This domain-specific data mainly includes ASR and OCR from product live streams and short videos, user comments, and textual information about the products themselves. We set \( p_{mar} = 0 \), \( p_{mrc} = 0.5 \), \( \alpha = 1 \), where MAR has already degraded into AR. This is because: 1. we need to enhance the model’s overall domain knowledge. 2. The downstream representation tasks only involve retrieval and classification, while MAR offers greater improvements for STS and clustering tasks.

The testing scenarios primarily focus on governance within the E-commerce Platform, with the ultimate goal of determining whether a given sample violates the platform’s regulatory guidelines. The tasks can be categorized into three main types. The language model task primarily involves generative QA, where for each sample, the model is directly asked whether it violates a specific set of regulations. The retrieval task focuses on identifying samples that share the same type of violation as a given problematic sample. The classification task involves using a large language model (LLM) to offline extract text representations of a sample, which serve as input for a classification head to categorize the sample’s violations.

5https://github.com/FlagOpen/FlagEmbedding
6https://huggingface.co/datasets/cfl/bge-full-data
G   Full Result

Table 5 and table 6 shows the detailed scores of models trained with UniMAE and UniMAE + SimCSE on all of the 56 MTEB datasets.
<table>
  <tr>
    <th>Task</th>
    <th>LLaMA3.2-1B</th>
    <th>LLaMA3.2-3B</th>
    <th>LLaMA3.1-8B</th>
  </tr>
  <tr><td>AmazonCounterfactualClassification.</td><td>76.58</td><td>80.66</td><td>75.01</td></tr>
  <tr><td>AmazonPolarityClassification.</td><td>78.49</td><td>74.89</td><td>76.79</td></tr>
  <tr><td>AmazonReviewsClassification.</td><td>40.26</td><td>39.98</td><td>40.34</td></tr>
  <tr><td>ArguAna.</td><td>48.28</td><td>50.38</td><td>53.87</td></tr>
  <tr><td>ArxivClusteringP2P.</td><td>45.54</td><td>46.99</td><td>47.45</td></tr>
  <tr><td>ArxivClusteringS2S.</td><td>34.58</td><td>36.62</td><td>38.63</td></tr>
  <tr><td>AskUbuntuDupQuestions.</td><td>55.73</td><td>57.02</td><td>55.00</td></tr>
  <tr><td>BIOSSES.</td><td>79.97</td><td>84.28</td><td>85.12</td></tr>
  <tr><td>Banking77Classification.</td><td>74.76</td><td>79.21</td><td>76.55</td></tr>
  <tr><td>BiorxivClusteringP2P.</td><td>36.41</td><td>35.89</td><td>35.86</td></tr>
  <tr><td>BiorxivClusteringS2S.</td><td>26.53</td><td>26.61</td><td>29.71</td></tr>
  <tr><td>CQADupstackTexRetrieval.</td><td>18.02</td><td>18.95</td><td>17.59</td></tr>
  <tr><td>ClimateFEVER.</td><td>21.99</td><td>19.03</td><td>22.36</td></tr>
  <tr><td>DBPedia.</td><td>26.13</td><td>22.56</td><td>23.78</td></tr>
  <tr><td>EmotionClassification.</td><td>44.82</td><td>48.18</td><td>42.86</td></tr>
  <tr><td>FEVER.</td><td>40.32</td><td>35.14</td><td>37.44</td></tr>
  <tr><td>FiQA2018.</td><td>24.50</td><td>26.62</td><td>25.70</td></tr>
  <tr><td>HotpotQA.</td><td>43.58</td><td>49.07</td><td>50.59</td></tr>
  <tr><td>ImdbClassification.</td><td>70.94</td><td>74.79</td><td>74.78</td></tr>
  <tr><td>MSMARCO.</td><td>21.33</td><td>20.29</td><td>19.64</td></tr>
  <tr><td>MTOPDomainClassification.</td><td>92.25</td><td>93.61</td><td>93.58</td></tr>
  <tr><td>MTOPIntentClassification.</td><td>73.41</td><td>74.20</td><td>71.03</td></tr>
  <tr><td>MassiveIntentClassification.</td><td>70.48</td><td>70.32</td><td>68.05</td></tr>
  <tr><td>MassiveScenarioClassification.</td><td>74.79</td><td>75.03</td><td>74.79</td></tr>
  <tr><td>MedrxivClusteringP2P.</td><td>29.50</td><td>26.77</td><td>29.28</td></tr>
  <tr><td>MedrxivClusteringS2S.</td><td>22.85</td><td>23.61</td><td>25.98</td></tr>
  <tr><td>MindSmallReranking.</td><td>29.92</td><td>29.14</td><td>32.91</td></tr>
  <tr><td>NFCorpus.</td><td>26.40</td><td>28.92</td><td>28.06</td></tr>
  <tr><td>NQ.</td><td>33.42</td><td>32.80</td><td>33.70</td></tr>
  <tr><td>QuoraRetrieval.</td><td>80.71</td><td>81.28</td><td>82.67</td></tr>
  <tr><td>RedditClustering.</td><td>43.36</td><td>43.28</td><td>48.13</td></tr>
  <tr><td>RedditClusteringP2P.</td><td>58.50</td><td>59.13</td><td>60.81</td></tr>
  <tr><td>SCIDOCs.</td><td>12.70</td><td>12.43</td><td>13.51</td></tr>
  <tr><td>SICK-R.</td><td>71.41</td><td>67.37</td><td>68.52</td></tr>
  <tr><td>STS12.</td><td>59.38</td><td>60.37</td><td>58.89</td></tr>
  <tr><td>STS13.</td><td>75.44</td><td>75.00</td><td>73.64</td></tr>
  <tr><td>STS14.</td><td>68.58</td><td>69.69</td><td>70.45</td></tr>
  <tr><td>STS15.</td><td>78.73</td><td>79.10</td><td>78.93</td></tr>
  <tr><td>STS16.</td><td>76.94</td><td>78.43</td><td>75.52</td></tr>
  <tr><td>STS17.</td><td>83.87</td><td>80.87</td><td>83.13</td></tr>
  <tr><td>STS22.</td><td>63.52</td><td>63.75</td><td>64.68</td></tr>
  <tr><td>STSBenchmark.</td><td>73.52</td><td>71.77</td><td>75.59</td></tr>
  <tr><td>SciDocRR.</td><td>72.72</td><td>74.21</td><td>78.48</td></tr>
  <tr><td>SciFact.</td><td>63.71</td><td>65.22</td><td>66.17</td></tr>
  <tr><td>SprintDuplicateQuestions.</td><td>81.65</td><td>90.50</td><td>87.58</td></tr>
  <tr><td>StackExchangeClustering.</td><td>57.92</td><td>63.54</td><td>63.80</td></tr>
  <tr><td>StackExchangeClusteringP2P.</td><td>32.58</td><td>33.96</td><td>33.36</td></tr>
  <tr><td>StackOverflowDupQuestions.</td><td>43.68</td><td>45.80</td><td>44.58</td></tr>
  <tr><td>SummEval.</td><td>29.93</td><td>29.25</td><td>30.40</td></tr>
  <tr><td>TRECCOVID.</td><td>44.99</td><td>40.94</td><td>44.71</td></tr>
  <tr><td>Touche2020.</td><td>11.16</td><td>10.19</td><td>10.74</td></tr>
  <tr><td>ToxicConversationsClassification.</td><td>69.44</td><td>72.38</td><td>69.21</td></tr>
  <tr><td>TweetSentimentExtractionClassification.</td><td>60.38</td><td>58.43</td><td>58.90</td></tr>
  <tr><td>TwentyNewsgroupsClustering.</td><td>34.79</td><td>32.05</td><td>40.71</td></tr>
  <tr><td>TwitterSemEval2015.</td><td>61.70</td><td>60.61</td><td>62.45</td></tr>
  <tr><td>TwitterURLCorpus.</td><td>84.37</td><td>83.82</td><td>84.65</td></tr>
  <tr><th>AVG</th><th>52.81</th><th>53.30</th><th>53.87</th></tr>
</table>

Table 5: Unsupervised results of UniMAE transformed models on MTEB.
<table>
  <tr>
    <th>Task</th>
    <th>LLaMA3.2-1B</th>
    <th>LLaMA3.2-3B</th>
    <th>LLaMA3.1-8B</th>
  </tr>
  <tr><td>AmazonCounterfactualClassification</td><td>72.39</td><td>78.01</td><td>67.97</td></tr>
  <tr><td>AmazonPolarityClassification</td><td>76.94</td><td>79.38</td><td>77.04</td></tr>
  <tr><td>AmazonReviewsClassification</td><td>37.80</td><td>40.27</td><td>40.27</td></tr>
  <tr><td>ArguAna</td><td>41.26</td><td>49.76</td><td>52.12</td></tr>
  <tr><td>ArxivClusteringP2P</td><td>47.43</td><td>48.03</td><td>49.04</td></tr>
  <tr><td>ArxivClusteringS2S</td><td>39.29</td><td>40.66</td><td>45.22</td></tr>
  <tr><td>AskUbuntuDupQuestions</td><td>57.32</td><td>59.21</td><td>59.41</td></tr>
  <tr><td>BIOSSES</td><td>83.21</td><td>85.32</td><td>87.23</td></tr>
  <tr><td>Banking77Classification</td><td>78.85</td><td>81.68</td><td>79.57</td></tr>
  <tr><td>BiorxivClusteringP2P</td><td>37.60</td><td>35.43</td><td>36.64</td></tr>
  <tr><td>BiorxivClusteringS2S</td><td>33.39</td><td>33.62</td><td>34.74</td></tr>
  <tr><td>CQADupstackTexRetrieval</td><td>16.13</td><td>21.40</td><td>19.48</td></tr>
  <tr><td>ClimateFEVER</td><td>15.81</td><td>19.54</td><td>22.67</td></tr>
  <tr><td>DBPedia</td><td>24.52</td><td>24.85</td><td>29.53</td></tr>
  <tr><td>EmotionClassification</td><td>48.30</td><td>50.47</td><td>45.34</td></tr>
  <tr><td>FEVER</td><td>24.78</td><td>37.28</td><td>55.28</td></tr>
  <tr><td>FiQA2018</td><td>23.43</td><td>27.64</td><td>28.81</td></tr>
  <tr><td>HotpotQA</td><td>33.02</td><td>46.53</td><td>53.72</td></tr>
  <tr><td>ImdbClassification</td><td>74.79</td><td>77.36</td><td>76.43</td></tr>
  <tr><td>MSMARCO</td><td>16.98</td><td>19.99</td><td>21.56</td></tr>
  <tr><td>MTOPDomainClassification</td><td>93.54</td><td>94.85</td><td>93.28</td></tr>
  <tr><td>MTOPIntentClassification</td><td>70.46</td><td>75.96</td><td>71.61</td></tr>
  <tr><td>MassiveIntentClassification</td><td>71.56</td><td>71.67</td><td>70.35</td></tr>
  <tr><td>MassiveScenarioClassification</td><td>76.89</td><td>78.67</td><td>77.22</td></tr>
  <tr><td>MedrxivClusteringP2P</td><td>30.06</td><td>29.83</td><td>30.58</td></tr>
  <tr><td>MedrxivClusteringS2S</td><td>29.34</td><td>28.56</td><td>30.46</td></tr>
  <tr><td>MindSmallReranking</td><td>32.83</td><td>31.90</td><td>33.36</td></tr>
  <tr><td>NFCorpus</td><td>27.36</td><td>31.99</td><td>27.99</td></tr>
  <tr><td>NQ</td><td>27.32</td><td>30.93</td><td>31.83</td></tr>
  <tr><td>QuoraRetrieval</td><td>85.61</td><td>85.71</td><td>85.97</td></tr>
  <tr><td>RedditClustering</td><td>54.29</td><td>57.34</td><td>57.39</td></tr>
  <tr><td>RedditClusteringP2P</td><td>57.13</td><td>59.43</td><td>61.87</td></tr>
  <tr><td>SCIDOCs</td><td>14.53</td><td>15.74</td><td>16.89</td></tr>
  <tr><td>SICK-R</td><td>73.99</td><td>70.22</td><td>72.73</td></tr>
  <tr><td>STS12</td><td>65.70</td><td>63.51</td><td>62.53</td></tr>
  <tr><td>STS13</td><td>81.08</td><td>79.88</td><td>79.11</td></tr>
  <tr><td>STS14</td><td>75.14</td><td>74.36</td><td>74.75</td></tr>
  <tr><td>STS15</td><td>83.15</td><td>82.50</td><td>81.50</td></tr>
  <tr><td>STS16</td><td>81.18</td><td>81.66</td><td>80.75</td></tr>
  <tr><td>STS17</td><td>88.88</td><td>86.53</td><td>85.55</td></tr>
  <tr><td>STS22</td><td>55.87</td><td>61.67</td><td>58.93</td></tr>
  <tr><td>STSBenchmark</td><td>81.32</td><td>79.68</td><td>78.99</td></tr>
  <tr><td>SciDocsRR</td><td>80.06</td><td>80.85</td><td>84.22</td></tr>
  <tr><td>SciFact</td><td>61.10</td><td>65.41</td><td>67.24</td></tr>
  <tr><td>SprintDuplicateQuestions</td><td>85.35</td><td>91.35</td><td>88.71</td></tr>
  <tr><td>StackExchangeClustering</td><td>66.02</td><td>68.15</td><td>71.84</td></tr>
  <tr><td>StackExchangeClusteringP2P</td><td>32.44</td><td>33.70</td><td>33.75</td></tr>
  <tr><td>StackOverflowDupQuestions</td><td>46.71</td><td>49.69</td><td>50.67</td></tr>
  <tr><td>SummEval</td><td>30.14</td><td>30.33</td><td>29.80</td></tr>
  <tr><td>TRECCOVID</td><td>51.78</td><td>52.27</td><td>51.29</td></tr>
  <tr><td>Touche2020</td><td>9.45</td><td>12.26</td><td>13.27</td></tr>
  <tr><td>ToxicConversationsClassification</td><td>69.14</td><td>71.11</td><td>67.81</td></tr>
  <tr><td>TweetSentimentExtractionClassification</td><td>60.45</td><td>59.74</td><td>58.39</td></tr>
  <tr><td>TwentyNewsgroupsClustering</td><td>48.39</td><td>46.73</td><td>55.11</td></tr>
  <tr><td>TwitterSemEval2015</td><td>65.35</td><td>65.59</td><td>67.52</td></tr>
  <tr><td>TwitterURLCorpus</td><td>83.42</td><td>85.75</td><td>84.24</td></tr>
  <tr><th>Average</th><th>54.11</th><th>56.11</th><th>56.60</th></tr>
</table>

Table 6: Unsupervised results of UniMAE+SimCSE transformed models on MTEB.

Reviewer paper summaries: This paper introduces a technique to take a pretrained, decoder-only LLM and continue finetuning with an MLM objective in order to produce a model that can *both* generate *and* produce a meaningful EOS embedding that is typically associated with encoders. This continued training is a combined loss of (1) a typical next-token prediction and (2) reconstructing the input (MLM) with a random attention mask.

In this paper, authors introduce UniMAE, a unsupervised training method designed to enable decoder-only large language models (LLMs) to perform both text generation and embedding tasks, traditionally handled by separate architectures. Inspired by Masked AutoEncoders in vision, UniMAE reformulates decoder-only LLMs into unidirectional masked auto-encoders by compressing semantic content into the [EOS] token and reconstructing input text through a small auxiliary decoder.

The paper introduces UniMAE, an unsupervised training method for turning a decoder-only LM into a masked auto-encoder, allowing it to be used as an encoder for embedding-based tasks such as MTEB without dropping performance on regular LM tasks. This is based on evaluation on both MTEB and a handful of traditional LM tasks. The authors justify  some core design choices (e.g., masking the input) through ablations, and also show their approach outperforms AR models in domain-specific (e-commerce) tasks.

This paper introduces UniMAE, an unsupervised method that compresses the semantic content of input text into the [EOS] (end-of-sequence) embedding of the last transformer layer, while preserving the capabilities of original LLMs. Comprehensive evaluations across 56 datasets demonstrate that it can achieve state-of-the-art results under unsupervised settings, establishing the first effective approach to unifying generation and representation learning in decoder-only architectures.
Summary of strengths: - Consistent positive gains on embedding tasks using this technique. Stated that generation performance was maintained but some scores missing (missing Figure 4?)
- Shows both quantitative (empirical scores on MTEB) and qualitative performance improvements. EOS encodings now nicely cluster based off domain which is desirable for many retrieval tasks.
- Thorough evaluation of retrieval-based tasks.
- Substantial consideration of alternative approaches.

S1: The core idea of separating Masked Auto-Regressive (MAR) and Masked Re-Construct (MRC) objectives, and then jointly optimizing them, is interesting. 

S2: The paper is well-structured and clearly written, making it easy to follow. The authors provide intuitive illustrations (e.g., Figure 1) and appropriate empirical validation across a wide range of tasks in the MTEB benchmark.

- The core approach is described well and straightforward to understand.
- The evaluation is done across a wide range of tasks and the model outperforms the proposed baselines by quite large margins.
- The model still remains solid in the AR evaluations performed.
- Choices like the masking ratio are ablated and well-justified.

- The paper is well presented and easy to understand.

- The paper addresses a crucial limitation for decoder-only LMs that dominate the NLP domain.

- The proposed methodology is novel and technically sound.

- The experiment covers a broad range of aspects, making the achieved result convincing.
Summary of weaknesses: - Requires additional parameters and data for training that are missing from main body. How big is the new decoder? Is it just a single layer? What effect does the size have? What training data is used for this (appendix says "randomly selected wiki data" but this is vague)? How much of it?
- Missing descriptions of baselines---particulalry MNTP and SimCSE---the latter of which is used in conjunction with this technique though it's not clear how. 
- MAR vs AR not particularly compelling. At most a ~0.02 improvement from 0.0 masking despite upward trend. This is likely a small normalization effect and doesn't show that the model would "memorize all the input tokens during training". I would be interested to see this chart next to a similar one to generation performance---do you see degradation with MAR over AR? Is there a tradeoff here that the reader would want to be aware of

W1: While the paper claims that UniMAE preserves both embedding and **generation** capabilities of decoder-only LLMs, the empirical evidence for generation ability preservation is relatively weak. The main experimental results (e.g., Tables 1 and 6) focus almost exclusively on embedding-oriented tasks, particularly those from the MTEB benchmark. No results for justifying the **generation** capabilities are preserved. 

W2: Although the authors mention in Section 3.3 and Appendix E that they evaluate the generative capability on 10 standard language modeling benchmarks, no quantitative results or detailed tables/figures are provided in the main paper or the appendix. Specifically, Figure 4, which is referred to in the main text, appears to be missing from both the main body and the appendix, and Appendix E only lists the evaluation datasets without presenting any results.

W3: The Masked Re-Construct (MRC) component introduces a “tiny decoder” that reconstructs the input from the [EOS] embedding. However, the architecture of this decoder is insufficiently described in the paper. Key details such as the number of layers, attention mechanism type, hidden dimensions, and whether any pre-trained weights are used are missing. One conceptual issue I found unclear is the design of the MRC (Masked Re-Construct) objective. While the MAR objective operates on a masked input (X̃), the MRC decoder appears to take unmasked inputs, combined with the [EOS] embedding, to reconstruct the original sequence (MRC's input). If the goal is to test whether the [EOS] token contains sufficient compressed semantic information, giving the decoder the full (or lightly corrupted) input may defeat the purpose. It’s unclear whether the decoder is genuinely relying on the [EOS] embedding for semantic reconstruction, or simply reusing the visible inputs.

- My largest concerns are around a lack of baselines: There is no comparison to models like GRIT-LM [1] or sGPT [2] (in particular, a weighted mean-pooling rather than regular average pooling). It would also be useful to show some comparable embedding-only approaches (e.g. maybe BGE [3]), even if you can’t use exactly the same base.
- I think the LM metrics used lean closer to common metrics in pretraining, and rely mostly on single-token prediction rather than long-form generations. As such, it is unclear if the proposed approach hurts the model’s ability to generate longer outputs.
- Perhaps the e-commerce tasks involve longer generations, but I couldn’t find a detailed description of those evaluations beyond what is in appendix F? Adding examples of the tasks would help in understanding what the e-commerce-centric evaluations look like and are testing.

Overall, I think the paper is solid for a short paper, but as said I worry that the authors are missing potentially relevant baselines and some further details on evaluation.

References:
- [1] Muennighoff, N., Su, H., Wang, L., Yang, N., Wei, F., Yu, T., Singh, A., & Kiela, D. (2024). Generative Representational Instruction Tuning. ArXiv, abs/2402.09906.
- [2] Muennighoff, N. (2022). SGPT: GPT Sentence Embeddings for Semantic Search. ArXiv, abs/2202.08904.
- [3] Xiao, S., Liu, Z., Zhang, P., Muennighoff, N., Lian, D., & Nie, J. (2023). C-Pack: Packed Resources For General Chinese Embeddings. Annual International ACM SIGIR Conference on Research and Development in Information Retrieval.

- The procedure for applying UniMAE to different tasks should be briefly described.

- The main experiment is only limited to LLaMA. Models from other LLM families may be tested to yield more convincing conclusions.

- The specific setting of random attention mask should be further elaborated in terms of, i.e., whether the same random attention matrix is applied on a single batch or each sample takes a different attention matrix.
Comments, suggestions and typos: You describe this technique as unsupervised, but it must require text---I believe semi-supervised is the correct verbiage.

You use the phrase "unidirectional" as a contrast to bidirectional however I think the term "causal mask" is more common.

Figure 1 is confusing. MAR/next-token prediction occurs *with* masks. Does this not mean that e1-->x2 and e2-->x3 are supposed to have identical inputs? Since the output of e1 is soley based on x1 and the output of e2 is based on x1 and [mask], aka *just* x1?

The MRC description is difficult to parse. The inline hidden state outputs are non-intuitive. An aligned MRC figure may aid understanding. 

Figures 2 and 3 are quite small--impossible to read on printed copy. Requires a lot of zooming for text size on pdf.

Typos:
- "an" Decoder only Line 010.
- missing period line 134
- In LLAMA 3.2-1.B section of Table 1, I believe your method is listed with MAE instead of UniMAE on the last line.

1:  Identical scores between “Uni. MEAN” and “Bi. MNTP” rows:
In multiple columns (e.g., LLaMA3.2-1B block), the performance metrics of “Uni. MEAN” and “Bi. MNTP” are exactly the same, which is unlikely and suggests a copy-paste error or mislabeling. This should be double-checked and corrected.

2. Inconsistent naming of UniMAE vs. MAE:
The method labeled as “Uni. UniMAE” in one row and “Uni. MAE + SimCSE” in another creates confusion—presumably both refer to the same base model (UniMAE), with and without SimCSE fine-tuning. For clarity, the method names should be consistent throughout the table and text (e.g., “UniMAE” and “UniMAE + SimCSE”).

- Line 199 says “Figure 4a”, but there is no figure 4 in the paper. Is this referring to Figure 3b?

Please refer to my comments above. Also, the title of Section 3.3 should end with a question mark, and title of Section 3.4 should be ''UniMAE optimizes Vector Space''
Overall assessment scores: 4

3.5

3.5

4
Metareview score: 4
Metareview: The paper proposes an unsupervised training technique (UniMAE) that converts a decoder-only LM into a text encoder for embedding-based tasks, while maintaining its performance as an autoregressive LM. The paper includes evaluation on MTEB and several traditional LM tasks. The paper includes a number of ablations and results on some domain-specific tasks related to e-commerce.
Summary of reasons to publish: * The approach is intuitively appealing and clearly described (for the most part--see first weakness below).

* The evaluation is thorough and the model clearly outperforms reasonable baselines on MTEB while retaining performance on LM tasks.

* Design choices are justified and evaluated well via ablations.