Thread

Xi Wang on "Full-bandwidth Transformer"

The lead author walks through the paper in eight posts: feeding the previous top-layer hidden state back into the input with the token embedding, what that opens across depth, why it costs almost nothing at inference, the multi-pass training it needs and the schedule that keeps the extra cost near 30% of the flops, the gain from extra prefill passes, and free-form generation comparable with models trained with twice the flops. Posted with the paper's first version: the figure and table numbers in the images are that version's.

The posts are embedded from X. The figure notes under them are written by this wiki.

  1. Figure. The paper's Figure 1 with its caption. Two diagrams of three tokens passing through a stack drawn with layers 0 to 3. Left, standard decoding: each token's output goes to the next token's input, and from the state being computed, at layer 1 of token 2, only the layer-0 states of tokens 0 and 1 are marked reachable; their states at layers 1 to 3 are marked unreachable. Right, latent feedback decoding: the top-layer state of each token is joined to the next token's embedding at the input by a symbol ⊗ labeled 'Dimension-preserving fusing', and every past state at every layer is marked reachable.

  2. Figure. Two listings of code side by side. Standard decoding: after prefill, the loop sets x = embed(tok), runs model_step(x, kv_cache) and samples the next token. Latent feedback decoding: the same loop with one line changed, shown in red: x = glu_cross(h_prev, embed(tok)). Above it is the function glu_cross(h, e), which returns (h @ W_u) * sigmoid(e @ W_g).

  3. Figure. A diagram of multi-pass training: three forward passes of an LLM over tokens 0 to 3, stacked from bottom to top. The first pass takes token embeddings and produces hidden states. In each later pass, the hidden state the previous pass produced at one position is fused, by a symbol ⊗, with the token embedding at the next position. Notes on the right: train with multiple forward pass, bottom to top; apply the NTP loss on each pass's hidden state; a practical efficiency trick, a scheduled number of passes: the first 75% of training on one pass only, then 22% on two passes, to get the model familiar with fused inputs, and 3% on three passes, to get it familiar with recurrently generated inputs; about 1.3 times the flops of standard training.

  4. Figure. Two panels captioned 'Figure 4: Feedback passes during prefilling improve non-generative performance.' Left, validation loss against the number of feedback passes for prefill, 0 to 3, for full-bandwidth models trained on 10B, 100B, 200B and 400B tokens, each with a dashed line for the standard model trained on the same tokens. Every solid line starts just above its dashed line at 0 passes, drops below it at 1 pass and is flat after. Right, 5-shot LM Eval average accuracy against 0 to 4 passes, with error bars, for the 100B, 200B and 400B models: each rises from 0 to 1 pass and is above its own dashed baseline from 1 pass on. A dashed black line for a standard model trained on 1T tokens lies above all of them.

  5. Figure. Four panels captioned 'Figure 5', for GSM8K (5-shot, Pass@1), Math500 (0-shot, Pass@1), HumanEval (0-shot, Pass@3) and MBPP (3-shot, Pass@3). Each plots a score under three decoding methods, Std., Soft and Fused, for full-bandwidth models trained on 100B, 200B and 400B tokens, with dashed lines for standard models at the same tokens and at 1T. On GSM8K, HumanEval and MBPP the solid lines mostly rise from Std. to Soft or Fused and lie at or above their own dashed baselines. On Math500 the 100B and 400B lines are below their baselines under all three methods, and the 400B line falls from Std. to Fused.

    Figure. A table captioned 'Table 1: Latent-feedback gains carry over through instruction tuning.' Scores in percent on four tasks, in the order GSM8K (Pass@1), MATH-500 (Pass@1), HumanEval (Pass@3), MBPP (Pass@3). Full-bandwidth at 200B under Standard, Soft and Fused decoding: 64.52, 67.93, 67.55; 43.80, 45.60, 45.60; 42.54, 45.06, 45.92; 38.39, 39.80, 41.22. Full-bandwidth at 400B: 67.90, 71.00, 71.80; 46.00, 45.40, 48.40; 46.50, 47.20, 47.60; 40.50, 40.60, 41.70. Standard transformer at 200B, 400B and 1T: 62.93, 68.39, 70.13; 42.40, 46.40, 47.40; 37.16, 44.85, 50.01; 38.61, 40.28, 41.93.