Overparameterized large-scale language models have impressive generalization performance of in-context few-shot learning. However, most language models allocate the same amount of parameters or computation to each token, disregarding the complexity or importance of the input data. We argue that in language model pretraining, a variable amount of computation should be assigned to different tokens, and this can be efficiently achieved via a simple routing mechanism. Different from conventional early stopping techniques where tokens can early exit at only early layers, we propose a more general method that dynamically skips the execution of a layer (or module) for any input token with a binary router. In our extensive evaluation across 24 NLP tasks, we demonstrate that the proposed method can significantly improve the 1-shot performance compared to other competitive baselines only at mild extra cost for inference.

翻译：过度参数化的大规模语言模型在上下文少样本学习方面具有出色的泛化性能。然而，大多数语言模型为每个标记分配相同数量的参数或计算量，忽略了输入数据的复杂性或重要性。我们认为，在语言模型预训练中，应为不同标记分配可变的计算量，这可以通过简单的路由机制高效实现。与传统的早期停止技术（标记只能在早期层提前退出）不同，我们提出了一种更通用的方法，该方法通过二元路由动态跳过任何输入标记的层（或模块）执行。在涵盖24个NLP任务的广泛评估中，我们证明，所提方法相比其他竞争基线仅需少量额外推理成本，即可显著提升1-shot性能。