Constraint-aware Ranking-distilled Pruning (ToP)

• The paper introduces ToP, which leverages ranking-distilled token importance and coarse-to-fine L0-regularized pruning to accelerate Transformer inference while maintaining accuracy.
• The methodology combines teacher-student ranking distillation with learnable binary masks optimized under a FLOPs-aware Lagrangian framework to enforce explicit compute budget constraints.
• Empirical results on benchmarks like GLUE, SQuAD, and 20News demonstrate up to 8× FLOPs reduction and 7× CPU latency improvement with negligible or positive accuracy impacts.

Constraint-aware and Ranking-distilled Pruning (ToP) is a Transformer inference acceleration framework that enables significant reductions in computation and latency by selectively pruning uninformative tokens within each layer, under explicit compute budget constraints, while maintaining or improving task accuracy. ToP addresses the challenge of suboptimal token-importance ranking in typical self-attention-based pruning by introducing a ranking-distillation technique that propagates more reliable token rankings from the deepest layer of a teacher model into the pruned student’s early layers. The approach combines ranking-aware distillation with a coarse-to-fine layer selection scheme, leveraging improved $L_0$ regularization for efficient, differentiable mask learning. ToP is applicable to pre-trained models such as BERT and demonstrates superior FLOPs and latency reduction on benchmarks including GLUE, SQuAD, and 20News, outperforming prior token pruning and structured compression methods (Li et al., 2023).

1. Design Principles and Core Components

ToP’s architecture is built around two principal mechanisms:

• Ranking-distilled token distillation: The key insight is that attention-derived token-importance scores in deeper Transformer layers are significantly more reliable than those computed in shallow layers. However, early pruning is needed to maximize acceleration. Therefore, ToP distills the token-importance ranking produced by the last layer in a reference “teacher” model into the early layers of the student using a ranking loss (LambdaLoss) optimized for NDCG. This ensures that pruning decisions in shallow layers better reflect the preservation needs identified by the deeper, more expressive layers.
• Coarse-to-fine pruning with $L_0$-regularized masks: ToP introduces two types of learnable binary masks per Transformer layer:
  • Gate masks that indicate which layers are allowed to prune tokens, alleviating the learning burden compared to enforcing pruning at every layer.
  • Ranking masks that specify, within selected layers, how many of the lowest-ranked tokens are pruned.

Both mask types are parametrized via a hard-concrete (continuous $L_0$) distribution and are jointly optimized with the Transformer’s weights under a FLOPs-aware Lagrangian penalty ($c(M)\approx C$), ensuring that the expected computational cost matches a user-specified budget.

2. Improved $L_0$ Regularization for Layerwise Pruning

The ToP framework employs an enhanced $L_0$ regularization scheme to optimize the binary mask variables differentiably. Specifically, all mask variables $M$ (the union of gate and ranking masks) are sampled as follows: $$m = \min\left(1, \max(0, \tilde{s})\right), \quad \tilde{s} = \sigma\left( \frac{1}{\beta}\left[\ln u - \ln(1-u) + \ln\alpha \right] \right)\cdot(r-l) + l$$ where $u\sim \mathrm{Uniform}(0,1)$, $\alpha$ is a learnable parameter, $L_0$0 determines sigmoid smoothness, and $L_0$1. This hard-concrete relaxation, adapted from Louizos et al. (2018), enables effective gradient-based updates. The full training optimization is: $L_0$2 where $L_0$3 is the task loss, $L_0$4 is the expected FLOPs given mask $L_0$5, and $L_0$6 is the FLOPs constraint.

The FLOPs computation per layer with retained token count $L_0$7 is: $L_0$8 where $L_0$9 is the hidden size, $L_0$0 is FFN intermediary size, and $L_0$1 is the number of attention heads.

At inference time, mask entries are deterministically computed (thresholded), yielding a fixed layerwise pruning schedule.

3. Ranking-distilled Importance Score Propagation

ToP corrects for unreliable shallow-layer pruning by distilling the ranking of final-layer importance scores—computed as: $L_0$2 with $L_0$3 denoting the headwise self-attention matrices—from an unpruned teacher into early student layers.

For each early layer $L_0$4, the student ranks $L_0$5 are aligned to the teacher’s final ranking $L_0$6 by minimizing a LambdaLoss term: $L_0$7 directly targeting NDCG and thus incentivizing the preservation of top-$L_0$8 tokens critical for downstream accuracy.

4. Training and Inference Workflow

The ToP system operates as follows:

Training Procedure (Algorithm 1):

1. Initialize weights ($L_0$9) and mask parameters $c(M)\approx C$0.
2. For each minibatch, sample gate/rank masks via the hard-concrete estimator.
3. The student model forwards input with pruning decisions guided by these masks.
4. Compute:
   • Downstream task loss.
   • FLOPs expectation and corresponding Lagrangian penalty.
   • Ranking-distillation loss from the teacher output.
5. Backpropagate the sum of these losses, updating network weights and mask parameters jointly with the (learned) Lagrange multipliers.

Inference Procedure (Algorithm 2):

1. For each input, propagate through layers.
2. At layers with active gate masks, drop the lowest-ranked tokens according to learned masks.
3. Output is computed from the tokens retained after the final layer, with no need for extra prediction modules.

5. Experimental Results and Resource Implications

ToP is evaluated on BERT_base, RoBERTa_base, and BERT_6 architectures across datasets:

• GLUE (eight tasks, up to 256 tokens)
• SQuAD v2.0 (384 tokens)
• 20News (512 tokens)

Key metrics include FLOPs reduction, real CPU inference latency (Intel Xeon), and GPU latency (V100):

• BERT_base on GLUE: ToP achieves average %%%%26$\alpha$27%%%% FLOPs reduction, matching or improving accuracy by $c(M)\approx C$30.5 points relative to full BERT.
• CPU Latency: Realized 2.9$c(M)\approx C$4 to 7.4$c(M)\approx C$5 speedup (e.g., 84 ms $c(M)\approx C$6 29 ms on MRPC; 347 ms $c(M)\approx C$7 47 ms on 20News) with negligible or positive accuracy impact.
• GPU Latency: Typical speedup of 1.2–1.3$c(M)\approx C$8 with simple PyTorch kernels; further gains are anticipated with production optimizations.

Comparison across baselines:

Baseline	FLOPs Reduction	Accuracy Impact	Latency Impact
PoWER-BERT	$c(M)\approx C$93–5$L_0$0	–2 to –4 pts	No auxiliary overhead
LTP	$L_0$17–9$L_0$2	–3 to –5 pts	No auxiliary overhead
Transkimmer	$L_0$37–12$L_0$4	Close to baseline (variable)	+30% overhead (pred. modules)
CoFi	%%%%40$L_0$41%%%%	–2 to –5 pts (small/long-tasks)	No auxiliary overhead
DistilBERT_6	%%%%42$L_0$43%%%%	Varies	Not directly comparable
ToP	%%%%44$\alpha$45%%%%	$L_0$10.5 pt (often +)	2–7$L_0$2 CPU, no aux. models

Among these, ToP is reported as the only framework that consistently reduces raw CPU inference time by up to %%%%48$u\sim \mathrm{Uniform}(0,1)$49%%%% while maintaining or improving accuracy, with no requirement for auxiliary prediction modules or specialized hardware (Li et al., 2023).

6. Comparative Analysis and Methodological Innovations

Attention-score-based pruning methods such as PoWER-BERT and LTP deliver acceleration at the expense of accuracy, often dropping 2–5 points on GLUE. Prediction-module methods (e.g., Transkimmer, TR-BERT) mitigate accuracy loss but introduce significant latency and complexity via auxiliary MLPs. Structured pruning combined with distillation (e.g., CoFi) achieves moderate speedup with non-trivial accuracy compromise on smaller or longer-sequence tasks. Knowledge-distillation-only approaches, exemplified by DistilBERT, reduce model size but offer limited inference acceleration.

ToP’s innovation is the combined deployment of ranking-distilled masking and coarse-to-fine $L_0$5-regularized pruning, under hard computational constraints. By distilling reliable importance rankings into early layers, and by learning which layers to prune and by how much, ToP maximizes both computational efficiency and accuracy retention—without external prediction modules or hardware customization.

The approach’s reliance on differentiable mask learning via the hard-concrete distribution, and direct optimization of resource-accuracy tradeoffs through learnable Lagrangian multipliers for expected FLOPs, constitutes a methodological advance over prior discrete-pruning strategies.

7. Significance, Limitations, and Context

ToP demonstrates that it is possible to prune aggressively within standard self-attention architectures—achieving up to %%%%51$\alpha$52%%%% FLOPs and up to %%%%53$u\sim \mathrm{Uniform}(0,1)$54%%%% CPU latency reductions—while aligning or slightly surpassing the task accuracy of the unpruned models. The approach preserves deployment simplicity, as it leverages pre-existing attention infrastructure and only augments the training stage with lightweight mask sampling and ranking loss computation.

Reported GPU speedups are more modest in prototypical software, suggesting a need or opportunity for future work in kernel optimization and hardware-software co-design. A plausible implication is that ToP’s high-level pruning strategy may generalize to other architectures and sequence-processing tasks, contingent on the reliability of deep-layer token-importance scores and the computability of resource constraints.

Constraint-aware and Ranking-distilled Pruning thus stands as a benchmark for inference-time token pruning in Transformers, providing a blend of theoretical rigor, practical utility, and empirical superiority on mainstream NLP tasks (Li et al., 2023).