A Training-Time Diagnostic for Generalization via the Log-Alignment Ratio

제목: 로그 정렬 비율을 통한 일반화에 대한 학습 시간 진단

문제 제기
신경망 모델과 데이터셋이 확장됨에 따라 학습 실행의 계산 비용이 증가하여, 학습 역학의 연구와 저렴한 진단 도구의 설계가 점점 더 가치 있게 되고 있습니다. 실무자들은 비용이 많이 드는 홀드아웃 검증 데이터에 의존하지 않고 학습 시간의 양만을 사용하여 과적합, 불안정성 또는 포화 상태와 같은 문제를 식별할 수 있는 지표를 필요로 합니다. 이전 연구들은 일반화를 이해하기 위해 가중치 행렬의 스펙트럼 특성 (예: 유효 차수, 안정적 차수) 이나 파라미터화 이론 (예: 최대 업데이트 파라미터화, $\mu$P) 을 탐구해 왔으나, 기존 지표들은 종종 가중치 또는 활성화 중 하나에만 집중하거나 완전한 특이값 분해 (SVD) 나 주성분 분석 (PCA) 과 같이 계산 비용이 많이 드는 연산을 필요로 합니다.

방법론
저자들은 가중치 행렬과 입력 활성화 간의 정렬 정도를 측정하기 위해 Everett 등 [2024] 이 파라미터화 이론에서 처음 도입한 **로그 정렬 비율 (Log-Alignment Ratio, LAR)**을 조사합니다. 본 논문은 LAR 을 이론적으로 재구성하여 두 번째, 이전에 탐구되지 않았던 의미를 밝혀냅니다: 즉, LAR 은 가중치 스펙트럼($p$) 과 활성화 스펙트럼($q$) 간의 중첩을 측정한다는 것입니다.

• 정의:
  • $W$를 특이값이 $s_i$인 행렬이라고 합시다. 가중치 스펙트럼 $p$는 정규화된 제곱 특이값으로 정의됩니다: $p_i = s_i^2 / \sum s_j^2$.
  • $X$를 입력 배치라고 합시다. 활성화 스펙트럼 $q$는 $W$의 특이 방향에 투영된 입력의 정규화된 제곱 투영으로 정의됩니다: $q_i = \sum_{x \in B} (v_i^\top x)^2 / \sum_{x \in B} \|x\|^2$.
• 재구성: 저자들은 LAR 을 다음과 같이 표현할 수 있음을 유도합니다:
  $$\text{LAR} = 1 + \frac{1}{2} \log_n \left( \sum_{i} p_i q_i \right)$$
  이 공식은 LAR 이 분포 $p$와 $q$ 간의 로그 중첩의 척도임을 보여줍니다. 이는 공분산과 피어슨 상관관계로 표현할 수도 있습니다:
  $$\text{LAR} = 1 + \frac{1}{2} \log_n \left( n \cdot \text{cov}(p, q) + \frac{1}{n} \right)$$
  $$\text{LAR} = 1 + \frac{1}{2} \log_n \left( n \sigma_p \sigma_q \text{corr}(p, q) + \frac{1}{n} \right)$$
  여기서 $\sigma_p$와 $\sigma_q$는 분포의 "퍼짐" 또는 집중도를 나타냅니다.
• 계산 효율성: LAR 은 순전파 중 사용 가능한 양 (행렬 노름) 에서 계산 가능하며, 홀드아웃 데이터나 완전한 SVD 를 필요로 하지 않고도 무시할 수 있는 오버헤드 ($O(|B|m + |B|n + mn)$) 로 계산할 수 있습니다.

주요 기여

1. 분포적 재구성: 본 논문은 LAR 이 단순한 정렬의 척도가 아니라 가중치와 활성화 스펙트럼 간의 중첩을 정량화한다는 것을 보여줍니다. 이는 지표의 행동을 기본 양의 관점에서 해석 가능하게 만드는 동등한 통계적 형태 (공분산, 상관관계) 를 제공합니다.
2. 유효 차수 예측: 저자들은 가중치 행렬이 $k$개의 특이 방향에 에너지를 집중시키고 활성화가 그것들과 정렬될 때, LAR 이 유효 차수 $k$와 $k \approx n^2(1 - \text{LAR})$를 통해 관련됨을 보여줍니다. 이 공식은 그로킹 (grokking) 실험에서 경험적으로 검증되었습니다.
3. 대규모 과적합 진단: 대규모 사전 학습 (30 억 파라미터 모델) 에서, 본 논문은 일반화 과정에서 언임베딩 (unembedding) LAR 이 안정화되지만 과적합이 접근함에 따라 급격히 감소함을 확립했습니다. 과적합이 아닌 기준선으로부터의 LAR 편차는 일반화 격차를 추적합니다.
4. 스펙트럼 퍼짐의 대리 지표: 연구 결과 최적화는 $p$와 $q$가 매우 높은 상관관계를 갖도록 유도한다는 것이 밝혀졌습니다. 결과적으로 LAR 은 이러한 분포의 분산 (퍼짐) 의 대리 지표 역할을 합니다. 사전 학습에서 상위 특이 방향이 지배적이므로, 상위 구성 요소 ($L_1$) 의 기여도가 과적합 전 조짐의 주요 지표가 됩니다.

실험 결과
저자들은 두 가지 다른 설정에서 이러한 발견을 검증했습니다:

1. 그로킹 실험 (소규모 알고리즘 작업):

   • 12 개의 이진 연산 작업에 대해 2 층 트랜스포머를 학습시켰습니다.
   • 관찰: 암기 단계 동안 LAR 은 낮습니다 ($p$와 $q$가 확산됨을 나타냄). 모델이 "그로킹"(일반화) 으로 전환됨에 따라 $p$와 $q$가 공유된 방향 집합에 집중되면서 LAR 이 증가합니다.
   • 차원성: LAR 이 예측한 차수 $k_{pred} = n^2(1 - \text{LAR})$는 최종 체크포인트에서 분산의 95% 를 설명하는 주성분의 수인 $k_{95}$와 낮은 제곱근 평균 제곱 오차 (RMSE) 로 밀접하게 일치합니다.
   • 구분: 그로킹 모델은 비그로킹 모델 ($\le 0.65$) 에 비해 더 높은 최종 LAR 값 ($>0.65$) 을 달성합니다.

2. 대규모 사전 학습 (30 억 파라미터 언어 모델):

   • 2 억에서 1000 억 토큰에 이르는 데이터셋에서 Adam 및 Muon 옵티마이저를 사용하여 Gemma-2 기반 모델 (30 억 파라미터) 을 학습시켰습니다.
   • 일반화 대 과적합: 과적합이 없는 영역 (예: 1000 억 토큰) 에서 LAR 은 정점에 도달한 후 서서히 감소하여 안정화됩니다. 과적합 영역 (소규모 데이터셋) 에서 LAR 은 과적합 단계가 접근함에 따라 급격히 감소합니다.
   • 일반화 격차: 특정 실행의 LAR 과 기준선 1000 억 토큰 실행의 LAR 간의 차이 (LAR 미분) 는 일반화 격차 (검증 손실에서 학습 손실을 뺀 값) 를 밀접하게 반영합니다.
   • 메커니즘: 과적합 이전의 LAR 감소는 주로 상위 특이 구성 요소 ($L_1$) 에 의해 주도됩니다. Adam 의 경우 이는 활성화의 확산 ($q_1$) 에 의해 주도되며, Muon 의 경우 가중치 ($p_1$) 와 활성화 모두에 관여합니다.

의의 및 주장
본 논문은 LAR 이 일반화에 대한 실용적이고 계산적으로 저렴하며 데이터가 불필요한 진단을 제공한다고 주장합니다.

• 효율성: SVD 를 필요로 하는 유효 차수나 다른 스펙트럼 측정법과 달리, LAR 은 순전파 중 무시할 수 있는 오버헤드로 계산할 수 있습니다.
• 예측 능력: 그로킹에서 암기에서 일반화로 전환을 성공적으로 추적하며, 대규모 사전 학습에서 과적합에 대한 조기 경고 신호로 작용하여 잠재적으로 평가가 불필요한 학습 방법을 가능하게 합니다.
• 해석 가능성: LAR 을 가중치 및 활성화 스펙트럼의 퍼짐과 연결함으로써, 이 지표는 최적화 역학 (스펙트럼의 집중 대 확산) 이 모델 일반화와 어떻게 관련되는지에 대한 통합된 관점을 제공합니다.

저자들은 한계를 인정하며, 그들의 증거는 상관관계적이며, 주로 언임베딩 레이어에 초점을 맞추고 있으며, 스펙트럼 퍼짐을 직접 정규화하는 것이 과적합을 방지하는지 여부는 아직 테스트하지 않았음을 지적합니다. 그들은 향후 연구가 다른 행렬 유형에 대한 LAR 과 다양한 배치 크기 하에서의 행동을 탐구해야 한다고 제안합니다.