arXiv:2502.01637cs.CLcs.LG2025-02NeurIPS被引 18

通过预计算词组嵌入,让小模型在低资源下超越大模型。

Scaling Embedding Layers in Language Models

  • 用频繁词组的上下文嵌入替代部分原词表,提升语义表达。
  • 10亿参数模型用一半算力和内存,性能超19亿参数基线。
  • 适合资源受限场景,如移动端或实时推理应用。

我们提出SCONE(可扩展、上下文化、卸载、n-gram嵌入),一种扩展输入嵌入层的新方法,以提升语言模型性能。为避免解码成本增加,SCONE保留原有词表,同时引入一组高频n-gram的嵌入表示。这些嵌入为每个输入词元提供上下文化表征,并通过独立模型在训练期间学习。训练完成后,嵌入被预先计算并存储于非加速器内存中;推理时查询嵌入的复杂度极低,对延迟影响微乎其微。SCONE支持两种新扩展策略:增加n-gram嵌入数量,以及扩大用于学习它们的模型规模,两者均在推理阶段保持固定的加速器使用量(以浮点运算次数和内存计)。实验表明,同时扩展这两个方面,使一个仅10亿参数驻留于加速器的模型,在多种语料上超越19亿参数基线,且推理时仅需约一半的浮点运算和加速器内存。

原文摘要 · Abstract (English)

We propose $SCONE$ ($S$calable, $C$ontextualized, $O$ffloaded, $N$-gram $E$mbedding), a new method for extending input embedding layers to enhance language model performance. To avoid increased decoding costs, $SCONE$ retains the original vocabulary while introducing embeddings for a set of frequent n-grams. These embeddings provide contextualized representation for each input token and are learned with a separate model during training. After training, embeddings are precomputed and stored in off-accelerator memory; during inference, querying them has minimal impact on latency due to the low complexity of embedding lookups. $SCONE$ enables two new scaling strategies: increasing the number of n-gram embeddings and scaling the model used to learn them, both while maintaining fixed accelerator usage during inference (in terms of FLOPS and memory). We show that scaling both aspects enables a model with 1B accelerator-resident parameters to outperform a 1.9B-parameter baseline across diverse corpora, while using only about half the FLOPS and accelerator memory during inference.

嵌入扩展模型压缩推理优化

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。