Multi-Head Attention Residuals: Transformer 잔차 흐름에 다중 헤드를 더하다

도입

Transformer의 잔차 연결은 단순히 학습을 안정화하는 장치로 보이지만, 실제로는 층을 가로질러 정보를 전달하는 핵심 통로입니다. 표준 구조에서는 각 sublayer가 주로 가장 최근의 hidden state를 읽고, 변환 결과를 다시 잔차 스트림에 더합니다. 논문 “Multi-Head Attention Residuals”는 이 깊이 방향 정보 흐름이 넓은 모델에서도 충분히 유연한지 묻습니다.

핵심 내용

• 문제는 하나의 공유 라우팅 분포: attention residuals는 learned softmax를 통해 각 sublayer가 이전 깊이 상태를 참조하게 합니다. 하지만 routing query가 전체 모델 폭에서 하나만 공유됩니다. 따라서 서로 다른 특징 하위공간도 같은 층 선택 분포를 써야 합니다.
• 모델이 넓어질수록 타협 비용이 커질 수 있음: 폭이 커지면 내부 특징 하위공간이 더 분화됩니다. 어떤 하위공간은 얕은 층의 정보를, 다른 하위공간은 더 깊은 층의 정보를 선호할 수 있습니다. 단일 softmax는 이런 차이를 하나의 절충안으로 압축합니다.
• MHAR는 잔차 라우팅을 다중 head로 분리: Multi-Head Attention Residuals는 routing query를 H개의 하위공간 head로 reshape합니다. 각 head는 깊이 이력에 대해 독립적인 softmax를 가지며, hidden dimension의 서로 다른 블록이 서로 다른 층을 읽을 수 있게 합니다.
• 추가 부담은 작음: 요약에 따르면 이 reshape는 파라미터를 늘리지 않고 계산 증가도 미미합니다. H=1이면 기존 attention residuals와 정확히 같아지므로, MHAR는 그 자연스러운 확장으로 볼 수 있습니다.
• 여러 규모에서 일관된 개선: 저자들은 중복 제거와 품질 필터링을 거친 Nemotron 기반 anneal corpus, 특히 STEM과 코드 비중이 높은 데이터로 처음부터 모델을 학습했습니다. 표준 Transformer 대비 100M, 350M, 1B에서 각각 -0.061, -0.149, -0.140의 검증 손실 개선을 보고했으며, 비교한 4개 방법 중 모든 설정에서 가장 좋은 결과를 냈다고 합니다.
• head 수는 중요한 설계 변수: H를 늘린다고 항상 좋아지는 것은 아닙니다. 검증 손실은 H에 대해 U자형을 보였고, H=4 또는 H=8 부근에서 비교적 평평한 최적점을 보였습니다. H=16에서는 일부 이득을 반납했습니다.

의미와 영향

MHAR의 장점은 Transformer의 큰 구조를 바꾸지 않고 잔차 스트림의 읽기 방식을 더 정교하게 만든다는 점입니다. 토큰 방향 attention에서 여러 head가 서로 다른 패턴을 보는 것은 이미 표준적인 설계입니다. MHAR는 같은 아이디어를 깊이 방향으로 가져와, 특징 블록마다 필요한 과거 층을 다르게 선택하게 합니다.

다만 이 결과는 특정 데이터 구성, 모델 규모, 구현 조건에서 나온 것입니다. 더 다양한 코퍼스와 아키텍처에서도 같은 경향이 유지되는지는 추가 검증이 필요합니다. 또한 실용성 측면에서는 커널 최적화가 중요합니다. 논문 요약은 fused Triton routing kernels를 통해 attention-residual 학습 처리량을 baseline의 0.2-0.5x에서 0.55-0.88x로 높이고, peak memory는 baseline에 가깝게 유지했다고 설명합니다.

전체적으로 MHAR는 잔차 흐름을 단일 수동 통로가 아니라 하위공간별 깊이 라우팅 메커니즘으로 보는 제안입니다. 보고된 확장성 추세가 더 넓게 재현된다면, LLM 아키텍처에서 품질 향상을 위한 가벼운 설계 축으로 주목받을 수 있습니다.

출처: Hugging Face Daily Papers