PaperlayerView exact revision on arXiv

Source-linked AI summary

Sparsifying Stochasticity, Not Capacity: Partial Stochasticity via Deep Weight Factorization of Prior Scales

Marius P Linhard, Maurizio Filippone

arXiv:2610.09886v1stat.MLcs.LG

TL;DR

Which parameters should be stochastic remains an open model-selection question, and the hybrid inference scheme’s target has not been established. The paper learns the split through deep weight factorization of prior scales and functional-prior matching; experiments report competitive UCI performance.

  • Problem

    Theory does not indicate which parameters to make stochastic, and the target approximated by the commonly used hybrid inference scheme has not been established.

    Source: §1 · p. 1

  • Method

    The method learns the stochastic/deterministic split by factorizing prior scales and matching the functional prior to a Gaussian process with an MMD objective.

    Source: §1 · p. 2, §6 · p. 8

  • Results

    On UCI benchmarks, the method is competitive; certificate repair changes at most 28.1 parameters per split on average and never more than 34.

    Source: §6 · p. 8, §5.5 · p. 8

  • Takeaways & Limitations

    A cheap-to-check certificate guarantees universal conditional density approximation, and a minimal repair restores the guarantee when it fails.

    Source: §6 · p. 8

    The tracking-error analysis assumes the optimizer path is independent of sampler noise, so it does not cover feedback between the two blocks in the algorithm.

    Source: §6.1 · p. 8

Abstract

from arXiv · show

Bayesian neural networks need not be fully stochastic to be universal conditional density approximators, but it remains open which parameters should be stochastic. We learn this split by applying deep weight factorization to the prior scales, which are the standard deviations of the parameter priors, while fitting the functional prior to a Gaussian process with a maximum mean discrepancy objective. A parameter whose prior scale falls below a cutoff becomes deterministic and is optimized during inference, so the regularizer sparsifies stochasticity rather than capacity. We give a certificate for universal conditional density approximation that is checkable in linear time, together with a minimal repair when it fails. We further show that the common hybrid scheme of sampling some parameters and optimizing the others is stochastic approximation for a type-II maximum a posteriori objective, and that coupled step sizes can leave a tracking error that does not vanish as the step size shrinks. On a bimodal target, the learned split stays close to an unconstrained reference across all budgets and is insensitive to the cutoff, while random masks that distribute the same prior scales across layers are worse by up to two orders of magnitude. On UCI benchmarks, our method performs on par with a fully stochastic network while keeping about half of its parameters deterministic.

1 INTRODUCTION

The paper addresses how to choose stochastic parameters and their priors in partially stochastic Bayesian neural networks. It learns the split through factorization of prior scales and contributes results on functional-prior fitting, approximation guarantees, and hybrid inference.

1 INTRODUCTION

  • Existing theory does not specify which parameters should be stochastic, and brute-force search over assignments is computationally impractical.Choosing sensible priors over the stochastic parameters is an additional challenge.Source: §1 · p. 1
  • Factorizing prior scales learns the stochastic/deterministic split; zero-scale parameters remain optimizable during inference, preserving network capacity.The method sparsifies stochasticity rather than pruning weights.Source: §1 · pp. 1–2
  • The paper fits the functional prior to a Gaussian process with an MMD objective and gives a linear-time certificate for universal conditional density approximation, with a minimal repair when it fails.The degree-two polynomial kernel reduces the objective to moment matching with kernel-fixed weights.Source: §1 · p. 2
  • The paper characterizes hybrid sampling-and-optimization inference as stochastic approximation for a type-II MAP objective and shows coupled step sizes can leave nonvanishing tracking error.The tracking bound is attained by a Gaussian example.Source: §1 · pp. 1–2

2 BACKGROUND

The background introduces parameter and functional priors, inference by stochastic-gradient Hamiltonian Monte Carlo, partially stochastic networks, and deep weight factorization. These concepts establish how parameter uncertainty induces function distributions and how factorized parameters are represented.

2 BACKGROUND

  • A parameter prior induces a functional prior; conditioned on parameters, the network-defined distribution over functions is deterministic.This shifts the perspective from parameters to functions.Source: §2 · p. 2

Inference

  • The posterior is intractable because the marginal likelihood is intractable, so the paper samples it with SGHMC, which scales to large datasets.The cited method is stochastic gradient Hamiltonian Monte Carlo (SGHMC).Source: Inference · p. 2

Partially and fully stochastic neural networks

  • Fully stochastic networks randomize all parameters, whereas partially stochastic networks contain both stochastic and deterministic parameters.The MLP layer definition uses weights, biases, and an activation function.Source: Partially and fully stochastic neural… · p. 2
  • Deep weight factorization extends the equivalence between L1 regularization and a two-factor Hadamard parameterization to depth D > 2, but increases objective dimensionality and makes optimization initialization-sensitive.The Hadamard product is element-wise.Source: Partially and fully stochastic neural… · p. 2

3 RELATED WORKS

Related work covers functional priors, partially stochastic networks, sparsity methods, and approaches to selecting stochastic parameters. The paper distinguishes its split learning from methods that select from a trained MAP estimate or reduce network capacity.

3 RELATED WORKS

  • Functional-prior methods impose priors on network-induced functions, and prior optimization to emulate a Gaussian process can improve performance substantially.The review points to Fortuin (2022) for functional priors in Bayesian deep learning.Source: §3 · p. 2
  • Sharma et al. (2023) established theoretical and empirical support for partially stochastic networks as universal conditional density approximators under mild topological constraints.The area has since attracted work on applications and properties (Calvo-Ordoñez et al., 2024; Allec and Ziatdinov, 2025).Source: §3 · pp. 2–3
  • Automatic relevance determination shrinks prior scales continuously, while spike-and-slab approaches prune parameters to zero; this paper instead makes the split explicit and preserves capacity.The paper learns prior scales and means, then optimizes the functional prior to prune scales.Source: §3 · p. 3
  • MAP-based selection methods train a full network before choosing a split, whereas this method learns its split from the functional prior using only inputs and no trained model.Sparse subspace variational inference reduces capacity; fixing a structured subset to break weight symmetries is described as complementary.Source: §3 · p. 3

4 METHODS

The methods learn which parameters are stochastic by factorizing prior scales, then assess whether the resulting mask preserves conditional-density expressiveness and characterize the hybrid inference procedure.

4.1 Learning the split

  • Prior-scale factorization learns the stochastic split; scales below a cutoff become deterministic, but those parameters remain optimizable during inference.The penalty is applied to prior scales rather than network weights, so it sparsifies stochasticity rather than fitted capacity.Source: §4 · p. 3, §4.1 · p. 3
  • For factorization depth D ≥3, exact stationary points have a gap between zero and nonzero scales, though approximate optimization does not guarantee that gap in practice.The fitted model is reported insensitive to the cutoff over more than an order of magnitude.Source: §4.1 · p. 4
  • The degree-two polynomial-kernel MMD matches means and second moments, with their relative weights fixed by the kernel.Measurement locations mix training inputs with uniform samples; the uniform component constrains the functional prior away from training data.Source: §4.1 · p. 4

4.2 Preserving expressiveness

  • A linear-time certificate checks whether a masked network remains a universal conditional density approximator, and a minimal repair restores the certificate when it fails.The certificate concerns expressiveness, not whether optimization finds the approximating parameters; on UCI datasets the repair changes at most 34 parameters, or 0.3% of the network.Source: §4.2 · pp. 4–5, §4.1 · p. 4

4.3 Inference

  • The hybrid inference scheme is stochastic approximation EM for a type-II MAP objective, with the sampler acting as a Monte Carlo E-step.The deterministic parameters are regularized around their prior means; the regularization parameter controls how far they may move.Source: §4.3 · p. 5
  • Coupled sampler and optimizer step sizes can leave a non-vanishing tracking error as the step size shrinks; decoupling them makes the bound vanish under the stated conditions.The authors therefore let the optimizer step size decay faster than the sampler step size.Source: §4.3 · p. 5

5 EXPERIMENTS

The experiments evaluate the learned split and certificate on a bimodal target, examine prior optimization on a one-dimensional Gaussian-process toy dataset, and compare methods on UCI benchmarks. They also test how the cutoff and deterministic-parameter drift affect results.

5.1 Learned Split

  • The learned split is stable across data splits, keeps the first and output layers largely deterministic, and leaves about half of the middle layer stochastic.The certificate's coverage condition holds in every run; where |T| ≥ d0 fails, repair changes 1–34 parameters.Source: §5.1 · p. 6

5.2 Conditions of the certificate

  • The certificate's bias-coverage, fully deterministic-unit, and activation conditions are jointly needed; it fails on high-input-feature datasets, where repair changes at most 34 parameters per split.On the bimodal target, one randomized bias in an admitted position comes close to the unconstrained reference, while output-layer randomization, removing the nonlinearity, and randomizing every weight with fixed scales fail.Source: §5.2 · p. 6, §5.5 · p. 8
  • The learned split stays near W1 ≈ 0.06 across 6.5%–99.6% deterministic parameters; cross-layer scale permutations are up to two orders of magnitude worse, while cutoff changes leave distance at 0.061–0.066.Within-layer permutations have no effect; the cutoff increases the deterministic share from 90.2% at c = 0.003 to 99.7% at c = 0.1 across the reported λ settings.Source: §5.2 · pp. 6–7

5.3 The role of γ

  • As γ grows, deterministic parameters move farther from their prior mean and RMSE and NLL improve monotonically; smaller γ trades some predictive accuracy for better calibration.ECE is lowest at γ = 0.03 and close to its regime-(iii) value at γ = 2.5, with differences within one standard deviation; every finite γ substantially improves on regime (i).Source: §5.3 · p. 7

5.4 Prior optimization

  • On the one-dimensional toy dataset, the prior objective and both moment terms converge, and the posterior gives sensible uncertainty away from the data despite prior sample paths differing from the GP in smoothness.The objective targets the mean and second moment, not the non-Gaussian part.Source: §5.4 · p. 7

5.5 UCI benchmarks

  • DWF-PSNN ranks at least as well as fully stochastic ALS-BNN in all six benchmark combinations while keeping about half its parameters deterministic.Table 3 reports mean ranks across classification and regression datasets.Source: §5.5 · p. 8, Table 3 · p. 8

6 CONCLUSION

The paper learns which parameters should be stochastic by factorizing prior scales, and its conclusion highlights an expressiveness certificate and an analysis of hybrid inference. It also identifies limits to the current analysis and describes the authors’ use of generative AI.

6 CONCLUSION

  • The method learns the stochastic/deterministic split from prior scales while fitting the functional prior to a Gaussian process, and experiments show competitive UCI performance.The certificate for universal conditional density approximation is cheap to check, with a minimal repair when it fails.Source: §6 · p. 8

6.1 Limitations and Future Work

  • The expressiveness analysis covers a convolutional or transformer head above a fixed stack, but requires a head wider than the feature dimension and excludes the shared-weight stack.The inference analysis assumes the optimizer path does not depend on sampler noise; the algorithm’s sample-dependent Adam updates leave feedback between the blocks outside the analysis.Source: §6.1 · p. 8

6.2 AI Use Statement

  • The authors report using generative AI for mathematical claims, proofs, methodology, experiments, implementation, and editing, and say they checked the assisted work.They state that they take responsibility for the final content.Source: §6.2 · pp. 8–9

A.1 Learning the split

  • For factorization depth D ≥3, the effective regularizer on the factorized scale u is an ℓ_(1/D) quasi-norm.At stationary points, non-zero prior scales are bounded away from zero.Source: §A.1 · p. 12
  • When stochastic parameters are fewer than measurement points, the output has no Lebesgue density and both KL divergences are infinite, motivating moment-based fitting.For a target that is a Gaussian mixture, the objective matches the moments of its Gaussian moment projection.Source: §A.1 · pp. 12–13
  • With the degree-two polynomial kernel, squared MMD equals a weighted discrepancy in means plus a discrepancy in second-moment matrices.The second-moment term uses the Frobenius norm.Source: §A.1 · p. 13

A.2 Preserving expressiveness

The section gives a certificate for universal conditional density approximation and a minimal repair for masks that fail it. Its construction uses first-layer stochasticity to recover input and noise, then applies downstream universal approximation.

A.2 Preserving expressiveness

  • Stochastic first-layer biases yield noise independent of the input, while deterministic units recover the input; together they let the first layer recover noise and input for downstream approximation.The downstream network’s universal approximation property then gives arbitrary approximation of the conditional target when it has a continuous generator.Source: §A.2 · pp. 13–14
  • A minimal repair makes the least-cost set of first-layer units deterministic to restore full rank; a greedy minimum-weight basis selects the units, with cost O(d1 log d1 + d1d2).The repair cost is measured by the number of parameters changed.Source: §A.2 · p. 14
  • The certificate requires at least dL stochastic biases and d0 fully deterministic first-layer units, and its conditions are checkable in O(|Θ(1)|) time.The repair first adds stochastic biases if needed, then makes the least costly eligible units fully deterministic to meet the unit-count condition.Source: §A.2 · pp. 13, 15

A.3 Inference

The inference analysis connects the hybrid sampler-and-optimizer scheme to a type-II MAP objective and studies how sampler tracking responds to step sizes. It also provides an error decomposition organized around modeling choices.

A.3 Inference

  • The parameter τ interpolates between freezing the deterministic block and type-II maximum likelihood: as τ decreases to zero the posterior concentrates at its prior mean, while as τ increases Jτ approaches log Z.For finite positive τ, Jτ is the log marginal posterior density up to a constant.Source: §A.3 · p. 15
  • The near-degenerate approximation identifies the hybrid objective with stochastic approximation for a type-II MAP objective under the stated regularity assumptions.The KL term is minimized by the conditional posterior, reducing minimization to maximizing Jτ.Source: §A.3 · pp. 15–16
  • With coupled sampler and optimizer step sizes, the tracking-error bound does not vanish as the step size shrinks; decoupling them can make the bound vanish.The coupled bound is attained in the stated Gaussian example, where the error is at least G(1 − e^-1) for all t ≥1.Source: §A.3 · p. 16
  • The error decomposition assigns terms to specific modeling choices, providing a way to relate those choices to the total error.The decomposition is stated for bounded, 1-Lipschitz test functions after burn-in.Source: §A.3 · p. 17

B DETAILED UCI RESULTS

The supplied UCI appendix passages identify dataset-characteristics and per-dataset results tables, but do not provide their cell values.

B DETAILED UCI RESULTS

C THE LEARNED SPLIT

The learned split is stable across data splits, with deterministic parameters concentrated in the first and output layers while the large middle layer remains less sparsified. The learned masks also satisfy the certificate’s coverage condition, and the minimal repair changes only some parameters on datasets where the other condition fails.

C THE LEARNED SPLIT

  • Standard deviations are at most 0.04 in deterministic-weight share across splits; the first and output layers are largely deterministic, while the middle layer keeps about half its weights stochastic.Because the middle layer contains most parameters, it dominates the overall deterministic share; classification datasets are slightly more sparsified than regression datasets.Source: C THE LEARNED SPLIT · p. 19, Table 14 · p. 19
  • The middle layer is sparsified least despite a larger regularization weight, plausibly because its larger fan-in makes its prior-scale objective gradients unlikely to share the other layers’ bounds.The qualification is that Remark 4.2 assumes comparable gradient bounds across layers.Source: C THE LEARNED SPLIT · p. 19
  • Every dataset and split retains 7–27 stochastic first-layer biases out of 100, satisfying the coverage condition U(S) ≥ dL with a wide margin.Before repair, |T| ≥ d0 holds on BANKNOTE, HTRU2, and nine of ten CONCRETE splits; elsewhere repair changes 5.5–28.1 parameters per split on average, with CREDIT requiring the most.Source: C THE LEARNED SPLIT · p. 19
  • Output-bias stochasticity varies from none of the BANKNOTE splits to all CREDIT and WINE splits, but the certificate imposes no output-bias condition.Thus this variation does not affect the masked network’s expressiveness; the reported values describe masks before repair.Source: C THE LEARNED SPLIT · p. 19

D THE CONDITIONS OF THE CERTIFICATE

The certificate conditions are tested by fitting conditional generators on a bimodal target, with an unconstrained reference reaching 0.053 against a 0.020 sampling floor. A first-layer stochastic weight can satisfy the count and position yet perform worse, while the experiments also examine scale thresholding and certificate repair.

D THE CONDITIONS OF THE CERTIFICATE

  • Figure 6 plots distance to the true conditional against randomized output-layer parameter count, with certified configuration and sampling floor as horizontal references.Figure 7 shows variants at three inputs; its blue configurations are those admitted by the certificate, and the vertical range is fixed by the target density.Source: Figure 6 · p. 20, Figure 7 · p. 21
  • The bimodal target excludes Gaussian conditionals; variants are fitted directly as conditional generators with energy score and evaluated by Wasserstein-1 distance.Each variant uses two hidden layers of width 64; the best of five restarts is selected by validation energy score, not the reported distance.Source: D THE CONDITIONS OF THE CERTIFICATE · p. 20
  • The unconstrained reference reaches 0.053 Wasserstein-1 distance to the true conditional against a 0.020 sampling floor, indicating that fitting is not the binding constraint.The table reports mean distance over nine held-out inputs and standard error over five data seeds; its sampling floor is the distance between independent true-conditional samples.Source: D THE CONDITIONS OF THE CERTIFICATE · p. 20, Table 15 · p. 20
  • A variant randomizing a first-layer weight reaches 0.284 despite satisfying the certificate’s count and position conditions, because its stochastic contribution vanishes at x = 0.The predictive distribution then degenerates there, motivating Definition 4.4’s focus on biases.Source: D THE CONDITIONS OF THE CERTIFICATE · p. 21
  • At λ = 10^-5, most factors fall many orders of magnitude below the 0.01 cutoff, whereas at λ = 3 × 10^-6 the scales span the cutoff.The accompanying histogram uses log10 |uθ| and notes that larger λ drives pruned products toward zero.Source: D THE CONDITIONS OF THE CERTIFICATE · p. 21, Figure 8 · p. 22

E EXPERIMENTAL DETAILS

  • In experimental details, learned masks satisfy the coverage condition in all runs; where |T| ≥ d0 fails on high-feature datasets, minimal repair changes 1–34 parameters per run.The reported implementation uses an MLP with two 100-unit hidden layers and tanh unless stated otherwise, and reports no extensive hyperparameter search.Source: E EXPERIMENTAL DETAILS · p. 22

E.1 Regression experiments

The regression experiments define their calibration measure and target Gaussian process, then specify prior-optimization and sampling configurations through the accompanying hyperparameter tables.

E.1 Regression experiments

  • Regression ECE is the sum of squared differences between confidence levels j/10 and empirical shares, computed from Gaussian summaries of predictive means and variances.The sum runs over ten confidence levels and is not an average, so its scale is not comparable to classification ECE.Source: §E.1 · p. 22
  • The target Gaussian process has zero mean, variance 1, length-scale √2 · d0, and noise variance 0.1, with standardized regression targets.The passage introduces additional hyperpriors but does not specify their values here.Source: §E.1 · p. 22
  • The setup uses factorization depth D = 3 and cutoff 0.01; Tables 16–18 give prior-optimization and model-specific sampling hyperparameters.The passage also points to a separate sampling configuration and sets γ = 2.5.Source: §E.1 · pp. 22–23, Table 16 · p. 23, Table 17 · p. 23

E.2 Classification experiment

The later experimental subsections specify classification and prior-optimization settings, then compare learned and randomized masks on a bimodal target. Random masks at the learned budget satisfy the certificate on few datasets or draws beyond BANKNOTE and HTRU2.

E.2 Classification experiment

  • The classification experiment reuses the regression Gaussian-process setup, uses prior-optimization batch size 128, and sets γ = 2.5.Tables 19–20 list hyperparameters for models other than SNI-PSNN and for SNI-PSNN, respectively.Source: §E.2 · pp. 23–24, Table 19 · p. 24, Table 20 · p. 24

E.3 Prior optimization

  • The prior-optimization subsection says hyperpriors remain the same and sets γ = 25, with configurations provided in Tables 21–22.The passages identify the tables as prior-optimization and sampling hyperparameters without giving their cell values.Source: §E.3 · pp. 24–25, Table 21 · p. 24, Table 22 · p. 25

E.4 Role of γ

  • The role-of-γ experiments use γ = 0.03 and 2.5 for one setting and γ = 10^13 for another.Table 23 is identified as the sampling-hyperparameter table.Source: §E.4 · p. 25, Table 23 · p. 25

E.5 Learned and random splits on the bimodal target

  • For the bimodal-target comparison, learned scales use three-factor deep weight factorization and a threshold of 0.01, followed by mean refitting with the mask fixed.The study uses seven λ values, layer weights in a 3:4:5 ratio, and compares learned masks with random permutations within layer and parameter type or across all parameters.Source: §E.5 · p. 25

E.6 Sensitivity to the cutoff

  • Cutoff sensitivity is tested at 0.003, 0.01, 0.03, and 0.1 for λ values 3 × 10^-6 and 10^-5, followed by mean refitting.The fits use five data seeds and 12,000 steps each.Source: §E.6 · p. 25

E.7 Certificate of random masks

  • At the learned budget, random first-layer masks satisfy the certificate on BANKNOTE and HTRU2, in 97% of CONCRETE draws, and in at most 9% of draws on other datasets.For each dataset, 2,000 random masks use the learned shares of deterministic weights and biases and are checked against both certificate conditions.Source: §E.7 · p. 26
Loading 2610.09886v1…