So et al. (2021; NeurIPS 2021) — ARXIV-2109.08668
Key takeaways
The two modifications (Primer-EZ).
- Squared ReLU in the feed-forward block,
max(x, 0)². It adds no parameters. It equals ReGLU when ReGLU’s two matrices are tied (§3). - Multi-DConv-Head Attention (MDHA). After each head’s Q, K and V projection, apply a causal 3×1 depthwise convolution along the sequence. Pointwise-then-depthwise works better than the usual separable order. Wider kernels and full convolutions “do not improve performance, but in several cases hurt it” (§3, Table 4).
The discovered Primer carries other changes: shared Q/K, a custom norm, 12× FFN, post-softmax gating and a pre/post norm swap. They are weaker, and some are harmful. Shared QK “generally hurts”. The 12× projection helps only near 35M parameters, and the gating does not work at variable length (§3, App. A.7).
The comparisons, and their units. “Speedup” is compute (accelerator hours) to reach the baseline’s final quality. Step-time costs are therefore already counted, and both changes do slow a step (§2).
- LM1B, 35M (App. A.6, several runs, ±): Primer-EZ is 2.34× vanilla in T2T on TPUv2, 2.03× on V100 and 1.75× in T5. Transformer++ is 1.37×, 1.54× and 1.33×. MDHA alone is 1.76× in T2T. “Vanilla” is the original Transformer of LIT-008, ReLU and layer norm outside the residual path, and the LM1B and C4/PG19 speedups are measured against reaching its final quality. Transformer++ is the strengthened baseline: RMSNorm plus the SwiGLU feed-forward of GLU Variants (LIT-030).
- C4 and PG19, 110M, T5 (Fig. 9): at least 1.8× vanilla at the end of training. Primer-EZ is 1.5× on Switch Transformer, and squared ReLU alone is 2.0× on Synthesizer. The Switch comparison adds MDHA and squared ReLU to a 550M Switch Transformer (LIT-189), in the T5 implementation Narang et al. used, and is worth 0.6 perplexity at matched compute (App. A.10).
- Full T5 replica, 537M (Table 1): 4.2× the original T5, trained with Raffel et al.’s (LIT-425) exact configuration applied to autoregressive LM, and 0.9 perplexity better at the same training cost. Against T5++ (SwiGLU + RMSNorm), Primer reaches 13.25 perplexity in 3.8K TPUv3-hours against 4.6K, about 1.2×, and 12.69 in 8.3K against 16.5K, about 2×.
- 1.9B, GPT-3-XL-like (Fig. 11, Table 6): full Primer against Transformer + GELU, not Transformer++. Primer matches it with ⅓ of the compute on pretraining perplexity and 27 one-shot tasks (5 better, 1 worse, 21 equal). The data is proprietary and the architecture has no SwiGLU arm.
Squared ReLU against the gated units (Fig. 5 right): one bar chart at 110M on C4, T5 codebase, 525K steps. Squared ReLU beats ReLU, GELU, Swish, ReGLU and SwiGLU. No variance is shown, and the paper does not say whether the gated variants were parameter-matched.
Traps
- “Gains increase with compute” has two meanings here, and the paper separates them. At compute-optimal sizes (Fig. 7), the saving is a constant factor. At a fixed model size the factor grows the longer you train, and the authors warn that this “can be ‘gamed’ by investing more compute than is necessary for baseline models” (§4.3). The 4.2× is the fixed-size kind, which is why they used Raffel et al.’s exact schedule.
- Hyperparameters are the Transformer’s, untuned for Primer. They are Adafactor, rsqrt decay and regularization disabled (App. A.8). That favours the baseline, so it is not a hidden advantage for Primer.
- Decoder-only. On T5 encoder–decoder masked LM (App. A.12, single runs), Primer-EZ “perform[s] only as well as Transformer++”.
- Squared ReLU is unbounded and grows quadratically. That is the property LIT-200 blames for SwiGLU’s activation outliers at low precision. Primer never tests narrow arithmetic.
Standing in the anthology
Filed from #290’s catalogue triage. It is the only paper in the record
that puts squared ReLU against SwiGLU. It supports no practice. It is now
one of SOTA-034’s contested_by papers, because it disputes SwiGLU’s quality
and not only its range. The evidence for that is one figure at one scale.
MDHA has no practice. It predates the short causal convolutions later
recurrent and hybrid layers put before their mixers. The record has not
filed that lineage.