# KL-optimal joint rare-event tilting

- ID: 2924
- Canonical URL: https://synthcore.org/idea/2924/kl-optimal-joint-rare-event-tilting
- API JSON: https://synthcore.org/api/idea/2924.json
- API Markdown: https://synthcore.org/api/idea/2924.md
- Verification status: unverified
- Source: [arXiv:2607.23134](https://arxiv.org/abs/2607.23134)
- Category: sampling
- Solves: sample-efficiency, accuracy, scalability
- ML areas: diffusion, diffusion-sampling, world-model, loss
- Math tags: probability, information-theory, stochastic-processes, optimization
- Ratings: usefulness 6/10; difficulty 5/10; novelty 4/10

## Idea description

Modify a diffusion sampler over a complete scenario trajectory z by exponentially tilting its prior toward a learned failure score s(z), rather than conditioning independently on environment and execution variables. The resulting sampler spends more evaluations in rare failure regions while preserving an explicit importance weight for estimating probabilities under the original distribution.

## Mathematical statement

The diffusion prior is p_0(z_0), with forward corruption q(z_t|z_{t-1}) = N(sqrt(alpha_t) z_{t-1}, (1-alpha_t)I), marginal q(z_t|z_0)=N(sqrt(bar_alpha_t)z_0,(1-bar_alpha_t)I), and bar_alpha_t=product_{s=1}^t alpha_s. Let s_phi(z_0) be a scalar failure or violation score learned from simulator labels, where larger values mean more likely or more severe failure, and let beta >= 0 be the tilt strength. The target distribution is p_beta(z_0)=p_0(z_0) exp(beta s_phi(z_0))/Z_beta, with normalizer Z_beta=E_{p_0}[exp(beta s_phi(z_0))]. It is the optimizer of minimizing KL(q||p_0) subject to a prescribed expected score E_q[s_phi]=c; equivalently, log Z_beta is the dual objective and d log Z_beta/d beta=E_{p_beta}[s_phi]. Its exact score is grad_{z_t} log p_beta(z_t)=grad_{z_t}log p_0(z_t)+grad_{z_t}log h_t(z_t), where h_t(z_t)=E[exp(beta s_phi(z_0))|z_t]. Approximate h_t with a learned noisy-state value model h_psi,t or score g_psi,t=grad_{z_t}log h_t. If epsilon_theta predicts forward noise, the prior score is approximately -epsilon_theta(z_t,t)/sqrt(1-bar_alpha_t), so the guided noise prediction is epsilon_guided=epsilon_theta-sqrt(1-bar_alpha_t)g_psi,t. Samples generated by the tilted chain receive importance weight w(z_0)=Z_beta exp(-beta s_phi(z_0)); self-normalized estimates can evaluate the original failure probability.

## Key formulas

- $$q(z_t\mid z_{t-1})=\mathcal{N}\!\left(\sqrt{\alpha_t}\,z_{t-1},(1-\alpha_t)I\right),\qquad q(z_t\mid z_0)=\mathcal{N}\!\left(\sqrt{\bar\alpha_t}\,z_0,(1-\bar\alpha_t)I\right),\quad \bar\alpha_t=\prod_{s=1}^{t}\alpha_s.$$
- $$p_\beta(z_0)=\frac{p_0(z_0)\exp\{\beta s_\phi(z_0)\}}{Z_\beta},\qquad Z_\beta=\mathbb{E}_{z_0\sim p_0}\left[\exp\{\beta s_\phi(z_0)\}\right].$$
- $$h_t(z_t)=\mathbb{E}_{p_0}\left[\exp\{\beta s_\phi(z_0)\}\mid z_t\right],\qquad \nabla_{z_t}\log p_\beta(z_t)=\nabla_{z_t}\log p_0(z_t)+\nabla_{z_t}\log h_t(z_t).$$
- $$\varepsilon_{\mathrm{guided}}(z_t,t)=\varepsilon_\theta(z_t,t)-\sqrt{1-\bar\alpha_t}\,g_{\psi,t}(z_t),\qquad g_{\psi,t}(z_t)\approx\nabla_{z_t}\log h_t(z_t).$$

## Implementation notes

Integrate the method into the reverse diffusion sampler at the noise-prediction or score-prediction line. Let z contain all jointly generated variables: environment parameters, initial conditions, disturbances, and optionally a rollout latent or action sequence. First train the ordinary diffusion prior p_0 with the forward schedule alpha_t. Then collect a small adaptive simulator dataset of pairs (z_0,y), where y is a binary failure label or a clipped continuous violation score. Train a value or guidance network v_psi(z_t,t) to estimate log h_t(z_t)=log E[exp(beta s_phi(z_0))|z_t]; in an MVP, train s_phi on clean z_0 and obtain noisy targets by sampling z_t from the known forward marginal, using a squared loss on log(exp(beta s_phi(z_0))+epsilon). At each reverse step, compute g=grad(z_t,v_psi(z_t,t)) with autodiff, obtain epsilon_guided=epsilon_theta-sqrt(1-bar_alpha_t)*g, and use the normal reverse Gaussian update with epsilon_guided. Clip g to a fixed norm and sweep beta in {0,0.5,1,2,4}; anneal beta from zero at high noise if guidance is unstable. Store s_phi(z_0) and estimate the prior failure rate with self-normalized weights exp(-beta s_phi), using log-sum-exp. What is learned from the paper's mathematics is the exponential tilt, conditional h-transform score, and importance correction; the diffusion schedule and value estimator are engineering approximations. First test on a 2D or 4D synthetic rare-event benchmark with a known Gaussian-mixture prior and failure region, then on a small diffusion trajectory model for a safety simulator. Compare unguided diffusion, ordinary classifier or reward guidance, and coordinate-wise conditional sampling at equal simulator calls and equal reverse steps. Pre-register these mechanism tests: (1) at fixed beta, the guided sample mean score must increase monotonically over beta and achieve at least a 2x failure rate at the same 1,000 simulator queries; (2) the empirical tilted score must satisfy d log Z_beta/d beta approximately equal to the sample mean score, with absolute discrepancy below 0.1 using finite differences; (3) importance-corrected failure estimates must agree with brute-force Monte Carlo on the synthetic problem within 10% relative error. Remove only g_psi for the ablation. The transfer is falsified if tilting increases nominal failure count but importance correction is biased by more than 10%, or if joint tilting gives no sample-efficiency advantage over ordinary conditional guidance at matched simulator cost.

## Disclaimer

AI-generated research hypothesis, automatically tested. Not peer-reviewed.
