Headless Language Models: Learning without Predicting with Contrastive Weight Tying

The paper proposes a simple method for pre-training LMs with contrastive loss. Instead of evaluating logits for all vocabulary, only embeddings from the batch are used. Such an approach allowed us to pre-train LMs more effectively by reducing memory constraints. Surprisingly, this approach also performed well for autoregressive language modeling. For both autoregressive and masked language modeling, the proposed method showed improved results on downstream tasks, as well as improved training efficiency with respect to training hours.

This paper proposes to replace the cross-entropy minimization by a new pre-training objective in language models: Contrastive Weight Tying (CWT). It consists in contrastively predicting the representation of the missing word. Instead of projecting the representation output by the model into the vocabulary, and using a softmax, the goal is to predict directly this static embedding of the reference word, which the model is trained for by (1) using a contrastive objective with the rest of the batch as negative examples (2) using input static embeddings as reference output embeddings, which is traditionally referred as weight tying. After a presentation of previous work, the paper introduces the method, consisting in weight tying, the contrastive objective with no projection, and how to adapt it for text generation, plus computational requirements. Experiments are performed on a Monolingual Encoder (GLUE), Monolingual Decoder (LAMBADA/LM Evaluation Harness), and Multilingual Encoder (XNLI Benchmark), with CWT showing improvements on all settings compared to the model trained classically. Lastly, the paper presents experiments on vocabulary size, as expanding the vocabulary is less costly and may provides some advantages with CWT.

This paper proposes replacing the standard cross-entropy-based language modeling objective (classification over token vocabulary) with a "contrastive weight tying" objective, which is more memory and compute efficient during training. Computing the full cross-entropy loss requires loading and multiplying the final model embeddings by a $d_{model} \times V$ matrix, where $V$ is the vocabulary size (often $\geq 30k$), which can be quite expensive (especially for small models, with few layers and $d_{model} \ll V$). Instead, this paper proposes using a contrastive objective, where the targets are the input token embeddings for the masked tokens ("weight tying"), and the negative samples are taken from the other "masked" tokens. It shows in experiments that the proposed method is generally more efficient, and attains better performance, than the standard cross-entropy-based training. Lastly, to allow for generating text (using next token prediction) with these "headless" models, the paper proposes performing fine-tuning with a LM head for a small number of tokens ($< 2$% of pre-training dataset).

This study introduces an innovative approach to self-supervised pre-training of language models. Instead of predicting probability distributions over token vocabularies, the method focuses on reconstructing input embeddings contrastively using Contrastive Weight Tying (CWT). This approach is applied to pretrain Headless Language Models in both monolingual and multilingual contexts, offering practical advantages. It is significantly more compute-efficient, data-efficient, and performant than classical predictive methods. The paper's contributions include the introduction of a new pretraining objective, the pretraining of encoder and decoder models for English and a multilingual encoder model, demonstrating the benefits of headless training, and exploring the effects of pretraining hyperparameters on downstream performance.

We thank all the reviewers for their very valuable feedback. We first address concerns that were shared by several reviewers:

On the lack of novelty

We are well aware that our CWT objective function is not novel nor creative as a learning objective with respect to general ML. What we wish to stress out in this article is that our work distinguishes itself from closely related literature by exploring a contrastive objective as a replacement for cross-entropy in the context of pre-training token-level language models in a self-supervised fashion. Recent concurrent work in this specific domain usually apply contrastive objectives as auxiliary losses that are added to cross-entropy (e.g. https://arxiv.org/pdf/2111.04198.pdf or https://arxiv.org/pdf/2202.06417.pdf). We believe our work is novel in that specific line of work. We will rephrase some statements made in the paper to clarify this point in an updated version.

On the additional cost of causal fine-tuning

We agree that the necessary fine-tuning of the LM classification head slightly complexifies the pre-training process from a technical/coding perspective. However, we only needed to use a small fraction of the training compute (1.7%) to obtain a performant language model. In Figure 5a and 5b, we include this overhead on the x-axis for the fine-tuned model, to demonstrate that the added computation complexity is actually negligible.

On understanding the performance gap

We attempt to answer this question in Appendix A. To summarize, we argue that cross-entropy pushes semantically close words apart at each training step, by incidentally using them as negative samples. Our CWT objective instead pushes close words apart only if they appear in the same batch, which is less likely. It results in a slightly higher average cosine-similarity for synonyms for headless models, which indicates more meaningful representations.

Another argument that is not directly mentioned in our paper is that our objective addresses the problem of updating rare and unused token embeddings at each step, which is a flaw of the cross-entropy loss (https://arxiv.org/abs/1907.12009).

On scale

We mention the limitations related with model scale in Appendix B. This paper was written in a context of limited computation capabilities, and extending a new pre-training approach to larger scales is extremely costly in terms of compute. Validating our method on Llama or ChatGLM, as suggested by Reviewer A6os, would have implied pre-training such models from scratch, which is inaccessible for our institution. For instance, the smallest Llama-2 model (7B parameters) took 184k A100 GPU hours to pretrain, which can be roughly estimated to cost ~200k$.

On ablation study

We want to clarify that we used classical weight-tying in all our “vanilla” models trained with cross-entropy, as we believe it is now standard in the literature. With regard to the other ablation Reviewer ccZo suggests, i.e. using an additional static embedding matrix for negative samples apart from the one used for input, we did not run this experiment as it results in adding a substantial count of parameters to the model, which could have biased the downstream performance comparison. Nevertheless, we are currently running the corresponding experiment at small scale, and we will communicate the results as soon as possible.

About trying different numbers of negative samples (as suggested by Reviewer ccZo), we think this idea is worth exploring, as we mentioned in the footnote of page 4. At submission time, we decided to leave it for future work due to space limitations, as it would have required a hyperparameter search. We now have preliminary results that indicate that it is possible to select fewer negative examples without loss of performance, up to a certain point.

Based upon the numerous comments and suggestions made by the reviewers, we updated our paper in several aspects:

• Ablation study : We added an ablation study for small-scale models, where we compare our approach with two variants (inspired from reviewer ccZo) : one where the loss is the addition of cross-entropy and CWT, and one contrastive-only approach that does not use weight-tying, i.e. that uses external static embeddings as targets (NCE wo/ WT). Here are the results, that can also be found on page 9 :

• New references : We updated the related work section (pages 2 and 3) by including most of the references suggested by revierwer CHGU. We contextualized our claims of novelty, and we also cited the appropriate papers along our article.

• Extended Appendix A (explanation for performance gain) : We tried to provide a deeper intuition about the data-efficiency gain of our method (see page 14).

We thank again all the reviewers for their suggestions that add great value and completeness to our article. We are open to additional suggestions and comments before the end of the discussion period.

Dear reviewers,

The reviewer-author discussion period will come to a conclusion in under 4 hours. If you have any further thoughts or comments, don't hesitate to share them with us. We remind you that we conducted additional work following your reviews, and we uploaded a revised version of the article. We sincerely thank you for your dedication and time in reviewing our submission.

Warmest regards, Author(s)

This paper proposes replacing the standard cross-entropy-based (large) language modeling objective (classification over token vocabulary) with a "contrastive weight tying" objective, which is more memory and compute efficient during training. Computing the full cross-entropy loss requires loading and multiplying the final model embeddings by a matrix, which can be quite expensive (especially for small models, with few layers, a small dimensionality and a large vocabulary). Instead, this paper proposes using a contrastive objective, where the targets are the input token embeddings for the masked tokens ("weight tying"), and the negative samples are taken from the other "masked" tokens. It shows in experiments that the proposed method is generally more efficient, and attains better performance, than the standard cross-entropy-based training. Lastly, to allow for generating text (using next token prediction) with these "headless" models, the paper proposes performing fine-tuning with a LM head for a small number of tokens (<2% of pre-training dataset).

Strengths

• The proposed method can reduce the number of trainable parameters meaningfully for small models with large vocabulary sizes (e.g., by my calculations, reduce # parameters of BERT-base by ~18%).
• The proposed method generally yields better accuracy on downstream tasks than standard LM training (across mono-lingual encoder experiments with BERT-base, mono-lingual decoder experiments with Pythia-70m, and multi-lingual encoder experiments with distilled multi-lingual BERT).
• This method decouples the training speed from the number of tokens in the vocabulary, which allows for choosing the vocabulary size that attains best performance.

Weaknesses

• The proposed method is not very novel --- straightforward application of contrastive loss function to LM training.
• The experiments are relatively small-scale, only considering models < 140m parameters. Unclear if any of the benefits of this method hold for larger models. The paper does not discuss that for larger models, the 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.
• 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 (although some were presented during the discussion)

Accept (poster)