通过非线性低秩分支提升Transformer训练速度,不增加太多参数开销。
NOBLE: Accelerating Transformers with Nonlinear Low-Rank Branches
- 在Transformer线性层中加入可学习的非线性低秩分支,永久嵌入架构中。
- 训练速度最高提升1.47倍,仅增4%参数,净时钟加速达1.22倍。
- 适合大规模语言模型、ViT等预训练任务,尤其关注训练效率提升。
我们提出NOBLE(Nonlinear lOw-rank Branch for Linear Enhancement),一种在Transformer线性层中添加非线性低秩分支的架构增强方法。与LoRA等参数高效微调方法不同,NOBLE专为从头预训练设计,其分支是架构的永久组成部分,而非冻结权重上的适配器。该分支计算σ(xWdown)Wup,其中σ为可学习非线性激活。我们测试多种激活函数,发现CosNet(两层余弦非线性,含可学习频率与相位,中间有瓶颈空间线性投影)表现最佳。实验表明,NOBLE在极小额外开销下实现显著提速:达到基线评估损失的训练步骤最多减少32%(最快提升1.47倍/步),额外参数仅增加4%,单步耗时增加7%,最终净墙钟速度提升达1.22倍。在LLM(250M和1.5B参数)、BERT、VQGAN和ViT上均一致提升训练效率。但发现混洗/裁剪增强(Mixup/CutMix)会干扰NOBLE在ImageNet分类中的优势,禁用后ViT性能也提升,可能因正则化鼓励平滑拟合,而NOBLE更擅长捕捉目标函数的尖锐特征。
原文摘要 · Abstract (English)
We introduce NOBLE (Nonlinear lOw-rank Branch for Linear Enhancement), an architectural augmentation that adds nonlinear low-rank branches to transformer linear layers. Unlike LoRA and other parameter-efficient fine-tuning (PEFT) methods, NOBLE is designed for pretraining from scratch. The branch is a permanent part of the architecture as opposed to an adapter for finetuning on top of frozen weights. The branch computes σ(xWdown)Wup where σ is a learnable nonlinearity. We evaluate several activation functions and find that CosNet, a two-layer cosine nonlinearity with learnable frequency and phase with a linear projection in between them in the bottleneck space, performs best. NOBLE achieves substantial improvements with minimal overhead: up to 1.47x step speedup to reach baseline eval loss (up to 32% fewer training steps), with as low as 4% additional parameters and 7% step time overhead, resulting in up to 1.22x net wallclock speedup. Experiments on LLMs (250M and 1.5B parameters), BERT, VQGAN, and ViT consistently show improved training efficiency. We identify one caveat: Mixup/CutMix augmentation interferes with NOBLE's benefits in Imagenet classification along with other stochastic augmentations, but when disabled, ViT also improves. This discrepancy is possibly explained by regularization techniques that encourage smoother fits to the target function while NOBLE may specialize more in sharper aspects of the target function.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。