inverse: explicit bond-angle diffusion with continuous triplet topology - #16
Open
crhysc wants to merge 4 commits into
Open
inverse: explicit bond-angle diffusion with continuous triplet topology#16crhysc wants to merge 4 commits into
crhysc wants to merge 4 commits into
Conversation
added 4 commits
August 29, 2026 17:08
Makes ALIGNN's three-body geometry an optional generative channel rather than only an input feature, and lets the line-graph topology vary continuously while the coordinates denoise. Everything is off by default: ALIGNNCSPDenoiser() with no arguments is the model that was there before, and a test asserts that enabling the angle head leaves eps_frac and eps_lattice bit-identical. Angular channel (FoldingDiff, doi:10.1038/s41467-024-45051-2). One extra head on the shared backbone's line-graph feature predicts the angular displacement the forward process introduced, trained with a wrapped smooth-L1 loss (beta = 0.1 pi). The angular latent reaches the coordinate and lattice heads through ALIGNN's ordinary angles -> bonds -> atoms path, so it is coupled rather than merely supervised. One deviation, documented rather than papered over: a crystal has no persistent theta_t to diffuse, because the triplet set is a function of the coordinates. The target is therefore computed on the triplet representation that exists at the current step, with both angles evaluated on the same periodic-image identity so it measures the corruption of one triplet rather than a change of neighbour. No new SDE or schedule was invented. Continuous topology (DimeNet arXiv:2003.03123, ReaxFF doi:10.1021/jp004368u). topology="radius" replaces the hard kNN rule with a radius candidate set, each pair weighted by the DimeNet envelope and each triplet by the product of its two bonds' weights. The envelope is reused from CutoffPolynomial, already in this repository. The weight multiplies the gate before both sums of the edge-gated average, which is the only placement making a zero-weight edge exactly equal to a deleted one. The graph is rebuilt from the current geometry every forward pass, so topology follows r_ij(t) with no annealing schedule, graph-temperature term or learned bondness network. Ablations are the deliverable, not the model: alignn/inverse/ablations.py defines A0-A4 and A6, run_angle_ablation.sh runs every arm on matched splits, budget and seeds, and the table in the README says which contrast answers which question (A5 is the A1-vs-A3 comparison, not a seventh configuration). Evaluation adds two mechanism metrics fixed before any run: generated vs held-out bond-angle distributions (FoldingDiff's diagnostic) and relaxation displacement to the nearest ALIGNN-FF minimum (MatterGen's proximity evaluation, doi:10.1038/s41586-025-08628-5). Not implemented, deliberately: DimeNet's spherical Fourier-Bessel basis (A6 substitutes the Fourier basis already in the repo), and every mechanism the brief rules out. Nothing has been trained. 32 new tests cover angular periodicity, the envelope and its first two derivatives at r_c, smooth pair and triplet fade-out, absence of a jump when a triplet enters or leaves the sparse graph, periodic-coordinate invariance, baseline equivalence, and that A4 provably cannot see the angular features. Claude-Session: https://claude.ai/code/session_01ULEhMd27yfBmKf5HXqmvFE
The model, training and sampling path was already pure PyTorch; this makes the evaluation helpers match rather than dropping into numpy for geometry. - evaluate.py computes bond angles with the same neighbour list and the same grouped-pair construction the line-graph builder uses, and does the histogram / KL / JS / Wasserstein arithmetic in torch. Results are identical to the numpy version: fcc still gives exactly 60/90/120/180 degrees and the Wasserstein distance is still calibrated in degrees. - angles.py drops an unused numpy histogram helper that duplicated evaluate.py's. - angle_eval.py follows. No new dependencies: setup.py and environment.yml are untouched, and the new modules import nothing beyond the standard library, torch and alignn itself (jarvis only lazily, for reading POSCARs). Adds a test pinning the evaluation geometry to fcc's known bond angles. Claude-Session: https://claude.ai/code/session_01ULEhMd27yfBmKf5HXqmvFE
The pieces were all in scripts/atombench already; what was missing was
anything tying them to a number in the paper. No multi-seed orchestration,
no aggregation with error bars, and no SLURM scripts anywhere in the repo.
Twelve tasks, each pinning the arguments to the existing scripts rather than
reimplementing them. A task is a list of units (one SLURM array element
each) and a unit is a list of stages: train -> generate -> symmetrize ->
score-nosym -> score-sym.
- Resumability is command-hashed, not file-existence based. Every finished
stage writes <rundir>/.stages/<name>.json with the exact argv, so bumping
--epochs retrains and rescores while re-submitting after a walltime kill
picks up where it stopped. A failed stage writes no marker.
- Stages declare their inputs, so a missing prerequisite is reported in a
second ("BLOCKED, missing .../pretrain_dft3d/best_model.pt (produced by
pretrain)") instead of crashing an hour in.
- Configurations shared between tasks are keyed by run directory and trained
once: bench-jarvis, arm A of ablation-linegraph and A0 of angle-ablation
are the same model.
- Both the unsymmetrised and symmetrised predictions are scored, because the
lattice columns are measured after symmetrisation and the force-field
paragraph is quoted before it. Each task declares which one --aggregate
reads by default.
Coverage is checked, not asserted. claims.py registers all 48 numbers the
inverse-design section prints -- both tables, the force-field paragraph, the
leakage fractions, the split sizes, the parameter match and the 2.4x step
cost -- against the task that regenerates each. `run_task.py verify` is a
to-do list in dependency order before you have run anything and a
measured-vs-published table afterwards, and exits non-zero if a claim ever
appears with no runner behind it. Parameter counts, training wall time, the
seed spread and the best individual run are derived from artefacts the
runner already writes, so the caption claims are measurable too.
Three cost settings, each in its own run tree so a cheap run can never
overwrite an expensive checkpoint: full (3000 epochs, 32 candidates, 3
seeds), --quick (300 / 8 / 2, whole test split, arms comparable to each
other but not to the published numbers), and --smoke (2 epochs, 4 targets,
plumbing only). --loss-only trains and stops, which needs no scoring
environment and no AtomBench clone; the denoising loss is the most
reproducible arm-vs-arm signal and a filter rather than a verdict, since the
line-graph loss gap reproduced across machines while match rate did not move
at all.
SLURM: cluster.env is the only site-specific file, and submit.sh sizes the
array from the arguments actually passed (--seeds 0,1 submits two elements)
and queues the aggregation with --dependency=afterok. The #SBATCH walltimes
are labelled placeholders; they have not been measured anywhere.
Two inconsistencies in the source material are flagged rather than papered
over. Table 4's "best single run" row pairs match 0.524 with RMSD 0.023,
which in the model registry belong to two different checkpoints, so
--aggregate names the run it is quoting. And the leakage paragraph says the
quoted results use no pretraining while the released csp_supercon_alex
behind the 0.485 row was fine-tuned from csp_pretrain_dft3d, so bench-alex
gained --from-scratch and pretrain-transfer runs both arms. Table 3's RMSD
spreads are the ones the inverse-design README calls a small-sample
artifact, so those tolerances are wide and the measured spread is printed.
EPOCHS["alex"] and EPOCHS["pretrain"] are not pinned by the manuscript and
say so; inspect_checkpoint.py reads the argument namespace out of a released
checkpoint to replace them with the published values.
data-jarvis was run end to end and reproduces the 847/105/103 split.
run_task.py doctor takes a real optimiser step rather than only importing
torch, which catches a mismatched install before a job is queued instead of
an hour into one.
Claude-Session: https://claude.ai/code/session_01RvR9sYTGsrRPBoUC41Vqir
Replaces the short overview of the angular channel with a reference section written for the group: an advisor and coauthors should be able to read it without having been briefed and know exactly what is in each arm. Adds, in order: a high-level overview of what angular denoising is, why it might matter, and the three falsifiable questions the suite asks — including what a negative result looks like, agreed in advance; a precise statement of what A0 already contains, since every arm is a delta against it; the mathematics of both additions with in-text citations separating what is taken from each source from what is ours; a paragraph on each arm; a confounds section; the protocol and metrics; a provenance table; and a bibliography. The three deviations from the literature are named and argued rather than buried: the angular state is induced rather than persistent (there is no fixed triplet list in a crystal), the target fixes the periodic-image identity, and the gate placement inside the edge-gated average is forced by the exact-deletion requirement rather than chosen. Also records a confound the earlier text did not: with gate_pair_messages on, the smooth arms truncate the pair interaction range as well as smoothing the topology. `--ablation A3 --gate-pair-messages 0` separates the two with no code change, and the section recommends running it. Numbers quoted from the graph construction were measured, not estimated (back-tracking triplet fraction, kNN vs radius triplet counts, envelope values at r_c/2). Three citation details are flagged inline as needing a check against the primary sources before they reach a manuscript. Claude-Session: https://claude.ai/code/session_01ULEhMd27yfBmKf5HXqmvFE
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
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
1 participant
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.
Makes ALIGNN's three-body geometry an optional generative channel rather than only an input feature, and lets the line-graph topology vary continuously while the coordinates denoise. Everything is off by default, and nothing has been trained — this is the machinery plus the ablation suite that can decide whether the idea works.
Branches cleanly off
develop(3f30f7a), no rebase needed. Merge whenever suits you.Motivation, from this repo's own data
The line-graph ablation already recorded in
alignn/inverse/README.mdis the reason this seemed worth building. Deleting the line graph and spending the same parameter budget on pair-graph depth costs a large, reproducible amount of denoising validation loss — 2.011 ± 0.018 against 2.351 ± 0.007, six models per arm, p = 0.002, reproducing on two machines — and leaves match rate exactly unchanged at 0.4709. Angular information demonstrably helps the network fit the score and demonstrably does not help it find the right structure more often. One reading is that angles are being used to interpolate rather than to constrain, in which case supervising them rather than merely observing them is the natural intervention. This PR is the test of that.Default behaviour is unchanged
ALIGNNCSPDenoiser()with no arguments is the model that is ondeveloptoday. A test loads the baseline's state dict into the angle-enabled model and assertseps_fracandeps_latticecome out bit-identical, so the angular channel is an addition rather than a perturbation. Released checkpoints keep loading: the weighted convolutions subclass the existing pure-torch layers with identical parameter names and shapes, andload_modelignores config keys that old checkpoints do not carry.What's new
Angular denoising channel. One extra head on the shared backbone's line-graph feature
zpredicts, per triplet, the angular displacement the forward process introduced, trained with a wrapped smooth-L1 loss. Target and loss form are FoldingDiff's (10.1038/s41467-024-45051-2). The angular latent reaches the coordinate and lattice heads through ALIGNN's ordinaryangles → bonds → atomspath, so the channel is coupled, not bolted on. No cross-attention, no second network, no Jacobian angle forces.Continuous topology.
topology="radius"replaces the hard kNN rule for triplet membership with a radius candidate set, each pair weighted by DimeNet's polynomial envelope (arXiv:2003.03123) and each triplet by the product of its two bonds' weights, which is ReaxFF's valence-angle switching (10.1021/jp004368u). The envelope is not re-derived — it is this repo's existingCutoffPolynomial. The weight multiplies the gate before both sums of the edge-gated average, the only placement that makes a zero-weight edge exactly equivalent to a deleted one (verified by test). Sincesis exactly zero beyondr_c, the radius graph is an exact sparsification rather than an approximation. The graph is rebuilt from the current geometry every forward pass, so topology followsr_ij(t)with no annealing schedule, graph-temperature term or learned bondness network.The deviation worth your eye
FoldingDiff can diffuse a persistent
Θ_tbecause a protein backbone's internal-coordinate list is fixed. A crystal has no such list — the triplet set is a function of the coordinates — so there is no independent angular variable to noise and noq(Θ_t | Θ_0)to write down. What is implemented is the closest well-defined thing: the angular target is computed on the triplet representation that exists at the current step, with both angles evaluated on the same periodic-image identity(i, j, k, n_ji, n_jk)so the difference measures the corruption of one triplet rather than a change of neighbour. Angles are an explicit denoising objective with their own head and loss; they are not an independently-noised variable, and any write-up should say "explicit angular supervision derived from the joint process" rather than "we diffuse angles".Also worth knowing: atom types are not diffused (the generator is conditioned on composition), so the state is
(F, L)and, with this branch,(F, L, Θ).Ablations
alignn/inverse/ablations.py;scripts/atombench/run_angle_ablation.shruns every arm on the same split, optimiser, budget and seed list.A0 vs A1 asks whether the objective helps on its own; A0 vs A2 whether a better-behaved noisy graph explains everything; A4 vs A3 whether the benefit is genuine coupling or would any auxiliary task on a shared trunk have done it; A1 vs A3 is the brief's A5 (hard kNN vs smooth radius), a comparison between existing arms rather than a seventh config. A4 keeps the angular features evolving, supervised, and back-propagating into the trunk — that is what auxiliary multitask learning is — but zeroes their forward contribution to the bond features, and a test confirms perturbing the angle embedding moves A3's coordinate score while leaving A4's at exactly zero.
Given the spread already recorded in the README (match rate 0.437–0.524 across fifteen models), one seed per arm settles nothing; the runner defaults to three.
Evaluation
Every AtomBench metric is preserved. Two mechanism metrics are added, fixed before any run: generated vs held-out bond-angle distributions (KL / JS / Wasserstein in degrees, calibrated to 0.001° on a synthetic shift — FoldingDiff's own diagnostic), and relaxation displacement to the nearest ALIGNN-FF minimum with volume change and energy drop (MatterGen's proximity evaluation, 10.1038/s41586-025-08628-5).
Confound flagged in the docs
With pair gating on, A2/A3/A4 differ from A0 in two ways at once — the topology is continuous and the interaction range is truncated at
r_c, whereas A0's dense pair graph lets every atom see every other atom in the cell.--ablation A3 --gate-pair-messages 0separates the two without any code change, and the README recommends running it before anything is written up.Pure PyTorch, no new dependencies
Verified rather than assumed: with
dgl,torch_geometric,torch_scatter,torch_sparse,networkx,matscipy,pymatgen,scipyandsympyall hard-blocked at import, the denoiser, training loop and reverse sampler import and run a full train step plus reverse sampling for every arm.setup.pyandenvironment.ymlare untouched; the new modules import only the standard library,torchandalignn(jarvis lazily, for reading POSCARs).Documentation
alignn/inverse/README.mdcarries a reference section written for readers who have not seen this work: the overview and the three falsifiable questions, what counts as a negative result agreed in advance, full mathematics for both additions with an in-text citation on every design choice saying what came from the source and what is ours, paragraph-level detail per arm, the confounds, a provenance table and a bibliography. It is marked as drafted by Claude, and three citation details are flagged in it as needing a check against the primary sources before they reach a manuscript.Tests
36 tests in
alignn/tests/test_inverse_angle_diffusion.py: angular periodicity of corruption and loss; the envelope and its first two derivatives vanishing atr_cby autograd; smooth pair and triplet fade-out; no jump in model output when a triplet enters or leaves the sparse graph (an atom swept through the cutoff, triplet count confirmed to change); zero-weight edge exactly equal to a deleted edge; periodic-coordinate and global-shift invariance for every arm; baseline equivalence; A4's severed coupling; forward/loss/backward/sampling for every arm; and the evaluation geometry pinned to fcc's known 60/90/120/180° angles.test_force_reduction.pystill passes. The four pre-existing collection errors intest_prop/test_eprop/test_alignn_ff/test_unified_calculator(cannot import name 'get_prediction' from 'alignn.pretrained') are already ondevelopand are untouched here.Also in this branch
task_runners/— one resumable runner per published inverse-design result, with arguments pinned, seeds handled and an sbatch script per task, including the A0–A6 suite above. Happy to split that into its own PR if you would rather keep them separate.