Exact merging and gradients for a low-rank adapter

Fix dimensions and matrices as in the definition of a low-rank adapter. Let $m\ge 1$, let $x_1,\dots ,x_m\in {\mathbb{R}}^k$, and let each ${\mathrm{\ell }}_j$ be a finite real-arithmetic loss program with scalar real output wherever its primitive calls are defined. Each primitive map is $C^1$ on a declared open domain or is piecewise $C^1$ with finitely many open branch regions and a $C^1$ formula on each region; when defined, a boundary primal value is specified separately. Here $C^1$ means totally differentiable with continuous derivative. Define the finite empirical loss on a weight matrix $U\in {\mathbb{R}}^{d\times k}$ by the following formula wherever all the programs return scalar values:

$\mathcal{L}(U)=\frac{1}{m}{\sum }_{j=1}^m{\mathrm{\ell }}_j(U{x}_j)\mathrm{.}$

At any $W=W_0+cBA$ where all primitive calls have values and every piecewise primitive call is on a unique strict branch, write $G={\mathrm{\nabla }}_W\mathcal{L}(W)$ for the Frobenius gradient, defined by $D\mathcal{L}(W)[H]=⟨G,H{⟩}_F$ for every $H\in {\mathbb{R}}^{d\times k}$, where $⟨X,Y{⟩}_F=\mathrm{tr}(X^{\mathsf{T}}Y)$. Then the factor gradients are

${\mathrm{\nabla }}_B\mathcal{L}(W_0+cBA)=cG{A}^{\mathsf{T}},{\mathrm{\nabla }}_A\mathcal{L}(W_0+cBA)=c{B}^{\mathsf{T}}G\mathrm{.}$

At every successful step $t$ of the finite-budget adaptation algorithm, write its ordered minibatch as $I_t=(i_{t,1},\dots ,i_{t,m_t})$, retaining repeated indices, and define

${\mathcal{L}}_t(U)=\frac{1}{m_t}{\sum }_{s=1}^{m_t}{\mathrm{\ell }}_{i_{t,s}}(U{x}_{i_{t,s}})\mathrm{.}$

With $W_t=W_0+c{B}_t{A}_t$ and $G_t={\mathrm{\nabla }}_W{\mathcal{L}}_t(W_t)$, all primitive calls in this batch objective have values and each piecewise primitive call is on a unique strict branch; branch choices may differ between successful steps. Reverse-mode automatic differentiation returns the exact real-arithmetic gradients ${\mathrm{\nabla }}_{B_t}{\mathcal{L}}_t(W_t)=c{G}_t{A}_t^{\mathsf{T}}$ and ${\mathrm{\nabla }}_{A_t}{\mathcal{L}}_t(W_t)=c{B}_t^{\mathsf{T}}{G}_t$. If $B=0$, then ${\mathrm{\nabla }}_A\mathcal{L}=0$ for every such empirical loss; if also $A=0$, then ${\mathrm{\nabla }}_B\mathcal{L}=0$. For every input $x\in {\mathbb{R}}^k$, the unmerged and merged layer outputs agree:

$W_{0}x+cB(Ax)=(W_0+cBA)x\mathrm{.}$

The factor matrices contain $rk+dr=r(d+k)$ trainable scalar coordinates, excluding the frozen entries of $W_0$; this coordinate count is not an assertion about identifiable degrees of freedom. All equalities are in real arithmetic and imply no floating-point, convergence, task-transfer, or hardware-performance guarantee.

Let $T\ge 0$ be the finite horizon of the procedure. The finite-budget procedure also has the following termination and status property. Suppose each primal value/status check terminates, and each defined primitive evaluation and each local derivative evaluation on a strict branch returns its specified real-arithmetic value or derivative in finite time. Then at the first step whose ordered primal evaluation encounters a primitive with no value or a defined piecewise primitive on a boundary or ambiguous branch, the procedure returns DOMAIN_ERROR or NONDIFFERENTIABLE, respectively, before updating the factors at that step. If no such evaluation occurs, it returns SUCCESS after exactly $T$ updates. For $T=0$, it returns immediately without evaluating a loss.

Hu et al., §4.1: factorized update, frozen base weights, and exact merge identity

MIT OCW, §§8.3–8.4: chain-rule accumulation and reverse-mode gradients of scalar objectives