Feed-Forward Steering in Transformer Residual Dynamics

기술 요약: 트랜스포머 잔차 역학에서의 피드포워드 스티어링(Feed-Forward Steering)

문제 정의
기존의 트랜스포머 역학에 관한 연속 이론은 주로 어텐션 메커니즘에 의해서만 구동되는, 구체(sphere) 위에 집합하는 입자 시스템으로서 잔차 스트림(residual stream)을 모델링한다. 이러한 "어텐션 전용(attention-only)" 프레임워크는 랭크 붕괴(rank collapse)나 클러스터링과 같은 현상을 설명하지만, FFN(Feed-Forward Network)의 역할을 설명하는 데는 실패한다. 동시에 해석 가능성(interpretability) 연구는 FFN을 키-값 메모리의 저장소이자 효과적인 편집 대상으로서 식별하고 있으나, 이러한 관점들은 역학적 시스템 관점과 연결되지 않은 채 남아 있다. 본 연구가 다루는 핵심 문제는 FFN을 토큰 방향의 기하학적 영향을 설명하고 순차적 블록을 병렬 블록으로 근사하는 것의 타당성을 결정하기 위해, 잔차 스트림 진화의 통합된 역학 모델로 어떻게 통합할 것인가이다.

방법론
저자들은 FFN을 토큰 상태에 작용하는 국소적 스티어링 필드(local steering field)로 포함함으로써 연속 어텐션 역학 프레임워크를 확장한다. 이 모델은 잔차 스트로 $x_i = r_i u_i$ (여기서 $u_i$는 단위 구 위의 방향, $r_i$는 크기)가 다음과 같이 진화한다고 간주한다:
$$\dot{u}_i = \frac{1}{s_i(t)} P_{u_i} \left[ A_i(U) + \Phi_i(u_i) \right]$$
여기서 $A_i(U)$는 어텐션으로부터 오는 비국소적 집합 필드(non-local aggregation field)를 나타내며, $\Phi_i(u_i)$는 FFN에 의해 제공되는 국소적 스티어링 필드이다. $P_{u_i}$는 접공간(tangent space)으로의 직교 투영으로, 벡터장의 접선 성분만이 방향을 변화시키고 반경 성분은 크기(및 그에 따른 각속도)에 영향을 미치도록 보장한다.

본 연구는 GPT-2, Pythia, Mistral, Llama 모델을 대상으로 네 가지 주요 실험 프로토콜을 수행한다:

1. 각 결손 진단(Angular-Deficit Diagnostic): 어텐션 전용 업데이트와 실제 블록 업데이트 사이의 정렬도를 측정하여 방향 변화에 대한 FFN의 기여도를 정량화한다.
2. 반경/접선 절제(Radial/Tangential Ablation): FFN 출력을 반경 성분($\langle u, \Phi(u) \rangle u$)과 접선 성분($P_u \Phi(u)$)으로 분해하여 어떤 성분이 방향 변화를 주도하는지 격리한다.
3. 집합 하의 다양성(Diversity under Aggregation): 어텐션 및 FFN 레이어 전후의 구형 다양성 통계를 분석하여 FFN이 어텐션에 의한 집중 현상을 어떻게 상쇄하는지 관찰한다.
4. 순차-병렬 결함(Sequential-to-Parallel Defect): 레이어별 민감도 점수($C_\ell$)를 추정하여, 순차적 블록($x + A(x) + \Phi(x+A(x))$)을 유의미한 성능 저하 없이 병렬 서로게이트($x + A(x) + \Phi(x)$)로 근사할 수 있는지 결정한다.

추가적인 진단으로는 레이어별 FFN 게이트 캘리브레이션, 임계점 기하학 분석(끌개(attractor), 안장점(saddle), 반발점(repeller) 식별), 그리고 프리픽스 민감도 테스트가 포함된다.

핵심 기여

1. 집합-스티어링 프레임워크(Aggregation-Steering Framework): 본 논문은 트랜스포머 잔차 역학을 어텐션은 비국소적 집합을 제공하고 FFN은 국소적 스티어링을 제공하는 구체 위의 집합-스티어링 흐름으로 공식화한다.
2. 접선 스티어링의 식별: 저자들은 FFN 필드의 접선 성분이 잔차 방향 공간에서의 운동에 필수적임을 수학적, 경험적으로 입증한다. 반경 성분은 주로 크기를 재조정하는 반면, 접선 성분은 방향을 결정한다.
3. 임계 방향 기하학(Critical Direction Geometry): 본 연구는 임계 잔차 방향이 비선형 투영 평형(nonlinear projective equilibria)에 해당함을 식별한다. FFN이 안정적인 끌개를 생성한다는 직관과 달리, 분석 결과 FFN 스티어링이 다양성을 보존하기 위해 국소 기하학을 안장점 및 회전 수송 모드로 능동적으로 재구성함을 보여준다.
4. 순차-병렬 결함 메트릭: 순차적인 어텐션-FFN 블록이 정확하게 병렬적인 가산 맵(additive map)으로 근사될 수 있는 시점을 특징짓는 직접적인 결함 메트릭을 도출한다. 이 결함은 리 브래킷(Lie bracket)뿐만 아니라 방향 도함수 항 $D_\Phi A$에 의해 지배된다.

결과

• 각도 불일치(Angular Misalignment): 어텐션 브랜치 단독으로는 실제 블록 업데이트에 비해 상당한 각도 정렬 결손을 보인다. 이 결손은 모델 규모가 커짐에 따라 GPT-2의 0.41에서 Llama-3-8B의 0.63까지 증가하며, 이는 FFN이 단순한 이차적 보정이 아니라 각도 업데이트의 주요 동인임을 나타낸다.
• 절제 연구 결과: 접선 FFN 성분만을 유지하면 대부분의 모델 품질이 보존되지만(손실 증가 0.07–1.72), 반경 성분만을 유지하면 성능이 붕괴되어(손실 증가 7.19–12.96) FFN을 완전히 제거한 것과 유사한 결과를 보인다.
• 다양성 보존: 어텐션은 일반적으로 잔차 방향을 수축시켜(다양성 감소) 집합을 유도하는 반면, FFN은 종종 이러한 다양성을 복구하거나 재형성한다. 접선 FFN 성분은 집합 압력 하에서도 과도한 집중을 방지하는 안장점 및 회전 수송 모드를 주입한다.
• 병렬화 가능성: 낮은 순차-병렬 결함(낮은 민감도 점수)을 가진 레이어들은 미미한 성능 저하와 함께 대략적으로 병렬화될 수 있다. 예를 들어, GPT-2-large에서 결함이 적은 18개 레이어를 병렬화하면 성능 저하 +0.02와 함께 약 1.24배의 속도 향상을 얻을 수 있는 반면, 결함이 높은 레이어를 병렬화하면 성능이 급격히 저하된다.
• 임계 기하학: 임서 방향 필드의 성분 절제 분석 결과, FFN 후보들은 주로 안장점 형태(혼합 부호 스펙트럼)를 띤다. FFN은 단순히 안정적인 끌개를 만드는 것이 아니라, 고차원 표현 기하학을 유지하기 위해 위상도(phase portrait)를 재구조화한다.

의의 및 주장
본 논문은 FFN 레이어가 트랜스포머 잔차 기하학을 형성하는 방향적 스티어링 필드라는 해석을 뒷받침한다고 주장한다. 저자들은 FFN이 블록 수준의 개입 가능성을 제어하고, 어텐션의 붕괴 압력에 맞서 표현 다양성을 유지하는 데 필수적이라고 주장한다.

저자들은 자신의 연구 결과를 역학 시스템 이론과 해석 가능성 사이의 가교로 위치시키며, "스티어링" 관점이 왜 FFN이 효과적인 편집 대상이며 왜 모델의 기하학적 안정성에 필요한지를 설명한다고 제안한다. 실질적인 의의는 추론 지연 시간을 줄이기 위해 안전하게 병렬화할 수 있는 레이어를 식별하는 데 있다.

본 논문은 범위에 대해 신중한 태도를 유지하며, 이론적 평균장(mean-field) 결과는 전체 어텐션을 설명하는 반면, 인과적 디코더 전용 모델은 프리픽스 인덱스 역학(prefix-indexed dynamics)을 필요로 한다는 점을 명시한다. 임계점 분석은 전체 잔차 공간에 대한 전역적 기술이라기보다 국소적 진단으로 제시된다. 또한, 실질적인 병렬화 결과는 최적화된 프로덕션 커널 성능에 대한 주장이 아니라 추정된 서브모듈 시간 속도 향상으로 프레임화된다. 저자들은 자신들의 발견이 진단적이며, 모든 아키텍처에 대한 안티-콜랩스(anti-collapse) 또는 완벽한 병렬화에 대한 보편적 보장을 주장하는 것이 아님을 명시적으로 밝힌다.