AFA-LoRA: Enabling Non-Linear Adaptations in LoRA with Activation Function Annealing

TL;DR

AFA-LoRA introduces Activation Function Annealing to inject temporary nonlinearity into LoRA adapters, enabling richer early learning while ensuring final linear mergeability for seamless deployment. The method formalizes a time-dependent activation $\sigma_{ ext{AFA}}(x; t) = \beta(t) \cdot \sigma(x) + (1 - \beta(t)) \cdot x$ and forward computation $F_{ ext{AFA}}(x; t) = W_2(t) \cdot \sigma_{ ext{AFA}}(W_1(t) x; t)$, with $eta(0)=1$, $eta(T)=0$. The authors demonstrate substantial performance gains across supervised fine-tuning, reinforcement learning (GRPO), and speculative decoding (Eagle), reducing the gap to full fine-tuning without sacrificing mergeability, and provide ablations showing a 30% decay schedule as a robust default. This work offers a practical path to more expressive yet deployable PEFT adapters, with broad implications for efficient adaptation in large language models. The approach is supported by theoretical formalization, optimization-landscape intuition, and diverse empirical validation.

Abstract

Low-Rank Adaptation (LoRA) is a widely adopted parameter-efficient fine-tuning (PEFT) method. However, its linear adaptation process limits its expressive power. This means there is a gap between the expressive power of linear training and non-linear training. To bridge this gap, we propose AFA-LoRA, a novel training strategy that brings non-linear expressivity to LoRA while maintaining its seamless mergeability. Our key innovation is an annealed activation function that transitions from a non-linear to a linear transformation during training, allowing the adapter to initially adopt stronger representational capabilities before converging to a mergeable linear form. We implement our method on supervised fine-tuning, reinforcement learning, and speculative decoding. The results show that AFA-LoRA reduces the performance gap between LoRA and full-parameter training. This work enables a more powerful and practical paradigm of parameter-efficient adaptation.

Figures (3)

• Figure 1: Architectural comparison of (a) Standard LoRA and (b) AFA-LoRA. The annealed activation function $\phi$ is placed beside the $A$ and $B$ matrices.
• Figure 2: Illustration of ReLU-to-linear activation annealing and the decay schedule for $\beta(t)$ during training.
• Figure 3: Llama3.1-8b training loss-epoch curves for Eagle-1 on the ShareGPT dataset