A Controlled Study of Attention-Only Transformers

기술 요약: 어텐션 전용 트랜스포머에 관한 통제 연구

문제 정의
피드포워드 네트워크(FFN)는 트랜스포머의 임베딩 제외 파라미터 중 약 3분의 2를 차지하며, 모델의 파라미터적 메모리(Geva et al. 2021, 2022; Meng et al. 2022; Dai et al. 2022)로 널리 해석된다. 그러나 FFN의 필요성은 통제된 조건 하에서 엄격하게 테스트된 적이 없다. FFN을 제거하려는 이전의 시도들은 FFN의 용량을 지속적 메모리 벡터(persistent memory vectors)를 통해 어텐션으로 다시 통합하거나(Sukhbaatar et al. 2019), 실제로는 사용되지 않는 잔차 연결(residual connections)이 없는 이론적 설정(Dong, Cordonnier, and Loukas 2021)을 분석함으로써 FFN의 필요성 문제를 미결 과제로 남겨두었다. 핵심적인 실험적 과제는 FFN을 제거할 때 파라미터 수, 토큰당 연산량(FLOPs), 그리고 비선형 합성 단계의 수(깊이)라는 세 가지 결정적인 변수가 동시에 변화한다는 점이다. 따라서 단일 축 제어만으로는 FFN의 진정한 비용을 분리해낼 수 없다.

방법론
저자들은 FFN이 제거된 트랜스포머인 단순 어텐션 네트워크(Simple Attention Networks, SANs)를 표준 트랜스포머와 대조하여 사전 학습함으로써 "필요성 테스트"를 수행한다. 혼란 변수를 제어하기 위해, 본 연구는 세 가지 별도의 매칭 전략(표 1)을 채택한다:

1. 등깊이(Iso-depth): 레이어 수를 일치시킨다(20개 레이어). 이는 표준 트랜스포퍼에서 72%의 파라미터를 제거한다.
2. 등파라미터(Iso-param): 총 임베딩 제외 파라미터 수를 일치시킨다. 확보된 예산은 어텐션 깊이로 재할당된다(20개의 SAN 레이어 대 4개의 표준 레이어).
3. 등연산량(Iso-FLOP): 훈련 연산량(토큰당 GFLOPs)을 일치시킨다. 어텐션은 시퀀스 길이에 대해 파라미터가 없는 이차 비용(quadratic cost)을 발생시키므로, 이 체제에서는 표준 트랜스포머가 더 많은 파라미터를 보유하게 된다.

하이퍼파라미터 공정성 및 제어:

• 하이퍼파라미터 공정성: 각 아키텍처 암(arm)은 5B 토큰 체크포인트에서 자체적인 학습률 스윕(0.5× ~ 16×)을 거쳐, 특정 암이 다른 암의 선호하는 하이퍼파라미터에서 경쟁하지 않도록 보장한다.
• 옵티마이저 제어: 훈련에는 AdamW 대조군에 대해 Mu-on을 사용하며, 옵티마이저 등가 제어를 통해 0.0015 nats의 보정된 노이즈 플로어를 설정하였다.
• 데이터 체제: 모델은 추론 밀도가 높은 합성 코퍼스인 SYNTH에서 최대 105B 토큰까지 사전 학습된다. 별도로 사전 등록된 테스트는 지식 밀도가 높은 Fineweb-edu에서 수행된다.
• 아키텍처: SAN은 Grouped Query Attention (GQA), RoPE, 그리고 특정 정규화 전략(ZCN, QK-normalization)을 사용하는 pre-norm 어텐션을 사용한다.

주요 결과
본 연구는 매칭 전략에 따른 SAN과 표준 트랜스포머 사이의 성능 격차를 단조적인 순서로 분해한다:

1. 제거 비용 (Iso-depth): 매칭된 깊이에서 FFN을 삭제하면 검증 손실(validation loss)에서 0.47 nats의 상당한 손실이 발생한다.
2. 연산 효율성 (Iso-FLOP): 매칭된 훈련 FLOPs에서 표준 트랜스포머가 여전히 0.26 nats 앞서는데, 이는 어텐션 전용 모델이 파라미터가 없는 이차 항에 더 많은 연산을 소비하여 용량을 위한 파라미터가 적어지기 때문이다.
3. 파라미터 효율성 (Iso-param): 확보된 파라미터 예산을 어텐션 깊이로 재할당했을 때, 격차는 깨끗한 시드 쌍(clean seed pairs)에서 매우 재현 가능한 +0.0055 ~ +0.0054 nats로 좁혀진다. 이 차이는 5B, 30B, 105B 토큰 예산 전반에 걸쳐 안정적이며, 임베딩 제외 파라미터 크기의 29배 범위 내에서 약 0.02 nats 수준을 유지한다.

잔여 격차의 국소화(Localization of the Residual Gap)
남아있는 성능 결손은 분산되어 있지 않고 엄격하게 국소화되어 있다:

• 토큰 영역: 결손은 완전히 저문맥 쿼리 토큰(라우팅할 정보가 거의 없는 곳)에 집중된다. 추론 궤적 및 답변 영역(암기 연습 포함)에서 SAN은 종종 FFN 모델과 대등하거나 오히려 능가한다. 105B 토큰에 도달하면 SAN은 모든 답변 영역에서 앞서며, 결손은 오직 쿼리 토크에서만 생존한다.
• 태스크 유형: 분포 외 회상(out-of-distribution recall)을 요구하는 벤치마크(예: Lambada)에서는 FFN 모델이 일관되게 승리한다. 답이 문맥 내에 제공되는 태스크(예: SciQ)에서는 SAN이 승리하며, 훈련이 진행됨에 따라 그 격차가 커진다.
• 분포 의존성: 지식 밀도가 높은 텍스트인 Fineweb-edu에 대한 사전 등록 테스트 결과, 0.040 nats의 격차가 측정되었으며, 이는 FFN의 이점이 저문맥, 지식 밀집 텍상에서 가장 크다는 예측을 확인시켜 준다. 반대로, Fineweb-edu로 훈련된 Lambada에서는 선호도가 역전되었는데, 이는 "저장 vs 라우팅"이 태스크의 고유한 속성이 아니라 태스크-분포 간의 매치(match)에 따른 속성임을 나타낸다.

메커니즘: 가중치 스펙트럼 역학(Weight-Spectrum Dynamics)
논문은 가중치 스펙트럼 분석을 통해 이러한 결과의 메커니즘을 식별한다:

• 라우팅 vs 콘텐츠: 라우팅 행렬(Query/Key)은 훈련의 첫 1/4 구간 내에 스펙트럼적으로 결정화(crystallize/freeze)된다. 콘텐츠 행렬은 안정적인 훈련 단계 동안 안정적인 랭크(stable rank)를 축적한다.
• 저장의 재배치: 표준 트랜스포머에서 콘텐츠 축적은 FFN 다운 프로젝션에서 일어난다. SAN에서 이 축적은 어텐션 출력 프로젝션($W_o$)으로 재배치되며, 이곳이 잔차 스트림(residual stream)으로 들어가는 유일한 쓰기 경로(write path)가 된다.
• 안정성: 연구는 깊은(48 레이어) 어텐션 스택을 훈련하는 데 있어 QK-normalization이 핵심 구성 요소임을 밝혀냈다. 이를 제거하면 즉각적인 발산(divergence)이 발생한다. 스칼라 잔차 게이트(scalar residual gates)는 성능 면에서는 중립적이지만 진단용으로는 유용하다는 것이 발견되었다.

의의 및 주장
본 논문은 트랜스포머 FFN에 대한 최초의 통제된 필요성 테스트를 제공한다고 주장한다. 주요 결론은 다음과 같다:

1. FFN은 별개의 계산 프리미티브가 아니다: FFN을 제거한다고 해서 모델이 깨지는 것이 아니라, 오히려 파라미터 용량을 쓰기 경로(write path)에 할당하는 특정 방식일 뿐이다.
2. 함수 형태 vs 용량: 추론 밀도가 높은 데이터에서 FFN의 함수 형태 자체는 크게 중요하지 않다. 중요한 것은 파라미터 용량이다. 이 용량을 어텐션 깊이로 재할당하면 거의 모든 성능을 회복할 수 있다.
3. 특정 결손: 남아있는 유일한 비용은 문맥이 정보를 제공하지 않아 파라미터적 메모리에 의존해야 하는 토큰에서의 작고 예측 가능한 결손이다.
4. 강건성: SAN은 학습률 스트레스 상황에서 표준 트랜스포머보다 더 큰 훈련 안정성을 보이며, FFN 실행 시 발생하는 불안정성과 달리 SAN 실행에서는 발산 이벤트가 관찰되지 않았다.

저자들은 테스트된 범위 내에서, 파라미터 예산이 일치하고 QK-normalization으로 아키텍처가 안정화된다면 "어텐션이 나머지를 수행한다"고 결론짓는다. 이 연구는 FFN을 필수적인 구조적 구성 요소가 아니라, 깊은 어텐션 레이어로 효과적으로 대체될 수 있는 지식을 저장하는 특정 메커니즘으로 재정의하며, 그 대가는 저문맥, 지식 집약적 예측 태스크에서의 국소적인 성능 저하라고 본다.