跳到正文
DEV Community· Prabhakar Chaudhary·· 3 小時前AI 評分47

彈性 Transformer 無需重新訓練:可微分軟剪枝為每個神經元及注意力頭排序

Elastic Transformers Without Retraining: Differentiable Soft Pruning Ranks Every Neuron and Head

AI 導讀

研究提出可微分軟剪枝,為 Transformer 編碼器的神經元和注意力頭評分排序,令單一預訓練模型可按推理運算預算調整,毋須重新訓練或更新權重。評估涵蓋 22M 至 8B 參數的 27 個視覺及文字編碼器;AugReg ViT-B/16 在 60% 稀疏度下,吞吐量最高達原密集模型的 1.80x,單次查詢延遲降低 1.26x。

正文

What the paper does

Most model families ship in a few fixed sizes: small, base, large. If you pick one for typical traffic, it is too slow at peak and wasteful when load is light. "Your Transformer Is Secretly Elastic" (Simoncini, Varghese, Burghouts and Snoek, University of Amsterdam and TNO, posted 14 October 2026) proposes turning a single pretrained transformer encoder into a model that can run at many compute budgets, chosen at inference time, with no retraining and no weight updates.

The approach is structured pruning with a better ranking. Every MLP neuron and every attention head gets an importance score. To run at a given budget, you keep the top-ranked units until their cumulative FLOP cost reaches the target fraction of the dense model, and drop the rest. The model weights never change. The entire contribution is in how the ranking is produced.

How the ranking is built

Step 1: initial scores from a Fisher approximation

Scores start from a second-order view of importance: the loss increase from removing a parameter is approximated by the diagonal of the Hessian, estimated as the empirical Fisher, meaning the mean of squared gradients over a small calibration set. Per-weight scores are averaged into per-unit scores, such as the input column and output row of an MLP neuron. For vision encoders the calibration loss is the self-supervised DINO objective, which works for supervised and self-supervised backbones alike. For text encoders it is a symmetrized in-batch InfoNCE loss. This follows earlier pruning work such as Kwon et al. (2022).

The authors list three weaknesses of these scores. The diagonal approximation ignores interactions between units, the proxy loss may not match the downstream task, and scores are not comparable across layers or between neurons and heads.

Step 2: differentiable soft pruning

The fix is to optimize the scores directly. Hard pruning is a threshold on a ranking and is not differentiable, so the paper relaxes it. Scores are stored as learnable log-values, centered by their median per unit type so neurons and heads share a common scale. For a sparsity target b, the method finds the threshold at which the kept units reach the FLOP budget, then replaces the hard keep-or-drop decision with a sigmoid of the distance from that threshold, scaled by a temperature. The soft mask multiplies MLP activations and attention head outputs in the forward pass.

The objective averages the masked model's loss across 14 sparsity targets, from 5% to 70% in 5% steps, using the same calibration batch. Only the log-scores are updated by gradient descent. Because mask gradients are largest near each threshold, different targets expose different parts of the ranking to useful gradients. After optimization, the ranking can be hard-pruned to any budget without further work.

An earlier method from the same group optimized one scaling factor per block and head using evolutionary search. That kept the neuron ordering inside each MLP block fixed, so keeping one important but low-scored neuron meant keeping every neuron ranked above it. Per-neuron scores remove that constraint.

Reported results

The evaluation covers 27 vision and text transformer encoders from 22M to 8B parameters, across image classification, semantic segmentation, text embedding and retrieval. Numbers below are from the paper:

  • Pruning DINOv3 ViT-H+/16 to 50% sparsity removes 427M parameters and costs 0.8% ImageNet-1k linear-probing accuracy, with no weight correction or retraining.
  • Degradation is smooth up to 60% sparsity, and the method beats one-shot baselines mostly under deep pruning. Under shallow pruning, baselines stay competitive.
  • On an A100, a 60%-sparse AugReg ViT-B/16 reaches up to 1.80x the dense throughput and 1.26x lower single-query latency.
  • In an ablation on ViT-B/16 k-NN accuracy, joint optimization over 14 targets beat a separate ranking per target at 20%, 40% and 60% sparsity, by 0.4, 1.9 and 3.6 points respectively.
  • Elastic models are built in minutes, and construction time grows sub-linearly with parameter count.

Why it matters for serving

The practical appeal is operational. Serving k budgets with ordinary pruning means producing, storing and switching among k separate models. For retrieval there is an additional cost, since embeddings from different pruned models may not be compatible, so you may need separate indices. A single ranking over one set of weights avoids both, and the authors note compatibility with a fixed dense index. A serving layer could then use the full model when load is low and a pruned configuration during a spike.

It also contrasts with other routes to elasticity. Nested designs such as MatFormer need a particular architecture and training from scratch, and other methods need substantial fine-tuning. Both are expensive when you only have a released checkpoint, or when you do not know its pretraining data. Skipping weight updates sidesteps that.

Tradeoffs and open questions

Scope is encoders

The results are for vision encoders and bidirectional text encoders. The authors state that deep pruning remains difficult for decoder-based text embedders, and they leave autoregressive language modeling and diffusion image generation to future work. Anyone hoping to apply this to a chat model should treat that as unproven.

Speedups are smaller than FLOP savings

A 1.80x throughput gain at 60% sparsity is real but well short of the FLOP reduction, and the single-query latency gain is only 1.26x. Structured pruning leaves smaller matrices, which GPUs do not always use efficiently, so measured gains depend on batch size and hardware. The paper includes a dedicated analysis of whether pruning translates to speedups, which is worth reading before planning capacity.

Calibration dependence

The method needs a calibration set and a loss. The paper uses DINO for vision and contrastive losses for text, and includes ablations on calibration set size and task-aligned calibration. A ranking tuned on one distribution may not hold under a very different one. Related work on importance estimation, such as Molchanov et al. (2019), has long noted that importance scores are sensitive to how they are measured.

Deployment complexity

Switching budgets at runtime means changing which units are active, which has its own kernel and memory-layout cost. The paper reports low budget-switching overhead, but a production stack with batching, caching and quantization adds interactions that a research benchmark does not test.

Closing reflection

The core idea is simple: pretrained transformers already contain a usable family of smaller models if you rank their units well. What is still unknown is whether the same ranking trick holds for decoder language models, where the authors saw sensitivity, and whether the speedups survive a real serving stack with mixed batch sizes. For this to matter in practice, someone would need to show the elastic ranking working on the models people actually deploy at scale, with latency gains that justify the added operational complexity.

Also see the paper's Hugging Face papers listing for discussion and linked resources.

來源:DEV Community · dev.to