Add tree normalizing flow for phylogenetic variational inference - #116
Open
christiaanjs wants to merge 13 commits into
Open
Add tree normalizing flow for phylogenetic variational inference#116christiaanjs wants to merge 13 commits into
christiaanjs wants to merge 13 commits into
Conversation
A normalising flow on the unconstrained per-internal-node coordinates that sit
underneath the node-height ratio transform, so a variational approximation can
put a flexible, tree-structured distribution on node heights while the existing
ratio machinery keeps enforcing the ordering and positivity constraints.
Each layer is a traversal sandwich: a postorder (tip-to-root) affine map, a
learnable elementwise nonlinearity, then a preorder (root-to-tip) affine map.
Up-then-down is the collect/distribute schedule of belief propagation, and it is
why a single layer couples every pair of nodes -- a node ends up depending on
the descendants of its ancestors, i.e. on everything. Down-then-up couples only
ancestor-descendant pairs until a second layer is stacked on; the order is
configurable and both patterns are checked by a test.
Invertibility is enforced by construction: every affine scale is a softplus
output, so each triangular map is non-singular; the nonlinearity is monotone and
analytically invertible; and the recursion weights are bounded (tanh, and
tanh/num_children) so a deep tree cannot amplify them. The log-det-Jacobian of
each affine map is sum_i log scale[i] -- no traversal -- and both inverses are a
single gather, so only the forward sweeps are actual traversals.
The nonlinearity is pluggable: a monotone rational-quadratic spline shared
across the tree with per-node affine conditioning (default), sinh-arcsinh
(unbounded, reshapes the tails), or none at all, which reduces the flow to a
Gaussian with tree-structured covariance in O(internal_node) parameters.
The flow is optionally conditional -- on auxiliary variables (clock rate,
population size, substitution parameters) and on per-node/per-branch variables --
through small conditioners whose output weights start at zero, so the whole flow
starts as the exact identity.
- treeflow/traversal/tree_affine.py: the affine maps on the existing preorder
and postorder traversal primitives, plus their gather-based inverses.
- treeflow/acceleration/native/cc/tree_affine_op.cc: native C++ forward sweeps
with analytic reverse-mode gradients, alongside the existing native ops.
- treeflow/bijectors/{tree_affine_bijector,elementwise_node_flow,
tree_normalizing_flow}.py: the TFP bijectors.
- treeflow/model/approximation/tree_flow.py: a variational approximation
coupling a mean-field/full-rank/IAF parameter block to the conditional tree
flow, block-triangular so the density stays exact.
- experiments/tree_normalizing_flow.ipynb: comparison against the root
full-rank approximation, with heights-only and full-model studies.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Jnv3damzcA9TVa4n6ziwLM
Executed experiments/tree_normalizing_flow.ipynb on the 71-taxon YFV dataset and wrote up the findings. Budget-matched at 2500 steps per family, the flow reaches 63 nats more ELBO than the root full-rank approximation with the other parameters fixed, and ~1195 nats more on the full model. Two results worth noting: a free 70x70 covariance (whole-tree full rank) gains only 1.6 nats over the root-only block, so the Gaussian families are limited by optimisation rather than by covariance structure; and removing the flow's nonlinearity entirely still recovers 62.9 of the 63.1 nats, so at this budget the tree-structured triangular maps are doing the work and the nonlinearity is a refinement. Adds experiments/tree_flow_boundary_mass.ipynb, which retains the comparison behind a design question: should the nonlinearity sit after the constraining sigmoid, so that the flow can represent near-polytomies (heavy mass at a height ratio of 0 or 1)? Mass near a ratio boundary is mass in the tail of the unconstrained coordinate, so what decides it is the tail class: - the spline is the identity outside its window, so it reproduces the Gaussian coordinate's boundary decay exactly -- it contributes nothing there; - a spline applied in constrained space instead is bi-Lipschitz on [0, 1], so it moves the constant (~40x more mass within 1e-2 of the boundary) but leaves the exponent alone, and it saturates numerically once 1-r falls below float64 resolution; - sinh-arcsinh does change the class: at tailweight 2 the local exponent is flat at ~1.15, i.e. genuine power-law mass against the boundary. The notebook also fits each option to a target with boundary mass, which orders them the same way. So the placement is not what limits boundary mass, and the nonlinearity stays in unconstrained space alongside the affine tree maps (which have to be there anyway, since an affine combination of ratios leaves (0, 1)). This is documented in the elementwise_node_flow module docstring. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Jnv3damzcA9TVa4n6ziwLM
Moves the samplers out of treeflow.vi, which they were only ever housed in, into a treeflow.mcmc package: hmc.py and random_walk.py over a shared util.py that maps the model's constrained variables -- including the tree, through the node-height ratio bijector -- to unconstrained space and back. treeflow.vi.hmc was the one import path in use (the CLI and its tests) and is updated; hmc.py otherwise keeps its behaviour, now built on the shared helper. The new sampler is random-walk Metropolis-Hastings, intended as a reference posterior for judging the variational approximations: it uses no gradients and none of the flow machinery, so it is an independent check rather than a variation on what is being measured. It runs the chain inside a tf.function with the proposal scale as a tensor argument, so the whole chain compiles once and is reused across adaptation chunks -- roughly a hundred times faster than the eager loop, which matters when a usable reference needs tens of thousands of steps. Supporting machinery, since TFP's step-size adaptation kernels do not apply to RandomWalkMetropolis: - Robbins-Monro tuning of log(scale) during burn-in towards the 0.234 target acceptance rate; - per-coordinate preconditioning from the chain's own spread partway through burn-in (proposal only, so the sampled distribution is unchanged); - thinning, so a long run's memory stays bounded; - effective sample size per variable and, across chains, R-hat, with check_effective_sample_size turning them into a pass/fail. This is not decoration: a random walk over 70 correlated node heights reaches an effective sample size of order 10 per few thousand draws, so the check is what stops the reference being quoted as ground truth when it has not mixed. The experiment notebook gains a reference-posterior section per study, and every size in it -- optimisation steps, ELBO samples, chain length, burn-in, chains, thinning -- now reads an environment variable defaulting to the notebook's own value, so a browser run is unchanged. experiments/run_tree_flow_experiment.py drives both experiment notebooks with those as command-line flags, streaming progress bars live, in the style of examples/run_example.py. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Jnv3damzcA9TVa4n6ziwLM
The root full-rank fit on the full model diverged on one run -- RobustOptimizer skips non-finite gradient steps but gives up after 100 consecutive ones, which aborted the whole notebook partway through study 2. The fit helper now retries once at a fifth of the learning rate, as one would by hand, and otherwise records the family as diverged rather than raising, so the rest of the comparison still runs and the divergence is reported instead of being fatal. Downstream cells iterate over the families that produced an approximation. Also records the boundary-mass notebook's executed results and its conclusions: against a target with real mass at the boundary the held-out KL orders the options logit-normal 0.0328 > unconstrained spline 0.0157 > constrained-space spline 0.0045 > sinh-arcsinh 0.0029 > sinh-arcsinh + spline 0.0006 nats. The tail-exponent measurements explain the ordering -- only sinh-arcsinh changes the rate at which boundary mass vanishes -- with the nuance that a spline acting in constrained space does fit this target appreciably better than one confined to a finite window, since none of the density is left untouched. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Jnv3damzcA9TVa4n6ziwLM
Executed both experiment notebooks (71 taxa, 654 sites; ten variational fits and two reference chains, a little over two hours on CPU) and wrote up the findings. Heights only: every flow variant beats every Gaussian family by about 26 nats of ELBO, and the node-height marginals are roughly twice as accurate in the mean and 3-8x more accurate in the spread against the random-walk reference -- the Gaussian families put the root height at 596-656 with a standard deviation of 129-220, against the reference's 424 and 51, while the flows land within about 1% of the mean. Two results worth noting beyond that: root full rank is *worse* than mean field in this setting, exactly as its construction predicts once there are no non-tree parameters to share its block with; and the affine-only ablation already scores +26.1 of the flow's +26.4, so nearly all of the gain is the tree-structured linear coupling rather than the nonlinearity. Full model: the tree flow with an autoregressive flow on the parameters is 1,207 nats better than root full rank, of which about 540 is attributable to conditioning the tree block on those parameters, and it is the only family that reproduces the clock-rate x root-height dependence. The reference chains are reported with their diagnostics rather than as ground truth: minimum effective sample sizes of 40 (heights only, R-hat 1.16) and 10 (full model, R-hat 1.67) both fail the check, which is the expected behaviour of a random walk in these dimensions. The differences discussed are far larger than that uncertainty, but the fine ordering among flow variants is not something these references can settle, and the findings say so. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Jnv3damzcA9TVa4n6ziwLM
Replaces the 71-taxon YFV alignment with data simulated by experiments/benchmarks/simulate.py -- the same simulator the benchmark pipeline uses, which builds the alignment through get_sequence_distribution, the code path that scores likelihoods. Twenty taxa by default, 500 sites, sampled serially over a five-time-unit window so the clock stays identifiable, under known values. The point is not only that it is faster. On the YFV data the random-walk reference reached a minimum effective sample size of 40 with R-hat 1.16 and failed its check, so the accuracy comparison rested on a reference that was itself flagged as untrustworthy. On the simulated data the same chain reaches a minimum effective sample size of 2153 with R-hat 1.00 (5122 at twelve taxa), and passes. Twenty taxa is the compromise: nineteen node heights is enough tree structure for the comparison to mean something while the reference still mixes comfortably. Simulating also supplies a second yardstick that owes nothing to the reference chain: the values that generated the data. The accuracy tables now report error against the truth alongside error against the reference, plus 95% credible interval coverage of the true node heights, and the marginal plots mark the true values. The dataset's size, sampling window and seed are environment variables like every other size in the notebook, exposed as --taxa, --sites, --sampling-window and --data-seed on the runner, so the previous large-dataset regime is one flag away. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Jnv3damzcA9TVa4n6ziwLM
… bars Raises the optimisation budget to the 60,000 steps at a 1e-3 learning rate these models are known to need, and adds the RelativeLossNotDecreasing criterion that advi_iaf.ipynb uses (rtol 1e-6, 1000-step window, 5000-step minimum, three consecutive) so that budget is a ceiling rather than a target: each fit stops when its own ELBO stops moving, and reports the steps it took and whether it converged or ran out of ceiling. On the simulated dataset a fit runs at 24-39 steps per second, so the ceiling is 25-42 minutes per family. The progress bars needed throttling to match. At 2,500 steps they already wrote 2,525 carriage-return refreshes into the executed notebook -- 141 kB of stream output that buries the results it is meant to accompany, and twenty times worse at the new budget. They now refresh every thirty seconds (configurable) with a compact format, and the optimisation bar updates every hundred steps rather than every ten, which brings a fit's recorded progress down to a few dozen lines. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Jnv3damzcA9TVa4n6ziwLM
…rly exit The smoke run exposed the problem: with a 5,000-step floor and a 1e-6 tolerance every fit stopped at 5,001 steps, its own minimum, which defeats the point of raising the budget. Tightens the tolerance to 1e-7 and raises the floor to 20,000 steps, so a fit runs between 20,000 and 60,000 steps and stops early only when its ELBO has genuinely flattened. Both are exposed as flags on the runner. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Jnv3damzcA9TVa4n6ziwLM
RelativeLossNotDecreasing fired within ~50 steps of whatever floor it was given -- at 5,000 and at 20,000 steps, at tolerances of 1e-6 and 1e-7 -- so every family "converged" at exactly the floor. A criterion that reports convergence wherever its floor is put is measuring the floor, so early stopping is now off by default and the fits run the full 60,000 steps; it can be switched back on with TREEFLOW_TREE_FLOW_CONVERGENCE_MIN_STEPS. Also adds the reference chain to the clock-rate x root-height correlation table, so that diagnostic is adjudicated against the sampler rather than left as a comparison between approximations with no yardstick. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Jnv3damzcA9TVa4n6ziwLM
…tion `linear_cross_coupling=True` adds a dense linear term carrying the parameter block into every tree coordinate, `y_tree += W @ y_parameters`. The flow's existing auxiliary conditioning models that dependence by modulating the flow's per-node parameters through a small tanh network -- flexible, but nonlinear and indirect. This instead adds plain linear correlation between every parameter and every node height, which is what a joint Gaussian would have, and what root_full_rank provides for the root height alone. Combined with `parameter_approximation="full_rank"`, `nonlinearity="affine"` and no auxiliary conditioner it makes the whole approximation a structured joint Gaussian: full covariance among the non-tree parameters, tree-structured covariance among the node heights, and dense cross-covariance between them -- linear correlation everywhere, in parameter^2 + parameter*node + O(node) parameters rather than the (parameter + node)^2 of a full covariance. The term is a shift by a function of the parameter block only, so the Jacobian stays block-triangular and the density stays exact, and it starts at zero so the identity initialisation is preserved. The notebook gains that hybrid as a study 2 family, along with a variant keeping the cross term alongside the flow's nonlinearity, and a diagnostic reporting the reference posterior's correlation in both height space and the unconstrained coordinates the approximations actually work in. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Jnv3damzcA9TVa4n6ziwLM
Looks at the dependence the tree flow is built to capture directly in the likelihood rather than through a fitted posterior: fix every parameter at its maximum (treeflow's ML routine maximises the unnormalised posterior, so strictly a MAP), then vary pairs of node heights on a grid and plot the surface, for several sequence lengths. Every surface is drawn twice: over the node heights, where combinations that are not trees at all are masked and the shape of that region is part of the geometry; and over the unconstrained coordinates the approximations actually parameterise, where every point is a valid tree and a diagonal ridge is dependence a mean-field family would miss. Grid ranges come from each node's own feasible interval -- bounded by its tallest child and its parent -- rather than a fixed percentage, so a node wedged between a close child and a close parent still gets a usable window. Each surface is summarised by the correlation of the Gaussian matching its curvature, so the pictures have a number attached. Pairs include a parent and child, the root with its child, the root with a deep node, and a genuinely non-adjacent pair (neither an ancestor of the other). The last section switches to a relaxed clock, where a branch length is a rate times a time: it plots node height against the rate on the node's own branch and against a distant branch, and reports the posterior correlations between all node heights and all branch rates from the random-walk reference sampler, split by whether the branch is one of the node's own. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Jnv3damzcA9TVa4n6ziwLM
Replaces the write-up, which still described the earlier 71-taxon dataset, with the results of the 60,000-step run on the 20-taxon simulation. Heights only: the flow leads by +0.95 nats over root full rank against Monte Carlo standard errors of 0.03-0.07 -- solid but small. The interesting part is where it comes from. The affine-only ablation takes +0.39 of that, so the nonlinearity supplies the larger share, which is the reverse of the 71-taxon run where affine-only took +26.1 of +26.4. The correlation diagnostic explains it: the ratio transform leaves mean |correlation| of 0.026 in the unconstrained coordinates against 0.101 in height space, so there is little linear structure left to model, and the whole-tree full-rank Gaussian gains only +0.30 with 190 covariance parameters where the structured affine gains +0.39 with O(node). Full model: the conditional flow leads by +0.88, more than half of which is the auxiliary conditioning rather than flexibility within the tree. The new linear cross-coupling reaches +0.23 as a purely Gaussian family and recovers most of the clock-root ridge, but adds nothing on top of the conditional flow (+0.83 against +0.88) -- an alternative route to the coupling, not a complement. Both sections record what the numbers do not settle: node-height means are slightly worse under the flow while standard deviations are twice as good, the full-model reference did not mix on the tree block, and the cross-size comparison is uncontrolled. Also fixes the geometry notebook's memory use, which is what killed the previous run: it evaluated a whole 900-point grid at once, materialising a 13 GB partials tensor. Grids are now evaluated in chunks sized from the alignment length against a fixed budget. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Jnv3damzcA9TVa4n6ziwLM
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Implements a tree normalizing flow — a flexible, tree-structured variational approximation for node heights in phylogenetic inference. The flow puts a learnable distribution on the unconstrained per-node coordinates that sit underneath treeflow's existing node-height ratio transform, enabling richer posterior approximations while preserving all tree constraints.
Key Changes
Core Flow Implementation
treeflow/bijectors/tree_normalizing_flow.py: Main bijector implementing a "traversal sandwich" architecture:treeflow/traversal/tree_affine.py: Reference implementations of the two structured affine maps with forward/inverse operations and log-det-Jacobian computationtreeflow/bijectors/elementwise_node_flow.py: Nonlinear layer with three pluggable options (monotone rational-quadratic spline, sinh_arcsinh, or affine-only), each analytically invertibletreeflow/bijectors/tree_affine_bijector.py: TFP bijector wrappers around the affine maps for composition with other bijectorsVariational Approximation
treeflow/model/approximation/tree_flow.py: Variational approximation combining:Native Acceleration
treeflow/acceleration/native/cc/tree_affine_op.cc: C++ custom ops for forward sweeps of both affine maps with analytic reverse-mode gradientstreeflow/acceleration/native/tree_affine.py: Python wrapper for native ops with fallback to pure TensorFlowTesting & Experiments
test/bijectors/test_tree_normalizing_flow.py: Tests for identity initialization, invertibility, dependence structure, and conditioningtest/bijectors/test_elementwise_node_flow.py: Tests for nonlinearity layerstest/traversal/test_tree_affine.py: Tests for affine map correctness against NumPy referencetest/acceleration/native/test_native_tree_affine.py: Tests for native ops against TensorFlow reference and finite differencestest/model/approximation/test_tree_flow.py: Integration tests for the full approximationexperiments/tree_normalizing_flow.ipynb: Comprehensive notebook comparing the flow with existing root full-rank approximation on two studies (heights-only and full model)Notable Implementation Details
O(n)and inverses pure gatherstanh(|w| < 1) and child weights are normalized to prevent numerical issues on deep treeshttps://claude.ai/code/session_01Jnv3damzcA9TVa4n6ziwLM