Walter Simoncini, Dheeraj Varghese, Gertjan J. Burghouts, Cees G.M. Snoek
9 min
Abstract
Scaling up transformers has produced models with broad capabilities, deployed across heterogeneous hardware and fluctuating workloads. Yet, model families offer only a few fixed sizes (e.g., S, B, L), forcing users to commit to a model based on hardware and latency constraints. This rigidity is often suboptimal under varying query loads. We introduce a retraining-free method that elastifies pre-trained transformers into a single model that adapts on-the-fly across a continuum of computational budgets. Our key contribution is a differentiable relaxation of the functional unit-ranking problem: unlike prior work, which optimizes a single importance weight per MLP block or attention head, we rank every individual MLP neuron and attention head jointly, on a common scale, by minimizing the model's loss under soft pruning across many sparsity targets, all at once. Across 27 vision and text transformer encoders ranging from 22M to 8B parameters, our method yields smooth degradation up to 60% sparsity and outperforms one-shot baselines under deep pruning. For example, pruning DINOv3 ViT-H+/16 to 50% sparsity removes 427M parameters while reducing ImageNet-1k linear-probing accuracy by only 0.8%, without the need for weight correction or retraining. At 60% sparsity, a pruned AugReg ViT-B/16 model achieves up to $1.80\times$ higher throughput and $1.26\times$ lower single-query latency on an A100 GPU. Elastic models are produced in minutes, with construction time scaling sub-linearly with parameter count.
Sam: But removing something is a yes-or-no decision. How do they learn scores from a decision that changes abruptly?
Alex: They use soft masks: smooth multipliers that scale each component’s output. Instead of immediately keeping or deleting a unit, the method partly turns it down. That makes the model’s loss differentiable with respect to importance scores, so gradient descent can improve them.
Sam: The phrase “without retraining” needs care here. Something is learning, so what stays fixed?
Alex: The pre-trained model weights never change. Only the importance scores change. They start from an estimate of how damaging removal might be, calculated using calibration data. Those initial estimates ignore interactions between units; optimizing the masked model’s loss lets the scores respond to units functioning together.
Sam: That explains learning one smaller model. What makes the result elastic rather than tied to a single size?
Alex: They optimize across fourteen sparsity targets at once. Sparsity here means the fraction of theoretical computation removed, not the fraction of weights set to zero. Each budget gets its own cutoff in the same ranking, and the losses are averaged to update the scores.
Sam: Why should multiple cutoffs help, beyond making the objective match deployment?
Alex: A soft mask is most sensitive near its cutoff. Far above or below that boundary, changing a score barely changes the output. Multiple cutoffs give different regions of the ranking useful gradients. The authors hypothesize that this broader signal improves learning.
Sam: And once that optimization is finished, does every new budget require another learning run?
Alex: No. At inference, they keep the highest-ranked units until the computational budget is filled, then remove the rest. They also test budgets between optimization anchors. In that test, performance declines smoothly, without an extra penalty for landing between the trained budget points.
Sam: Let’s separate the appealing mechanism from the evidence. How broad was the evaluation?
Alex: They evaluate twenty-seven vision and text encoders. Tasks include image classification, patch-level semantic segmentation, text embedding, and retrieval. The main comparisons leave weights uncorrected for every method, isolating which units each method selects.
Sam: Does the new method win even when very little is removed, or mainly when pruning gets aggressive?
Alex: Mainly when pruning gets aggressive. Existing baselines often remain competitive under shallow pruning, including SnapViT. As removal deepens, the new ranking generally degrades more smoothly. But the amount of retained performance depends strongly on the model and task.
Sam: Give us one concrete result, with the evaluation attached. Otherwise “preserves performance” could hide a lot.
Alex: For the large DINO version three vision encoder, removing half the theoretical computation costs only about eight-tenths of a percentage point in ImageNet linear-probing accuracy. That evaluation trains a linear classifier on frozen representations. It does not mean every downstream task suffers equally little.
Sam: Do the ablations support the specific claim that individual rankings matter, rather than just spending more effort optimizing?
Alex: They compare progressively finer score groupings under equal optimization compute. Accuracy rises as groups get finer, with independent neuron and head scores performing best. They also compare joint budget optimization against separate rankings learned for individual targets. Joint optimization wins at every evaluated target in that comparison.
Sam: The retrieval promise seems harder. Matching a fixed full-model index requires more than preserving a standalone benchmark score.
Alex: They test that asymmetric setting explicitly: the full model encodes the index, while pruned models encode queries. Their approach better preserves compatibility under deep pruning than the evaluated baselines. It does not eliminate the cost: substantial retrieval losses remain for some models.
Sam: And less theoretical computation is not automatically less waiting. What happened on actual hardware?
Alex: On an A100 GPU, an AugReg vision model at sixty percent sparsity reaches up to one-point-eight times the full model’s throughput. That uses compilation and GPU-friendly pruning. Single-query latency improves more modestly, and only with compilation settings that reduce launch overhead.
Sam: So the hardware conditions belong in the claim. Are there other costs hidden behind switching sizes?
Alex: Switching is fast once compiled graphs are cached, but first compilation is excluded from that switching measurement. They also need extra buffers. And hardware-friendly neuron groupings can reduce accuracy under deep pruning. Fewer operations alone are not a deployment guarantee.
Sam: What is the biggest boundary on generalizing this beyond the tested encoders?
Alex: Decoder-based text embedders remain more sensitive to deep pruning than bidirectional encoders. Larger text models are not consistently more robust either. And the authors leave autoregressive language generation and diffusion-based image generation to future work. These results do not establish elasticity for those tasks.
Sam: Who should read the full document, and where should they start?
Alex: If you deploy image or text encoders, start with Elastic Deployment, then the appendix on practical speedups. Those sections connect representation compatibility to actual execution costs. If you study pruning, read Differentiable Soft Pruning and the ablations on score granularity and joint budget optimization.
Sam: And for researchers who only need the central idea?
Alex: A model can learn what to leave out without relearning what it knows. Here, that turns a fixed encoder into a choice of operating sizes—but the usable range still depends on the task and hardware.
Sam: Keep the flexibility, and check the trade-off.