Geometry-Guided Layerwise FFN Width Allocation in Transformers

기술 요약: 트랜스포머의 기하학적 가이드를 통한 레이어별 FFN 너비 할당

문제 정의
표준 트랜스포머 아키텍처에서 피드포워드 네트워크(FFN)의 은닉 너비(hidden width)는 일반적으로 모든 레이어에 걸쳐 일정하게 유지되지만, 레이어의 기능과 표현 역학(representation dynamics)이 깊이에 따라 변화한다는 증거가 존재함에도 불구하고 이를 따르고 있다. 최근 연구(Tapered Language Models)는 깊은 레이어에서 얕은 레이어로 용량을 재할당하는 것이 성능을 향ano 개선한다는 것을 보여주었으나, 기존 방식들은 수기로 설계된 단조로운 스케줄(예: 코사인 테이퍼)에 의존한다. 이러한 규정된 형태는 특정 모델이나 데이터 분포의 구체적인 계산 요구 사항에 적응하지 못한다. 본 논문은 고정된 아키텍처 사전 지식에 의존하는 대신, 순전파(forward pass) 중 발생하는 활성화의 기하학적 동작으로부터 데이터 의존적이고 레이어별인 너비 할당을 직접 도출할 수 있는지 묻는다.

방법론
저자들은 FFN 레이어를 토큰 표현 클라우드(token representation clouds)의 운송자(transporters)로 취급하는 프레임워크를 제안한다. 핵심 방법론은 세 단계로 구성된다:

1. 기하학적 작업 측정 (Geometric Work Measurement):
   본 논문은 "기하학적 작업"을 FFN 레이어가 토큰 표현 클라우드에 유도하는 변화로 정의한다. 이 지표가 잔차 스케일링(residual scaling)의 인위적 결과가 아닌 FFN의 실제 기여도를 반영하도록 하기 위해, 저자들은 세 가지 서로 다른 기하학적 프레임워크 하에서 세 가지 별개의 작업 통계량을 평가한다:

   • Raw Geometry (원시 기하학): 정규화되지 않은 상태(unnormalized states)에 대한 유클리드 거리를 사용한다. 이는 스케일에 민 phép하며 제어 집단 역할을 한다.
   • Spherical Geometry (구형 기로학): 단위 정규화된 상태와 각도 거리를 사용한다.
   • Hyperbolic Geometry (쌍곡 기하학): 단위 정규화된 상태를 푸앵카레 볼(Poincaré ball)로 매핑하여 각도 거리에 대한 비선형적 재가중치(reweighting) 역할을 수행하게 한다.

   각 기하학에 대해 세 가지 작업 통계량이 계산된다:

   • Shift (이동): 토큰의 평균 변위(대응 관계 보존).
   • Metric Distortion (메트릭 왜곡): 그로모프-와서스타인(Gromov-Wasserstein, GW) 거리를 통해 측정되며, 쌍별 관계가 어떻게 보존되거나 변경되는지를 정량화한다.
   • Topological Change (위상적 변화): 1차 지속성 호몰로지 다이어그램(degree-one persistent homology diagrams, Vietoris–Rips filtration) 사이의 2-와서스타인(2-Wasserstein) 거리를 통해 측정되며, 루프(loops)의 생성 또는 파괴를 포착한다.

2. 근사 대리 모델 및 최적화 (Approximation Surrogate and Optimization):
   저자들은 기하학적 작업을 최적의 레이어 너비로 매핑하는 대리 모델을 도입한다. 이들은 FFN 레이어의 근사 오차가 고유 차원($d_{int}$)과 손실 민감도($\lambda$)에 따라 너비에 따라 스케일링된다고 가정한다.

   • 가정: (1) 손실 효과는 국소적으로 가산적(locally additive)이다; (2) 오차는 $w^{-a}$ (여기서 $a$는 고유 차원에 의존)로 스케일링된다; (3) 측정된 기하학적 작업은 근사의 어려움을 나타내는 대리 지표이다.
   • 최적화 도구: 고정된 총 파라미터 예산 하에서 가중치 적용된 근사 오차의 합을 최소화하는 제약 최적화 문제를 해결한다. 이를 통해 하한값(lower bounds)과 깊이 의존적 지수를 포함하는 각 레이어의 최적 너비 $w_\ell$에 대한 폐쇄형 해(closed-form solution)를 도출한다.

3. 스케줄 구축 (Schedule Construction):
   참조 모델(균일 너비)을 실행하여 활성화 클라우드를 수집한다. 기하학적 작업 프로파일을 추정, 평활화(smoothing)한 후 도출된 최적화 도구를 사용하여 너비 스케줄로 매핑한다. 이 스케줄은 새로운 모델을 처음부터 학습시키는 데 적용된다.

주요 기여

1. 기하학적 작업의 정의: 원시, 구형, 쌍곡 메트릭 하에서의 세 가지 상호 보완적인 순전파 작업 통계량(이동, GW 왜곡, 위상적 변화)을 공식화하여, 단순 크기와 스케일 정규화된 구조적 변화를 구분하였다.
2. 정확한 예산 최적 솔버: 고정된 파라미터 예산 하에서 레이어별 너비 할당을 위한 고유한 연속 최적화 도구를 도출하였으며, 하한값과 레이어별로 변화하는 고유 차원을 처리하였다.
3. 포괄적 진단: 7개의 사전 학습된 모델(Llama 3.1/3.2, Gemma 2/3, Mistral)에 대해 스케일 민감도, 국소 손실 가산성, 고유 차원, 그리고 섭동 기반 레이어 민감도와의 상관관계를 평가하였다.
4. 실증적 검증: 세 가지 모델 규모(128M, 256M, 440M)에 대해 기하학 유도 스케줄을 균일 너비, 수기 설계된 코사인 테이퍼, 그리고 "안티-토폴로지컬(anti-topological)" 제어 모델과 비교하였다.

결과

• 스케일 민감도: 원시 기하학적 작업(유클리드)은 주로 잔차 노름(residual norm)의 성장을 추종하는 반면, 정규화된 메트릭(구형/쌍곡)은 초기 레이어에서 가장 높은 기하학적 작업을 보이는 "전방 집중형(front-loaded)" 프로파일을 드러낸다.
• 민감도 vs 할당: 그로모프-와서스타인 작업은 섭동 기반 레이어 민감도와 더 강한 상관관계를 보였으나, 위상적/쌍곡적 할당이 실제 학습 시 가장 낮은 평균 검증 손실을 기록했다. 이는 민감도와 근사 난이도가 서로 다른 요소임을 시사한다.
• 학습 성능:
  • 128M 및 256M 실험에서, 정규화된 기하학 스케줄은 균일 너비 및 안티-토폴로지컬 제어 모델보다 일반적으로 우수한 성능을 보였다(단, 제한된 시드로 인해 탐색적 결과임).
  • 5개의 시드를 사용한 440M 실험에서, 위상적/쌍곡적 스케줄은 균일 너비 대비 평균 검증 손실을 0.019 감소시켰다. 이는 TLM 스타일의 코사인 테이퍼가 달성한 감소치(0.003)보다 약 6.3배 높은 수치이다.
  • 안티-토폴로지컬/원시 제어 모델은 균일 너비보다 성능이 낮게 나타나, 할당의 방향성이 중요함을 확인시켜 주었다.

의의 및 주장
본 논문은 기하학적 작업이 전역적인 FFN 너비 상수를 레이어별 스케줄로 대체할 수 있는 원칙적이고 데이터 중심적인 방법을 제공한다고 주장한다. 저자들은 다음을 확립하였다:

• FFN 용량은 균일하다고 가정해서는 안 되며, 레이어 업데이트는 깊이에 따라 역할이 변한다.
• 고정된 단조 테이퍼만이 효과적인 대안은 아니다; 활성화 기반 프로파일은 비단조적 구조를 포착할 수 있으며 동일한 조건에서 설계된 스케줄보다 더 큰 성능 향상을 달착할 수 있다.
• 정규화된 기하학적 변화(특히 위상적/쌍곡적 변화)는 효율적인 모델링을 위한 전이 가능한(transferable) 전방 집중형 프로파일을 드러낸다.

저자들은 결과가 기술적인 평균 비교임을 명시하며 결론을 신중하게 내리고 있다(쌍체 차이 불확실성이 완전히 보고되기 전까지). 이들은 특정 기하학(예: 위상 vs GW)의 보편적 우월성을 증명하려 하는 것이 아니라, 정규화된 기하학적 변화가 유용한 신호임을 보여주는 데 목적을 둔다. 이 연구는 활성화 기반 할당을 테이퍼링의 실용적인 확장으로 위치시키며, 기하학적 작업을 레이어별 근사 요구량에 직접적으로 캘리브레이션하는 향후 연구의 동기를 부여한다.