[AI System] Efficient On-Device Training - Part 1

"Efficient On-Device Training - Part 1"

[Bottlenecks During Training]

· Computation : Backpropagation에서 행렬곱 연산이 많이 필요

    → Backpropagation FLOPs > Forward Propagation FLOPs    ∴ 학습은 추론보다 훨씬 무거움

· Parameter Memory : 학습할 때는 단순히 Weight만 저장하는 것 X

    → Model Parameter : n / Gradient : n / SGD : 0 / Momentum SGD : N / Adam : 2N

· Activation Memory : 중간 Activation을 저장해야 Backpropagation 가능

    → Activation Size = d_in x n    ∴ Batch Size 증가 시 Activaton 증가

▶ 효율적인 학습 방법의 목적은 위 3가지 병목을 완화하는 것

[LoRA(Low-Rank Adaptation)]

· 최근 LLM 학습/미세조정에서 널리 쓰이는 방법으로 큰 모델 전체를 다 학습시키는 대신, 적은 수의 추가 파라미터만 학습 시킴

[Revisiting Model Training]

<특정 Layer의 Weight Matrix>

<Revisiting Model Training>

· 학습 전의 상태 : 이미 학습된 Model의 Weight(W) 보유

· 학습 후의 상태 : Target Task 데이터로 SGD/Adam 같은 Optimizer를 돌려 학습된 Weight(W') 보유

→ ΔW = W' - W = 업데이트 행렬 = 학습 과정 전체에서 누적된 변화량

∴ 추론 시 입력 x가 들어오면 z = W'x = (W + ΔW)x = Wx + ΔWx ▶ Pre-Trained Weight는 고정시키고, Update만 학습하는 방식

    (Wx = 이미 학습된 Model이 내는 출력 / ΔWx = 새 Task에 맞추기 위한 보정량)

→ 학습의 핵심은 전체 W를 다 만지는 것이 아니라 좋은 보정항 ΔW를 얻는 것

[High-Level Idea of LoRA]

<z = Wx + ΔWx = Wx + BAx>

· 기존 Weight(W)는 고정

· 직접 큰 ΔW를 학습하지 않고, 작은 두 행렬 A, B를 학습해서 그 곱으로 Update를 표현

① 입력 x에 대해 먼저 A를 통과하여 낮은 차원으로 변경

② B로 다시 확장해서 원래 출력에 합산

** 핵심 조건은 r << d₁, d₂

> ΔW를 직접 학습 : 필요한 Parameter 수 = d₁d₂

> LoRA 적용 : A = rd₂, B = d₁r    ∴ 필요한 Parameter 수 = (d₁ + d₂)r → r이 매우 작으면 감소폭이 매우 커짐

<LoRA를 여러 Layer에 적용>

· LoRA는 모든 Layer에 적용 가능 또는 일부 Layer에만 적용 가능

[LoRA Intuition - Rank, Matrix Decomposition]

※ Rank : 행렬 A가 m x n 행렬일 때, Rank는 열 Vector 중 서로 선형독립인 최대 개수

    (Rank↓ = 중복된 정보↑ / Rank↑ = 다양한 방향 정보)

** Rank = Column Rank = Row Rank

※ Full Rank : Rank(C) = min(d₁, d₂) → Ex. 3x5 행렬의 최대 Rank = 5x3 행렬의 최대 Rank = 3

▣ Non-Full-Rank 행렬

· 행렬 C의 Rank = r & r < min(d₁, d₂) → C = BA = 낮은 Rank의 큰 행렬은 2개의 작은 행렬 곱으로 압축 표현 가능

C =  U∑V^T = U_r∑_rV_r^T = BA

· U = d₁ x d₁ / ∑ = r개를 제외한 나머지는 0인 대각행렬 / V^T = d₂ x d₂

· U_r = d₁ x r / ∑_r = r개인 대각행렬 = I_r / V_r^T = r x d₂

· A = ∑_rV_r^T / B = U_r

→ 큰 신경망에는 엄청난 양의 Parameter 존재

→ 많은 고차원 행렬은 Intrinsic Low-Rank Structure를 가질 수 있음

→ Update 행렬이 완전환 Full-Rank일 필요 없이 저차원 구조로도 충분히 Task Adaptation을 표현 가능

[LoRA Implementation]

① Pre-Trained W & A, B 초기화

② Forward Propagation

③ Backpropagation (Gradient는 A와 B에 대해서만 계산/Update)

④ Update

** LoRA는 Full Fine-Tuning과 비교 시 매우 적은 학습 Parameter로 강력한 성능을 낼 수 있음

→ But, 항상 Full Fine-Tuning보다 좋은 것은 아니며, Rank r이 너무 낮으면 성능 저하 가능

▣ LoRA Inference

· 학습 완료 후, 최종 Weight(W') = W + B'A' → LoRA는 학습할 때만 분리해서 다루고, 추론할 때는 합쳐서 사용

    = Wx', B'A'x' 분리해서 계산하지 않고, W'x' 한 번에 계산하므로 추가 Latency X

** 학습 : W와 BA 분리 / 추론 : W와 BA 합산

[LoRA Efficiency Analysis]

· Full Fine-Tuning : Task마다 별도의 W_k 전체를 저장

· LoRA : 공통 기반 모델 W 하나는 공유하고, Task마다 작은 A_k, B_k만 따로 저장

▣ Parameter Memory

> Original Full Model Update

  · Model Weights = W → d₁d₂

  · Gradients for W = W → d₁d₂

  · Optimize State 2개 (Moment1, Moment2) → 2d₁d₂

  ∴ 4d₁d₂ (Adam)

> LoRA

  · Original Model Weights = W → d₁d₂

  · LoRA Weights → r(d₁ + d₂)

  · Gradients for A, B → r(d₁ + d₂)

  · Optimize State 2개 (Moment1, Moment2) → 2r(d₁ + d₂)

  ∴ d₁d₂ + 4r(d₁ + d₂) (Adam)    ** r << d₁, d₂

▣ Activation (Layer Output/Input) Memory

> Original Full Model Update

  · W의 Gradient를 계산하려면, Forward Propagation 때의 입력 x를 저장해야 함

  · 입력 x의 size = d₂ x n (n = Mini-Batch Size) → Batch가 커질수록 Activation Memory가 선형적으로 증가

> LoRA

  · A의 Gradient를 계산하려면 A의 입력인 x, B의 Gradient를 계산하려면 B의 입력인 Ax를 저장해야 함

  · 입력 x의 size = d₂ x n (n = Mini-Batch Size), 입력 Ax의 size = r x n

  · 전체 Activation Memory = (d₂ + r)n → Batch가 커질수록 Activation Memory가 선형적으로 증가

∴ LoRA는 Parameter/Optimizer 메모리 절감에는 강하지만, Activation Memory 절감은 제한적

    → But, A, B만 Update하면 되기 때문에 Backpropagation FLOPs는 감소

[Other Variants of LoRA]

▣ LoRA+

· LoRA : A와 B를 같은 Learning Rate로 Update

· LoRA+ : A용 Learning Rate / B용 Learning Rate → 추가 자원을 거의 요구 X & 성능↑ & 학습 속도↑

∴ 구조 자체를 바꾸지 않아도 Optimizer 설정(A와 B에 서로 다른 Learning Rate를 사용)으로 성능/속도 개선 가능

▣ LoRA-FA(Frozen A)

· LoRA : A, B 둘 다 학습 → 각각에 대한 입력 Activation 저장 → Activation Memory를 충분히 줄이지 못함

· LoRA-FA : A는 랜덤 초기화 후 고정 & B만 학습

    → A : 입력 Feature를 뽑아내는 역할 / B : 해당 Feature를 원하는 출력 쪽으로 Mapping하는 역할 = A와 B는 대칭적인 역할 X

        ∴ B를 Tuning하는 것이 A를 Tuning하는 것보다 본직적으로 더 효과적일 수 있음(Asymmetry in LoRA)

> Parameter Memory

· LoRA : d₁d₂ + 4r(d₁ + d₂)

· LoRA-FA : d₁d₂ + 4rd₁ + rd₂ → A를 학습하지 않아 그에 해당하는 Gradient와 Optimizer State 감소

> Activation Memory

· LoRA : (d₂ + r) x n → A, B 둘 다 Gradient 계산이 필요해서 두 입력을 저장

· LoRA-FA : r x n → B만 Gradient 계산하면 되므로 B 입력만 저장

∴ A를 고정하면 Trainable Parameter 수와 Activation Memory를 줄일 수 있고, 성능 저하는 크지 않음

▣ VeRA(Vector-Based Random Matrix Adaptation)

· LoRA : A, B 전체 학습

· VeRA : 큰 랜덤 행렬을 만들어두고, 실제 학습은 그 위에 곱해지는 작은 Vector p, q 같은 값만 조정

> LoRA

z = Wx + BAx

· A = r x d₂ / B = d₁ x r → # of Trainable Parameters = (d₁ + d₂)r

→ 작은 행렬 자체를 학습

> VeRA

z = Wx + Λ_pBΛ_qAx

· Λ_p = d₁ x d₁ 대각 행렬 / Λ_q = r x r 대각 행렬 → # of Trainable Parameters = r + d₁

→ 이 대각 행렬이 행/열을 Scale 하거나 비활성화하는 역할

→ 이미 있는 랜덤 변환을 얼마나 살릴지 조절

∴ A, B를 무작위로 초기화하고 고정한 후, A, B의 요소들을 확대/축소/비활성화하는 2개의 작은 Vector만으로 학습

    → 학습 가능한 Parameter의 수가 크게 감소

▣ NoLA

# of Trainable Parameters

· Full Fine-Tuning : d₁d₂

· LoRA : (d₁ + d₂)r → Rank가 1이어도 최소 (d₁ + d₂) Parameter 필요

① 랜덤 Model 여러 개 생성 후 고정

② 전체 모델을 선형 결합의 합으로 표현

③ 계수 α만 학습(Model 자체 학습 X)

→ # of Trainable Parameters = k (Model 크기와 완전히 독립적)

> LoRA에 적용

· 기존 LoRA : ΔW = BA

① A₁, A₂, …, A_k / B₁, B₂, …, B_l

② A = ∑α_iA_i / B = ∑β_jB_j

③ ΔW = (∑β_jB_j) x (∑α_iA_i) & α₁, α₂, …, α_k / β₁, β₂, …, β_l만 학습

→ # of Trainable Parameters = k + l (Model 구조와 무관 & Rank와 무관)

∴ LoRA 행렬을 랜덤 Basis의 선형 결합으로 표현해서, 학습 Parameter 수를 더 줄일 수 있음

** LoRA 계열 기법들은 Large-Scale Model을 Cloud Server뿐만 아니라 Device에서도 효율적으로 학습/적용 가능

    + LoRA는 현재 Large Foundation Model Fine-Tuning의 표준에 가까움

[Reference]

· Efficient On-Device Training Part 1 (AI System : Week 3) - Dongjun Han