Full-bandwidth transformer

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.

한 줄 요약

Full-bandwidth transformer는 토큰 단위 피드백 대신 전체 히든 상태를 재사용해 추론 성능을 향상시키며, 1.5배 더 많은 토큰으로 학습한 모델과 유사한 결과를 낸다.

핵심 기여도

**Latent feedback decoding**을 도입하여, 토큰이 아닌 전체 히든 상태를 다음 입력으로 재사용함.
**Gated linear unit**을 통해 이전 토큰 임베딩과 상위 레이어 히든 상태를 결합.
**Scheduled multi-pass objective**를 사용해 학습 시 병렬 teacher forcing 유지.
1B 파라미터 모델을 400B 토큰까지 학습하여, **5-shot 평가, 수학 및 코드 생성, 지시어 조정 성능** 개선 확인.

핵심 아이디어

기존 autoregressive transformer는 토큰 하나만 다음 입력으로 전달하며, 이 과정에서 풍부한 히든 상태 정보가 손실된다. 이에 반해, Full-bandwidth transformer는 **latent feedback decoding**을 통해, **이전 상위 레이어 히든 상태**를 다음 입력으로 재사용함으로써, **비언어화된 계산 결과**가 다시 모델의 하위 레이어에서 재처리될 수 있도록 한다. 이는 **gated linear unit**을 사용해 토큰 임베딩과 히든 상태를 결합한 후, 다음 입력으로 전달하는 방식이다. 이 접근법은 기존의 KV 캐시와 언어 모델링 목적 함수를 유지하면서도, **depth budget**을 재할당해 더 깊은 계산을 가능하게 한다.

기술적 접근법

**Latent feedback decoding**: 이전 상위 레이어 히든 상태와 토큰 임베딩을 **gated linear unit**으로 결합.
**Scheduled multi-pass objective**: 학습 시 **sequential recurrence**를 피하기 위해, **progressive schedule**과 **prefix mixin**을 사용.
**NorMuon** 최적화기 사용 (학습률 1×10⁻², weight decay 0.01), **Adam** (학습률 5×10⁻⁴, weight decay 없음).
**WSD 학습률 스케줄** 적용 (warmup 200, cooldown 25%, z-loss 계수 1×10⁻⁵).
**Jitter noise** (σ=0.02)를 학습 시 추가.

주요 결과

**GSM8K, Math500, HumanEval, MBPP** 데이터셋에서 **latent feedback decoding**이 표준 decoding 방식 대비 성능 향상.
**2-pass** 학습 시, **1.5× 더 많은 토큰으로 학습한 모델**과 유사한 성능 달성.
**3-pass** 학습 시, 일부 작업에서 **5× 더 많은 토큰으로 학습한 모델**과 유사한 성능.
**5-shot 평가, 수학 및 코드 생성, 지시어 조정**에서 개선 확인.
**추론 토큰당 오버헤드는 거의 없음**.

의의 및 한계

Full-bandwidth transformer는 기존 transformer 아키텍처를 유지하면서도, **추론 시 계산 효율성과 정확도를 동시에 향상**시킬 수 있다. 특히, **비언어화된 계산**을 재사용함으로써, **더 짧은 추론 추적**을 유지하면서도 **동일 또는 더 높은 정확도**를 달성하는 것이 가능하다. 그러나, 학습 시 **sequential recurrence**를 피하기 위해 복잡한 **multi-pass objective**가 필요하며, 이는 학습 과정에서 **안정성**을 유지하기 위해 **3% 이상의 3-pass 배치**가 필요하다는 한계가 있다. 또한, **latent feedback decoding**은 **모델 구조 변경 없이도 가능**하지만, **추가적인 forward pass**가 필요하다는 점에서 **계산 비용**이 발생할 수 있다.

실용적 활용

Full-bandwidth transformer는 **수학 문제 풀이, 코드 생성, 지시어 기반 작업** 등에서 **더 높은 정확도와 효율성**을 요구하는 산업 및 연구 분야에 적용 가능하다. 특히, **짧은 추론 추적**과 **높은 정확도**가 필요한 **대규모 언어 모델의 추론 최적화**에 유용하며, **추가 토큰 없이도 성능 향상**이 가능한 점에서 **데이터 효율성** 측면에서도 활용 가능하다.