AdaPreLoRA: Adafactor 선형 전처리 기반 Low-Rank Adaptation

LoRA의 factor-space 최적화는 JG의 차원 축약으로 인해 역전이 불가능한 특이성을 내재한다. 이로 인해 W-space preconditioner Ft를 적용하더라도 (JG의) 역연산은 고유하게 정의되지 않는 문제가 생긴다. AdaPreLoRA는 gradient-statistics-aware Ft와 Ht-imbalance 규칙을 결합해 JG의 해공간에서 W 업데이트에 가장 근접한 해를 닫힌 형태로 얻고, 메모리는 O((m+n)r)로 유지한다. 이 방식은 GPT-2(E2E), Mistral-7B, Qwen2-7B 및 diffusion 모델 개인화에서 기존 LoRA 계열과 경쟁하거나 우수한 성능을 보여주며 LoRA 프로토콜의 메모리 오버헤드를 증가시키지 않는다.

출발점: LoRA는 W0를 고정하고 W = BA로 재표현하므로 G([B,A]) = BA이고 JG([B,A])는 gauge 재매개화(B,A) → (BC, C^{-1}A)로 인해 랭크가-deficient하다. 이로 인해 Ft가 주어져도 J*GFtJG의 역은 일반적으로 고유하지 않으며, 역연산이 ill-defined해진다. 이 문제를 해결하기 위해 프레임워크는 (i) Ft를 선택하는 방법과 (ii) JG의 해집합에서 어떤 원소를 선택할지 결정하는 규칙의 두 축으로 구성된다. 해결의 핵심은 JG[ΔB, ΔA] = Get의 해가 존재하고, 해의 모양은 r^2 차원의 affine 공간이다라는 사실이다. AdaPreLoRA는 이 affine 공간에서의 유일한 해를 선택하는 규칙으로, Ht-imbalance를 최소화하는 Xt를 고정해 닫힌-form 업데이트를 얻는다. 이를 통해 W 업데이트는 Ht-orthogonal projection된 Get onto Tt에 대한 근사이며, Ht-노름에서의 projection으로 정의된다.

단계 요약(입력-계산-결과-의미의 흐름으로 제시) :

1. 입력: [Bt, At], Gt = ∇Wt L(W0 + Wt)와 Ft = Empirical Fisher on W, Ft = Lt ⊗ Rt, Ht = F^{1/2}_t. JG([Bt,At]) 및 J*G([Bt,At])를 구성한다. [어떤 값을 입력으로] → [연산] → [결과] → [의미]
2. 시스템 구성: JGFtJG [ΔBt, ΔAt] = JG(Gt) 를 만족하는 해를 구한다. 이때 역은 non-unique하며 해집합은 ker(JG)의 affine 공간이다. [입력] → [정의] → [확인] → [결과]
3. 해의 일반적 표현: 해는 ΔBt(Xt), ΔAt(Xt)로 나타나며 Xt ∈ R^{r×r}으로 매개된다. (13)
4. 해의 선택: Solution 2의 Ht-imbalance 최소화를 통해 Xt를 고정하여 단일 해를 선택한다. (14)
5. AdaPreLoRA의 닫힌-form 업데이트: Δopt Bt, Δopt At의 최종 형태는 (14)로 주어지며 P_eBt, Q_eAt는 각각 Ft-가중 투영 연산자이다. 이로써 W 업데이트는 Tt에 가장 가까운 방향으로 수렴한다. (Theorem 3.2)
6. 요지: AdaPreLoRA의 업데이트는 gradient-통계 기반 Ft와 LoRA의 저랭크 구조를 유지하면서, W-space 방향의 preconditioned 방향에 가장 가까운 LoRA 해를 선택하는 방식이다. 구현은 O((m+n)r) 메모리로 가능한 closed-form 해를 제공한다.

주요 벤치마크 결과:

• GPT-2 small(랭크 r=4, SGD/AdamW): AdaPreLoRA SGD가 69.5(BLEU)로 최고치, MET 8.77, ROUGE-L 46.5, CIDEr 71.5를 기록. AdaPreLoRA AdamW 역시 최고치에 근접하며 70.0(BLEU), 8.84, 46.3, 71.3로 확인.
• GPT-2 medium(랭크 r=4, SGD/AdamW): AdaPreLoRA SGD가 70.3(BLEU), 8.84, 46.9, 71.7으로 best에 근접. AdaPreLoRA AdamW도 70.3, 8.84, 46.7, 71.8로 비슷한 성능을 보임.
• 7B 규모(Mistral-7B, Qwen2-7B), 랭크 r=8: AdaPreLoRA AdamW가 RTE 89.5, CoLA 71.4, MRPC 90.0, ARC 91.0, GSM8K 76.4로 상위 성능. LoRA-Pro AdamW와 비교해도 우수한 편.
• Mix-of-Show diffusion personalization: AdaPreLoRA AdamW가 클립(CLIP) 점수 및 FID에서 최상 혹은 근접한 성능을 보임. 0.7 스케일에서 CLIP 31.47, FID 30.17, 1.0 스케일에서 CLIP 31.58, FID 28.18(또는 31.47/29.01–27.13 범주). 대체로 우수한 품질의 생성 이미지를 달성.
• 효율성: LoRA 프로토콜 대비 메모리 오버헤드가 증가하지 않으며, LoRA-Pro의 mn 규모 gradient/moments 방식보다 메모리 부담이 낮다. 예를 들어 Mistral-7B에서 AdaPreLoRA AdamW의 peak memory는 약 26.0 GB로 Scaled AdamW의 수준에 가깝고 LoRA-Pro AdamW는 약 50.4 GB를 필요로 한다.

단락 1: 전체 아키텍처 구조

• LoRA는 W0를 고정하고 W = BA로 재표현하며, G([B,A]) = BA로 정의된 생성자 G의 도함수 JG는 gauge 불변성(B, A) -> (BC, C^{-1}A)으로 인해 rank-deficient다. 이로 인해 Ft가 주어져도 J*GFtJG의 역은 일반적으로 존재하지 않는다.
• Ft는 W-space의 gradient-통계 기반 역전파 가중치를 제공하며, 이 Ft에 의해 J*GFtJG가 정의된다. Ft의 구조에 따라 기존 LoRA 최적화기가 구성된다. 단락 2: 핵심 메커니즘의 수학적 기반
• JG([B,A])[P,Q] = P A + B Q, JG([B,A])(C) = [C A^T, B^T C]. JG JG ([B,A])(C) = C A^T A + B B^T C.
• Ft = empirical Fisher in W-space, Ft = (1/N) sum vec(Gt)^⊤ vec(Gt) 형태. Ft가 SPD일 때 J*GFtJG은 singular하지만 해집합은 일관적이다.
• 문제 (7): JGFtJG [ΔB, ΔA] = JG(Gt)이며, 해집합은 ker(JG)의 affine 공간으로 구성된다. Solution 1은 이 시스템의 최소-노름 해를 찾고, Solution 2는 Xt를 통해 ker(JG) 방향의 차이를 Ht-imbalance를 최소화하는 방식으로 제거한다. 단락 3: Prior work 대비 차별점
• 대다수의 기존 LoRA 옵티마이저는 Ft를 I로 두거나 block-diagonal 근사를 사용하여 JG의 선형 시스템을 우회한다. AdaPreLoRA는 Adafactor diagonal Kronecker Ft를 사용하고, Xt를 통해 ker(JG) 편향을 해소한다. 이로써 W-update는 Ht-노름에서 Get의 Tt으로의 projection의 최적 근사이며, 업데이트는 닫힌-form으로 계산된다. 단락 4: 구현 및 학습 세부사항, 이론적 분석
• Adafactor Ft = L^{1/2}_t ⊗ R^{1/2}_t를 사용하고, Ht = F_t^{1/2}에 의해 Ht-노름 ⟨·,·⟩_Ht를 정의한다. (9)
• Solution 1: 문제 (11)의 최소-제곱 해를 얻고, Xt 매개변수화로 ΔBt(Xt), ΔAt(Xt)를 얻는다. (13)
• Solution 2: Xt를 최적화 문제로 두고, Xt^opt = −1/2 (B_t^T L^{1/2}_t B_t)^{-1} B_t^T G_t A^T_t (A_t R^{1/2}_t A^T_t)^{-1}로 해를 구한 뒤 (13)에 대입해 Theorem 3.2의 닫힌 해를 얻는다. (B.1–B.2)
• 최종 업데이트: Δopt Bt, Δopt At / PeBt, QeAt를 이용한 공식으로 주어진다. (Equ. 14)
• 알고리즘 구현: Algorithm 1은 SGD, Algorithm 2는 Momentum 버전의 AdaPreLoRA를 제시한다. 시간 복잡도는 per-iteration에서 O(mn + (m+n)r^2 + r^3)이며, 메모리는 O((m+n)r)이다.