크로스 레이어 저랭크 잔차 네트워크로 효율적인 대규모 언어 모델 학습

초록

본 논문은 인접 레이어 활성화 차이가 저랭크 구조를 가진다는 새로운 관찰을 기반으로, 저랭크 파라미터와 잔차 연결을 결합한 CR‑Net을 제안한다. CR‑Net은 파라미터 효율성을 높이고, 연산량 및 활성화 메모리를 크게 절감하면서도 기존 저랭크 방법보다 우수한 성능을 달성한다.

상세 분석

CR‑Net의 핵심 아이디어는 “인접 레이어 활성화 잔차(ΔY)는 저랭크 특성을 강하게 보인다”는 경험적 발견이다. 기존 저랭크 접근법은 파라미터 자체를 저랭크 행렬로 제한하거나, 그래디언트를 저랭크 공간에 투사하는 방식으로 메모리와 연산을 절감했지만, (1) 성능 저하, (2) 추가적인 SVD 연산 등으로 인한 연산 오버헤드, (3) 활성화 저장량 감소가 제한적이었다는 한계가 있었다. CR‑Net은 이러한 한계를 극복하기 위해 두 가지 설계를 도입한다. 첫째, 각 선형 레이어의 출력 Yₚˡ을 이전 레이어 출력 β₀·Yₚˡ⁻¹와 현재 레이어의 저랭크 잔차 LRᵣ(ΔYₚˡ)를 선형 결합하는 형태로 재구성한다. 여기서 Aₚˡ∈ℝ^{h_in×r}, Bₚˡ∈ℝ^{r×h_out}는 학습 가능한 저랭크 매트릭스이며, βₚˡ은 학습 가능한 스케일링 파라미터이다. βₚˡ이 0에 가까우면 저랭크 잔차가 주도하고, 1에 가까우면 이전 레이어의 고랭크 신호가 크게 반영된다. 이렇게 하면 첫 번째 레이어만 전통적인 풀‑랭크 가중치를 유지하면서도, 이후 레이어에서는 고랭크 정보를 손실 없이 저랭크 파라미터만으로 재현할 수 있다. 둘째, 활성화 재계산 전략을 맞춤 설계했다. 일반적인 Gradient Checkpointing은 O(L²) 비용이 발생하지만, CR‑Net은 선택된 레이어 집합 A에 대해 입력 Xₗ과 선형 결과를 저장하고, 나머지 레이어는 이전 레이어의 저장된 출력과 저랭크 매트릭스만을 이용해 재계산한다. 이는 재계산 오버헤드를 선형 수준으로 낮추면서도 메모리 사용을 크게 줄인다. 실험에서는 60M7B 규모 모델에 대해 동일한 학습 단계에서 LoRA, GaLore, RSO 등 최신 저랭크 기법보다 낮은 파라미터 수와 메모리 사용량을 보였으며, 특히 7B 모델에서는 검증 perplexity가 12% 개선되었다. 또한, 학습 속도는 기존 저랭크 방법과 거의 동일하거나 약간 향상되었으며, 이는 저랭크 잔차 연산이 기존 선형 연산에 비해 가벼운 행렬 곱셈으로 구현되기 때문이다. 전체적으로 CR‑Net은 (L1) 성능 저하 최소화, (L2) 연산 오버헤드 감소, (L3) 활성화 메모리 절감이라는 세 가지 기존 저랭크 방법의 한계를 동시에 해결한다는 점에서 의미가 크다.