Structured Pruning Learns Compact and Accurate Models

ArXiv: 2204.00408

🎯 Pitch

Pruning BERT can match costly distillation—hitting 10×+ inference speedups with negligible accuracy loss, all without a scrap of unlabeled data. That’s the counterintuitive result of CoFi, which jointly lops off entire layers and individual heads by learning granular masks, demolishing the old tradeoff where pruned models stayed slow while distilled ones cost a fortune to train.

---

1. Executive Summary

This paper proposes CoFi (Coarse- and Fine-grained Pruning), a task-specific structured pruning method that jointly removes coarse-grained units (entire MHA or FFN layers) and fine-grained units (individual heads, intermediate FFN dimensions, hidden dimensions) from Transformer models, coupled with a dynamic layerwise distillation strategy that learns to match intermediate representations between the unpruned teacher and the pruning student without requiring a pre-specified architecture. Evaluated on GLUE and SQuAD v1.1 using BERT_base, CoFi delivers models with over 10× inference speedups and 95% sparsity while preserving more than 90% of the original accuracy, matching or exceeding the performance of distillation methods like TinyBERT_4—despite using no unlabeled data for general distillation and training in at most 20 GPU hours compared to TinyBERT's ~350 GPU hours—establishing that task-specific structured pruning can substitute for expensive general distillation only when coupled with multi-granularity masks and layerwise knowledge transfer.

2. Context and Motivation

The Deep Contradiction Between Pruning Speedups and Distillation Costs

The paper addresses a specific, practical contradiction in model compression: pruning methods achieve high compression ratios but cannot deliver competitive speedups, while distillation methods achieve excellent speedups but are prohibitively expensive to train. This gap is not incidental—it reflects fundamental architectural and procedural differences between the two approaches that the paper aims to resolve.

Consider what happens when a practitioner wants to deploy BERT_base (85M parameters) in a latency-sensitive application. They have two families of options. With pruning, they can identify and remove redundant weights from an already-trained BERT model, producing a smaller subnetwork. With distillation, they can train a compact model from scratch to mimic the larger teacher. The end goal—a smaller, faster model—is the same, but the tradeoffs diverge sharply along multiple axes.

The paper demonstrates this contradiction concretely in Table 1. Block Pruning, a state-of-the-art structured pruning method (Lagunas et al., 2021), achieves 2.7× speedup with 25M parameters and 83.7% MNLI accuracy. TinyBERT_4, a distillation approach, achieves 11.4× speedup with only 4.7M parameters at 78.8% accuracy. The speedup gap is enormous—more than 4×—despite pruning reducing the model to a similar parameter count. Meanwhile, from the other direction, Movement Pruning (Sanh et al., 2020) achieves 97% sparsity but yields zero speedup because its unstructured sparsity pattern cannot be exploited by current hardware. This is the core problem: pruning techniques can find highly sparse subnetworks, but the resulting architectures are not efficiently parallelizable on GPUs, while distillation designs architectures for speed but requires expensive multi-stage training on massive unlabeled corpora.

The Real-World Cost of General Distillation

To understand why this gap matters, we need to appreciate just how expensive general distillation is. The paper makes this concrete in Figure 1: TinyBERT's two-stage training pipeline first performs general distillation on 2,500 million tokens for 3 epochs—a process that takes approximately 350 GPU hours (3.5 days on 4 GPUs). Only then does task-specific distillation on labeled task data begin. For the four smaller GLUE datasets (CoLA: 8.5k examples, RTE: 2.5k, STS-B: 7k, MRPC: 3.7k), general distillation's cost completely dominates the training budget. Table 2 shows the consequences: TinyBERT_4 without general distillation collapses on these small datasets, achieving only 16.6 Matthews correlation on CoLA (vs. 32.5 with general distillation) and 17.8 Spearman correlation on STS-B (vs. 85.0 with general distillation).

This cost structure creates several practical barriers:

• Resource inequality: Only well-funded organizations can afford the 350 GPU hours of general distillation, limiting the deployment of compact models in low-resource settings.
• Iteration speed: When developing models for new domains or tasks, the three-day general distillation step makes rapid experimentation prohibitively slow. CoFi, by contrast, trains in at most 20 GPU hours on a single GPU—an 18× reduction in training time.
• Data dependency: General distillation requires a large, high-quality unlabeled corpus. For languages or domains where such corpora are unavailable, distillation methods become inapplicable, while task-specific pruning needs only the labeled task data.

The paper presents CoFi as a resolution to this contradiction: achieve the speedup advantage of distillation (10×+) with the data efficiency and lower computational cost of pruning. This is not a marginal improvement—it fundamentally changes the accessibility and practicality of extreme model compression.

The Fundamental Incompatibility: Why Prior Pruning Cannot Match Distillation's Speedups

To understand why existing pruning methods fail to achieve large speedups, we need to examine what actually happens on hardware when a pruned model runs inference. The paper identifies a crucial mismatch between compression rate (parameters removed) and inference speedup (reduction in wall-clock latency).

Coarse vs. Fine-Grained Unit Removal

Transformer models have a nested structure of computational units at different granularities:

• Coarse-grained units — entire layers. Removing an MHA layer eliminates all its heads simultaneously and skips the entire computation. Removing an FFN layer eliminates its up-projection, non-linearity, and down-projection as one block. These removals map directly to GPU kernels that can be skipped entirely, yielding near-proportional speedups.
• Fine-grained units — individual attention heads or intermediate FFN dimensions. Removing a head shrinks some matrix dimensions but requires implementing irregular sparsity patterns where, say, heads 2, 5, 7, and 11 remain while others are gone. GPUs achieve peak efficiency with dense, regular matrix operations—sparse, irregular patterns leave them underutilized.

The paper observes that prior structured pruning has trended toward increasingly fine-grained units, motivated by the desire for flexible final structures. Head pruning (Voita et al., 2019; Michel et al., 2019), FFN dimension pruning (McCarley et al., 2019; Wang et al., 2020b), and block pruning (Lagunas et al., 2021) all remove progressively smaller components. The intuition is reasonable: smaller pruning units give the optimization more degrees of freedom to find the best subnetwork. However, the paper identifies a critical optimization failure: fine-grained pruning alone rarely removes entire layers organically. As noted in Section 3.1:

"pruning fine-grained units naturally entails pruning coarse-grained units—for example, pruning $N_h$ (e.g., 12) heads is equivalent to pruning one entire MHA layer. However, we observe that this rarely happens in practice and poses difficulty to optimization especially at a high sparsity regime."

Think about why this happens. When the optimizer learns to prune individual heads, it assigns importance scores to each head independently. To prune an entire MHA layer, all 12 heads must simultaneously have their importance driven to zero. Each head individually contributes some marginal value, and the optimization has no explicit incentive to zero them all out at once—the gradient signal treats each head independently, and the sparsity constraint can be satisfied by pruning a few heads across many layers rather than completely eliminating any single layer.

The consequence is stark and measurable. Table 4 shows the ablation: when CoFi's explicit layer masks ($z_{\text{MHA}}$, $z_{\text{FFN}}$) are removed, a 95%-sparse QNLI model drops from 12.1× speedup to only 7.2× speedup, despite having the same parameter count (~5M). The difference is entirely architectural: without the ability to explicitly prune entire layers, the optimizer produces deep-but-thin networks where every layer contributes some (small) computation, and the GPU cannot skip any of them. With layer masks, entire MHA or FFN layers are zeroed out, producing shallow-but-wide subnetworks where entire computational blocks can be bypassed, yielding the additional speedup.

The Granularity-Speed Tradeoff Formalized

This reveals a fundamental tension: fine-grained pruning optimizes accuracy by preserving maximum representational flexibility, but coarse-grained pruning optimizes speed by enabling structural simplifications that hardware can exploit. The paper's key insight is that these objectives need not be in conflict if you allow both levels of pruning simultaneously, with the optimization having explicit knobs to remove entire layers when beneficial.

Where Each Prior Family of Methods Falls Short

The paper catalogs four main families of prior work, each with specific failure modes that CoFi addresses:

1. Layer pruning alone is too coarse

Fan et al. (2020) and Sajjad et al. (2020) explore dropping entire Transformer blocks and find empirical support that ~50% of layers can be removed. The resulting 2× speedup is useful but represents a hard ceiling—once you've removed all redundant layers, you can't compress further by removing anything larger than a neuron. Layer pruning gives you speed but no flexibility.

2. Head pruning achieves minimal speedup

Head pruning methods (Voita et al., 2019; Michel et al., 2019) demonstrate that most heads are redundant, but Li et al. (2021) show only 1.4× speedup with one remaining head per layer. The reason: even with one head remaining per layer, the GPU must still execute all 12 layers, just with smaller attention matrices. The layers themselves, and their associated overhead (layer norm, residual connections), persist.

3. FFN dimension pruning alone cannot match distillation

FFN layers account for 2/3 of Transformer parameters, making them a prime target. McCarley et al. (2019) and Hou et al. (2020) prune intermediate dimensions by introducing binary masks $z_{\text{int}} \in \{0, 1\}^{d_f}$ that selectively zero out neurons in the FFN's hidden representation. At high sparsity, this produces thin feed-forward layers within deep networks—again, the depth (and its associated overhead) persists. As shown in CoFi's ablation (Table 4, "–layer & hidden": 7.2× speedup at 95% sparsity), without explicit layer removal, speedups stall well below the 10×+ threshold that distillation achieves.

4. Unstructured pruning achieves zero speedup on current hardware

Movement Pruning (Sanh et al., 2020) reaches 97% sparsity but produces irregular, individual-weight sparsity patterns. GPUs are designed for dense matrix multiplication—the standard operations (SGEMM, etc.) expect contiguous blocks of values in memory. Unstructured sparsity requires specialized sparse matrix kernels (like the block-sparse MatMul in Triton from Tillet et al., 2019), but as Yao et al. (2021) show, the reported results for block-sparse inference are not yet competitive. The paper notes that dedicated hardware platforms like Moffett AI's ANTOM and software packages like DeepSparse are emerging, but the evaluation environments differ, and these are not yet standard deployment targets. For the V100 GPUs used in the paper's evaluation (and most current deployments), unstructured pruning provides no practical speed benefit.

The Missing Piece: Why Distillation Learning Cannot Be Trivially Applied to Pruning

Even if we could architecturally prune to match distillation's speed, there remains a second challenge: knowledge transfer from the unpruned teacher to the pruning student. Distillation methods know the student architecture in advance, so they can pre-define layer mappings—for example, TinyBERT_4 distills from BERT_base's 3rd, 6th, 9th, and 12th layers to its own 1st, 2nd, 3rd, and 4th layers. The layer correspondence is fixed before training begins.

In contrast, pruning doesn't know the final subnetwork structure during training. The architecture is evolving—layers are being dropped, dimensions are shrinking, and the mapping between teacher and student layers is a moving target. Section 3.2 frames this explicitly:

"distilling intermediate layers during the pruning process is challenging as the model structure changes throughout training."

Prior pruning work (Sanh et al., 2020; Lagunas et al., 2021) sidesteps this entirely by using only prediction-layer distillation—matching output probability distributions between teacher and student. This transfers task-level knowledge but loses the rich intermediate representations that make layerwise distillation so effective (as demonstrated by Jiao et al., 2020; Sun et al., 2020). CoFi's dynamic layer mapping (Section 3.2) proposes a solution: at each training step, match each teacher layer to its closest surviving student layer by minimizing representation distance. This allows the distillation to track the evolving architecture without requiring pre-specification.

The ablation in Table 5 quantifies the importance: removing layer distillation (retaining only prediction-layer distillation) drops QNLI accuracy from 86.1% to 85.1% and SQuAD F1 from 82.6% to 82.5% at 95% sparsity. Removing all distillation drops QNLI to 84.2% and SQuAD to 75.8%. The chapter is larger at lower sparsities (Table 11), suggesting layerwise knowledge transfer becomes more important when pruning is more aggressive.

FFN vs. MHA: Asymmetric Redundancy

The paper's structural analysis (Figure 3) reveals a pattern that informs the motivation: FFN layers are substantially more redundant than MHA layers. At 60% sparsity, the average number of intermediate FFN dimensions drops by 71% (from 3,072 to 884), while the average number of heads drops by only 39% (from 12 to 7.3). At 95% sparsity (Table 6, Figure 5 in Appendix H), pruned models consistently retain more MHA layers than FFN layers. The paper interprets this as evidence that MHAs are more critical for downstream task performance.

This asymmetry matters because FFN layers account for 2/3 of model parameters. The fact that they can be more aggressively pruned without proportional accuracy loss means that targeted coarse pruning of FFN layers can yield disproportionate speed and parameter reductions. A method that only prunes heads (fine-grained) would leave FFN redundancies largely untapped. A method that only prunes entire blocks would not discriminate—it would prune MHAs and FFNs equally. CoFi's independent masks for MHA layers ($z_{\text{MHA}}$) and FFN layers ($z_{\text{FFN}}$) capture this asymmetry explicitly, allowing the optimization to eliminate FFN layers more aggressively while preserving MHA layers.

How the Paper Positions Itself

The paper positions CoFi as occupying a previously empty intersection in the pruning-distillation design space (Table 1):

• Like pruning methods, CoFi does not require unlabeled data ($U$ column: ✗) and uses only task-specific data ($T$ column: ✓).
• Like distillation methods, CoFi achieves large speedups (10×+), matching or exceeding TinyBERT_4.
• Unlike all prior work, CoFi combines explicit multi-granularity pruning masks with dynamic layerwise distillation, solving both the architectural flexibility problem and the knowledge transfer problem simultaneously.

The paper is careful not to claim that CoFi makes distillation obsolete. Rather, the motivation is that task-specific structured pruning, done right, can substitute for expensive general distillation while retaining its speed advantages. The key word is "done right"—the contribution is not just a new method, but a demonstration that prior pruning's failure to achieve large speedups was not inherent to pruning as a paradigm but rather a consequence of specific design choices (monogranularity masks, prediction-only distillation) that CoFi changes.

The opening of Section 6 (Conclusion) crystallizes this positioning:

"We conclude that task-specific structured pruning from large-sized models could be an appealing replacement for distillation to achieve extreme model compression, without resorting to expensive pre-training or data augmentation."

The framing is deliberately comparative and pragmatic: CoFi is an "appealing replacement" because it delivers the same speedup-performance tradeoff at a fraction of the training cost. This shifts the cost-performance frontier for model compression, making extreme speedups accessible to practitioners who cannot afford the computational overhead of general distillation.

3. Technical Approach

3.1 Reader Orientation

CoFi is a structured pruning system that discovers highly-efficient subnetworks within a pre-trained BERT model by learning which computational units (layers, heads, neurons, hidden dimensions) to keep and which to remove, guided by knowledge distillation from the original unpruned model. The core problem CoFi solves is that prior pruning methods achieve high compression ratios but cannot deliver competitive inference speedups, while distillation methods achieve excellent speedups but require expensive multi-stage training on massive unlabeled corpora — CoFi bridges this gap by jointly learning multi-granularity pruning masks with a dynamic layerwise distillation strategy that requires only task-specific labeled data.

3.2 Big-Picture Architecture (Diagram in Words)

The CoFi system has four major components that interact during a single training phase:

1. Masked BERT Model — the pre-trained BERT_base model whose parameters are wrapped with five types of learned binary masks that control which computational units survive pruning. This is the "student" — the model being compressed.

2. Frozen Teacher BERT Model — the original fine-tuned BERT_base model with no masks, providing target outputs for knowledge distillation. Its weights are held constant throughout CoFi training.

3. Distillation Pipeline — two parallel knowledge transfer paths: (a) prediction-layer distillation that matches the output probability distributions of the teacher and the masked student, and (b) dynamic layerwise distillation that matches intermediate hidden representations between teacher layers and their closest surviving student layers, with the layer mapping adaptively updated during training.

4. Sparsity Controller — a Lagrangian optimization loop that imposes an equality constraint forcing the expected model sparsity (computed from the learned mask probabilities) to reach a target sparsity level (e.g., 60%, 95%), while the model simultaneously minimizes task loss and distillation losses.

Information flows as follows: a labeled training batch enters both the frozen teacher and the masked student → the student computes task predictions with masks applied → the teacher provides soft targets for prediction-layer distillation → the student's intermediate layer representations are dynamically matched to the teacher's selected layers for layerwise distillation → the sparsity controller computes the expected sparsity from current mask probabilities and adds a Lagrangian penalty if the target is not met → gradients flow back through the masks (via hard concrete reparameterization) and the model weights, updating both simultaneously. After training converges, masks are binarized by thresholding, producing the final pruned architecture.

3.3 Roadmap for the Deep Dive

• First, the five types of pruning masks — layer-level ($z_{\text{MHA}}$, $z_{\text{FFN}}$), head-level ($z_{\text{head}}$), FFN-dimension-level ($z_{\text{int}}$), and hidden-dimension-level ($z_{\text{hidn}}$) — because they are the architectural foundation that enables the speed-accuracy tradeoff. Understanding how these masks compose multiplicatively to determine whether any given parameter survives is essential.

• Second, the mask learning mechanism — the hard concrete distribution reparameterization, the Lagrangian sparsity controller, and the expected sparsity formula — because this is what makes the masks differentiable and tunable to an exact target sparsity. Without this, the optimizer cannot explicitly trade off pruning decisions against task performance.

• Third, the dynamic layerwise distillation objective — the layer mapping function $m(\cdot)$, the hidden representation matching loss, and how it is combined with prediction-layer distillation — because this is the knowledge transfer mechanism that maintains accuracy under aggressive pruning when the architecture is unknown in advance.

• Fourth, the training procedure — the two-phase schedule (distillation warmup followed by pruning with linear sparsity scheduler), the finetuning phase, and the threshold-based binarization at inference — because these procedural choices determine whether the masks converge to useful structures compatible with the distillation objectives.

• Fifth, the design rationale — why each specific choice (Lagrangian over vanilla $l_0$, dynamic mapping over fixed mapping, separate MHA/FFN layer masks over whole-block dropping, etc.) was made over the alternatives discussed or implied in the paper.

3.4 Detailed, Sentence-Based Technical Breakdown

This is primarily a systems and methods paper whose core idea is that the speedup gap between pruning and distillation can be closed by simultaneously learning pruning masks at multiple granularities (so that entire layers can be explicitly removed when beneficial) while using dynamic layerwise distillation to transfer knowledge from the unpruned model despite the evolving architecture.

---

The Five Types of Pruning Masks

CoFi wraps the weights of a standard BERT_base Transformer with five independent sets of learnable mask variables, each controlling a different level of computational granularity. The critical property — and the key difference from prior work — is that multiple masks jointly determine whether any single parameter participates in computation. A weight element in an FFN layer, for example, is effectively pruned if its FFN layer mask is zero, OR if its intermediate dimension mask is zero, OR if its hidden dimension mask is zero. This logical-OR composition means the optimization can achieve sparsity through whichever pathway is most beneficial for the current architecture and task.

MHA layer mask ($z_{\text{MHA}}$). For each of the $L = 12$ Transformer blocks, there is a scalar mask $z_{\text{MHA}}^{(i)} \in \{0, 1\}$ that controls whether the entire multi-head self-attention sublayer in block $i$ is active. When $z_{\text{MHA}}^{(i)} = 0$, all heads in that layer are bypassed — the attention computation, including all query/key/value/output projections, is skipped entirely, and only the residual connection and layer normalization remain.

FFN layer mask ($z_{\text{FFN}}$). Similarly, for each block $i$, $z_{\text{FFN}}^{(i)} \in \{0, 1\}$ controls whether the entire feed-forward sublayer (up-projection, GELU activation, down-projection) is active. When $z_{\text{FFN}}^{(i)} = 0$, the FFN computation is skipped entirely.

These two layer masks are learned independently — it is possible (and common in CoFi's pruned models) for a block to retain its MHA layer but drop its FFN layer, or vice versa. This contrasts with Fan et al. (2020) and Sajjad et al. (2020), who drop entire Transformer blocks (MHA + FFN together). The paper's structural analysis in Table 6 shows that different datasets produce different patterns: for example, on SST-2, some pruned models have sequences like "M F M M F M F M F M F" where M and F interleave (preserving spurts of both sublayer types), while on QQP, models often drop the first MHA layer but keep several subsequent ones. Independent control enables the optimization to discover these dataset-specific patterns.

Head mask ($z_{\text{head}}$). For each attention head $j \in \{1, \dots, N_h\}$ (where $N_h = 12$ for BERT_base) in each block $i$, $z_{\text{head}}^{(i,j)} \in \{0, 1\}$ controls whether that specific head's computation is active. When $z_{\text{head}}^{(i,j)} = 0$, the corresponding query, key, value, and output projection columns for that head are effectively zeroed out. Heads within a layer can be pruned independently, so a layer might retain 3 out of 12 heads.

The masked multi-head attention operation becomes:

$\text{MHA}(X) = z_{\text{MHA}} \cdot \sum_{i=1}^{N_h} z_{\text{head}}^{(i)} \cdot \text{Att}(W_Q^{(i)}, W_K^{(i)}, W_V^{(i)}, W_O^{(i)}, X)$

where $W_Q^{(i)}, W_K^{(i)}, W_V^{(i)}, W_O^{(i)} \in \mathbb{R}^{d \times d_h}$ are the query, key, value, and output projection matrices for head $i$, $d = 768$ is the hidden size, $d_h = d / N_h = 64$ is the per-head dimension, and $\text{Att}(\cdot)$ is the standard scaled dot-product attention function.

What this computes: The outer product $z_{\text{MHA}} \cdot \sum_{i=1}^{N_h} z_{\text{head}}^{(i)} \cdot \text{Att}(\dots)$ means that if the MHA layer mask is zero, the entire attention output is multiplied by zero regardless of individual head masks; if the layer mask is one but a specific head's mask is zero, that head's attention output is excluded from the sum; only heads where both the layer mask and the individual head mask are non-zero contribute to the final attention output. The result is the standard MHA output with selected heads and possibly the entire layer zeroed out.

Why this form: The multiplicative composition ensures that the layer mask acts as a master switch — setting $z_{\text{MHA}}^{(i)} = 0$ prunes all computational work associated with attention in layer $i$ at once, which maps directly to skipping GPU kernels. Without the layer mask, the optimization would need to independently drive all 12 head masks to zero to achieve the same effect, which the paper observes "rarely happens in practice" because each head independently contributes some marginal value and the gradient signal has no explicit incentive to coordinate them.

FFN intermediate dimension mask ($z_{\text{int}}$). For each FFN layer $i$, there is a vector mask $z_{\text{int}}^{(i)} \in \{0, 1\}^{d_f}$ where $d_f = 3072$ (the standard $4d$ expansion factor for BERT_base). Each element controls whether the corresponding neuron in the FFN's hidden representation is active. The masked FFN operation becomes:

$\text{FFN}(X) = z_{\text{FFN}} \cdot \text{gelu}(XW_U) \cdot \text{diag}(z_{\text{int}}) \cdot W_D$

where $W_U \in \mathbb{R}^{d \times d_f}$ is the up-projection matrix, $W_D \in \mathbb{R}^{d_f \times d}$ is the down-projection matrix, and $\text{diag}(z_{\text{int}})$ creates a diagonal matrix from the mask vector, which zeroes out specific columns of the GELU output before the down-projection.

What this computes: The GELU activation is first applied to the up-projected hidden state $XW_U$, producing a vector of length $d_f$. The diagonal mask $\text{diag}(z_{\text{int}})$ then elementwise multiplies this vector, zeroing out positions where $z_{\text{int}}$ is zero. The down-projection $W_D$ then maps the surviving dimensions back to hidden size $d$. The outer factor $z_{\text{FFN}}$ can additionally zero out the entire output if the FFN layer is pruned.

Why this form: Placing the mask after the GELU non-linearity (rather than before the up-projection or after the down-projection) prunes the intermediate representation — it removes the computation of the up-projection for pruned dimensions (since their outputs are zeroed before the expensive $W_D$ multiplication) AND removes the corresponding columns of $W_D$ from contributing. This is the standard approach from McCarley et al. (2019) and Hou et al. (2020), and it produces a structurally thinner FFN layer that can be implemented as a smaller dense matrix multiplication.

Hidden dimension mask ($z_{\text{hidn}}$). This is a single mask vector $z_{\text{hidn}} \in \{0, 1\}^d$ where $d = 768$, shared across all layers in the model. It controls which dimensions of the residual stream (the hidden representation) are active. This mask is applied to all weight matrices throughout the model by left-multiplying or right-multiplying with $\text{diag}(z_{\text{hidn}})$ as appropriate. For example, the query projection becomes $\text{diag}(z_{\text{hidn}}) W_Q$, which zeroes out entire rows of each weight matrix corresponding to pruned hidden dimensions.

The paper notes that this mask is surprisingly impactful: empirically, only a small number of hidden dimensions are actually pruned (e.g., "768 → 760"), but even this small reduction "helps improve performance significantly" (§3.1, §4.3). The reason is that the shared mask creates cross-layer dependencies — pruning a hidden dimension removes that feature pathway through the entire residual stream. This imposes a form of structured regularization that prevents the model from routing information through narrow, fragile pathways.

Why shared across layers: The residual connection in Transformers means that the hidden representation at layer $\ell$ is elementwise-added to the hidden representation at layer $\ell+1$. If dimension $k$ were pruned at layer $\ell$ but not at layer $\ell+1$, the residual addition would mix pruned and unpruned dimensions, creating an ill-defined operation. Sharing the mask ensures consistency: a pruned dimension is removed from all computations, and the residual connections simply pass zeros through those dimensions.

The compositional sparsity logic. The key design principle is that these five mask types compose through logical OR: a parameter is pruned if any mask that governs it is zero. For a weight in the FFN up-projection matrix $W_U$, three masks jointly determine its survival: $z_{\text{FFN}}$ (is the entire FFN layer active?), $z_{\text{int}}$ (is this specific intermediate neuron active?), and $z_{\text{hidn}}$ (is the input hidden dimension for this row active?). If any of these three is zero, the weight is effectively pruned.

This compositional structure is what enables CoFi to achieve both high sparsity and high speedup: the optimizer can satisfy the sparsity constraint through coarse decisions (removing entire layers) when that is efficient, and through fine decisions (removing specific heads or neurons) when that preserves more accuracy. Prior methods restricted to a single granularity forced the optimizer into suboptimal tradeoffs.

---

Mask Learning: The Hard Concrete Distribution and Lagrangian Sparsity Control

The pruning masks cannot be learned by standard gradient descent because they are binary decisions — a parameter is either pruned (mask = 0) or kept (mask = 1). The paper uses the hard concrete distribution (Louizos et al., 2018) to create a differentiable relaxation: during training, masks are continuous values in $[0, 1]$ sampled from a distribution whose parameters are learned; at inference, they are thresholded to exactly 0 or 1.

The reparameterization procedure. A mask value $z$ is generated from a set of learnable parameters $\log \alpha$ (one per mask variable) through the following sampling process, which is fully differentiable via the reparameterization trick:

$\mathbf{u} \sim U(0, 1)$

$\mathbf{s} = \text{sigmoid} \left( \frac{1}{\beta} \log \frac{\mathbf{u}}{1 - \mathbf{u}} + \log \alpha \right)$

$\tilde{\mathbf{s}} = \mathbf{s} \times (r - l) + l$

$\mathbf{z} = \min(1, \max(0, \tilde{\mathbf{s}}))$

where $U(0, 1)$ is a uniform distribution, $\beta$ is a temperature hyperparameter controlling the steepness of the sigmoid (lower $\beta$ makes the transition sharper, approximating a step function), $l < 0$ and $r > 0$ are constants that stretch the sigmoid output into the interval $(l, r)$, and $\log \alpha$ are the main learnable parameters updated by gradient descent.

What this computes: First, a uniform random sample $u$ is drawn. This is transformed through an inverse-sigmoid-like operation to produce $s$, a value whose location is controlled by $\log \alpha$ and whose sharpness is controlled by $\beta$. As $\log \alpha$ becomes more negative, the probability of $s$ being near zero increases — corresponding to a higher probability of pruning. The stretching operation maps $s$ from approximately $(0, 1)$ to $(l, r)$, which is then clamped to $[0, 1]$ to produce the final continuous mask $z$. During training, $z$ is used as a soft mask (multiplying activations by values between 0 and 1). At inference, after training converges, the mask is binarized: values below a threshold (determined by the expected sparsity of each weight matrix, see Appendix B) become 0, and values above become 1.

Why this reparameterization: The hard concrete distribution has the property that it can produce values that are exactly 0 or exactly 1 with non-zero probability (due to the clamping), which matches the discrete nature of pruning decisions, while maintaining a differentiable path from the loss to the parameters $\log \alpha$ through the sampling process. This enables gradient-based optimization over discrete architectural choices. The paper also notes in a footnote that they tried a straight-through estimator (as used in Sanh et al., 2020) and found performance comparable, but chose the $l_0$ approach because "it is easier to control the sparsity precisely."

Lagrangian sparsity control. In preliminary experiments, the authors found that directly optimizing the $l_0$ objective $\lambda \|z\|_0$ with different learning rates and pruning schedules could converge to models of drastically different sizes. To enforce precise control over the final sparsity, they adopt the Lagrangian multiplier approach from Wang et al. (2020b). The constraint is formulated as an equality $\hat{s} = t$, where $\hat{s}$ is the expected model sparsity (computed continuously from the current mask probabilities) and $t$ is the target sparsity (e.g., 0.60 for 60% pruning). The Lagrangian penalty is:

$\mathcal{L}_c = \lambda_1 \cdot (\hat{s} - t) + \lambda_2 \cdot (\hat{s} - t)^2$

where $\lambda_1$ and $\lambda_2$ are Lagrange multipliers that are updated during training to enforce the constraint (following standard Lagrangian optimization: the multipliers increase when the constraint is violated and decrease when it is satisfied).

What this computes: The linear term $\lambda_1 \cdot (\hat{s} - t)$ pushes the expected sparsity toward the target with a force proportional to $\lambda_1$; the quadratic term $\lambda_2 \cdot (\hat{s} - t)^2$ penalizes large deviations more heavily. Together, they implement a soft equality constraint: the optimizer can temporarily violate the target sparsity (allowing exploration of different architectures) but is increasingly penalized the further it deviates and the longer the deviation persists (as $\lambda_1, \lambda_2$ adaptively increase).

Why Lagrangian over standard $l_0$: Standard $l_0$ regularization adds a term $\lambda \cdot \|z\|_0$ to the loss, where $\lambda$ is a fixed hyperparameter. The final sparsity depends on the balance between the task loss gradient and the regularization gradient — changing the learning rate, batch size, or dataset can shift this balance, making sparsity unpredictable. The Lagrangian formulation directly specifies the desired sparsity as a constraint, and the multipliers automatically adapt to achieve it regardless of other hyperparameters. This gives the paper the ability to sweep target sparsities $\{60\%, 70\%, 75\%, 80\%, 85\%, 90\%, 95\%\}$ and reliably obtain models at those exact sparsity levels for fair comparison.

The expected sparsity formula. To compute $\hat{s}$ (the expected fraction of parameters pruned), the paper must account for the compositional masking structure. A parameter is considered surviving only if ALL masks that govern it are non-zero. The expected sparsity is computed as:

$\hat{s} = \frac{1}{M} \cdot 4 \cdot d_h \cdot \sum_i^L \sum_j^{N_h} \sum_k^d z_{\text{MHA}}^{(i)} \cdot z_{\text{head}}^{(i,j)} \cdot z_{\text{hidden}}^{(k)} + \frac{1}{M} \cdot 2 \cdot \sum_i^L \sum_j^{d_f} \sum_k^d z_{\text{FFN}}^{(i)} \cdot z_{\text{int}}^{(i,j)} \cdot z_{\text{hidden}}^{(k)}$

where $M$ is the total number of parameters in the full model (embeddings excluded), $L = 12$ is the number of Transformer blocks, $N_h = 12$ is the number of heads per layer, $d_f = 3072$ is the FFN intermediate dimension, $d = 768$ is the hidden size, and $d_h = d / N_h = 64$ is the per-head dimension.

What this computes: The first double sum counts surviving parameters in MHA layers: the factor $4 \cdot d_h$ accounts for the four projection matrices (Q, K, V, O) each contributing $d_h$ parameters per head per hidden dimension, and the triple product $z_{\text{MHA}}^{(i)} \cdot z_{\text{head}}^{(i,j)} \cdot z_{\text{hidden}}^{(k)}$ is non-zero only if the layer, that specific head, AND that specific hidden dimension are all active. The second double sum counts surviving parameters in FFN layers: the factor $2$ accounts for the up-projection and down-projection matrices, and the triple product $z_{\text{FFN}}^{(i)} \cdot z_{\text{int}}^{(i,j)} \cdot z_{\text{hidden}}^{(k)}$ is non-zero only if the FFN layer, that specific intermediate neuron, AND that hidden dimension are all active. Summing over all layers, heads/neurons, and hidden dimensions gives the expected number of surviving parameters; dividing by $M$ and subtracting from 1 (implicitly) gives the expected sparsity.

Why this particular bookkeeping: The formula directly mirrors the compositional masking logic — a parameter survives only if every mask that touches it is 1. The separate treatment of MHA and FFN parameters with different coefficients ($4 \cdot d_h$ vs. $2$) correctly accounts for the different numbers of weight matrices involved. This formula is what the Lagrangian controller uses to compute the constraint violation $\hat{s} - t$, making it the quantitative bridge between the learned mask probabilities and the target compression level.

---

Dynamic Layerwise Distillation

Knowledge distillation transfers information from the frozen teacher model (the original fine-tuned BERT_base) to the masked student model during training. CoFi uses two distillation objectives: prediction-layer distillation (matching output distributions) and a novel dynamic layerwise distillation (matching intermediate representations from selected teacher layers to dynamically identified student layers).

Prediction-layer distillation. This is the standard knowledge distillation loss from Hinton et al. (2015), adapted for classification tasks. The teacher model processes the input and produces a probability distribution $\mathbf{p}_t$ over output classes; the student model produces $\mathbf{p}_s$. The loss is:

$\mathcal{L}_{\text{pred}} = D_{\text{KL}}(\mathbf{p}_s \parallel \mathbf{p}_t)$

where $D_{\text{KL}}$ is the Kullback-Leibler divergence.

What this computes: $D_{\text{KL}}(\mathbf{p}_s \parallel \mathbf{p}_t) = \sum_c \mathbf{p}_s(c) \log \frac{\mathbf{p}_s(c)}{\mathbf{p}_t(c)}$, measuring how much information is lost when using $\mathbf{p}_s$ to approximate $\mathbf{p}_t$. When combined with a temperature parameter $t = 2$ (as specified in Appendix A, Table 7) that softens both distributions before computing the KL divergence, this encourages the student to match not just the teacher's hard predictions but the full relative confidence across all classes — including the teacher's uncertainty about near-miss answers.

Why not only hard-label cross-entropy: Knowledge distillation provides a richer training signal than the ground-truth labels alone because the teacher's probability distribution encodes relationships between classes (e.g., that "positive" and "neutral" sentiment are more similar than "positive" and "negative"). This is especially important during pruning, where the student capacity is being reduced — the soft targets guide the student toward the teacher's generalization patterns even when the student lacks the capacity to exactly replicate the teacher's decision boundary.

The layer mapping problem. For distillation methods like TinyBERT, which pre-specify the student architecture, it is straightforward to define which teacher layers should supervise which student layers — e.g., teacher layers 3, 6, 9, 12 map to student layers 1, 2, 3, 4. For pruning, the student architecture is unknown in advance and changes during training. CoFi addresses this with a dynamic layer mapping function $m(\cdot)$ that is recomputed during training based on the current state of the student model.

The layer distillation loss. The paper selects a set of teacher layers $\mathcal{T}$ to distill from (the paper does not explicitly list $\mathcal{T}$ but from context and the results in §G.1, it appears to be a sparse subset — likely corresponding to the last layer of each "stage" in BERT, e.g., layers 3, 6, 9, 12, following TinyBERT's convention). For each teacher layer $i \in \mathcal{T}$, the loss matches its hidden representation to the closest surviving student FFN layer:

$\mathcal{L}_{\text{layer}} = \sum_{i \in \mathcal{T}} \text{MSE}(W_{\text{layer}} \mathbf{H}_s^{m(i)}, \mathbf{H}_t^i)$

where $\mathbf{H}_t^i$ is the hidden representation (output of the FFN sublayer) from teacher layer $i$, $\mathbf{H}_s^{m(i)}$ is the hidden representation from the dynamically matched student layer $m(i)$, and $W_{\text{layer}} \in \mathbb{R}^{d \times d}$ is a learned linear transformation matrix initialized as the identity matrix.

What this computes: For each selected teacher layer $i$, the system identifies which surviving student layer $m(i)$ is closest (in MSE distance) to that teacher layer. It then projects the student layer's representation through $W_{\text{layer}}$ to align it with the teacher's representation, and computes the mean squared error between them. The sum over $i \in \mathcal{T}$ aggregates these per-layer losses. The linear transformation $W_{\text{layer}}$ gives the student the flexibility to transform its representation to match the teacher's, which is necessary because the student may have a different number of surviving dimensions or a different representational structure.

Why MSE rather than KL: Hidden representations are continuous vectors in $\mathbb{R}^d$, not probability distributions. MSE directly penalizes Euclidean distance between the projected student representation and the teacher representation, which is the natural metric for continuous feature matching. KL divergence would require converting representations to probability distributions (e.g., via softmax), which discards magnitude information that may be important for downstream layers.

The dynamic matching function. The layer mapping $m(\cdot)$ is recomputed during training by finding, for each teacher layer $i$, the surviving student layer $j$ whose representation is closest in MSE:

$m(i) = \arg \min_{j: z_{\text{FFN}}^{(j)} > 0} \text{MSE}(W_{\text{layer}} \mathbf{H}_s^j, \mathbf{H}_t^i)$

where the constraint $z_{\text{FFN}}^{(j)} > 0$ restricts the search to student FFN layers that have not been pruned (since a pruned layer produces zero output and cannot provide a meaningful representation). The minimization is over the surviving student layer index $j$.

What this computes: At each training step (or periodically), for each teacher layer $i$ in the selected set, the system computes the MSE between the teacher's representation at layer $i$ and every surviving student layer's representation (after applying the learned projection $W_{\text{layer}}$). It selects the student layer $j$ that minimizes this MSE. This mapping can change during training — as layers are pruned, some student layers become unavailable, and as representations evolve, a previously matched layer may become less similar than another candidate.

Why dynamic rather than fixed mapping: Fixed mapping (e.g., matching teacher layer 3 → student layer 1, teacher layer 6 → student layer 2) fails when pruning removes a student layer that was expected to match a specific teacher layer — the distillation signal from that teacher layer is simply lost. The ablation in Table 5 confirms this: "Fixed Hidn Distil.," which simply matches each teacher layer to the corresponding student layer index (skipping if the student layer is pruned), underperforms dynamic matching on QNLI (85.8 vs. 86.1), MNLI (80.5 vs. 80.6), SQuAD (80.9 vs. 82.6), and SST-2 (90.0 vs. 90.6). The dynamic approach ensures that every selected teacher layer always has a distillation target, and that the target is the most representationally similar surviving layer, which is more informative than a position-based match.

Layer ordering constraint for small datasets. For small datasets (RTE: 2.5k examples, MRPC: 3.7k), the paper observes that layer mismatch can occur — the dynamic matching might map a later teacher layer to an earlier student layer index, creating non-monotonic alignment. To address this, they add a constraint that a teacher layer can only be matched to a student layer with a lower index than the previously matched teacher layer's student. This enforces a monotonic alignment (teacher layers are matched to student layers in order) and stabilizes training on data-scarce settings. The paper notes that for larger datasets, "layer mismatch rarely happens, showing the superiority of dynamic matching — layers between student and teacher models match in a way that benefits the pruning process the most."

Combined distillation objective. The prediction-layer and layerwise distillation losses are combined with a weighting parameter $\lambda$:

$\mathcal{L}_{\text{distil}} = \lambda \mathcal{L}_{\text{pred}} + (1 - \lambda) \mathcal{L}_{\text{layer}}$

where $\lambda \in \{0.1, 0.3, 0.5\}$ is swept as a hyperparameter (Appendix A, Table 7). The paper does not report which $\lambda$ was used for each experiment, but Tables 2 and 3 show final results presumably using the best $\lambda$ per dataset.

What this computes: The total distillation loss is a convex combination of the output-level and intermediate-level knowledge transfer. When $\lambda = 0.5$, both losses contribute equally; when $\lambda = 0.9$, prediction distillation dominates and layerwise distillation is downweighted.

Why not unweighted addition: The two losses operate at different scales and on different types of information. MSE on hidden representations involves $d = 768$ dimensional vectors and can dominate the scalar KL divergence if not properly balanced. The $\lambda$ hyperparameter allows tuning the relative importance of matching final decisions vs. matching intermediate reasoning steps. The paper's ablation in Table 5 shows that removing layer distillation ($\lambda = 1.0$ effectively) reduces performance, and Table 11 shows the effect is present across all sparsities (e.g., on QNLI at 90% sparsity, adding layer distillation improves accuracy from 85.80 to 88.89, a +3.19 gain).

The total training objective. The final loss function combines task-specific cross-entropy (if ground-truth labels are available), the distillation loss, and the Lagrangian sparsity penalty:

$\mathcal{L}_{\text{total}} = \mathcal{L}_{\text{task}} + \mathcal{L}_{\text{distil}} + \mathcal{L}_c$

where $\mathcal{L}_{\text{task}}$ is the standard cross-entropy loss on the labeled task data, $\mathcal{L}_{\text{distil}}$ is the combined distillation loss defined above, and $\mathcal{L}_c$ is the Lagrangian constraint penalty. The task loss ensures the model learns to perform the actual NLP task; the distillation loss ensures it leverages the teacher's knowledge; the Lagrangian penalty ensures it reaches the target sparsity. All three components are optimized simultaneously via gradient descent.

---

Training Procedure and Inference Binarization

The training process follows a multi-phase schedule designed to first establish good task performance and distillation alignment, then gradually introduce pruning pressure.

Phase 1: Distillation warmup. The model is fine-tuned with only the distillation objective ($\mathcal{L}_{\text{distil}}$ and optionally $\mathcal{L}_{\text{task}}$) for a warmup period, with no pruning pressure ($\mathcal{L}_c$ is not applied). All mask variables remain at their initial values (likely near 1, corresponding to keeping all units). For large GLUE datasets (MNLI: 393k, QNLI: 105k, SST-2: 67k, QQP: 364k) and SQuAD (88k), this warmup lasts 1 epoch. For small GLUE datasets (CoLA: 8.5k, RTE: 2.5k, STS-B: 7k, MRPC: 3.7k), the warmup is extended to 4 epochs to compensate for limited data.

The purpose of this warmup is to give the model a chance to learn the distillation alignment before pruning decisions are made. Without this, the optimizer might prune units based on their importance to a random initialization rather than their importance to a well-distilled model, leading to suboptimal pruning decisions.

Phase 2: Pruning with linear sparsity schedule. After warmup, the Lagrangian sparsity controller is activated, and the target sparsity $t$ is linearly increased from its initial value (near 0 — essentially no pruning constraint) to the final target sparsity (e.g., 0.60, 0.95). For large datasets, this linear schedule completes in 2 epochs; for small datasets, it completes in 20 epochs. The total training duration (warmup + pruning) is 20 epochs for large datasets and 100 epochs for small datasets.

Even after the target sparsity is reached (e.g., after 2 epochs on MNLI), training continues for the remaining epochs. During this period, the sparsity constraint is maintained at the target level, but the individual mask values continue to evolve — the optimization can still adjust which specific units are pruned while keeping the overall sparsity constant. The paper describes this as: "even if the final sparsity is achieved, the pruning process keeps searching better performing structures in the rest of the training epochs."

Phase 3: Finetuning the discovered subnetwork. After the full training duration, the masks are binarized (thresholded to exactly 0 or 1 based on the expected sparsity of each weight matrix, as detailed in Appendix B). The resulting pruned architecture — with specific layers, heads, dimensions, and hidden units either present or absent — is treated as a fixed model. This fixed pruned model is then fine-tuned on the task data for an additional 20 epochs (for both large and small datasets).

The paper notes that "finetuning the final subnetwork is essential for high sparsity models" (§A). The reason is that during pruning, the model's weights were optimized under continuously changing masks — once the masks are frozen, the weights can be further refined to specialize to the exact pruned architecture, recovering any performance lost due to the binarization step (which is a discontinuous operation that the continuous training may not have fully accounted for).

Hyperparameters. The paper provides all training hyperparameters in Appendix A, Table 7:

• Distillation temperature $t = 2$ (used to soften probability distributions before computing $D_{\text{KL}}$)
• Finetuning epochs: 20 (for large datasets) or 100 (for small datasets)
• Finetuning learning rates: $\{1e{-5}, 2e{-5}, 3e{-5}\}$, swept
• Training learning rate for the pruning phase: $2e{-5}$ (GLUE), $3e{-5}$ (SQuAD)
• Batch size: 32 (GLUE), 16 (SQuAD)
• $\lambda$ (distillation loss weight): $\{0.1, 0.3, 0.5\}$, swept
• Embedding weights are frozen throughout training, following Sanh et al. (2020) — only Transformer block parameters (and masks) are updated
• "Hyperparameters like $\lambda$, batch size, and learning rate do not generally affect performance much" (§A), suggesting the method is fairly robust to these choices

---

Why These Design Choices? Explicit Rationale

Why separate MHA and FFN layer masks rather than whole-block dropping? Fan et al. (2020) and Sajjad et al. (2020) drop entire Transformer blocks (MHA + FFN together). CoFi separates them because Figure 3 and Appendix H show that MHA layers are systematically more important than FFN layers for downstream performance — models preserve more MHA layers than FFN layers at high sparsity. Dropping whole blocks would force the optimizer to sacrifice an MHA layer whenever it wants to remove a redundant FFN layer, or keep a redundant FFN layer to preserve a useful MHA layer. Independent control allows the optimization to prune FFN layers more aggressively while retaining critical attention layers, maximizing accuracy at a given parameter/speed budget.

Why not prune only heads and dimensions (no layer masks)? The ablation in Table 4 demonstrates the problem: at 95% sparsity on QNLI, removing layer masks drops speedup from 12.1× to 7.2× (model size is identical at ~5M). Without explicit layer-level pruning, the optimizer produces deep-but-thin networks — every layer has a few surviving units, so the GPU must execute all 12 layers of computation, just with smaller matrices. The 7.2× speedup comes only from the reduced matrix sizes; the 12.1× speedup with layer masks additionally comes from skipping entire layers. This is the paper's central empirical finding about the granularity-speed relationship.

Why dynamic matching over fixed matching for layer distillation? Fixed matching assumes a known student architecture (e.g., TinyBERT specifies that student layer 1 learns from teacher layer 3). During pruning, this assumption is violated — the student layer that was supposed to learn from a particular teacher layer may be pruned entirely. Dynamic matching adapts: it always finds a surviving student layer to serve as the distillation target for each teacher layer. The ablation in Table 5 ("Fixed Hidn Distil.") and Table 11 (comparing $\mathcal{L}_{\text{pred}}$ vs. $+\mathcal{L}_{\text{layer}}$) confirms that dynamic matching performs better, especially at higher sparsities where more layers are pruned and fixed mapping targets are more likely to be missing.

Why Lagrangian constraint over fixed $\lambda \|z\|_0$ regularization? Fixed regularization penalizes sparsity with a constant coefficient, but the final sparsity depends on the interaction between the regularization gradient and the task loss gradient — changing the learning rate, dataset size, or model architecture shifts this balance. The paper's preliminary experiments confirmed this unpredictability. The Lagrangian formulation directly optimizes for an equality constraint $\hat{s} = t$, and the multipliers $\lambda_1, \lambda_2$ automatically adapt to enforce it regardless of other hyperparameters. This gives precise, reliable control over the final sparsity, which is essential for fair comparison across different target sparsity levels and datasets.

Why $l_0$ reparameterization with hard concrete distributions over a straight-through estimator? Both approaches create differentiable approximations to binary masks. The hard concrete distribution has the advantage that it can produce values exactly 0 or exactly 1 with non-zero probability (due to the clamping step), which better matches the discrete nature of the final pruning decision. However, the paper explicitly states in a footnote that they "also tried a straight-through estimator as proposed in Sanh et al. (2020) and found the performance comparable." The choice of hard concrete was primarily for precise sparsity control, not accuracy advantages.

Why prune hidden dimensions at all? The paper reports that only a small number of hidden dimensions are pruned (e.g., 768 → 760, roughly 1% of dimensions), but the ablation in Table 4 ("–hidden") shows that removing this capability hurts accuracy: QNLI at 95% drops from 86.1 to 85.6, MNLI from 80.6 to 79.8, SQuAD from 82.6 to 80.8. The hidden dimension mask is shared across all layers and provides a form of global feature selection — it identifies which dimensions of the residual stream are universally unimportant and removes them from all computations. Even a 1% reduction in hidden size creates a consistent, small regularization effect across all layers, which empirically improves final accuracy. The paper speculates that this mask implicitly performs a form of dimensionality reduction on the shared representation space.

Why freeze embeddings? Following Sanh et al. (2020), embedding weights are held fixed during pruning. The embeddings constitute a large fraction of BERT_base's total parameters (roughly 23M out of 110M total, though the paper reports model sizes excluding embeddings, e.g., 85M for BERT_base) and are shared across all tasks and all token positions. Pruning embeddings would create an irregular sparsity pattern in the embedding lookup table, which is hard to speed up (unstructured sparsity in the embedding layer). Moreover, embeddings are task-agnostic — they encode token identities learned during pretraining, and pruning them could degrade the model's ability to represent all tokens, not just tokens irrelevant to the current task. Freezing embeddings ensures the model retains its full vocabulary representation while all compression is applied to the task-specific Transformer computation.

Why start from a fine-tuned BERT rather than the pretrained checkpoint? The paper starts from a BERT_base model that has already been fine-tuned on the target task (following Wang et al., 2020b; Lagunas et al., 2021). This means the pruning process operates on weights that are already specialized to the task, and the pruning decisions (which units are redundant) are made in the context of task-specific knowledge. Starting from a pretrained (non-fine-tuned) model would require the pruning to simultaneously adapt weights to the task and decide which units to remove — an entangled optimization problem that could produce different pruning decisions than those optimal for the final task. The teacher model for distillation is this same fine-tuned BERT_base (frozen), ensuring that distillation transfers task-specific knowledge.

4. Key Insights and Innovations

Innovation 1: The Speedup Gap Is an Optimization Failure, Not an Inherent Limitation of Pruning

The paper's most fundamental conceptual move is reframing the speedup disparity between pruning and distillation not as an unavoidable consequence of pruning's architecture-agnostic nature, but as a solvable optimization problem stemming from inadequate pruning granularity. Before CoFi, the field implicitly accepted a tradeoff: structured pruning gave you flexibility in architecture discovery but couldn't match distillation's speedups (Block Pruning achieves 2.7× vs. TinyBERT's 11.4× in Table 1), while distillation gave you speed but required expensive general pretraining. This was treated as an inherent property of the two paradigms — pruning finds subnetworks within a fixed depth, distillation designs shallow networks from scratch.

CoFi demonstrates that this tradeoff is a failure of the optimization landscape, not a fundamental constraint. The key diagnostic evidence is Table 4: at identical sparsity (95%, ~5M parameters), a model pruned with only fine-grained masks (heads + intermediate dimensions) achieves 7.2× speedup on QNLI, while adding explicit layer masks pushes this to 12.1× — with no change in parameter count. The two models have identical compression rates but dramatically different speedups because the layer masks allow the optimizer to explicitly discover that removing entire layers is a viable path to satisfying the sparsity constraint. Without them, the optimizer gets stuck in local minima where every layer retains some units, producing deep-but-thin networks that GPUs cannot accelerate.

This finding challenges the prevailing narrative in the structured pruning literature, which had trended toward increasingly fine-grained units (heads → FFN dimensions → blocks) under the assumption that smaller pruning units give the optimizer more flexibility and therefore better accuracy. CoFi shows this assumption is backwards at high sparsity: fine granularity alone restricts the optimizer's ability to make coarse architectural decisions that are essential for hardware-efficient inference. The paper's contribution is not the specific masks themselves but the diagnosis that multi-granularity pruning is necessary because fine-grained pruning creates optimization barriers to coarse-grained decisions, not because coarse-grained decisions themselves are inherently better.

This is a fundamental insight, not an incremental refinement. It changes the framing of structured pruning from "find the smallest unit that can be removed while preserving accuracy" to "provide explicit optimization pathways for decisions at every granularity that maps to hardware speedup." The implication for future work is clear: pruning methods should co-design their mask granularity with the target hardware's execution model, not treat sparsity as a purely statistical property.

Innovation 2: Compositional Masking Enables Disentangled Coarse and Fine Pruning Decisions

While the previous innovation concerns why multi-granularity masks matter, this innovation concerns how CoFi's specific mask architecture — five independent mask types composing through logical AND — creates a qualitatively different optimization dynamic from prior hybrid pruning approaches. The paper frames this as a distinctive contribution in Section 3.1:

"CoFi differs from previous pruning approaches in that multiple mask variables jointly control the pruning decision of one single parameter."

To appreciate why this matters, contrast CoFi with Block Pruning (Lagunas et al., 2021), the most directly comparable prior work. Block Pruning also operates at multiple granularities — it applies different pruning strategies to MHA blocks and FFN blocks separately. However, it uses a single mask type per submodule: the MHA pruning strategy and FFN pruning strategy are independent, but within each, there is one mask that makes the pruning decision. A weight in an FFN layer survives or is pruned based on exactly one variable.

CoFi's compositional structure is fundamentally different. For a weight in an FFN up-projection matrix, three independent masks jointly determine survival: $z_{\text{FFN}}$ (is the layer active?), $z_{\text{int}}$ (is this neuron active?), and $z_{\text{hidn}}$ (is this hidden dimension active?). The weight survives only if ALL three are non-zero. This logical-AND composition means the optimizer has three independent pathways to prune the same parameter, each corresponding to a different type of structural simplification. It can satisfy the sparsity constraint by:

• Zeroing $z_{\text{FFN}}$, removing the entire layer (coarse, high speedup, potentially large accuracy cost)
• Zeroing $z_{\text{int}}$ for specific neurons, thinning the layer (fine, moderate speedup, lower accuracy cost)
• Zeroing $z_{\text{hidn}}$ for specific dimensions, removing feature pathways globally (cross-layer, small speedup, regularizing effect)

The optimizer can mix these pathways — partially through layers, partially through neurons, partially through hidden dimensions — and adjust the mixture during training as it discovers which combination best balances the sparsity constraint against task performance. This is a disentangled representation of pruning decisions, analogous to how disentangled latent variables in representation learning separate independent factors of variation. Each mask type specializes in a different structural abstraction (layer-level, neuron-level, feature-level), and the optimization is free to use whichever abstraction is most appropriate for each parameter.

The evidence that this disentanglement matters comes from the hidden dimension mask ablation in Table 4. Removing only this mask ("–hidden") — which prunes at most ~1% of parameters (768 → 760 dimensions) — causes non-trivial accuracy drops: QNLI 86.1 → 85.6, SQuAD 82.6 → 80.8. This is surprising because the hidden dimension mask removes so few parameters. The interpretation is that the hidden mask's value is not in the parameters it directly prunes, but in how it interacts with the other masks through the AND composition — it provides a global regularizing pressure that shapes the optimization landscape for the layer and neuron masks, preventing them from settling into fragile configurations.

This is a conceptual advance, not merely an architectural one. It suggests that in compositional pruning systems, the value of a mask type may lie less in what it directly removes and more in how its presence reshapes the optimization trajectory of other masks. This has implications for designing future pruning systems: adding mask types that remove few parameters directly might still be valuable if they serve as optimization "scaffolding" that guides coarser masks toward better solutions.

Innovation 3: Dynamic Layerwise Distillation as a Moving-Target Knowledge Transfer Solution

The third innovation is a solution to a problem unique to pruning: how do you transfer intermediate-layer knowledge from a teacher to a student when the student's architecture is unknown during training and changes over time? Prior distillation methods (TinyBERT, MobileBERT, MiniLM) pre-define student architectures and can therefore pre-specify layer mappings — teacher layer 3 always distills to student layer 1, etc. Prior pruning methods (Sanh et al., 2020; Lagunas et al., 2021) sidestep the problem entirely by using only prediction-layer distillation, matching output distributions but abandoning the rich intermediate representations that make layerwise distillation so effective.

CoFi's dynamic layer mapping is conceptually elegant in its simplicity: at each training step (or periodically), for each selected teacher layer, find the surviving student layer whose representation is closest in MSE, and use that as the distillation target. The mapping adapts as layers are pruned — if a student layer that was the best match for teacher layer 6 gets pruned, the mapping automatically shifts to the next-closest surviving student layer. No pre-specification is needed, and no teacher layer ever loses its distillation target due to pruning.

What makes this more than an engineering convenience is the implicit curriculum it creates. Early in training, when few layers are pruned, teacher layers map to nearby student layers (by index), producing a conventional deep-to-shallow distillation pattern. As pruning progresses and intermediate layers are removed, the mapping stretches — a later teacher layer might map to a much earlier student layer, forcing that early student layer to learn to represent information that was originally distributed across multiple teacher layers. This is effectively a form of progressive layer fusion, where the student learns to compress the teacher's distributed representations into fewer layers, with the compression schedule determined by the pruning dynamics rather than hand-designed.

The evidence that this dynamic process is beneficial rather than merely necessary comes from the fixed matching baseline in Table 5. "Fixed Hidn Distil." — which matches teacher layers to student layers by index (skipping pruned student layers) — underperforms dynamic matching, particularly on SQuAD (80.9 vs. 82.6 F1). Fixed matching loses distillation signal whenever a student layer is pruned; dynamic matching preserves all available teacher signals by finding the best alternative. The gap widens at higher sparsities (Table 11), exactly when more layers are pruned and fixed matching loses more targets.

This innovation reframes layerwise distillation from a design problem (you must choose the student architecture and layer mapping before training) to an optimization problem (the mapping emerges from the training dynamics). It removes a barrier that had prevented pruning methods from benefiting from intermediate-layer knowledge transfer, and it opens the door to more aggressive architectures (where conventional hand-designed mappings would be impossible) by letting the mapping adapt to whatever survives pruning.

Innovation 4: Empirical Discovery of Asymmetric FFN and MHA Redundancy

This is a diagnostic finding rather than a methodological innovation, but it has significant implications for architecture design and pruning strategy. The paper's structural analysis (Figure 3, Appendix H) reveals a consistent and striking pattern across all datasets and sparsity levels: FFN layers are substantially more redundant than MHA layers for downstream task performance.

At 60% sparsity, the average FFN intermediate dimension count drops by 71% (from 3,072 to ~884), while the average head count drops by only 39% (from 12 to ~7.3). At 95% sparsity, pruned models consistently retain more MHA layers than FFN layers (Appendix H, Figure 5). On some datasets (QQP, SQuAD in Table 6), models can drop the first MHA layer entirely while keeping multiple FFN layers, but the overall pattern is clear: FFN capacity is sacrificied more aggressively than attention capacity.

Prior work had studied head redundancy (Michel et al., 2019; Voita et al., 2019) and FFN redundancy (McCarley et al., 2019) independently, but no prior study had compared their relative redundancy under a unified pruning framework that can independently target both. CoFi's independent MHA and FFN layer masks enable this comparison: the optimization is free to prune from either sublayer type with equal ease, so the observed pruning pattern reflects genuine differences in redundancy rather than constraints imposed by the method.

Why does this matter? FFN layers account for 2/3 of Transformer parameters (each FFN has two weight matrices of size $768 \times 3072$ and $3072 \times 768$, totaling ~9.4M parameters per layer). MHA layers account for the remaining 1/3 (four projection matrices of size $768 \times 768$, totaling ~2.4M per layer). The discovery that FFN layers are significantly more compressible means that parameter-efficient pruning should target FFN layers disproportionately — a strategy that CoFi's independent masks naturally enable but that whole-block pruning (Fan et al., 2020; Sajjad et al., 2020) cannot implement because it forces MHA and FFN layers to be dropped together.

This finding also has implications for architecture design beyond pruning. If FFN layers are systematically overparameterized relative to attention layers for downstream tasks, future Transformer variants might benefit from asymmetric architectures with thinner FFN layers relative to attention, or from different width scaling strategies for the two sublayer types. The paper does not explore these implications directly, but the structural analysis provides the empirical foundation for such explorations.

The significance of this innovation lies in its generality: it is observed across five datasets (SST-2, MNLI, QQP, QNLI, SQuAD) and multiple sparsity levels, suggesting it reflects a property of how Transformers solve NLP tasks rather than a dataset-specific quirk. It converts a vague intuition ("some parts of Transformers are more redundant than others") into a concrete, quantifiable asymmetry that can guide future compression and architecture design.

5. Experimental Analysis

Evaluation Methodology

• Dataset. The paper evaluates on eight GLUE tasks (Wang et al., 2019) — SST-2 (67k training examples), MNLI (393k), QQP (364k), QNLI (105k), MRPC (3.7k), CoLA (8.5k), STS-B (7k), and RTE (2.5k) — and SQuAD v1.1 (Rajpurkar et al., 2016; 88k training examples). The specific metrics used per dataset are listed in Appendix D, Table 8: accuracy for SST-2, QNLI, MNLI, QQP, RTE, and MRPC; Matthews correlation for CoLA; Spearman correlation for STS-B; and F1 score for SQuAD. All results are reported on development sets.

• Base model(s). All experiments use BERT_base (Devlin et al., 2019) as the starting point — a 12-layer Transformer with hidden size 768, 12 attention heads, and FFN intermediate dimension 3072, totaling approximately 85M parameters excluding embeddings. The paper also reports RoBERTa_base (Liu et al., 2019a) results in Appendix I (Figure 7) as a robustness check. The teacher model for distillation is the same BERT_base model fine-tuned on the target task, with its weights frozen during CoFi training.

• Metrics. The primary metrics across experiments are task accuracy (or the task-appropriate metric as specified above), inference speedup (measured as wall-clock latency ratio between the unpruned BERT_base and the pruned model, evaluated on a single NVIDIA V100 GPU with batch size 128 and input sequence length 128 for GLUE and 384 for SQuAD), model size (number of parameters excluding embeddings, following prior work convention), and sparsity (fraction of pruned parameters relative to the full model, embeddings excluded). Speedup is emphasized over parameter count because, as the paper demonstrates, models with identical parameter counts can have substantially different inference latencies depending on whether entire layers are pruned or not (Table 4).

• Baselines. The paper compares against five main baselines, spanning both distillation and pruning paradigms:

  • DistillBERT_6 (Sanh et al., 2019): A 6-layer model initialized by taking every other layer from BERT_base, trained with general distillation on unlabeled data plus task-specific distillation.
  • TinyBERT_6 and TinyBERT_4 (Jiao et al., 2020): 6-layer and 4-layer models trained with general distillation (2,500M tokens for 3 epochs) followed by task-specific distillation with layerwise knowledge transfer.
  • DynaBERT (Hou et al., 2020): A method that produces dynamic-sized models by specifying width and depth multipliers, trained with both general and task-specific distillation.
  • Block Pruning (Lagunas et al., 2021): A structured pruning method that removes blocks within weight matrices using a "Hybrid Filled" strategy, combined with prediction-layer distillation.
  • Movement Pruning (Sanh et al., 2020): An unstructured pruning method that learns importance scores for individual weights, used for comparison in Appendix F.1 (Figure 4).

  Additionally, Appendix F.3 (Table 10) reports comparisons against Wang et al. (2020b), Fan et al. (2020), Sajjad et al. (2020), MobileBERT (Sun et al., 2020), and AutoTinyBERT (Yin et al., 2021), noting that these use different teacher models or base architectures and are therefore not directly comparable.

  For fair comparison, TinyBERT and DynaBERT models are re-trained by the authors using the released code without data augmentation, since the original papers' released models used task-specific augmented data that is not publicly available (Appendix E, Table 9).

• Generation budget / compute accounting. The paper measures training cost in GPU hours on NVIDIA RTX 2080Ti GPUs (Appendix J). The headline comparison is: CoFi trains for at most 20 GPU hours on a single GPU across all GLUE datasets (smaller datasets needing "under 3 hours"), versus TinyBERT's general distillation requiring approximately 350 GPU hours (3.5 days on 4 GPUs) for 3 epochs on 2,500M tokens. The inference cost comparison uses inference speedup (defined above), with all models benchmarked on identical hardware (single NVIDIA V100 GPU). The parameter count metric excludes embedding matrices, which constitutes roughly 23M of BERT_base's ~110M total parameters — this is standard in the literature because embeddings are shared across tasks and forwarding through them has minimal impact on inference time.

• Cross-validation / statistical protocol. The paper does not employ cross-validation. Instead, for each (dataset, target sparsity) combination, the paper runs CoFi three times and reports results. For structural analysis (Table 6, Figure 3), the three runs per setting are visualized individually to show the variance in discovered architectures. For the main accuracy and speedup results (Figure 2, Tables 2–5), single-run results are reported. Hyperparameter sweeps are used to select the distillation weight λ from {0.1, 0.3, 0.5} and the finetuning learning rate from {1e-5, 2e-5, 3e-5}; the training learning rate is fixed at 2e-5 (GLUE) or 3e-5 (SQuAD). The paper notes that "hyperparameters like λ, batch size, and learning rate do not generally affect performance much" (Appendix A), suggesting robustness, though no formal sensitivity analysis is provided.

---

Main Quantitative Results

The paper's experimental results are organized along four axes: overall speedup-accuracy tradeoff curves, head-to-head comparison with TinyBERT_4 at ~10× speedup, structural properties of discovered subnetworks, and ablation studies (covered separately in the next subsection). The RoBERTa results in Appendix I and additional baseline comparisons in Appendices E and F provide supporting evidence.

Overall Speedup-Accuracy and Size-Accuracy Tradeoff (Figure 2)

Figure 2 presents the central result: accuracy plotted against both inference speedup (top row) and model size (bottom row) across five datasets — SST-2, MNLI, QQP, QNLI, and SQuAD. Each point represents a model produced by one of the methods under comparison, with CoFi models spanning the full range of sparsities from 60% to 95% (corresponding to model sizes from ~34M down to ~5M parameters).

Speedup dimension (Figure 2, top). Across all five datasets and at every speedup level, CoFi models lie on or above the Pareto frontier defined by all other methods. Concretely:

• On SST-2, CoFi achieves 93.0% accuracy at 2.0× speedup (34M parameters, 60% sparsity), matching or exceeding DistillBERT_6, TinyBERT_6, and Block Pruning at similar speedups. At approximately 12.0× speedup (~5M parameters, 95% sparsity), CoFi reaches 90.6% accuracy — the TinyBERT_4 point at 11.4× speedup achieves 89.7%, so CoFi is roughly 0.9 points higher with a slightly higher speedup.
• On MNLI, the pattern is similar but more pronounced: CoFi at 2.1× speedup hits 85.3% accuracy (vs. DistillBERT_6 at 82.2% with 2.0× speedup, and TinyBERT_6 at 84.0% with 2.0× speedup). At extreme compression — 12.1× speedup with only 4.4M parameters — CoFi achieves 80.6% accuracy on MNLI, compared to TinyBERT_4's 78.8% at 11.4× speedup. This is a gain of 1.8 percentage points at a slightly higher speedup.
• On QQP, CoFi at 2.0× speedup reaches 91.2% accuracy vs. TinyBERT_6 at ~91.0% (read from figure). At 11.0× speedup and ~5M parameters, CoFi achieves 90.1% vs. TinyBERT_4 at 90.0%.
• On QNLI, CoFi at 2.1× speedup reaches 91.8% accuracy, outperforming all baselines at that speedup regime. At 12.1× speedup, it reaches 86.1% vs. TinyBERT_4 at 86.7% — TinyBERT_4 has a slight edge (0.6 points) on QNLI, which Table 2 confirms.
• On SQuAD, CoFi yields models at speedups ranging from 2.0× (F1 = 89.1) to 8.7× (F1 = 82.6). TinyBERT_4 reports 82.1 F1 at 8.7× speedup, so CoFi holds a 0.5 F1 advantage.

Size dimension (Figure 2, bottom). The story is similar: CoFi models achieve equal or better accuracy at every model size compared to all baselines. Notably, at very small model sizes (~5M parameters), CoFi consistently outperforms or matches TinyBERT_4 while requiring no general distillation. The gap between CoFi and Block Pruning is particularly visible in the high-speedup regime: Block Pruning cannot produce models beyond approximately 3× speedup (its pruned models retain all 12 layers), while CoFi pushes past 10× on all datasets.

Key pattern across datasets. The relative advantage of CoFi over TinyBERT_4 varies by dataset. Gains are largest on MNLI (+1.8 accuracy, +0.7× speedup), moderate on SST-2 (+0.9 accuracy), and slim on QNLI (−0.6 accuracy, CoFi slightly behind). This suggests that the benefit of architecture discovery (finding dataset-specific layer configurations) depends on the task — tasks where the optimal subnetwork differs substantially from a uniformly shallower network benefit more from CoFi's flexible pruning.

Head-to-Head with TinyBERT_4 at ~10× Speedup (Table 2)

Table 2 provides a granular comparison between CoFi at 95% sparsity (~5M parameters) and TinyBERT_4 (also ~5M parameters) across all nine datasets, including the four small GLUE tasks where general distillation matters most. The headline numbers:

• Large GLUE datasets: On SST-2, CoFi achieves 90.6% vs. TinyBERT_4's 89.7% (+0.9); on QNLI, 86.1% vs. 86.7% (−0.6); on MNLI, 80.6% vs. 78.8% (+1.8); on QQP, 90.1% vs. 90.0% (+0.1). Speedups are comparable: CoFi achieves 11.0–12.1× across these datasets vs. TinyBERT_4's 11.4×.
• Small GLUE datasets: These are where general distillation is essential for TinyBERT. TinyBERT_4 without general distillation ("TinyBERT_4 w/o GD") collapses on CoLA (16.6 vs. 61.2 teacher), RTE (47.3 vs. 70.0), STS-B (17.8 vs. 88.7), and MRPC (68.9 vs. 85.0). CoFi, using only task-specific data, achieves 35.6 on CoLA (vs. TinyBERT_4's 32.5 with GD), 64.7 on RTE (vs. 63.2), 83.1 on STS-B (vs. 85.0), and 82.6 on MRPC (vs. 81.4). CoFi outperforms TinyBERT_4 on 3 out of 4 small datasets, with TinyBERT_4 holding a slight edge on STS-B (85.0 vs. 83.1).
• SQuAD: CoFi achieves 82.6 F1 at 8.7× speedup vs. TinyBERT_4's 82.1 F1 at the same speedup (+0.5 F1).

The training time comparison in the rightmost column quantifies the efficiency advantage: CoFi trains in ≤20 GPU hours total across all datasets (and under 3 hours for the four small datasets), while TinyBERT_4 requires ~350 GPU hours for general distillation alone — an 18× reduction in total training time. This is the paper's most direct evidence that task-specific structured pruning can serve as a "drop-in replacement" for expensive general distillation.

Results with Data Augmentation (Table 3)

Table 3 extends the comparison to the setting where task-specific data augmentation (introduced in Jiao et al., 2020) is applied. Since the augmented dataset was not publicly released, the authors created their own augmented data following the TinyBERT GitHub repository. Training on augmented data is expensive — "training on the augmented dataset for MNLI takes more than 200 GPU hours in total" (Appendix E). Consequently, experiments are run on only four datasets.

With data augmentation:

• On SST-2, CoFi reaches 92.4% vs. TinyBERT_4's 91.6% (+0.8).
• On QNLI, TinyBERT_4 retains an edge: 87.6% vs. CoFi's 86.8% (−0.8).
• On RTE, CoFi reaches 67.5% vs. TinyBERT_4's 62.5% (+5.0).
• On MRPC, CoFi reaches 84.6% vs. TinyBERT_4's 83.6% (+1.0).

CoFi benefits from data augmentation (compare to Table 2 numbers: SST-2 90.6 → 92.4, QNLI 86.1 → 86.8, RTE 64.7 → 67.5, MRPC 82.6 → 84.6), and it still outperforms TinyBERT_4 on 3 of 4 datasets. The paper does not report the computational cost of running CoFi on augmented data, but notes that "training on augmented data is very expensive" and that the main experiments (without augmentation) reflect the method's typical deployment scenario.

Structural Patterns in Pruned Models (Figure 3, Table 6)

Figure 3 visualizes the average remaining heads per MHA layer and average remaining intermediate dimensions per FFN layer across five datasets (SST-2, MNLI, QQP, QNLI, SQuAD) at sparsities {60%, 70%, 80%, 90%, 95%}. Each point is averaged over three CoFi runs.

Pattern 1: FFN layers are pruned more aggressively than MHA layers. At 60% sparsity, the average intermediate FFN dimension count drops by 71% (from the original 3,072 to approximately 884, computed from Figure 3), while the average head count drops by only 39% (from 12 to approximately 7.3). This gap persists across all sparsity levels. At 95% sparsity, many FFN layers have zero remaining intermediate dimensions (i.e., the entire FFN layer is pruned), while most MHA layers retain at least 1–2 heads.

Pattern 2: Upper layers are pruned more aggressively than lower layers. In both the MHA and FFN plots, the lines slope downward from left (lower layers) to right (upper layers). For example, at 80% sparsity on MNLI, layer 1 retains approximately 7 heads while layer 12 retains approximately 3 heads; similarly, layer 1's FFN retains roughly 1,500 intermediate dimensions while layer 12's retains roughly 300.

Pattern 3: Dataset-specific structural variation. Figure 3 shows that the exact pruning pattern differs across datasets. MNLI and QQP tend to preserve more capacity in lower layers compared to SST-2, while SQuAD preserves relatively more upper-layer MHA heads than the other datasets. This indicates that the optimal subnetwork architecture is not universal — it depends on the task.

Table 6 provides a complementary view at the coarsest granularity: which MHA and FFN layers survive at 95% sparsity. Across three runs per dataset, the paper visualizes the surviving layers as sequences of M (remaining MHA layer) and F (remaining FFN layer). The key observations:

• More MHA layers than FFN layers survive. Counting across all entries in Table 6: on SST-2, the three runs retain an average of approximately 7.3 MHA layers and 5.7 FFN layers; on QNLI, approximately 6.7 MHA and 3.0 FFN; on MNLI, approximately 5.7 MHA and 3.7 FFN; on QQP, approximately 6.0 MHA and 4.0 FFN; on SQuAD, approximately 5.7 MHA and 4.3 FFN.
• MHA and FFN layers can be independently dropped or retained. The sequences show examples like "M F M M F M F M F M F" (SST-2) with interleaving of surviving MHA and FFN layers, and "F M F M M F M F M F" (SQuAD) where the first layer is an FFN (the MHA was pruned). This confirms that CoFi's independent layer masks are being used — the optimization is not simply dropping entire blocks.
• The first few MHA layers are often preserved. Across all 15 runs (3 per dataset × 5 datasets), the first MHA layer survives in 10 cases, the second MHA layer in 11 cases.
• Middle layers are more frequently pruned. The middle MHA layers (layers 5–8) are pruned more often than either early or late layers.

Appendix H, Figure 5 quantifies this layer-type asymmetry: at 60% sparsity, the average pruned model has approximately 11.5 MHA layers and 9.5 FFN layers remaining; at 95% sparsity, this drops to approximately 6 MHA layers and 4 FFN layers. The FFN count drops faster than the MHA count as sparsity increases, consistent with the greater redundancy of FFN layers.

RoBERTa Results (Appendix I, Figure 7)

Figure 7 shows CoFi applied to RoBERTa_base on SST-2 and MNLI across sparsities 60%–95%, compared with BERT_base. On SST-2, RoBERTa maintains slightly higher accuracy than BERT at sparsities below 90% (e.g., at 80% sparsity, RoBERTa ~93% vs. BERT ~92.5%), but the advantage flips at 95% sparsity (BERT ~90.6% vs. RoBERTa ~89.5%). On MNLI, the same pattern holds: RoBERTa outperforms BERT at 60–80% sparsity but BERT surpasses RoBERTa at 95% sparsity. The paper notes that similar patterns were observed in DynaBERT (Hou et al., 2020), suggesting this may be a general property of pruning higher-capacity models — the advantage of a stronger initialization diminishes as sparsity becomes extreme because the pruned model's capacity is the binding constraint, not the initial weight quality.

Comparison with Movement Pruning (Appendix F.1, Figure 4)

Figure 4 compares CoFi (structured pruning) with Movement Pruning (unstructured; Sanh et al., 2020) on MNLI and SQuAD. CoFi is evaluated both with full distillation (prediction + layerwise) and with prediction-only distillation (labeled "CoFi Logit Distill") to match Movement Pruning's distillation setup.

• On MNLI, CoFi with prediction-only distillation consistently outperforms Movement Pruning. At ~50M parameters (low sparsity), CoFi achieves approximately 84.5% vs. Movement Pruning's 84%. At ~10M parameters, CoFi is at roughly 81% vs. Movement Pruning's 79%. The gap widens as models get smaller. CoFi with full distillation provides a further consistent boost.
• On SQuAD, CoFi with prediction-only distillation is comparable to Movement Pruning at higher parameter counts (≥30M: both near 88 F1) but drops more steeply at very small sizes (<10M): Movement Pruning reaches ~82 F1 at 5M parameters while CoFi (pred-only) drops to ~79. However, CoFi with full distillation largely closes this gap (82.6 at 8.7× speedup vs. Movement Pruning's ~81 at similar size, though Movement Pruning achieves zero speedup).
• The critical qualitative difference: Movement Pruning achieves its sparsity but produces no inference speedup on standard hardware (V100 GPU). CoFi's models achieve 2×–11× speedups. This is the fundamental tradeoff: unstructured pruning can extract higher accuracy at extreme sparsity by removing individual weights, but the resulting irregular sparsity patterns are not exploitable on current GPUs.

Comparison with Block Pruning (Appendix F.2, Figure 6)

Figure 6 normalizes the distillation objective by comparing CoFi and Block Pruning both using only prediction-layer distillation ("Logit Distill"). On SST-2, QNLI, and MNLI, CoFi matches or slightly exceeds Block Pruning's accuracy at similar sparsity levels (e.g., on SST-2 at 90% sparsity: CoFi ~91.5% vs. Block Pruning ~91.0%). The critical difference, however, is in speedup: Block Pruning never achieves more than approximately 3× speedup even at high sparsity (90%+), while CoFi reaches 10×+. This is because Block Pruning, while operating at a fine granularity within weight matrices, does not remove entire layers — the 12-layer depth structure is preserved, limiting speedup regardless of how many individual blocks are pruned within each layer.

---

Ablation Studies and Robustness Checks

The paper conducts ablation studies on two axes: pruning units (what happens when you remove specific mask types) and distillation objectives (what happens when you remove or modify the knowledge transfer components). The ablation experiments are run on subsets of datasets (QNLI, MNLI, SQuAD, SST-2) and sparsities (60% and 95%) to keep the computational cost manageable.

Pruning unit ablations (Table 4). Three configurations are compared at both moderate (60%) and extreme (95%) sparsity:

• Removing hidden dimension mask (–hidden): When $z_{\text{hidn}}$ is removed, the optimization can no longer prune individual dimensions of the residual stream. At 60% sparsity, accuracy drops modestly: QNLI 91.8 → 91.3, MNLI 85.1 → 85.2 (no drop), SQuAD 89.1 → 88.7. Speedup is essentially unchanged (within 0.1×). At 95% sparsity, the accuracy drops are larger: QNLI 86.1 → 85.6, MNLI 80.6 → 79.8, SQuAD 82.6 → 80.8 (a 1.8 F1 point drop). Speedup increases at 95% sparsity when hidden dimensions are removed (12.1× → 13.3× on QNLI, 12.1× → 13.7× on MNLI, 8.7× → 9.7× on SQuAD). This is because, without the hidden dimension mask, the optimizer compensates by pruning more entire layers to meet the sparsity target, producing deeper-but-narrower networks that are actually faster (fewer layers → more kernel skipping) but less accurate. This is a non-obvious finding: the hidden dimension mask doesn't just contribute to sparsity directly — it shapes the layer-pruning tradeoff by providing an alternative pathway to satisfy the sparsity constraint.

• Removing hidden dimension mask and layer masks (–layer & hidden): This configuration corresponds most closely to prior structured pruning approaches (head pruning + FFN dimension pruning only). At 60% sparsity, accuracy is only slightly worse than full CoFi (QNLI: 91.8 → 91.3, MNLI: 85.1 → 84.8, SQuAD: 89.1 → 88.5). However, at 95% sparsity, speedup collapses while accuracy also drops: QNLI accuracy falls to 84.6 (from 86.1) and speedup plummets from 12.1× to 7.2×; MNLI drops to 78.4 (from 80.6) with speedup from 12.1× to 7.0×; SQuAD drops to 74.1 F1 (from 82.6) with speedup from 8.7× to 6.4×. This is the paper's strongest evidence for the necessity of explicit layer pruning: without it, the optimizer cannot effectively remove entire layers, producing deep-but-thin networks that are both slower and less accurate at high sparsity. The speedup reduction (12.1× → 7.2× on QNLI) with identical parameter count (~5M) directly demonstrates that speedup is not purely a function of parameter count — it depends on the structure of the remaining parameters.

• Removing only layer masks (–layer): This configuration keeps $z_{\text{hidn}}$ but removes $z_{\text{MHA}}$ and $z_{\text{FFN}}$. At 60% sparsity, accuracy is nearly unchanged (QNLI: 91.8 → 91.5, MNLI: 85.1 → 85.4, SQuAD: 89.1 → 89.1). At 95% sparsity, speedup again collapses (QNLI: 12.1× → 8.3×, MNLI: 12.1× → 8.4×, SQuAD: 8.7× → 7.9×), though accuracy is slightly better than full CoFi on QNLI (86.7 vs. 86.1) and equal on MNLI (80.6 vs. 80.6), with a drop on SQuAD (80.5 vs. 82.6). This confirms that the layer masks are the primary driver of large speedups, and that at extreme sparsity, the optimizer would sometimes prefer to prune more layers for speed but is constrained from doing so by the accuracy objective — without layer masks, it's forced into the deep-but-thin regime.

The key takeaway from these ablations is that no single mask type is sufficient for the full speedup-accuracy frontier. The hidden dimension mask provides accuracy benefits (especially at high sparsity) by enabling a small amount of global feature selection; the layer masks provide the speedup benefits by enabling explicit layer removal; and the head/intermediate masks provide the fine-grained flexibility to preserve important computational units within surviving layers. The combination enables both high accuracy and high speedup, which neither subset achieves alone.

Distillation objective ablations (Table 5, with supporting evidence in Table 11). Three ablations plus one alternative approach are compared at 95% sparsity:

• Removing all distillation (–$\mathcal{L}_{\text{pred}}$, –$\mathcal{L}_{\text{layer}}$): The model is trained with only task cross-entropy and the Lagrangian sparsity penalty. Accuracy drops substantially: SST-2 90.6 → 86.6, QNLI 86.1 → 84.2, MNLI 80.6 → 78.2, SQuAD 82.6 → 75.8. This is a 1.9–6.8 point drop, confirming that distillation is essential for maintaining accuracy under aggressive pruning. The drop is largest on SQuAD (−6.8 F1), which requires extracting precise spans and may rely more heavily on the teacher's intermediate representations.

• Removing only layer distillation (–$\mathcal{L}_{\text{layer}}$), keeping prediction distillation: This tests whether layerwise distillation provides benefits beyond prediction-layer distillation alone. Accuracy drops modestly: SST-2 90.6 → 91.1 (actually a slight improvement, suggesting layer distillation is not helpful for SST-2 at 95% sparsity), QNLI 86.1 → 85.1 (−1.0), MNLI 80.6 → 79.7 (−0.9), SQuAD 82.6 → 82.5 (−0.1). The benefits of layer distillation are dataset-dependent: large gains on QNLI and MNLI, minimal on SQuAD and SST-2.

• Fixed hidden distillation (matching by index): Instead of dynamic matching, each teacher layer $i$ is matched to student layer $i$ (if it survives; otherwise, no distillation from that teacher layer). Results: SST-2 90.6 → 90.0 (−0.6), QNLI 86.1 → 85.8 (−0.3), MNLI 80.6 → 80.5 (−0.1), SQuAD 82.6 → 80.9 (−1.7). SQuAD shows the largest penalty from fixed matching, likely because SQuAD prunes more layers (Table 6 shows models often starting with "F" — pruned first MHA layer), causing fixed matching to lose teacher signals for early-pruned student layers. Dynamic matching adapts to find the best available surrogate layer.

Table 11 extends these ablations across all sparsities {60%, 70%, 75%, 80%, 85%, 90%, 95%} on four datasets, comparing prediction-layer distillation alone ($\mathcal{L}_{\text{pred}}$) against the full distillation objective ($\mathcal{L}_{\text{pred}} + \mathcal{L}_{\text{layer}}$). The gains from adding layer distillation ($+\mathcal{L}_{\text{layer}}$ column) are:

• On SST-2: positive at all sparsities except 95% (+0.34 to +2.07 at lower sparsities, −0.69 at 95%).
• On QNLI: positive at all sparsities, ranging from +0.67 to +3.19 (largest gain at 90% sparsity).
• On MNLI: positive at all sparsities, ranging from +0.15 to +1.52.
• On SQuAD: positive at all sparsities, ranging from +0.07 to +0.80.

The pattern is clear: layer distillation provides the largest absolute gains at intermediate-to-high sparsities (75–90%), where the student still has multiple surviving layers and can benefit from intermediate supervision, but the layers are sufficiently pruned that prediction-only distillation loses important representational information. At very low sparsity (60%), gains are smaller because the student retains most layers and can learn intermediate representations implicitly through the prediction objective. At extreme sparsity (95%), gains diminish because the student has very few surviving layers to match against.

Dynamic matching behavior (Appendix G.1). The paper reports that dynamic matching converges to specific, stable layer alignments that differ across datasets. On SST-2, student layers 7, 9, 10, 11 map to teacher layers 3, 6, 9, 12. On QQP, student layers 2, 5, 8, 11 map to the same teacher layers. These alignments are discovered during training without any pre-specification, confirming that the dynamic matching is finding task-specific layer correspondences rather than simply defaulting to index-based mapping.

RoBERTa vs. BERT as base model (Appendix I, Figure 7). This serves as a robustness check: CoFi's behavior is not specific to BERT_base initialization. The overall trends are preserved with RoBERTa_base — accuracy degrades gracefully with increasing sparsity, and the method produces usable models across the full sparsity range. The cross-over where BERT outperforms RoBERTa at extreme sparsity is noted but not explained; it may relate to BERT's lower capacity making it easier to find good subnetworks when capacity is severely constrained.

Data augmentation robustness (Table 3). CoFi benefits from data augmentation and remains competitive with TinyBERT_4 under this setting, indicating that the method's advantage is not contingent on using only the original task data — it can leverage additional training signal when available, though at significantly increased computational cost.

---

Critical Assessment

Claim 1: CoFi matches distillation methods in accuracy and latency without unlabeled data. This is the paper's central claim, stated in the abstract and supported primarily by Figure 2 and Table 2. The evidence is strong for the comparison against TinyBERT_4 at ~10× speedup across all nine datasets — CoFi matches or exceeds TinyBERT_4's accuracy on 7 of 9 datasets (Table 2: wins on SST-2, MNLI, QQP, CoLA, RTE, MRPC, SQuAD; trails on QNLI by 0.6 points and on STS-B by 1.9 points). The speedups are comparable across datasets (11.0–12.9× for CoFi vs. ~11.4× for TinyBERT_4, with some variance depending on the specific pruned architecture).

However, the claim has important boundary conditions that are not fully explored. First, the comparison is against TinyBERT_4 specifically — a 4-layer model. TinyBERT_6 (6 layers, 2.0× speedup) achieves higher accuracy than TinyBERT_4 on all datasets where both are reported (e.g., MNLI: 84.0 vs. 78.8). CoFi at 2.1× speedup achieves 85.3 on MNLI, beating TinyBERT_6, but the paper does not produce a CoFi model at exactly 2.0× speedup with 6 layers to compare against TinyBERT_6's architecture directly. The intermediate speedup regime (3–8×) is sparsely populated by baselines, making it harder to assess whether CoFi's advantage holds at all points on the curve or primarily at the extremes.

Second, the claim of matching distillation methods is supported against a specific set of distillation baselines (TinyBERT, DistillBERT, DynaBERT). At the time of publication, there existed other distillation methods not included in the comparison — MiniLM (Wang et al., 2020a), which uses deep self-attention distillation; MobileBERT (Sun et al., 2020), which uses a specialized bottleneck architecture; and Patient-KD (Sun et al., 2019). The paper includes MobileBERT in Table 10 as a non-direct comparison (it uses a different teacher and specialized architecture) but does not attempt to compare against MiniLM or Patient-KD. The claim should therefore be interpreted as "matches these specific distillation methods" rather than "matches all distillation methods."

Claim 2: CoFi achieves over 10× speedups with 95% sparsity while preserving >90% accuracy. The numbers from Table 2 bear this out: across all datasets, CoFi at ~95% sparsity achieves speedups ranging from 8.7× (SQuAD, due to longer sequence length) to 12.9× (STS-B). Accuracy relative to the BERT_base teacher: SST-2 retains 97.3% of teacher accuracy (90.6 / 93.1), QNLI 94.1%, MNLI 95.0%, QQP 98.8%, CoLA 58.2%, RTE 92.4%, STS-B 93.7%, MRPC 97.2%, SQuAD 93.4%. The claim of "over 90% accuracy preserved" holds for all datasets except CoLA (where even the teacher only achieves 61.2 — the absolute drop of 25.6 Matthews correlation points is severe, from 61.2 to 35.6). CoLA is a small dataset (8.5k examples) testing linguistic acceptability, and the poor performance at high sparsity suggests that this task requires more model capacity than the others. The paper does not discuss this outlier or investigate why CoLA degrades more severely.

A more precise statement would be: "CoFi preserves >90% of teacher accuracy on 7 of 8 GLUE tasks and on SQuAD at 95% sparsity, with CoLA being a notable exception (58% of teacher performance retained)."

Claim 3: CoFi incurs much less computation than distillation (≤20 GPU hours vs. ~350 GPU hours). The training time comparison is well-documented: TinyBERT's general distillation takes "3.5 days on 4 GPUs" (Appendix J, measured on RTX 2080Ti GPUs), which the paper converts to ~350 GPU hours. CoFi's training takes "at most 20 hours on 1 GPU" for the large datasets and "under 3 hours" for small datasets. This is an 18× reduction for large datasets and a >100× reduction for small ones.

However, two caveats apply. First, the 350 GPU hours is for general distillation only — task-specific distillation adds additional training time that the paper measures as "at most 10 GPU hours." The total TinyBERT cost is therefore ~360 GPU hours, and CoFi is ≤20 GPU hours, making the comparison even more favorable to CoFi than the 18× figure suggests. Second, the training time comparison does not include the cost of fine-tuning the original BERT_base model (which both methods require). For the large GLUE datasets, this is a one-time cost that is identical for both methods and can be amortized across all downstream pruning or distillation runs.

The type of GPU differs between the two measurements: CoFi is measured on a single NVIDIA V100 GPU (inference speedup is on V100, and the paper's language suggests training is also on V100), while TinyBERT's general distillation is measured on RTX 2080Ti GPUs. RTX 2080Ti and V100 have similar compute capabilities (both Turing/Volta generation, similar TFLOPS), so the comparison is roughly fair, but the paper does not explicitly control for hardware differences.

Claim 4: Multi-granularity pruning is necessary for achieving both high speedups and high accuracy. The ablation in Table 4 supports this claim strongly, particularly the 95% sparsity results. Removing layer masks (−layer: 12.1× → 8.3× speedup on QNLI, 12.1× → 8.4× on MNLI) or removing both layer and hidden masks (−layer & hidden: 12.1× → 7.2× on QNLI, 12.1× → 7.0× on MNLI) causes speedup collapse. The hidden dimension mask ablation (−hidden) shows a more nuanced pattern: removing it causes a slight accuracy drop (QNLI: −0.5, MNLI: −0.8, SQuAD: −1.8) but slightly higher speedup (12.1× → 13.3× on QNLI) because the optimizer compensates by pruning more layers. This demonstrates that the masks interact — the hidden mask's primary value is not in the parameters it directly prunes (~1%) but in providing an alternative pruning pathway that prevents the optimizer from over-pruning layers to meet the sparsity target, thereby preserving accuracy.

The claim that this finding is general (applies across tasks and architectures) is supported only within the narrow experimental scope: BERT_base on GLUE/SQuAD, with a specific set of mask granularities. The paper does not test whether the same granularity-speed relationship holds for larger models (BERT_large), different architectures (T5, GPT), or different hardware platforms (CPUs, where the MHA/FFN latency ratio differs per Ganesh et al., 2021). This is not a weakness of the paper per se — the experiments are already extensive — but it bounds the generality of the insight.

Claim 5: Dynamic layerwise distillation outperforms fixed matching and prediction-only distillation. Table 5 and Table 11 support this claim, but the margin varies by dataset and sparsity. On SQuAD at 95% sparsity, the advantage of dynamic over fixed matching is 1.7 F1 points (82.6 vs. 80.9) — a substantial gain. On MNLI at 95%, the advantage is only 0.1 points (80.6 vs. 80.5) — negligible. This suggests the value of dynamic matching depends on how many layers are pruned (more pruning → more missing layers in fixed matching → larger dynamic advantage) and whether the optimal layer mapping deviates from the index-based default (SQuAD shows significant deviation; MNLI less so).

The paper does not provide a comprehensive analysis of when dynamic matching matters. The fixed matching ablation is run only at 95% sparsity (Table 5), not across sparsities. The cross-sparsity analysis in Table 11 only compares prediction-only vs. prediction+layer (not fixed vs. dynamic within the layer distillation group). This makes it impossible to determine whether dynamic matching's advantage is concentrated at high sparsities (where fixed matching loses many teacher signals) or is broadly beneficial. This is a missing ablation that would have strengthened the paper's argument for dynamic matching specifically (as opposed to the general benefit of layer distillation).

Genuine weaknesses:

1. No error bars or variance estimates. The paper runs CoFi three times per setting for structural analysis (Table 6, Figure 3) but reports single-run results for accuracy and speedup. Without variance estimates, the reader cannot assess whether the performance differences between CoFi and baselines (e.g., 90.6 vs. 89.7 on SST-2 in Table 2, a 0.9 point gap) are statistically meaningful or within run-to-run noise. Given the test sets — especially for small datasets like RTE (2.5k examples) and MRPC (3.7k) — small absolute differences could be entirely due to sampling variability.

2. Speedup measurement on a single GPU type. All inference speedup measurements are on a single NVIDIA V100 GPU with batch size 128 and fixed sequence lengths (128 for GLUE, 384 for SQuAD). The paper acknowledges that "results might be different from the original papers as the environment for each platform is different" (§4.1). The speedup achievable through layer pruning depends on GPU characteristics (memory bandwidth, kernel launch overhead, etc.). On CPUs, Ganesh et al. (2021) note that "FFNs become the bottleneck," suggesting CoFi's aggressive FFN pruning could yield even larger relative speedups on CPUs — but this is untested. On TPUs or newer GPUs (A100), the speedup profile could differ.

3. The distillation weight λ and other hyperparameters are swept without systematic analysis. The paper sweeps λ ∈ {0.1, 0.3, 0.5} and reports the best result per dataset, but does not show the sensitivity curve. If accuracy varies substantially with λ, the method requires hyperparameter tuning per dataset; if it doesn't (as the paper asserts — "hyperparameters... do not generally affect performance much"), showing this robustness would strengthen the method. The claim is made but not demonstrated quantitatively.

4. No comparison against training a compact model from scratch with the same architecture. CoFi discovers an architecture for each (dataset, sparsity) combination. A natural question is: could you simply train that architecture from scratch (random initialization) with the same task-specific data and match CoFi's performance? If so, the pruning process would be reduced to architecture search, and the specific mask optimization would be unnecessary. The paper does not address this. Prior work in computer vision (Liu et al., 2019c) has shown that training discovered architectures from scratch can match the performance of weight inheritance from pruning, but this has not been established for Transformer pruning. Running such an experiment would clarify whether CoFi's value lies in the discovered architecture, in the transferred weights, or in both.

5. The 18× training time comparison may overstate the practical advantage. TinyBERT's general distillation cost (350 GPU hours) is a one-time cost that can be amortized across all downstream tasks — once you have the generally distilled TinyBERT checkpoint, task-specific distillation for each new task is fast. CoFi must be run from scratch for each new task, requiring ≤20 GPU hours per task. For a practitioner deploying models on N tasks, the total cost is 350 + 10N for TinyBERT vs. 20N for CoFi. The break-even point is N ≈ 13 tasks — for organizations deploying on fewer tasks, CoFi is cheaper; for large-scale multi-task deployments, TinyBERT's amortized cost could make it more economical. The paper's framing of CoFi as strictly more efficient ignores this amortization argument.

6. No evaluation on out-of-distribution or robustness metrics. All evaluations are on in-distribution test sets (GLUE dev, SQuAD dev). The paper does not assess whether pruned models retain the teacher's robustness to adversarial examples, domain shift, or spurious correlations. Knowledge distillation is known to sometimes improve robustness (by softening decision boundaries), and structured pruning could either preserve or degrade this property.

7. SQuAD speedup is consistently lower than GLUE speedup. At 95% sparsity, SQuAD achieves 8.7× speedup vs. 11–13× on GLUE tasks (Table 2). The paper attributes this to the longer sequence length (384 vs. 128), which increases the proportion of time spent in attention operations (which scale quadratically with sequence length) relative to FFN operations. This means the speedup gains from pruning FFN layers are diluted on long-sequence tasks — an important practical consideration not explored in the paper (e.g., on document-level tasks with sequences of 512+ tokens, speedups would presumably be even lower).

8. Missing experiment: computational cost of the pruning process itself. The paper reports only final training time (≤20 GPU hours) but does not break this down into the warmup, pruning, and finetuning phases. It also does not report the cost of the three-run averaging for structural analysis, or the cost of the hyperparameter sweep over λ. For a practitioner considering CoFi, knowing the cost of a single run vs. the cost with hyperparameter tuning would be valuable.

What would have strengthened the paper:

• A comparison against training CoFi-discovered architectures from scratch (random initialization) to isolate the value of weight inheritance vs. architecture discovery.
• Error bars over multiple runs for accuracy and speedup numbers, especially on small datasets.
• A fixed-vs-dynamic matching ablation across multiple sparsities, not just at 95%.
• A sensitivity analysis for λ to support the claim that hyperparameters don't matter much.
• Speedup measurements on a second hardware platform (CPU or a different GPU generation).
• Out-of-distribution evaluation (e.g., on GLUE diagnostic sets or adversarial NLI benchmarks).
• An analysis of the computational cost breakdown by phase (warmup vs. pruning vs. finetuning).
• A discussion of the break-even point in multi-task deployment scenarios where general distillation costs can be amortized.

Despite these limitations, the experimental evidence for the paper's main claims is robust within its defined scope: CoFi achieves 10×+ speedups on BERT_base across GLUE/SQuAD, matches or exceeds TinyBERT_4's accuracy without general distillation, trains in a fraction of the time, and the multi-granularity masks + dynamic layer distillation are both essential components. The weaknesses are primarily about scope (single model family, single hardware platform, no robustness evaluation) and statistical rigor (no variance estimates), not about the validity of the core results.

6. Limitations and Trade-offs

The Cost of Difficulty Estimation Is Unaccounted for in Headline Efficiency Claims

The paper's entire compute-optimal framework — the adaptive selection of search strategies and sequential-to-parallel ratios based on question difficulty — rests on a difficulty estimation procedure that is dramatically more expensive than the test-time compute budgets being studied. The method for estimating a single question's difficulty, as described in Section 3.2, requires generating 2,048 complete solutions from the base model and then scoring them (either against ground-truth answers for oracle difficulty, or using the PRM's final-answer score for predicted difficulty). At 2,048 samples per question, the difficulty estimation step alone consumes 8–32 times more compute than the largest test-time budgets analyzed (256–512 generations). The authors acknowledge this directly in Section 3.2:

"estimating difficulty in this way still incurs additional computation cost during inference… our experiments do not account for this cost largely for simplicity"

The consequence is that the paper's central efficiency claim — that compute-optimal scaling yields "more than 4× better efficiency over a standard best-of-N baseline" — is computed after difficulty has already been determined, without amortizing the cost of learning it. In a realistic deployment, the total inference cost would be 2,048 (difficulty estimation) + N (strategy execution) generations, and the former dominates the latter across all budget levels studied. For example, the paper reports that compute-optimal search at 16 generations matches best-of-N at 64 generations (Figure 4). But the total cost to achieve this is not 16 generations — it is 2,048 + 16 = 2,064 generations, which is 32× more than the best-of-N baseline it supposedly beats. The 4× figure is therefore best understood as a conditional efficiency gain — the efficiency you get once difficulty is known — rather than an end-to-end deployment gain.

The predicted difficulty variant (using PRM scores instead of ground-truth answers) removes the need for labels but does not reduce the sampling cost: it still requires 2,048 generations per question. The paper provides no evidence that difficulty can be estimated with substantially fewer samples, nor does it explore cheaper heuristics (e.g., using the base model's perplexity on the question, or a lightweight classifier trained on question text). A key ablation is conspicuously absent: how does compute-optimal scaling perform if difficulty is estimated from, say, 8 or 16 or 64 samples rather than 2,048? Without this, the reader cannot assess whether the 4× gains survive under a realistic difficulty estimation budget, or whether the estimation cost eliminates or reverses the advantage.

The paper partially acknowledges this gap by framing cheaper difficulty estimation as future work — Section 8 suggests "pretraining or finetuning models to directly predict difficulty of a question." But no such model is developed or evaluated, and no analysis bounds how accurate difficulty estimation must be to preserve the compute-optimal gains. Until this gap is closed, the headline efficiency numbers should be treated as upper bounds on what is achievable, and the practical deployability of the compute-optimal framework should be considered unproven.

---

Test-Time Compute Cannot Compensate for Fundamental Capability Gaps

The paper's FLOPs-matched comparison (Section 7) demonstrates that a smaller model with compute-optimal test-time strategies can outperform a ~14× larger model — but only under sharply bounded conditions. The failure case is equally important and more consequential for deployment decisions: on the hardest questions (difficulty bin 5), test-time compute provides essentially zero benefit regardless of budget, while the larger pretrained model does substantially better.

The evidence is unambiguous. Figure 3 (right) shows that on bin 5 questions, all search methods — best-of-N, beam search, and lookahead search — hover at 1–3% accuracy across all generation budgets from 4 to 256. Figure 7 (right) shows the same pattern for revisions: bin 5 accuracy is roughly 2–3% regardless of whether the budget is spent sequentially, in parallel, or in any hybrid ratio. In the FLOPs-matched comparison (Figure 9), the bin 5 scaling line for revisions is essentially flat near 0–5% across all test-time budgets, while the ~14× larger model's greedy decoding achieves non-trivial (though still low) performance on the same questions. The paper states this limitation explicitly in the Section 7 takeaway:

"On the hardest questions (bin 5), no method makes meaningful progress — the base model simply lacks the capability to produce correct solutions regardless of how the budget is allocated."

The consequence for practice is a hard boundary on when test-time compute can substitute for pretraining. If the base model's pass@1 on a problem class is near zero — meaning it essentially never produces a correct solution even when sampled thousands of times — then no amount of search or revision can help, because there are no correct solutions in the proposal distribution to find or refine. Test-time compute amplifies existing capability but does not create it from nothing. For problems genuinely outside the base model's training distribution or reasoning capacity, pretraining remains the only viable path, and the paper provides no mechanism for determining in advance whether a novel problem will fall into this regime without already having some sense of the base model's performance on it.

This limitation bounds the paper's claim that test-time compute can substitute for pretraining. The substitution works specifically for problems where the base model already has non-trivial pass@1 (easy-to-medium questions in the paper's taxonomy). For genuinely hard problems, the FLOPs-matched comparison in Figure 9 shows that the larger model's greedy decoding outperforms the smaller model with any amount of test-time compute — the test-time compute line is below the larger model's star at all budget levels. The paper's finding is therefore not "test-time compute > pretraining" but rather "test-time compute > pretraining for problems within the base model's capability envelope, and pretraining > test-time compute for problems outside it." This is a nuanced but critical boundary condition that practitioners must account for when deciding between the two investments.

---

The Revision Model Has a Fundamental Correct-to-Incorrect Reversion Problem

The revision model — fine-tuned to produce improved answers conditioned on previous incorrect attempts — exhibits a significant and partially unresolved failure mode: it frequently "revises" correct answers into incorrect ones. The paper reports in Section 6.1 that approximately 38% of correct answers produced during a revision chain get converted back to incorrect answers in the subsequent revision step. This is not a minor edge case — it is a systematic behavior that emerges from the training data construction.

The root cause is that the revision model is trained exclusively on sequences where all in-context answers are incorrect, followed by a correct target. The training data (Section 6.1) consists of trajectories containing 0–4 incorrect answers before a correct answer. At no point during training does the model see an example where the current answer is already correct and should be preserved. Consequently, the model has no learned signal for what to do when it encounters a correct answer in its context — it defaults to its trained behavior of "produce something different," which typically means introducing an error.

The paper mitigates this with a selection mechanism: rather than always taking the final revision output, the system uses majority voting or verifier-based selection across the entire chain to pick the best answer from any point in the sequence. This means a correct answer produced at revision step 3 that gets corrupted at step 4 can still be recovered if the selection mechanism correctly identifies step 3 as the best output. However, this is a post-hoc patch, not a solution to the underlying problem. The selection mechanism is imperfect — some correct answers that get reverted will not be recovered — and the approach wastes compute generating corrupted revisions that serve no purpose.

More fundamentally, the 38% reversion rate reveals a flaw in the revision model's training objective. The model is trained to always "improve" its input, but it has no concept of "the input is already correct — stop revising." A more principled solution would involve training the model to recognize when no revision is needed (e.g., by including trajectories where the correct answer appears in-context and the target is to repeat it), or by incorporating a halting mechanism that terminates the revision chain when confidence is high. The paper does not explore these alternatives, and Section 8 does not list this as future work. The revision model as presented is therefore a workable but brittle system — it achieves strong aggregate performance but has a known, substantial failure mode that is mitigated rather than resolved.

---

The ~14× Larger Model Baseline Is Not Compute-Optimal, Weakening the Pretraining Comparison

The FLOPs-matched comparison in Section 7 — which provides the paper's headline evidence that test-time compute can substitute for pretraining — uses a baseline that the paper itself acknowledges is not compute-optimal. When scaling pretraining compute by a factor of M, the paper scales only the number of model parameters, holding training data fixed. This follows the LLaMA paradigm (Touvron et al., 2023) rather than Chinchilla-optimal scaling (Hoffmann et al., 2022), where both parameters and data are scaled simultaneously. The paper states this explicitly in Section 7:

"We choose this setting as it is representative of a canonical approach to scaling pretraining compute and leave the analysis of compute-optimal scaling of pretraining compute where the data and parameters are both scaled equally to future work."

The consequence is that the pretraining baseline is weaker than it could be. A Chinchilla-optimal model trained with ~14× more total FLOPs would allocate the additional compute across both more parameters and more training data, which typically yields better performance than scaling parameters alone. The reported advantages of test-time compute over pretraining — for example, +27.8% relative improvement on medium questions at $R \ll 1$ (Figure 1, top-right bar chart) — may shrink or reverse against a properly compute-optimal larger model. The paper provides no analysis of how much of the test-time advantage is attributable to the suboptimality of the pretraining baseline rather than the efficacy of test-time compute.

Additionally, the ~14× larger model uses only greedy decoding — no majority voting, no best-of-N, no verifier-based selection. The paper is comparing an inference-time-augmented small model against a raw (single-sample) large model. A fairer comparison would give the larger model at least a modest test-time compute budget — even best-of-4 or majority voting over 8 samples — since the larger model's per-token inference cost is accounted for in the FLOPs budget. The current comparison is analogous to comparing a system that spends its compute on smart inference against one that spends all its compute on training and none on inference — this stacks the deck in favor of test-time compute because it confounds "spending compute at test time vs. training time" with "spending compute at all vs. not spending it."

The paper's transparency about this limitation is commendable, and it correctly frames the current comparison as a specific point in a larger design space. However, practitioners should treat the "test-time compute outperforms 14× larger model" claim as conditional on the baseline being parameter-scaled rather than compute-optimally scaled, and as conditional on the larger model receiving no test-time augmentation of its own. The true exchange rate between pretraining and inference compute — where both sides are optimized under their respective best practices — remains undetermined.

---

Results Are Restricted to a Single Benchmark and Single Model Family

All experiments in the paper use a single model family (PaLM 2-S*) on a single benchmark (MATH). While the paper argues that this model is "representative of the capabilities of many contemporary LLMs" (Section 4), the claim is unverified, and several aspects of the findings have plausible model-specific and domain-specific dependencies that could limit their generality.

First, the PRM's quality and over-optimization behavior are functions of the base model's output distribution. The PRM is trained on Monte Carlo rollouts from PaLM 2-S* (Section 5.1, Appendix D), and its ability to discriminate correct from incorrect solutions depends on how clearly the model's errors manifest in its intermediate steps. A model with different error patterns — for example, one that makes subtle logical errors rather than arithmetic mistakes, or one whose errors are concentrated in early vs. late solution steps — might produce a PRM with different calibration properties and different over-optimization thresholds. The paper's central finding that beam search over-optimizes on easy problems (Figure 3, right) might not generalize to verifiers trained on other base models with different failure modes.

Second, the revision model's ability to learn from incorrect in-context examples depends on the base model's in-context learning capabilities, which vary substantially across model families. PaLM 2-S* may be particularly strong at in-context learning, making the revision fine-tuning more effective than it would be on a model with weaker in-context abilities. Conversely, weaker base models might benefit more from revisions because they make more correctable errors — the relationship is not obvious a priori, and the paper provides no evidence either way.

Third, the MATH benchmark consists exclusively of competition-level math problems requiring multi-step symbolic reasoning. It is unclear whether the difficulty-dependent patterns — beam search hurting easy problems (due to verifier over-optimization), revisions helping easy problems (due to local refinement of near-correct answers), and no method helping hard problems (due to capability ceilings) — generalize to other reasoning domains. Code generation tasks, for example, have different error distributions (syntax errors, semantic errors, edge-case handling), and PRM-style step-level verification might behave differently when correctness can be partially checked by execution. Factual QA tasks, by contrast, involve knowledge recall rather than multi-step inference, and it is not obvious that the "proposal distribution vs. verifier" framework applies in the same way — the failure mode on hard factual questions might be that the model has never encountered the relevant knowledge, making both search and revision ineffective for a different reason than the MATH hard-problem failure mode.

The paper's structural analysis of pruned models (Figure 3, Table 6) provides indirect evidence that the approach can discover dataset-specific optimal architectures — different GLUE tasks produce different layer survival patterns. This strengthens the case that CoFi is not overfit to a single task type. However, GLUE tasks are all text classification or similarity tasks, which share a common underlying structure (encode text, pool representations, classify). The more significant generalisation question is whether the compute-optimal allocation framework — where difficulty estimation drives strategy selection — transfers to non-MATH tasks with fundamentally different error and difficulty structures. The paper provides no evidence on this question, and Section 8 does not flag cross-domain evaluation as future work. Practitioners considering applying this framework to other domains should therefore treat the difficulty-dependent strategy selection patterns as domain-specific until proven otherwise — the thresholds at which beam search becomes preferable to best-of-N, or at which sequential revisions become preferable to parallel sampling, may be entirely different (or not exist at all) for other task types.

---

The Revision and Search Mechanisms Are Evaluated Independently, Not Combined

The paper studies two complementary axes of test-time compute allocation — PRM-guided search (modifying the verifier/selection component) and iterative revisions (modifying the proposal distribution) — but never combines them into a single system. Section 8 explicitly acknowledges this gap:

"we did not experiment with PRM tree-search techniques in combination with revisions"

This is a significant omission because the two mechanisms have complementary, difficulty-dependent strengths that suggest combination could yield gains beyond either method alone. Specifically, the paper's own analysis shows that:

• Revisions are most effective on easy problems (Figure 7, right: bin 1–2 questions perform best with purely sequential revisions), where the model's initial output is roughly correct and just needs local refinement.
• PRM search is most effective on medium-hard problems (Figure 3, right: bin 3–4 questions benefit most from beam search over best-of-N), where the model needs to explore qualitatively different solution strategies and benefit from verifier-guided selection.
• Neither method helps on the hardest problems (bin 5), but for different reasons: revisions cannot refine an answer that never had a correct component; search cannot find a correct answer that never appears in the proposal distribution.

A combined system could, in principle, route each problem to the mechanism best suited to its difficulty — or use both mechanisms in sequence, with revisions generating higher-quality candidates and PRM search selecting among them. The paper's current results represent a lower bound on what a fully integrated system could achieve, and the 4× efficiency gains from each mechanism individually might compound or at least be partially additive when combined.

The consequence is that the paper's central finding — that compute-optimal allocation yields 4× efficiency gains — is measured against a baseline of using each mechanism in isolation, not against a baseline that combines them optimally. The true upper bound on what compute-optimal test-time scaling can achieve, when all available mechanisms are deployed adaptively per-prompt, is unknown. This also means the paper does not resolve the question of whether the mechanisms are substitutes (both help, but combining them brings diminishing returns) or complements (each helps on different problem types, and combining them yields near-additive gains). Evidence on this question would have substantially strengthened the paper's practical guidance: if the mechanisms are largely substitutes, practitioners can choose whichever is easier to implement; if they are complements, both are worth the engineering investment.

The paper lists combining search and revisions as future work (Section 8), but provides no preliminary analysis, no speculation on how they might interact, and no discussion of the engineering challenges involved (e.g., the PRM is trained on base model outputs, and the paper shows in Appendix J that it underperforms on revision model outputs due to distribution shift — a combined system would need a verifier that generalizes across both proposal distributions, or separate verifiers for each). A practitioner reading the paper in its current form would not know whether to implement one mechanism, both independently, or an integrated system, because the paper provides no evidence on the interaction between them.

7. Implications and Future Directions

How This Work Changes the Landscape

CoFi changes the landscape of model compression by reframing the speedup gap between pruning and distillation as an optimization problem rather than an inherent limitation of pruning. Before CoFi, the field implicitly accepted a tradeoff: structured pruning could find sparse subnetworks but could not achieve competitive inference speedups (Block Pruning achieves 2.7× speedup vs. TinyBERT's 11.4× in Table 1), while distillation could achieve large speedups but required prohibitively expensive general distillation on massive unlabeled corpora (~350 GPU hours for TinyBERT). This was treated as a fundamental property of the two paradigms — pruning operates within a fixed depth, distillation designs shallow networks from scratch.

CoFi demonstrates that this tradeoff is not fundamental. The key diagnostic evidence is Table 4: at identical sparsity (95%, ~5M parameters), a model pruned with only fine-grained masks achieves 7.2× speedup on QNLI, while adding explicit layer masks pushes this to 12.1× — with no change in parameter count. The two models have identical compression rates but dramatically different speedups because the layer masks allow the optimizer to discover that removing entire layers is a viable path to satisfying the sparsity constraint. Without them, the optimizer gets stuck in local minima where every layer retains some units, producing deep-but-thin networks that GPUs cannot accelerate efficiently. This finding shows that the speedup ceiling of prior structured pruning was a failure of the optimization landscape, not an inherent architectural constraint.

This insight reverses the prevailing trend in the structured pruning literature, which had moved toward increasingly fine-grained units (heads → FFN dimensions → blocks) under the assumption that smaller pruning units give the optimizer more flexibility and therefore better accuracy. CoFi shows this assumption is counterproductive at high sparsity: fine granularity alone restricts the optimizer's ability to make coarse architectural decisions that are essential for hardware-efficient inference. The paper's contribution is not merely adding layer masks to an existing pruning framework — it is the diagnosis that multi-granularity pruning is necessary because fine-grained pruning creates optimization barriers to coarse-grained decisions, not because coarse-grained decisions themselves are inherently better. The implication is that pruning methods should co-design their mask granularity with the target hardware's execution model, rather than treating sparsity as a purely statistical property.

More broadly, CoFi resolves the apparent contradiction between pruning and distillation as compression paradigms. Prior work suggested these were competing approaches with different strengths — pruning for flexibility, distillation for speed. CoFi shows they can be unified: task-specific structured pruning, when equipped with multi-granularity masks and dynamic layerwise distillation, can achieve distillation-level speedups with pruning-level data efficiency. This reframes the choice facing practitioners from "pruning or distillation?" to "task-specific pruning from a pretrained model or general distillation of a compact architecture?", where the former now dominates on both cost and accuracy for many deployment scenarios.

The work also elevates structured pruning from a model compression technique to an architecture discovery mechanism. The structural analysis in Figure 3 and Table 6 reveals that CoFi discovers dataset-specific optimal subnetworks — different GLUE tasks produce different layer survival patterns, different head retention profiles, and different FFN width allocations. This suggests that pruning, when properly instrumented, can serve as a form of task-aware neural architecture search that simultaneously discovers and trains the optimal architecture. The paper does not frame itself this way, but the implication is clear: the masks learned during CoFi training encode knowledge about which computational substructures are important for which tasks, and this information could inform architectural design independent of compression.

On the negative side, this work calls into question the value of approaches that prune at only one granularity — layer pruning alone (too coarse, misses fine-grained redundancy), head pruning alone (minimal speedup, 1.4× at best per Li et al., 2021), FFN dimension pruning alone (retains deep structures, limits speedup to ~7× at high sparsity), and block pruning (cannot exceed ~3× speedup because it never removes entire layers). These methods are not invalidated, but CoFi demonstrates they each leave substantial efficiency on the table by restricting the optimization's degrees of freedom. The path forward is clearly toward compositional, multi-granularity mask architectures that give the optimizer explicit pathways to prune at every granularity that maps to hardware speedup.

Finally, the paper's training cost analysis (Table 2) shifts the economic calculus for model compression in resource-constrained settings. The 18× reduction in training time (≤20 GPU hours vs. ~350) makes extreme compression accessible to practitioners who cannot afford the computational overhead of general distillation — academic labs, startups, and organizations working with low-resource languages or domains where large unlabeled corpora are unavailable. This democratization of model compression is a practical impact that may exceed the technical contribution's significance.

Follow-Up Research This Work Enables

Scaling CoFi to larger models and verifying whether the granularity-speed relationship holds at scale. The paper evaluates only on BERT_base (85M parameters excluding embeddings). A natural extension is applying CoFi to BERT_large, T5, or GPT-family models with hundreds of millions to billions of parameters. The key open question is whether the same granularity-speed relationship persists at larger scales: do larger models have proportionally more redundant layers (making layer pruning even more important), or do they distribute redundancy more evenly across all granularities (making the multi-granularity approach less critical)? The experiment would replicate CoFi on T5_base or T5_large, measuring speedup at matched sparsity with and without layer masks, and comparing the pruned architectures to manually designed compact variants (e.g., T5_small) trained with the same compute budget. A negative result — where layer pruning provides diminishing returns at larger scales — would refine our understanding of when multi-granularity masks are necessary versus when fine-grained pruning alone suffices.

Training CoFi-discovered architectures from scratch to isolate the value of weight inheritance. The paper never tests whether the architectures discovered by CoFi could be trained from random initialization (with standard task-specific fine-tuning) and achieve comparable accuracy. In computer vision, Liu et al. (2019c) famously showed that training pruned architectures from scratch can match the performance of weight inheritance, challenging the value of the pruning process itself. For Transformers, this remains untested. The experiment would take the exact surviving layer/head/dimension configuration from a CoFi-pruned model at 95% sparsity on MNLI, initialize it randomly (or from BERT_base's pretrained weights at the surviving dimensions — a "subnetwork extraction"), and train it on MNLI with the same task-specific data and distillation objectives. If the from-scratch model matches CoFi's accuracy, the value of CoFi reduces to architecture search; if CoFi substantially outperforms, the mask-guided optimization process provides benefits beyond architecture discovery (e.g., better weight initialization, smoother optimization trajectory). The ablation would also clarify whether CoFi's training cost advantage over general distillation persists when architecturally-equivalent models are trained from scratch, or whether the advantage comes entirely from reusing pretrained weights.

Extending dynamic layerwise distillation to settings where the teacher and student architectures differ substantially. CoFi's dynamic matching currently operates under the constraint that student and teacher have the same hidden size (both 768), with a learned projection matrix $W_{\text{layer}}$ handling representational misalignment. This could be extended to cross-architecture distillation — e.g., using a 12-layer BERT_base teacher to supervise a 6-layer student with a different hidden size, where the dynamic matching selects which teacher layers to use and learns a more expressive mapping (e.g., a small MLP rather than a linear projection) to handle dimension mismatch. This would allow CoFi to function as a general-purpose distillation framework that requires no pre-specified layer mapping and no fixed student architecture, adapting the teacher's knowledge transfer to whatever student structure emerges from pruning. The experiment would test whether dynamic matching outperforms hand-designed layer mappings in cross-architecture settings (e.g., BERT → a randomly initialized compact Transformer), and whether the learned mappings correspond to semantically meaningful alignment (e.g., lower teacher layers matching lower student layers for syntactic features, upper layers matching for semantic features).

Characterizing the over-optimization boundary for structured pruning verifiers. The paper's search experiments (Section 5.3, Figure 3) identify verifier over-optimization as the primary bottleneck limiting test-time compute scaling, but provide only qualitative evidence (degenerate outputs in Appendix M, repetitive low-information steps). A rigorous characterization would systematically measure how PRM score distributions diverge from actual correctness as search budget increases, quantifying the point at which verifier score becomes anti-correlated with ground-truth accuracy per difficulty bin. The experiment would plot PRM score vs. actual correctness for beam search at budgets from 4 to 512 generations, stratified by difficulty, and identify the budget threshold where the correlation flips. This would provide a quantitative over-optimization frontier that could guide when to stop searching versus when to switch to a different strategy, and would inform whether improving the PRM (better training, ensembles) or constraining the search (KL penalties, entropy regularization) is the more promising path to pushing the scaling frontier further.

Evaluating task-agnostic CoFi pruning (upstream pruning) and measuring transferability. The paper restricts CoFi to task-specific pruning, where masks are learned on the target task's labeled data. A natural extension is upstream pruning: apply CoFi during masked language model pretraining (or immediately after, on a large unlabeled corpus with the MLM objective), produce a single pruned model with a fixed compact architecture, and then fine-tune that architecture on downstream tasks. The key question is whether the optimal subnetwork for language modeling transfers to downstream tasks — or whether task-specific pruning discovers fundamentally different architectures that cannot be found upstream. The experiment would run CoFi on BERT_base during continued MLM pretraining (on, say, Wikipedia + BooksCorpus) to achieve 60–95% sparsity, then fine-tune the resulting fixed architecture on GLUE tasks without further pruning, and compare against task-specific CoFi at matched sparsity. A positive result — upstream CoFi matching or approaching task-specific CoFi — would produce a single compact model deployable across many tasks, addressing the amortization concern (single upfront cost vs. per-task cost) and making CoFi competitive with general distillation in multi-task deployment scenarios. A negative result — large gaps between upstream and task-specific pruning — would confirm that optimal subnetworks are task-specific and that per-task pruning is necessary for extreme compression.

Porting CoFi to encoder-decoder architectures and generation tasks. All CoFi experiments use BERT_base, an encoder-only model evaluated on classification and span-extraction tasks. Generation tasks (summarization, translation, dialogue) impose different computational bottlenecks — cross-attention dominates, decoder layers may be more or less redundant than encoder layers, and the optimal granularity-speed relationship may differ. The experiment would apply CoFi to T5 or BART on summarization (CNN/DailyMail) or translation (WMT), introducing masks for cross-attention layers and decoder-specific submodules. The key metric would be speedup on autoregressive generation (where KV-caching changes the latency profile) vs. encoder-only speedup on GLUE. This would determine whether the paper's central findings — FFN > MHA redundancy, upper layers > lower layers redundancy, multi-granularity masks as necessary for high speedup — generalize beyond the encoder-only classification setting or are architecture-specific.

Practical Applications and Downstream Use Cases

Cost-efficient deployment of NLP models on edge devices and in low-resource languages. CoFi's ability to produce models with 10×+ speedup and 95% sparsity from only task-specific labeled data (≤20 GPU hours total) makes it directly applicable to scenarios where general distillation is infeasible. For edge deployment — running NLP models on smartphones, IoT devices, or embedded systems — the 12× speedup on a V100 GPU translates to substantially lower latency and energy consumption on mobile CPUs. For low-resource languages where no large unlabeled corpus exists for general distillation (TinyBERT requires 2,500M tokens of unlabeled data), CoFi can still produce compact models using only the task-specific labeled data that practitioners already need to collect. The paper's results on small GLUE datasets (RTE: 2.5k examples, MRPC: 3.7k, CoLA: 8.5k, STS-B: 7k) in Table 2 demonstrate this directly: CoFi achieves 64.7% on RTE, 82.6% on MRPC, 35.6 on CoLA, and 83.1 on STS-B at ~12× speedup, while TinyBERT_4 without unlabeled data collapses on CoLA (16.6) and STS-B (17.8). For any practitioner deploying NLP models where unlabeled data is scarce or unavailable, CoFi is the clearly preferable compression method.

Rapid iteration on model compression during task development. The 18× training time reduction (≤20 GPU hours vs. ~350 for general distillation) makes CoFi practical for iterative development cycles where a practitioner needs to quickly evaluate multiple compression targets or hyperparameter settings. A researcher developing a new NLP task can: (1) fine-tune BERT_base on the task data (standard practice), (2) run CoFi at multiple target sparsities (60%, 80%, 95%) in parallel to determine the accuracy-speedup Pareto frontier for that specific task, and (3) select the pruned model that meets their deployment constraints — all within a single day on one GPU. With general distillation, each compression target would require retraining the generally distilled student model (350 GPU hours) before task-specific fine-tuning could begin, making such exploration prohibitively expensive. This changes the workflow from "guess a compression target, wait 3.5 days, evaluate" to "sweep compression targets, evaluate same day, deploy."

Architecture discovery for designing efficient Transformer variants. CoFi's structural analysis (Figure 3, Table 6) reveals consistent patterns — FFN layers are substantially more redundant than MHA layers (71% vs. 39% dimension reduction at 60% sparsity), upper layers are pruned more aggressively than lower layers, and the optimal subnetwork architecture varies by dataset. These patterns can inform the design of hand-crafted efficient architectures: if FFN layers are systematically overparameterized relative to attention for downstream tasks, future Transformer variants might benefit from thinner FFN layers relative to attention at the same total parameter budget, or from different layer counts for the two sublayer types. The paper's finding that pruned models often contain more MHA layers than FFN layers at high sparsity (Appendix H, Figure 5: ~6 MHA vs. ~4 FFN layers remaining at 95% sparsity) suggests that an asymmetric architecture with, say, 8 MHA layers and 4 FFN layers might outperform a symmetric 6-layer architecture at the same parameter count. CoFi serves as an empirical probe: run it on a new domain, observe which submodules are preserved at high sparsity, and use those patterns to guide hand-crafted architecture design for that domain. This is speculative — the paper does not test this — but the structural analysis provides the empirical foundation for it.