Composing Linear Layers from Irreducibles

기술 요약: 기약 요소(Irreducibles)로부터 선형 레이어 구성하기

문제 정의
현대의 대규모 언어 모델(LLM)은 풍부한 기능을 갖춘 모듈로 구성되는 저수준 프리미티브(low-level primitives)에 의존하고 있음을 시사하는 행동을 보입니다. 그러나 이러한 선형 변환의 근본적인 구성 요소에 대해서는 여전히 이해가 부족합니다. 기존 연구들은 매개변수 효율적인 근사법(예: Low-Rank Adaptation, Block-Hadamard)이나 모듈형 구성(예: Mixture-of-Experts)을 탐구해 왔으나, 복잡한 선형 동작을 생성하는 저수준 프리미티브의 *기계적 구성(mechanistic composition)*에 대한 이해는 결여되어 있습니다. 저자들은 구체적인 문제를 제기합니다: 일반적으로 수백만 개의 매개변수($O(d^2)$)를 포함하는 선형 레이어의 기능을 최소한의 기약 기하학적 객체 세트로부터 합성할 수 있는가?

방법론
저자들은 선형 변환을 분해하고 재구성하기 위해 **클리포드 대수(Clifford algebra)**를 사용하는 대수적 프레임워크를 제안합니다.

1. 대수적 정식화:

   • 표준 선형 사상은 전통적으로 밀집 행렬(dense matrices)로 표현됩니다. 저자들은 대신 선형 변환을 클리포드 대수 $Cl(n)$ 내의 바이벡터(bivectors)(방향이 있는 평면을 인코딩하는 기하학적 객체)의 합성으로 표현합니다.
   • 이들은 특수 직교 군 $SO(n)$의 이중 피복(double cover)인 스핀 군(Spin group) $Spin(n)$을 활용합니다. 로터(rotors)($r$)라고 불리는 이 그룹의 원소들은 샌드위치 곱 $x \mapsto rxr^\dagger$를 통해 멀티벡터(multivectors)에 작용합니다.
   • 임의의 멀티벡터는 $2^n$개의 매개변수를 필요로 하지만, 저자들은 로터로 구성된 바이벡터로 제한하여 구성을 수행합니다. 로터 $r$은 바이벡터 $b$의 지수 함수입니다: $r = \exp(b)$.

2. 매개변수 효율성:

   • 밀집 $d \times d$ 행렬은 $O(d^2)$의 매개변수를 필요로 합니다.
   • 변환을 부분 공간에 작용하는 로터들의 합성으로 표현함으로써, 필요한 스칼라 매개변수의 수는 $O(\log_2 d)$로 스케일링됩니다. 구체적으로 너비가 $w$일 때, 매개변수 수는 약 $w \binom{\log_2 d}{2}$입니다. $d=2048$인 경우, 이는 수백만 개의 매개변수를 단 몇 백 개로 줄입니다.

3. 미분 가능한 분해 알고리즘:

   • 일반적인 바이벡터에 대한 지수 맵은 무한 급수를 포함하므로 계산적으로 까다롭습니다. 그러나 단순 바이벡터(하나의 회전 평면을 나타내는 것)의 경우, 지수 맵은 폐쇄형 해(closed-form solution)를 가집니다: $\exp(b) = \cos(\|b\|) + \sin(\|b\|)\frac{b}{\|b\|}$.
   • 일반적인 바이벡터를 처리하기 위해, 저자들은 모든 바이벡터를 최대 $k = \lfloor n/2 \rfloor$개의 서로 교환 가능한 직교 단순 바이벡터의 합으로 표현하는 불변 분해(invariant decomposition)(Roelfs와 Keninck [2021] 기반)를 사용합니다.
   • 저자들은 클리포드 대수에 적응된 크릴로프 부분 공간(Krylov subspace) 기반의 거듭제곱 반복법(power iteration method)을 사용하여 이러한 단순 바이벡터를 반복적으로 추출하는 미분 가능한 불변 분해 알고리즘(알고리즘 1 및 2)을 도입합니다. 이는 명시적인 고유값 분해를 피하여 경사 하강법 기반 최적화를 위한 수치적 안정성과 미분 가능성을 보장합니다.

4. 일반화된 로터 가젯(Generalized Rotor Gadget):

   • 입력/출력 차원이 2의 거듭제곱이 아닌 임의의 차원($d_{in}, d_{out}$)을 처리하기 위해, 저자들은 서로 다른 부분 공간에서 작동하는 다수의 로터-샌드위치 모듈을 집합(aggregate)한 후 풀링 연산자를 적용하여 일반화된 변환 $\psi$를 구축합니다.

주요 기여

1. 대수적 표현: 저자들은 로터 샌드위치 곱을 사용하여 기하학적 프리미티브(바이벡터)의 합성으로 선형 변환을 표현함으로써, 매개변수 복잡도를 $O(d^2)$에서 $O(\log_2 d)$로 줄였습니다.
2. 미분 가능한 분해: 임의의 바이벡터를 불변 분해를 통해 그에 대응하는 로터로 매핑하는 폐쇄형의 미분 가능한 알고리즘을 제안하여, autograd 프레임워크와의 통합을 가능하게 했습니다.
3. 경험적 검증: 이 방법은 LLM 어텐션 레이어의 Key, Query, Value 투영에 적용되었습니다. 로터 기반 레이어는 밀집 레이어(layer-wise)를 모방하도록 학습되었으며, FMNIST 데이터셋에 대해 엔드투엔드(end-to-end)로 학습되었습니다.

결과

• 매개변수 효율성: LLaMA-3.2 1B 및 Qwen-2.5 1.5B 실험에서 로터 레이어는 엄청난 매개변수 감소를 달성했습니다. 예를 들어, LLaMA-3.2 1B의 쿼리 투영은 밀집 레이어의 약 4.19M 매개변수에서 로터 레이어의 $\le 896$개 매개변수로, 약 4700배 감소했습니다.
• 성능: 극단적인 압축에도 불구하고, 로터 기반 레이어는 강력한 베이스라인인 Low-Rank(LR1, LR4) 및 Block-Hadamard(BH1) 근사법과 대등하거나 이를 능가하는 성능을 보였습니다.
  • Wikitext2 (LLaMA-3.2 1B)에서 단일 어텐션 레이어를 교체했을 때, 로터는 2.629의 log-PPL을 기록하여 최상의 베이스라인(BH1, 2.636)보다 우수한 성능을 보였습니다.
  • Arc Challenge에서 로터는 55.31%의 정확도를 달려 BH1(54.83%)과 LR4(53.84%)를 앞질렀습니다.
• 강건성(Robustness): 로터의 성능은 다양한 레이어 조합(1, 2, 또는 3개 레이어 교체) 및 다양한 데이터셋(Wikitext2, C4, PTB, Arc, HellaSwag) 전반에 걸쳐 일관되게 유지되었습니다.
• 엔드투엔드 학습: FMNIST에서 처음부터 학습했을 때, 로터 기반 MLP는 밀집 베이스라인(89.67%)과 유사한 최종 정확도(88.36%)로 수렴했으나, 수렴 속도는 약간 더 느렸습니다.

의의 및 주장
본 논문은 소수의 기하학적 빌딩 블록으로부터 압축된 세트의 선형 레이어를 합성하는 것이 가능함을 입증한다고 주장합니다. 저자들은 단순히 밀집 레이어(현재 하드웨어 최적화 커널 부족으로 인해 실제 추론 시간 측면에서 불리함)를 대체하는 것이 목적이 아니라, 저수준 프리미티브가 어떻게 고수준 함수로 구성되는지에 대한 대수적 관점을 제공하는 것이 목표임을 강조합니다.

이러한 결과는 딥 모델의 복잡한 선형 레이어 기능이 수백만 개의 독립적인 가중치가 아니라, 수백 개의 기하학적 객체(로터)의 구조화된 합성을 통해 실현될 수 있음을 시사합니다. 이는 다음과 같은 새로운 경로를 제시합니다:

• 해석 가능성(Interpretability): 선형 레이어를 방향이 있는 평면들의 합성으로 이해함.
• 매개변수 효율성: 성능을 유지하면서 매개변수 수를 획기적으로 감소시킴.
• 미래 하드웨어: 기하학적 매개변수로부터 온칩(on-chip)에서 가중치를 합성함으로써, 대규모 모델 추론의 메모리 대역폭 병목 현상을 해결하기 위해 연산량을 메모리 대역폭과 교환할 수 있음.

저자들은 본 연구의 범위가 시스템 수준의 통합 및 하드웨어 인식 구현을 넘어선다는 점을 언급하며, 실질적인 이점을 얻기 위해서는 추가적인 작업이 필요함을 명시하며 신중한 태도를 유지했습니다.