cs.LGSep 26, 2026

HamiFormer: Dual-Expert Diffusion Fields with Affine Symplectic Maps

Authors: Haoxiang Huang, Xiang Liu, Shuwei Wang, Jingheng Ma, Sen Cui

Organizations: University of Science and Technology of China · Tsinghua University · Chinese Academy of Sciences

Read PDF
Ask AI
Cite

Copy citation

DOINot available for this paper

Abstract

Predicting smooth dynamics and collisions requires modeling continuous evolution and abrupt state changes. We introduce HamiFormer, a dual-expert diffusion field combining whole-window denoising with residual-corrected Hamiltonian propagation. Their mixed-state feedback attenuates the direct contribution of inherited autoregressive error: each mixed state guides subsequent propagation within the jointly refined window. Parallel Local Affine Scan (PLAS) amortizes iterative refinement across rectified-flow steps and evaluates derivatives in parallel across physical time. PLAS's affine symplectic maps achieve lower solver error and runtime than sequential explicit Euler in our evaluation. A Regime Model Tree specializes residuals and routing to balance typical-state accuracy against large tail errors. Our analysis gives conditions for physically consistent refinement and warm-start tracking, and finite-window error bounds under diffusion feedback. In 192-step evaluations, HamiFormer reduces normalized phase-space MSE by 26.3% against PhysiFormer on HamiBalls-1 and 21.4% against DiT on HamiBalls-2, with comparable model capacities. Disjoint-interval comparisons show the lowest late-horizon position and momentum errors among baselines on both datasets. Project page: https://hamiformer.github.io/.

Figures & tables

Figure 1: HamiFormer overview: schematic hybrid dynamics (top), with ground-truth and predicted HamiBalls-2 states at the same step and normalized phase-space MSE over three disjoint intervals spanning 192 steps (bottom).
Figure 2: Overview of HamiFormer. At each RF field evaluation, a whole-window diffusion expert and a residual-corrected Hamiltonian expert produce candidates for componentwise mixing. PLAS constructs affine maps in parallel from cached anchors and scans the mixed history. The Regime Model Tree conditions the residual and router readouts.
MethodTotalContinuousContact
zzqqppqqppqqpp
HamiBalls-1
PhysiFormer0.340290.346890.333690.347100.328540.333780.65186
HG-DPF1.179661.298541.060771.296731.058971.410081.17228
HamiFormer (Ours)0.250710.255960.245460.256130.239850.245290.59206
HamiBalls-2
Table 1: Pooled 192-edge normalized MSE ( ↓\downarrow ), with ground truth supplied only at initialization. Green, yellow, and red mark the best, second-best, and third-best results per dataset and metric.
Figure 3: Two HamiBalls-2 trajectories: ground truth, PhysiFormer, and Ours at selected steps. Higher gloss denotes higher restitution; red circles mark large object-wise prediction errors.
Method1–48 edges49–96 edges97–192 edges
qqppqqppqqpp
HamiBalls-1
PhysiFormer0.020760.108870.234680.299400.566060.46324
HG-DPF0.230830.703191.295461.046991.833931.24646
HamiFormer (Ours)0.014510.077190.160780.214960.424270.34484
HamiBalls-2
Table 2: Disjoint-interval normalized MSE ( ↓\downarrow ) without resets; colors follow Table 1 .
Figure 4: HamiBalls-2 phase-space MSE averaged over valid objects and realizations at each edge. Axes use separate linear ranges; the dashed line marks the diffusion window boundary at edge 144. Transformer-AR is competitive initially but exhibits stronger error growth later. Ours maintains lower error across most later edges, consistent with the interval averages over the evaluated horizon.
Figure 5: HamiBalls-1 integrator ablation. Left/center: cumulative-mean q/pq/p RMSE through each RF field. Right: pooled phase-space RMSE versus synchronized B64 ×\times 48 sampling time; the dashed line connects nondominated configurations. See Table 3 and Appendix E.1 for values and protocol.
SolverRMSE ( ×10−6\times 10^{-6} ) ↓\downarrowTime (s) ↓\downarrow
qqppzz 
Explicit Euler685.2617.34700484.5800.8875
SymEuler20.049550.056610.053191.2752
PLAS10.60938.920389.801270.7940
PLAS-S12.13609.7428411.00470.7427
PLAS-DC0.050310.056940.053731.0057
Table 3: Solver RMSE against FP64 Newton on HamiBalls-1 (512 sources, two noises). Time: median of 11 warmed, synchronized B64 ×\times 48 runs with a shared comparison interface (Appendix E.1 ).
MethodHamiBalls-1HamiBalls-2
qqppzzqqppzz
H+r candidate only0.0226010.0985540.0605770.3916731.0695610.730617
Output-only mixing0.0208020.0840920.0524470.0793740.4267520.253063
Ours0.0144720.0772400.0458560.0433070.3782030.210755
Table 4: Expert feedback ablation: normalized MSE ( ↓\downarrow ), ranked by color within each column.
IDStructureObjectiveError reduction (%) ↑\uparrowCell win (%) ↑\uparrow
zzqqpp 
ASharedBasic−55.96-55.96−61.89-61.89−54.82-54.8218.4718.47
BTreeBasic11.6911.6912.0212.0211.6311.6354.1454.14
CSharedFull8.428.428.838.838.348.3461.6061.60
DTreeFull29.27\mathbf{29.27}30.19\mathbf{30.19}29.09\mathbf{29.09}64.41\mathbf{64.41}
GLGlobal leafFull7.037.038.658.656.726.7246.6346.63
Table 5: HamiBalls-1, 48 edges (%, ↑\uparrow ). Error reductions are relative to PhysiFormer. Cell wins compare noise-mean phase-space MSE per source–edge–object against the same baseline.
Appendix figures & tables20 assets

Supplementary material from the paper’s appendix.

Appendix

RadiusPositive cost (%)95% interval (%)Maximum MSE increase
10−510^{-5}0.0416[0.0221, 0.0684]0.00770
10−410^{-4}0.2098[0.1528, 0.2786]0.02501
10−310^{-3}0.9998[0.8310, 1.1957]0.04176
Appendix
Table 6: Matched-target 48-edge collision perturbations on 510 sources. Positive cost is the mean casewise positive MSE increase divided by baseline MSE; intervals resample sources with paired signs and noises retained.
Theoretical ingredientProof or direct checkComplementary empirical evidence
PLAS-map symplecticity and implicit regularitySymplecticity proof; measured B\mathbf{B} spectra and affine defects6,881,280 deployed edge maps; learned-H local approximation order
Regular local refinementFixed-point and quadratic-remainder proof; nonlinear mixed controlMulti-start refinement and independently checked point Jacobians
Smooth-model and correction budgetsConditional consistency proof; physical-step refinement controlFirst-order convergence with exact H, vanishing model error, and quadratic corrections
Single-refinement deploymentK=1K=1 tracking bound; finite-schedule output decompositionN=8,12,20N=8,12,20 checks; 512-source, 192-edge refinement comparisons
Finite-window feedback and memoryNonnegative response expansion with a conditional D budgetPhysical-coordinate and internal-memory responses; complete RF interventions
Regular collision sequencesFinite-event composition bound510 contact onsets; matched perturbed targets and event-boundary tails
Appendix
Table 7: Correspondence between theoretical ingredients and supporting evidence. Analytic controls test smooth-system hypotheses; deployed-model diagnostics assess finite-step behavior and complete-sampling effects.
SettingBudget scalingObserved order
Exact HamiltonianϵH=ηh=0\epsilon_{H}=\eta_{h}=00.9956–1.0060
Vanishing gradient errorϵH=O(h)\epsilon_{H}=O(h)0.9973–1.0065
Quadratic step correctionηh=O(h)\eta_{h}=O(h)0.9953–1.0066
Appendix
Table 8: Physical-step convergence in the analytic Hamiltonian control. Slopes fit the maximum trajectory error against hh over five resolutions; ranges span all three initial states.
NNSourcesConditionsWarm startAll starts
88112108107
128112111111
20512716870307015
Appendix
Table 9: Fixed-condition checks within 64 refinements. Counts include both noises. All starts requires all three initializations to meet the numerical criterion.
NNMedian qqMaximum qqMedian ppMaximum pp
88.55×10−88.55\times 10^{-8}1.10×10−71.10\times 10^{-7}8.51×10−88.51\times 10^{-8}1.23×10−71.23\times 10^{-7}
128.45×10−88.45\times 10^{-8}1.21×10−51.21\times 10^{-5}8.54×10−88.54\times 10^{-8}4.56×10−64.56\times 10^{-6}
208.41×10−88.41\times 10^{-8}3.06×10−73.06\times 10^{-7}8.02×10−88.02\times 10^{-8}2.09×10−72.09\times 10^{-7}
Appendix
Table 10: Normalized RMS implicit-equation residuals after the second refinement. This separate panel uses eight sources for N=8,12N=8,12 and 32 for N=20N=20 , with two noises and seven fields.
CenterRadiusMedian gainMaximum gainGain ≥1\geq 1
Warm10−210^{-2}0.0002210.22040/84
Warm10−310^{-3}0.0022280.005180/84
Warm10−410^{-4}0.0242922.04012/84
Terminal10−210^{-2}0.0002250.22040/84
Terminal10−310^{-3}0.0021860.004460/84
Terminal10−410^{-4}0.0229322.04892/84
Appendix
Table 11: Paired anchor-distance gain after eight refinements, with 84 source–field–direction cases per row.
SourcesCountK=1K=1 MSEK=16K=16 MSEPhysiFormerPositive cost (%)
All5120.2507320.2506350.3402890.1903
Flagged970.3225100.3215080.4127140.4744
Other4150.2339550.2340690.3233610.0990
Appendix
Table 12: 192-edge complete predictions with two paired noises. Flagged sources have at least one unmet local criterion at 64 or extended refinement, cross-start disagreement, or exact tree threshold. Positive cost is 100 E[(ℓ1−ℓ16)+]/E[ℓ16]100\,\mathbb{E}[(\ell_{1}-\ell_{16})_{+}]/\mathbb{E}[\ell_{16}] , taking the positive part before averaging source–noise cases.
Routing policySup gain >10>10 (%)Maximum sup gainResponse energy
Live0.59014601.55×10−61.55\times 10^{-6}
Fixed live gates0.28626042.61×10−72.61\times 10^{-7}
All H+rH+r2.35284483.26×10−53.26\times 10^{-5}
Appendix
Table 13: Finite-field reset-response tails. Sup gain divides maximum coordinate response by the injected coordinate amplitude. Response energy is the mean squared subsequent normalized-state response.
Perturbed blockMaximum RMSAggregate MSE change (%)Maximum MSE increase
Recurrent hidden state0.000142+0.0000121.64×10−71.64\times 10^{-7}
Previous gate0.000143+0.0000173.30×10−73.30\times 10^{-7}
Previous H candidate0.000143+0.0000822.67×10−62.67\times 10^{-6}
Contact-provenance context0.000142+0.0000241.72×10−71.72\times 10^{-7}
Cumulative reset0.023837+0.0036411.26×10−41.26\times 10^{-4}
Momentum reset energy0.023837+0.0034451.26×10−41.26\times 10^{-4}
Appendix
Table 14: Internal-memory perturbations followed by complete RF sampling. Each block has 96 paired cases. Response RMS compares the resulting 48-edge predictions with the unperturbed predictions. MSE changes use a common unperturbed aggregate MSE of 0.0371873; maxima retain individual tails.
ComponentArchitecture or stateCount
Diffusion expertWidth 128; 4 blocks; 4 heads; 2 registers; feed-forward width 2561,159,172
Hamiltonian modelWidth 16; 2 object-interaction blocks; 4 heads; scalar readout4,464
Shared residualSiLU MLP, 154→32→32→4154\to 32\to 32\to 46,148
Router12-dimensional projection and physical-time GRU; width-16 q/pq/p offset heads; leaf-affine corrections7,588
Neural subtotalAll four components1,177,372
Ridge readoutsFitted linear readouts1,056
Appendix
Table 15: HamiBalls-1 architecture and learned-size accounting. Counts for the router include its leaf-affine terms. Linear coefficients and other learned state are reported separately from neural parameters.
ComponentArchitecture or stateCount
Diffusion expertWidth 168; 12 blocks; 8 heads; 2 registers; feed-forward width 3365,615,598
Hamiltonian modelWidth 32; 2 object-interaction blocks; 4 heads; pair-energy readout22,252
Shared residual and routerToken-conditioned SiLU residual; physical-time GRU with q/pq/p offset heads and leaf-affine corrections32,729
Neural subtotalAll HamiFormer neural components5,670,579
Ridge readoutsFitted linear readouts6,960
Combined scalarsNeural parameters plus fitted linear coefficients5,677,539
Appendix
Table 16: HamiBalls-2 architecture and learned-size accounting. Totals include neural parameters and fitted ridge coefficients; tree structure, thresholds, and normalization statistics are additional state.
SettingH1H2
Base updates; batch size50,000; 6450,000; 64
Optimizer; weight decay; grad clipAdamW; 10−410^{-4} ; 1.0AdamW; 10−410^{-4} ; 1.0
Peak learning rate: D / H3×10−43\times 10^{-4} / 10−410^{-4}3×10−43\times 10^{-4} / 10−410^{-4}
Peak LR: residual / scalar / final q/p3×10−43\times 10^{-4} / 3×10−43\times 10^{-4} / 3×10−33\times 10^{-3}10−310^{-3} / 3×10−33\times 10^{-3} / 3×10−33\times 10^{-3}
RF denominator floor0.050.05
Downstream stages4×2004\times 200 updates4×2004\times 200 updates
Appendix
Table 17: Training settings for HamiBalls-1 (H1) and HamiBalls-2 (H2).
Trainable componentUpdatesObjective and carrier policy
Diffusion ++ Hamiltonian base50,000LD+LH\mathcal{L}_{D}+\mathcal{L}_{H} ; whole-window RF (batch 64) and local relation minibatches
Shared residual200Eq. ( D.8 ); random-gate carriers; refresh every 8 updates
First scalar router200Global sequence projection regret; fixed mixed- NN pool
Second scalar router200Projection regret plus non-degradation penalty; fixed pool from the preceding scalar router
Model tree—Fit once after scalar stages outside gradient-based optimization
Dual-channel/leaf router200Eq. ( D.16 ); 25 refreshes of 8 updates; cumulative ridge fitting
Appendix
Table 18: Training sequence. The diffusion and Hamiltonian losses form one 50,000-update base stage. A refresh regenerates rollout carriers. Fixed previous-stage policies generate the scalar-stage pools.
Integrator48-edge zz -MSE192-edge zz -MSEWall-clock (s)
Explicit Euler0.0459100.2539740.8875
SymEuler20.0458460.2506471.2752
PLAS0.0458500.2507070.7940
Appendix
Table 19: HamiBalls-1 prediction accuracy and sampling time ( ↓\downarrow ). MSE uses 512 sources and two noises; wall-clock is per B64 ×\times 48 sampling call under the shared timing protocol of Table 3 .
Implementationzz MSEqq MSEpp MSEWin (%)
Ours0.045849450.014505590.0771933264.4263
Ours (PLAS-S)0.045841860.014500820.0771828964.4173
Appendix
Table 20: Prediction accuracy with full and low-frequency Hessian refresh on HamiBalls-1 over 48 edges; paired timing is described in Appendix E.4 . Win rates use noise-mean pointwise errors against PhysiFormer.
MethodRF steps NNSeconds / batch192-edge MSE
PhysiFormer200.4020.34029
PhysiFormer400.8090.32767
Ours (PLAS)200.7860.25071
Appendix
Table 21: Inference cost and prediction accuracy on HamiBalls-1, evaluated on 512 sources and two noise realizations. Latency is the median seconds per batch of 64 48-edge windows; normalized phase-space MSE covers all 192 predicted edges.
Figure 6: Object-count scaling on an A800. Left: median inference time after three warm-up calls, using five timed calls. Right: peak allocated GPU memory after warm-up, measured in separate processes for each method and size. Both vertical axes are logarithmic. Batch size is 8, window length is 48, and RF uses 20 steps.
SettingHamiBalls-1HamiBalls-2
Domain[−1,1]2[-1,1]^{2}[−1,1]2×[0,2][-1,1]^{2}\times[0,2]
Objects5 disks5–10 spheres
Mass; radiusU[0.5,1.5]\mathcal{U}[0.5,1.5] ; U[0.06,0.10]\mathcal{U}[0.06,0.10]Same ranges
Initial speedU[0.35,0.90]\mathcal{U}[0.35,0.90]U[0.35,0.90]\mathcal{U}[0.35,0.90]
Body restitutionU[0.4,0.9]\mathcal{U}[0.4,0.9]U[0.90,0.98]\mathcal{U}[0.90,0.98]
Boundary restitution0.80.81.01.0
Appendix
Table 22: Dataset generation settings. U[a,b]\mathcal{U}[a,b] denotes uniform sampling; speed ranges apply before the HamiBalls-1 center-of-mass adjustment. Restitution entries are per-body material coefficients.
 PhysiFormerOurs
Physical seedzzqqppzzqqpp
400.067180.022630.111730.045570.015200.07594
410.065820.020460.111180.051460.016860.08606
420.067180.024200.110160.044310.015500.07311
430.059070.015740.102400.042060.010450.07366
Appendix
Table 23: HamiBalls-1 absolute 48-edge normalized MSE by physical data-generation seed.
 PhysiFormerOurs
Physical edgeszzqqppzzqqpp
1–480.064810.020760.108870.045850.014510.07719
49–960.267040.234680.299400.187870.160780.21496
97–1440.468410.499810.437020.339720.363460.31598
145–1920.560890.632300.489470.429380.485070.37370
Appendix
Table 24: HamiBalls-1 absolute normalized MSE for disjoint 48-edge continuation chunks.
More in this topic

Physics-Informed Diffusion Models

Physics-informed diffusion models use physical laws and measurements to guide generative sampling, reconstruction, and forecasting.

See all 50 papers

Explore similar work

CardsList

  1. Learning Hamiltonian Flow Maps: Mean Flow Consistency for Large-Timestep Molecular Dynamics

    Jan 29, 2026Winfried Ripken, Michael Plainer, Gregor Lied +5Dynamical SystemsFlow Map Learning

    Technical University Berlin · BIFOLD · Free University of Berlin +6

  2. PIDM-DP: Physics-Informed Diffusion with Dormand-Prince Integration for Chaotic System Identification and State Reconstruction across Multiple Dynamical Regimes

    May 26, 2026Shailendra DabralPhysics-Informed Diffusion ModelsNonlinear System Identification

    Indian Institute of Technology Indore

  3. CoSynFlow: Conformal Symplectic Neural Flows for Cross-System Prediction of Dissipative Hamiltonian Dynamics

    Aug 1, 2026Baige Xu, Takaharu YaguchiNeural Operators

    RIKEN AIP · Motooka Nishi-ku, Fukuoka, Japan · IMI, Kyushu University