COALA: 안정적·고속 컨텍스트 인식 저랭크 근사 프레임워크

초록

COALA는 입력 활성화 행렬을 가중치로 활용한 컨텍스트‑인식 저랭크 근사에서 발생하는 그램 행렬 계산·역전 문제를 완전히 회피한다. TSQR 기반 QR 분해와 정규화 기법을 결합해 메모리·수치 안정성을 확보하고, 제한된 데이터·대규모 캘리브레이션 상황에서도 정확한 근사와 빠른 실행 시간을 제공한다.

상세 분석

본 논문은 대규모 신경망 압축·미세조정에 널리 사용되는 컨텍스트‑인식 저랭크 근사(weight ≈ U V) 문제의 핵심 난제인 ‘그람 행렬(XXᵀ)의 형성·역전’이 초래하는 수치 불안정을 근본적으로 해결한다는 점에서 의의가 크다. 기존 연구들은 보통 (1) XXᵀ를 직접 계산하고 Cholesky·SVD 등으로 정규화된 행렬 S를 구한 뒤, W′ = U Σ_r Vᵀ S⁻¹ 형태의 해를 도출한다. 그러나 XXᵀ가 거의 특이(singular)하거나 부동소수점 한계에서 매우 작은 특이값을 가질 경우, S⁻¹ 연산 자체가 큰 오차를 유발한다. 특히 LLaMA3‑8B와 같이 수천 차원의 활성화 행렬을 수백 개 샘플에 대해 처리할 때 메모리 초과와 연산 비용이 급증한다는 실험적 증거를 제시한다.

COALA는 두 가지 핵심 아이디어로 이 문제를 회피한다. 첫째, ‘역전 없는( inversion‑free )’ 접근법으로, 목표 함수 ‖WX − W′X‖_F 를 직접 최소화한다. 이를 위해 W X를 먼저 QR 분해(QRᵀ = Xᵀ)하고, 상삼각 행렬 R만을 이용해 W Rᵀ의 SVD를 수행한다. Proposition 1에 따르면, W′ = U_r U_rᵀ W (U_r은 W Rᵀ의 앞 r개의 좌측 특이벡터) 가 최소해가 되며, 여기서는 어떠한 그람 행렬도 필요하지 않다. 둘째, 메모리 효율성을 위해 Tall‑Skinny QR(TSQR) 알고리즘을 적용한다. TSQR은 큰 X를 여러 청크로 나누어 순차적으로 QR를 수행하고, 최종 R을 재귀적으로 결합함으로써 GPU 메모리 한계를 넘어서는 데이터도 처리 가능하게 만든다. 실험에서는 TSQR 기반 방법이 기존 SVD‑LLM 대비 2배 이상 빠르면서도 상대 오차를 10⁻³ 수준으로 크게 감소시켰다.

정규화 측면에서는 ‖W − W′‖_F²에 μ ‖·‖_F² 항을 추가한 형태(문제 (4))를 제시한다. Proposition 3에 의해 이 정규화 문제는 X′ =