Source-linked AI summary
Primer: Searching for Efficient Transformers for Language Modeling
David R. So, Wojciech Mańke, Hanxiao Liu, Zihang Dai, Noam Shazeer, Quoc V. Le
TL;DR
Transformer scaling has made training increasingly expensive, motivating a search for more efficient language-model architectures. The paper searches low-level TensorFlow primitives and identifies Primer, whose simple modifications reduce compute requirements while preserving or improving autoregressive language-model quality. The reported gains transfer across scales and established codebases, though evidence is concentrated on decoder-only models and remains below full state-of-the-art model scale.
Problem
Scaling Transformers improves performance but has made training costs prohibitively expensive, motivating more efficient Transformer language models.
Method
The paper uses evolutionary search over TensorFlow computation-graph primitives to discover efficient decoder-only autoregressive language-model architectures.
Results
Primer’s compute savings increase with training scale; it reaches T5 quality with 4.2X less compute at 500M parameters and matches Transformer one-shot performance with 3X less training compute at 1.9B parameters.
Takeaways & Limitations
Primer’s gains are mostly attributed to squaring ReLUs and adding depthwise convolutions after Q, K, and V projections, and the modifications transfer across codebases without additional changes.
Takeaways & Limitations
The study focuses primarily on decoder-only models and evaluates models much smaller than full-scale GPT-3; encoder-decoder comparisons are not studied in significant depth.
Abstract
from arXiv · showhide
Large Transformer models have been central to recent advances in natural language processing. The training and inference costs of these models, however, have grown rapidly and become prohibitively expensive. Here we aim to reduce the costs of Transformers by searching for a more efficient variant. Compared to previous approaches, our search is performed at a lower level, over the primitives that define a Transformer TensorFlow program. We identify an architecture, named Primer, that has a smaller training cost than the original Transformer and other variants for auto-regressive language modeling. Primer's improvements can be mostly attributed to two simple modifications: squaring ReLU activations and adding a depthwise convolution layer after each Q, K, and V projection in self-attention. Experiments show Primer's gains over Transformer increase as compute scale grows and follow a power law with respect to quality at optimal model sizes. We also verify empirically that Primer can be dropped into different codebases to significantly speed up training without additional tuning. For example, at a 500M parameter size, Primer improves the original T5 architecture on C4 auto-regressive language modeling, reducing the training cost by 4X. Furthermore, the reduced training cost means Primer needs much less compute to reach a target one-shot performance. For instance, in a 1.9B parameter configuration similar to GPT-3 XL, Primer uses 1/3 of the training compute to achieve the same one-shot performance as Transformer. We open source our models and several comparisons in T5 to help with reproducibility.
1 Introduction
Primer is an efficiently searched Transformer alternative for autoregressive language modeling, designed to reduce training costs while maintaining or improving quality. Its gains are attributed mainly to two simple architectural modifications and remain robust across scales, datasets, hardware, codebases, and model families.
- Motivation: The study targets prohibitively expensive Transformer training costs created by scaling, while noting that inference costs also matter.The paper frames Transformer scaling as improving performance but increasing the cost of training larger models.
- Search objective: Primer searches over TensorFlow computation-graph primitives to find lower-cost alternatives to Transformer language models.The search uses evolution over TensorFlow programs and focuses on decoder-only autoregressive language modeling.
- Core findings: Primer provides higher quality at fixed training cost, lower training cost at target quality, and lower inference cost at target quality.These benefits are reported across model sizes from 20M to 1.9B parameters and training scales from 10 to 105 accelerator hours.
- Reproducibility: The authors provide open-source models and T5 comparisons to support reproducibility.The comparisons span multiple codebases, hardware platforms, datasets, and model families.
- Quantitative results: 4.2X less compute lets 500M-parameter Primer reach the T5 baseline’s quality, while matching quality with 3X less training compute at 1.9B parameters.The 500M comparison reports a 0.9 perplexity improvement at equal training cost; the 1.9B comparison covers pretraining perplexity and downstream one-shot tasks.
- Architectural changes: Primer’s improvements are mostly attributed to squaring ReLU activations and adding depthwise convolutions after Q, K, and V projections.The modifications can be inserted into existing Transformer codebases for autoregressive language modeling.
2 Search Space and Search Method
The search space represents stackable decoder blocks as low-level TensorFlow programs encoded by evolutionary-search DNAs. Models receive a fixed compute budget and are selected by validation perplexity, with conceptual initialization addressing the prevalence of unstable random programs.
- Search space: Each search program maps input tensors X ∈ R^n×d to same-shaped outputs representing next-token prediction embeddings when blocks are stacked.The programs specify model architectures, excluding embeddings, preprocessing, and optimization.
- Program representation: A DNA stores indexed subprograms and instructions that call TensorFlow primitives or higher-index subprograms, allowing low-level operations to construct high-level blocks without cycles.SUBPROGRAM 0 is the execution entry point, while the remaining subprograms form a callable bank.
- Instruction arguments: Instruction arguments draw from a union of possible fields, including tensor-input indices, constants, and transformation dimension sizes.Examples include SIN for a tensor input, ADD for two tensor inputs, MAX for constants, and CONV 1X1 for projections.
- Compilation and execution: Programs become TensorFlow code one instruction at a time, and the final instruction’s tensor is used as the subprogram output.With deferred execution, graph nodes that do not contribute to the output are not executed.
- Evolutionary search: Evolutionary search evaluates each architecture with a fixed 24 TPUv2-hour budget and uses LM1B validation perplexity as fitness.This implicit efficiency objective rewards sample-efficient modifications even when they increase training step time.
- Search efficiency: Hurdles reduce search cost by 6.25X by allocating equal compute across bands, with 7 training hours used as a proxy for a full day.The hurdle configuration uses a 50th-percentile passing bar.
- Initialization: 78% of random programs of Transformer-equivalent length fail to train beyond five minutes because of numerical instability.The search therefore initializes its population with Transformer copies through conceptual initialization.
3 Primer
Primer’s gains come mainly from two simple modifications—squared ReLU activations and depthwise convolutions after attention projections—that transfer across codebases. Other discovered changes are less consistently useful, sometimes harmful, or excluded because they do not support the intended setting.
- Primer-EZ: Primer-EZ combines squared feed-forward ReLUs with depthwise convolution in attention projections as an easy-to-adopt variant.These two modifications produced the most robust improvements across T5 and Tensor2Tensor.
- Squared ReLU: Squared ReLU replaces ReLU in the Transformer feed-forward block and is identified as the most effective modification.The paper compares its asymptotics with common nonlinearities and reports better quality than evaluated GLU variants.
- Multi-DConv-Head Attention: MDHA adds 3x1 depthwise convolutions after the query, key, and value projections in self-attention.The convolutions operate over the spatial dimension of each dense projection’s output.
- Other Modifications: Several additional modifications are not retained: shared Q/K representations generally hurt performance, while pre/post normalization helps only in some cases.Primer also uses custom normalization, but the paper reports that it is not always effective.
- Other Modifications: The 12X bottleneck projection helps at smaller model sizes but is less effective for larger models, so it is excluded from Primer and Primer-EZ.It uses d_model = 384 and d_ff = 4608 instead of the baseline’s 512 and 2048.
- Other Modifications: Post-softmax spatial gating improves perplexity for fixed-length sequences but cannot be applied to variable sequence lengths, so it is excluded from experiments.The discovered model also contains extraneous modifications that produce no meaningful quality difference and are omitted from the main discussion.
4 Results
Primer outperforms Transformer baselines across search, scaling, transfer, inference, and large-scale T5 experiments. Its compute savings increase with scale and transfer across codebases, datasets, hardware, model families, and downstream tasks.
- 4.1 Search Task Comparison: 1.7X or more speedup: Primer reaches the vanilla Transformer’s final LM1B quality faster across codebases and hardware.The search used approximately 35M-parameter models and compared Tensor2Tensor and T5 on TPUv2s and V100 GPUs.
- 4.1 Search Task Comparison: Primer’s compute savings over Transformer follow a power law with quality when models are optimally sized.The compared architecture lines are roughly parallel in log space, indicating quality-dependent compute savings.
- 4.2 Transferability: Primer improves inference at a majority of target LM1B qualities, with train-step and inference timing showing 0.97 correlation.Forward-pass timing is used as an inference proxy because decoding methods have different compute costs.
- 4.2 Transferability: 1.8X or more efficiency: Primer and Primer-EZ transfer to PG19 and C4 in T5, with savings increasing as more compute is invested.The modifications also transfer to Switch Transformer and Synthesizer; Primer-EZ reduces Switch Transformer compute by 1.5X.
- 4.3 Large Scale T5 Training: 4.2X more compute efficient: Primer reaches target C4 language-model quality than original T5 at 537M parameters.Primer is also 2X more efficient than the strengthened Transformer++ baseline.
- 4.4 Downstream One-Shot Tasks: 3X less compute: Primer matches Transformer+GELU pretraining perplexity and one-shot downstream performance.With equal compute, Primer outperforms Transformer+GELU on 15 of 27 tasks, underperforms on 2, and ties on 10.
5 Conclusion
The paper recommends Primer and Primer-EZ for auto-regressive language modeling because their gains are simple, robust, and transferable. It also identifies smaller-than-state-of-the-art scale and limited encoder-based evaluation as study boundaries.
- Limitations: The study evaluates models substantially smaller than full-scale GPT-3 and focuses primarily on decoder-only architectures.Encoder-decoder masked-language-modeling comparisons are limited in depth, and Primer only matches Transformer++ there.
- Recommendations and Future Directions: Primer and Primer-EZ are recommended for auto-regressive language modeling because of their performance, simplicity, and robustness to codebase and hyperparameter changes.The authors report gains after dropping the modifications into established codebases without changes.
- Recommendations and Future Directions: Small activation changes are identified as a promising direction for further efficient-Transformer research.The paper specifically highlights more efficient training from simple activation-function modifications.
A.1 TensorFlow Primitives Vocabulary
The TensorFlow primitives vocabulary maps named search operations to TensorFlow functions and specifies how instruction arguments are converted into function arguments.
- TensorFlow Primitives Vocabulary: Each vocabulary entry gives an operation name, its TensorFlow function, and the mapping from DNA arguments to TensorFlow arguments.The vocabulary is constructed largely from low-level TensorFlow operations needed to create Transformers.
- TensorFlow Primitives Vocabulary: The vocabulary defines the low-level operations used to construct Transformer computation graphs.These primitives form the operation set for the search space.
A.2 Constructing TensorFlow Graphs
TensorFlow graphs are constructed from DNA programs whose instructions control dimensions, shared values, causal masking, branching, and dimension compatibility.
- Relative Dimensions: Relative dimensions allow searched models to be resized within parameter limits using a fixed vocabulary of dimension multipliers.The vocabulary ranges from 1 to 64, and its values were not tuned.
- Values Bank: Shared banks let multiple instructions reuse and change constant or dimension-size values simultaneously.The constants bank contains 2 values and the dimension-size bank contains 6 values.
- Causal Masking: Causal masking prevents each position from accessing the token it must predict or information from future positions.Operations that move information spatially receive masking overhead, including convolutions.
- Branching: Branching executes an instruction or subprogram in parallel and concatenates the resulting tensors along their embedding axes.This enables multi-head capabilities and supports branching values from 1 through 16.
- Resolving Dimension Mismatches: Dimension mismatches are resolved deterministically by pseudorandomly setting one tensor dimension to match the other.This permits programs with incompatible binary-operation dimensions to be made valid.
A.3 Halving Hurdles
Halving hurdles allocates evolutionary-search training across progressively filtered compute budgets, using fitness-based survival at each hurdle. Correlated performance across budgets reduces total search cost.
- Hurdle configuration: Halving hurdles retain the top 50% of individuals by fitness at each hurdle.The search uses hurdles at 812.9s, 2438.7s, 5690.3s, and 12,193.5s.
- Hurdle configuration: The four hurdles distribute roughly one-fifth of the compute budget to each successive training band.The bands cover progressively smaller surviving populations trained for longer durations.
- Hurdle configuration: Four hurdles are used because a fifth would place the first hurdle below ten minutes, which was empirically too noisy.The authors state that this configuration requires only the number of hurdles as a hyperparameter.
- Search-cost reduction: 4064s average training time per model reduces search compute by approximately 6.2X.The figure attributes the reduction to correlated model performance across different training budgets.
- Search-cost reduction: Unlike Successive Halving and Hyperband, the method integrates halving with a changing evolutionary population rather than a static population.The comparison concerns how candidates are maintained during search.
A.4 Evolution Search Details
Evolutionary search generates new TensorFlow-program candidates through randomized mutations, while rejecting mutations that leave the compute graph unchanged.
- Search procedure: Regularized Evolution uses a population of 100 and tournament selection of size 10, with these values not tuned.The search applies mutations to create new candidates from the evolving population.
- Mutation operators: Each candidate is created by uniformly selecting a parent and applying one mutation.Five mutation types are used: delete, insert, delete-and-insert, mutate-field, and swap.
- Mutation operators: Swap mutations exchange instruction positions and also swap their input tensors, preserving the net instruction switch in the compute graph.The operator changes ordering within a randomly selected subprogram.
- Mutation operators: A light equivalence check triggers another mutation when the candidate’s compute graph is exactly equivalent to its parent’s.This prevents unchanged graph structures from remaining as newly generated candidates.
A.5 Transformer and Primer Program Comparisons
The paper compares Transformer and Primer by resolving their subprograms into constituent primitives and highlighting Primer’s architectural differences. The comparison spans main, attention, feed-forward, projection, softmax, normalization, residual, and activation components.
- Subprogram comparisons: Figures 13–22 depict model subprograms with Primer changes highlighted in orange.The figures separately cover main, attention, feed-forward, multi-head projection, softmax, normalization, z-score, scale-shift, residual, and activation subprograms.
- Primitive-level programs: Table 3 maps operation names to graph symbols used in the subsequent program graphs.This key supports reading the primitive-resolved comparisons.
- Subprogram comparisons: The normalization, residual, and scale-shift figures identify downstream realization, a functional no-op, and no change, respectively.Figure 18 attributes normalization changes to downstream changes in S6; Figure 20 reports no scale-shift changes; Figure 21 calls the residual change essentially a functional no-op.
- Subprogram comparisons: The activation-function figure marks architectural changes in orange, alongside the other subprogram-level comparisons.The supplied caption identifies the highlighted changes but does not specify their individual operations.
- Primitive-level programs: Figures 23–25 resolve the full Transformer and Primer programs into constituent primitives and list their corresponding instructions.Figure 23 compares the full programs, while Figures 24 and 25 list Transformer and Primer instructions.
A.6 Exact LM1B Numbers
Table 4 reports exact LM1B search-task comparisons across hardware platforms and libraries using their default hyperparameters.
- Comparison scope: Table 4 compares auto-regressive LM1B results across TPUv2s and V100 GPUs, and Tensor2Tensor and T5.The table contains the precise numbers underlying Figure 6, but the supplied passage does not include those values.
- Metric definition: Speedup is defined as the fraction of compute needed to match the vanilla Transformer baseline trained with the full compute budget.The supplied passage identifies the metric definition without reporting its numerical entries.
A.7 Ablation and Insertion Studies
Insertion and ablation studies evaluate which Primer modifications transfer across model settings and codebases. The studies identify squared ReLUs and MDHA-related attention changes as especially important, while broader gains and compute savings vary by configuration.
- Study design: Insertion studies add each Primer modification to a vanilla Transformer, while ablation studies remove each modification from Primer one at a time.Both studies use normalized perplexity deltas to compare whether modifications help or hurt performance.
- Modification effects: Three techniques are beneficial across all tested scenarios: 12X proj, MDHA, and squared ReLUs.12X proj helps smaller models but is not useful at larger sizes; MDHA and squared ReLUs define the simpler Primer-EZ model.
- Transfer to T5: Primer improves all transferred T5 baselines on larger language-modeling datasets, while Primer-EZ reaches parity with Primer and sometimes surpasses it.Switch Transformer and Synthesizer also benefit from the Primer-EZ modifications.
- Downstream evaluation: At 1.9B parameters, Primer achieves roughly Transformer-level one-shot performance with 1/3 of the pretraining compute and stronger performance on most tasks at full compute.The comparison uses five pretraining checkpoints and reports 95% confidence intervals.
- Transferability: Primer modifications transfer across codebases without additional tuning, but their compute savings depend on hardware, library operation speeds, and task-specific sample efficiency.The T5 replica demonstrates 4.2X savings, while the authors expect results to vary across configurations.
- Scope: The paper focuses on autoregressive language modeling; masked language modeling and encoder-decoder results are limited, and Primer is not obviously better than Transformer++ there.The authors caution that modifications effective for autoregressive modeling may be less effective for masked language modeling.