ICML 2024: Paper Review #8

Fabian Gloeckle, Badr Youbi Idrissi, Baptiste Rozière, David Lopez-Paz, Gabriel Synnaeve

This paper introduces a novel approach to LLM pre-training, challenging the conventional next-token prediction paradigm. The authors propose a non-autoregressive multi-token prediction loss, aiming to enhance model performance whilst maintaining parallelisability. Intuitively such a multi-token loss should help as it makes the pre-training objective more similar to the autoregressive text generation in downstream tasks.

The research examines models ranging from 0.3 to 13 billion parameters, revealing that the multi-token loss becomes increasingly beneficial as model size grows. Notably, this approach yields significant improvements in coding benchmarks, such as MBPP and HumanEval. The fact that smaller models (under 1 billion parameters) experience performance degradation, even on coding benchmarks, with this method may explain why it hasn’t been previously explored.

Beside these gains on coding benchmarks, I think their findings on byte-level training are particularly promising, as the multi-token loss nearly bridges the performance gap between byte and standard tokenisers. This development could potentially pave the way for ‘tokeniser-free’ models in the future.

Beside those empirical gains, the paper proposes two intuitions that illustrate why a multi-token loss should indeed help with auto-regressive text generation which I found helpful.

It’s worth noting that the claimed reduction in inference time refers to wall-clock time rather than FLOPs as the gains stem from increased parallelisability and speculative decoding, not from fundamental computational efficiency improvements.

Alexander Hägele, Elie Bakouch, Atli Kosson, Loubna Ben Allal, Leandro Von Werra, Martin Jaggi

Determining the optimal training length for a given model is a crucial aspect of generative AI research. This line of inquiry, culminated (in the public domain) with the Chinchilla paper, which demonstrated that the learning rate schedule is an important factor. ^[1] Until recently, cosine annealing was the scheduler of choice. Despite its strong performance, it has a major drawback: one must specify the decay scale upfront, and proportional to the number of training steps to get state of the art performance. This makes it difficult to reuse model checkpoints in different experiments.

This paper investigates a simple learning rate schedule: constant followed by cool-down. The results are promising, with the final performance of models being at least as good as with cosine annealing. In addition, the checkpoints from the constant learning rate phase can be reused for experiments of different durations. Interestingly, the Llama 3.1 paper, released during ICML, describes a very similar approach. Though the opted for a linear learning rate decay during cool-down rather than the 1-sqrt version proposed in the paper. ^[2]

Beyond scaling law experiments, the new scheduler enables checkpoint sharing for investigations into different data mixtures in a curriculum learning setup. Again, Llama 3.1 leveraged this too and explored different data mixtures during the cooldown period, ie. re-using the last checkpoint from the constant learning rate phase.

The authors also analysed Stochastic Weight Averaging and the recently proposed schedule free optimiser, which promise optimal performance throughout the training. Whilst these yield strong performance during all training iterations, the paper finds that they cannot quite match the performance of the proposed constant and cool-down method. Additionally, these alternatives would only offer compute savings for scaling law runs, not for experiments investigating data mixtures or curriculum learning stages.

^{[1] Training Compute-Optimal Large Language Models}

^{[2] The Llama 3 Herd of Models}