전대역폭 트랜스포머

초록

자기회귀 트랜스포머는 두 축을 따라 계산을 수행한다: 생성된 토큰을 가로지르는 수평 축과 모델 깊이를 관통하는 수직 축이다. 밀집 어텐션은 각 토큰이 과거에 대해 넓은 수평적 접근을 갖게 하지만, 디코딩 단계 사이의 수직 피드백 채널은 여전히 좁다. 즉, 샘플링된 토큰만 스택의 최하단으로 돌아가고, 최상위 계층의 은닉 상태는 폐기된다. 우리는 잠재 피드백으로 이 채널을 넓히는 전대역폭 트랜스포머를 제안한다. 각 디코딩 단계에서 이전 최상위 계층의 은닉 상태는 게이트 선형 유닛을 통해 샘플링된 토큰 임베딩과 융합된 후 다음 입력으로 피드백된다. 잠재 피드백은 언어화되지 않은 계산이 새로운 깊이 예산을 가지고 스택에 다시 들어갈 수 있게 하면서도, 표준 트랜스포머 아키텍처, KV 캐시, 언어모델링 목적 함수를 유지한다. 병렬 교사 강제를 잃지 않으면서 전대역폭 트랜스포머를 학습시키기 위해, 우리는 사전 학습 후반에 잠재 피드백을 도입하고 안정성을 위해 소량의 더 깊은 피드백 패스를 혼합하는 스케줄된 다중 패스 목적 함수를 사용한다. 우리는 1B 파라미터의 전대역폭 트랜스포머를 최대 400B 토큰까지 학습시켰고, 잠재 피드백이 검증 손실, 5샷 언어모델 평가, 수학 및 코딩 생성, 명령어 튜닝 성능을 향상시킨다는 것을 발견했다. 무시할 수 있는 토큰당 디코딩 오버헤드로, 전대역폭 트랜스포머는 약 1.5배 더 많은 토큰으로 학습된 표준 트랜스포머와 일치하거나 그에 접근하며, 동일하거나 더 나은 정확도로 더 짧은 추론 궤적을 생성한다.

English

Autoregressive transformers compute along two axes: horizontally across generated tokens, and vertically through model depth. Dense attention gives each token broad horizontal access to the past, but the vertical feedback channel between decoding steps remains narrow: only the sampled token returns to the bottom of the stack, while the top-layer hidden state is discarded. We introduce the full-bandwidth transformer, which widens this channel with latent feedback: at each decoding step, the previous top-layer hidden state is fused with the sampled token embedding through a gated linear unit and fed back as the next input. Latent feedback lets non-verbalized computation re-enter the stack with a renewed depth budget, while preserving the standard transformer architecture, KV cache, and language-modeling objective. To train full-bandwidth transformers without losing parallel teacher forcing, we use a scheduled multi-pass objective that introduces latent feedback late in pretraining and mixes a small fraction of deeper feedback passes for stability. We train 1B-parameter full-bandwidth transformers up to 400B tokens and find that latent feedback improves validation loss, 5-shot language-model evaluation, math and coding generation, and instruction-tuned performance. With negligible per-token decoding overhead, full-bandwidth transformers match or approach standard transformers trained with roughly 1.5times more tokens, and manage to produce shorter reasoning traces at equal or better accuracy.