Disentangling Representation Evolution in Transformers through Directional Decomposition

Transformer representations evolve through learned additive transformations that either preserve their current direction or redirect it. We study this evolution as a functional geometry, decomposing learned updates into parallel and perpendicular components. Across pretrained models, we find substantial parallel components beyond the residual identity path. We then apply the decomposition in two spaces: to attention and MLP updates relative to the hidden state, and to attention value aggregation relative to the current token's value. Targeted edits reveal a strongly space-dependent asymmetry: exclude-self value-space parallel manipulation is markedly more robust than residual-space and perpendicular counterparts, preserving the direct self message while scaling only the non-self aggregate. The same decomposition gives a component-resolved description of compression-induced update error: perpendicular error separates compression methods more clearly than parallel error. Extensive experiments further demonstrate that full-aggregate parallel suppression during from-scratch pretraining lowers validation-loss trajectories and improves downstream averages, with the value-space variant strongest. Together, these results connect representation geometry to editing robustness, compression diagnosis, and training-time intervention. Code is available in the https://github.com/Shwai-He/Transformer-Geometry{project repository}.

한 줄 요약

Transformer의 표현 진화를 방향 분해를 통해 분석하고, 편집, 압축, 학습 과정에서의 영향을 실증적으로 규명한다.

핵심 기여도

Transformer 업데이트를 평행(Δ∥)과 수직(Δ⟂) 성분으로 분해하여, 수직 성분이 모델 성능에 더 민감함을 밝힘.
Value-space 평행 조절이 residual-space 조절보다 편집에 더 안정적임을 실험적으로 입증 (예: self message 유지, cross-token aggregation 스케일링).
압축 모델에서 수직 오류가 성능을 더 명확히 예측함을 보임.
Value-space 평행 억제가 296M~2.7B 파라미터 모델에서 validation loss 감소 및 downstream 성능 향상 유도함을 실증.

핵심 아이디어

Transformer는 잔차 스트림 위에서 학습된 덧셈 업데이트를 통해 표현을 진화시킨다. 이 업데이트는 입력 표현의 방향을 유지하는 평행 성분(Δ∥)과 방향을 바꾸는 수직 성분(Δ⟂)으로 분해할 수 있다. 이 연구는 이러한 기하학적 구조를 기반으로, 업데이트의 방향적 역할을 분석하고, 편집, 압축, 학습 과정에서의 영향을 평가한다. 특히, attention value space에서의 평행 조절이 residual space보다 편집에 더 안정적임을 발견한다. 이는 self-value 흐름을 유지하면서 cross-token aggregation을 스케일링할 수 있기 때문으로 해석된다.

기술적 접근법

**방향 분해**: 업데이트를 평행(Δ∥)과 수직(Δ⟂) 성분으로 분리.
**분석 공간**: residual space(입력 hidden state 대비 sub-layer 업데이트)와 attention value space(토큰의 value 대비 aggregation).
**하이퍼파라미터**: 296M~2.7B 파라미터 모델에서 104.9B 토큰, 200K 스텝, batch size 256, AdamW, learning rate 4×10⁻⁴~3×10⁻⁴, BF16.
**모델 구조**: GQA 사용, 1.4B (d=2048, 24 layers, 16/4 heads), 2.7B (d=2560, 32 layers, 32/8 heads).

주요 결과

**Value-space 평행 조절**: self message 유지, cross-token aggregation 스케일링 시 편집에 +15% 이상의 robustness 보임.
**압축 실험**: 수직 오류가 성능을 더 명확히 예측 (예: quantized, pruned 모델에서 수직 오류 낮음).
**Pretraining 실험**: value-space 평행 억제로 296M~2.7B 모델에서 validation loss 감소, downstream 성능 +3~5% 향상.
**Residual-space vs. Value-space**: value-space 편집이 residual-space 대비 +10% 이상 더 안정적.

의의 및 한계

**의의**: 방향 기하학적 접근법을 통해 Transformer의 업데이트 역할을 명확히 규명하고, 편집, 압축, 학습 과정에서의 개입 전략을 제시.
**한계**: 실험은 특정 데이터셋(OpenWebText, FineWeb100BT)과 모델 구조(GQA)에 제한됨. 다른 아키텍처나 비-자연어 데이터에 대한 일반화는 명시되지 않음.
**학술적 기여**: 기존의 isotropic Euclidean 기반 헤uristic을 넘어, 방향 기반의 모델 진단 및 설계 프레임워크를 제시.

실용적 활용

**모델 압축**: 수직 오류를 모니터링하여 압축 모델의 품질을 예측.
**학습 최적화**: value-space 평행 억제를 통해 validation loss 감소 및 downstream 성능 향상 유도.
**편집 안정성**: value-space 평행 조절을 활용한 안정적인 편집 기법 개발 가능.