Keeping Score: Efficiency Improvements in Neural Likelihood Surrogate Training via Score-Augmented Loss Functions
Authors: Alexander Shen, Mikael Kuusela
Summary
The authors improve simulation-based inference (SBI) for models with expensive likelihoods by augmenting the standard training loss with exact score information—the gradient of the log-likelihood with respect to parameters. SBI trains a neural network to approximate the likelihood using simulated data, but this is expensive; the authors show that when you can compute the score (even if the full likelihood is intractable), adding it to the loss function drastically improves the quality of the surrogate without needing much more training data. In their experiments, the method matches the performance of a 10× data increase with only a 1.1× increase in training time.
Main takeaways:
- Standard SBI treats the data-generating process as a black box and trains a likelihood surrogate via binary classification (real vs. simulated).
- This paper relaxes the black-box assumption: if you can compute the score (gradient of log p(x | θ)), you can add it as an auxiliary loss term.
- The score-augmented loss is combined with adaptive weighting based on loss gradients to balance the two objectives.
- Experiments on network dynamics and spatial processes show much better surrogate quality and downstream inference at a fraction of the simulation cost.
- Practical upshot: you get the benefit of ~10× more training data for ~1.1× more compute.
Relevance
Not directly related to my persona/midtraining work—this is about parameter inference for stochastic simulators. Included because it's a useful technique for training surrogate models more efficiently when you have partial analytical access to the target.
Abstract
arXiv:2605.12118v1 Announce Type: new Abstract: For stochastic process models, parameter inference is often severely bottlenecked by computationally expensive likelihood functions. Simulation-based inference (SBI) bypasses this restriction by constructing amortized surrogate likelihoods, but most SBI methods assume a black-box data generating process. While these surrogates are exact in the limit of infinite training data, practical scenarios force a strict tradeoff between model quality and simulation cost. In this work, we loosen the black-box assumption of SBI to improve this tradeoff for structured stochastic process models. Specifically, for neural network likelihood surrogates trained via probabilistic classification, we propose to augment the standard binary cross-entropy loss with exact score information $\nabla_\theta \log p(x \mid \theta)$ and adaptive weighting based on loss gradients. We evaluate our approach on case studies involving network dynamics and spatial processes, demonstrating that our method improves surrogate quality at a drastically lower computational cost than generating more training data. Notably, in some cases, our approach achieves downstream inference performance equivalent to a 10x increase in training data with less than a 1.1x increase in training time.