ICML 2024: Paper Review #5

Jesse Farebrother, Jordi Orbay, Quan Vuong, Adrien Ali Taiga, Yevgen Chebotar, Ted Xiao, Alex Irpan, Sergey Levine, Pablo Samuel Castro, Aleksandra Faust, Aviral Kumar, Rishabh Agarwal

Scaling up models has been less straightforward in Reinforcement Learning compared to its overwhelming success in supervised settings. In previous work, several examples have been pointed out where increasing model capacity or training iterations eventually starts to decrease performance.

This paper investigates replacing the usual MSE loss on the scalar value function with cross-entropy loss, after quantizing the value into a number of bins. A previous approach distributed the probability mass into the two closest bins such that the expectation matches the original value (“Two-Hot encoding”). The authors propose smearing with a Gaussian and initially attempt adapting its width to the bin size (essentially proposing a form of Six-Hot encoding) but discover that the absolute variance of the Gaussian, rather than the number of bins in its effective support, is the relevant hyperparameter.

The results demonstrate improved performance over alternate distributions as well as standard regression and achieve monotonic scaling in the cases where regression did not. Ablation demonstrates that the use of the cross-entropy loss is critical compared to just lifting the representation from a scalar to a distribution.

Zeyuan Allen-Zhu, Yuanzhi Li

This paper is part of a series on how large language models (LLMs) acquire various capabilities, using controlled experiments with purely synthetic data.

It studies the conditions under which LLMs learn knowledge in a way that, for example, allows them to answer questions, rather than just repeat sentences verbatim.

One key finding is that it is necessary for a fact to appear in the training data in multiple variations (for example, sentence permutations or translations). The understanding of a fact about X is demonstrated not only through a model’s ability to answer questions about X but also through probing that shows that such facts are only directly associated with X when they appear with variations, otherwise the association is with the whole sentence.

The paper also shows that instruction-fine tuning with questions cannot recover knowledge that hasn’t been presented with variations during pretraining. Conversely it shows that adding questions during pretraining can improve performance. The authors therefore recommend this approach, along with augmenting the pretraining data with sentences rewritten by an auxiliary model.

The controlled “physics-like” approach is a fascinating way to study the behaviour of LLMs.