PaperlayerView exact revision on arXiv

Source-linked AI summary

Density Ratio Estimation with Stein Displacement Fields

Song Liu

arXiv:2610.12437v1stat.MLcs.LG

TL;DR

Density ratios and displacement fields are complementary views of distribution shift, yet are usually estimated separately. The paper models the log-ratio through a Stein displacement field and iterates fitting and movement for sampler correction and transformation fitting. Its applications report posterior correction on par with discriminator guidance while retaining an exact density, and improved pretrained ICA models.

  • Problem

    Density ratios and displacement fields are usually estimated separately, while ratio estimation can be fragile when distributions are far apart because samples from one may say little about the other.

    Source: §1 · p. 1, §2.1 · p. 2

  • Method

    SDRE models the log density ratio as minus the base distribution’s Stein operator applied to a displacement field, and repeated fit-and-move steps yield push-forward and pull-back algorithms.

    Source: §1 · pp. 1–2

  • Results

    In SBI, correction is on par with discriminator guidance while retaining an exact density; in nonlinear ICA, SDRE improves mean MCC in all six settings.

    Source: §1 · p. 2, §5.1 · p. 7, §5.2 · p. 8

  • Takeaways & Limitations

    The framework supports correcting a pretrained sampler without retraining and refining pretrained ICA models through pull-back.

    Source: §1 · p. 2

    Separated modes may require complex displacement fields; restricting the field class can limit transport, while highly flexible unregularized fields can destabilize high-dimensional ratio estimation.

    Source: §6 · p. 8

Abstract

from arXiv · show

Density ratios quantify distribution shift from a probability-mass point of view, whereas displacement fields describe, from a dynamical point of view, how one distribution is transported onto another. Although both offer complementary insights, they are usually estimated separately, and converting one into the other requires post-processing. In this paper, we estimate the density ratio between a target and a base distribution by parametrizing it through a displacement field acting on the base: the log-ratio is modeled as minus the Stein operator of the base applied to the field, up to a normalizing constant. This gives both statistical and dynamical descriptions of the distribution shift through a single convex optimization problem. Iterating this estimate-and-move step gives two inference algorithms: push-forward moves the model and corrects a pretrained sampler without retraining it, whereas pull-back moves the data closer to the base and fits a transformation model one layer at a time. Applications to distribution shift in simulation-based inference and to nonlinear independent component analysis illustrate the benefits and limitations of the approach.

1 INTRODUCTION

Density ratios and displacement fields offer complementary views of distribution shift, but are commonly estimated separately, and ratio estimation can be fragile when distributions are far apart. The paper connects the views by modeling the log-ratio through a Stein displacement field, then uses repeated fit-and-move steps for sampler correction and transformation fitting.

1 INTRODUCTION

  • Distribution comparisons can use divergences that measure differences in probability mass or transport methods that learn how to move one distribution onto another.The paper notes that divergences also serve as generative-model training objectives and that their gradient flows move distributions.Source: §1 · p. 1
  • Density-ratio methods can be fragile when distributions are far apart because samples from one may say little about the other, while transport methods’ prescribed paths need not follow a specified divergence’s gradient flow.The paper describes the first difficulty as the “density chasm.”Source: §1 · p. 1
  • The log-ratio is modeled as minus the base distribution’s Stein operator applied to a displacement field, enabling a convex fit of the ratio and a lower bound on KL(q ∥p) from target samples and the base score.The model uses no base-distribution samples.Source: §1 · p. 1
  • At a population optimum attaining the DV bound, moving the base along the estimated field changes its density to first order along the Fisher–Rao gradient flow of KL(q ∥p).The density update agrees with multiplying q by the estimated ratio, up to normalization and to first order.Source: §1 · pp. 1–2
  • Repeating the fit-and-move step yields push-forward, which appends displacement layers to correct a pretrained model without retraining, and pull-back, which moves data toward a parametric family for fitting.The paper associates these approaches with model updating and latent independent source recovery, respectively.Source: §1 · p. 2

2 BACKGROUND

This background introduces density-ratio estimation and Stein’s method as the two ingredients combined by the paper’s estimator. It also reviews a Stein-based ratio model and the limits of its out-of-sample use.

2 BACKGROUND

  • The estimator combines density-ratio estimation from samples with Stein’s method.Source: §2 · p. 2

2.1 Density ratio estimation

  • KLIEP fits a nonnegative basis expansion of p/q through a convex objective, while log-linear KLIEP parallels the DV representation; ratio estimation can be difficult when p and q are far apart.The density chasm arises because samples from one distribution may say little about the other, and reviewed bridging methods also require samples from both distributions.Source: §2.1 · p. 2

2.2 Stein’s method

  • The Langevin–Stein operator maps a vector field to a scalar whose expectation vanishes under q, and it depends on q through its score.In SVGD, maximizing the operator’s expectation over a kernel ball supplies a particle descent direction for KL(ρ ∥q).Source: §2.2 · pp. 2–3

2.3 Stein density ratio estimation

  • Liu et al. [49] form a linear ratio model from Stein-transformed feature gradients, so its normalization constraint holds at the population level without q samples; lack of guaranteed positivity away from observed samples limits out-of-sample evaluation.Source: §2.3 · p. 3

3 STEIN DISPLACEMENT RATIO ESTIMATION

Stein displacement ratio estimation models the log density ratio through the base distribution’s Stein operator, using target samples and the base score. The resulting field describes both the ratio and a transport step that moves the base toward the target.

3 STEIN DISPLACEMENT RATIO ESTIMATION

  • SDRE estimates log r as −A_q v plus a constant, using target samples and the score of a smooth, positive base distribution without requiring base samples.The population bound is tight when −A_q v equals log r up to a constant.Source: §3 · p. 3

3.1 The regularized estimator

  • The regularized finite-sample estimator minimizes an empirical objective with a ridge penalty, and its optimization is convex when the vector-field class is linear.The ridge level is selected using validation samples; the empirical diagnostic is not itself a guaranteed lower bound.Source: §3.1 · p. 3

3.2 One estimate, two perspectives on distribution shift

  • At the population optimum, moving the base along the estimated field changes its log-density in the direction of the Fisher–Rao gradient flow of KL(q ∥ p), to first order.The displacement update agrees with an explicit Euler step up to O(η²).Source: §3.2 · p. 4
  • One convex problem yields both a ratio estimate and a displacement, while invertible-map updates move existing mass; restricted fields and small steps can limit movement between separated modes.The field can also be restricted to encode known dynamical structure, such as linear or low-rank transformations.Source: §3.2 · p. 4

4 INFERENCE ALGORITHMS

The inference algorithms repeatedly estimate a ratio and construct a displacement layer, moving either a pretrained model toward data or data toward a simple base. Their choice depends on the inference objective and changes how scores and validation are handled.

4 INFERENCE ALGORITHMS

  • The inference loop estimates a field for the current base, constructs a candidate layer, and repeats after validation; it can move the model toward data or pull data toward an interpretable base.The two directions are called push-forward and pull-back, respectively.Source: §4 · pp. 4–5

4.1 Push-forward algorithm

  • Push-forward appends each accepted displacement layer to a pretrained invertible sampler, keeping its generator fixed while correcting it toward target data.The current model score is back-traced through the sampler and prior layers; validation detects discrepancy before accepting another layer.Source: §4.1 · p. 5

4.2 Pull-back algorithm

  • Pull-back transforms latent samples through each inverse displacement and composes the layer with the forward map, fitting the data representation toward a parametric prior.When the prior has unknown parameters, it is refit on current latent samples between layers, and its score is evaluated directly.Source: §4.2 · p. 5

4.3 Choosing between push-forward and pull-back

  • Push-forward suits model correction and recomputes scores through the full sampler; pull-back suits fitting an interpretable prior and evaluates its score directly, with validation rules matched to each objective.Push-forward accepts layers when a discrepancy is detected, whereas pull-back keeps a layer when it improves validation likelihood.Source: §4.3 · p. 6

5 APPLICATIONS

The applications test push-forward on posterior correction in simulation-based inference and pull-back on source recovery in nonlinear ICA. Results show correction of a frozen sampler under several shifts and gains from refining pretrained source models.

5.1 Correcting a posterior sampler in simulation-based inference

  • C2ST fell from 0.724 to 0.523 under a unimodal prior change, near retraining’s 0.511, using existing simulations.Table 1 reports C2ST against reference posteriors, where 0.5 means indistinguishable.Source: Table 1 · p. 7, §5.1 · p. 7

5.2 Recovering sources in nonlinear ICA

  • Mean MCC for OSF on synthetic nc = 1000 rose from 0.762 to 0.887 after SDRE refinement; refinement improved every pretrained model across all six settings.Table 2 compares standalone estimators with pretrained models refined by SDRE.Source: Table 2 · p. 8, §5.2 · p. 8
  • 39 of 44 paired runs favored SDRE refinement over four-times extended training, supporting gains beyond simply adding training steps.Table 3 reports seed means and counts the runs where refinement beats extended training.Source: Table 3 · p. 8, §5.2 · p. 8
  • From scratch, SDRE achieved the highest mean MCC among evaluated non-oracle estimators in five of six settings; oracle-structured OSF led all six.OSF uses oracle knowledge of the mixing architecture.Source: §5.2 · p. 8, Table 2 · p. 8

6 CONCLUSION AND LIMITATIONS

The conclusion presents SDRE as a convex way to estimate a density ratio and displacement field together, then use repeated updates in two inference algorithms. It also notes transport, numerical-stability, and computational limits.

6 CONCLUSION AND LIMITATIONS

  • SDRE models the log density ratio as minus the base Stein operator applied to a displacement field, and a convex fit estimates the ratio and a KL lower bound.At population-bound equality, the field gives a first-order Fisher–Rao update.Source: §6 · p. 8
  • Repeated fit-and-update steps yield push-forward and pull-back algorithms, evaluated for posterior correction in SBI and source recovery in nonlinear ICA.Push-forward appends layers to a pretrained sampler; pull-back transforms data toward a prior.Source: §6 · p. 8
  • Separated modes may require complex fields; restricted fields can limit transport, while flexible unregularized fields can destabilize high-dimensional ratio estimation.Push-forward also costs more as layers accumulate because it back-traces scores through prior corrections.Source: §6 · p. 8
  • The discussion relates pull-back’s greedy layer-wise estimation to end-to-end nonlinear ICA methods such as GIN [68] and compares transport-based Fisher–Rao approximations with kernel approaches.The paper also connects its layer-wise objective to likelihood-based normalizing flows [74].Source: §6 · p. 8

A PROOFS

The appendices justify the first-order transport interpretation, describe conditional use and implementation choices for both algorithms, and detail the field classes used in experiments.

A PROOFS

  • The proofs characterize equality in the population bound and show that, under stated smoothness and step-size conditions, the displacement map is invertible and gives a first-order Fisher–Rao update.The first-order agreement is up to O(η^2); the proof also establishes the equality condition for the bound.Source: A PROOFS · p. 13

B CONDITIONAL DISTRIBUTIONS

  • For conditional distributions, the objective averages conditional divergences; one target sample per condition suffices empirically, but a shared restricted field need not attain the conditional bounds.The conditional Fisher–Rao conclusion requires equality in the conditional bound at each condition.Source: B CONDITIONAL DISTRIBUTIONS · p. 14

C DETAILS OF THE INFERENCE ALGORITHMS

C.1 Push-forward

  • Push-forward’s sample-based step-size rule does not guarantee global invertibility, and using the same validation samples to select ridge level makes its stopping test somewhat optimistic.The candidate step size uses the maximum field-gradient norm over target samples.Source: §C.1 · p. 15

C.2 Pull-back

  • Pull-back selects candidate layers by validation likelihood, but likelihood alone cannot identify latent sources for an unconditional base with sufficiently flexible transports.Conditional structure and restrictions on the shared map can resolve the ambiguity under additional identifiability assumptions.Source: §C.2 · p. 15

D BACK-TRACED SCORE

  • Push-forward back-traces scores through the sampler’s invertible maps, computing inverses one at a time; samples whose inversion fails to converge are dropped before field fitting.The recursion keeps one Jacobian in memory at a time.Source: D BACK-TRACED SCORE · p. 15

E PULL-BACK: LIKELIHOOD AND INVERSES

  • Pull-back tracks latent samples and accumulated inverse log-determinants, evaluates candidate-layer likelihoods, and uses fixed-point or Newton iterations for inversion.Nonlinear candidates are discarded when sampled Jacobians are close to singular.Source: E PULL-BACK: LIKELIHOOD AND INVERSES · p. 16

F FIELD CLASSES

  • Finite-basis fields make the objective a strictly convex log-sum-exp plus quadratic penalty; the conditional push-forward field combines a global affine correction with a local gated component.For linear ICA, a first-order pull-back update recovers the relative-gradient update for maximum-likelihood ICA.Source: F FIELD CLASSES · pp. 16–17

G PUSH-FORWARD ON A DIFFUSION SAMPLER

The appendix describes push-forward corrections interleaved with a deterministic diffusion sampler’s predictor steps. It also clarifies the notation used for the nonlinear ICA experiments.

G PUSH-FORWARD ON A DIFFUSION SAMPLER

  • At each noise level, push-forward corrects the predictor-step sampler against target samples noised to that same level, then returns a model with an exact log-density.The procedure allows up to K correction layers per noise level.Source: G PUSH-FORWARD ON A DIFFUSION SAMPLER · p. 17

H NONLINEAR ICA: EXPERIMENTS AND ADDITIONAL RESULTS

  • In the nonlinear ICA experiments, “+ SDRE” denotes pull-back refinement of a pretrained model, while “SDRE from scratch” starts from the identity transformation.Source: H NONLINEAR ICA: EXPERIMENTS AND… · p. 18

H.1 Data and evaluation

The experiments test SDRE on synthetic and speech mixtures, where speech sources are only approximately segment-wise Gaussian. The methods use conditional bases and structured transport fields, with pull-back applied either from identity or to pretrained flows.

H.1 Data and evaluation

  • Speech sources are heavy-tailed and only approximately segment-wise Gaussian, so the conditional Gaussian model is misspecified for those benchmarks.The experiments also use synthetic data with five Gaussian sources and segment-wise means and standard deviations.Source: §H.1 · p. 18

H.2 Models and estimation

  • Pull-back uses a conditional base with segment-specific parameters and a field shared across segments; sources are recovered by applying the inverse fitted transformation.The base is Gaussian or smooth super-Gaussian, and restrictions on the field encode assumptions about the mixing.Source: §H.2 · p. 18
  • SDRE fits transformations layer by layer using convex problems, either from identity or by refining pretrained models with tractable likelihoods.From-scratch and refinement configurations differ in initialization, candidate field classes, and, for speech, base family.Source: §H.2 · pp. 18–19
  • The tested field classes include linear, coordinatewise, joint, rank-one, and rank-d forms; for linear fields, one first-order step is a relative-gradient update for maximum-likelihood ICA.Solving the layer problem exactly gives a convex Fisher-scoring-type update.Source: §H.2 · p. 18
  • Validation likelihood selects among pretrained flow configurations, while implementations solve the layerwise convex problems with L-BFGS or damped Newton.A full 40–100-layer refinement takes 6–29 seconds on one RTX 4090 with the JAX implementation, or 1.5–6 minutes on one CPU core.Source: §H.2 · p. 19

H.3 Main results and training controls

Across the six settings, SDRE refinement improves every pretrained model, and refining the tanh flow gives the best result in each setting. These gains generally exceed those from extended training.

H.3 Main results and training controls

  • Refining the tanh flow gives the best result in all six settings, and SDRE refinement improves the mean MCC of every pretrained model in every column.The refined tanh flow also has a smaller seed-to-seed spread than its pretrained version in every speech column.Source: §H.3 · p. 19
  • For the true-architecture flow, extended training changes mean MCC by at most 0.06, whereas refinement adds 0.02–0.18.Table 4 reports validation MCC over mixing seeds and the number of seeds where refinement beats extended training.Source: §H.3 · p. 20, Table 4 · p. 20
  • Refinement is better on 58 of 66 seeds across three pretrained models, while extended training changes mean MCC by −0.022 to +0.094.The largest extended-training gain, 0.762 → 0.856 for the tanh flow on synthetic nc = 1000, remains below refinement at 0.887.Source: §H.3 · p. 20

H.4 Additional analyses

Additional analyses examine pull-back behavior in nonlinear ICA and describe the simulation-based inference correction setup. They show that likelihood-based selection has boundaries, and that comparisons vary by task and evaluation setting.

H.4 Additional analyses

  • Stopping pull-back when the discrepancy falls below κ m/(2n) is counterproductive: all 40 layers average MCC 0.770, versus 0.736 or 0.717 for thresholds κ = 1 or 4.An oracle stopping at the best layer reaches 0.777; later layers can help despite being individually insignificant.Source: §H.4 · p. 20, Table 5 · p. 20
  • Selecting the refined pretrained model by final validation likelihood gives mean MCC 0.857, versus 0.855 for always refining the tanh flow, and is never worse across 28 seeds.Selecting by pretrained likelihood instead chooses GIN in 7 of 28 cases and yields 0.810 mean MCC.Source: §H.4 · p. 20
  • On one 20-second ARCTIC seed, refining the true-architecture flow raises validation likelihood from −1.34 to +0.90 while MCC collapses from 0.698 to 0.502.Refined tanh flow remains best at MCC 0.865 and 0.791 for d = 3 and 5, respectively; the comparison is confounded by changing speaker subsets.Source: §H.4 · p. 21, Table 6 · p. 21
  • On light-tunnel physical mixing, all nonlinear methods, including SDRE, remain below linear FastICA: MCC is 0.878 and 0.815 for FastICA versus at most 0.798 and 0.623.The sources are strongly sub-Gaussian within segments and change little across segments, while the tested Gaussian and super-Gaussian bases are misspecified.Source: §H.4 · p. 21, Table 7 · p. 21
  • With a log-cosh base, 23 SDRE layers reach validation log-likelihood −4.83 per day on stock returns, versus −4.92 and −4.98 for tanh and leaky-ReLU structured flows.The recovered components comprise one market factor correlated 0.99 with the S&P 500 and four stock-specific factors correlated 0.69–0.79 with their stock.Source: §H.4 · pp. 21–22

I SIMULATION-BASED INFERENCE: EXPERIMENTS AND COMPARISONS

  • The SBI experiments adapt a frozen NPSE sampler to prior changes, simulator misspecification, or partial coverage, using target pairs and a conditional displacement field across noise levels.Prior-change pairs are resampled from existing simulations; simulator-change pairs come from a modified simulator, and partial-coverage pairs cover only negative-bias neurons.Source: I SIMULATION-BASED INFERENCE:… · p. 23, §I.1 · p. 23

I.2 Sampler and corrections

  • Correction fields are fitted on target pairs and accepted by validation discrepancy; accuracy is measured by C2ST against reference samples, with chance at 0.5.The comparisons include guidance, classifier-based ratio estimation, fine-tuning, and retraining baselines.Source: §I.2 · p. 23, §I.3 · pp. 23–24

I.4 Statistical comparisons and limitations

Statistical comparisons find settings where SDRE outperforms fine-tuning and where another correction has lower mean C2ST. The section also explains why push-forward can stop short when target mass lies in separated modes, then situates SDRE among related methods.

I.4 Statistical comparisons and limitations

  • Under simulator change, SDRE outperforms fine-tuning at both target-pair counts; under partial coverage, it outperforms fine-tuning in both covered and uncovered regions.PriorGuide with Langevin has lower mean C2ST in all three prior-change settings, but its advantage is statistically significant only for the bimodal prior.Source: §I.4 · p. 24
  • Push-forward can stop short when the target requires mass to move between separated modes: its small invertible local moves yield C2ST 0.573 versus 0.510 for retraining.Langevin refinement is described as better suited to this case because it is not limited to small invertible moves.Source: §I.4 · p. 24

J RELATED WORK

  • Unlike methods that need samples from both distributions, SDRE estimates the ratio from target samples and the base score, while jointly estimating a displacement field.Related methods include direct ratio estimators and flow-based bridges; push-forward and pull-back re-estimate the ratio after each step rather than traversing a prescribed path.Source: J RELATED WORK · pp. 24–25
  • Pull-back resembles greedy Gaussianization and normalizing-flow methods, but each layer solves a convex problem and simultaneously estimates a ratio with a re-estimated, potentially conditional base.For nonlinear ICA, SDRE uses a conditional-base exact-likelihood flow model class like GIN, but replaces end-to-end training with convex layerwise estimation.Source: J RELATED WORK · p. 25
Loading 2610.12437v1…