Paper deep dive
LaPrune: Controllable Differentiable Sparsity at Million Scale
Jakub Antczak, Joanna Wojciechowicz, Łukasz Struski, Jacek Tabor
Intelligence
Status: succeeded | Model: Gemma-4-26B-A4B | Prompt: intel-v1 | Confidence: 93%
Last extracted: 8/9/2026, 2:41:40 AM
Summary
The paper introduces LaPrune, a differentiable sparse-selection layer that provides controllable sparsity at million-scale. Unlike previous methods that couple mask hardness to selection mass, LaPrune uses a normalized second-moment constraint to independently control the budget (k) and the hardness (gamma) of the mask. It preserves the exact selection mass via a LapSum barrier while allowing the mask to transition smoothly from a dense equal-mass allocation to a hard top-k selection. The method includes theoretical analysis of mask saturation and demonstrates robust scale-invariance and efficient implementation.
Entities (7)
Relation Signals (8)
Jacek Tabor → affiliatedwith → Jagiellonian University
confidence 95% · Jacek Tabor... Jagiellonian University
Jakub Antczak → affiliatedwith → Wrocław University of Science and Technology
confidence 95% · Jakub Antczak... Wrocław University of Science and Technology
Joanna Wojciechowicz → affiliatedwith → Wrocław University of Science and Technology
confidence 95% · Joanna Wojciechowicz... Wrocław University of Science and Technology
Łukasz Struski → affiliatedwith → Jagiellonian University
confidence 95% · Łukasz Struski... Jagiellonian University
LaPrune → controls → Normalized Hardness Parameter
confidence 95% · The normalized hardness parameter is invariant to score scale
LaPrune → uses → Normalized Second Moment
confidence 92% · controls the normalized second moment while preserving the selected mass.
LaPrune → extends → LapSum
confidence 90% · LaPrune, a differentiable sparse-selection layer that extends the exact LapSum soft top-k operator
LaPrune → enables → Top-k Selection
confidence 85% · Top-k selection determines which components of a sparse model remain active.
Cypher Suggestions (0)
No Cypher suggestions yet.
Abstract
Abstract:Top-$k$ selection determines which components of a sparse model remain active. Hard selection blocks gradients, while continuous relaxations often couple mask hardness to the selected mass. We introduce LaPrune, a mathematically exact-budget differentiable layer that controls the normalized second moment while preserving the selected mass. A LapSum barrier preserves the selection mass, and a normalized second-moment constraint moves the mask from a dense equal-mass allocation toward hard top-$k$ at each budget. We derive a population prediction of the saturated fraction, a near-binary limiting law, and a tight worst-case guarantee on the near-zero fraction. The normalized hardness parameter is invariant to score scale, while a fixed LapSum temperature is not.
Tags
Links
- Source: https://arxiv.org/abs/2608.04057v1
- Canonical: https://arxiv.org/abs/2608.04057v1
PDF not stored locally. Use the link above to view on the source site.
Full Text
58,627 characters extracted from source content.
Expand or collapse full text
LaPrune: Controllable Differentiable Sparsity at Million Scale Jakub Antczak, Joanna Wojciechowicz11footnotemark: 122footnotemark: 2, Łukasz Struski, Jacek Tabor33footnotemark: 3 These authors contributed equally to this work.Wrocław University of Science and Technology, Wrocław, Poland. Jakub Antczak: 268745@student.pwr.edu.pl; Joanna Wojciechowicz: 255747@student.pwr.edu.pl.Faculty of Mathematics and Computer Science, Jagiellonian University, Kraków, Poland. Łukasz Struski: lukasz.struski@uj.edu.pl; Jacek Tabor: jacek.tabor@uj.edu.pl. Abstract Top-k selection determines which components of a sparse model remain active. Hard selection blocks gradients, while continuous relaxations often couple mask hardness to the selected mass. We introduce LaPrune, a mathematically exact-budget differentiable layer that controls the normalized second moment while preserving the selected mass. A LapSum barrier preserves the selection mass, and a normalized second-moment constraint moves the mask from a dense equal-mass allocation toward hard top-k at each budget. We derive a population prediction of the saturated fraction, a near-binary limiting law, and a tight worst-case guarantee on the near-zero fraction. The normalized hardness parameter is invariant to score scale, while a fixed LapSum temperature is not. 1 Introduction Sparse computation allows neural systems to increase capacity without activating every component. Mixture-of-experts models route each token to a few experts, long-context models retain a limited number of memories, retrieval systems return a short list, and sparse autoencoders activate a small set of latents. Despite their architectural differences, these systems share the same primitive: top-k selection. At inference, this primitive is a discrete selection. During training, however, a hard top-k gate blocks gradients at the boundary where the model must learn which items to retain. Continuous gates restore gradients, but they must also preserve the operational constraints of the sparse system. Two independent controls are required. The budget specifies a selection mass k, corresponding in the hard limit to k active slots such as expert slots, memory entries, retrieved documents, or retained pixels. The hardness determines how closely the continuous mask approximates a binary mask. A dense mask can distribute gradients broadly early in training, whereas a near-binary mask can approximate the deployment path later in training. Varying hardness should not change the budget. Most differentiable top-k relaxations do not provide this separation. Some return soft ranks or soft permutations instead of a directly budgeted mask. Others introduce a temperature or entropy parameter that changes sparsity, but without fixing the selection mass. Some recent linear-time operators obtain their speed by relaxing the budget constraint itself, and in the sparse regime this can produce masks whose mass is far from the intended k. We introduce LaPrune, a differentiable sparse-selection layer that extends the exact LapSum soft top-k operator (Struski et al. 2025). LapSum takes scores r∈ℝnr ^n, a budget k, and a user-specified temperature t, then returns pi=σ((ri−b)/t)p_i=σ((r_i-b)/t) with a barrier b chosen to satisfy ∑ipi=k _ip_i=k. The temperature controls mask hardness, but its value depends on the scale and distribution of the scores. LaPrune replaces this scale-dependent control with a normalized parameter γ∈[0,1]γ∈[0,1]. For 0<γ<10<γ<1, the layer determines a finite t>0t>0 and a barrier b so that both ∑ipi=k _ip_i=k and ∑ipi2=[a+(1−a)γ]k _ip_i^2=[a+(1-a)γ]k hold, where a=k/na=k/n. We define the endpoints directly: γ=0γ=0 gives the equal-mass mask pi=ap_i=a, whereas γ=1γ=1 gives the hard top-k indicator. Differentiability through the implicit finite-temperature solution applies to the interior range. The first constraint preserves the selection budget, while the second controls how concentrated the mask is. Thus, unlike LapSum with a manually chosen temperature, LaPrune provides normalized hardness control without changing the budget. Our contributions are: • LaPrune, a mathematically exact-budget differentiable layer whose normalized parameter γ controls mask hardness through a second-moment constraint, • a comprehensive theoretical analysis of mask saturation, providing both a predictive model for typical score distributions and a strict, distribution-free lower bound on the fraction of near-zero entries, • empirical evidence demonstrating that the normalized hardness parameter is robustly scale-invariant and improves feature recovery in sparse networks, • a highly efficient implementation that strictly enforces the selection budget while scaling seamlessly up to n=107n=10^7. 2 Related Work Differentiable Ranking and Top-k. Prior work addresses several parts of differentiable selection. NeuralSort (Grover et al. 2019), SoftSort (Prillo and Eisenschlos 2020), differentiable sorting networks (Petersen et al. 2021), and permutahedron-based operators (Blondel et al. 2020) approximate ranks or permutations. Methods designed specifically for top-k selection use entropic optimal transport (Xie et al. 2020), convex isotonic optimization (Sander et al. 2023), perturbation (Berthet et al. 2020), subset reparameterization (Xie and Ermon 2019), or sorting networks (Petersen et al. 2022). These methods primarily construct differentiable approximations of a discrete selection. Our objective is to modulate mask hardness while keeping the selected mass fixed. Sparsemax (Martins and Astudillo 2016) and entmax (Peters et al. 2019) produce sparse probability vectors on the simplex. They enforce ∑ipi=1 _ip_i=1 and have data-dependent support, so neither operator accepts an exact user-specified budget k. A matched exact-budget comparison would require an additional calibration that changes their training objective. Scalar-threshold methods are closer to our setting. LapSum (Struski et al. 2025) enforces the budget through a Laplace-CDF barrier. DFTopK (Zhu et al. 2025) obtains a closed-form threshold by relaxing ∑ipi=k _ip_i=k. The resulting budget drift is measured in the Supplementary Materials. LaPrune adds a second-moment constraint to LapSum and converts temperature into a normalized hardness parameter. Where Budgeted Selection Is Used. The meaning of a sparsity budget depends on the model. Mixture-of-experts systems limit the number of experts assigned to each token (Shazeer et al. 2017; Fedus et al. 2022). Network pruning fixes the fraction of retained weights or channels (Han et al. 2015; Frankle and Carbin 2019), and token-pruning methods limit the tokens processed by later layers (Rao et al. 2021). Sparse autoencoders constrain the number of active latent features (Makhzani and Frey 2014; Gao et al. 2025). Despite using different architectures, these applications face the same trade-off. Hard top-k enforces the required budget at the cost of blocking gradients. Soft selection permits gradient-based training, but it may not preserve the requested selection mass for each sample. LaPrune preserves the budget for every selection while controlling how close the mask is to hard top-k. Background: Exact-Budget Soft Top-k. LapSum (Struski et al. 2025) takes a score vector r∈ℝnr ^n, a budget k∈1,…,n−1k∈\1,…,n-1\, and a temperature t>0t>0. Writing a=k/na=k/n, it returns a mask p∈(0,1)np∈(0,1)^n whose entries are determined by a scalar barrier b, such that pi(t) p_i(t) =σ(ri−b(t)t), =σ\! ( r_i-b(t)t ), 1n∑i=1npi=a 1n _i=1^np_i=a ⟺∑i=1npi=k, _i=1^np_i=k, where σ(u)=12euσ(u)= 12e^u for u≤0u≤ 0 and 1−12e−u1- 12e^-u for u>0u>0 is the standard Laplace CDF. For fixed r, k, and t, the map b↦∑iσ((ri−b(t))/t)b _iσ((r_i-b(t))/t) is continuous and strictly decreasing from n to 0. It therefore has a unique solution satisfying ∑ipi=k _ip_i=k. For each chosen temperature t, LapSum computes the corresponding barrier b to preserve the budget. The temperature determines the width of the transition between rejected and selected scores. Our method replaces the direct choice of t with a normalized hardness parameter and determines both t and the corresponding barrier. Geometric Interpretation of LapSum. Equivalently, if Zi∼Laplace(ri,t)Z_i (r_i,t), then pi=Pr(Zi>b)p_i= (Z_i>b). The budget equation chooses the shared barrier so that the total tail mass is k. Scores near the barrier receive fractional mask values and scores far from it receive values close to zero or one. Let g∈ℝng ^n denote an upstream gradient. Differentiating the mask with respect to the scores while accounting for the dependence of b on r yields the analytical vector–Jacobian product (∇rp)⊤g ( _rp) g =κ⊙(g−⟨g,q⟩), =κ (g- g,q 1 ), (1) qi q_i =κi∑jκj,κi=12te−|ri−b|/t. = _i _j _j, _i= 12te^-|r_i-b|/t. Here κi _i is the Laplace density at the barrier for score rir_i, and q is its normalization across scores. The correction term ⟨g,q⟩ g,q 1 accounts for the change in b required to preserve ∑ipi=k _ip_i=k. Consequently, the gradient of each mask value depends on all scores through the shared barrier. LapSum preserves ∑ipi=k _ip_i=k for every t>0t>0. For integer k and pairwise distinct scores, its mask converges to the hard top-k indicator as t→0t→ 0 (Struski et al. 2025, Theorem 3.3). 3 LaPrune: Exact-Budget Hardness Control LapSum enforces an exact selection budget for every temperature. However, the temperature is not an interpretable measure of mask hardness, because its numerical value depends on the scale and distribution of the input scores. LaPrune replaces this scale-dependent control with a normalized hardness parameter. Given a target hardness, the layer determines the corresponding temperature, while the LapSum barrier continues to enforce the exact budget. Let the score vector, target budget, and budget fraction be defined by r∈ℝn,k∈1,…,n−1,a=kn.r ^n, k∈\1,…,n-1\, a= kn. For every positive temperature, the mask entries are pi(t)=σ(ri−b(t)t),i=1,…,n,p_i(t)=σ ( r_i-b(t)t ), i=1,…,n, where σ is the Laplace cumulative distribution function and the barrier b(t)b(t) is the unique solution of ∑i=1npi(t)=k. _i=1^np_i(t)=k. (2) Second Moment as a Measure of Hardness. The budget constraint determines the first moment of the mask ∑i=1npi=k. _i=1^np_i=k. It does not, however, specify how the mass is distributed among the entries. The same budget can be realized by a uniform mask, in which all entries are equal, or by a binary mask containing exactly k active entries. We quantify this distinction using the second moment M2(p):=∑i=1npi2.M_2(p):= _i=1^np_i^2. A small second moment indicates that the mask entries are concentrated near their common mean, a diffuse mask. A large second moment indicates that the entries are concentrated near zero and one, a hard, approximately binary mask. Feasible Range of the Second Moment. The budget and box constraints imply ak=k2n≤∑i=1npi2≤k.ak= k^2n≤ _i=1^np_i^2≤ k. (3) The lower bound is Cauchy–Schwarz and is attained only by the uniform mask pi=ap_i=a. The upper bound follows from pi2≤pip_i^2≤ p_i and is attained by binary masks with exactly k ones. Normalized Hardness Parameter. The feasible interval permits a scale-free normalization of the second moment. We define γ:=∑i=1npi2−ak−ak=∑i=1npi2−ak(1−a),γ∈[0,1].γ:= _i=1^np_i^2-akk-ak= _i=1^np_i^2-akk(1-a), γ∈[0,1]. Equivalently, setting the normalized hardness imposes the second-moment constraint ∑i=1npi2=βk,β=a+(1−a)γ,β∈[a,1]. _i=1^np_i^2=β k, β=a+(1-a)γ, β∈[a,1]. Jointly rescaling the scores, barrier, and temperature leaves the mask, and therefore γ, unchanged. Because LaPrune solves for the temperature, a rescaling of the input scores produces the corresponding rescaled solution. A fixed numerical temperature does not have this invariance. Temperature Determined by the Hardness Constraint. For every positive temperature, let the barrier be defined by Equation (2), and define M2(t):=∑i=1n(σ(ri−b(t)t))2.M_2(t):= _i=1^n (σ ( r_i-b(t)t ) )^2. Theorem 1 (Unique Temperature) For every nonconstant score vector, M2(t)M_2(t) is continuous and strictly decreasing in t. If the scores are pairwise distinct, then M2(t)→kM_2(t)→ k as t↓0t 0 and M2(t)→akM_2(t)→ ak as t→∞t→∞. Consequently, every γ∈(0,1)γ∈(0,1) determines a unique temperature tγt_γ satisfying the target second moment. The proof is given in the Supplementary Materials section "Monotonicity and Uniqueness of the Temperature." Increasing γ decreases tγt_γ. The finite-temperature family is defined for γ∈(0,1)γ∈(0,1). Its limits as γ→0γ→ 0 and γ→1γ→ 1 are the uniform and hard top-k masks, respectively. We therefore define the endpoints explicitly as pi=ap_i=a at γ=0γ=0 and the hard top-k indicator at γ=1γ=1. Computing LaPrune: Joint Solves and Implicit Gradients. For 0<γ<10<γ<1, LaPrune must determine both the barrier b, which preserves the selection mass, and the temperature t, which realizes the requested hardness. We set τ=logtτ= t and determine the barrier and temperature simultaneously by solving F1(b,τ) F_1(b,τ) :=∑i=1npi−k=0, := _i=1^np_i-k=0, F2(b,τ) F_2(b,τ) :=∑i=1npi2−βk=0, := _i=1^np_i^2-β k=0, (4) where pi=σ(ri−bt),β=a+(1−a)γ,a=kn.p_i=σ\! ( r_i-bt ), β=a+(1-a)γ, a= kn. To solve the system in (4), let zi=ri−bt,di=σ′(zi),z_i= r_i-bt, d_i=σ (z_i), and define the four reductions S0 S_0 :=∑idi, = _id_i, S1 S_1 :=∑idizi, = _id_iz_i, U0 U_0 :=∑ipidi, = _ip_id_i, U1 U_1 :=∑ipidizi. = _ip_id_iz_i. The Jacobian of the constraints with respect to (b,τ)(b,τ) is J(b,τ)F=[−S0/t−S1−2U0/t−2U1].J_(b,τ)F= bmatrix-S_0/t&-S_1\\ -2U_0/t&-2U_1 bmatrix. (5) Using the Jacobian in (5), each Newton iteration computes the increment (Δb,Δτ)( b, τ) from J(b,τ)F[ΔbΔτ]=−[F1F2].J_(b,τ)F bmatrix b\\ τ bmatrix=- bmatrixF_1\\ F_2 bmatrix. (6) The parameters are then updated as b←b+Δb← b+ b and τ←τ+Δτ←τ+ τ. Each iteration requires only elementwise operations, four reductions over the scores, and the solution of a 2×22× 2 linear system. We initialize the barrier using an estimate of the (1−a)(1-a)-quantile of the score distribution. Under a Gaussian approximation, the initial barrier is b0=μ^+s^Φ−1(1−a),b_0= μ+ s\, ^-1(1-a), where μ μ and s s are the row-wise mean and standard deviation. Using the near-binary approximation, we initialize the temperature as t0≈4a(1−a)(1−γ)3ρ^(b0),t_0≈ 4a(1-a)(1-γ)3 ρ(b_0), where ρ^(b0) ρ(b_0) denotes the corresponding density estimate at the initial barrier. Large Newton updates are clipped, and a nested bisection solver provides a robust fallback when the residuals do not reach the prescribed tolerance. The endpoints γ=0γ=0 and γ=1γ=1 bypass the finite-temperature solver and return the uniform and hard top-k masks, respectively. For the backward pass, gradients are obtained by differentiating the converged solution directly, without unrolling the Newton iterations. For any scalar input x∈r1,…,rn,a,γx∈\r_1,…,r_n,a,γ\, the implicit function theorem gives J(b,τ)F[∂b/∂x∂τ/∂x]=−∂F∂x|b,τ.J_(b,τ)F bmatrix∂ b/∂ x\\ ∂τ/∂ x bmatrix=- . ∂ F∂ x |_b,τ. (7) The system in (7) reuses the same 2×22× 2 constraint Jacobian. The resulting vector–Jacobian product is evaluated analytically, so the backward pass does not need to store the solver trajectory and retains only the score vector and the quantities computed at the converged solution. 4 Theoretical Analysis of the Hardness Parameter The hardness parameter γ moves the mask from a dense allocation toward a near-binary top-k mask. We characterize this transition through the fraction of mask entries that saturate to 0 or 11, and we do so in two complementary ways: a population prediction under a score model and a lower bound that holds for any score distribution. Throughout, let the scores follow a density ρ with CDF F, and let μ(x)=σ((x−b)/t)μ(x)=σ((x-b)/t) denote the population mask, the value the layer assigns to a score equal to x. With a=k/na=k/n and β=a+(1−a)γβ=a+(1-a)γ, the budget and hardness constraints fix the first two population moments of the mask at a and βaβ a. Mean-Field Prediction. The mean-field prediction replaces the n random scores by their population density ρ and works with the expected moments, which the empirical moments approach as n grows. The pair (t,b)(t,b) then solves the two-moment fixed point M¯1(t,b) M_1(t,b) :=∫σ(x−bt)ρ(x)x=a, = σ\! ( x-bt )ρ(x)\,dx=a, (8) M¯2(t,b) M_2(t,b) :=∫(σ(x−bt))2ρ(x)x=βa. = (σ\! ( x-bt ) )^2ρ(x)\,dx=β a. An entry is treated as numerically zero when μ(x)<εμ(x)< , where we use ε=10−3 =10^-3. For 0<ε<1/20< <1/2 and x≤bx≤ b, the Laplace CDF satisfies σ(u)=12euσ(u)= 12e^u. Therefore, μ(x)<ε⇔12e(x−b)/t<ε⇔x<b+tln(2ε).μ(x)< 12e^(x-b)/t< x<b+t (2 ). Similarly, μ(x)>1−ε⇔x>b−tln(2ε).μ(x)>1- x>b-t (2 ). Thus, under the population model, the fractions of entries close to zero and one are frac0 _0 =F(b+tln(2ε)), =F\! (b+t (2 ) ), (9) frac1 _1 =1−F(b−tln(2ε)), =1-F\! (b-t (2 ) ), sε s_ =frac0+frac1. =frac_0+frac_1. The two conditions above place the nonsaturated values, those with ε≤μ(x)≤1−ε ≤μ(x)≤ 1- , in the band [b−t|ln(2ε)|,b+t|ln(2ε)|] [\,b-t (2 ) ,\ b+t (2 ) \, ], of width 2t|ln(2ε)|2t (2 ) . Increasing γ decreases t, which narrows this band and moves the mask values closer to 0 and 11. We solve the two-moment system (8) numerically for (t,b)(t,b). The Gaussian experiments use numerical quadrature. Closed-form Laplace moments and implementation details are given in the Supplementary Materials section "Numerical and Analytical Details." Near-Binary Limit Under Regularity. Assume that F has a continuous and positive density in a neighborhood of the unique (1−a)(1-a)-quantile q⋆:=q1−a.q_ :=q_1-a. As γ→1γ→ 1, we have t→0t→ 0. The budget constraint M¯1=a M_1=a then reduces to 1−F(b)=a1-F(b)=a, so b→q⋆b→ q_ . Moreover, M¯1−M¯2=∫ℝσ(x−bt)[1−σ(x−bt)]ρ(x)x. M_1- M_2= _Rσ\! ( x-bt ) [1-σ\! ( x-bt ) ]ρ(x)\,dx. After the change of variables u=(x−b)/tu=(x-b)/t, the first-order expansion is M¯1−M¯2=tρ(q⋆)I+o(t),I:=∫ℝσ(u)(1−σ(u))u. M_1- M_2=tρ(q_ )I+o(t), I:= _Rσ(u)(1-σ(u))\,du. For the unit Laplace CDF, I=34I= 34. Since M¯1−M¯2=a−βa=a(1−a)(1−γ), M_1- M_2=a-β a=a(1-a)(1-γ), we obtain t≈a(1−β)Iρ(q⋆)=43a(1−a)(1−γ)ρ(q1−a).t≈ a(1-β)Iρ(q_ )= 43 a(1-a)(1-γ)ρ(q_1-a). (10) The nonsaturated fraction is the probability mass inside the transition interval: 1−sε=F(b−tln(2ε))−F(b+tln(2ε)).1-s_ =F\! (b-t (2 ) )-F\! (b+t (2 ) ). For small t, 1−sε≈2t|ln(2ε)|ρ(q⋆).1-s_ ≈ 2t (2 ) ρ(q_ ). Substituting Equation (10) cancels the density term and gives 1−sε≈83|ln(2ε)|a(1−a)(1−γ).1-s_ ≈ 83 (2 ) a(1-a)(1-γ). (11) Therefore, near the binary limit, the nonsaturated fraction decreases linearly with 1−γ1-γ. Its leading-order coefficient depends on a and ε , but not on the local score density. This result holds for continuous score distributions with a positive density at q1−aq_1-a. We make no density-independent claim for discrete score distributions, atoms at the quantile, or a vanishing density at q1−aq_1-a. Worst-Case Guarantee for Any Distribution. The mean-field analysis describes typical, large-sample behavior. We next establish a worst-case counterpart: a deterministic lower bound on the fraction of near-zero mask entries that holds for every score distribution and depends only on the first two mask moments. After sorting the mask values, let f:[0,1]→[0,1]f:[0,1]→[0,1] denote the resulting nondecreasing rank profile. It satisfies ∫01f(x)x=a,∫01f(x)2x=a(a+(1−a)γ). _0^1f(x)\,dx=a, _0^1f(x)^2\,dx=a (a+(1-a)γ ). Theorem 2 (Worst-Case Sparsity Floor) Let f:[0,1]→[0,1]f:[0,1]→[0,1] be nondecreasing and satisfy ∫01f(x)x _0^1f(x)\,dx =a, =a, ∫01f(x)2x _0^1f(x)^2\,dx =aβ, =aβ, β β =a+(1−a)γ. =a+(1-a)γ. Then, for every ε∈(0,β] ∈(0,β], |x∈[0,1]:f(x)<ε|≥H⋆(ε), |\x∈[0,1]:f(x)< \ |≥ H_ ( ), where H⋆(ε):=(1−a)max0,1−a(1−γ)ε.H_ ( ):=(1-a) \0,1- a(1-γ) \. This bound is tight over the stated moment class. The proof and extremizing constructions are given in the Supplementary Materials section "Proof of the Worst-Case Sparsity Floor." The lower bound is zero when ε≤a(1−γ) ≤ a(1-γ), because the moment constraints can then be satisfied while keeping every mask value at least ε . The bound increases with γ. As γ→1γ→ 1, it converges to 1−a1-a for every fixed ε>0 >0, recovering the fraction of zeros in a hard top-k mask. 5 Empirical Evaluation We test five claims. First, LaPrune should retain LapSum’s favorable empirical scalability. Second, γ should control a normalized second moment while the budget remains fixed and provide a monotone soft-to-hard path at each budget. Third, the soft-to-hard path should improve recovery when the selected features are known. Fourth, budget-preserving training should remain useful in a sparse model with real activations. Fifth, we numerically verify that the generated masks respect the distribution-free worst-case sparsity floor established in Theorem 2. The budget equation is exact mathematically and is enforced to numerical tolerance in the implementation. The normalized budget and second-moment residuals are approximately 10−710^-7 in our diagnostic sweep. Additional results appear in the Supplementary Materials sections "Additional Hardness Diagnostics," "Feature-Selection Budget Sweep," "Sparse-Autoencoder Sweeps," "CIFAR-100 Top-k Classification," and "Numerical and Analytical Details." Computational Scaling and Memory Figure 1: Combined forward-and-backward runtime (top) and peak memory (bottom) against score dimension n on an A100, with k=n/16k=n/16 and γ=0.9γ=0.9. LaPrune and LapSum reach n=107n=10^7 while preserving per-sample mass to numerical tolerance. We measure combined forward and backward runtime and peak GPU memory for dimensions up to n=107n=10^7 at k=n/16k=n/16 and γ=0.9γ=0.9 on an NVIDIA A100. LaPrune and LapSum reach n=107n=10^7 while enforcing the per-sample mass constraint to numerical tolerance. Sorting relaxations exhaust memory at substantially smaller dimensions. At n=107n=10^7, LaPrune takes 10.7510.75 ms and 305 MB, compared with 5.385.38 ms and 544 MB for LapSum. The second constraint therefore adds a factor of two in runtime while reducing measured peak memory. Sander-PAV and optimal transport reach the same dimension but are 1818–22×22× slower. Each point uses a single float32 score vector, one warm-up call, and the median of 3–15 synchronized repetitions in an isolated subprocess. LaPrune uses its analytical implicit backward pass, whereas the baseline backward implementations follow their respective reference or local implementations. The full timing protocol and implementation details are given in the Supplementary Materials section "Scaling Protocol and Implementations." Figure 2: Normalized second-moment control on Gaussian scores with n=2000n=2000, averaged over five seeds. The vertical quantity is the fraction of entries within ε=0.01 =0.01 of zero or one. Left: LaPrune provides a monotone path between the equal-mass and binary endpoints at each fixed budget k/nk/n. Equal γ need not yield equal saturation across budgets. Right: LaPrune keeps the same discretized fraction when scores are rescaled by c, while fixed or once-tuned LapSum temperatures drift. Normalized Hardness and Budget Preservation Figure 2 tests whether γ provides a normalized second-moment control that cannot be reproduced by a fixed LapSum temperature. We draw n=2000n=2000 Gaussian scores, average over five seeds, and define an entry as discretized when pi<0.01p_i<0.01 or pi>0.99p_i>0.99. The points at γ=0γ=0 and γ=1γ=1 are evaluated using the explicit equal-mass and hard top-k definitions, not a finite-temperature solve. At every tested budget, these endpoints give an equal-mass mask and a binary mask, respectively. At γ=0.9γ=0.9, the discretized fraction ranges from 0.740.74 at k/n=0.5k/n=0.5 to 0.980.98 at k/n=0.02k/n=0.02. Thus, equal values of γ do not imply equal saturation across budgets. Instead, γ parameterizes a monotone path between the two endpoints for each fixed budget. The right panel rescales the same scores by c∈[0.1,100]c∈[0.1,100]. LaPrune holds the discretized fraction at 0.4810.481 for γ=0.5γ=0.5 and 0.9050.905 for γ=0.9γ=0.9. LapSum with t=1t=1 changes from 0.0000.000 to 0.9870.987. A temperature tuned to match LaPrune at c=1c=1 also drifts away from the target. This experiment isolates the role of the normalized second moment: the solved t follows score scale, while γ retains its meaning. Table 1: Feature selection with 200 features, 10 informative features, and exact hard top-10 deployment. Values are mean ± standard deviation over 20 seeds. The near-binary fraction is the fraction of training-time mask entries within ε=10−3 =10^-3 of zero or one. Method Recovery F1 Deploy ACC Near-Binary Fraction LaPrune (γ: 0→0.9γ:\,0\!→\!0.9) (ours) 0.855 ± 0.097 0.872 ± 0.028 0.869 ± 0.029 LaPrune (γ=0.9γ=0.9) (ours) 0.800 ± 0.118 0.882 ± 0.028 0.908 ± 0.030 LapSum (t=1t=1) 0.795 ± 0.116 0.807 ± 0.037 0.000 ± 0.000 hard top-k (STE) 0.375 ± 0.155 0.743 ± 0.097 1.000 ± 0.000 DFTopK 0.535 ± 0.106 0.852 ± 0.048 0.000 ± 0.000 Two checks separate hardness control from numerical error. On n=104n=10^4 scores, the normalized budget and second-moment residuals remain below approximately 10−710^-7, so LaPrune and LapSum preserve the requested budget to numerical tolerance. On n=2×105n=2× 10^5 scores with ε=10−3 =10^-3, the population prediction in Equations (8)–(9) matches the empirical saturated fraction with maximum error 0.00220.0022. By contrast, DFTopK reaches 3939 times the target mass for k=n/100k=n/100 at τ=5τ=5. The corresponding figures are in the Supplementary Materials section Additional Hardness Diagnostics." Differentiable Feature Selection We construct 200-dimensional data with exactly ten informative features, using 1500 training and 500 test samples. A learned score for each feature passes through the tested mask into the same linear classifier. Every trained model is evaluated with an exact hard top-10 mask. Feature-recovery F1 directly measures whether the selected set matches the known informative set. Table 1 reports mean and standard deviation across 20 paired seeds. The annealed schedule increases γ linearly by epoch from 0 at epoch 0 to 0.90.9 at epoch 149 of 150. This schedule gives F1 0.855±0.0970.855± 0.097, compared with 0.795±0.1160.795± 0.116 for LapSum. The paired difference is 0.060±0.0660.060± 0.066 with Wilcoxon p=0.005p=0.005. The gain over fixed γ=0.9γ=0.9 is 0.055±0.0670.055± 0.067 with p=0.007p=0.007. DFTopK obtains recovery F1 0.535±0.1060.535± 0.106 after hard top-10 deployment. Table 2: Sparse autoencoders on ResNet-18/CIFAR-100 activations at dictionary size m=4096m=4096 and deployment budget k=32k=32. Values are mean ± standard deviation over eight seeds. Method FVU Probe ACC (%) LaPrune (γ: 0→0.9γ:\,0\!→\!0.9) (ours) 0.3587 ± 0.0006 56.59 ± 0.36 LaPrune (γ=0.9γ=0.9) (ours) 0.3462 ± 0.0002 57.19 ± 0.40 LapSum (t=1t=1) 7.906 ± 0.123 45.35 ± 0.33 hard top-k (STE) 1.155 ± 0.010 46.89 ± 0.39 DFTopK 0.5825 ± 0.0005 51.03 ± 0.15 Sparse Autoencoders on Real Activations Figure 3: Direct numerical check of the worst-case sparsity floor at budget fraction a=1/10a=1/10 on n=2×105n=2× 10^5 Gaussian scores, averaged over five seeds. Open markers show the measured near-zero fraction of LaPrune masks and solid lines show the distribution-free lower bound H⋆(ε)H_ ( ) from Theorem 2. Left: the threshold ε varies at fixed hardness. Right: the hardness γ varies at fixed threshold. Every measured point lies at or above its corresponding floor. We train overcomplete sparse autoencoders on frozen ResNet-18 activations from CIFAR-100. The activation dimension is 512, the dictionary has m=4096m=4096 latents, and the deployment budget is k=32k=32. Each model is trained for 30 epochs on 50,000 activations using mean-squared reconstruction error without an auxiliary dead-latent loss. The annealed schedule increases γ linearly from 0 at epoch 0 to 0.90.9 at epoch 23, then holds it fixed through epoch 29. We report eight-seed results on 10,000 held-out activations. Reconstruction uses fraction of variance unexplained (FVU). A linear probe trained on the exact hard top-32 codes measures retained class information. Table 2 shows that both LaPrune schedules outperform the other measured budget-preserving operators. Fixed γ=0.9γ=0.9 obtains FVU 0.3462±0.00020.3462± 0.0002 and probe accuracy 57.19±0.40%57.19± 0.40\%. Hard top-k reaches FVU 1.155±0.0101.155± 0.010 and probe accuracy 46.89±0.39%46.89± 0.39\%. LapSum remains diffuse at t=1t=1, leading to FVU 7.906±0.1237.906± 0.123. DFTopK obtains lower FVU than these two baselines. Simplex operators are not included because their training masks have mass one and data-dependent support. Forcing those models through hard top-32 deployment would confound operator quality with a change in mask scale and support size. In paired comparisons, fixed LaPrune improves both FVU and probe accuracy in all eight runs against LapSum, hard top-k, and DFTopK (p=0.0078p=0.0078 for each Wilcoxon test). Worst-case sparsity floor We directly check Theorem 2 on masks produced by the implemented operator. We draw n=2×105n=2× 10^5 Gaussian scores, fix the budget fraction at a=1/10a=1/10, and average the near-zero fraction over five seeds. Figure 3 compares the measured fraction Pr(pi<ε) (p_i< ) with the distribution-free floor H⋆(ε)=(1−a)max0,1−a(1−γ)ε.H_ ( )=(1-a) \0,1- a(1-γ) \. The measured fraction remains at or above the floor for every tested threshold and hardness value. The Gaussian masks need not attain equality because the theorem is tight over the full moment class and not for every score distribution. 6 Conclusion This work develops a mathematical framework for controlling the softness of differentiable top-k masks without changing their prescribed mass. LaPrune expresses budget and hardness as separate constraints, links the normalized second moment to mask saturation, and characterizes the transition from equal allocation to discrete selection. The resulting theory includes a population description, an asymptotic law near the binary limit, and a distribution-free sparsity guarantee. The experiments are designed as controlled tests of these mathematical claims, not as an effort to maximize performance on individual benchmarks. They show that the proposed parameter retains its interpretation under score rescaling, that the constrained masks can support learning in synthetic and real-feature settings, and that the operator remains computationally practical at large dimensions. Their role is therefore to establish that the theoretical construction is numerically realizable and relevant to sparse learning. Limitations. The present formulation uses a Laplace-CDF mask and controls hardness through one normalized moment. Its guarantees describe mask geometry, not the complete optimization dynamics or downstream generalization of models trained with the layer. Numerical conditioning also becomes more difficult close to the binary endpoint. Extending the analysis to other mask families, alternative concentration constraints, and training-dependent score distributions may yield a broader theory of differentiable sparsity. Improving endpoint solvers and studying the framework in larger sparse architectures are natural directions for future work. References L. Berrada, A. Zisserman, and M. P. Kumar (2018) Smooth loss functions for deep top-k classification. External Links: 1802.07595, Link Cited by: Appendix A. Q. Berthet, M. Blondel, O. Teboul, M. Cuturi, J. Vert, and F. Bach (2020) Learning with differentiable perturbed optimizers. In Advances in Neural Information Processing Systems (NeurIPS), External Links: 2002.08676, Link Cited by: §2. M. Blondel, O. Teboul, Q. Berthet, and J. Djolonga (2020) Fast differentiable sorting and ranking. In International Conference on Machine Learning (ICML), External Links: 2002.08871, Link Cited by: §2. W. Fedus, B. Zoph, and N. Shazeer (2022) Switch transformers: scaling to trillion parameter models with simple and efficient sparsity. Journal of Machine Learning Research (JMLR). External Links: 2101.03961, Link Cited by: §2. J. Frankle and M. Carbin (2019) The lottery ticket hypothesis: finding sparse, trainable neural networks. In International Conference on Learning Representations (ICLR), External Links: 1803.03635, Link Cited by: §2. L. Gao, T. Dupré la Tour, H. Tillman, G. Goh, R. Troll, A. Radford, I. Sutskever, J. Leike, and J. Wu (2025) Scaling and evaluating sparse autoencoders. In International Conference on Learning Representations (ICLR), External Links: 2406.04093, Link Cited by: §2. A. Grover, E. Wang, A. Zweig, and S. Ermon (2019) Stochastic optimization of sorting networks via continuous relaxations. In International Conference on Learning Representations (ICLR), External Links: 1903.08850, Link Cited by: §2. S. Han, J. Pool, J. Tran, and W. J. Dally (2015) Learning both weights and connections for efficient neural networks. In Advances in Neural Information Processing Systems (NeurIPS), External Links: 1506.02626, Link Cited by: §2. A. Makhzani and B. Frey (2014) K-sparse autoencoders. In International Conference on Learning Representations (ICLR), External Links: 1312.5663, Link Cited by: §2. A. Martins and R. Astudillo (2016) From softmax to sparsemax: a sparse model of attention and multi-label classification. In International conference on machine learning, p. 1614–1623. Cited by: §2. B. Peters, V. Niculae, and A. F. Martins (2019) Sparse sequence-to-sequence models. In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, p. 1504–1519. Cited by: §2. F. Petersen, C. Borgelt, H. Kuehne, and O. Deussen (2021) Differentiable sorting networks for scalable sorting and ranking supervision. In International Conference on Machine Learning (ICML), External Links: 2105.04019, Link Cited by: §2. F. Petersen, H. Kuehne, C. Borgelt, and O. Deussen (2022) Differentiable top-k classification learning. In International Conference on Machine Learning, p. 17656–17668. Cited by: §2. S. Prillo and J. M. Eisenschlos (2020) SoftSort: a continuous relaxation for the argsort operator. In International Conference on Machine Learning (ICML), External Links: 2006.16038, Link Cited by: §2. Y. Rao, W. Zhao, B. Liu, J. Lu, J. Zhou, and C. Hsieh (2021) DynamicViT: efficient vision transformers with dynamic token sparsification. In Advances in Neural Information Processing Systems (NeurIPS), External Links: 2106.02034, Link Cited by: §2. M. E. Sander, J. Puigcerver, J. Djolonga, G. Peyré, and M. Blondel (2023) Fast, differentiable and sparse top-k: a convex analysis perspective. In International Conference on Machine Learning (ICML), External Links: 2302.01425, Link Cited by: §2. N. Shazeer, A. Mirhoseini, K. Maziarz, A. Davis, Q. Le, G. Hinton, and J. Dean (2017) Outrageously large neural networks: the sparsely-gated mixture-of-experts layer. In International Conference on Learning Representations (ICLR), External Links: 1701.06538, Link Cited by: §2. Ł. Struski, M. B. Bednarczyk, I. T. Podolak, and J. Tabor (2025) LapSum - one method to differentiate them all: ranking, sorting and top-k selection. In Proceedings of the 42nd International Conference on Machine Learning, p. 56990–57007. Cited by: §1, §2, §2, §2. S. M. Xie and S. Ermon (2019) Reparameterizable subset sampling via continuous relaxations. arXiv preprint arXiv:1901.10517. Cited by: §2. Y. Xie, H. Dai, M. Chen, B. Dai, T. Zhao, H. Zha, W. Wei, and T. Pfister (2020) Differentiable top-k operator with optimal transport. In Advances in Neural Information Processing Systems (NeurIPS), External Links: 2002.06504, Link Cited by: §2. Y. Zhu, Z. Zhang, Y. Wang, Z. Wang, Y. Li, R. Zhou, S. Wen, P. Jiang, C. Lin, and J. Yang (2025) Differentiable fast top-k selection for large-scale recommendation. External Links: 2510.11472, Link Cited by: §2. Appendix A Supplementary Materials Additional Hardness Diagnostics Mean-field agreement. We draw n=2×105n=2× 10^5 Gaussian scores and sweep γ at budget fractions a∈1/2,1/10a∈\1/2,1/10\. An entry is saturated when pi<εp_i< or pi>1−εp_i>1- , with ε=10−3 =10^-3. Figure 4 overlays finite-sample measurements with the population fixed-point and saturation equations from the main paper. The maximum absolute discrepancy is 0.00220.0022 for a=1/2a=1/2 and 0.00060.0006 for a=1/10a=1/10. Figure 4: Population prediction and finite-sample saturated fraction on n=2×105n=2× 10^5 Gaussian scores with ε=10−3 =10^-3. Lines show the population model and open markers show simulated masks. Sensitivity to ε . Figure 5 repeats the population and simulation comparison for ε∈10−4,10−3,10−2 ∈\10^-4,10^-3,10^-2\. The threshold changes where a finite mask entry is counted as saturated, but the monotone transition and agreement between theory and simulation persist. Figure 5: Population prediction and simulation for three numerical saturation thresholds. Each curve uses LaPrune masks with the budget equation enforced to numerical tolerance at a=1/10a=1/10. Budget drift. Figure 6 measures the realized mass of DFTopK against budget-preserving operators on n=104n=10^4 Gaussian scores. For k=n/100k=n/100, the DFTopK mass grows from 1.031.03 times the target at τ=0.05τ=0.05 to 38.7538.75 times the target at τ=5τ=5. LaPrune and LapSum remain at the requested mass to numerical tolerance. Figure 6: Realized soft-mask mass against DFTopK softness τ on n=104n=10^4 Gaussian scores. DFTopK overshoots sparse budgets as τ increases. LaPrune and LapSum enforce the budget equation. Scaling Protocol and Implementations The scaling benchmark uses one float32 score vector per call (batch size one), with k=n/16k=n/16 and γ=0.9γ=0.9, on an NVIDIA A100-SXM4-40GB. We use Python 3.9.25, PyTorch 2.8.0 with CUDA 12.8, difftopk 0.2.0, NumPy 2.0.2, and Numba 0.60.0. Every method–dimension pair runs in a fresh subprocess with a 240-second timeout. After one untimed warm-up call, we reset PyTorch’s peak-memory counter. CUDA events measure 15 repetitions for n≤105n≤ 10^5, seven for 105<n≤10610^5<n≤ 10^6, and three for larger n. We synchronize the device before reading the events and report the median. Peak memory is torch.cuda.max_memory_allocated over the timed calls. LaPrune uses the compiled pure-PyTorch Newton forward and the custom analytical implicit VJP in this paper. LapSum uses its sort–logcumsumexp barrier solve and custom one-constraint implicit VJP. NeuralSort, SoftSort, and DiffSort use difftopk and PyTorch autograd. The Sander Dykstra and PAV curves use local PyTorch ports of the authors’ isotonic formulation with a custom block-averaging backward. The optimal-transport curve uses a local two-bin Sinkhorn implementation with 200 iterations. The first 199 are detached and the final iteration uses autograd. Thus, each curve includes the backward implementation intended for that operator. The lower measured peak memory of LaPrune (305 MB) than LapSum (544 MB) is an implementation effect, not a consequence of adding fewer constraints. LaPrune is sort-free. Each Newton pass performs elementwise reductions, and its custom backward stores only the score vector and solved scalars. The measured LapSum implementation sorts the full vector and materializes sorted values, indices, and prefix and suffix logcumsumexp work arrays while locating the barrier. These temporary arrays determine its higher peak, whereas LaPrune’s 30 Newton passes reuse fixed-size buffers without accumulating per-iteration state. Feature-Selection Budget Sweep Figure 7 varies the deployment budget over k∈3,5,10,20,40k∈\3,5,10,20,40\. Each method is evaluated with the same exact hard top-k rule. Annealed LaPrune remains competitive with the fixed-γ variant and at or above the measured non-LaPrune budget-preserving baselines throughout the sweep. Figure 7: Feature-recovery F1 against deployment budget k for 200-dimensional synthetic data with ten informative features. Curves are means over eight seeds and bands show one standard deviation. The vertical line marks the true feature count. Sparse-Autoencoder Sweeps Architecture, training, and evaluation. Let x∈ℝ512x ^512 be a standardized frozen ResNet-18 activation. The SAE computes u u =Wenc(x−bpre)+benc, =W_enc(x-b_pre)+b_enc, h h =p(u)⊙u, =p(u) u, x x =Wdech+bpre. =W_dech+b_pre. where p(u)p(u) is the method-specific training mask. We minimize the mean squared reconstruction error ∥x^−x∥22/512 x-x _2^2/512 with no sparsity penalty or dead-latent auxiliary loss. We initialize the columns of Wdec∈ℝ512×4096W_dec ^512× 4096 from independent standard Gaussians and normalize each to unit norm, set Wenc=Wdec⊤W_enc=W_dec only at initialization, and initialize both biases to zero. The encoder and decoder are then untied. After every Adam update, we renormalize each decoder column to unit norm. Training uses all 50,000 componentwise standardized training activations, Adam with learning rate 10−310^-3, batches of 4096, and 30 epochs. For the annealed model at epoch e, the exact schedule is γe _e =min(0.9,0.9e⌊0.8(30−1)⌋). = \! (0.9,\, 0.9e 0.8(30-1) ). so γ reaches 0.90.9 at epoch 23 and remains there. Evaluation replaces the training mask with hard top-32 and computes FVU on all 10,000 test activations. FVU =∑j∥x^j−xj∥22∑j∥xj−x¯test∥22. = _j x_j-x_j _2^2 _j x_j- x_test _2^2. The linear probe is trained on hard top-32 codes from all 50,000 training activations and evaluated on the 10,000 test codes. It is a linear 4096→1004096\!→\!100 softmax classifier with PyTorch’s default initialization, trained for 60 epochs with cross-entropy, Adam at learning rate 10−210^-2, weight decay 10−410^-4, and batches of 8192. Features are extracted from the avgpool output of an ImageNet-1K-pretrained ResNet-18 and standardized using training-set means and standard deviations. The main SAE experiment fixes m=4096m=4096 and k=32k=32. Figure 8 changes the deployment budget at fixed dictionary size. Figure 9 changes dictionary size at fixed budget. Figure 8: Sparse-autoencoder FVU and linear-probe accuracy against deployment budget k at fixed dictionary size m=4096m=4096. Curves show means over three seeds and bands show one standard deviation. Figure 9: Sparse-autoencoder FVU and linear-probe accuracy against dictionary size m at fixed deployment budget k=32k=32. Curves show means over three seeds and bands show one standard deviation. Table 3: CIFAR-100 design ablation under the pure top-5 objective. All entries use the same ResNet-18 and 200-epoch protocol. Configuration ACC@1 ACC@5 fixed γ=0.3γ=0.3 72.8 82.4 fixed γ=0.5γ=0.5 71.8 86.0 fixed γ=0.7γ=0.7 67.4 88.0 fixed γ=0.9γ=0.9 1.0 5.0 fixed γ=0.99γ=0.99 2.0 27.1 anneal 0.1→0.80.1→ 0.8, clip norm 5 72.0 87.4 anneal 0.1→0.80.1→ 0.8, no clipping 71.9 86.9 anneal 0.1→0.90.1→ 0.9, clip norm 5 1.0 5.0 CIFAR-100 Top-k Classification We also evaluate each operator with a pure top-5 cross-entropy loss using a ResNet-18 trained from scratch on CIFAR-100. For LaPrune, γ increases linearly by epoch from 0.10.1 at epoch 0 to 0.80.8 at epoch 199. Table 4 reports a single-seed comparison using the competitor softness values specified by the LapSum benchmark. The Smooth top-k row uses the loss of Berrada, Zisserman, and Kumar (Berrada et al. 2018). Across two seeds, LaPrune achieves ACC@1 71.93±0.0771.93± 0.07 and ACC@5 87.13±0.3087.13± 0.30. Table 4: CIFAR-100 classification with a ResNet-18, 200 training epochs, and pure top-5 loss. All methods use the same random seed. Method ACC@1 ACC@5 Smooth top-k 57.3 90.3 NeuralSort 13.4 86.3 SoftSort 60.4 89.8 SinkhornSort 62.4 90.2 DiffSortNets 62.4 90.6 LapSum 65.9 92.3 LaPrune (γ: 0.1→0.8γ:\,0.1\!→\!0.8) 72.0 87.4 Table 3 separates the effects of hardness, clipping, and the endpoint of the schedule. Fixed γ exposes an ACC@1–ACC@5 trade-off. Training collapses near the binary endpoint. Gradient clipping at norm 5 has little effect for the 0.1→0.80.1→ 0.8 schedule, while extending the schedule to 0.9 collapses under this pure top-5 objective. Reproducibility Details All reported random sweeps use fixed integer seeds and store per-seed results before aggregation. The feature-selection models train for 150 epochs with Adam at learning rate 10−210^-2. The table uses 20 seeds and the budget sweep uses eight. The SAE models train for 30 epochs with Adam at learning rate 10−310^-3 and batch size 4096. Linear probes train for 60 epochs with Adam at learning rate 10−210^-2 and weight decay 10−410^-4. The SAE main table uses eight seeds. Both SAE sweeps use three. CIFAR-100 uses Adam at learning rate 10−3.2510^-3.25, batch size 100, cosine decay, 200 epochs, and gradient clipping at norm 5 for the main LaPrune schedule. Scaling results use combined forward and backward passes on an NVIDIA A100-SXM4-40GB. The solver stopping, fallback, and warm-start rules are given below. Numerical and Analytical Details Solver Tolerances and Clipping. The Newton solve for the joint budget and second-moment system defined in the main paper performs at most 30 iterations and freezes a batch row once |F1| |F_1| <10−7max(k,1), <0^-7 (k,1), |F2| |F_2| <10−7max(βk,1). <0^-7 (β k,1). We report absolute residuals e1=|F1|e_1=|F_1| and e2=|F2|e_2=|F_2|, and relative residuals e1/ke_1/k and e2/(βk)e_2/(β k). After 30 iterations, a row enters the robust fallback if either residual exceeds 10−410^-4 times the corresponding target clipped below at one. Let srs_r be the row-wise score standard deviation. Each Newton increment is clipped to Δb b ∈[−(8sr+1), 8sr+1], ∈[-(8s_r+1),8s_r+1], Δτ τ ∈[−1.5,1.5]. ∈[-5,5]. and the updated log-temperature is clipped to τ∈[−12,8]τ∈[-12,8]. The barrier itself has no global clipping interval. Bisection Fallback. The fallback uses 48 outer and 48 inner bisection iterations. The outer log-temperature bracket is τ∈[−12,8]τ∈[-12,8]. At a candidate t=eτt=e^τ, the inner barrier bracket is blo b_lo =rmin−40(t+sr)−1, =r_ -0(t+s_r)-1, bhi b_hi =rmax+40(t+sr)+1. =r_ +0(t+s_r)+1. The inner solve moves the bracket according to the sign of ∑ipi−k _ip_i-k. The outer solve moves it according to the sign of ∑ipi2−βk _ip_i^2-β k. Both use their fixed iteration limits without a separate stopping tolerance. The default is 48 outer and 48 inner iterations. The CIFAR-100 training runs use 20 and 20, the discretization-control experiment uses 40 and 40, and the supplementary ε -sensitivity sweep uses 24 and 24. The remaining reported experiments retain the default. In a documented diagnostic with 100 independent Gaussian rows, n=100n=100, and a=0.05a=0.05, the fallback was required for 0%0\%, 27%27\%, and 97%97\% of rows at γ=0.3γ=0.3, 0.90.9, and 0.990.99, respectively. This increase reflects the ill-conditioning as t↓0t 0. Interior values are restricted to the numerical bracket above, so values very near an endpoint may invoke the fallback. Asymptotic Warm Start. For the Gaussian approximation described in the main-paper discussion “Computing LaPrune: Joint Solves and Implicit Gradients,” the implementation estimates the density at b0b_0 using the row standard deviation s s and za=Φ−1(1−a)z_a= ^-1(1-a): ρ^(b0)=ϕ(za)max(s^,10−9). ρ(b_0)= φ(z_a) ( s,10^-9). We clip ρ^(b0) ρ(b_0) below at 10−410^-4 and initialize t0 t_0 =max10−3,4a(1−a)(1−γ)3ρ^(b0). = \10^-3, 4a(1-a)(1-γ)3 ρ(b_0) \. This is a Gaussian plug-in density estimate and therefore has no kernel bandwidth. The optional distribution-free large-n initializer instead draws m=⌈n⌉m= n scores, uses their empirical (1−a)(1-a)-quantile b b, and applies a rectangular-window estimator ρ^(b^) ρ( b) =12hm∑j=1m|rj−b^|<h, = 12hm _j=1^m1\|r_j- b|<h\, h h =12m∑j=1m|rj−median(r)|+10−6. = 12m _j=1^m|r_j-median(r)|+0^-6. again clipped below at 10−410^-4. Analytical Implicit VJP. The implementation evaluates the vector–Jacobian product from the implicit-differentiation system in the main-paper discussion “Computing LaPrune: Joint Solves and Implicit Gradients” analytically, using the reductions defined there. For upstream derivatives gi=∂ℒ/∂pig_i= /∂ p_i, let G0 G_0 =∑igidi, = _ig_id_i, G1 G_1 =∑igidizi, = _ig_id_iz_i, D D =S0U1−S1U0. =S_0U_1-S_1U_0. Then ∂ℒ∂ri ∂ r_i =dit(gi−Ci), = d_it(g_i-C_i), Ci C_i =G0(U1−S1pi)D+G1(S0pi−U0)D. = G_0(U_1-S_1p_i)D+ G_1(S_0p_i-U_0)D. For completeness, with ca=2a+(1−2a)γc_a=2a+(1-2a)γ, ∂ℒ∂a ∂ a =−n2DQa, =- n2DQ_a, Qa Q_a =(S1ca−2U1)G0 =(S_1c_a-2U_1)G_0 +(2U0−S0ca)G1, +(2U_0-S_0c_a)G_1, ∂ℒ∂γ ∂γ =−na(1−a)2D(S1G0−S0G1). =- na(1-a)2D(S_1G_0-S_0G_1). These derivatives apply only to the interior finite-temperature solution. Closed-Form Moments for Laplace Scores. For ρ=Laplace(0,1)ρ=Laplace(0,1) and b≥0b≥ 0, direct integration gives M¯1 M_1 =e−b−t2e−b/t2(1−t2)=a, = e^-b-t^2e^-b/t2(1-t^2)=a, (12) M¯2 M_2 =(4+t)e−b2(1+t)(4−t2)−t2e−2b/t4(4−t2)=βa. = (4+t)e^-b2(1+t)(4-t^2)- t^2e^-2b/t4(4-t^2)=β a. (13) The apparent singularities at t=1t=1 and t=2t=2 are removable and are evaluated by continuity. For b<0b<0, symmetry gives M¯1(b) M_1(b) =1−M¯1(−b), =1- M_1(-b), M¯2(b) M_2(b) =1−2M¯1(−b)+M¯2(−b). =1-2 M_1(-b)+ M_2(-b). Monotonicity and Uniqueness of the Temperature Proof of the Unique-Temperature Theorem. Let zi(t) z_i(t) =ri−b(t)t, = r_i-b(t)t, wi(t) w_i(t) =σ′(zi(t)), =σ (z_i(t)), z¯w(t) z_w(t) =∑iwi(t)zi(t)∑iwi(t). = _iw_i(t)z_i(t) _iw_i(t). Differentiating the budget constraint ∑iσ(zi(t))=k _iσ(z_i(t))=k gives ∑iwi(t)zi′(t)=0 _iw_i(t)z_i (t)=0. Since zi′(t)=−b′(t)+zi(t)t,z_i (t)=- b (t)+z_i(t)t, we obtain b′(t)=−z¯w(t),zi′(t)=−zi(t)−z¯w(t)t.b (t)=- z_w(t), z_i (t)=- z_i(t)- z_w(t)t. Consequently, dM2(t)dt dM_2(t)dt =2∑iσ(zi(t))wi(t)zi′(t) =2 _iσ(z_i(t))w_i(t)z_i (t) =−2t∑iwi(t)σ(zi(t))(zi(t)−z¯w(t)) =- 2t _iw_i(t)σ(z_i(t)) (z_i(t)- z_w(t) ) =−2t(∑iwi(t))Covw(z(t),σ(z(t))). =- 2t ( _iw_i(t) )Cov_w (z(t),σ(z(t)) ). Because σ is strictly increasing, the weighted covariance is positive whenever the scores are nonconstant, hence M2′(t)<0M_2 (t)<0. Continuity follows from continuity of the barrier solution. For pairwise distinct scores, the LapSum endpoint limits give M2(t)→kM_2(t)→ k as t↓0t 0 and M2(t)→akM_2(t)→ ak as t→∞t→∞. The intermediate value theorem and strict monotonicity therefore give a unique tγt_γ for every γ∈(0,1)γ∈(0,1). □ Proof of the Worst-Case Sparsity Floor Proof of the Worst-Case Sparsity Floor Theorem. Let U be uniform on [0,1][0,1] and set Y=f(U)Y=f(U). Then 0≤Y≤1,Y=a,Y2=aβ.0≤ Y≤ 1, =a, ^2=aβ. Write D:=[Y(1−Y)]=a(1−β)=a(1−a)(1−γ).D:=E[Y(1-Y)]=a(1-β)=a(1-a)(1-γ). For any y∈[0,1]y∈[0,1] and ε>0 >0, y<ε≥(1−y)(1−yε):1\y< \≥(1-y) (1- y ): the right-hand side is at most one when y<εy< and is nonpositive when y≥εy≥ . Taking expectations gives Pr(Y<ε)≥1−a−Dε=(1−a)(1−a(1−γ)ε). (Y< )≥ 1-a- D =(1-a) (1- a(1-γ) ). (14) Combining Equation (14) with the trivial lower bound zero proves the inequality. For tightness, set c=a(1−γ)c=a(1-γ). If 0<ε≤c0< ≤ c, take a two-level nondecreasing step function with values c and 11, assigning mass λ=a−c1−cλ= a-c1-c to the value 11. Its moments are a and aβaβ, and it has no value strictly below ε , so the zero branch is tight. For c≤ε≤βc≤ ≤β, assign masses mε m_ =a(1−β)ε(1−ε), = a(1-β) (1- ), m1 m_1 =a−εmε, =a- m_ , m0 m_0 =1−mε−m1 =1-m_ -m_1 to the ordered values 0,ε,10, ,1. These masses are nonnegative on this interval and satisfy εmε+m1=a,ε2mε+m1=aβ. m_ +m_1=a, ^2m_ +m_1=aβ. Moreover, m0=(1−a)(1−a(1−γ)ε),m_0=(1-a) (1- a(1-γ) ), so equality holds in Equation (14). Hence both branches are tight. □ Extension to β<ε≤1β< ≤ 1. For completeness, let V=Var(Y)=a(1−a)γV=Var(Y)=a(1-a)γ. Cantelli’s inequality gives the tight lower bound Pr(Y<ε)≥(ε−a)2(ε−a)2+V. (Y< )≥ ( -a)^2( -a)^2+V. Equality is attained by a nondecreasing two-level step function with values u=a−Vε−aandε.u=a- V -a . The condition ε≥β ≥β is exactly what ensures u≥0u≥ 0.