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

Jun 4, 2026

Equivariant Neural Belief Propagation

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).

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

Ask PDF

Mar 19, 2025

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

Proteins move and deform to ensure their biological functions. Despite significant progress in protein structure prediction, approximating conformational ensembles at physiological conditions remains a fundamental open problem. This paper presents a novel perspective on the problem by directly targeting continuous compact representations of protein motions inferred from sparse experimental observations. We develop a task-specific loss function enforcing data symmetries, including scaling and permutation operations. Our method PETIMOT (Protein sEquence and sTructure-based Inference of MOTions) leverages transfer learning from pre-trained protein language models through an SE(3)-equivariant graph neural network. When trained and evaluated on the Protein Data Bank, PETIMOT shows superior performance in time and accuracy, capturing protein dynamics, particularly large/slow conformational changes, compared to state-of-the-art diffusion and flow-matching approaches, as well as traditional physics-based models. Our code and protocols are available at https://github.com/PhyloSofS-Team/PETIMOT.

Valentin 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)

Topic pending

Ask PDF

Sep 24, 2026

Predicting Transmembrane Protein Topology from 3D Structure

This paper presents a novel approach to infer protein topology using the state-of-the-art graph neural network (GNN), SchNet. The model is trained on the same dataset used to develop the recent DeepTMHMM model with 5-fold cross-validation. Unlike the conventional approaches based on using only the protein sequences or the

α

-carbons as features, we have decoded our classifier in this way, so all atom-level embeddings are used. Without applying any pre-trained weight, the final results have shown great potential that GNNs can be used for topological predictions.

Sitong Chen, Xiaopeng Mao

Technical University of Denmark

Topic pending

Ask PDF

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