Sub-Gaussian Score Matching

Integrate this into the score-network training loss, immediately after computing the predicted score tensor and the teacher or denoising target score. For each minibatch sample \(i\), compute the residual vector \(e_i=s_\theta(x_i,t_i)-s^\star(x_i,t_i)\), flatten over spatial and channel dimensions, and use its squared Euclidean norm; optionally divide by the number of dimensions so the threshold is resolution-independent. Use \(\beta=2\), a small positive \(r\) such as 0.1, 0.5, or 1.0, and evaluate the log-mean-exp stably as \((m+\log\sum_i\exp(r\|e_i\|^2-m)-\log B)/r\), where \(m=\max_i r\|e_i\|^2\). Add the squared hinge penalty shown in the adapted formula, or directly minimize the MGF quantity with a Lagrange multiplier. Keep the usual denoising score-matching loss unchanged and ramp \(\lambda\) from zero during the first 10% of training to avoid early optimization instability. The mathematical quantity computed exactly is the empirical estimator of the exponential moment; \(\varepsilon\), \(r\), and \(\lambda\) are engineering hyperparameters. Estimate per-noise-level thresholds by maintaining an exponential-moving percentile of residual norms, or use a fixed threshold calibrated from a baseline model. First test on a 2D mixture of Gaussians and CIFAR-10 with a small NCSN or DDPM. Compare ordinary DSM against DSM plus the MGF penalty using identical compute, then corrupt scores only in rare low-density regions at sampling time. Measure FID or sliced Wasserstein distance, reverse-process divergence, and worst-decile residual norms. Success means similar training loss but substantially lower tail residuals and improved sample quality or stability under rare-region corruption; an initial target is a 20% reduction in maximum or 99th-percentile score error with no more than 5% training overhead.