Paper: Headless Language Models: Learning without Predicting with Contrastive Weight Tying (ICLR-2024)

Reviewer comment (Novelty):
The main issue with this paper is a huge lack of bibliographic work, which keeps the paper from proper contextualization; a proper related work section would have greatly helped in making the experiments less shallow.

Author response:
- *On missing references*

We are truly grateful to the reviewer for pointing out these very interesting references. We had indeed missed to what extent this part of the literature was relevant to our work. When we wrote the Related Work section, we wanted to focus on the specific context of (pre-)training models in a self-supervised fashion. Hence, we specifically cited concurrent works that used contrastive learning in that specific context.

We still advocate that our work is novel *in that specific context*, but we will adjust some phrasings in our article to avoid conveying any impression of over-statement. In the next few days, we will upload a revised version of our paper with an updated Related Work section that engages better with the line of work you refer to. 

We would also like to underline **the differences between our work and the references the reviewer suggested**. A significant difference between our approach and importance sampling methods [1,2,4] is that our regularization is done on the hidden representation distribution and not on its support. As a matter of fact, we sample the negative samples uniformly from the batch, which means several instances of a given token may be found among them. The sorting objective of [3] is indeed an application of the NCE loss for the language modeling task, but they do not use the input embeddings as targets, resulting in a memory overhead. While [5] uses a tied input/output approach, they freeze the input embeddings and use trained static embeddings as targets. In [6], the authors also use various trained non-contextual embeddings to anchor their predictions. Although it could be interesting to see how the EBM approach from [7] works with our encoder, it is not clear how it could solve the problem of calibration in the decoder case, as it is used to sort sequences sampled from the pretrained LM, while our causal headless model needs to be adapted for token-level generation.

Reviewer comment (Novelty):
The lack of reference to relevant previous work keeps the paper from being well-positioned within the literature, especially when it comes to comparison with the numerous pre-existing similar methods.

Author response:
- *On missing references*

We are truly grateful to the reviewer for pointing out these very interesting references. We had indeed missed to what extent this part of the literature was relevant to our work. When we wrote the Related Work section, we wanted to focus on the specific context of (pre-)training models in a self-supervised fashion. Hence, we specifically cited concurrent works that used contrastive learning in that specific context.

We still advocate that our work is novel *in that specific context*, but we will adjust some phrasings in our article to avoid conveying any impression of over-statement. In the next few days, we will upload a revised version of our paper with an updated Related Work section that engages better with the line of work you refer to. 

We would also like to underline **the differences between our work and the references the reviewer suggested**. A significant difference between our approach and importance sampling methods [1,2,4] is that our regularization is done on the hidden representation distribution and not on its support. As a matter of fact, we sample the negative samples uniformly from the batch, which means several instances of a given token may be found among them. The sorting objective of [3] is indeed an application of the NCE loss for the language modeling task, but they do not use the input embeddings as targets, resulting in a memory overhead. While [5] uses a tied input/output approach, they freeze the input embeddings and use trained static embeddings as targets. In [6], the authors also use various trained non-contextual embeddings to anchor their predictions. Although it could be interesting to see how the EBM approach from [7] works with our encoder, it is not clear how it could solve the problem of calibration in the decoder case, as it is used to sort sequences sampled from the pretrained LM, while our causal headless model needs to be adapted for token-level generation.

Reviewer comment (Experiments):
The lack of bibliographic references also keeps the paper from investigating relevant experimental settings, for example, comparison with similar objectives and ablation studies which could shed some light on what makes CWT work better than some previous methods.

Author response:
- *On missing references*

We are truly grateful to the reviewer for pointing out these very interesting references. We had indeed missed to what extent this part of the literature was relevant to our work. When we wrote the Related Work section, we wanted to focus on the specific context of (pre-)training models in a self-supervised fashion. Hence, we specifically cited concurrent works that used contrastive learning in that specific context.

We still advocate that our work is novel *in that specific context*, but we will adjust some phrasings in our article to avoid conveying any impression of over-statement. In the next few days, we will upload a revised version of our paper with an updated Related Work section that engages better with the line of work you refer to. 

We would also like to underline **the differences between our work and the references the reviewer suggested**. A significant difference between our approach and importance sampling methods [1,2,4] is that our regularization is done on the hidden representation distribution and not on its support. As a matter of fact, we sample the negative samples uniformly from the batch, which means several instances of a given token may be found among them. The sorting objective of [3] is indeed an application of the NCE loss for the language modeling task, but they do not use the input embeddings as targets, resulting in a memory overhead. While [5] uses a tied input/output approach, they freeze the input embeddings and use trained static embeddings as targets. In [6], the authors also use various trained non-contextual embeddings to anchor their predictions. Although it could be interesting to see how the EBM approach from [7] works with our encoder, it is not clear how it could solve the problem of calibration in the decoder case, as it is used to sort sequences sampled from the pretrained LM, while our causal headless model needs to be adapted for token-level generation.

Reviewer comment (Novelty):
On the missing related work — the claim "This work stands as the first attempt to train language models using an explicit contrastive loss" is inaccurate, as the first known attempt to the reviewer would be [1], with many other more recent references, especially those making use of Noise Contrastive Estimation, such as [2], which has been adapted into a loss resembling the one in the paper in [3], and the first use of a direct softmax-like objective with negative examples coming from the batch appearing in [4].

Author response:
- *On missing references*

We are truly grateful to the reviewer for pointing out these very interesting references. We had indeed missed to what extent this part of the literature was relevant to our work. When we wrote the Related Work section, we wanted to focus on the specific context of (pre-)training models in a self-supervised fashion. Hence, we specifically cited concurrent works that used contrastive learning in that specific context.

We still advocate that our work is novel *in that specific context*, but we will adjust some phrasings in our article to avoid conveying any impression of over-statement. In the next few days, we will upload a revised version of our paper with an updated Related Work section that engages better with the line of work you refer to. 

We would also like to underline **the differences between our work and the references the reviewer suggested**. A significant difference between our approach and importance sampling methods [1,2,4] is that our regularization is done on the hidden representation distribution and not on its support. As a matter of fact, we sample the negative samples uniformly from the batch, which means several instances of a given token may be found among them. The sorting objective of [3] is indeed an application of the NCE loss for the language modeling task, but they do not use the input embeddings as targets, resulting in a memory overhead. While [5] uses a tied input/output approach, they freeze the input embeddings and use trained static embeddings as targets. In [6], the authors also use various trained non-contextual embeddings to anchor their predictions. Although it could be interesting to see how the EBM approach from [7] works with our encoder, it is not clear how it could solve the problem of calibration in the decoder case, as it is used to sort sequences sampled from the pretrained LM, while our causal headless model needs to be adapted for token-level generation.

Reviewer comment (Novelty):
While input/output embeddings weight sharing was common in most early language models prior to 2017, the paper combines this approach with a contrastive objective to eliminate the output vocabulary, which is a nice idea; however, related prior work such as [5], which proposes a dedicated loss for predicting static embeddings, and [6], which studies how to choose the target embeddings (relevant due to the use of weight tying), should be considered.

Author response:
- *On missing references*

We are truly grateful to the reviewer for pointing out these very interesting references. We had indeed missed to what extent this part of the literature was relevant to our work. When we wrote the Related Work section, we wanted to focus on the specific context of (pre-)training models in a self-supervised fashion. Hence, we specifically cited concurrent works that used contrastive learning in that specific context.

We still advocate that our work is novel *in that specific context*, but we will adjust some phrasings in our article to avoid conveying any impression of over-statement. In the next few days, we will upload a revised version of our paper with an updated Related Work section that engages better with the line of work you refer to. 

We would also like to underline **the differences between our work and the references the reviewer suggested**. A significant difference between our approach and importance sampling methods [1,2,4] is that our regularization is done on the hidden representation distribution and not on its support. As a matter of fact, we sample the negative samples uniformly from the batch, which means several instances of a given token may be found among them. The sorting objective of [3] is indeed an application of the NCE loss for the language modeling task, but they do not use the input embeddings as targets, resulting in a memory overhead. While [5] uses a tied input/output approach, they freeze the input embeddings and use trained static embeddings as targets. In [6], the authors also use various trained non-contextual embeddings to anchor their predictions. Although it could be interesting to see how the EBM approach from [7] works with our encoder, it is not clear how it could solve the problem of calibration in the decoder case, as it is used to sort sequences sampled from the pretrained LM, while our causal headless model needs to be adapted for token-level generation.

Reviewer comment (Novelty):
Language modelling without an output vocabulary (contrastive objective or not) has been extensively explored in the past; the specific problem addressed in subsection 3.3 with fine-tuning has had many proposed solutions, with [5] and [6] using nearest neighbour approaches, and [7] proposing an energy-based model to repurpose a well-trained encoder (outputting embeddings) into a generative model.

Author response:
- *On missing references*

We are truly grateful to the reviewer for pointing out these very interesting references. We had indeed missed to what extent this part of the literature was relevant to our work. When we wrote the Related Work section, we wanted to focus on the specific context of (pre-)training models in a self-supervised fashion. Hence, we specifically cited concurrent works that used contrastive learning in that specific context.

We still advocate that our work is novel *in that specific context*, but we will adjust some phrasings in our article to avoid conveying any impression of over-statement. In the next few days, we will upload a revised version of our paper with an updated Related Work section that engages better with the line of work you refer to. 

We would also like to underline **the differences between our work and the references the reviewer suggested**. A significant difference between our approach and importance sampling methods [1,2,4] is that our regularization is done on the hidden representation distribution and not on its support. As a matter of fact, we sample the negative samples uniformly from the batch, which means several instances of a given token may be found among them. The sorting objective of [3] is indeed an application of the NCE loss for the language modeling task, but they do not use the input embeddings as targets, resulting in a memory overhead. While [5] uses a tied input/output approach, they freeze the input embeddings and use trained static embeddings as targets. In [6], the authors also use various trained non-contextual embeddings to anchor their predictions. Although it could be interesting to see how the EBM approach from [7] works with our encoder, it is not clear how it could solve the problem of calibration in the decoder case, as it is used to sort sequences sampled from the pretrained LM, while our causal headless model needs to be adapted for token-level generation.

Reviewer comment (Experiments):
The reviewer believes that consulting and incorporating insights from this previous work would help to improve the design of the experiments and identify directions for improvement.

Author response:
- *On missing references*

We are truly grateful to the reviewer for pointing out these very interesting references. We had indeed missed to what extent this part of the literature was relevant to our work. When we wrote the Related Work section, we wanted to focus on the specific context of (pre-)training models in a self-supervised fashion. Hence, we specifically cited concurrent works that used contrastive learning in that specific context.

We still advocate that our work is novel *in that specific context*, but we will adjust some phrasings in our article to avoid conveying any impression of over-statement. In the next few days, we will upload a revised version of our paper with an updated Related Work section that engages better with the line of work you refer to. 

We would also like to underline **the differences between our work and the references the reviewer suggested**. A significant difference between our approach and importance sampling methods [1,2,4] is that our regularization is done on the hidden representation distribution and not on its support. As a matter of fact, we sample the negative samples uniformly from the batch, which means several instances of a given token may be found among them. The sorting objective of [3] is indeed an application of the NCE loss for the language modeling task, but they do not use the input embeddings as targets, resulting in a memory overhead. While [5] uses a tied input/output approach, they freeze the input embeddings and use trained static embeddings as targets. In [6], the authors also use various trained non-contextual embeddings to anchor their predictions. Although it could be interesting to see how the EBM approach from [7] works with our encoder, it is not clear how it could solve the problem of calibration in the decoder case, as it is used to sort sequences sampled from the pretrained LM, while our causal headless model needs to be adapted for token-level generation.

Reviewer comment (Evaluation):
On experiments — what is the hypothesis for why CWT achieves better performance than the vanilla model with the same number of training tokens on most GLUE tasks?

Author response:
**Questions**

- *On better data-efficiency*

We attempt to explain the better performance of our approach in **Appendix A**, as detailed in the General rebuttal. Our argument is also used by [4] (section 3.1.1), which we will mention accordingly.

Reviewer comment (Experiments):
When fine-tuning a causal LM for text generation, is the supplementary amount of compute received by the fine-tuned model compared to the vanilla model significant?

Author response:
**Questions**

- *Supplementary compute*

See *General rebuttal / On the additional cost of causal fine-tuning*. The additional fine-tuning costs 1,7% of the pre-training compute.

Reviewer comment (Evaluation):
The paper does not discuss that for larger models, the $d_{model} \times V$ classification matrix corresponds to a tiny percentage of the total number of model parameters (and thus, total memory and computation during training). For example, it corresponds to roughly 2.5% and 0.4% of the Llama2 (7b) and Llama2 (70b) model parameters, respectively.

Author response:
We sincerely thank the reviewer for its thorough review, and for its careful technical considerations.

**Weaknesses**
- *Scaling of the input embeddings proportion*

We entirely agree that as monolingual models scale, the benefits of our approach in terms of memory and compute efficiency will progressively decrease. However, we argue that the benefits of our approach in terms of data-efficiency, i.e. the final performance of a model trained on the same data, can be expected to scale. We addressed this point in **Appendix B**.

Moreover, we believe that the computational benefits theoretically hold for large language models with larger vocabularies. For mGPT (1.3B parameters, 100000 tokens vocabulary) and Bloom-7B1 (250880 tokens vocabulary), input embeddings represent ~15% of the model parameters. If one used the XLM-V vocabulary (1M tokens) to train the architecture of Llama-7B, the input embeddings would represent ~60% of the trainable parameters. We advocate that these use cases could greatly benefit from our approach.

Reviewer comment (Experiments):
The paper does not perform careful ablations to help understand why the proposed method attains better model performance than standard cross-entropy (CE) LM training. For example, it could have considered weight-tying with the CE loss, or a contrastive loss without weight tying. This latter option could still be implemented in a compute-efficient way during training (since only a subset of columns of the classification matrix would need to be used). It could have also tried using different numbers of negative samples for the contrastive loss function (assuming # negative samples was decoupled from # masked samples), such that on one extreme (# neg samples = vocab size) the loss would be equivalent to the cross-entropy loss function (modulo the weight tying). Overall, the reviewer was confused by why the proposed method would yield better-performing models than cross-entropy.

Author response:
We sincerely thank the reviewer for its thorough review, and for its careful technical considerations.

**Weaknesses**
- Points 1, 2, 4 and 5 are addressed in the *General Rebuttal*
- *On the lack of ablations*

We thank the reviewer for this very relevant question. We addressed most of its concerns in *General Rebuttal / On ablation study*. As a clarification, we would like to point out that increasing the number of negative samples would not make the CWT loss equivalent to the cross-entropy loss function. The negative samples would indeed follow the Zipfian distribution of the data instead of the uniform distribution implicitly used in the cross-entropy loss. We hypothesize that regularizing based on the unigram distribution instead of the uniform vocabulary distribution may also explain the performance improvement of CWT. 

We also invite the reviewer to read **Appendix A** that gives some insight on the reasons of the effectiveness of our method.

Reviewer comment (Novelty):
The reviewer was torn between assigning "marginal accept" and "marginal reject," ultimately choosing "marginal accept" given the efficiency benefit for smaller models and the intriguing possibility of improved model quality. The reasons for considering "marginal reject" were that the method was not very novel, only led to meaningful efficiency gains for small models without experiments on medium/large models, and the paper did not adequately explain why the method performs better.

Author response:
We sincerely thank the reviewer for its thorough review, and for its careful technical considerations.

**Weaknesses**
- Points 1, 2, 4 and 5 are addressed in the *General Rebuttal*
- *Scaling of the input embeddings proportion*

We entirely agree that as monolingual models scale, the benefits of our approach in terms of memory and compute efficiency will progressively decrease. However, we argue that the benefits of our approach in terms of data-efficiency, i.e. the final performance of a model trained on the same data, can be expected to scale. We addressed this point in **Appendix B**.

Moreover, we believe that the computational benefits theoretically hold for large language models with larger vocabularies. For mGPT (1.3B parameters, 100000 tokens vocabulary) and Bloom-7B1 (250880 tokens vocabulary), input embeddings represent ~15% of the model parameters. If one used the XLM-V vocabulary (1M tokens) to train the architecture of Llama-7B, the input embeddings would represent ~60% of the trainable parameters. We advocate that these use cases could greatly benefit from our approach.

- *On the lack of ablations*

We thank the reviewer for this very relevant question. We addressed most of its concerns in *General Rebuttal / On ablation study*. As a clarification, we would like to point out that increasing the number of negative samples would not make the CWT loss equivalent to the cross-entropy loss function. The negative samples would indeed follow the Zipfian distribution of the data instead of the uniform distribution implicitly used in the cross-entropy loss. We hypothesize that regularizing based on the unigram distribution instead of the uniform vocabulary distribution may also explain the performance improvement of CWT. 

We also invite the reviewer to read **Appendix A** that gives some insight on the reasons of the effectiveness of our method.

Reviewer comment (Evaluation):
I'm curious about the experimental results in Table 2. I wonder why there is a significant improvement on all test sets except the CB test set. Why does the Headless model perform so much worse than the vanilla model on the CB test set? Could the authors provide a detailed analysis and explanation for this? Is there anything special about the CB test set?

Author response:
We analyzed the reason for this discrepancy more thoroughly. We found that the CB validation set we used for our analysis **only counts 56 examples**. While the standard deviation of the performance of each model is close to 1, making the score gap meaningful across seeds, it only requires a correct classification on two more samples to explain such a gap with this sample size. We believe that this explains the observed discrepancy.

Reviewer comment (Presentation):
In Table 3, why is there no output for the Headless model's perplexity (ppl) on the validation set?

Author response:
We did not report perplexity for headless models during our initial training, as it would have required us adapting the forward procedure to use the transposed input embeddings matrix as a projection head during validation steps only. Moreover, we did not realize *before training* that we would need to put the model through a fine-tuning phase after pre-training. Finally, as shown by the evaluation on the Lambada dataset, the perplexity level of the headless model before fine-tuning is not really relevant for comparison with the other models.

We would very much like to compute a perplexity score now to answer your question, but the OpenWebText2 dataset is not hosted anymore (https://huggingface.co/datasets/the_pile_openwebtext2/discussions/5), which does not let us do so.