NOBLE: Accelerating Transformers with Nonlinear Low-Rank Branches

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 σ(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.

• The paper demonstrates that integrating a nonlinear low-rank CosNet branch within Transformer layers accelerates pretraining by capturing high-frequency residuals.
• Empirical results report a 1.34–1.37x reduction in training steps and a 1.22x wallclock speedup with only a 4–12% parameter increase.
• The approach permanently augments the model's linear pathway, enhancing representational capacity while introducing a fixed inference cost and sensitivity to aggressive regularization.

NOBLE: Accelerating Transformers with Nonlinear Low-Rank Branches

Introduction

The NOBLE (Nonlinear lOw-rank Branch for Linear Enhancement) architecture introduces an augmentation to Transformer linear layers by embedding a nonlinear low-rank computational branch, designed expressly for pretraining from scratch rather than parameter-efficient fine-tuning. Unlike LoRA and its derivatives that adapt frozen weights post-pretraining for downstream transfer, NOBLE permanently integrates the nonlinear branch into the model, expanding its representational capacity and complementing the main linear pathway with minimal overhead. The principal mathematical innovation is a bypass branch that applies a learnable nonlinearity—most notably, a novel CosNet activation—over a low-rank subspace, with the intent of capturing feature variations that are not efficiently expressible through conventional affine transformations alone.

Architectural Formulation

NOBLE generalizes a standard linear layer $f(x) = xW + b$ to

$$f_\text{NOBLE}(x) = xW + b + \sigma(xW_\text{down})W_\text{up}$$

where $$W_\text{down} \in \mathbb{R}^{d_\text{in} \times r}$$, $W_\text{up} \in \mathbb{R}^{r \times d_\text{out}}$, $r \ll \min(d_\text{in}, d_\text{out})$, and $\sigma$ is a nonlinear function. The architectural premise is the inclusion of a nonlinear channel (the NOBLE branch), parameterized by the CosNet nonlinearity, in all linear projections—attention (Q, K, V, out) and feedforward layers.

Figure 1: Architectural comparison of NOBLE with CosNet and related low-rank adaptation techniques.

CosNet employs a two-layer cosine nonlinearity with learnable frequency and phase in the bottleneck, coupled via a small mixing matrix, surpassing alternatives such as ReLU, GELU, and their variants. The main branch and the bypass branch are initialized so that the latter has negligible influence at initialization (i.e., near-zero initialization for $W_\text{up}$), then trained jointly end-to-end with task-specific learning rate scaling for stability and rapid adaptation.

Empirical Results—Language Modeling and Pretraining Efficiency

NOBLE yields substantial reduction in training steps to reach target eval loss relative to vanilla Transformers, with only modest increases in parameter count and per-step computation. On LLM pretraining (OpenWebText, 250M and 1.5B parameter settings), models reach baseline's best eval loss ($2.51$) in $143$–$154$k steps, as opposed to $196$k for the baseline—a 1.34–1.37$\times$ step speedup.

Figure 2: Eval loss curves for 1.5B LLM pretraining on OpenWebText, demonstrating that NOBLE attains baseline loss substantially faster.

This efficiency persists across scale and rank hyperparameters (ranks 64–256). The wallclock speedup net of overhead is up to 1.22$\times$. At all tested ranks, NOBLE achieves lower converged eval loss (improvements of $0.02$–$0.07$). These findings are robust to masked language modeling (BERT), image token autoregressive modeling, and persist in unconditional tasks when strong data augmentations are disabled.

Figure 3: Eval loss curves for autoregressive language modeling at two scales, with NOBLE outperforming the baseline throughout training.

Vision and Image Token Modeling

For vision, NOBLE enhances both ViT-Small classification (ImageNet-1k) and generative image token modeling. The architectural improvement, however, is dependent on the augmentation regime: with aggressive label smoothing augmentations (Mixup/CutMix), the NOBLE branch’s benefit vanishes on validation accuracy and is subdued on training loss. Disabling these augmentations reinstates a 5% lower training loss, but also exposes a tendency for overfitting, with minimal gain or slight regression on eval loss.

Figure 4: ViT-S training loss and validation accuracy, with and without Mixup/CutMix. NOBLE improves optimization when augmentations are off but does not generalize better.

In autoregressive image token modeling (VQGAN tokens over ImageNet-1k), NOBLE confers a consistent benefit by lowering eval loss at all tested ranks.

Nonlinearity and Activation Function Design

An extensive ablation on the bottleneck nonlinearity reveals the necessity of periodic, symmetric, and non-saturating activations in the low-rank context.

Figure 5: Comparison of eval loss curves for bottleneck activations; CosNet outperforms all linear and nonlinear alternatives.

CosNet, the recommended activation, is a two-layer sandwich of cosine functions with learnable frequency and phase, separated by a mixing matrix, achieving the deepest training curve reductions. ReLU-like and tanh activations induce moderate or minimal gains, with their inherent asymmetry or saturation limiting expressiveness in a low-rank subspace. CosNet’s learnable frequency modulation confers the ability to fit high-frequency residual components that complement the main linear trend.

Theoretical Implications

NOBLE’s benefit is interpreted spectrally: the primary dense linear pathway fits low-frequency, smooth targets, while the nonlinear cosine branch targets high-frequency, non-smooth residuals. When regularization or data augmentation enforces smoothness (label smoothing, Mixup, CutMix, dropout, etc.), the high-frequency signal is suppressed or excised, leading to diminished NOBLE efficacy. In canonical pretraining and generative scenarios (where label/target structure is unperturbed and less regularized), the architecture yields its maximal gain.

The computational tradeoff is favorable: $\sim$4–12% parameter and $\sim$7–12% per-step cost yields a step reduction of 21–32%. However, unlike PEFT adapters, this induces a permanent inference cost.

Practical and Future Directions

NOBLE is a robust, easy-to-implement architectural augmentation for accelerating pretraining of Transformers, particularly effective when the target distribution retains non-smooth structure and where inference cost is not tightly constrained. Its positive empirical profile across language and image pretraining tasks, and explicit independence from PEFT strategies, make it relevant for foundational model training. Future work may explore selective application (layer or blockwise), scaling laws beyond 1.5B parameters, interactions with a broader set of augmentations and optimizers, and extension to structured prediction tasks in vision and beyond.

Conclusion

NOBLE demonstrates that permanent nonlinear low-rank branches, particularly with periodic CosNet activation, can accelerate transformer pretraining by over 1.4$\times$ in steps, and 1.2$\times$ on wallclock, for only moderate overhead. These gains extend to both language and vision settings where high-frequency structure in the target function is present. Limitations include permanent inference cost and diminished returns under aggressive regularization or input/label smoothing. The results invite further explorations into hybrid linear–nonlinear architectural designs for high-capacity, efficient pretraining beyond the fine-tuning-centric paradigm.

---

Reference: "NOBLE: Accelerating Transformers with Nonlinear Low-Rank Branches" (2603.06492)