Refined Topology-Pass Scheduling: Damped Trust-Region Compensation for Dynamic k-NN Graphs in SE(3)-Equivariant Structure Prediction

Jiankai Huang; Hanwen Shen

cs.LGOct 13, 2026

Refined Topology-Pass Scheduling: Damped Trust-Region Compensation for Dynamic k-NN Graphs in SE(3)-Equivariant Structure Prediction

Authors: Jiankai Huang, Hanwen Shen

Abstract

Equivariant structure-prediction pipelines alternate coordinate updates with the construction of a k-nearest-neighbour (kNN) graph whose membership is a discrete, piecewise-constant function of the coordinates. Rebuilding that graph at every block boundary (PB) is the expensive part of an epoch, whereas reusing one frozen graph (F) is fast but leaves a topology gap. We study a two-pass design in which the cheap pass records a small amount of state (margins, kNN radii, measured per-boundary displacements) and a second pass compensates conservatively. We formulate the choice of the graph-construction basis as a damped fixed-point iteration whose unique fixed point is exactly PB, and derive (i) exact non-flip certificates with the two-sided constant $m_i/4$ and one-sided constant $m_i/2$ , (ii) the minimum-variance (BLUE) replay gain $γ^*$ and the resulting feasibility crossover, (iii) an $\ell_\infty$ trust-region projection that preserves convergence, and (iv) an analytic near-tie spacing law $E[m] \approx r_k/(3k)$ that matches measurements to 1%. Three compensation levels are implemented in C++ (L1: frozen set with live geometry; L2: margin-gated local repair; L3: damped, projected basis) and validated under a matched-conditions protocol (identical seed, MSA depth and trajectory family, only the topology mode differs). L2 reproduces the PB edge set exactly ( $D = 0.0000$ at every boundary), L1 stays within 0.6% and L3 interpolates monotonically: L2 reproduces the PB edge set exactly when the trust region is removed, and L3 interpolates between the frozen and the PB topology. The recorded trajectory itself yields a measured $1.5\times$ Pass-1 and $1.76\times$ epoch speed-up with no measurable loss gap. An open-source implementation of the runtime and of all the experiments reported here is available at https://github.com/AlbertHuang96/ProteinPredictMachineLearning.

Figures & tables

Level	Modified object	Cost / property
$\mathsf{L}_{1}$	only geometry: $D(B^{t}),W(B^{t})$ (edge set stays $E(x^{0})$ )	$O(\|\mathcal{E}\|)\approx 7$ k $\Rightarrow$ free; removes the $0.29$ – $0.4$ Å geometry term
$\mathsf{L}_{2}$	edge set locally: $E_{t}$ in ( 35 ), evaluated inside $C_{i}$	$O(\sum_{i\in\Omega}\|C_{i}\|\log\|C_{i}\|)$ with $\|C_{i}\|\leq 1+k+k^{2}$ (two-hop); $\Omega$ closed under frozen neighbours
$\mathsf{L}_{3}$	the basis $B^{t}=\Pi_{\kappa}(x^{0}+\sum_{s<t}\gamma_{s}\Delta^{s})$ (predictor only)	storage $O(LT)$ ( $T{=}8$ : $\approx 2.5$ k floats); needs $\sigma_{p}\!<\!\sigma_{d}$
$\mathsf{L}_{4}$	full rebuild every boundary ( $\mathsf{PB}$ )	$1.8$ s/block; reference

TABLE 1: What each level touches in ( 35 ); costs from the shared cost model.

Family	$L$	boundaries	$\|\mathcal{E}\|$	Jaccard
P62805 (degenerate)	103	6/6	7105–7107	1.0000
P62891	51	6/6	2550 (ident.)	1.0000
P69905	142	3/3	9089 $=64L{+}1$	1.0000
P62241	208	3/3	13588, 13595	1.0000
Runtime mode			$\mathcal{D}$ ( $t{=}1$ / $t{=}2$ )	$J(t,\text{frozen})$
$\mathsf{L}_{1}$ (frozen set, live geometry)			0.0057 / 0.0068	1.0000

TABLE 2: Correctness verification. Top: family-reconstruction audit (offline rebuild of make_graph from the recorded coordinates must reproduce the dumped edge set exactly). Bottom: runtime verification in the matched-conditions protocol (P69905,

L{=}142

, seed 1234, MSA depth 3, 3 boundaries;

\mathcal{D}

measured against a

\mathsf{PB}

reference of the same family).

Sample	$L$	collisions	mean $C_{\alpha}$ dist.	margin median	drift (Å)
P69905	142	$0.01\%$	$19.2$ Å	$0.0787$ Å	$0.479$ ( $0.074$ )
P62241	208	$0.014\%$	$29.6$ Å	$0.118$ Å	$0.387$ ( $0.043$ )
P62805 (deg.)	103	$86.8\%$	$1.63$ Å	$\equiv 0$ (103/103)	$1.026$ ( $0.263$ )

TABLE 3: Sample characterisation. Collisions: fraction of point pairs closer than

10^{-3}

Å. Drift:

C_{\alpha}

RMSD of the last recorded boundary against

x^{0}

(in parentheses: mean

|\Delta d|

over that boundary’s own edge set).

Item	baseline	optimised
epoch wall time (two-pass vs $\mathsf{PB}$ )	$34.1$ s	$19.4$ s ( $1.76\times$ )
Pass-1 stage	—	$1.5\times$
loss ( $\mathsf{F}$ vs $\mathsf{PB}$ , 3 seeds, $\pm 1.3$ )	$15.91$	$15.49$
kNN rebuild per block	$1.8$ s	$\mathsf{L}_{1}$ : $0$ ; $\mathsf{L}_{2}$ / $\mathsf{L}_{3}$ : $O(\|\Omega\|)$
geometry refresh (frozen set)	$O(\|\mathcal{E}\|)$	free (held live)

TABLE 4: Measured optimisations (from the original pipeline and from this implementation). Costs follow the shared cost model: full-set refresh (kNN rebuild)

\approx 1.8

s/block, frozen-set geometry refresh

O(|\mathcal{E}|)

(free), local rebuild

O(\sum_{i\in\Omega}|C_{i}|\log|C_{i}|)

with

|C_{i}|\leq 1+k+k^{2}

(the implementation uses partial_sort , i.e.

O(\sum_{i\in\Omega}|C_{i}|\log k)

), extra pass

\approx 8

–

11

s/epoch.

Sample	setting	$\kappa_{\text{med}}$ (Å)	$\mathcal{D}_{w}$	reduction	$\|\Omega\|$
P69905	frozen ( $\gamma{=}0$ )	—	0.0578	—	0
P69905	$\sigma{=}0.5$	0.020	0.0515	$-11\%$	133
P69905	$\sigma{=}1$	0.039	0.0458	$-21\%$	116
P69905	$\sigma{=}2$	0.079	0.0401	$-31\%$	80
P69905	$\sigma{=}4$	0.157	0.0274	$-53\%$	47
P69905	$\sigma{=}8$	0.315	0.0227	$-61\%$	26

TABLE 5: Offline oracle replay (

\gamma=1

) on both samples: weighted topology distance

\mathcal{D}_{w}

against the

\mathsf{PB}

edge set at the same boundary, and its reduction against the frozen basis (

\gamma=0

B\equiv x^{0}

|\Omega|

is the risk-set size at

t=1

for both samples, out of

L

residues (identically

0

when

B\equiv x^{0}

); the five-point

\sigma

sweep is complete on both.

Protein Structure Prediction

Protein structure prediction uses computational models to infer three-dimensional protein shapes and complexes from sequences and other biological data.

See all 40 papers

Equivariant Neural Belief Propagation
Jun 4, 2026Zehua Cheng, Wei Dai, Jiahao SunEquivariant Representation Learning Factor Graph Optimization
Department of Computer Science, University of Oxford, Oxford, United Kingdom · FLock.io, London, United Kingdom
PETIMOT: A Novel Framework for Inferring Protein Motions from Sparse Data Using SE(3)-Equivariant Graph Neural Networks
Mar 19, 2025Valentin Lombard, Julien Nguyen Van, Sergei Grudinin +1
aDepartment of Computational, Quantitative, and Synthetic Biology (CQSB), UMR 7238 IBPS, Sorbonne Universit´e, CNRS Paris, 75005, France · cUniv. Grenoble Alpes, CNRS, Grenoble INP, LJK 38000 Grenoble, France. · bInstitut universitaire de France (IUF)
Predicting Transmembrane Protein Topology from 3D Structure
Sep 24, 2026Sitong Chen, Xiaopeng Mao
Technical University of Denmark

Jun 4, 2026J/K move · Enter open · S save

Equivariant Neural Belief Propagation

Zehua Cheng, Wei Dai, Jiahao Sun

Department of Computer Science, University of Oxford, Oxford, United Kingdom · FLock.io, London, United Kingdom

Equivariant Representation Learning Factor Graph Optimization Gaussian Mixture Models Bayesian Inference Equivariant Neural Networks Probabilistic Graphical Models+2

Open paper AskPDF arXiv

Probabilistic inference over spatially embedded variables requires beliefs that respect

SE(3)

symmetry, yet existing equivariant networks produce only scalars and vectors -- not the rank-2 precision tensors needed for anisotropic uncertainty, and single-component messages collapse multi-modal energy landscapes to physically meaningless averages. We introduce Equivariant Neural Belief Propagation (ENBP), a factor-graph framework whose messages are equivariant Gaussian mixture models with sufficient statistics that transform exactly under

SE(3)

. Rank-2 precision matrices are synthesised via equivariant outer products, ingested through differentiable spectral decomposition, and kept tractable by a greedy KL-based mixture reduction that provably commutes with

SE(3)

. On GEOM-QM9 and GEOM-Drugs, ENBP achieves 98.9% conformational coverage at 0.090

\mathring{A}

error with sub-second latency -- over

100\times

faster than diffusion baselines at higher accuracy. On multi-body robotic inference, vanilla loopy BP diverges at 15+ agents while ENBP converges with near-zero collision rates and machine-precision equivariance error (

{\sim}10^{-7}

vs.\

10^{-1}

for augmented baselines).

Abstract

Protein Structure Prediction

Explore similar work

Equivariant Neural Belief Propagation

PETIMOT: A Novel Framework for Inferring Protein Motions from Sparse Data Using SE(3)-Equivariant Graph Neural Networks

Predicting Transmembrane Protein Topology from 3D Structure