PaperlayerView exact revision on arXiv

Source-linked AI summary

Theoretical Guarantees for One-Shot Magnitude Pruning and Compute-Adaptive Early Exit

Erdem Koyuncu

arXiv:2609.12337v1cs.LG

TL;DR

The paper addresses the limited principled explanation for compute savings from conditional computation. It develops theory for one-shot magnitude pruning and early exit, finding compute-accuracy tradeoffs tied to pruning distortion and exit-to-final alignment.

  • Problem

    Conditional computation can save compute while maintaining accuracy, but principled explanations of when and why remain limited.

    Source: §1 · p. 1

  • Method

    The paper analyzes one-shot pruning and early exit in single-neuron and deep-network models, including conditional perceptrons and infinite-width Gaussian coupling.

    Source: §1 · pp. 1–2

  • Results

    The theory characterizes pruning distortion and its depth-wise accumulation, while early-exit excess error decays as a power of compute gap with an exponent controlled by alignment.

    Source: §1 · pp. 1–2, §7 · p. 9

  • Takeaways & Limitations

    The analyses provide compute-accuracy tradeoffs for one-shot pruning and early exit, with early-exit exponents increasing as partial and full computations align.

    Source: §1 · pp. 1–2

    IID weights and infinite width provide idealized reference settings, not literal models of trained finite networks.

    Source: §7 · p. 9

Abstract

from arXiv · show

We study compute reduction in neural networks through a unified partial versus full computation view, captured by one-shot magnitude pruning in the static regime and early exit in the adaptive regime. In an asymptotic single-neuron model, we prove a concentration theorem for one-shot magnitude pruning with explicit rates. We also introduce the conditional perceptron for early exit and show that its excess generalization error decays as a power of the compute gap, with an exponent that grows to infinity as the alignment between partial and full computations tends to one. We then extend the analysis to deep networks, characterizing how pruning-induced distortions accumulate with depth and deriving a corresponding compute-accuracy tradeoff for frozen-backbone early exit under a neural network Gaussian process model. Numerical simulations corroborate the predicted scaling laws.

1. Introduction

Principled explanations of compute–accuracy tradeoffs in conditional computation remain limited despite compute savings with maintained accuracy. The paper analyzes partial and full computations across static pruning and adaptive early exit, from single neurons to deep networks.

1. Introduction

  • The paper models conditional computation as partial computation used either for every input in static pruning or selectively before full computation in adaptive early exit.It analyzes these regimes first for a single neuron and then for multi-layer networks.Source: §1 · p. 1
  • For single neurons, one-shot magnitude pruning has a pruning-rate-dependent limiting similarity with explicit convergence rates, while conditional perceptrons have an alignment-controlled power-law compute–error tradeoff.The excess generalization error exponent diverges as partial and full computations become aligned.Source: §1 · pp. 1–2
  • For deep networks, the analysis tracks how pruning misalignment accumulates with depth and derives an alignment-controlled early-exit tradeoff through an effective correlation parameter.The early-exit analysis considers frozen backbones with distilled intermediate exits.Source: §1 · p. 2
  • The theory isolates tractable mechanisms rather than aiming to describe every trained architecture quantitatively, and experiments test whether the predicted trends persist in pretrained models.The stated mechanisms include alignment loss, depth accumulation, and activation dependence.Source: §1 · p. 2

2. Static Neurons

For a single neuron, one-shot magnitude pruning produces a deterministic, pruning-rate-dependent similarity limit in high dimensions, with explicit finite-dimensional convergence rates. Renormalizing the pruned vector removes its norm loss while retaining this directional similarity.

2. Static Neurons

  • One-shot pruning removes the ⌈qn⌉ smallest-magnitude coordinates of w; renormalization rescales the resulting vector to match the original weight norm.The partial computation may use either the pruned vector or its renormalized version.Source: §2 · p. 2
  • Under isotropic weight models, the cosine similarity between w and its pruned counterparts concentrates around a deterministic pruning-rate-dependent limit, with an effective finite-n deviation bound and explicit convergence rates.The theorem states convergence in probability to τq for the associated order-statistic ratio.Source: §2 · p. 2
  • At pruning rate q = 0.4, keeping 60% of coordinates gives wᵀŵp ≈ 0.982, corresponding to an angle of about 10.8°.The limit for the renormalized vector is √τq; for the unrenormalized pruned vector, the inner product limit is τq.Source: §2 · p. 3
  • The concentration framework also extends to blockwise semi-structured N:M pruning, with limiting similarities derived in closed form.Those limiting similarities are given in Appendix A.5.Source: §2 · p. 3

3. Adaptive Neurons

The adaptive single-neuron analysis uses a partial computation as a confidence test and invokes the full computation when that test fails. In a student–teacher setting, it establishes an alignment-dependent compute–generalization tradeoff for this conditional perceptron.

3.1. Conditional Perceptrons

  • The conditional perceptron first computes the renormalized pruned field and exits on its sign when its magnitude reaches the threshold; otherwise, it computes the full decision.The early-exit branch costs 2(n − k) + 1 FLOPs, while the fallback costs 2n + 2 FLOPs.Source: §3.1 · p. 3

3.2. Learning on the Unconditional Perceptron

  • The student is evaluated against teacher labels in a student–teacher setting; for a uniform-sphere test input, the unconditional student’s error depends on its angle to the teacher.The unconditional error is given by arccos(sim(w, t))/π; the passage also describes a max-margin student and its large-sample setting.Source: §3.2 · p. 4

3.3. Analyzing the Conditional Perceptron

  • As normalized compute approaches the unconditional baseline, the conditional perceptron’s excess error decays polynomially in the compute gap, with an exponent that diverges as alignment approaches one.The result is an explicit achievability bound, not a tight characterization.Source: §3.3 · p. 4
  • The bound approaches full-compute performance at a sharper rate when pruning preserves stronger alignment and the unconditional student already generalizes well.The stated bound separates teacher–student overlap from student–pruned-student overlap using an angular triangle inequality; directly characterizing their joint overlap is difficult.Source: §3.3 · p. 4
  • The theorem provides an achievability scaling law rather than a tight characterization because pruning support depends on student coordinates correlated with the teacher through training.A sharper result would require controlling the joint law of student coordinates, teacher coordinates, and the magnitude-selected support.Source: §3.3 · p. 4

4. Static Networks

The section extends one-shot magnitude-pruning analysis to fully connected deep networks, where pruning-induced mismatch can compound with depth. In the wide limit, deterministic kernel recursions characterize this distortion and show how it depends on the activation.

4. Static Networks

  • Even mild pruning mismatch at one layer can propagate through nonlinearities and compound across layers in networks pruned without retraining.The analysis targets neuronwise magnitude pruning.Source: §4 · p. 4, §4.2 · p. 5

4.1. Pruning with Renormalization in Networks

  • The network model uses IID, zero-mean, unit-variance layer weights, while renormalizing pruned weights by 1/√τq restores signal scale and isolates directional mismatch.The weights are assumed to have finite fourth moments; the readout is IID, zero-mean, and independent of the layer weights.Source: §4.1 · pp. 4–5

4.2. Main Result with Applications

  • With widths tending to infinity at fixed depth, the network’s empirical kernels converge with probability tending to one to deterministic limits governed by layerwise recursions.The result covers kernels and cross-kernels for fixed inputs.Source: §4.2 · p. 5
  • Renormalized-pruned and unpruned networks share limiting self-kernels, so their output distortion comes from the cross-kernel gap rather than shrinking signal energy.Under common unit self-kernel normalization, the limiting MSE reduces to 2(1 − αL).Source: §4.2 · pp. 5–6
  • The same pruning ratio produces activation-dependent depth accumulation: the sign recursion is highly sensitive near perfect alignment, while normalized ReLU has a smoother map and milder per-layer alignment loss.The paper reports numerical validation of these predictions.Source: §4.2 · pp. 5–6

5. Adaptive Networks

The section develops frozen-backbone early exit from a distilled intermediate score, then derives a deep compute–accuracy tradeoff whose power-law exponent is governed by exit-to-final alignment. In the infinite-width random-backbone regime, the alignment and exponent have explicit depth-gap-dependent forms.

5.1. Early-Exit Model

  • The early-exit model freezes the backbone, trains a linear intermediate head by least-squares distillation to the final score, and exits when the intermediate score’s magnitude meets a threshold.The task remains binary classification; the score magnitude is a margin-like confidence statistic, and task-specific alternatives include maximum softmax probability and top-two logit margin.Source: §5.1 · p. 6

5.2. Deep compute–accuracy tradeoffs

  • The deep tradeoff uses within-class tail-control assumptions rather than an exact Gaussian model; approximate linearity or changed tail-bound constants alter the bound’s constants while preserving its qualitative power law.The scalar alignment parameter ρ_ℓ governs the exponent.Source: §5.2 · p. 7
  • A compute gap yields excess generalization error that decays as a power of the required exit fraction, with exponent A′ = 0.5/(1 − ρ_ℓ^2) increasing sharply as within-class alignment ρ_ℓ approaches one.The deep theorem’s alignment parameter depends on the frozen backbone and the gap between exit depth and final depth.Source: §5.2 · p. 7
  • In the infinite-width random-backbone limit, the distilled exit-to-final correlation is χ_ϕ^(L−ℓ), giving exponent A′ = 0.5/(1 − χ_ϕ^(2(L−ℓ))).For χ_ϕ between zero and one, earlier exits have lower correlation and a less steep tradeoff; exits nearer the final layer can have a large exponent.Source: §5.2 · p. 8
  • The infinite-width result gives a closed-form characterization of the true within-class alignment, with a limit independent of class label and concentration of the distilled correlation.The explicit exponent follows by substituting this alignment into the deep tradeoff theorem.Source: §5.2 · pp. 7–8

6. Numerical Results

The experiments report the predicted scaling patterns across pruning and early exit: pruning distortion concentrates and varies with activation, while excess error in early exit follows approximate power laws. Deeper exits show sharper decay in both the DistilBERT and ViT experiments.

Measure concentration

  • At q = 0.9, simulated PDFs concentrate around √τ_0.9 ≃ 0.66279 as dimension increases from 10 to 10,000, consistent with the theoretical limit.The simulations use 100,000 trials.Source: Figure 2 · p. 8, §6 · p. 8

Early exit on DistilBERT

  • The DistilBERT early-exit curves are approximately linear on log-log axes, and theory and experiment both show steeper slopes for more aligned exits, especially the deepest exit.For ℓ = 5, the predicted and empirical slopes are 11.00 and 11.75, respectively.Source: Figure 4 · p. 8, §6 · p. 8

Multi-layer pruning on CNNs

  • In the CIFAR-10 CNN experiments, signum networks distort rapidly and lose accuracy at mild sparsity, while ReLU VGGNet7 maintains accuracy through moderate sparsity and MobileNetV2 degrades earlier.The observed activation ordering matches the analytical estimate; without renormalization, ReLU-type outputs are also attenuated at high sparsity.Source: §6 · p. 8, Figure 3 · p. 9

Early exit on ViT-Base

  • For ViT-Base on CIFAR-100, excess error scales approximately as a power of the compute gap, with deeper exits showing sharper decay.The analytical exponents for exits ℓ = 3, 6, 9, 11 are approximately 0.53, 0.60, 0.77, 1.24; the ℓ = 11 empirical slope is unstable at finite test-set resolution.Source: Figure 5 · p. 9, §6 · p. 9
  • The ViT theory–experiment slope gap is larger than in the binary DistilBERT case, and the authors suggest a multiclass correction may be needed for a confidence rule based on 100-class maximum softmax.The empirical ℓ = 11 slope is especially unstable because finite test-set resolution quantizes the curve’s left edge.Source: §6 · p. 9

7. Conclusions

The paper develops a tractable theory of compute reduction across one-shot pruning and early exit, from single neurons to deep networks. Its assumptions define idealized reference settings, and the authors identify extensions to trained backbones, multiple exits, and alternative pruning rules as future work.

7. Conclusions

  • IID weights and infinite width provide idealized reference ensembles, not literal models of trained finite networks.One-shot pruning, frozen backbones, and distilled exit heads specify regimes of practical interest.Source: §7 · p. 9
  • The Gaussian-type score coupling is a surrogate tail model, not an exact distributional claim.The authors assess it with within-class diagnostics.Source: §7 · p. 9
  • The theory links one-shot pruning and early exit to alignment loss, depth accumulation, and activation dependence across single neurons and deep networks.Experiments support the qualitative predictions.Source: §7 · p. 9
  • Future work includes trained backbones, multiple exits, and comparisons of alternative one-shot pruning rules under the same concentration framework.Source: §7 · p. 9

A.1. Some auxiliary results on Gamma random variables

The appendix assembles auxiliary probability results used in the pruning analysis, including bounds for truncated Gamma variables and concentration around a Gamma quantile.

A.1. Some auxiliary results on Gamma random variables

  • The appendix defines truncated Gamma variables by conditioning on values below or above a threshold, and uses their moments in auxiliary bounds.Source: §A.1 · p. 13
  • For truncated Gamma variables, the upper conditional variance lies between min{s, 1}θ^2 and max{1, s}θ^2 for shape s and scale θ.Source: §A.1 · p. 16
  • A corollary bounds a linear combination of truncated-Gamma variances when y is sufficiently close to τ, for large n.Source: §A.1 · p. 16

A.2. Other auxiliary results

  • For k = ⌈nq⌉, the appendix establishes effective concentration of the Gamma order statistic G(k) around the quantile ωq, using Hoeffding’s inequality on binomial counts.Source: §A.2 · p. 18

A.3. Proof of the Theorem 2.1

The appendix proves the pruning concentration result and extends its scope to normalized spherical vectors and blockwise N:M pruning.

A.3. Proof of the Theorem 2.1

  • The proof bounds deviations on both sides of a finite-n proxy threshold, then transfers the bounds to the target threshold.Source: §A.3 · pp. 20–21

A.4. Proof of Corollary 2.2

  • For a uniformly random spherical vector, pruning by retaining the largest-magnitude coordinates yields inner-product convergence through Theorem 2.1.Source: §A.4 · p. 22

A.5. Semi-structured N :M pruning

  • With fixed block size M and N retained entries per block, blockwise N:M pruning has its own constant τN:M, analogous to the unstructured-pruning constant τq.Source: §A.5 · p. 23
  • When block size grows proportionally with ambient dimension, the semi-structured limit returns to the unstructured constant τq.For 2:4 pruning, the stated retained-energy constant is τ2:4 ≈ 0.86755 and the renormalized similarity limit is √τ2:4 ≈ 0.93142.Source: §A.5 · pp. 24–25

B. Proof of Theorem 3.1

The proof analyzes conditional-perceptron generalization error by separating teacher-student alignment, student-exit alignment, and the exit criterion. It also derives normalized compute accounting for early exit.

B. Proof of Theorem 3.1

C.1. Proof of Theorem 4.1

The proof couples unpruned and renormalized-pruned networks through joint Gaussian limits, then tracks their kernels across layers. The appendix also gives activation-specific examples and contrasts renormalized with non-renormalized pruning.

C.1. Proof of Theorem 4.1

  • The proof tracks paired pre-activations to obtain a joint Gaussian limit, then uses kernel updates and induction to establish the deep-network limits.The induction applies to a fixed finite input set as widths grow with depth fixed.Source: §C.1 · pp. 27–29

C.2. Proof of Example 4.1; C.3. Proof of Example 4.2

  • For sign and normalized ReLU activations, the appendix derives example-specific alignment recursions and gives asymptotic MSE 2(1 − α_L).The sign case uses the correlated-Gaussian sign identity; the ReLU case uses the Gaussian ReLU kernel formula.Source: §C.2 · p. 30, §C.3 · p. 30

C.4. Non-renormalized pruning recursions

  • Without renormalization, pruning attenuates the pruned self-kernel by τ_q per layer and changes the unnormalized row-overlap coefficient to τ_q.With renormalization, the self-kernel stays on the original scale and the overlap coefficient is √τ_q.Source: §C.4 · p. 31
  • For sign activation, positive-rescaling invariance leaves the one-dimensional sign recursion unchanged without renormalization.The non-renormalized covariance still has τ_q in its cross entry, but converting to Gaussian correlation recovers the same recursion.Source: §C.4 · p. 31

Define the normalized cross-alignment

  • Normalizing cross-alignment removes output-scale attenuation, leaving the alignment loss caused by pruning as the distortion measure.The non-renormalized pruned output variance is attenuated by τ_q^L, unlike the unit-variance renormalized model.Source: Define the normalized cross-alignment · p. 31

D. Threshold calibration under a compute budget

The section calibrates early-exit thresholds to meet a target average-compute budget, including an analytic rule for the conditional perceptron. Its feasibility depends on the sparsity and requested compute gap.

D. Threshold calibration under a compute budget

  • Average compute equals 1 − (1 − μ_ℓ) times the exit probability, so meeting a target budget reduces to requiring a minimum exit frequency.The population compute-matching threshold is the largest threshold satisfying that frequency; in practice, a held-out validation quantile estimates it.Source: D. Threshold calibration under a… · pp. 31–32
  • For the conditional perceptron, the calibration map has an analytic inverse, and the feasible sparsity range is q ∈ [ϵ, 1].The confidence statistic is |v₁|, and the compute-matching threshold is obtained by enforcing the compute constraint at equality.Source: D. Threshold calibration under a… · p. 32
  • As sparsity q increases, the threshold τ_cal(q, ϵ) rises for fixed compute gap ϵ; q < ϵ is infeasible because even universal early exit cannot meet the gap.Figure 6 plots compute-matching thresholds across sparsities and compute gaps.Source: D. Threshold calibration under a… · p. 32, Figure 6 · p. 32

E. Proofs of Results on Adaptive Networks

  • The appendix begins proofs of the adaptive-network results after the calibration discussion.The supplied section passage identifies this transition but gives no further proof result.Source: E. Proofs of Results on Adaptive… · p. 32

E.1. Proof of Theorem 5.2

The proof bounds early-exit excess error by controlling exit frequency and disagreement with the full-compute prediction. It also relates exit-to-final score correlation to depth and checks finite-dimensional pruning concentration numerically.

E.1. Proof of Theorem 5.2

  • The proof splits early-exit errors into full-compute errors and disagreements on exited inputs, then controls high-confidence disagreement using tail assumptions.The exit-rate tail is bounded using a sub-Gaussian assumption, while disagreement is controlled through the residual between standardized exit and final scores.Source: §E.1 · p. 33
  • The derived bound is ϵ_c − ϵ_uc ≤ (8π)^A′ r^A′, relating excess error to the exit rate r and exponent A′.The proof simplifies the bound using r_ε ≤ 1 and A′ ≥ 1/2.Source: §E.1 · p. 34

E.2. Proof of Proposition 5.3

  • Across the L − ℓ remaining random layers, the degree-one component retains a multiplicative factor χ_ϕ in L2 at each layer.The proof uses this iterated decomposition to relate final-score variance to the span of layer-ℓ features.Source: §E.2 · p. 35

E.3. Proof of Corollary 5.4

  • Applying Theorem 5.2 conditional on the realized backbone transfers its bound using the within-class score correlation from Proposition 5.3.The proof states that the resulting bound holds with probability approaching one over the backbone draw.Source: §E.3 · p. 35

F.1. Pruning and concentration of measure

  • In Figure 7, distributions at q = 0.4 and q = 0.9 concentrate around √τ_0.4 ≃ 0.98217 and √τ_0.9 ≃ 0.66279 as dimension increases.The mean absolute deviation decays approximately as 1/d and the standard deviation as 1/√d; sharp proofs of these observed rates are left for future work.Source: §F.1 · p. 36

F.2. Semi-structured pruning at the neuron level

The experiments validate predicted semi-structured pruning constants at the neuron level and test their use in finite-width networks. They show that fixed-block patterns with the same keep ratio can have different limiting distortions.

F.2. Semi-structured pruning at the neuron level

  • For practical 2:4 pruning at n = 4096, empirical raw and renormalized means are 0.86754 and 0.93142, close to predicted τ_2:4 = 0.86755 and √τ_2:4 = 0.93143.The empirical standard deviations decay approximately as n^-1/2, with fitted slopes −0.4989 and −0.4993.Source: §F.2 · p. 36
  • For N = 0.125n and M = 0.25n, empirical means 0.92840 and 0.96354 closely match predicted τ_q = 0.92867 and √τ_q = 0.96368 at q = 0.5.The fitted standard-deviation slopes −0.5103 and −0.5111 are consistent with n^-1/2 concentration.Source: §F.2 · p. 37, Figure 9 · p. 37
  • At the same keep ratio N/M = 1/2, fixed-block patterns have different limiting constants, so distortion depends on block structure as well as retained fraction.The compared patterns are 1:2, 2:4, 4:8, and 8:16; the constants increase with block size.Source: §F.2 · p. 37, Table 1 · p. 37

F.3. Semi-structured pruning in random networks

  • Finite-width network experiments show that sign-network alignment and symmetric NMSE closely follow the 2:4 recursion prediction, while ReLU runs show larger deviations but the same qualitative block-size trend.For normalized ReLU networks at depth 5 and keep ratio 1/2, larger blocks yield larger output alignment and smaller distortion.Source: §F.3 · p. 38, Table 2 · p. 38

F.4. Compute constrained tradeoffs for the conditional perceptron

Compute-constrained sweeps show power-law-like excess-error decay for the conditional perceptron, while deep-pruning predictions vary by activation and diagnostics support increasingly aligned, lighter-tailed exit scores at greater depth.

F.4. Compute constrained tradeoffs for the conditional perceptron

  • Across tested configurations, excess generalization error falls approximately as a power of the compute gap, and conditional predictors approach unconditional error baselines quickly as compute increases.The sweeps use compute-constrained operating points obtained by binary search.Source: Figure 10 · p. 38, §F.4 · pp. 38–39
  • Optimized-frontier tangents are local rather than one fixed-q power law; theoretical and empirical slopes are steeper at smaller compute gaps.The frontier re-optimizes the operating point for each target compute gap.Source: Figure 11 · p. 39, §F.4 · p. 39

F.5. Multi-layer pruning on CIFAR-10, additional theory comparison

  • Renormalized predictions isolate alignment loss: sign loses alignment fastest, normalized ReLU slower, and ReLU6 accumulates distortion faster with saturation and greater depth.Without renormalization, ReLU-type outputs also attenuate; at extreme sparsity, raw MSE approaches 1 rather than the variance-normalized decorrelation value 2.Source: §F.5 · p. 40

F.6. Empirical diagnostic for Assumption 5.1 on DistilBERT (SST-2)

  • On DistilBERT, within-class conditional means are close to linear, alignment increases with depth, and residual tails decay approximately sub-Gaussian over a moderate deviation range.The diagnostic supports Assumption 5.1 as an upper-bound surrogate, not an exact Gaussian model.Source: §F.6 · p. 41

F.7. ViT diagnostics (CIFAR-100): conditional mean and residual tails

  • In ViT diagnostics, within-class correlation rises from ρ≈0.199 at ℓ=3 to ρ≈0.825 at ℓ=11, while residual kurtosis falls from about 4.725 to 1.736.The trends support the coupling surrogate as a useful upper-bound model, not an exact distributional claim.Source: Figure 14 · p. 42, §F.7 · p. 42

G. Related Work

  • Prior early-exit analyses focus on regret or threshold optimization, whereas this work targets compute–generalization relationships; its pruning analysis isolates one-shot distortion before recovery.The authors identify extending the theory to prune-and-retrain settings as an important next step.Source: G. Related Work · p. 43
Loading 2609.12337v1…