Lightweight Surrogate-Assisted Language Model Pretraining

Keywords: Transformers, Pretraining, Knowledge Distillation, Loss Objectives, Efficiency

TL;DR: We propose a novel loss objective that leverages small teacher reference models for substantial pretraining efficiency gains.

Abstract: Standard causal language model pretraining uses a single-label cross-entropy objective that ignores the existence of multiple valid next-token continuations, resulting in sample inefficiency. Our initial analyses of language model outputs reveal that only a handful of tokens carry meaningful signal compared to the rest of the vocabulary. Predicated on this finding, we introduce a multi-label pretraining objective that modifies the loss to append a small set of context-conforming auxiliary tokens selected by a lightweight surrogate language model. Distinct from existing knowledge distillation methods, the surrogate is used only for token selection and suggestion rather than full distribution matching. We empirically show that adopting this multi-target stance achieves superior performance on benchmarks with notably less FLOPs utilization and tokens in our tested setting.

Submission Number: 60