Brainformers: Trading Simplicity for Efficiency

ArXiv: 2306.00008

🎯 Pitch

By breaking free from uniform transformer blocks, Brainformer achieves a startling 5× faster per-step training time and 2× faster convergence versus GLaM at 8B scale. It then delivers 3% higher SuperGLUE accuracy without any extra activated parameters, simply by evolving a heterogeneous mix of sparse feed-forward, dense feed-forward, and attention layers.

---

1. Executive Summary

This paper introduces the Brainformer, a complex, non-uniform transformer block architecture discovered via regularized evolutionary search that replaces the standard interleaved attention and feed-forward layers with a diverse sequence of sparsely gated feed-forward layers, dense feed-forward layers, and attention layers, coupled with tunable gating mechanisms (expert choice vs. top-2 token routing). Evaluated on a 1.6-trillion-token corpus from GLaM and downstream benchmarks including SuperGLUE, the Brainformer at 8B activated parameters demonstrates 2× faster training convergence and 5× faster step time compared to its GLaM counterpart, while achieving a 3% higher SuperGLUE score with fine-tuning and outperforming the NAS-derived Primer dense model on fewshot generative tasks—establishing that trading architectural regularity for block-level heterogeneity substantially improves both training efficiency and downstream quality only when the search jointly optimizes layer composition, gating strategy, and model dimensions under a fixed wall-clock training budget.

2. Context and Motivation

The Core Problem: The Uniform Transformer Block Is a Bottleneck for Efficiency

The fundamental question this paper tackles is architectural: does the standard, uniform stacking of transformer layers — where attention and feed-forward blocks strictly alternate — represent the most compute-efficient way to build large language models? The paper's hypothesis is that the answer is no, and that deliberately introducing non-uniformity into the block structure — varying the type, width, and ordering of layers within a repeated block — can yield substantial gains in training speed, inference throughput, and model quality.

This matters because the dominant paradigm for building scaled-up transformers, from BERT to GPT-3 to GLaM, has been to design a single, regular block (e.g., attention + dense FFN, or attention + MoE + dense FFN in GLaM's interleaved pattern) and then stack it $L$ times to increase depth. This uniformity simplifies implementation and reasoning about the model, but it imposes a rigid constraint: every token sees exactly one attention layer and one feed-forward layer in every block, regardless of whether that allocation is optimal for the task, the training stage, or the hardware. The Brainformer paper questions whether this convenience cost is worth paying.

Why This Problem Is Important

Training cost dominates the economics of large models. Modern LLMs can cost millions of dollars in compute to train. If architectural heterogeneity can deliver a 2× improvement in training convergence (as the paper claims for Brainformer vs. GLaM at the 8B scale) and a 5× improvement in per-step throughput, those compound into enormous dollar and carbon savings. For an organization deciding whether to train a 100B+ parameter model, a more efficient block design directly translates to fewer TPU/GPU hours, lower energy consumption, and faster iteration cycles.

Inference throughput shapes deployability. The paper reports Brainformer achieving 1.96 steps/sec at the 8B64E scale compared to GLaM's 0.39 steps/sec — a 5× difference. For production systems serving real-time queries, this is the difference between a model that can run on a manageable cluster and one that requires prohibitive hardware. Architectural choices that improve inference efficiency without degrading quality have outsized practical impact because inference costs often dominate the total cost of ownership once a model is deployed.

The uniform-block assumption is largely untested. Despite the transformer's dominance, the field has invested relatively little effort in systematically exploring whether different layer types should appear at different frequencies, in different orders, or with different widths. Most work on efficient transformers has focused on replacing components (e.g., replacing attention with linear approximations in Linformer or Performer) rather than rearranging them. The Brainformer paper argues that the arrangement itself — the "architecture of the architecture" — is a rich optimization surface that prior work largely ignored.

Sparse models amplify the need for good architecture. In dense transformers, the block structure is simple and the cost is distributed uniformly. But sparse Mixture-of-Experts (MoE) models introduce a new axis of complexity: the gating mechanism, the number of experts, the capacity factor, and where MoE layers are placed relative to dense layers. A bad architectural choice in an MoE model can produce load imbalance (some experts overused, others idle), communication bottlenecks, or experts that never train on enough data. The Brainformer paper's approach of searching over block composition, gating type, and capacity jointly is motivated by the observation that these choices interact — the optimal gating function depends on the surrounding layer types and the block's position in the network.

Where Existing Approaches Fall Short

The paper identifies several distinct limitations in prior work, each of which motivates a different aspect of the Brainformer design.

1. Dense transformer efficiency work focuses on attention, but feed-forward layers dominate cost at typical sequence lengths.

A large body of work targets attention-layer efficiency — Linformer introduces low-rank approximations, Performer uses kernel-based methods, Synthesizer removes attention entirely. However, the paper notes (Section 1) that "recent work has also identified that dense feed-forward layers constitute most of the computational cost for common sequence lengths ($\leq 2048$), particularly when the model is large." This is a critical observation: for models like GLaM or GPT-3 operating on typical context windows, the feed-forward network (FFN) layers consume more FLOPs than the attention layers because the hidden dimension $d_{\mathrm{ffn}}$ is typically $4\times$ the model dimension $d$. Optimizing attention helps at very long sequences, but at short-to-medium sequences, FFN cost dominates. This means the low-hanging fruit for efficiency at current deployment scales lies in making FFN layers cheaper, not just attention layers.

The paper uses this observation to motivate introducing sparsity in the FFN layers via MoE — the sparsely gated feed-forward layer replaces one large, dense matrix multiplication with multiple smaller expert FFNs, activated conditionally. Unlike prior MoE work that still placed MoE layers in a uniform interleaved pattern, Brainformer explores where and how many MoE layers to place.

2. MoE architectures (GLaM, Switch Transformer) use uniform, manually-crafted blocks.

GLaM (Du et al., 2022) interleaves a dense transformer block with a sparse transformer block: every other layer is an MoE layer, and the gating function is fixed to top-2 token routing. Switch Transformer uses top-1 routing with a uniform block of alternating attention and MoE layers. These are reasonable hand-designed patterns, but they expose no degrees of freedom: the ratio of attention-to-dense-to-MoE layers is fixed, the layer widths are uniform, and the gating mechanism is baked in.

The paper argues that this uniformity is suboptimal because different positions in the network benefit from different layer types. Early layers might benefit from broader spatial mixing (attention), middle layers from specialized knowledge stored in different experts (MoE), and later layers from dense nonlinear transformations (dense FFN). By searching over the composition, Brainformer discovers patterns that deviate substantially from the uniform interleaving baseline — for example, the searched Brainformer block 1 contains 8 sub-layers (a mix of attention, MoE, and dense FFN) in a specific, non-alternating order that is repeated as a unit.

3. Prior architecture search work (Primer) searched within the uniform-block paradigm.

Primer (So et al., 2021) applies neural architecture search (NAS) to find efficient transformer variants, but it operates within the assumption of a uniform, repeating block structure — it searches over activation functions, normalization placements, and attention mechanisms, but not over which layer types appear at each position or in what ratio. The paper explicitly positions Brainformer against Primer in the fewshot evaluation (Table 5), where Brainformer outperforms Primer on all tasks except one while being faster in steps/sec (1.37 vs. 1.50 for Primer). The implication is that the additional degrees of freedom — block-level heterogeneity and gating mechanism choice — matter more than the intra-block optimizations Primer explored.

4. Routing mechanisms (token-based vs. expert-based) were studied in isolation from architecture.

Token-based routing (used in GShard, GLaM, Switch Transformer) and expert-based routing (Expert Choice, from Zhou et al., 2022) were developed and evaluated in separate contexts, on fixed architectures. The Brainformer paper argues that the optimal routing mechanism likely depends on the architecture — a block with many narrow experts might benefit more from expert-choice routing (which guarantees load balance), while a block with few wide experts might work better with top-2 token routing. By including both gating functions in the search space (Table 1), the paper aims to discover which routing mechanism pairs best with which block composition, rather than assuming one is universally better.

5. Fair comparison between model families is poorly defined.

The paper devotes Section 3.3 entirely to the problem of fair comparison. Prior scaling studies typically fix one dimension of the comparison — e.g., compare at equal parameter count, or equal training tokens, or equal FLOPs. Each choice favors different model families:

• Equal parameters: Unfair to sparse models, which have many more total parameters but the same computational cost per token.
• Equal training tokens: Unfair to models that train faster — a model that can process more tokens in the same wall-clock time is penalized if tokens are capped.
• Equal FLOPs: Better, but still ignores architectural differences in step time — two models with the same FLOPs can have different throughput due to memory access patterns, communication overhead, or hardware utilization.

The paper proposes fixed wall-clock training time plus constrained inference step time as the fairest comparison framework. This means the search algorithm is rewarded for finding architectures that (a) train quickly (so they can process more training data within the fixed time budget) and (b) are fast at inference (so they are useful in deployment). This is a more holistic fairness criterion than prior work used.

6. The interaction between layer width, gating, and architecture was unexplored.

GLaM fixes the FFN hidden dimension at $4\times$ the model dimension, following the standard transformer recipe. But in an MoE model, where the effective FFN capacity is distributed across $E$ experts, the optimal expansion ratio might be different — each expert sees only a fraction of tokens, so it might benefit from being narrower (less overfitting, faster computation) or wider (more capacity per expert for better specialization). The Brainformer search space makes the hidden dimensions of both dense and sparse FFN layers tunable independently (Table 1: $d_{\mathrm{moe}}$ and $d_{\mathrm{ffn}}$ with values 1536–4096), allowing the search to discover configurations where MoE layers are narrower than the standard 4× expansion but compensate through more experts or different placement.

How This Paper Positions Itself

The Brainformer paper positions itself at the intersection of three research threads, aiming to unify and extend them:

From EfficientNet (vision): The idea of layer-wise compound scaling and non-uniform architectures. EfficientNet showed that scaling CNN width, depth, and resolution together — and doing so with layer-specific multipliers — outperforms uniform scaling. Brainformer imports this philosophy into language models, but replaces EfficientNet's compound coefficient scaling with an evolutionary search over block composition.

From MoE literature (GLaM, Switch Transformer, Expert Choice): The insight that conditional computation via sparsely gated experts can dramatically increase capacity without proportionally increasing FLOPs. Brainformer treats the MoE layer as a pluggable component whose placement, gating function, and capacity factor are all subject to optimization rather than fixed by hand.

From neural architecture search (Primer, EfficientNet): The methodology of automated architecture discovery via evolutionary algorithms. Brainformer's key innovation over prior NAS for transformers is the block-wise search space — instead of searching over a single block type and then repeating it, Brainformer searches over the composition of a heterogeneous block (which can contain any sequence of attention, MoE, and dense FFN sub-layers) and then stacks that complex block to create models at different scales. The block-wise approach makes scaling practical, since only the number of block repetitions needs to change, not the internal structure.

The paper's central thesis is that jointly optimizing layer composition, gating strategy, and model dimensions under a fixed wall-clock training budget produces architectures that substantially outperform both hand-designed sparse transformers and NAS-discovered dense transformers — and that the gains come specifically from trading away the simplicity of architectural uniformity for the efficiency of a heterogeneous, fitness-maximized block structure.

3. Technical Approach

3.1 Reader orientation

The Brainformer is a non-uniform transformer block architecture discovered through evolutionary search, where a single repeated block contains a heterogeneous sequence of sub-layers — sparsely gated feed-forward (MoE) layers, dense feed-forward layers, and self-attention layers — with independently tunable widths and gating mechanisms. The core problem it solves is that standard transformers enforce a rigid alternating pattern of attention and feed-forward layers, which is computationally convenient but architecturally suboptimal; the Brainformer solution trades away this regularity by allowing the search algorithm to discover block compositions that are highly asymmetric and non-alternating, achieving better quality per unit of training time and per inference step than either hand-designed sparse transformers (GLaM) or NAS-discovered dense transformers (Primer).

3.2 Big-picture architecture (diagram in words)

The Brainformer system has four major components that operate in sequence:

1. Block Search Space: A space of permissible sub-layer sequences, where each position in a block can be an attention layer, a dense feed-forward layer, or a sparsely gated MoE layer, with tunable dimensions and gating functions. This is the set of all possible block architectures the search algorithm can propose.

2. Evolutionary Search Controller: A regularized evolutionary algorithm that samples candidate block architectures from the search space, instantiates them as small proxy models (100M activated parameters), trains them under a fixed wall-clock time budget with early pruning, and selects top candidates based on validation perplexity subject to an inference step-time constraint.

3. Block Stacking and Scaling: The top-k discovered blocks are scaled up by (a) multiplying model dimensions by 2× and 4× factors and (b) increasing the number of times the block is repeated, producing model variants at 100M, 1B, and 8B activated parameter scales for final evaluation.

4. Final Evaluation Pipeline: The scaled Brainformer models are trained on the full 1.6T-token GLaM corpus and evaluated on pre-training perplexity, fine-tuning performance on SuperGLUE, and one-shot performance on five generative QA tasks, with all comparisons made against GLaM and Primer baselines under equal wall-clock training time or equal inference step-time constraints.

3.3 Roadmap for the deep dive

• First, the sub-layer type space (Equation 3): what each sub-layer is, how many types exist, and what hyperparameters each type exposes — because everything downstream is built from these primitives.

• Second, the block composition formalism (Equations 1, 2, 4): how sub-layers are composed into a block, how the block is treated as a single architectural unit, and how the search objective is formally defined — because this is the abstraction that makes search and scaling tractable.

• Third, the search space table and the evolutionary search algorithm (Algorithm 1): what exactly is being searched over (Table 1), how candidates are sampled and evaluated under a fixed wall-clock budget, and how early stopping prunes unpromising candidates — because the search procedure's choices directly determine which architectures the paper discovers and evaluates.

• Fourth, the fair comparison framework (Section 3.3): why standard equal-parameter or equal-FLOP comparisons are inadequate, and how the paper defines a fixed-training-time-plus-inference-constraint comparison that allows architectures to trade off capacity for speed — because this framework is what makes the Brainformer's claimed gains legitimate rather than artifacts of an unfair baseline.

• Fifth, the gating mechanisms (Section 4, Figure 6): how token-based routing and expert-based routing differ computationally, why both are included in the search space, and how the choice of gating interacts with the surrounding block architecture — because the gating function is a first-class architectural decision, not a fixed hyperparameter.

• Sixth, the scaling protocol (Algorithm 1, lines 13–17): how a discovered block is scaled to larger models by dimension multiplication and repeated stacking — because this is what distinguishes the Brainformer approach from one-shot architecture search and makes the block reuse practical.

3.4 Detailed, sentence-based technical breakdown

This is primarily an architecture design and search paper whose core idea is that transformer blocks should not be architecturally uniform, and that evolutionary search under a fixed wall-clock training budget can discover heterogeneous block compositions that outperform both hand-designed sparse transformers and NAS-discovered dense transformers.

---

Sub-Layer Types and Hyperparameters (Equation 3)

The fundamental building block of the Brainformer architecture is a generic sub-layer function $\mathcal{F}_i$ that transforms an input tensor $X_i$ into an output tensor $Y_i$, with the sub-layer type drawn from a set of three possibilities:

$\mathcal{F}_i = \begin{cases} \mathcal{F}_i^{d, h, a}, \text{if } \mathcal{F}_i = \mathcal{F}_{attn} \\ \mathcal{F}_i^{d, d_{ffn}, a}, \text{else if } \mathcal{F}_i = \mathcal{F}_{ffn} \\ \mathcal{F}_i^{d, d_{moe}, g, c, a}, \text{otherwise } \mathcal{F}_i = \mathcal{F}_{moe} \end{cases}$

where $\mathcal{F}_i$ is the sub-layer at position $i$ in the block, $d$ is the model dimension (the width of the residual stream), $h$ is the number of attention heads, $a$ is the activation function, $d_{ffn}$ is the hidden dimension of a dense feed-forward layer, $d_{moe}$ is the hidden dimension of each expert within the MoE layer, $g$ is the gating function type, and $c$ is the capacity factor controlling how many tokens each expert can process.

What Equation 3 computes: it defines a conditional specification of sub-layer hyperparameters based on the sub-layer type. If the sub-layer is attention ($\mathcal{F}_{attn}$), it is parameterized by model dimension $d$, number of heads $h$, and activation $a$. If it is a dense FFN ($\mathcal{F}_{ffn}$), it is parameterized by model dimension $d$, hidden dimension $d_{ffn}$, and activation $a$. If it is a sparsely gated MoE layer ($\mathcal{F}_{moe}$), it is parameterized by model dimension $d$, per-expert hidden dimension $d_{moe}$, gating function $g$, capacity factor $c$, and activation $a$.

Why this form: separating the parameterization by sub-layer type allows the search to independently tune the width of dense FFN layers versus MoE layers versus the attention layer dimensions. In a standard transformer, $d_{ffn}$ is typically $4 \times d$, and the same hidden dimension applies uniformly. Brainformer allows the MoE hidden dimension $d_{moe}$ to differ from the dense FFN hidden dimension $d_{ffn}$, and both to differ from the model dimension $d$. This is grounded in the observation from Section 3.1 and Figure 3 that MoE layers effectively factorize a large matrix multiplication into multiple smaller ones — if the factorization is good, the individual experts can be narrower than the equivalent dense matrix and still achieve the same representational capacity. Without independent dimension parameters, the search could not discover configurations where, for example, MoE layers use a narrow per-expert hidden dimension to maximize throughput while dense FFN layers use a wide hidden dimension to compensate for the absence of sparsity.

The input tensor $X_i$ has shape $\{B, L, H\}$ where $B$ is the batch size, $L$ is the sequence length, and $H$ is the model dimension at that sub-layer. Crucially, $H$ is allowed to vary across sub-layers within a single block, taking values from $\{\frac{3}{4}, 1, \frac{3}{2}\} \times H_{\text{model\_dim}}$, where $H_{\text{model\_dim}}$ is a global model dimension baseline. This means a sub-layer can be narrower or wider than its neighbors, creating a network topology with variable-width sub-layers — analogous to the funnel-shaped or bottleneck-shaped architectures common in computer vision but rare in transformer language models. The motivation, stated in Section 3.2, is "to enable more flexible network topologies with various factorization methods," such as adding wide experts or narrow experts at different positions in the block.

---

Block Composition Formalism (Equations 1, 2, and 4)

A Brainformer block $\mathcal{N}$ is defined as the sequential composition of $k$ sub-layer functions:

$\mathcal{N} = \mathcal{F}_k \odot ... \odot \mathcal{F}_2 \odot \mathcal{F}_1(X_1) = \bigodot_{j=1...k} \mathcal{F}_j(X_1)$

where $\mathcal{F}_j$ is the sub-layer at position $j$ drawn from $\{\mathcal{F}_{attn}, \mathcal{F}_{moe}, \mathcal{F}_{ffn}\}$, $\odot$ denotes functional composition (the output of one sub-layer becomes the input to the next), and $X_1$ is the input to the first sub-layer of the block. The entire block $\mathcal{N}$ is treated as a single architectural unit — the search discovers the sequence $\mathcal{F}_1, \mathcal{F}_2, ..., \mathcal{F}_k$ that defines what the block contains, and then this block is stacked $N$ times to form a full model.

What Equation 1 computes: it defines a block as an ordered list of $k$ sub-layers applied sequentially to the input. The block is a function that accepts a tensor $X_1$ and produces a transformed tensor $\mathcal{N}(X_1)$. The composition operator $\bigodot$ means each sub-layer receives the output of the previous one, exactly as in a standard residual network but without the constraint that the sub-layer types alternate in a fixed pattern.

Why this form: treating the block as a single atomic unit of search and scaling is the paper's key architectural abstraction. Instead of searching over an entire model's layer-by-layer configuration (which would be combinatorially explosive — with three sub-layer types and 64 layers, the search space would be $3^{64}$), the search only needs to discover the composition of one block (with $k$ sub-layers, where the discovered blocks have $k=8$), and then the model is constructed by stacking that block. This is analogous to the convolutional block designs in ResNet or EfficientNet, where a bottleneck block is repeated many times with varying channel dimensions. In the Brainformer context, the block is repeated 3 times for 100M-scale models, 6 times for 1B-scale models, and 8 times for 8B-scale models, with the model dimension scaled multiplicatively at each scale.

The search objective is formally stated as:

$\min_{\mathcal{F}_{1:k}, d, d_{moe}, d_{ffn}, h, g, c, a} \mathcal{L}(\mathcal{N}(\mathcal{F}_{1:k}, d, d_{moe}, d_{ffn}, h, g, c, a))$

where $\mathcal{L}$ is the pre-training validation cross-entropy loss, and the minimization is over all structural and dimensional hyperparameters of the block.

What Equation 2 computes: it frames architecture discovery as a constrained optimization problem — find the block composition and all associated hyperparameters that minimize validation perplexity, subject to the inference step-time constraint given in Equation 5.

Why this form is important: unlike standard neural architecture search which maximizes a quality metric in isolation, this objective explicitly recognizes that the optimization runs under a fixed wall-clock training budget (Section 3.4). The cross-entropy minimization is coupled with an implicit constraint: architectures that train too slowly will see fewer training steps within the fixed time budget and will achieve worse loss even if their per-step capacity is higher. The search algorithm rewards architectures that (a) converge quickly per training step and (b) have fast per-step throughput, since both contribute to lower loss at the end of the fixed time window. This coupling between architecture quality and architecture speed is what distinguishes the Brainformer search from prior NAS work on transformers.

The block architecture constraint is repeated in the optimization formulation:

$\text{s.t. } \mathcal{N}(\mathcal{F}_{1:k}, d, d_{moe}, d_{ff}, h, g, c, a) = \bigodot_{i=1...k} \mathcal{F}_i(X_1)$

and the inference step-time constraint is:

$\text{Step\_Time}(\mathcal{N}) \leq \text{baseline\_step\_time}$

which ensures that any discovered architecture is not slower at inference than the GLaM baseline it aims to replace.

---

The Search Space (Table 1)

Table 1 enumerates the full search space from which the evolutionary algorithm samples candidate block architectures. The baseline against which the search operates is a 100M-parameter, 12-layer dense transformer with model dimension $H_{\text{model\_dim}} = 768$. The search items and their ranges are:

• Layer Type ($\mathcal{F}_i$): $\{\mathcal{F}_{attn}, \mathcal{F}_{moe}, \mathcal{F}_{ffn}\}$ — the three permissible sub-layer types. A block can contain any sequence of these, including contiguous repetitions (e.g., two MoE layers in a row without intervening attention).

• Model Dimension ($d$): $\{512, 768, 1024\}$ — the width of the residual stream at each sub-layer. The base model uses 768; the search can narrow or widen individual sub-layers within this range. Allowing $d=512$ enables the discovery of narrow layers that save computation, while $d=1024$ enables the discovery of wide layers that add capacity at critical positions.

• MoE Hidden Dimension ($d_{moe}$): $\{1536, 2048, 3072, 4096\}$ — the hidden dimension of each expert within an MoE layer. The standard transformer expansion ratio is $4\times$, which would be 3072 given a model dimension of 768. The search includes values both below (1536, 2048) and at (3072) and above (4096) this default, enabling the discovery of narrower experts that reduce per-token computation.

• FFN Hidden Dimension ($d_{ffn}$): $\{1536, 2048, 3072, 4096\}$ — the hidden dimension of a dense FFN layer, with the same set of values as the MoE hidden dimension but independently tunable. This means a block can have a dense FFN with one hidden dimension and an MoE layer with a different hidden dimension, unlike GLaM where all FFN layers share the same expansion factor.

• Attention Heads ($h$): $\{12, 16, 20\}$ — the number of attention heads in an attention sub-layer. More heads enable finer-grained attention patterns but increase computation; the search can allocate more heads to some positions and fewer to others.

• Gating Function ($g$): $\{\text{Top-2, Expert Choice}\}$ — the two routing mechanisms explored. Top-2 is token-based routing where each token selects the two highest-scoring experts (following GShard and GLaM). Expert Choice is expert-based routing where each expert selects the top-$k$ tokens (following Zhou et al., 2022), guaranteeing perfect load balance.

• Capacity Factor ($c$): $\{1, 2, 3, 4\}$ — controls how many tokens each expert can process, expressed as a multiplier on the evenly-divided load. A capacity factor of 1 means each expert processes exactly $\text{tokens}/\text{experts}$ tokens; a factor of 2 means it can process twice that many, providing a buffer for uneven routing at the cost of additional computation.

• Activation Function ($a$): $\{\text{Gated GeLU, ReLU, GeLU}\}$ — the activation used within the feed-forward sub-layers. Gated GeLU is a variant that applies a gating mechanism (a learned gate that multiplies the activation output), which the paper includes as an option alongside the standard ReLU and GeLU.

The total combinatorial size of the search space is enormous: for a block with $k=8$ sub-layers, each position can be one of 3 layer types, and each type has multiple independently tunable hyperparameters. The evolutionary algorithm's job is to efficiently explore this space and identify architectures that perform well within the proxy training budget.

---

Evolutionary Search Algorithm (Algorithm 1)

The search procedure operates in two phases — a discovery phase where block architectures are sampled and evaluated at small scale under a proxy training task, and a scaling phase where the top candidates are scaled to target sizes for final evaluation.

Phase 1: Proxy Training with Early Stopping (Lines 1–11)

The search maintains a population of candidate block architectures $\mathcal{B}$. In each iteration $t$ (up to $T_0$ total iterations):

1. Sample candidates: $p$ block architectures are sampled from the population using the evolutionary algorithm's selection mechanism. The paper does not specify the exact selection mechanism (tournament selection, fitness-proportional, etc.), but the high-level reward function $\mathcal{R}^{(i)}$ (line 8) drives selection toward architectures with low perplexity and fast step time.

2. Instantiate proxy models: each sampled block $\mathcal{B}^{(i)}$ is stacked three times to create a small proxy model $\mathcal{G}^{(i)}$ (line 3: $\mathcal{G}^{(i)} \leftarrow \text{StackThreeTimes}(\mathcal{B}^{(i)})$). This produces a model with approximately 100M activated parameters and 32 experts, matching the smallest GLaM baseline configuration. Using a proxy task at this small scale makes the search computationally feasible — training a full-scale 8B model for every search trial would be prohibitively expensive.

3. Early stopping check (lines 4–5): at 25% of the maximum training steps, the proxy model is evaluated against two criteria: the inference time constraint ($\text{Step\_Time}(\mathcal{G}^{(i)}) \leq \text{baseline\_step\_time}$) and a perplexity threshold relative to the GLaM baseline. If the model violates either constraint, it receives a reward of $-1$ and its training is terminated early — it will not consume further compute. This early pruning is critical for efficiency: the search can explore many architectures by quickly discarding those that are too slow or too poor in quality, rather than training every candidate to completion.

4. Full training and reward computation (lines 7–8): models that pass early stopping are trained for the full $T_{max}$ steps. Their final accuracy $\mathcal{A}^i$ (validation perplexity) and training step time $\mathcal{T}^i$ are recorded, and a reward $\mathcal{R}^{(i)}$ is computed as a function of both: $\mathcal{R}^{(i)} \leftarrow f(\mathcal{A}^i, \mathcal{T}^i)$. The paper does not specify the exact form of $f$, but it is described as a function that trades off accuracy and speed — architectures with better perplexity and faster step time receive higher rewards. Since the training is under a fixed wall-clock time budget, a faster architecture can process more training steps within the same time window, potentially achieving better perplexity despite lower per-step capacity. This creates an implicit pressure toward architectures that are throughput-efficient.

Why this reward structure: fixing wall-clock time rather than number of training steps is the key fairness mechanism (Section 3.4). If the search instead fixed training steps, it would favor architectures with higher per-step capacity but slower throughput, which might not be desirable in production. By fixing time, the search naturally discovers architectures that make optimal use of hardware — they can be "compensated with more training steps" (Section 3.4) if they train quickly. The paper states: "We empirically find that fixing training wall clock time while meeting an inference time constraint yields models with faster training convergence and higher quality."

Phase 2: Scaling Top Candidates (Lines 12–17)

After the evolutionary search completes (after $T_0$ iterations), the top-$k$ block architectures with the highest rewards are selected (line 12: $\mathcal{G}_{topk} \leftarrow \text{TopK}(\{\mathcal{G}^{(i)}, \mathcal{R}^{(i)}\})$). Each of these top candidate blocks is then scaled to target model sizes:

1. Dimension scaling (line 14): the model dimension is multiplied by 2× and 4×, following the scaling factors used in GLaM, to create blocks targeting 1B and 8B activated parameter scales.

2. Depth scaling (line 15): the scaled block is stacked $N$ times, where $N$ is determined mathematically to reach the target total activated parameters. For the 1B64E variant, the block is stacked 6 times; for the 8B64E variant, it is stacked 8 times.

3. Full-scale training and evaluation (lines 16): each scaled model is trained on the full 1.6T-token corpus and evaluated for final perplexity, downstream task performance, and step time.

This two-phase approach — search at small scale, scale to large scale — is computationally essential. Searching directly at the 8B scale would cost orders of magnitude more and would be infeasible within reasonable compute budgets. The assumption is that a block architecture that performs well at the 100M proxy scale will transfer its advantages to larger scales, which the paper validates empirically in Section 5.

---

The Fair Comparison Framework (Section 3.3)

The paper argues that standard model comparison methodologies — fixing total parameter count, or fixing training tokens, or even fixing total FLOPs — are inadequate when comparing across model families with different architectures. Instead, it proposes comparing under fixed wall-clock training time plus constrained inference step time.

Why fixing total parameters is unfair: sparsely activated models have many more total parameters than dense models with the same computational cost. A 1B/64E GLaM model has 27B total parameters but only 1.88B activated parameters per token; comparing it to a 1.7B dense model on total parameters would penalize the sparse model heavily despite similar per-token FLOPs. The GLaM paper addressed this by comparing on activated parameters (which approximates FLOPs per token), and Brainformer follows this convention — all comparisons in Table 2 group models by activated parameters, not total parameters.

Why fixing training tokens is also unfair: if models are compared after processing the same number of training tokens, a faster model that could have processed more tokens in the same wall-clock time is unfairly constrained. The paper states this explicitly: "comparing models with a fixed amount of training tokens may still also not be fair as some smaller models can benefit more from additional training data and outperform a bigger model with the same total training cost." This is the Chinchilla insight applied to architecture comparison — the optimal model at a given compute budget balances capacity against data quantity.

Why fixing FLOPs is insufficient: two models with identical FLOPs can have different throughput (steps per second) due to memory access patterns, communication overhead, or hardware utilization efficiency. The paper's inference time constraint ($\text{Step\_Time}(\mathcal{N}) \leq \text{baseline\_step\_time}$) directly addresses this: an architecture that uses the same FLOPs more efficiently (lower latency per step) is rewarded in the search.

The fixed wall-clock time approach: the search objective (Equation 2) minimizes perplexity subject to a step-time constraint, with all training runs within the search having the same wall-clock duration. This means a model that trains faster can see more data and achieve lower loss — the search naturally discovers this Pareto-optimal tradeoff curve between per-step quality and per-step throughput. The paper notes that this is "the first to suggest compute-efficient scaling, which varies both model capacity and training tokens at a fixed computational cost" in the context of architecture comparison, not just pretraining scaling laws.

---

Gating Mechanisms: Token-Based vs. Expert-Based Routing (Section 4, Figure 6)

The paper explores two fundamentally different routing strategies for the sparsely gated MoE layers, treating the gating function as a tunable architectural choice rather than a fixed design decision.

Token-based routing (Top-2): each token computes an affinity score with every expert via a learned gating matrix $W_g$, and is routed to the top-2 highest-scoring experts. This is the approach used in GShard and GLaM. The paper follows the top-2 variant (rather than top-1 as in Switch Transformer) because "top-2 has demonstrated stronger empirical performance than top-1 gating."

The computation is: given token embedding $x$, compute un-normalized scores $s = W_g x$ where $W_g$ is a learned matrix mapping from the model dimension to the number of experts. Normalize along the expert dimension (to avoid causal leakage in autoregressive decoding — normalization along the token dimension would require knowing future tokens' scores). Select the two experts with the highest normalized scores for each token.

Expert-based routing (Expert Choice): each expert selects the top-$k$ tokens based on the same token-expert affinity scores, but the selection is performed from the expert's perspective. The paper follows Zhou et al. (2022)'s Expert Choice gating. The key property is that "perfect load balance is achieved" by construction — each expert processes exactly $k$ tokens (determined by the capacity factor), eliminating the load imbalance problem that token-based routing addresses with auxiliary loss functions. However, some tokens might not be selected by any expert (if they score low across all experts), which is the tradeoff.

Why both are included in the search space: the paper argues in Section 4 that "routing strategy can change the optimal model architecture when sparsely activated layers are introduced." The intuition is that an architecture with many narrow experts might benefit from expert-choice routing (guaranteed load balance prevents experts from being underutilized), while an architecture with few wide experts might prefer top-2 token routing (tokens have more experts to choose from, reducing the risk of tokens being dropped). By including both in the search space, the evolutionary algorithm can discover which gating function pairs best with which block composition, without the paper committing to one approach a priori.

Normalization along the expert dimension: the paper specifies that for both routing strategies, the affinity scores should be "normalized either along the token dimension or the expert dimension," and recommends "normalizing along the expert dimension for both token-based routing and expert-based routing" to "avoid causal leakage in decoding mode." Causal leakage would occur if the normalization for a given token depended on other tokens in the sequence — during autoregressive generation, future tokens are not yet available, making token-dimension normalization invalid. Expert-dimension normalization (softmax over experts for each token individually) is causal by construction.

---

Discovered Brainformer Blocks (Figures 8 and 9, Section 6.1)

The paper presents two top-performing discovered blocks (Brainformer Block 1 and Block 2) in Figures 8 and 9 and discusses their properties in Section 6.1. These visualizations show the concrete architectural patterns that the evolutionary search discovered.

Brainformer Block 1 (Figure 8) contains 8 sub-layers. This block is repeated 3 times (100M scale), 6 times (1B scale), or 8 times (8B scale). The search selected:

• A larger model dimension of 1024 (compared to the baseline's 768)
• A smaller expansion factor in the dense FFN and MoE hidden dimensions (compared to the standard 4× of 3072 given 768 model dim)
• The Expert Choice gating function with a capacity factor of 1, meaning each token is routed to a single expert on average — a very sparse configuration
• This sparsity makes the block fast in step time, allowing more training steps (and thus more data) within the fixed wall-clock budget

Brainformer Block 2 (Figure 9) is an alternative top candidate with a larger capacity factor in the MoE layers. This block is "lightly slower in step time, but takes fewer training steps to get good accuracy, thus is more data efficient." The tradeoff between these two blocks illustrates the Pareto frontier that the evolutionary search explores — Block 1 maximizes throughput (more steps, more data), while Block 2 maximizes per-step quality (better use of each training token), and both achieve competitive final perplexity under the fixed-time constraint.

Ablation finding on simplification (Section 6.2): the paper reports that the ratio of layer types (how many attention vs. dense FFN vs. MoE layers appear in the block) is critical to quality — replacing one layer type with another degrades performance. However, the network is relatively insensitive to layer order within the block — swapping any two adjacent sub-layers does not significantly affect performance. This means the block composition (the counts of each layer type) matters much more than the specific sequence. For practitioners seeking to simplify the architecture, this suggests that layers of the same type can be grouped contiguously or interleaved without major quality loss, as long as the total count of each type is preserved.

---

Design Choices and Their Justifications

Block-level search rather than layer-level search: reduces the combinatorial search space from $3^L$ (where $L$ is total layers, e.g., 64) to $3^k$ (where $k$ is sub-layers per block, discovered to be 8), and enables stacking-based scaling without re-optimizing the architecture for each scale. The assumption is that a good block at small scale transfers to large scale, which is validated empirically.

Evolutionary algorithm over gradient-based NAS: the search space includes discrete choices (layer type, gating function type) that are not differentiable. Evolutionary search handles discrete-continuous mixed spaces naturally.

Early stopping at 25% of max steps: prunes unpromising candidates early (violating step-time or perplexity constraints), enabling the search to evaluate more architectures within a fixed total compute budget. 25% is chosen empirically as a point early enough to save compute but late enough for quality trends to be visible.

Fixed wall-clock time budget: creates a level playing field between architectures that trade off capacity for speed. This is the core fairness mechanism — the search rewards architectures that make good use of hardware, not just architectures with high per-step capacity.

Scaling by dimension multiplication and block repetition: follows the EfficientNet compound scaling philosophy. Model dimensions are multiplied by 2× and 4× (matching GLaM scaling factors), and depth is increased by stacking the block more times. This is simpler than layer-wise compound coefficient search (used in EfficientNet) and is motivated by the block-level abstraction making scaling straightforward.

4. Key Insights and Innovations

Innovation 1: Uniformity Is an Assumption, Not a Requirement — and It's Costly

The paper's most fundamental conceptual move is to question a design choice so deeply embedded in transformer research that it is rarely acknowledged as a choice at all: the assumption that transformer layers should alternate in a fixed, repeating pattern. From the original "Attention Is All You Need" through BERT, GPT-3, GLaM, and Switch Transformer, every major transformer variant stacks identical blocks — whether [attention → dense FFN] or [attention → MoE → dense FFN] in GLaM's interleaved pattern — and repeats that block L times. This uniformity is a convenience, not a requirement, and the Brainformer paper treats it as a design constraint to be relaxed, not a law to be obeyed.

What makes this insight genuinely novel is not just that the paper tries non-uniform architectures — prior work like Sandwich Transformer (Press et al., 2019) experimented with layer reordering and showed mixed results. Rather, the innovation is the systematic argument that uniformity carries an efficiency tax that prior work never quantified, and that this tax is large enough to justify paying the complexity cost of heterogeneity. Concretely, a uniform block forces every token to encounter attention layers and feed-forward layers at exactly the same ratio regardless of where it sits in the depth of the network, what information has already been extracted, or what hardware the model runs on. Early layers might benefit from more spatial mixing (attention), middle layers from specialized expert retrieval (MoE), and later layers from dense nonlinear transformations — but a uniform block allocates the same mix everywhere.

Prior NAS work in the transformer space (Primer, So et al., 2021) operated entirely within the uniform-block paradigm: it searched over activation functions, layer normalization placements, and attention variants, but always produced a single block type that was then uniformly repeated. The Brainformer paper shows that relaxing the block-uniformity constraint — allowing different layer types to appear at different ratios within a repeated heterogeneous block — yields gains that exceed what intra-block optimization alone can deliver. The comparison with Primer in Table 5 is the smoking gun: Primer found a highly optimized dense transformer block through NAS, yet Brainformer's heterogeneous sparse block outperforms it on fewshot tasks while being comparably fast (1.37 vs. 1.50 steps/sec). The gap is attributable specifically to the additional degrees of freedom that Primer's uniform-block assumption denied it.

This is not an incremental refinement. It is a fundamental reframing of the architecture design problem for transformers: from "design the best layer, then stack it" to "design the best block composition, where the block itself is heterogeneous and the ratios of layer types are tunable." The fact that the discovered Brainformer blocks contain 8 sub-layers (Figures 8, 9) — substantially more complex than the traditional 2-sublayer transformer block — and that the ratio of layer types matters critically while the ordering is largely irrelevant (Section 6.2 ablation) reinforces that the composition is the key variable, not the sequence.

---

Innovation 2: Wall-Clock Time, Not FLOPs, Should Be the Fair-Comparison Basis Across Model Families

The paper makes a methodological contribution that extends beyond its own architecture: a reframing of what "efficiency" means in model comparison. Prior scaling studies compared models at equal parameter count (penalizing sparse models), equal training tokens (penalizing fast models), or equal FLOPs (better, but ignoring hardware utilization differences). The Brainformer proposes a more holistic criterion: compare at fixed wall-clock training time with an inference step-time constraint.

This matters because it fundamentally changes what architectures the optimization surface rewards. Under a fixed-training-tokens comparison, an architecture that processes tokens 2× faster gets no credit for its speed — it finishes training sooner but its final perplexity is compared to a slower model that trained for the same number of tokens at higher cost. Under a fixed-wall-clock-time comparison, that same speed advantage translates directly into more training data seen within the budget, potentially yielding lower perplexity despite lower per-step capacity. The paper calls this out explicitly: "fixing training wall clock time while meeting an inference time constraint yields models with faster training convergence and higher quality."

This framing connects to the Chinchilla scaling laws (Hoffmann et al., 2022) but extends them from the pretraining domain — where the tradeoff is between model size and data quantity at a fixed FLOP budget — to the architecture search domain, where the tradeoff is between per-step capacity and per-step throughput at a fixed wall-clock budget. An architecture that is faster but slightly less capable per step can win if its speed lets it process enough additional data to overtake its slower-but-smarter competitor. The Brainformer search discovers precisely this tradeoff: Block 1 maximizes throughput (capacity factor 1, narrow experts, Expert Choice routing for load balance) and wins by data volume, while Block 2 has higher per-step accuracy but slower throughput and wins on data efficiency (Section 6.1).

Prior NAS work in NLP — including Primer — did not explicitly incorporate wall-clock time into the reward function. The search objective typically optimized validation loss alone, with step time or inference latency treated as a post-hoc filter or secondary metric. The Brainformer paper elevates step time to a first-class constraint in the optimization (Equation 5), enabling the search to discover architectures that are genuinely Pareto-optimal on the quality-throughput frontier rather than architectures that optimize quality in a vacuum and hope to be fast enough.

This is a fundamental methodological advance for architecture search, not just a convenience. It means the Brainformer search space is explored under a fitness criterion that reflects real-world deployment constraints — training budget in TPU/GPU hours, not just abstract FLOP counts — making the discovered architectures more likely to be useful in practice. The evidence that this matters comes from the 8B-scale results (Table 3): Brainformer achieves 1.96 steps/sec vs. GLaM's 0.39 steps/sec — a 5× throughput gap — while also achieving better perplexity (1.99 vs. 2.12). These gains are not independent; the throughput advantage gave Brainformer more training steps within the fixed wall-clock budget, contributing to its quality advantage. A FLOPs-only comparison would have missed this coupling.

---

Innovation 3: Gating Strategy Is Not a Fixed Hyperparameter — It Interacts with Architecture

Prior MoE work treated gating mechanism selection as a design decision made once at the start and held constant across the entire model. GShard and GLaM committed to top-2 token routing. Switch Transformer committed to top-1 token routing with an auxiliary load-balancing loss. Expert Choice (Zhou et al., 2022) proposed a fundamentally different approach — expert-based routing — and demonstrated its benefits, but did so on a fixed, hand-designed architecture. The implicit assumption across all this work was: pick the best gating function for your use case, apply it uniformly, and let the model adapt.

The Brainformer paper challenges this assumption by including gating function as a tunable architectural parameter within the search space, alongside layer composition, model dimensions, and capacity factors (Table 1). The paper argues — and the evolutionary search validates — that the optimal gating mechanism depends on the surrounding architecture. The discovered Brainformer Block 1 selected Expert Choice routing with a capacity factor of 1 (very sparse: one expert per token on average), while Block 2 selected a larger capacity factor with presumably the same or a different gating function. These choices were not made by the paper's authors; they emerged from the search as the best-fitting pairings for each block's specific layer composition and width configuration.

This is significant as a diagnostic insight: the field's prior approach of comparing gating mechanisms on a fixed architecture conflates the quality of the gating function with the suitability of the architecture it was tested on. A routing method that underperforms on a dense, uniform baseline might excel when paired with a heterogeneous block composition discovered by search. By making gating a searchable dimension, the paper demonstrates that routing strategy and architecture are co-adapted — the right combination matters more than finding the universally best routing method in isolation.

The downstream implication is practical: teams building MoE models should not simply adopt the latest published gating method and bolt it onto their existing architecture. They should search jointly over architecture and gating, or at minimum test multiple gating variants per candidate architecture rather than fixing gating early in development and optimizing architecture around it. The fact that Brainformer Block 1 uses Expert Choice routing — a relatively recent method at the time — with a very sparse capacity factor of 1, and that this pairing contributes to the block's 5× step-time advantage at the 8B scale, suggests that the interaction between sparsity, routing, and architecture is where significant efficiency gains are hiding.

This is a conceptual contribution (gating and architecture are co-adaptive) backed by empirical discovery (the search found different optimal gate-architecture pairings for different block candidates), making it a fundamental reframing of how MoE design should be approached rather than an incremental finding.

---

Innovation 4: Block-Level Heterogeneity Is More Important Than Layer Ordering

In Section 6.2, the paper reports a finding from an ablation study that is arguably more surprising than the main results: the ratio of layer types matters critically, but the order in which those layers appear within the block does not. Swapping any two adjacent sub-layers within the discovered Brainformer block does not significantly degrade quality, while replacing one layer type with another (changing the ratio) does.

This is a diagnostic insight with direct practical implications. It means the compositional constraint — "a Brainformer block needs roughly X attention layers, Y dense FFN layers, and Z MoE layers" — is what carries the architectural signal, not the specific interleaving pattern. The field's obsession with whether attention should precede or follow feed-forward layers (a debate that spans from the original Transformer paper through Sandwich Transformer and Primer) may be largely misdirected. What matters is not the order but the ratios and the widths.

This finding also explains why the evolutionary search discovered good architectures in a manageable number of trials (the paper notes "the search identified better model architecture within as early as 500 trials" in Section 8): the effective search space is smaller than it appears because many orderings are approximately equivalent in quality. The search mainly needs to discover the right layer-type counts and dimension assignments, after which ordering is forgiving. This is a conceptual simplification of the architecture design problem masquerading as a negative result — it tells future practitioners that they can interleave layer types arbitrarily or group same-type layers contiguously without much loss, greatly reducing the complexity of implementing Brainformer-like architectures.

For the broader architecture search community, this finding suggests that search spaces should focus on compositional ratios rather than fine-grained ordering. Including ordering in the search space when it doesn't matter wastes compute exploring equivalent configurations and slows convergence. Future NAS work building on Brainformer could constrain the search to discover ratios first, then test orderings only on top candidates — a two-stage approach that would be more sample-efficient than the current joint search.

This is a fundamental empirical finding (not a theoretical advance or an incremental refinement) that changes how one should think about the architecture design space: the "what" (which layer types, at what widths) dominates the "where" (in what sequence).

5. Experimental Analysis

Evaluation Methodology

• Dataset. All models are trained on a high-quality 1.6-trillion-token corpus originally assembled for GLaM (Du et al., 2022), consisting of a filtered subset of webpages combined with smaller corpora of books, Wikipedia, conversations, forums, and news. The exact mixture weights follow the GLaM paper. Training uses a maximum sequence length of 1024 tokens with a SentencePiece subword tokenizer of vocabulary size 256K.

• Base model(s). The search operates against a 100M-parameter 12-layer dense transformer with model dimension 768, serving as the baseline from which the search space is defined (Table 1). For final evaluation, Brainformer models are compared against two families: (1) GLaM sparse MoE models at 100M/32E, 1B/64E, and 8B/64E scales with manually crafted interleaved dense-and-sparse transformer blocks and fixed top-2 token routing (Du et al., 2022); and (2) Primer, a NAS-discovered dense transformer (So et al., 2021) evaluated at the 1B scale. All comparisons are grouped by activated parameters per token, which approximates computational cost per inference step (Table 2).

• Metrics. The primary training metric is pre-training validation perplexity (cross-entropy loss), reported with ± standard deviation across multiple runs. For downstream evaluation, the paper reports: (1) fine-tuning accuracy on 11 selected classification tasks from GLUE and SuperGLUE, aggregated as an average score; (2) one-shot accuracy on five generative question-answering tasks (Natural Questions, TriviaQA, Web Questions, Squadv2, Lambada); and (3) training throughput measured in steps per second on 64 or 512 Cloud TPU-V4 chips. The paper does not report standard deviations for downstream task scores, only for perplexity.

• Baselines. Three distinct baselines are used across experiments: (1) GLaM — the hand-designed sparse MoE architecture with interleaved dense and sparse transformer blocks, top-2 token routing, and uniform hidden dimensions (Du et al., 2022); (2) Primer — a dense transformer discovered via neural architecture search that optimizes activation functions, normalization, and attention mechanisms within a uniform repeating block (So et al., 2021); and (3) Search-w-Top2 — an intermediate ablation that uses the same evolutionary search procedure as Brainformer but with the gating function fixed to top-2 token routing rather than including Expert Choice in the search space. This isolates the contribution of gating flexibility from the contribution of architectural heterogeneity.

• Generation budget / compute accounting. Compute is measured in two complementary ways. For training comparisons, the paper uses fixed wall-clock training time — all models within a scale group are trained for the same duration on the same hardware (64 TPU-V4 chips for 100M and 1B scales, 512 TPU-V4 chips for 8B scale), meaning faster models process more training steps and thus more data. For inference comparisons, activated parameters per token serves as the primary cost metric (Table 2), supplemented by direct step time measurements (steps per second) on identical hardware. The search itself operates under a fixed wall-clock time constraint per trial (Section 3.4), with early stopping at 25% of maximum steps for architectures that violate the inference step-time constraint or a perplexity threshold.

• Cross-validation / statistical protocol. The evaluaton section does not describe cross-validation for the downstream task results. Pre-training perplexity is reported with ± standard deviation (Table 3), but downstream fine-tuning scores (Table 4) and fewshot scores (Table 5) are reported as point estimates without confidence intervals. The evolutionary search uses a population-based approach with top-k selection from the final population, but no statistical testing between final candidates is described.

---

Main Quantitative Results

Pre-Training Convergence and Throughput (Table 3, Figure 7)

At every scale, Brainformer achieves lower perplexity with faster step time and similar or fewer activated parameters compared to GLaM and Search-w-Top2.

At the smallest scale (100M/32E) , the numbers from Table 3 are:

Brainformer-1 achieves a perplexity improvement of 0.16 over GLaM (2.73 → 2.57) — roughly 5.9% relative reduction — while using comparable activated parameters (156M vs. 145M) and slightly faster step time (2.03 vs. 1.92 steps/sec). The Search-w-Top2 baseline improves over GLaM (2.67 vs. 2.73) but falls short of Brainformer-1, confirming that architectural heterogeneity alone (with fixed top-2 gating) helps, but the full gain requires jointly optimizing architecture and gating.

At the 1B/64E scale:

Here the result is more nuanced: Brainformer-1 matches GLaM's perplexity (2.25) while using 1.38B activated parameters vs. GLaM's 1.88B — that is, it achieves equivalent quality with 27% fewer activated parameters per token. It is also 63% faster in step time (2.00 vs. 1.23 steps/sec). Brainformer-2 achieves slightly better perplexity (2.23) with even fewer activated parameters (1.31B) but slower step time than Brainformer-1 (1.76 vs. 2.00). The Search-w-Top2 baseline outperforms GLaM in perplexity (2.21 vs. 2.25) but at the cost of substantially more activated parameters (3.05B vs. 1.88B) — making it less efficient overall. This is the clearest evidence that fixing gating to top-2 while searching architecture produces models with worse compute-quality tradeoffs than joint optimization.

At the largest scale (8B/64E) , the paper's headline result:

Brainformer-1 achieves a perplexity of 1.99 — 0.13 lower than GLaM's 2.12 — while using 7.4B activated parameters vs. 9.8B (a 24% reduction), and running at 1.96 steps/sec vs. 0.39 (a 5.0× speedup). The Expert-based Gating baseline (which takes the GLaM hand-designed architecture but swaps in Expert Choice routing) improves over standard GLaM in both perplexity (2.03 vs. 2.12) and step time (0.50 vs. 0.39), confirming that Expert Choice routing helps even on a fixed architecture. However, Brainformer-1 still substantially outperforms this baseline: 1.99 perplexity vs. 2.03, and 1.96 steps/sec vs. 0.50 (a 3.9× speedup). This means the gap between Expert-based Gating on GLaM's architecture and Brainformer-1 is attributable to architectural heterogeneity — the non-uniform block composition, the tuned layer widths, and the discovered sub-layer ratios that the search identified.

Figure 7 reinforces these results visually: (a) at the 100M32E scale, Brainformer (solid line) trains to lower perplexity than both GLaM (dashed) and Search-w-Top2, with the gap widening as training progresses; (b) at the 8B64E scale, Brainformer (solid) trains to lower perplexity than Expert Choice on GLaM's architecture (dashed), again with a growing gap.

The paper describes these results as demonstrating "2× faster training convergence" — meaning Brainformer reaches a given perplexity threshold in roughly half the number of steps that GLaM requires, though this is inferential since the paper trains for fixed wall-clock time (so all models train for the same duration; the claim is that Brainformer converges faster within that duration). The 5× step time improvement is directly measured.

---

Fine-Tuning Results on GLUE/SuperGLUE (Table 4)

Brainformer-1 substantially outperforms GLaM on fine-tuned classification tasks at both the 100M/64E and 1B/64E scales, with average score improvements of 2.1 and 4.0 percentage points respectively.

At the 100M/64E scale (Table 4):

Brainformer-1 wins on 8 of 11 tasks, ties on 1 (SST2), and loses on exactly 1: QQP, where GLaM scores 0.907 and Brainformer-1 scores 0.812 — a substantial 9.5-point deficit. This is the only clear negative result in the fine-tuning evaluation and is notable because QQP (Quora Question Pairs) is a pairwise semantic equivalence task that may benefit from a different allocation of attention vs. feed-forward capacity than what the search discovered. The paper does not comment on this specific result, but it is the only task-level evidence that Brainformer's architectural gains are not uniform across all task types.

At the 1B/64E scale:

Brainformer-1 wins on 10 of 11 tasks and ties on 1 (CB, both 0.938). The largest individual gains are on WNLI (+0.172) and RTE (+0.083) — both natural language inference tasks that require understanding logical relationships between sentences. The QQP deficit observed at 100M disappears at 1B, where Brainformer edges ahead (0.917 vs. 0.911), suggesting the issue may be scale-dependent. The average score improvement of 4.0 percentage points (0.833 → 0.873) is substantial — roughly the gap between a strong baseline and a new state of the art on these benchmarks.

The paper trains all models for a fixed wall-clock time before fine-tuning, meaning Brainformer's pre-training advantage (more data processed due to faster step time) partially explains the downstream gains. The paper does not ablate whether the architectural advantage persists when controlling for total training tokens; this is a genuine limitation discussed in the Critical Assessment.

---

Fewshot Results (Table 5)

Brainformer 1B/64E outperforms both GLaM 1B/64E and Primer 1B on four of five generative QA tasks in the one-shot setting, while being substantially faster than GLaM.

Table 5 reports exact match or F1 scores (the metric varies by task, following each benchmark's standard):

Brainformer wins on TriviaQA (43.4 vs. 41.8 GLaM, 24.7 Primer), WebQA (12.0 vs. 10.8 GLaM, 6.50 Primer), Squadv2 (49.5 vs. 46.2 GLaM, 49.2 Primer), and Lambada (25.7 vs. 25.2 GLaM, 22.6 Primer). It loses to GLaM on Natural Questions (8.23 vs. 9.14) — the second task-level deficit observed for Brainformer across all evaluations, alongside the QQP fine-tuning loss at 100M.

The comparison with Primer is particularly informative because Primer represents the state of the art in NAS-discovered dense transformers. Brainformer outperforms Primer on every single task, in some cases by large margins (TriviaQA: 43.4 vs. 24.7; WebQA: 12.0 vs. 6.50), while being only marginally slower in step time (1.37 vs. 1.50 steps/sec). The Primer result is the paper's strongest evidence that block-level heterogeneity with sparsity outperforms dense architectures optimized via NAS under the uniform-block paradigm.

The GLaM comparison at the 1B scale shows a mixed efficiency picture: Brainformer achieves slightly better quality on most tasks but at much higher throughput (1.37 vs. 0.55 steps/sec — a 2.5× speedup). However, all models in this comparison were trained with 200B training tokens (as noted in the Table 5 caption), not with the fixed wall-clock protocol used in Table 3. This means Brainformer's pre-training throughput advantage did not translate into more training data for this particular comparison — the architectural quality gain is isolated from the training data advantage.

---

Search Space Ablation: Search-w-Top2 vs. Full Brainformer (Table 3, Figure 7)

Architectural search with fixed top-2 gating (Search-w-Top2) improves over GLaM but consistently underperforms Brainformer with flexible gating, confirming that gating-architecture co-optimization is the key driver, not architecture search alone.

At all three scales from Table 3:

• 100M/32E: Search-w-Top2 achieves 2.67 PPLX (vs. 2.73 GLaM, a 0.06 improvement) at 2.03 steps/sec, while Brainformer-1 achieves 2.57 (an additional 0.10 improvement) at the same step time. The gap between Search-w-Top2 and Brainformer (0.10 PPLX) is larger than the gap between Search-w-Top2 and GLaM (0.06 PPLX), suggesting gating flexibility contributes more than architecture heterogeneity alone at this scale.

• 1B/64E: Search-w-Top2 achieves 2.21 PPLX at 1.27 steps/sec with 3.05B activated parameters — better perplexity than Brainformer-1's 2.25 but at 2.2× the activated parameters and 37% slower step time. This is the scale where fixing gating to top-2 most clearly hurts compute-efficiency: Search-w-Top2 finds a quality-competitive architecture but one that is expensive to run.

• 8B/64E: Search-w-Top2 is not evaluated at this scale. The Expert Choice baseline on GLaM's fixed architecture (2.03 PPLX, 0.50 steps/sec) provides an approximate comparison — Brainformer-1's 1.99 PPLX at 1.96 steps/sec represents an enormous combined quality and speed improvement over what gating flexibility on a fixed architecture can achieve.

---

Ablation Studies and Robustness Checks

• Gating function search vs. fixed top-2 (Search-w-Top2 ablation, Table 3, Figure 7). Including gating function (Top-2 vs. Expert Choice) in the search space yields substantially better compute-efficiency than searching architecture alone with fixed top-2 gating. At 100M/32E, Search-w-Top2 achieves 2.67 PPLX while Brainformer-1 achieves 2.57; at 1B/64E, Search-w-Top2 requires 3.05B activated parameters for 2.21 PPLX while Brainformer-1 uses only 1.38B activated parameters for 2.25 PPLX — a 2.2× reduction in inference cost for comparable quality. This demonstrates that gating strategy and architecture are co-adapted and should not be selected independently.

• Expert Choice routing on fixed GLaM architecture (Table 3, 8B/64E row). Replacing GLaM's top-2 token routing with Expert Choice routing while keeping the interleaved dense-sparse block design fixed improves perplexity from 2.12 to 2.03 and step time from 0.39 to 0.50 steps/sec. This is a meaningful gain (0.09 PPLX improvement, 28% faster step time) but is dwarfed by Brainformer-1's combined architecture + gating improvement (1.99 PPLX, 1.96 steps/sec). The residual gap between Expert Choice on GLaM's architecture and Brainformer-1 quantifies the architectural heterogeneity benefit at ~0.04 PPLX and a 3.9× step time improvement.

• Brainformer block 1 vs. block 2 (Table 3, Section 6.1). The search discovers a Pareto frontier. Block 1 uses Expert Choice routing with capacity factor 1, achieving maximum throughput and winning on total quality under fixed wall-clock time by processing more training data. Block 2 uses a larger capacity factor, achieving competitive perplexity with fewer training steps but slower throughput. At 100M/32E, Block 2 has worse perplexity (2.59 vs. 2.57) despite more activated parameters (266M vs. 156M) but slightly faster step time (2.16 vs. 2.03); at 1B/64E, Block 2 achieves better perplexity (2.23 vs. 2.25) with fewer activated parameters (1.31B vs. 1.38B) but slower step time (1.76 vs. 2.00). Both blocks outperform GLaM on all metrics, confirming the search discovers multiple valid solutions on the quality-throughput frontier.

• Layer type ratio vs. layer order (Section 6.2, no table). The paper reports a qualitative ablation: replacing a layer of one type with another (changing the ratio of attention-to-MoE-to-dense-FFN) degrades quality, but swapping the position of any two adjacent layers within the block does not significantly affect performance. No quantitative results are reported for this claim — it is stated as a finding from "an ablation study on block simplification" without supporting numbers, which weakens it considerably. If true, it has important implications for practitioners (layer order is forgiving), but the lack of numerical evidence makes it difficult to assess the magnitude of invariance.

• Scaling from 100M proxy to 1B and 8B target models (Tables 3 and 4). The paper's core methodological assumption — that block architectures discovered at the 100M proxy scale transfer to larger scales — is validated implicitly by the strong performance at 1B and 8B. However, there is no ablation testing whether re-running the search directly at the 1B or 8B scale would produce different (better) architectures. The assumption of transferability is convenient but untested — architectures optimal at 100M might be suboptimal at 8B (e.g., the optimal attention-to-FFN ratio might shift with scale), and the paper provides no evidence either way.

---

Critical Assessment

Claim from Section 1: Brainformer demonstrates "2× faster training convergence and 5× faster step time compared to its GLaM counterpart."

The 5× step time claim is strongly supported by Table 3: at 8B/64E, Brainformer-1 achieves 1.96 steps/sec vs. GLaM's 0.39 steps/sec, which is a factor of 5.02×. This is a direct, hardware-controlled measurement and is unambiguous.

The 2× training convergence claim requires more careful interpretation. The paper states this in the abstract and introduction, but Table 3 trains all models for a fixed number of steps (0.5M, 1.0M, or 1.5M depending on scale) and reports final perplexity. Faster convergence typically means reaching a target perplexity in fewer steps, but the paper does not show learning curves annotated with the number of steps to reach specific perplexity thresholds. What Table 3 shows is that at 8B/64E, Brainformer-1 reaches 1.99 PPLX after 1.5M steps while GLaM reaches 2.12 after the same 1.5M steps. But GLaM processes steps at 0.39 steps/sec while Brainformer processes at 1.96 steps/sec, so after the same wall-clock time, Brainformer has processed 5× more steps. If GLaM were allowed to train for proportionally longer (7.7M steps to match Brainformer's total FLOPs or wall-clock time), would it catch up? The paper doesn't answer this because the fixed-step comparison controls for steps, not time, and the fixed-wall-clock-time protocol (Section 3.4) applies only during the search phase, not during final evaluation. The 2× convergence claim is therefore inferential and only partially supported; the paper would be stronger with perplexity-vs-wall-clock-time curves at each scale.

Claim from Section 1: Brainformer "demonstrates a 3% higher SuperGLUE score with fine-tuning compared to GLaM with a similar number of activated parameters."

Supported by Table 4 but with qualifications. At 100M/64E, Brainformer-1 achieves 0.840 average vs. GLaM's 0.819 — a difference of 2.1 percentage points (not 3%). At 1B/64E, the difference is 4.0 percentage points (0.873 vs. 0.833). The "3%" claim appears to be a rough average or refers specifically to the 1B scale. The "similar number of activated parameters" qualification holds roughly: at 100M/64E, Brainformer-1 has 156M vs. GLaM's 145M (7% more); at 1B/64E, Brainformer-1 has 1.38B vs. GLaM's 1.88B (27% fewer — Brainformer actually has substantially fewer activated parameters, making the quality gain more impressive). The paper does not control for pre-training data volume between these comparisons (Brainformer was trained for the same wall-clock time and thus processed more data, which partially explains downstream gains).

Claim from Section 1: Brainformer "largely outperforms a Primer dense model derived with NAS with similar computation per token on fewshot evaluations."

Strongly supported by Table 5. Brainformer wins on all five tasks, in some cases by large margins (TriviaQA: 43.4 vs. 24.7, WebQA: 12.0 vs. 6.50). The throughput comparison is favorable: Brainformer at 1.37 steps/sec is slightly slower than Primer at 1.50 steps/sec (9% slower), but the quality gap is much larger than the speed gap, meaning Brainformer is unambiguously on a better point of the quality-throughput Pareto frontier. The caveat is that this is a single scale (1B) and both models were trained on only 200B tokens, which may not represent their behavior at convergence or at other scales.

Genuine weaknesses in the experimental design:

1. No wall-clock-time-controlled training comparisons at final evaluation. The search phase uses fixed wall-clock time (Section 3.4), but the final evaluation in Table 3 uses fixed training steps, not fixed time. This means the comparison does not directly measure what the search optimized for — it compares architectures at equal steps, not at equal cost. Since Brainformer is faster per step, it actually achieved its reported perplexity with less wall-clock time than GLaM, but the paper never quantifies this advantage in hours or TPU-days.

2. Pre-training data volume is confounded with architecture quality in downstream results. All fine-tuning and fewshot comparisons use models trained for equal wall-clock time, meaning Brainformer processed more tokens during pre-training (due to faster throughput). The downstream quality advantage could come from (a) better architecture, (b) more pre-training data, or (c) a combination. Without a comparison where GLaM is trained on the same number of tokens as Brainformer (or Brainformer is capped to GLaM's token count), the causal attribution is impossible. The paper does not run this ablation.

3. Single dataset for pre-training (GLaM's 1.6T-token corpus). While large and high-quality, this is a single training distribution. Whether the discovered Brainformer block transfers to other corpora (e.g., more code-heavy, multilingual, or domain-specific data) is untested. The architecture search was performed on this corpus and the proxy task uses this corpus — it's possible that the optimal block composition is dataset-specific.

4. No statistical significance or confidence intervals on downstream results. Table 4 reports fine-tuning scores to three decimal places without standard deviations. Given that some task-level differences are small (e.g., WiC at 1B: 0.720 vs. 0.711, a 0.009 gap), it's impossible to assess whether these differences are reliable without variance estimates. The QQP result at 100M (0.812 vs. 0.907, a 9.5-point deficit) is large enough to likely be significant, but the paper does not discuss it, and without error bars it's unclear whether it's a genuine task-specific weakness or within noise.

5. No comparison to Switch Transformer or other MoE architectures beyond GLaM. The paper compares only to GLaM (hand-designed sparse) and Primer (NAS-discovered dense). Other sparse architectures (Switch Transformer, BASE Layers, Hash Layers) are cited in related work but not evaluated. The claim that Brainformer "consistently outperforms the state-of-the-art dense and sparse Transformers" (abstract) overstates the evidence, since only one sparse baseline (GLaM) and one dense NAS baseline (Primer) are tested.

6. The layer-order invariance claim (Section 6.2) is stated without quantitative evidence. This is an important finding if true — it would simplify Brainformer implementation dramatically — but the paper provides no table, figure, or numerical results to support it. The claim reads as an informal observation from an unpublished ablation study.

7. The search is computationally expensive and the cost is not amortized over different model families. The paper reports using 512 TPU-V4 chips for one week for the search (Section 8). Whether this cost is justified depends on how many downstream models benefit from the discovered block. If the block is general enough to apply to many model sizes and tasks, the cost amortizes; if a new search is needed for each significantly different setting (different data, different scale, different hardware), the cost may be prohibitive. The paper provides no evidence on transferability beyond scaling within the same data regime, which limits the practical applicability of the method.

Missing experiments that would strengthen the paper:

• A perplexity-vs-wall-clock-time plot for the 8B models, showing Brainformer-1 reaching the same perplexity as GLaM in a fraction of the wall-clock time (or lower perplexity at equal time) — this would directly substantiate the 2× convergence claim.
• A token-matched pre-training comparison: train GLaM for the same number of tokens as Brainformer-1 processed within the fixed time budget, to isolate the architectural quality gain from the data volume gain.
• Downstream evaluation with error bars or multiple fine-tuning runs to assess whether task-level differences are statistically reliable.
• A test of the discovered block on an out-of-distribution pre-training corpus to assess transferability.
• Evaluation at the 8B scale on downstream tasks, not just pre-training perplexity — the paper's strongest efficiency results are at 8B (Table 3), but all downstream comparisons (Tables 4 and 5) are at 100M and 1B scales. Without 8B downstream results, the paper's strongest architectural claim is limited to pre-training.
• A quantitative ablation on layer ordering vs. layer ratio with perplexity numbers, to substantiate the Section 6.2 claim that ordering is irrelevant.

In summary, the experiments strongly support the paper's core architectural claims — Brainformer blocks discovered via joint architecture-and-gating search under fixed wall-clock time yield faster training and inference throughput with competitive or better quality at multiple scales. However, the downstream quality advantage is partially confounded with pre-training data volume, the 2× convergence claim requires more careful supporting evidence than is provided, and the generalizability of the discovered architecture to other data distributions and model families remains unproven. The paper makes a compelling case for block-level heterogeneity as an underexplored design dimension, but the evidence that the specific discovered blocks are optimal (rather than merely good) is limited by the absence of comparisons to a broader set of sparse baselines and by the lack of token-controlled pre-training ablations.

6. Limitations and Trade-offs

Pre-Training Data Volume Confounds Architectural Quality Gains in Downstream Evaluation

The paper uses a fixed wall-clock training budget during the search phase and a fixed-step protocol during final pre-training evaluation, but all downstream fine-tuning and fewshot comparisons are made using models that were pre-trained for equal wall-clock time. Because Brainformer trains faster than GLaM (5× faster step time at the 8B scale, Table 3), models evaluated in Tables 4 and 5 have processed substantially more pre-training data than their baselines. The paper never controls for this: there is no comparison where GLaM is allowed to train on the same number of tokens as Brainformer, or where Brainformer is capped to GLaM's token count. This means the downstream quality gains attributed to Brainformer's architecture cannot be separated from the data volume effect. For a practitioner deciding whether to adopt Brainformer, this distinction matters enormously: architectural gains are permanent (you get them regardless of training budget), while data volume gains can be replicated by simply training any model longer. If a substantial fraction of Brainformer's 3–4% SuperGLUE average improvement (Table 4) comes from seeing more data rather than from architectural efficiency, then a well-funded team training GLaM for proportionally longer might close the gap without adopting the complexity of heterogeneous blocks. The paper's own protocols acknowledge the importance of time-based fairness during search (Section 3.4: "We empirically find that fixing training wall clock time while meeting an inference time constraint yields models with faster training convergence"), but the downstream evaluation does not match this standard. This is partially mitigated by the fewshot results (Table 5), where all models were trained on exactly 200B tokens — isolating the architectural quality contribution, and where Brainformer still outperforms GLaM and Primer on most tasks. However, the token-controlled comparison exists only at the 1B scale for generative tasks; the classification results in Table 4, which carry the paper's headline 3% SuperGLUE claim, remain confounded.

The Discovered Architecture Is Untested Outside the GLaM Data Distribution

Brainformer's block architecture was discovered via evolutionary search on a single pre-training corpus: the 1.6-trillion-token GLaM dataset, a filtered mixture of webpages, books, Wikipedia, conversations, forums, and news articles (Section 5, Dataset). The search's proxy task, early-stopping criteria, and reward function all use this same data distribution. There is no experiment testing whether the discovered block composition — its specific ratio of attention to dense FFN to MoE layers, its chosen model dimensions, its Expert Choice gating with capacity factor 1 — transfers to meaningfully different corpora (e.g., code-heavy data, multilingual data, domain-specific scientific or medical text, or a different language). The paper acknowledges this implicitly in its Limitations section (Section 8): "Our empirical results are primarily on NLP domain, thoroughly on a wide range of NLU and NLG tasks. However, we leave it to future work to apply Brainformer to computer vision." But the scope limitation is deeper than domain: even within NLP, the optimal block composition might be data-dependent. A corpus with longer documents might benefit from a different attention-to-FFN ratio than a corpus of short-form text; a multilingual corpus might require different MoE expert counts or gating strategies than an English-dominated corpus. The consequence is that a practitioner with a data distribution that differs substantially from GLaM's cannot assume the Brainformer block will transfer effectively without re-running the search, which costs 512 TPU-V4 chips for a week (Section 8). The paper provides no evidence on transferability to other corpora even at small scale, leaving adoption risk unquantified.

The Search Cost Makes the Method Impractical for All But the Largest Compute Budgets

The paper reports using 512 TPU-V4 chips for one week to complete the evolutionary search that discovered the Brainformer block (Section 8). This is an enormous compute expenditure — approximately 86,000 TPU-V4-hours — placed before any model is trained at target scale. The authors acknowledge this: "Another limitation can be large resource consumption. In the Brainformer search, we used 512 TPU v4 for a week to arrive at the best solutions" (Section 8). They argue the cost is partially mitigated because "the search identified better model architecture within as early as 500 trials" and "practically, the resource consumption can be small if we only need to identify better but suboptimal models." However, the paper provides no evidence on how much earlier-than-final architectures perform relative to the final discovered blocks: does a model found at trial 500 achieve 90% of the final gains? 50%? Without this, the cost-amortization argument is speculative. The consequence for practitioners is that Brainformer-style architecture search is viable only for organizations that can afford to burn thousands of TPU/GPU hours on search before training their actual model — which restricts the method to a small set of industrial labs. The paper's own cost comparison (a 2× training convergence improvement and 5× step-time improvement on a model trained on 1.6T tokens) suggests the search pays for itself at the 8B scale through downstream training savings, but requires the upfront capital to spend a week of 512-TPU-V4 time on search — a chicken-and-egg problem for smaller teams. The paper does not discuss whether the discovered block transfers across model families (e.g., does Brainformer Block 1 work well with a different tokenizer, a different initialization scheme, or a different base architecture like a decoder-only vs. encoder-decoder design?), which would be necessary to amortize the search cost across multiple projects. The authors suggest that "this will be mitigated when we use a smaller model size and smaller number of experts in the MoE layers," but no such experiments are reported.

No Downstream Evaluation at the 8B Scale Where Efficiency Gains Are Strongest

The paper's most dramatic efficiency results — 5× faster step time, 2× training convergence, and 1.99 vs. 2.12 perplexity improvement over GLaM — all come from the 8B/64E scale (Table 3). Yet every downstream task evaluation (SuperGLUE fine-tuning in Table 4, fewshot generative tasks in Table 5) is conducted at the 100M/64E and 1B/64E scales only. The paper never reports whether the pre-training perplexity improvements at 8B translate to downstream quality improvements, nor whether the architecture's task-level performance characteristics (the QQP deficit at 100M, the Nqs deficit at 1B in Table 5) persist, reverse, or amplify at scale. For a practitioner considering adopting Brainformer for a production model at the multi-billion-parameter scale, this is a critical missing data point: pre-training perplexity improvements do not always predict downstream performance (as the Chinchilla paper and others have documented), and the paper provides no evidence that Brainformer's 0.13 perplexity advantage at 8B (1.99 vs. 2.12) would yield meaningful gains on practical tasks. The omission is particularly striking given that the paper's strongest architectural claims (5× speedup, 2× convergence) are made at 8B. The paper's own resource constraints likely explain this gap — training and fine-tuning an 8B-scale model on the full SuperGLUE suite would be expensive — but the consequence is that the evaluation is incomplete: the paper demonstrates architectural advantages at the scale where they matter most on the metric that matters least for applications (pre-training perplexity), and demonstrates task-level advantages at small scales where the efficiency gains are modest.

The Layer-Order Invariance Claim Is Central to the Method's Practicality but Unsupported by Evidence

Section 6.2 reports a finding from an "ablation study on block simplification" that "the ratio of different layer types is critical to model quality" while "the network is relatively insensitive to layer order, such that swapping any two layers would not affect performance much." This is a practically important claim — if true, it means practitioners implementing Brainformer-style blocks can freely interleave or group layer types for hardware efficiency without worrying about precise ordering, dramatically simplifying deployment. However, the paper provides no quantitative evidence for this claim: no table, no figure, no perplexity comparisons between order-shuffled variants of the same block, not even a specification of how many orderings were tested or at what scale. The claim is stated as a qualitative observation from an unspecified ablation. The consequence is that a practitioner cannot assess the magnitude of the claimed invariance — does "not affect performance much" mean a 0.01 perplexity change (negligible) or a 0.05 change (which at the 8B scale could erase the gap between Brainformer and GLaM)? Without numbers, the paper's strongest architectural insight (that ratios matter more than order) remains an untested hypothesis rather than an established finding. This is not a minor omission — it's the paper's primary claim about the structure of the architecture search space and the key takeaway for simplifying Brainformer implementation, left entirely unsubstantiated. The paper does not flag this as a limitation or propose future work to quantify it; the result stands as an assertion.

Comparison to Only One Sparse Baseline (GLaM) Limits the Generality of "State-of-the-Art" Claims

The paper's abstract claims Brainformer "consistently outperforms the state-of-the-art dense and sparse Transformers," but the sparse comparison is exclusively against one baseline: GLaM (Du et al., 2022) with top-2 token routing and interleaved dense-sparse blocks. Other sparse architectures from the MoE literature — Switch Transformer (Fedus et al., 2021) which introduced top-1 routing, BASE Layers (Lewis et al., 2021) which proposed alternative gating formulations, Hash Layers (Roller et al., 2021), and the Expert Choice paper itself (Zhou et al., 2022) — are all cited in the related work (Section 2) but never evaluated. The paper does test Expert Choice routing, but only as a gating variant applied to GLaM's fixed architecture (Table 3, 8B/64E row) — it never compares against a model that was architected from the ground up for Expert Choice, or against Switch Transformer's simpler top-1 approach at comparable scale. The consequence is that the "state-of-the-art" claim is stronger than the evidence supports. Brainformer may outperform GLaM specifically, but GLaM is only one point in the larger MoE design space. A practitioner who has already adopted Switch Transformer or Expert Choice on a custom architecture cannot infer from this paper whether Brainformer would outperform their current setup — the comparison set is too narrow. The paper partially addresses this by including the Primer comparison (dense NAS baseline) and the Expert Choice-on-GLaM ablation, which shows that gating improvements alone do not close the gap. But the absence of alternative sparse baselines means the paper has demonstrated superiority over one sparse design pattern, not over the general category of sparse transformers. This limitation is structural — the cost of training competitors at the 8B scale would be prohibitive — but it constrains the generalizability of the paper's strongest claims.

7. Implications and Future Directions

How This Work Changes the Landscape

This paper causes a methodological reframing rather than a paradigm shift: it does not invent a fundamentally new layer type, training objective, or scaling law, but it demonstrates that the architectural regularity the field has taken for granted since "Attention Is All You Need" — uniform, alternating blocks stacked L times — is a self-imposed constraint that carries a quantifiable efficiency tax. By relaxing this constraint and allowing an evolutionary search to discover heterogeneous block compositions jointly with gating strategies and layer widths, the paper extracts gains (5× step-time improvement, 2× training convergence at the 8B scale, Table 3) that no prior work on efficient transformers achieved through intra-block optimization alone. This changes the landscape in three specific ways.

First, it redirects architecture search effort from "find a better layer" to "find a better block composition." Prior NAS work for transformers — most prominently Primer (So et al., 2021) — searched within the uniform-block paradigm: discover an optimized attention mechanism, activation function, or normalization placement, then stack that improved layer uniformly. Brainformer shows that the additional degrees of freedom from varying layer types and ratios within a heterogeneous block yield improvements that exceed what intra-layer optimization alone can deliver. The comparison with Primer in Table 5 is the direct evidence: Brainformer 1B/64E outperforms Primer 1B on every fewshot task while being only marginally slower (1.37 vs. 1.50 steps/sec), and Primer already represented the state of the art in NAS-optimized dense transformers. The takeaway is not that attention gating or activation functions don't matter — it's that the architectural composition (how many attention layers relative to FFN layers, where MoE layers are placed, at what widths) is a higher-leverage optimization surface than intra-layer details. Future NAS for transformers should expand their search spaces to include block-level heterogeneity; work that continues to optimize only within a uniform repeating block is leaving the larger efficiency gains on the table.

Second, it establishes that gating strategy and network architecture are co-adapted and must be searched jointly, not selected independently. The Search-w-Top2 ablation in Table 3 is the key evidence: when the evolutionary search is given architectural freedom but with gating fixed to top-2 token routing, it discovers models that improve over GLaM (2.67 vs. 2.73 PPLX at 100M/32E) but require substantially more activated parameters at the 1B scale (3.05B vs. 1.38B for Brainformer-1 at comparable quality). This means architecture search with a suboptimal, fixed gating function can produce models that are actively worse on the compute-efficiency Pareto frontier — they find expensive ways to compensate for a gating function poorly matched to the architecture. Prior work in MoE treated gating mechanism selection as a pre-experiment design choice (GShard and GLaM committed to top-2, Switch Transformer to top-1), and the Expert Choice paper (Zhou et al., 2022) demonstrated its method's benefits on a fixed architecture. Brainformer reframes the problem: the right question is not "which gating function is best?" but "which gating function pairs best with which architecture?" The Expert Choice-on-GLaM ablation (Table 3, 8B row: 2.03 PPLX at 0.50 steps/sec vs. GLaM's 2.12 at 0.39) confirms that Expert Choice routing helps even on a fixed architecture, but Brainformer's 1.99 PPLX at 1.96 steps/sec shows that the combination of architecture heterogeneity with Expert Choice routing is where the largest gains live. This implies that future MoE research should report results with gating as a tunable hyperparameter, not a fixed experimental condition, and that cross-gating comparisons on a single architecture are of limited informativeness.

Third, the paper's fixed-wall-clock-time comparison framework, while applied only during the search phase, exposes a methodological gap in how the field evaluates efficient architectures. The paper argues (Section 3.3) that comparing models at equal parameters, equal training tokens, or even equal FLOPs is systematically unfair because each choice favors different model families, and proposes that comparing at fixed wall-clock training time with an inference step-time constraint is the most holistic fairness criterion. The search phase operationalizes this: architectures that train faster get more steps within the budget and are rewarded for doing so. However, the paper then evaluates its final models at equal training steps (Table 3), not equal wall-clock time — a disconnect between the search objective and the evaluation protocol that the paper never resolves. This inconsistency weakens the paper's methodological contribution but also highlights a genuine tension: fixed-time comparisons are expensive to run for final evaluation (they require training all baselines for the same duration regardless of step count), and the field lacks standardized protocols for doing so. Future benchmarking efforts that adopt time-based or cost-based comparison frameworks would address a real gap this paper identifies but only partially fills.

The paper also reconciles a latent tension between the EfficientNet-style compound scaling philosophy (that architecture should be optimized jointly with scale) and the transformer community's uniform-block assumption. EfficientNet showed for CNNs that scaling width, depth, and resolution together — and doing so with per-layer heterogeneity — yields better Pareto curves than uniform scaling. The transformer community largely imported the scaling lesson (Chinchilla's joint scaling of parameters and data) but not the heterogeneity lesson. Brainformer demonstrates that the EfficientNet insight transfers to transformers, but through block composition rather than layer-wise coefficient multipliers. The implication is that transformer scaling laws (Kaplan et al., 2020; Hoffmann et al., 2022) should ideally be parameterized not just by model size and data quantity, but also by architectural heterogeneity — a larger, more heterogeneous model at the same activated parameter count may outperform a uniform model, and scaling studies that fix architecture miss this degree of freedom.

Research directions that become more attractive: heterogeneous block-level architecture search for transformers, joint optimization of gating and architecture in MoE models, time-aware NAS reward functions that reflect real-world cost constraints, and systematic study of which architectural regularities are necessary (and which are merely convenient) across model scales and domains.

Research directions that become less attractive: NAS within the uniform-block paradigm (the Primer approach is shown to leave significant efficiency on the table), and comparisons of gating mechanisms on single fixed architectures (the Brainformer evidence suggests such comparisons are of limited generalizability since the optimal gating depends on the architecture).

Follow-Up Research This Work Enables

Token-controlled downstream evaluation to isolate architectural quality from data volume effects. The paper's fine-tuning results (Table 4) compare models pre-trained for equal wall-clock time, meaning Brainformer processed more tokens than GLaM. A follow-up study should pre-train Brainformer and GLaM at the 1B and 8B scales on exactly the same number of tokens (matching Brainformer's throughput-advantaged token count to GLaM's, or vice versa), then evaluate on SuperGLUE and the five fewshot tasks from Table 5. This would cleanly separate the architectural quality contribution from the data volume contribution. If Brainformer's 3–4% SuperGLUE advantage persists under token-matched conditions, the architectural case is ironclad. If the gap closes substantially, it would reveal that the paper's headline downstream gains are primarily a throughput story (Brainformer is better because it lets you train on more data for the same cost, not because the architecture is inherently more capable per token). Either outcome is scientifically informative; the current paper conflates the two effects and the token-matched comparison would resolve the ambiguity.

Direct search at target scale to test the proxy-to-scale transferability assumption. The paper's core methodological assumption is that block architectures discovered at the 100M proxy scale transfer effectively to 1B and 8B. A follow-up should re-run the evolutionary search directly targeting the 1B scale (using a proportionally larger proxy model, or accepting the higher cost of searching at scale) and compare the discovered architecture against a Brainformer block scaled up from the 100M proxy following the paper's protocol. If the directly-searched 1B architecture substantially outperforms the scaled-up proxy architecture, it means the optimal block composition is scale-dependent and the paper's scaling protocol is leaving efficiency on the table. If the architectures converge to similar designs, the proxy-transfer assumption is validated and future work can confidently use small-scale search. The cost of such an experiment is high (likely thousands of TPU-hours), but the finding would determine whether Brainformer-style search is a one-time cost amortizable across all scales or must be repeated per target scale — a critical practical question the paper leaves open.

Cross-corpus and cross-domain transferability of the discovered block. The Brainformer block was discovered on GLaM's 1.6T-token English-heavy web corpus. A follow-up should take the exact Brainformer Block 1 architecture (Figure 8) and train it from scratch on a meaningfully different data distribution — for example, a code-heavy corpus (the Stack), a multilingual corpus (mC4), or a domain-specific scientific corpus (PubMed) — and compare its perplexity and downstream performance against a GLaM baseline trained on the same data for the same wall-clock time. This tests whether the discovered block composition (specific ratio of attention to dense FFN to MoE layers, chosen widths, Expert Choice routing with capacity factor 1) is an artifact of the GLaM data mixture or a genuinely general architectural improvement. If Brainformer transfers well, it becomes a reusable architectural template; if it degrades, then Brainformer-style search must be re-run per data distribution, substantially limiting the approach's practicality. The paper's Section 6.2 finding that layer ratios matter more than ordering suggests the block might be robust, but the data-dependence of the optimal ratio is entirely untested.

Quantitative characterization of the layer-order invariance hypothesis. Section 6.2 claims — without numerical evidence — that swapping adjacent layers within the Brainformer block does not affect quality, while changing the ratio of layer types does. A follow-up should systematically test this at the 100M scale: take Brainformer Block 1, generate all possible permutations of its 8 sub-layers (or a representative sample if 8! = 40,320 is infeasible), train each permutation from scratch for a fixed number of steps, and report perplexity distributions. The specific measurements of interest are: (a) the variance in final perplexity across permutations (quantifying how much ordering matters), (b) whether certain layer types have sensitive positions (e.g., is there a penalty for placing attention layers at the very beginning or end of the block?), and (c) whether the invariance holds at all scales or breaks down at larger model sizes. This experiment would convert an unsubstantiated claim into an established finding with clear practical implications: if ordering truly doesn't matter, practitioners can freely group same-type layers for hardware efficiency without quality loss; if it matters in specific ways (e.g., attention must not be the first layer), those rules become implementation guidelines.

Downstream evaluation at the 8B scale where efficiency gains are strongest. The paper's most dramatic results are at 8B/64E (5× step time, 1.99 vs. 2.12 perplexity), but all downstream evaluation (Tables 4, 5) is at 100M and 1B scales. A follow-up should fine-tune the 8B Brainformer and GLaM models on SuperGLUE and evaluate fewshot performance on the Table 5 tasks, following the same protocols used at smaller scales. The specific questions are: (a) does the 0.13 perplexity advantage at 8B translate to downstream gains, and if so, at what magnitude? (b) Do the task-level patterns observed at smaller scales (Brainformer's QQP deficit at 100M, Nqs deficit at 1B) persist, reverse, or disappear at 8B? and (c) Is the 5× step-time advantage maintained after fine-tuning, or does the fine-tuned model's computational profile change? This experiment is expensive but essential — without it, the paper's strongest claims are limited to pre-training, and a practitioner deciding whether to adopt Brainformer for a production 8B-scale model has no evidence about downstream behavior at that scale.

Comparison against a broader set of sparse baselines at matched activated parameters. The paper compares Brainformer only to GLaM among sparse architectures. A follow-up should train Switch Transformer, BASE Layers, and Expert Choice (on its own architecture, not just as a gating swap onto GLaM's) at the same activated parameter scales (100M, 1B, 8B) on the same 1.6T-token corpus under the same fixed-step or fixed-time protocol, and compare perplexity, step time, and downstream task performance. This would establish whether Brainformer's advantages are specific to the GLaM comparison or whether the heterogeneous block approach genuinely represents a state-of-the-art point on the MoE Pareto frontier. The paper's "Search-w-Top2" baseline partially addresses this by showing that architecture search with top-2 gating underperforms joint search, but the baseline is still Brainformer's own search procedure with a constraint — it's not an independent architecture from a different research group. External baselines would strengthen the generalizability claim considerably.

Time-aware benchmarks for efficient architecture comparison. The paper identifies a genuine methodological problem — existing benchmarks don't support fair comparison across model families with different throughput characteristics — but only partially solves it. A follow-up methodological contribution could develop standardized benchmarking protocols that report quality as a function of total wall-clock training time (not steps or tokens), enabling fair comparison between fast-low-capacity and slow-high-capacity architectures. Such a benchmark would need to specify: (a) a fixed hardware configuration for all comparisons, (b) reporting of both training time and inference throughput alongside quality metrics, (c) a standard way to trade off training cost against inference cost in the comparison. The Brainformer paper's search protocol provides a template but is too domain-specific (tied to the GLaM corpus and TPU-V4 hardware). A general-purpose time-aware benchmark would make the paper's methodological contribution actionable for the broader community.

Practical Applications and Downstream Use Cases

Cost-efficient training of large-scale MoE language models in industrial settings. The most direct application of this work is for organizations that regularly train large sparse language models, such as cloud providers building foundation models or large tech companies training internal LLMs. The paper provides a concrete recipe: run an evolutionary search at the 100M activated-parameter scale (which costs less than training a single 8B model from scratch) to discover a heterogeneous block architecture, then scale it to the target size by multiplying dimensions and stacking the block. At the 8B scale, the discovered Brainformer block delivers a 5× step-time improvement and 2× training convergence improvement over the GLaM baseline (Table 3) — meaning an organization that previously needed 512 TPU-V4 chips for 6 weeks to train an 8B-scale GLaM-quality model could train a Brainformer model of equivalent or better quality (1.99 PPLX vs. 2.12) in approximately 3 weeks on the same hardware, or train on half the hardware for the same duration. For a training run costing millions of dollars in cloud compute, the search cost (one week of 512 TPU-V4 chips) is small relative to the downstream savings. The paper's finding that suboptimal but useful architectures emerge early in the search ("within as early as 500 trials," Section 8) further reduces the barrier: organizations with tighter budgets can run a shorter search and still expect meaningful efficiency gains over a hand-designed GLaM-style block.

Inference-optimized deployment of MoE models in latency-sensitive production systems. The 5× step-time improvement at the 8B scale (1.96 vs. 0.39 steps/sec, Table 3) translates directly to inference throughput: a Brainformer model serving real-time queries can handle 5× the request volume of a GLaM model with the same activated parameter count on the same hardware. For a production system fielding millions of queries per day, this is the difference between running on a 512-TPU cluster and running on a ~100-TPU cluster, with proportional savings in hardware cost, energy consumption, and operational complexity. The Expert Choice gating with capacity factor 1 — which the search selected for Brainformer Block 1 — is particularly well-suited for inference because it guarantees perfect load balance without the overhead of auxiliary loss computation or the stochasticity of variable token-to-expert assignments. A team deploying an MoE model behind an API could use the Brainformer recipe specifically to optimize for inference throughput: run the search with a tight inference step-time constraint (Equation 5), select the block with the fastest step time among high-quality candidates (Block 1 rather than Block 2), and scale to the target size. The paper's Section 6.2 finding that layer ordering is largely irrelevant further simplifies deployment, since layers of the same type can be grouped contiguously to maximize hardware utilization without quality loss.

Architecture discovery for domain-specific or resource-constrained MoE models. While the paper focuses on NLP at large scale, the search methodology is domain-agnostic and can be applied to discover efficient architectures for settings where the standard uniform transformer block is known to be poorly matched to the constraints. Examples include: on-device models where the optimal attention-to-FFN ratio depends on the hardware's relative cost of matrix multiplications vs. attention operations; vision transformers where the spatial structure of image patches might benefit from different layer type distributions than text tokens; or multimodal models where different modalities share a common backbone but might benefit from different block compositions. In each case, a practitioner would define a search space analogous to Table 1 (replacing NLP-specific components with domain-appropriate primitives), run the evolutionary search at a small proxy scale under a fixed time budget, and scale the discovered block to the target model size. The key practical insight the paper provides for these applications is not the specific Brainformer block architecture (which may not transfer), but the methodology: (1) treat the block as a heterogeneous, searchable unit rather than a fixed sequence, (2) include gating/sparsity strategy in the search space jointly with architecture, (3) use a reward function that trades off quality against real-world cost (wall-clock time, inference latency), and (4) scale by dimension multiplication and block repetition rather than re-searching at each scale.