xHC: Transformer 잔차 스트림을 N=16까지 확장하는 방법

도입

대규모 언어 모델의 확장은 보통 파라미터 수, 깊이, 폭, 데이터, 계산량을 중심으로 논의됩니다. xHC는 그보다 더 내부적인 구조인 Transformer의 잔차 스트림에 주목합니다. Hyper-Connections(HC)는 잔차 스트림을 N개의 병렬 스트림으로 확장해 모델 내부의 메모리 용량을 늘리는 방식으로 볼 수 있습니다. 하지만 기존 HC 계열 방법은 대체로 N=4에서 멈췄습니다. xHC는 잔차 스트림 폭을 더 키워도 대규모 사전학습에서 실제 이득을 얻을 수 있는지 묻습니다.

핵심 내용

• 가능성은 있었지만 확장은 어려웠습니다. HC는 N=1에서 N=4로 갈 때 큰 성능 향상을 보여, 잔차 스트림 확장이 유망한 스케일링 축일 수 있음을 시사했습니다. 그러나 저자들의 실험에 따르면 mHC를 그 이상으로 키우면 성능 이득은 줄어들고 훈련 비용은 빠르게 증가합니다.
• 병목은 두 가지로 정리됩니다. 첫째, 스트림 수가 많아질수록 각 스트림에 충분한 정보를 다시 써 넣기 어렵습니다. 둘째, 잔차 혼합 생성 비용이 N에 대해 세제곱으로 증가해 큰 N에서 비효율이 커집니다.
• xHC는 전체를 보되 일부만 갱신합니다. temporal feature augmentation을 사용해 write-back에 필요한 정보를 더 풍부하게 만들고, N=16 스트림 중 매번 k=4개만 업데이트하는 희소 잔차 스트림 구조를 사용합니다. 동시에 전체 잔차 상태에 대한 밀집 접근은 유지합니다.
• 대규모 MoE에서 효과가 보고됐습니다. 18B MoE 모델에서 xHC는 DeepSeek mHC 대비 최종 훈련 loss를 1.776에서 1.758로 낮추고, 평균 다운스트림 점수를 44.8에서 48.8로 높였습니다. vanilla baseline 대비 추가 훈련 FLOPs는 4.1%에 그쳤습니다. 28B MoE에서도 xHC는 mHC보다 평균 3.1점 높은 결과를 냈습니다.
• 효율성도 중요한 주장입니다. scaling-law 실험에 따르면 같은 loss에 도달하기 위해 vanilla와 mHC는 각각 xHC의 1.50배, 1.19배 계산량이 필요합니다. 또한 xHC-Flash는 서브레이어당 메모리 트래픽을 73.5C에서 40C로 줄여, N=4 mHC의 34C에 가까운 수준을 제시합니다.

의미와 영향

xHC의 핵심은 단순히 잔차 스트림을 많이 늘렸다는 점이 아닙니다. mHC가 N=4 이후 왜 효율을 잃는지 분석하고, 희소 업데이트로 큰 N을 다루기 쉽게 만들었다는 데 있습니다. 이 접근은 잔차 스트림 확장을 구조적 아이디어에서 실제 대규모 훈련에 적용 가능한 설계로 한 단계 끌어올립니다.

논문의 결과가 더 다양한 아키텍처와 훈련 조건에서도 재현된다면, 잔차 스트림 폭은 모델 깊이, 폭, MoE 전문가 수, 데이터 규모와 함께 새로운 확장 차원이 될 수 있습니다. 특히 MoE 사전학습에서는 훈련 FLOPs를 크게 늘리지 않고 표현 용량을 높이는 선택지가 될 가능성이 있습니다. 다만 현재 근거는 논문 내 실험에 기반하므로, 독립 재현과 다른 모델군에서의 안정성 검증은 여전히 필요합니다.

출처: Hugging Face Daily Papers