arXiv:2601.10639cs.LG2026-01被引 12

用嵌入模块替代部分前馈网络,实现高效高容量模型训练。

STEM: Scaling Transformers with Embedding Modules

  • 用静态嵌入查表替代前馈网络上投影层,减少计算开销。
  • 在350M和1B模型上提升3-4%准确率,减少约三分之一的参数访问。
  • 支持可解释的知识编辑与注入,适合需要可控推理的场景。

细粒度稀疏性在不增加每标记计算量的前提下提升参数容量,但常面临训练不稳定、负载不平衡和通信开销问题。我们提出STEM(基于嵌入模块的Transformer扩展),一种静态的、基于标记索引的方法:将前馈网络的上投影替换为局部嵌入查找,同时保持门控和下投影稠密。该方法消除了运行时路由,支持异步预取的CPU卸载,并将容量与每标记浮点运算及跨设备通信解耦。实验证明,即使在极端稀疏条件下,STEM仍能稳定训练。其下游性能优于稠密基线,同时降低每标记浮点运算和参数访问次数(约消除1/3的前馈网络参数)。STEM学习具有大角度分布的嵌入空间,增强了知识存储能力。更有趣的是,这种增强的知识容量带来了更好的可解释性:因嵌入具有标记索引特性,无需修改输入文本或额外计算即可简单实现可解释的知识编辑与注入。此外,STEM提升了长序列表现:随着序列长度增长,更多不同参数被激活,实现实际测试时的容量扩展。在3.5亿和10亿参数规模下,整体准确率提升约3-4%,在知识与推理密集型基准(ARC-Challenge、OpenBookQA、GSM8K、MMLU)上表现尤为突出。总体而言,STEM是一种有效扩展参数记忆的同时提升可解释性、训练稳定性与效率的方法。

原文摘要 · Abstract (English)

Fine-grained sparsity promises higher parametric capacity without proportional per-token compute, but often suffers from training instability, load balancing, and communication overhead. We introduce STEM (Scaling Transformers with Embedding Modules), a static, token-indexed approach that replaces the FFN up-projection with a layer-local embedding lookup while keeping the gate and down-projection dense. This removes runtime routing, enables CPU offload with asynchronous prefetch, and decouples capacity from both per-token FLOPs and cross-device communication. Empirically, STEM trains stably despite extreme sparsity. It improves downstream performance over dense baselines while reducing per-token FLOPs and parameter accesses (eliminating roughly one-third of FFN parameters). STEM learns embedding spaces with large angular spread which enhances its knowledge storage capacity. More interestingly, this enhanced knowledge capacity comes with better interpretability. The token-indexed nature of STEM embeddings allows simple ways to perform knowledge editing and knowledge injection in an interpretable manner without any intervention in the input text or additional computation. In addition, STEM strengthens long-context performance: as sequence length grows, more distinct parameters are activated, yielding practical test-time capacity scaling. Across 350M and 1B model scales, STEM delivers up to ~3--4% accuracy improvements overall, with notable gains on knowledge and reasoning-heavy benchmarks (ARC-Challenge, OpenBookQA, GSM8K, MMLU). Overall, STEM is an effective way of scaling parametric memory while providing better interpretability, better training stability and improved efficiency.

Transformer稀疏性嵌入模块可解释性

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