The Geometry of Sequential Learning: Lie-Bracket Prediction of Transfer Order

기술 요약: 순차적 학습의 기하학 (The Geometry of Sequential Learning)

문제 정의
다음 토큰 사전 학습 도메인 적응(next-token pretraining domain adaptation)부터 지시어 미세 조정(SFT) 및 직접 선호 최적화(DPO)에 이르는 순차적 학습 파이프라인은 본질적으로 순서 의존적이다. 모델을 $N$개의 후보 소스 도메인에서 타겟 도메인으로 적응시킬 때, $N!$개의 가능한 훈련 커리큘럼이 존재한다. 중간 규모의 $N$에 대해서도 이 공간을 전수 조사하는 것은 계산적으로 불가능하다. 핵심 과제는 두 가지 전체 훈련 스케줄을 모두 실행하지 않고도, 도메인 $A$를 학습한 후 $B$를 학습하는 것($A \to B$)이 그 반대 순서($B \to A$)보다 낮은 타겟 손실(target loss)을 산출할지 여부를 예측하는 것이다.

방법론
본 논문은 순서 의존성이 그래디언트 업데이트 연산자의 비가환성(non-commutativity)에서 발생하는 구조적인 기하학적 현상이라고 상정한다. 저자들은 최적의 전이 순서를 예측하기 위해 국소적이고 계산 가능한 기하학적 양인 **리 브래킷 교환자(Lie-bracket commutator)**를 유도한다.

1. 이론적 기초:
   도메인 $D$에서의 손실을 $L_D(\theta)$, 그래디언트를 $g_D$, 헤시안(Hessian)을 $H_D$라고 하자. 단일 그래디언트 단계는 연산자 $U_D(\theta) = \theta - \eta g_D(\theta)$이다. 두 소스 $A$와 $B$에 대해, 서로 반대되는 순서로 두 단계를 거친 후의 파라미터 공간 차이는 다음과 같다:
   $$\theta_{AB} - \theta_{BA} = \eta^2 (H_B g_A - H_A g_B) + O(\eta^3)$$
   이때 $b_{AB} := H_B g_A - H_A g_B$ 항은 그래디언트 흐름 벡터장의 리 브래킷(Lie bracket)으로 식별된다. 이 벡터가 순서 의존성을 유발한다.

2. 쌍별 예측 (The Planner):
   타겟 도메인 $E$에 대한 손실을 최소화하는 순서를 예측하기 위해, 이 방법은 브래킷 벡터를 타겟 그래디언트 $g_E$에 투영한다. 방향성 점수는 다음과 같다:
   $$\sigma^{(E)}_{AB} = \langle g_E, b_{AB} \rangle$$
   만약 $\sigma^{(E)}_{AB} < 0$이면, $A \to B$가 더 낮은 손실을 낼 것으로 예측되며, 그렇지 않으면 $B \to A$가 선호된다.

   • 드리프트 매칭 추정 (Drift-Matched Estimation): 무한소 영역을 넘어 정확도를 높이기 위해, 타겟 그래디언트 $g_E$는 초기화 지점 $\theta_0$가 아니라 "트로터 참조점(Trotter reference point)" $\theta_{ref} = \theta_0 - \eta(g_A + g_B)$에서 평가된다. 이 지점은 두 순서의 공통된 1차 드리프트(first-order drift)를 근사하여 그래디언트 추정 오차를 줄인다.
   • 계산 효율성: 이 방법은 그래디언트와 헤시안-벡터 곱(HVP)만을 필요로 한다. 전체 헤시안 행렬을 직접 구현하지 않는다. 실제 적용 시 계산은 훈련 가능한 파라미터 부분 집합(예: LLM의 최종 트랜스포머 레이어)으로 제한되어, 쌍당 약 3번의 그래디언트 평가와 2번의 HVP 비용이 발생한다.

3. 확장 가능한 스케줄링 (Lie-Bracket Tournaments):
   $N > 2$인 소스의 경우, 쌍별 점수는 가중치 토너먼트를 형성한다. 저자들은 $N!$개의 순열을 평가하지 않고 커리큘럼 순서를 결정하기 위해 보다/행 합계 순위 지정(Borda/row-sum ranking) 방식을 제안한다.

   • 핵심 최적화: 모든 쌍에 대해 공유된 타겟 그래디언트 참조($\theta_0$)를 사용함으로써, 쌍별 에지 가중치는 소스 그래디언트와 사전 계산된 HVP($u_i = H_i g_E$)의 내적을 통해 계산될 수 있다.
   • 복잡도: 이를 통해 스케줄링 문제를 $O(N)$의 내적과 $O(N \log N)$의 정렬 문제로 축소하여, $O(N^2)$의 전체 쌍 행렬 구현을 피한다.

주요 기여

1. 전이 순서의 교환자 이론: 본 논문은 순서 의존적 타겟 손실 차이와 방향성 브래킷 점수 $\sigma^{(E)}_{AB}$ 사이의 이론적 연결 고리를 확립하고, 이 부호가 최적의 쌍별 순서를 예측함을 증명한다.
2. 배포 가능한 쌍별 플래너: 드리프트 매칭된 트로터 추정기를 도입하여, 그래디언트와 HVP를 사용하여 순서 효과를 높은 정확도로 예측하며, "스테이크(stakes, 효과 크기)" 및 "신뢰도(confidence, 정렬도)" 지표를 포함한다.
3. 리 브래킷 토너먼트: $N!$의 커리큘럼 탐색을 랭킹 문제로 변환하여 수십 개의 도메인을 효율적으로 정렬할 수 있는 확장 가능한 스케줄링 방법을 제공한다.
4. 자동 스텝 사이즈 선택: 파일럿 데이터로부터 스텝 사이즈를 자동으로 선택하여 감지 가능성과 고차 안정성 제약 사이의 균형을 맞추는 이론 기반의 "오토파일럿($\eta$-autopilot)"을 제시한다.
5. 광범한 검증: 본 프레임워크는 지시어-SFT, DPO, 다음 토큰 사전 학습 적응, 디퓨전 모델 적응에 걸쳐 검증되었다.

실험 결과

• 쌍별 정확도:
  • 지시어-SFT: $k=1$ (도메인당 1단계)에서 **98.1%**의 부호 정확도를, $k=20$에서 **73.1%**를 달성했다.
  • DPO: $k=1$에서 98.9%, $k=20$에서 **72.2%**의 정확도를 달s했다.
  • 사전 학습/디퓨전: 4개의 LLM과 1개의 디퓨전 모델에 대해 높은 정확도를 유지했다 (예: Llama-3.2-1B에서 92.0%, 디퓨전에서 91.1%). 이는 랜덤 베이스라인을 크게 상회하는 수치이다.
• 장기 호라이즌 일반화 (Long-Horizon Generalization): 초기화 시 계산된 예측기는 긴 훈련 호라이즌 동안에도 효과적으로 유지되었다. Llama-3.2-1B의 경우, $k=20$에서 81.5%, $k=50$ (총 100번의 업데이트)에서 **65.3%**의 정확도를 유지하며, 랜덤 순서 대비 상당한 후회(regret) 감소를 보여주었다.
• 다중 도메인 스케줄링:
  • 작은 $N$ ($N=3$): 이 방법은 모든 $3!$개의 스케줄 중 최적의 스케줄을 **87.5%**의 시행에서 찾아냈다.
  • 큰 $N$: 56개의 MMLU 주제에 대해, 이 방법은 99.0~99.6번째 샘플 백분위수에 도달하여, 1번째 백분위수 미만으로 떨어진 하강 그래디언트 노름(descending gradient-norm) 베이스라인을 압도했다. 85개의 Stack 프로그래밍 언어 도메인에 대해서는 Python 타겟에 대해 99번째 백분위수에 도달했다.
• 효율성: 플래너는 도메인당 $k \approx 3$ 단계의 스텝을 사용할 때 두 가지 순서를 모두 직접 평가하는 브루트 포스 방식보다 계산 비용이 저렴해진다.

의의 및 주장
본 논문은 순차적 학습을 기하학적 토너먼트 문제로 재정의한다. 저자들은 리 브래킷 교환자가 국소적인 쌍별 순서 정보와 다중 도메인 스케줄 구축을 위한 확장 가능한 원 primitives를 모두 제공한다고 주장한다. 결과는 그래디언트 업데이트의 비가환성이 커리큘럼 설계의 지배적인 요인이며, 이러한 기하학적 구조를 활용하여 전수 조사를 피할 수 있음을 시사한다. 저자들은 이 방법이 특히 파라미터 효율적 미세 조정(예: 최종 레이어 또는 어댑터 튜닝) 시나리오에서 효과적이며, 쌍별 기하학적 프리미티브가 다양한 훈련 목적(SFT, DPO, 사전 학습) 및 아키텍처(Transformer, UNet)에 걸쳐 일반화된다는 점을 강조한다 임의의 커리큘럼에 대한 전역 최적화를 해결한다고 주장하는 것이 아니라, 국소적인 기하학적 휴리스틱이 최소한의 계산 오버헤드로 최적에 가까운 스케줄을 복구할 수 있음을 입증한다.