MAGNET: Augmenting Generative Decoders with Representation Learning and Infilling Capabilities

Thank you for thoroughly reviewing our work. We are glad to hear that you appreciated the deployed attention mechanism and our analysis of the repetition problem. We have tried to address each of your concerns point by point in this rebuttal.

(Weakness 1) What is the purpose of the unified framework? Could each task benefit from each other?

Ans. In Appendix C of our paper, we perform an ablation analysis demonstrating the impact of unified training on the model's performance on the representation learning tasks. Notably, we observe that while token-level representation learning is purely promoted by the MNTP loss, adding MSG (which is a generative loss) boosts the results across all token-level classification tasks.

Secondly, the results in Tables 1, 2, and 3, demonstrate the superior performance of our approach over LLM2Vec and LLM2Vecᴹᴺᵀᴾ [1]. This also underscores the beneficial interaction between different objectives in a unified framework. This is because LLM2Vecᴹᴺᵀᴾ uses the same MNTP objective as our approach, but in isolation. By simultaneously training for multiple objectives, we demonstrate superior performance on the representation learning tasks.

Besides the beneficial synergy between losses, training unified models for multiple tasks offers compelling practical and performance benefits too. For instance, these models require less storage and computational overhead compared to deploying multiple specialized models. Training a unified model is also more resource-efficient and it can ensure consistency across related tasks, which is crucial for real-world applications where multiple capabilities must work together seamlessly. Thus, owing to these benefits, our goal in this work was to explore whether and how a single LLM can effectively perform representation learning and text generation (infilling and open-ended).

We acknowledge that these points could have been emphasized more effectively. As a result, we have revised our draft to highlight them more clearly, with the most significant updates in Sections 4.1 and 4.2.

(Weakness 2) The proposed three training objectives are not new. I suggest the author to explore a more sophisticated loss function, rather than a simple linear combination of these three.

Ans. We do not claim the novelty of the loss function itself. In our view, the innovation of the paper lies in combining three widely-used self-supervised loss functions for language modeling with a nuanced attention mechanism, demonstrating the possibility of a unified model that can perform both representation learning and text infilling. To the best of our knowledge, no prior work has attempted to simultaneously augment a pretrained LLM with these capabilities and evaluate them within a single framework.

We adopt the formulation and naming of the MNTP task as described in [1]. Appendix E provides some justification for this naming convention and explores the potential application of the standard MTP objective for token-level representation learning using decoder-only models.

(Weakness 3 and Question 2) Compare with stronger baselines on word-level and STS tasks like XLNet and StructBERT.

Ans. We present a performance comparison between the proposed approach, XLNet, and StructBERT in the table below.

Following the approach outlined in the paper, we train a linear classifier for token-level tasks and perform zero-shot evaluation on sentence-level tasks. We observe that LLaMa-2-7B, adapted using our method, outperforms the text encoders across all tasks by a significant margin.

For XLNet, we experiment with both the last token and mean pooled token representations and report the best results. For StructBERT, we use the CLS token representation. The results on the STS tasks demonstrate that these strong encoder models do not perform well for zero-shot sentence-level tasks and require additional fine-tuning for effective performance. In contrast, LLMs adapted using MAGNET can effectively handle these tasks. Additionally, fine-tuning XLNet and StructBERT with contrastive loss may negatively impact their performance on word-level tasks, whereas our method offers a unified framework that maintains strong performance across both word- and sentence-level tasks.

[1] BehnamGhader et al., LLM2Vec: Large Language Models Are Secretly Powerful Text Encoders, COLM 2024

continued

Dear Reviewer Ygdb,

We are following up on our previous responses addressing your concerns. We are eager to understand if our revisions meet your requirements or if you need any additional clarifications. Your specific feedback will help us finalize and improve our submission.

We look forward to your reply.

Thank you for taking the time to review our work. We appreciate your constructive suggestions and have made efforts to address them in this response.

Regarding comparison with other unified pre-training approaches

We compare the model adapted using our approach with XLNet on word-level representation learning tasks, sentence-level representation learning tasks, and a text generation task.

• Word-level representation learning: Following the approach outlined in the paper, we train a linear classifier for token-level tasks of Chunking, NER, and POS-tagging. We find that our approach outperforms XLNet on all these tasks.

• Sentence-level representation learning: We perform zero-shot evaluations on the STS tasks. For XLNet, we evaluated both the last token and mean pooling strategies for generating sentence encoding and selected the best of two numbers for comparison. The results demonstrate that XLNET does not perform well for zero-shot sentence-level tasks and requires additional fine-tuning for effective performance.

• Text generation: We also evaluate XLNet on the MMLU benchmark, and our evaluation reveals severe limitations in its ability to handle the formal evaluation setup for such tasks. Specifically, when presented with the standard input format of questions and multiple choices, XLNet consistently fails to generate responses in the required format (which is not the case with our model). To ensure a thorough assessment, we also attempt an alternative evaluation approach by analyzing the model's prediction scores for each option and selecting the highest-scoring choice. Even with this modified methodology, XLNet-Large achieves only 21.4% accuracy compared to our approach's 43.7%. These results underscore XLNet's fundamental limitations in two critical areas: the inability to generate responses in the required format, and insufficient language understanding/reasoning capability to perform effectively on comprehensive benchmarks like MMLU.

While our evaluation of XLNet provides useful insights, we acknowledge that it may not be the most relevant baseline given our research objectives. Our work focuses on augmenting powerful LLMs with additional capabilities such as representation learning and infilling, rather than developing a new pre-training framework from scratch. Therefore, we believe that more appropriate comparisons would be with other adaptation methods that modify pre-trained LLMs for specific tasks, as well as with the base LLMs that we adapt. These baselines would better align with our goal of enhancing, rather than replacing, existing language models.

The authors should present a comprehensive related work section.

Thank you for the feedback! In the revised version, we have significantly expanded our related work section to better position our work in the context of the existing approaches that unify language understanding and generation.

Dear Reviewer qxxH,

We are following up on our previous responses addressing your concerns. We are eager to understand if our revisions meet your requirements or if you need any additional clarifications. Your specific feedback will help us finalize and improve our submission.

We look forward to your reply.

Thank you for your thorough review of our work. We are glad that you found the proposed approach innovative and most of the experiments satisfactory. Below, we address the weaknesses and questions you raised.

(Weakness) Lack of experiments on text generation tasks.

Ans. Thank you for bringing this to our attention! To substantiate the model's text generation capabilities post-adaptation, we conduct evaluations using the MMLU benchmark. We compare our adapted model against the base LLaMa-2-7B model by fine-tuning on two datasets:

• Wikitext: The post-adaptation results we present in the paper use the Wikitext-103 dataset. When the base model is adapted using only the Wikipedia data, its performance drops by 1.6% on average. This degradation can be attributed to the dataset's encyclopedic nature, potentially introducing a knowledge bias that narrows the model's generalization capabilities.

• SlimPajama: To mitigate the limitations observed in the Wikitext-based adaptation, we also adapt the model by fine-tuning on the SlimPajama dataset, which comprises a mixture of texts from more diverse data sources like Commoncrawl, C4, GitHub, Books, ArXiv, Wikipedia, and StackExchange. Notably, using this dataset, the performance drop on MMLU is only 0.4% on average, with some categories (like 'Others') showing comparable or even slightly improved performance.

So, all in all, we demonstrate the text generation capabilities of the adapted model using the following:

1. We show that the model retains the knowledge and generation capabilities of the base model through experiments on the MMLU benchmark (as discussed above).
2. We show that, compared to other text encoders (like BERT) and LLM-adaptation methods (like LLM2Vec), the models adapted using our approach are significantly less prone to repeating words and phrases (Section 4.4).
3. Through human evaluation, we demonstrate that our method adapts LLM to effectively perform text infilling, which is inherently a generative task (Section 4.3).
4. We also include qualitative examples in Table 11 of the paper, showcasing the text generated from a seed prompt by our model and the baselines.

(Question) Have you tried other methods of using Llama features (for sentence encoding)?

Ans. We compare different strategies for using the LLaMa-2-7B features for the STS task. Specifically, we evaluate three approaches for encoding a sentence: (1) mean pooling across all token representations, (2) using the last token representation, and (3) introducing a learnable CLS token. We use SimCSE fine-tuning to assess each pooling strategy's effectiveness in generating meaningful sentence embeddings, and the results are presented in the table below.

We find that the last token and mean pooling produce comparable results, outperforming the learned CLS token. In our setup, we prefer using the last token representation, as it enables the separation of tokens used for word-level and sentence-level representation learning tasks (please see Sections 3.2.1 and 3.2.2 for more details).

Thank you for your thorough review. We are pleased you appreciated the writing and the proposed concept. We value your constructive feedback and address your concerns point-by-point in this response.

(Weakness 1) Substantiating the claim that the model retains its text generation capabilities post-adaptation with our approach.

Ans. To substantiate the model's text generation capabilities post-adaptation, we conduct evaluations using the MMLU benchmark. We compare our adapted model against the base LLaMa-2-7B model by fine-tuning on two datasets:

• Wikitext: The post-adaptation results we present in the paper use the Wikitext-103 dataset. When the base model is adapted using only the Wikipedia data, its performance drops by 1.6% on average. This degradation can be attributed to the dataset's encyclopedic nature, potentially introducing a knowledge bias that narrows the model's generalization capabilities.

• SlimPajama: To mitigate the limitations observed in the Wikitext-based adaptation, we also adapt the model by fine-tuning on the SlimPajama dataset, which comprises a mixture of texts from more diverse data sources like Commoncrawl, C4, GitHub, Books, ArXiv, Wikipedia, and StackExchange. Notably, using this dataset, the performance drop on MMLU is only 0.4% on average, with some categories (like 'Others') showing comparable or even slightly improved performance.

So, all in all, we demonstrate that the model retains its text generation capabilities post-adaptation in the following ways:

1. We show that the model retains the knowledge and generation capabilities of the base model through experiments on the MMLU benchmark (as discussed above).
2. Compared to other text encoders and LLM-adaptation methods (like LLM2Vec), the models adapted using our approach are significantly less prone to repeating words and phrases (Section 4.4).
3. Through human evaluation, we demonstrate that our method adapts LLM to perform text infilling, which is inherently a generative task (Section 4.3).
4. We also include qualitative examples in Table 10, showcasing the text generated from a seed prompt by our model and the baselines.

We have included the MMLU experiment in the revised version of our paper and hope these findings substantiate our claim that the model largely retains its text-generation capabilities post-adaptation.

(Weakness 2) It's unclear whether the results reported in Section 4 reflect zero-shot performance or task-specific fine-tuning.

Ans. All results, except for token-level classification (Table 1), are obtained in a zero-shot manner. For token-level tasks, we train a linear layer to enable the model to interpret class labels, which we believe is standard practice for evaluating such tasks.

(Question) In the "Overall Loss" section of Appendix A, it's unclear why a two-stage training approach is necessary and why only two objectives are used in the first stage.

Ans. We begin with the MNTP and MSG tasks because they help the model incorporate future context, a capability that the base model lacks. This approach aligns with prior work [1], which adapts LLMs for representation learning.

We regret not addressing this point in the original version of our paper and have now updated it accordingly. Thank you for bringing this to our attention!

[1] BehnamGhader et al., LLM2Vec: Large Language Models Are Secretly Powerful Text Encoders, COLM 2024.

Dear Reviewer V5cp,

Thank you for providing valuable feedback on our submission. We have submitted our response to your concerns, along with a revised version of the paper. With only a few days remaining until the rebuttal deadline, we would greatly appreciate it if you could review our response and let us know if you have any additional questions or concerns. We are eager to ensure all your points are thoroughly addressed and look forward to your feedback.

Best regards,

Authors of Submission3348

Dear Reviewer V5cp,

We are following up on our previous responses addressing your concerns. We are eager to understand if our revisions meet your requirements or if you need any additional clarifications. Your specific feedback will help us finalize and improve our submission.

We look forward to your guidance.