Analogies between Transformer Layers and Power Method

기술적 요약: 트랜스포머 레이어와 파워 메서드 간의 유사성

문제 제기
이 논문은 깊은 트랜스포머에서 토큰 표현이 어떤 특정 벡터 방향으로 수렴하는지에 대한 이해 부족을 다룹니다. 이 현상은 '과평활화 (oversmoothing)', '랭크 붕괴 (rank collapse)', 또는 '토큰 균일화 (token-uniformity)'로 알려져 있습니다. 깊은 트랜스포머가 모든 토큰이 동일해지는 합의 상태 (consensus state) 로 향한다는 점은 확립되었으나, 이 합의 벡터의 본질과 이를 주도하는 메커니즘은 여전히 불분명합니다. 기존 문헌은 종종 자기주의 (self-attention) 를 연속 시간 시스템으로 모델링하거나 단일 헤드 시나리오에 집중하여, GPT2, BERT, ALBERT 와 같은 실제 모델의 이산 시간 (discrete-time), 다중 헤드 행동을 덜 탐구했습니다. 구체적으로, 이 논문은 토큰을 트랜스포머 레이어를 통과시키는 반복 과정이 주 고유벡터를 찾기 위한 알려진 수치 방법과 유사한지 확인하고자 합니다.

방법론
저자들은 레이어 인덱스 $t$를 시간으로 간주하는 이산 시간 동역학 시스템으로서 트랜스포머를 모델링합니다. 그들은 피드포워드 네트워크 (FFN) 를 명시적으로 배제하고, 어텐션 및 프로젝션 메커니즘의 기하학적 효과를 분리하기 위해 표준 레이어 정규화 (LayerNorm) 를 대신하여 $L_2$ 정규화를 사용하는 단일 레이어를 통한 토큰 벡터 $x_i(t)$의 업데이트 법칙에 초점을 맞춥니다.

방법론의 핵심은 각 레이어 $t$에 대해 출력 가중치 행렬 ($O$) 과 어텐션 헤드의 연결된 값 가중치 행렬 ($V$) 의 곱으로 구성된 '레이어 행렬' $F_t$를 정의하는 것입니다:
$$F_t = O_t V_t$$
업데이트 규칙은 동역학 시스템의 한 단계로 분석됩니다:
$$x_i(t+1) = \text{LayerNorm}\left(x_i(t) + \text{AttentionOutput}\right)$$
저자들은 이 업데이트와 행렬의 주 고유벡터를 계산하는 데 사용되는 반복 알고리즘인 파워 메서드 (Power Method) 간의 직접적인 유사성을 제시합니다. 파워 메서드에서 벡터는 행렬에 의해 반복적으로 곱해지고 정규화되며, 스펙트럼 갭이 존재할 경우 지배적인 고유벡터로 수렴합니다.

분석은 두 가지 주요 시나리오로 나뉩니다:

1. 공유 가중치 (시간 불변): 모든 레이어에서 가중치가 동일한 ($F_t = F$) ALBERT 모델을 사용하여 분석됩니다. 저자들은 주 고유값이 실수이고 엄격하게 지배적인 경우, 복소수인 경우, 또는 지배적이지 않은 경우에 대한 엄격한 분석적 특성을 제공합니다.
2. 레이어별 가중치 (시간 가변): BERT, GPT2, GPTNeo 를 사용하여 분석됩니다. 여기서 $F_t$는 각 레이어마다 변합니다. 저자들은 목표가 변함에도 불구하고 토큰이 여전히 현재 레이어의 주 고유벡터 $v_{1,t}$와 정렬되는 경향을 보이는지 조사합니다.

이론을 검증하기 위해 저자들은 사전 훈련된 모델에 대한 수치 실험을 수행하여 토큰 벡터와 레이어 행렬의 주 고유벡터 간의 상관관계를 측정합니다. 또한 최종 레이어를 수정하여 출력을 '조종 (steering)'하는 방법을 제안합니다.

주요 기여

1. 파워 메서드와의 유사성: 이 논문은 FFN 을 무시한 트랜스포머 레이어의 연산이 파워 메서드의 단일 단계와 유사하게 작용함을 확립합니다. 토큰은 투영되고 정규화되어 레이어 행렬 $F_t$의 주 고유벡터 방향으로 기울어지는 경향이 있습니다.
2. 수렴의 분석적 특성화:
   • 공유 가중치 트랜스포머 (예: ALBERT) 의 경우: 저자들은 레이어 행렬 $F$가 엄격하게 지배적인 실수 고유값을 가지면, 시스템이 모든 토큰이 주 고유벡터 $v_1$과 정렬되는 합의 상태로 수렴함을 증명합니다.
   • 안정성 조건을 규명하여, 주 고유값이 복소수인 경우 토큰이 고유벡터의 실수부와 허수부가 span 하는 2D 부분 공간에서 극한 주기 (limit cycle) 로 수렴함을 지적합니다.
   • '이분 합의 (bipartite consensus)'(토큰이 $\pm v_k$와 정렬된 두 그룹으로 나뉘는 것) 가 이론적으로 가능하지만, 다중 헤드 설정에서는 경험적으로 드물고 불안정함을 보여줍니다.
3. 실제 모델에 대한 경험적 검증:
   • ALBERT에 대한 실험은 토큰이 공유 레이어 행렬의 주 고유벡터로 수렴함을 확인합니다. 그러나 수렴 속도는 종종 작은 스펙트럼 갭 ($\lambda_1 \approx \lambda_2$) 이나 복소수 고유값에 의해 제한됩니다.
   • **레이어별 가중치 모델 (BERT, GPT2, GPTNeo)**에 대한 실험은 토큰이 변하는 $F_t$로 인해 단일 전역 고유벡터로 정렬이 누적되지는 않지만, 각 단계에서 여전히 현재 레이어의 주 고유벡터 $v_{1,t}$와의 상관관계가 통계적으로 유의미하게 증가함을 보여줍니다. 이는 시간 가변 시스템에서도 '파워 메서드' 행동이 국소적으로 유지됨을 확인합니다.
4. 랭크 -1 섭동에 의한 토큰 조종: 이 논문은 임의의 원하는 방향으로 토큰을 조종하기 위한 구성적 방법을 제안합니다. 최종 레이어의 출력 행렬에 충분히 큰 랭크 -1 섭동을 추가함으로써 큰 스펙트럼 갭을 생성할 수 있습니다. 이는 최종 레이어 행렬이 새로운 사용자 정의 주 고유벡터를 갖게 하여 모든 토큰이 이 특정 방향으로 정렬되게 합니다.

결과

• ALBERT: 공유 가중치를 가진 모델에서 토큰은 빠르게 합의에 수렴합니다. ALBERT-base 및 large 의 경우 주 고유값은 실수이지만 두 번째 고유값과 가까워 첫 번째와 두 번째 고유벡터의 span 내에서 수렴합니다. xlarge 및 xxlarge 의 경우 복소수 고유값으로 인해 토큰이 2D 부분 공간에서 회전합니다.
• 레이어별 가변 모델: BERT 및 GPT 변형 모델에서 현재 레이어의 주 고유벡터와의 정렬 ($\gamma_t$) 은 양수이지만 일반적으로 약합니다. 이 약함은 실제 트랜스포머의 작은 스펙트럼 갭 ($\lambda_{1,t} - \lambda_{2,t}$) 에 기인합니다. 그러나 각 단계에서 정렬이 증가할 확률 ($\rho_t > 0.5$) 은 일관되게 관찰되어 유사성을 검증합니다.
• 조종: 최종 레이어를 랭크 -1 섭동으로 수정하면 스펙트럼 갭이 커지고, 토큰은 모델 아키텍처와 관계없이 새로 유도된 주 고유벡터와 거의 완벽한 정렬을 이룹니다.

의의 및 주장
이 논문은 파워 메서드의 관점을 통해 트랜스포머를 바라보는 것이 토큰 과평활화 및 합의에 대한 원칙적인 설명을 제공한다고 주장합니다. 합의의 '방향'이 임의적인 것이 아니라 값 (value) 과 출력 가중치 행렬의 곱의 스펙트럼 속성에 의해 결정됨을 시사합니다.

저자들은 표준 트랜스포머에서의 자연스러운 정렬이 작은 스펙트럼 갭으로 인해 약하지만, 그 메커니즘은 활용하기에 충분히 견고하다고 강조합니다. 제안된 조종 방법은 모델 출력을 제어하기 위한 기존 기법 (프롬프트 엔지니어링이나 활성화 편집 등) 에 대한 원칙적인 대안을 제공합니다. 이는 최종 레이어의 기하학적 투영 속성을 직접 조작하기 때문입니다. 논문은 FFN 을 무시하더라도 이 유사성이 성립한다고 결론지으며, 어텐션 메커니즘과 레이어 정규화가 이 고유벡터 정렬 행동의 주요 동인임을 시사합니다.