Triangular transport adaptation by aseyboldt · Pull Request #355 · pymc-devs/nutpie

aseyboldt · 2026-10-02T17:07:15Z

This (unfortunately gigantic) PR adds a new HMC preconditioning method, that allows fixing many issues in common posteriors, like many funnels and other more complicated posterior geometry.

It builds on the ideas of the previous coupling flow implementation, but makes use of more model specific information: It computes the sparsity structure of the Hessian, which allows us to find the Markov blanket of each variable. We use this to build a sparse factorization of the model. During warmup we then fit families of non-linear functions that precondition the posterior. Just as with the linear preconditioning and the coupling flows, we use the Fisher loss to find those functions. The optimization problem is solved using Levenberg-Marquardt, where the linear model is solved using conjugate gradient with block Jacobi preconditioning.

Using this in a model requires three steps:

First we compute the sparsity structure of the Hessian. For models that use Jax this is done using asdex. For Stan models we use Hessian vector products to infer the sparsity at individual sample points.
In a second step we then define a factorization order. By default, this is chosen to minimize fill-in, Just as sparse linear Cholesky decompositions do. Users can override this and specify that important hyperparameters should be at the front of the factorization so that they can more easily influence following parameters.
We then sample the model as usual with flow preconditioning.

model = """
data {
  int N;
}
parameters {
  real log_sd;
  vector[N] x;
}
model {
  log_sd ~ normal(0, 1);
  x ~ normal(0, exp(log_sd));
}
"""

compiled = (
    nutpie
    .compile_stan_model(code=model, cache=True, ad_hessian=True)
    .with_data(N=10)
    .with_hessian_sparsity(num_points=5, seed=5)
    .with_factorization(front=["log_sd"])
    .with_transform_adapt(
        nn_width=0,
        transformer="asymmetric_sas",
        transformer_layers=1,
    )
)

# Optionally plot the factorization
compiled.factorization.plot_spy()

trace = nutpie.sample(
    compiled,
    adaptation="flow",
    maxdepth=6,
    walnuts=True,
    tune=3000,
)

There are (currently way too many) settings for the preconditioner:

Often the most important is front=[list_of_variable_names]. Important hyper parameters of the model should usually be near the beginning of the factorization, so that following parameters can be reparameterized based on those values. The automatic order will often do this already, but in some models it helps to intervene. Inspect the spy plot of the factorization for details.

Other parameters of the preconditioner are set in the with transform adapt method:

What the flow can represent

nn_width (4): hidden units in each variable's MLP conditioner. 0 means the conditioners are only linear, i.e. a constant plus the skips.
scale_skip (True): adds a linear map from the parents to each variable's scale. Without it the scale can only depend on the parents through the MLP. If nn_width=0, scale_skip=False and `transformer="affine", the preconditioner is a sparse linear map, and can fit correlations but no non-linear features of the posterior.
transformer ("sas"): the 1-D map applied to each variable. The choices are "affine" (Gaussian conditionals), "sas" (sinh-arcsinh, for skew and tails), "tangent_sas" and "asymmetric_sas" (separate tails on each side).
transformer_layers (1): how many transformer layers to stack (not allowed for "affine").
scale_squash (3.0): turns the scale coefficient into a log scale, c·asinh(s/c). Use None for scales that are log-linear in the parents, such as funnels; the default keeps the scale from exploding away from the draws.

Diagnostics

verbose: (0): Setting to 1 gives some output about each new fit, setting to two also give info about each LM step during the optimization.
debug_save_bijection: Store each fit in nutpie.triangular_adapter._BIJECTION_TRACE. Horrible placeholder interface, this will change in a future release without notice!. It is used in for diagnostic plots in nutpie.flow_diagnostics.

Which draws the fit sees and when a fit is accepted

window_size (512): the maximum number of draws per fit, split into train and validation. When there are more draws, they are thinned.
val_fraction (0.3), val_block_size (16): the share of draws held out, taken in blocks of consecutive draws. The held-out draws decide whether a refit replaces the current flow and when LM stops early. 0 turns validation off.
forget_fraction (0.5), initial_skip (120): ignore the oldest half of all draws and the first 120 draws. Both protect against draws from before the chain converged.
num_diag_windows (9): how many early windows fit only a diagonal affine before the first triangular flow is built.
linear_fits (1): the first N fits of a new flow use only the Gaussian part, i.e. the constant plus the location skip, without the scale skip or MLPs.
recency (1.0): how strongly thinning favours newer draws. 0 thins evenly.
refit_diag (False): on each refit, re-standardise the diagonal affine to the current window and move the conditioners so the flow keeps its meaning.

Regularisation

These are all ridge penalties on a sum over draws, so they matter less as windows grow.

lm_shape_ridge (10): pulls each transformer towards affine, which keeps the SAS shape parameters near identity.
lm_mlp_ridge (1): ridge on the MLP weights.
lm_scale_skip_ridge (1), lm_skip_ridge (0.1): ridges on the linear scale and location skips. The location skip's penalty is measured relative to the conditional's width.
log_gamma_bounds (None): bounds the tail exponent of sas and asymmetric_sas layers, e.g. (-1, 1).
input_squash (None): the MLPs see c·asinh(y/c) of each parent, so their outputs grow logarithmically away from the draws. Skips are unaffected.

The LM optimiser

lm_patience (5), solver_rtol (0.05): stop when the last patience steps improved the loss (or the validation loss, with early stopping) by less than this relative amount.
early_stopping (True): stop on the validation loss and return the step with the best one.
max_epochs (60): the maximum number of LM steps per fit.
lm_min_loss (e⁻³): stop once the Fisher divergence is below this. Separately, a hard-coded threshold skips refitting entirely when the current flow already has log F below −3.
lm_linear_steps (80): the maximum number of CG iterations per LM step.
lm_max_exact_block_size (512): the largest Gauss–Newton sub-block that gets factored exactly for the preconditioner.
lm_forcing ("model"): the rule for when CG stops: "model" (Nash–Sofer), "residual" or "rho".
lm_lmp_size (8), lm_lmp_tol (1e-8): the limited-memory preconditioner built from recent CG directions, and its eigenvalue cutoff.

Rarely matters

activation (softplus, or "gelu"): the MLP activation.
zero_init (True): new MLP output weights start at 1e-3 times their usual scale, so a new flow starts near the identity.
tangent_sas_fix_b (False): only for tangent_sas.
conditioners (None): an explicit list of conditioner dicts. When given, it overrides nn_width, activation, input_squash, scale_skip, zero_init and the ridges.
stop_event: diagnostics and control only.

Trying it out

Steps to run it locally:

Install a rust toolchain with rustup
Check out the PR branch
Bridgestan needs libclang, so on ubuntu for instance apt install libclang-dev
In a python env with maturin installed, run maturin develop --release

ricardoV94 · 2026-10-21T15:16:08Z

The CI is slightly red

jessegrabowski

LLM-assisted review. I validated the findings but ya know

jessegrabowski · 2026-10-21T17:24:47Z

 pymc = ["pymc >= 5.20.1", "numba >= 0.60.0"]
 pymc-jax = ["pymc >= 5.20.1", "jax >= 0.4.27"]
-nnflow = ["flowjax >= 17.1.0", "equinox >= 0.11.12"]
+nnflow = ["flowjax >= 17.1.0", "equinox >= 0.11.12", "optimistix >= 0.1.0"]


The default adaptation="flow" path imports asdex, pymetis and scipy, and none of them are declared, so CI fails on every platform. Add them here, and bound jax below 0.12, since asdex 0.5.2 doesn't import with jax 0.12.

jessegrabowski · 2026-10-21T17:24:47Z

+            self, _hessian_sparsity=sparsity, _factorization=None
+        )
+
+    def with_factorization(self, order="metis", *, front=None):


Default to "amd", or fall back to it when pymetis is missing.

jessegrabowski · 2026-10-21T17:24:47Z

+                    )
+                except Exception as err:
+                    raise RuntimeError(
+                        "The triangular flow needs the Hessian sparsity of the "


coupling_type="masked" then fails on rust_flow, and rust_flow=False fails on scale_squash. Name the combination that works (coupling_type="masked", rust_flow=False, scale_squash=1.0) or drop the hint.

jessegrabowski · 2026-10-21T17:24:47Z

+            if sparsity is None:
+                try:
+                    sparsity = check_hessian_sparsity(
+                        self._detect_hessian_sparsity(), self.n_dim


seed=None, so the detected structure isn't reproducible from sample(seed=...). Pass the sample seed.

jessegrabowski · 2026-10-21T17:24:47Z

+                        "`coupling_type`."
+                    ) from err
+            factorization = self._factorize(sparsity)
+        return {"order": factorization.order, "sparsity": factorization.filled}


This overwrites a user's with_transform_adapt(order=...). Let the user value win or raise on the conflict.

jessegrabowski · 2026-10-21T17:24:49Z

+    fn step(&mut self) -> Result<StepInfo> {
+        let s = self.settings.clone();
+        let lam = self.lam;
+        self.ridge = Ridge::at(&self.problem, self.lin());


The ridge weights depend on theta and are refreshed every step without their derivative, so F and F_out aren't the same objective across steps. Freeze them for the whole fit.

jessegrabowski · 2026-10-21T17:24:49Z

+        .map_init(
+            || (vec![0f64; n], vec![0f64; n]),
+            |(vector, out), job| {
+                interrupt()?;


interrupt runs on rayon workers, where check_signals does nothing, so Ctrl-C can't stop Stan detection. Poll on the main thread and pass an AtomicBool.

jessegrabowski · 2026-10-21T17:24:49Z

+    let text: String = obj
+        .py()
+        .import("json")?
+        .call_method1("dumps", (obj,))?


json.dumps writes NaN/Infinity, which serde rejects, and raises on numpy ints. For example max_energy_error=inf doesn't round-trip.

jessegrabowski · 2026-10-21T17:24:49Z

+        let mut y = vec![0.0; positions.len()];
+        let mut g = vec![0.0; positions.len()];
+        for ((pos, grad), (y, g)) in positions
+            .chunks_exact(n)


chunks_exact(0) panics for draws with zero columns. Reject n_var == 0.

jessegrabowski · 2026-10-21T17:24:49Z

+/// How to walk the DAG.
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Deserialize)]
+#[serde(rename_all = "snake_case")]
+pub enum Schedule {


Only Serial is used. Dataflow overflows the stack on deep chains, and none of Levels, Dataflow or Auto is tested. Delete them.

jessegrabowski · 2026-10-21T17:27:07Z

+SCHEDULES = ("auto", "serial", "levels", "dataflow")
+
+
+def compile_transform(


compile_transform has no caller. Remove it, together with PySparseTriangularTransform.

jessegrabowski · 2026-10-21T17:27:07Z

+                raw = costs.mean()
+                normalized = raw / self.target_norm
+
+            # stick the landing


Remove, if False:.

jessegrabowski · 2026-10-21T17:27:07Z

+
+            return jnp.log(raw)
+
+            return normalized + jax.lax.stop_gradient(raw - normalized)


Unreachable after the return on 972.

jessegrabowski · 2026-10-21T17:27:07Z

+
+
+# Remat toggle for the per-draw residual, see `FisherLoss.residuals`.
+CHECKPOINT_RESIDUAL = False


CHECKPOINT_RESIDUAL is never set. Remove.

jessegrabowski · 2026-10-21T17:27:07Z

-    def __init__(self, gamma=None, log_inside_batch=False):
-        self._gamma = gamma
-        self._log_inside_batch = log_inside_batch
+def _huberise(residuals, delta):


_huberise is never called and huber_delta is never read.

jessegrabowski · 2026-10-21T17:27:09Z

@@ -492,13 +493,13 @@ fn update_nuts_from_nested_dict(
 ) -> PyResult<()> {
    match inner {
        NutsSettingsKind::Diag(settings) => {
-            *settings = depythonize(value).map_err(|err| PyValueError::new_err(err.to_string()))?;
+            *settings = crate::pyjson::from_py(value).map_err(|err| PyValueError::new_err(err.to_string()))?;


nit: run cargo fmt.

jessegrabowski · 2026-10-21T17:27:09Z

@@ -510,13 +511,13 @@ fn update_mclmc_from_nested_dict(
 ) -> PyResult<()> {
    match inner {
        MclmcSettingsKind::Diag(settings) => {
-            *settings = depythonize(value).map_err(|err| PyValueError::new_err(err.to_string()))?;
+            *settings = crate::pyjson::from_py(value).map_err(|err| PyValueError::new_err(err.to_string()))?;


nit: cargo fmt.

jessegrabowski · 2026-10-21T17:27:09Z

@@ -837,7 +838,7 @@ impl PyNutsSettings {
            "adaptation": adaptation,
            "settings": settings,
        });
-        let obj = pythonize(py, &value).map_err(|err| PyValueError::new_err(err.to_string()))?;
+        let obj = crate::pyjson::to_py(py, &value).map_err(|err| PyValueError::new_err(err.to_string()))?;


nit: cargo fmt.

jessegrabowski · 2026-10-21T17:27:09Z

@@ -893,7 +894,7 @@ impl PyMclmcSettings {
            "adaptation": adaptation,
            "settings": settings,
        });
-        let obj = pythonize(py, &value).map_err(|err| PyValueError::new_err(err.to_string()))?;
+        let obj = crate::pyjson::to_py(py, &value).map_err(|err| PyValueError::new_err(err.to_string()))?;


nit: cargo fmt.

jessegrabowski · 2026-10-21T17:27:10Z

+        if n_in > 0 {
+            pad.seed_y[..self.num_params].copy_from_slice(&pad.dy_dtheta);
+            pad.seed_l[..self.num_params].copy_from_slice(&pad.dld_dtheta);
+            let (out_y, out_l) = if degree == 0 {


nit: cargo fmt.

aseyboldt force-pushed the prepare-0.17.0 branch 12 times, most recently from d794720 to 78ecea8 Compare October 16, 2026 10:15

aseyboldt force-pushed the prepare-0.17.0 branch from 78ecea8 to e3673d0 Compare October 16, 2026 10:39

aseyboldt added 17 commits October 16, 2026 14:39

feat(flow): implement triangular transform

8a8ab74

feat(flow): implement triangular transform in numba and rust

4c3c47f

feat(flow): add linear shift shortcut in triangular

787ce0b

feat(flow): add line search in triangular LM

561f7ff

feat(flow)d axact preconditioning blocks in LM

6003471

feat(flow): use rust triangular in leapfrogs and clean docs

82f8ec8

feat: implement stan sparsity detection

c486e02

feat: sparsity and factorization in compiled models

a99e9c5

feat(flow): better verbose output and optional cholesky jitter

ea98d12

feat: support KeyboardInterrupt in hessian calculation

99cde3b

feat: forward keyboard interrupts to flow training

bba99a9

feat: diagnostic plots for the fitted flows

daa978b

feat: describe new flows in verbose mode

718492e

feat: improve flow_influence estimation

3e25625

style: some ruff fixes

82bf133

feat: add a little more verbose optimizer output

6cb6635

feat: add some optional diagnostics at verbose = 3

a04ef64

aseyboldt added 13 commits October 16, 2026 14:39

feat(flow): first fit of flow is linear only

d463a1c

feat(flow): fit triangular flows without jax

7e7c726

style: some ruff fixes

d5ab18f

perf(flow): reduce max memory usage

10c00f6

feat: remove fisher regularization

c99310d

feat(flow): improve ridge penalty

a61eea3

feat(flow): improve verbose output

d73e1cb

feat(flow): allow multiple linear fits and fix nonlinear-parameters

b1d2c4f

feat(flow): restructure transformers and add asymmetric_sas

f667f3a

feat(flow): add linear skip to scale parameter

e5030ef

feat(flow): add refit_diag option

bbbd451

refactor(flow): traits for Conditioner and Transformer

7a22c76

perf(flow): use the new conditioner trait for sampling

ce9eee6

aseyboldt force-pushed the prepare-0.17.0 branch from e3673d0 to c9ecb81 Compare October 16, 2026 12:43

aseyboldt changed the title ~~WIP Triangular transport adaptation~~ Triangular transport adaptation Oct 16, 2026

aseyboldt added 3 commits October 16, 2026 14:52

fix: better error message on stan usage without stan support

b59a4d3

chore: update pyo3 and drop pythonize dependency

fd3c9d4

fix: use rust flow in triangular diagnostics

7113d6f

aseyboldt force-pushed the prepare-0.17.0 branch 2 times, most recently from e78b2ef to 8050b27 Compare October 16, 2026 16:09

aseyboldt added 5 commits October 16, 2026 18:24

test: more sampler accuracy tests

19f3cbd

feat(flow): implement scale_squash parameter

623e512

fix: manually enable empty rayon pool in wasm

9e964d6

ci: mark flow tests

76ceae7

chore(release): prepare release with changelog

13e1543

aseyboldt force-pushed the prepare-0.17.0 branch from 8050b27 to 13e1543 Compare October 16, 2026 16:34

ricardoV94 approved these changes Oct 21, 2026

View reviewed changes

jessegrabowski reviewed Oct 21, 2026

View reviewed changes

		SCHEDULES = ("auto", "serial", "levels", "dataflow")


		def compile_transform(


		return jnp.log(raw)

		return normalized + jax.lax.stop_gradient(raw - normalized)



		# Remat toggle for the per-draw residual, see `FisherLoss.residuals`.
		CHECKPOINT_RESIDUAL = False

Conversation

aseyboldt commented Oct 2, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What the flow can represent

Diagnostics

Which draws the fit sees and when a fit is accepted

Regularisation

The LM optimiser

Rarely matters

Trying it out

Uh oh!

ricardoV94 commented Oct 21, 2026

Uh oh!

jessegrabowski left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

This branch has not been deployed

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

aseyboldt commented Oct 2, 2026 •

edited

Loading