NOBLE: Accelerating Transformers with Nonlinear Low-Rank Branches

TL;DR

NOBLE introduces nonlinear low-rank branches into transformer layers, significantly improving training efficiency and speed with minimal additional parameters across various models, including LLMs, BERT, VQGAN, and ViT.

Contribution

This work presents NOBLE, a novel architectural augmentation that adds nonlinear low-rank branches for pretraining transformers from scratch, outperforming existing PEFT methods in efficiency.

Findings

Achieves up to 1.47x step speedup with minimal overhead.

Improves training efficiency across multiple model types.

Discovered that certain augmentations interfere with NOBLE's benefits.

Abstract

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 {\sigma}(xWdown)Wup where {\sigma} 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.