NOBLE: 비선형 저랭크 분기를 이용한 트랜스포머 가속화

NOBLE: Accelerating Transformers with Nonlinear Low-Rank Branches

본 논문에서는 NOBLE (Nonlinear lOw-rank Branch for Linear Enhancement)이라는 새로운 아키텍처 확장 기법을 소개합니다. NOBLE은 트랜스포머의 선형 레이어에 비선형 저랭크 분기를 추가합니다. 기존의 LoRA 및 기타 효율적인 파라미터 튜닝 (PEFT) 방법과는 달리, NOBLE은 처음부터 사전 훈련하는 것을 목표로 설계되었습니다. 이 분기는 기존 가중치를 동결하고 그 위에 어댑터를 추가하는 방식이 아니라, 아키텍처의 고정된 부분입니다. 이 분기는 σ(xWdown)Wup 형태의 계산을 수행하며, 여기서 σ는 학습 가능한 비선형 함수입니다. 다양한 활성화 함수를 평가한 결과, CosNet이라는 두 계층의 코사인 비선형 함수 (bottleneck 공간에서 선형 투영을 포함하며, 학습 가능한 주파수와 위상을 가짐)가 가장 우수한 성능을 보였습니다. NOBLE은 최소한의 오버헤드로 상당한 성능 향상을 제공합니다. 구체적으로, 기준 평가 손실에 도달하는 데 최대 1.47배의 속도 향상을 얻을 수 있으며 (최대 32% 더 적은 훈련 단계), 추가 파라미터는 4% 미만, 훈련 단계 시간 오버헤드는 7% 미만이며, 결과적으로 최대 1.22배의 전체 훈련 속도 향상을 달성했습니다. LLM (250M 및 1.5B 파라미터), BERT, VQGAN 및 ViT에 대한 실험 결과, 훈련 효율성이 일관되게 향상되는 것을 확인했습니다. 한 가지 주의사항은 Mixup/CutMix 증강 기법이 ImageNet 분류 작업에서 NOBLE의 장점을 저해하며, 다른 확률적 증강 기법도 유사한 문제를 일으킬 수 있다는 것입니다. 하지만 이러한 증강 기법을 비활성화하면 ViT의 성능도 향상됩니다. 이러한 차이는 정규화 기술이 목표 함수에 더 매끄러운 피팅을 장려하는 반면, NOBLE은 목표 함수의 더 뚜렷한 부분에 더 특화될 수 있기 때문일 수 있습니다.

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.