Keeping Score: Efficiency Improvements in Neural Likelihood Surrogate Training via Score-Augmented Loss Functions

2026-05-12

AI Summary

This work addresses the challenge of parameter inference in stochastic processes, where conventional simulation-based inference methods suffer from computationally expensive likelihood evaluations and struggle to balance surrogate model accuracy against simulation cost under limited data. Breaking from the black-box assumption, this study introduces a novel approach that incorporates exact score information and a loss-gradient–based adaptive weighting mechanism into a probabilistic classification framework for neural likelihood surrogates, optimizing binary cross-entropy loss. The proposed method substantially enhances both the efficiency and accuracy of the surrogate model. Across multiple stochastic process benchmarks, it achieves downstream inference performance equivalent to using ten times more training data while requiring only 1.1× the original training time, effectively alleviating the data–cost trade-off bottleneck.

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_θ\log p(x \mid θ)$ 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.

Problem

• simulation-based inference
• likelihood surrogate
• stochastic process models
• parameter inference
• computational cost

Innovation

• score-augmented loss
• simulation-based inference
• neural likelihood surrogate
• stochastic process models
• gradient-based weighting