HamFormer: A Hamiltonian Flow-based SE(3)-Transformer for Conformation Optimization

Minjie Cheng; Zhenghao Liu; Yitian Wang; Fanmeng Wang; Anqi Dong; Hongteng Xu

cs.LGOct 31, 2026

HamFormer: A Hamiltonian Flow-based SE(3)-Transformer for Conformation Optimization

Authors: Minjie Cheng, Zhenghao Liu, Yitian Wang, Fanmeng Wang, Anqi Dong, Hongteng Xu

Organizations: Gaoling School of Artificial Intelligence, Renmin University of China · Laboratory for Human-Oriented Mobility Eco-system, École Polytechnique Fédérale de Lausanne · Beijing Key Laboratory of Research on Large Models and Intelligent Governance · Engineering Research Center of Next-Generation Intelligent Search and Recommendation, MOE

Abstract

The Transformer architecture plays a central role in modern deep learning, and its SE(3)-equivariant variants have been widely applied to molecular conformation modeling and optimization. Given its importance, many methods attempt to interpret and improve the Transformer through the lens of dynamics. However, they mainly focus on a single attention head rather than the full Transformer architecture and seldom consider SE(3)-equivariance in their proposed models. In this study, we propose an improved SE(3)-Transformer architecture, called HamFormer, which revisits the design of SE(3)-Transformer through the lens of Hamiltonian flow. In particular, we define a parametrized Hamiltonian flow in the phase space of position and velocity, corresponding to optimizing an energy functional that combines an optimal transport-based potential with a damped momentum. The Lie-Trotter splitting framework of the Hamiltonian flow leads to the proposed HamFormer architecture, whose feedforward computation achieves the optimization step of the proposed energy functional. Furthermore, we combine HamFormer with Equivariant Graph Neural Network (EGNN), build a new model for molecular conformation optimization. Extensive experiments on molecular ground-state conformation prediction show that HamFormer provides a strong and physically interpretable surrogate for classic SE(3)-Transformer and generally outperforms state-of-the-art competitors.

Figures & tables

Figure 1: An illustration of the proposed molecule model based on HamFormer.

Figure 2: HamFormer architecture.

Component	Form	Physics	State	Molecular Role
OT Attention	$-\nabla_{\bm{X}}V_{\theta}$	Force	$\bm{P}$	Atom Interactions
Kinetic Map	$\nabla_{\bm{P}}K_{\theta}$	Velocity	$\bm{X}$	Feature Motion
Retention	$-\tanh(\alpha_{\ell})\bm{P}$	Signed Carry	$\bm{P}$	Refinement Memory
Relations	$\bm{T}_{h}^{(\ell)}\mapsto\bm{T}_{h}^{(\ell+1)}$	Coupling	$\bm{\mathsf{T}}$	Pair Information
EGNN	$\bm{c}_{i}-\bm{c}_{j}$	Direction	$\bm{C}$	Coordinate Refinement

Table 1: Physical roles of HamFormer components.

Model	Reference dynamics	Interaction	Momentum parametrization	Pair memory	SE(3)
Sinkformers ( 2022 )	Wasserstein PDE: $\partial_{t}\rho=\nabla\cdot(\rho\nabla\mathcal{F}^{\prime}(\rho))$	OT	–	$\times$	$\times$
WGFormer ( 2025 )	Wasserstein PDE: $\partial_{t}\rho=\nabla\cdot(\rho\nabla\mathcal{F}^{\prime}(\rho))$	OT	–	$\checkmark$	$\checkmark$
YuriiFormer ( 2026 )	Nesterov: $\Delta\bm{X}^{k+1}=\beta_{k}\Delta\bm{X}^{k}-\eta_{k}\nabla E(\bm{X}^{k}+\alpha_{k}\Delta\bm{X}^{k})$	Softmax	Learned coefficients	$\times$	$\times$
SympFormer ( 2026 )	Hamiltonian PDE: $\partial_{t}\rho=\mathcal{H}^{\prime}_{S},\hskip 9.24994pt\partial_{t}S=-\mathcal{H}^{\prime}_{\rho}-\gamma(t)S$	Softmax	Metric, damping	$\times$	$\times$
HamFormer (ours)	Hamiltonian PDE: $\partial_{t}\mu+\{\mu,h_{\mu}\}=0$	OT	Kinetic map, signed retention	$\checkmark$	$\checkmark$

Table 2: Reference dynamics under the respective theoretical assumptions and features of the reported architectures. OT denotes entropic optimal transport, and pair memory denotes relations retained across layers. The SympFormer row reports its softmax variant.

Dataset	Model	Validation			Test
Dataset	Model	D-MAE $\downarrow$	D-RMSE $\downarrow$	C-RMSD $\downarrow$	D-MAE $\downarrow$	D-RMSE $\downarrow$	C-RMSD $\downarrow$
Molecule3D (Random)	GINE(2D)	0.590	1.014	1.116	0.592	1.018	1.116
	GATv2(2D)	0.563	0.983	1.082	0.564	0.986	1.083
	GPS(2D)	0.528	0.909	1.036	0.529	0.911	1.038
	GTMGC(2D)	0.432	0.719	0.712	0.433	0.721	0.713
	SE(3)-Transformer	0.466	0.712	0.800	0.467	0.774	0.802

Table 3: Performance comparison on Molecule3D and QM9 datasets.

\downarrow

indicates that lower values are better. Bold indicates the best performance, and underline indicates the second best. Baseline results follow Wang et al. (2025) .

Figure 3: Layer-wise evolution of the latent kinetic energy, potential energy, and total Hamiltonian evaluated on 1,000 random molecules from the Molecule3D test set. Solid curves denote the mean, and shaded regions represent one standard deviation across molecules. The dashed line marks the transition between Stage 1 (layers H1–H3) and Stage 2 (layers H4–H6).

Panel of Figure 4 — Figure 4: Qualitative conformation visualizations and quantitative ablation studies on Molecule3D. Left: Quantitative ablation on the test set comparing the full HamFormer with zeroed inference momentum ( $\rho_{l}=0$ ) and the EGNN baseline. Right: Visual comparison of predicted 3D molecular geometries against DFT references for RDKit, WGFormer, and HamFormer.

Appendix figures & tables6 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 8

Setting	QM9	Molecule3D random
Peak learning rate	$8\times 10^{-4}$	$4\times 10^{-4}$
Scheduled epochs	500	40
Warmup fraction	0.06	0.05
Evaluated checkpoint epoch	294	40
Training GPUs	8	2
Batch size per GPU	32	64

Appendix

Table 4: Full-model training settings.

	Feature	$P\&T$	QM9			Molecule3D random
Variant	block	branch	D-MAE	D-RMSE	C-RMSD	D-MAE	D-RMSE	C-RMSD
Full model	HamFormer	$\checkmark$	0.218	0.419	0.198	0.387	0.656	0.654
EGNN-only	$\times$	$\times$	0.233	0.427	0.239	0.399	0.666	0.697

Appendix

Table 5: Module ablation on the test sets. “Feature block” denotes the module applied before each EGNN stage.

	Validation			Test
Inference setting	D-MAE	D-RMSE	C-RMSD	D-MAE	D-RMSE	C-RMSD
Learned $\rho_{l}$	0.386	0.653	0.652	0.387	0.656	0.654
$\rho_{l}=0$	0.418	0.684	0.712	0.420	0.716	0.714

Appendix

Table 6: Inference-time momentum-retention ablation on Molecule3D random.

Figure 5: Layer-wise evolution of the latent kinetic energy, potential energy, and total Hamiltonian evaluated on 1,000 randomly selected molecules from the QM9 test set.

Figure 6: The visualizations for the ground-state conformations predicted by RDKit, WGFormer, and HamFormer.

Transformer

Transformers are neural architectures that use attention to model relationships among sequence elements in language, vision, and scientific data.

See all 551 papers

EquiformerV3: Scaling Efficient, Expressive, and General SE(3)-Equivariant Graph Attention Transformers
Apr 10, 2026Yi-Lun Liao, Alexander J. Hoffman, Sabrina C. Shen +3Equivariant GNNs Graph Transformer
Massachusetts Institute of Technology (MIT) · Mirror Physics · Entalpic
Geometric Flow Matching for Molecular Conformation Generation via Manifold Decomposition
May 25, 2026Yunqing Liu, Yi Zhou, Wenqi FanFlow Matching Generative Modeling
The Hong Kong Polytechnic University Hong Kong, China
VEDA: Generation of 3D Molecules via Variance-Exploding Diffusion with Annealing
Nov 11, 2025Peining Zhang, Jinbo Bi, Minghu Song3D Molecular Generation Diffusion Models
University of Connecticut, Storrs, Connecticut 06269, USA · Institute of Health and Medicine, Hefei Comprehensive National Science Center, Hefei 230601, China

Apr 10, 2026J/K move · Enter open · S save

EquiformerV3: Scaling Efficient, Expressive, and General SE(3)-Equivariant Graph Attention Transformers

Yi-Lun Liao, Alexander J. Hoffman, Sabrina C. Shen, Alexandre Duval+2

Massachusetts Institute of Technology (MIT) · Mirror Physics · Entalpic

Equivariant GNNs Graph Transformer Machine Learning Interatomic Potentials Computational Materials Science

Open paper AskPDF arXiv

SE(3)

-equivariant graph neural networks mature as a core tool for 3D atomistic modeling, improving their efficiency, expressivity, and physical consistency has become a central challenge for large-scale applications. In this work, we introduce EquiformerV3, the third generation of the

SE(3)

-equivariant graph attention Transformer, designed to advance all three dimensions: efficiency, expressivity, and generality. Building on EquiformerV2, we have the following three key advances. First, we optimize the software implementation, achieving

1.75\times

speedup. Second, we introduce simple and effective modifications to EquiformerV2, including equivariant merged layer normalization, improved feedforward network hyper-parameters, and attention with smooth radius cutoff. Third, we propose SwiGLU-

S^2

activations to incorporate many-body interactions for better theoretical expressivity and to preserve strict equivariance while reducing the complexity of sampling

S^2

grids. Together, SwiGLU-

S^2

activations and smooth-cutoff attention enable accurate modeling of smoothly varying potential energy surfaces (PES), generalizing EquiformerV3 to tasks requiring energy-conserving simulations and higher-order derivatives of PES. With these improvements, EquiformerV3 trained with the auxiliary task of denoising non-equilibrium structures (DeNS) achieves state-of-the-art results on OC20, OMat24, and Matbench Discovery.

Abstract

Transformer

Explore similar work

EquiformerV3: Scaling Efficient, Expressive, and General SE(3)-Equivariant Graph Attention Transformers

Geometric Flow Matching for Molecular Conformation Generation via Manifold Decomposition

VEDA: Generation of 3D Molecules via Variance-Exploding Diffusion with Annealing