MATRYOSHKALORA: LLM 파인튜닝을 위한 정확한 계층적 저순위 표현 학습

기존 LoRA는 학습 시 설정한 고정된 Rank(r)에서만 동작하여 최적의 효율을 찾기 위해 반복적인 실험이 필요했다. 이 논문은 하나의 모델 학습만으로 다양한 Rank를 자유롭게 선택해 사용할 수 있게 하여 연산 비용을 획기적으로 줄이고 가변적인 하드웨어 환경에 즉각 대응할 수 있게 한다.

LoRA는 거대 모델의 가중치를 직접 수정하는 대신, 두 개의 작은 행렬(A, B)의 곱으로 변화량을 근사한다. 이때 행렬의 크기를 결정하는 Rank(r)가 성능과 비용의 균형을 결정하는데, 기존 방식은 학습이 끝난 후에는 이 Rank를 바꿀 수 없다는 한계가 있다.

MatryoshkaLoRA는 임베딩 공간의 차원을 계층적으로 구성하는 Matryoshka Representation Learning 개념을 LoRA에 이식한다. 핵심 원리는 학습 시 모든 가능한 Rank(1, 2, 4, ... R)의 출력을 동시에 고려하여 가중치를 업데이트하는 것이다. 이를 통해 낮은 Rank의 파라미터가 높은 Rank의 파라미터 내부에 핵심 정보로 응축되도록 만든다.

결과적으로 사용자는 모델을 한 번만 학습시키면 된다. 추론 시점에 자원이 부족하면 낮은 Rank의 앞부분 파라미터만 잘라서 쓰고, 고성능이 필요하면 전체 파라미터를 사용하는 식으로 유연한 배포가 가능해진다.

전체 접근 방식은 기존 LoRA의 순전파 식 Y = x(W0 + sRAB)를 확장하여, 모든 서브 Rank r의 기여도를 합산하는 구조를 취한다. 수식적으로는 Y = x(W0 + Σ srArBr) 형태를 가지며, 이는 모든 Rank의 손실값이 파라미터 업데이트에 반영되도록 설계된 것이다.

구현 효율성을 위해 각 Rank마다 별도의 연산을 수행하는 대신, 고정된 대각 행렬 P를 도입한다. [입력 벡터 x와 행렬 A의 곱 결과값에] → [대각 행렬 P의 요소들을 각 차원별로 곱하는 연산을 수행하여] → [스케일링된 중간 표현을 얻고] → [이를 다시 행렬 B와 곱함으로써] 단 한 번의 행렬 연산으로 모든 계층적 Rank의 학습 효과를 통합한다.

대각 행렬 P의 각 요소 pi는 해당 차원이 포함되는 모든 Rank의 스케일링 인자(sr)들의 합으로 계산된다. [사용하고자 하는 Rank 집합 S와 최대 Rank R이 주어질 때] → [알고리즘 1을 통해 각 차원별 누적 가중치를 계산하여] → [R 차원의 벡터 P를 생성하고] → [이 벡터가 학습 시 그래디언트의 흐름을 조절하는 역할]을 수행하게 된다.

Llama-3.2-1B 모델을 GSM-8K 데이터셋으로 파인튜닝한 결과, MatryoshkaLoRA는 모든 Rank 구간에서 기존 LoRA 및 DyLoRA를 압도하는 성능을 보였다. 특히 Rank=2와 같은 매우 낮은 설정에서도 MatryoshkaLoRA는 약 35% 이상의 정확도를 유지한 반면, 기존 방식들은 32% 수준에 머물렀다.

AURAC 지표 기준으로는 MatryoshkaLoRA가 38.4%를 기록하여 LoRA(34.5%) 및 DyLoRA(34.9%) 대비 약 3.5%p 이상의 성능 향상을 입증했다. 이는 단일 모델 내의 모든 서브 Rank들이 독립적으로 학습된 모델들보다도 더 높은 효율성을 가짐을 의미한다.

Llama-3.1-8B 모델 실험에서도 동일한 경향이 확인되었다. ARC-C 및 HellaSwag 벤치마크에서 Rank가 증가함에 따라 성능이 계단식으로 확실히 상승하는 계층적 특성을 보였으며, 최대 Rank에서는 베이스라인 대비 3%p 이상의 정확도 이득을 얻었다.

MatryoshkaLoRA는 DyLoRA의 확률적 Rank 샘플링 방식이 가지는 데이터 비효율성을 결정론적 대리 손실 함수(Deterministic Surrogate Objective)로 해결한다. DyLoRA가 매 스텝 하나의 Rank만 업데이트하여 그래디언트 신호가 희소해지는 것과 달리, 본 논문은 모든 Rank의 기여도를 가중 합산하여 한 번의 스텝에서 전체 계층을 동시에 최적화한다.

수학적으로는 다중 Rank 목적 함수 L_multi를 1차 테일러 전개를 통해 근사하여, 중간에 대각 가중치 행렬 P가 삽입된 단일 Lo래 구조로 변환 가능하다는 것을 증명했다. 이는 추가적인 연산 오버헤드 없이도 다중 Rank 학습이 가능함을 이론적으로 뒷받침한다.

학습 시 어댑터 스케일링 파라미터 sk를 1로 설정하는 것이 1/r이나 1/sqrt(r)을 사용하는 것보다 안정적인 학습률 그리드를 형성함을 실험적으로 확인했다. 이는 상위 Rank가 하위 Rank의 정보를 포함하면서도 추가적인 세부 특징을 효과적으로 학습할 수 있게 돕는다.

성능 측정을 위해 각 Rank별로 개별적인 평가(Evaluation)를 수행해야 하므로, 전체 평가 시간이 Rank 개수에 비례하여 증가하는 단점이 있다. 또한, 모든 레이어에 동일한 Rank k를 적용하는 시나리오에 집중되어 있어 레이어별 민감도를 고려한 가변 Rank 적용 연구는 향후 과제로 남아있다.