FLARE: Flow Matching with Local Axis-Angle Representations for Stochastic Micromagnetic Evolution

Pengyu Li; Renjie Tong; Xuanlue Jiang; Jianmin Li; Yuanyuan Zhou

Item	Value	Item	Value
Grid	$256\times 256$	Trainable parameters	65,260,803
Image channels	19	Base width / multipliers	64 / $[1,2,4,8]$
Residual blocks	2 per level	Bottleneck attention	8 heads
Global condition width	512	Group normalization	8 groups
Stage–1 updates	50,000	Global batch	128
Optimizer	AdamW	Learning rate	$2{\times}10^{-4}\rightarrow 2{\times}10^{-6}$

Appendix

Table 3: Compact implementation specification. Batch size is global across two training devices. The reported checkpoint is the final EMA state rather than a test-selected checkpoint.

Quantity	Value or realized support
Film and discretization	$256\times 256\times 1$ cells of $2\times 2\times 1$ nm 3 ; $512\times 512\times 1$ nm square film; open boundaries; edge smoothing 8
Material constants	$M_{s}=5.8\times 10^{5}$ A m -1 , $A_{\rm ex}=1.5\times 10^{-11}$ J m -1 , $K_{u1}=8.0\times 10^{5}$ J m -3 , and $\alpha=0.30$ in every base
DMI content	Interfacial Dzyaloshinskii–Moriya interaction in every base, $D_{\rm ind}=3.25\times 10^{-3}$ J m -2 ; no zero-DMI or bulk-DMI cases
Temperature	$30$ K (334), $150$ K (333), or $300$ K (333), with an independent MuMax 3 thermal seed for each repeat
Applied field	Uniform $B_{z}\in\{0,8,16,24,32\}$ mT (200/200/200/201/199)
SOT current	Slonczewski SOT on a static $8\times 8$ spatial mask; signed nonzero drive amplitude $-4.487\times 10^{12}$ to $4.413\times 10^{12}$ A m -2 , with $\|J\|\in[0.404,4.487]\times 10^{12}$ A m -2 ; $J=0$ after the pulse

Appendix

Table 4: Definition and realized support of the custom x5 MuMax 3 corpus. Counts in parentheses are numbers of the 1,000 base conditions; stochastic thermal repeats inherit their base parameters. Fixed material entries are singletons, not omitted ranges.

Collection	Items	Base clusters	Draws per item	Segment role	Use
Training / validation / test trajectories	6,227 / 786 / 757	split by base	–	both	Optimization and held-out evaluation
Single segment	1,493 cases	100	4	within segment	Single-segment evaluation
Complete-path prediction distribution	33 conditions	33	5	drive $\rightarrow$ predicted handoff $\rightarrow$ relax	Distribution Score and angular energy distance
MuMax 3 repeat reference	34 conditions	17	30	17 drive + 17 relax	Two-repeat angle, five-versus-five distances, and driven exact-anchor score
Exact-anchor complete paths (Table 1 )	165 anchors	33	5 stochastic / 1 deterministic	drive $\rightarrow$ predicted handoff $\rightarrow$ relax	Fair energy score
Multi-stage rollouts	736 rollouts	97	fixed	two or more	Teacher-forced and autoregressive evaluation

Appendix

Table 5: Data and statistical units used by the frozen evaluation. Pixels, draws, repeats, and segments are never treated as independent confidence-interval samples.

Figure 4: Exact-control complete-path distribution diagnostics. (a) Distribution Score at the final post-relaxation endpoint after prediction-only handoff. (b) Repeat-paired angle versus angular energy distance; lower-left is better and the star marks FLARE. The two axes need not rank a stochastic model identically. (c) Paired Distribution-Score ratios for FLARE relative to each 50,000 learned-dynamics baseline with a reported Distribution Score; ratios above one favor FLARE and bars are 95% base-cluster intervals. The plotted comparison uses five outcomes for every method and 33 complete-path conditions from Table 1 .

Figure 5: Conditional outcome variability under identical physical inputs. Each row fixes the exact magnetization anchor, physical controls, and target horizon, and compares 30 independent MuMax 3 thermal outcomes, 30 fresh Stage 1 FLARE ODE–10 draws, and one deterministic Poseidon-T prediction. Six representative fields from each stochastic ensemble are selected using the same deterministic medoid-based rule, while all 30 outcomes are included in the observable scatter plot. (a) Texture and topological-charge variability. (b) Domain-pattern variability. (c) Topology-count variability. (d) An additional case with texture and topological-charge variability. Across these cases, FLARE generates diverse future states that reproduce substantial portions of the MuMax 3 outcome spread in both field space and observable space, whereas the deterministic baseline provides only a single prediction for each fixed input.

Setting	FLARE	Best comparator	FLARE/comparator [95% CI]
Bandwidth $0.5\times$	0.3274	PDEArena 0.2330	1.405 [1.256, 1.590]
Bandwidth $1\times$	0.8050	PDEArena 0.7362	1.094 [1.059, 1.132]
Bandwidth $2\times$	0.9543	PDEArena 0.9287	1.028 [1.016, 1.040]
No translation	0.8318	PDEArena 0.7769	1.071 [1.039, 1.107]

Appendix

Table 6: Single-segment Distribution Score sensitivity. The bandwidth multiplier rescales only the MuMax 3 nearest-neighbor bandwidth; “no translation” removes local translation matching while retaining the nominal bandwidth rule. The last column is a paired base-group bootstrap ratio; these values are not the primary complete-path results in Table 1 .

Metric	Combined [95% CI]	Driven [95% CI]	Post-relaxation [95% CI]	MuMax 3 repeat protocol
Distribution Score $\uparrow$	0.901 [0.866, 0.930]	0.941 [0.927, 0.953]	0.862 [0.801, 0.915]	Disjoint five versus five
Two-repeat angle ( ∘ ) $\downarrow$	28.86 [22.84, 34.82]	27.47 [21.51, 33.37]	30.25 [23.89, 36.66]	All distinct trajectory pairs
Angular energy distance ( ∘ ) $\downarrow$	11.60 [9.17, 14.07]	10.99 [8.62, 13.35]	12.20 [9.62, 14.86]	Disjoint five versus five
Fair energy score $\downarrow$	–	0.283 [0.230, 0.335]	–	Same exact anchor; disjoint forecasts/targets

Appendix

Table 7: Independent MuMax 3 –MuMax 3 reference levels for absolute interpretation of the distribution metrics. The audit uses 30 thermal repeats for each of 17 held-out base conditions; intervals resample the 17 bases and retain both segment roles. Distinct-fold rows average all ordered pairs of disjoint five-repeat folds. These are simulator-only finite-sample references, not algebraic zeros or paired estimates on the 33 main bases.

Figure 6: Metric, sample-budget, and translation-matching stress tests. (a) The frozen FLARE checkpoint on the independent x30 audit, using five MuMax 3 targets and

N\in\{1,2,5,10,20\}

model outcomes. (b) Disjoint MuMax 3 sets at 5-versus-5, 10-versus-10, and 15-versus-15. Fair energy score is multiplied by 30 only for display on the angular-distance axis. (c) Probability that FLARE ranks first in the frozen eight-method single-segment subset under base-cluster resampling. (d) The complete

3\times 4

patch-grid/window response for three analytic sphere-valued perturbations. (e) The nominal

4\times 4

grid as a function of search radius. (f) A topology-changing counterexample: local matching can reduce the field discrepancy but cannot certify

Q

. Bars and bands are 95% base-cluster intervals.

Source comparison	Budget	Distribution Score $\uparrow$	Angular energy dist. ( ∘ ) $\downarrow$	Fair energy score $\downarrow$
FLARE–MuMax 3	$1$ vs. $5$	0.620 [0.474, 0.781]	36.03 [29.85, 41.91]	–
FLARE–MuMax 3	$2$ vs. $5$	0.658 [0.516, 0.816]	22.06 [18.37, 25.74]	0.3105 [0.2592, 0.3594]
FLARE–MuMax 3	$5$ vs. $5$	0.691 [0.553, 0.836]	13.53 [11.42, 15.60]	0.3095 [0.2615, 0.3565]
FLARE–MuMax 3	$10$ vs. $5$	0.694 [0.556, 0.841]	10.70 [9.08, 12.37]	0.3094 [0.2631, 0.3577]
FLARE–MuMax 3	$20$ vs. $5$	0.694 [0.553, 0.845]	9.31 [7.87, 10.75]	0.3092 [0.2620, 0.3562]
MuMax 3 –MuMax 3	$5$ vs. $5$	0.908 [0.886, 0.930]	11.53 [9.15, 13.95]	0.2831 [0.2306, 0.3357]

Appendix

Table 8: Finite-ensemble pressure test on the independent x30 audit. The FLARE rows use

N

model outcomes against five MuMax 3 outcomes; the reference rows use two disjoint MuMax 3 sets of the indicated size. Fair energy score is reported only for driven branches, where independent thermal outcomes share the exact anchor. At

N=1

the finite-ensemble correction is undefined; the corresponding ordinary one-member energy score is 0.592 [0.498, 0.679].

Perturbation	Pixel angle	Pixel MSE	$\|\Delta Q\|$	$D_{0}$	$D_{8}$	$D_{16}$	$D_{32}$
Skyrmion displacement	$3.71^{\circ}$	0.02485	0.000	0.03727	0.00395	0.00000	0.00000
Topology disappearance	$5.57^{\circ}$	0.04043	0.983	0.06064	0.04790	0.03712	0.03032
Domain-wall shift	$9.84^{\circ}$	0.06456	0.000	0.09683	0.02517	0.00000	0.00000

Appendix

Table 9: Response of the nominal

4\times 4

local patch grid to analytic perturbations.

D_{w}

is the translation-matched distance with a

\pm w

-pixel search. Figure 6 d reports all

2\times 2

,

4\times 4

, and

8\times 8

grids. Pixel MSE averages over both cells and the three magnetization components, and therefore equals

e_{\mathrm{MSE}}/3

in Eq. 13 .

(a) Physical-observable distribution error
Method	$W_{1}(Q)$	$W_{1}(\overline{m}_{z})$	$W_{1}(E_{\rm ex})$
FLARE ODE-10	0.48	0.0283	0.00352
Poseidon-T	1.88	0.0458	0.03598
CNO-FM	1.64	0.0470	0.03703
DPOT-Ti	0.77	0.0796	0.03831
MPP-AViT-Ti	0.83	0.0827	0.03743

Appendix

Table 10: Single-segment physical-observable distribution errors and near-antipodal rotation statistics. Panel (a) reports equal-condition empirical Wasserstein-1 errors [ 70 ] between five model and five MuMax 3 outcomes; lower is better. Panel (b) uses the implementation threshold.

Figure 7: Heun-step sensitivity on 1,493 held-out transitions. (a) Angular error with base-cluster intervals. (b) Absolute topological-charge error. (c) Amortized latency per transition at batch size 32. Increasing from 5 to 20 steps is nearly linear in cost but does not reduce angular error, indicating that the learned vector field, rather than integration truncation, dominates current error. ODE–10 is kept fixed across single-segment, cache, and rollout evaluation.

Figure 8: Complete-path quality and representative compute tradeoffs. (a) Complete-path Distribution Score versus standardized fixed two-segment 5-ns latency at each method’s fastest reported batch; marker area encodes parameter count. (b) Parameter count versus representative 5-ns latency. Parameter count alone does not determine runtime because temporal interfaces range from two direct-horizon predictions to forty velocity evaluations. FLARE remains non-dominated in the quality–throughput plane.

Method / batch $B$	1	2	4	8	16	32	64	128	256	512	1024	2048
FLARE ODE-10	995.2	554.6	353.2	245.8	191.6	163.9	150.3	145.7	–	–	–	–
Poseidon-T	61.91	32.61	16.90	8.451	4.140	2.223	1.268	0.8515	0.7038	0.6689	0.6499	0.7609
CNO-FM	72.89	42.73	21.44	10.76	5.415	3.824	3.569	3.429	3.377	3.354	3.382	3.406
DPOT-Ti	91.13	49.80	28.55	17.65	12.32	9.906	9.231	9.591	9.448	–	–	–
MPP-AViT-Ti	613.4	332.4	194.2	147.1	135.0	123.7	121.4	119.6	–	–	–	–
PDEArena U-Net	40.17	20.53	14.15	11.58	10.58	10.12	9.871	9.758	9.892	10.61	10.94	–

Appendix

Table 11: Standardized fixed two-segment 5-ns (2-ns drive plus 3-ns relaxation, with one predicted-state handoff) single-H200 batch scaling. Entries are mean milliseconds per output; bold marks the lowest measured latency per method. A dash denotes an unsupported or unrecorded setting. MuMax 3 has no ensemble-batch interface and is shown only at its native batch one. Every learned entry uses the reported 50,000 checkpoint.

Method	Work in one fixed 2+3-ns timed output	Train H200-h	Train peak (GiB)	$B=1$ peak (GiB)	Fastest $B$ / peak (GiB)
FLARE ODE-10	20 Heun steps; 40 velocity evaluations	15.20	$96.9^{\ast}$	0.539	128 / 18.54
Poseidon-T	Two lead-time-conditioned endpoint calls	3.71	1.15	0.182	1024 / 65.62
CNO-FM	Two lead-time-conditioned endpoint calls	6.83	4.61	0.493	512 / 27.44
DPOT-Ti	20 recurrent 0.25-ns calls	6.91	2.38	0.350	64 / 18.61
MPP-AViT-Ti	20 recurrent 0.25-ns calls	6.77	12.29	0.887	128 / 105.97
PDEArena U-Net	20 recurrent 0.25-ns calls	0.953	1.58	0.270	128 / 12.55

Appendix

Table 12: Compute and memory accounting for Table 1 ’s 50,000-update checkpoints. Training time is accelerator device-hours, so FLARE’s two-device wall time is multiplied by two. Training peaks for the external endpoint models are PyTorch allocated-memory high-water marks; FLARE’s starred value is the maximum 10-s process-memory sample per device and is therefore a more inclusive measurement. Inference peaks use the fixed two-segment representative 5-ns

B=1

query; the final column shows the memory cost at each method’s fastest measured batch.

Figure 9: Conditioning on the requested physical horizon. Dotted lines mark unseen 1.5- and 2.5-ns interpolation queries. (a) Angular error. (b) magnetization MSE. (c) angle relative to Stage 1 on the same cases; whiskers at the two unseen horizons are paired 95% base-condition bootstrap intervals. Removing the explicit target-time channels consistently worsens the matched prediction, including at both unseen horizons.

Figure 10: Complete supported-horizon diagnostics. Circles at 1, 2, and 3 ns lie on the training duration grid; 1.5 and 2.5 ns are held-out interpolation queries. (a) Repeat-paired angular error and angular energy distance. (b) Distribution Score. (c) Empirical Wasserstein-1 error for

Q

, mean

m_{z}

, and the normalized exchange-texture energy proxy. (d) Ratio of model to MuMax 3 within-ensemble angular spread; the dashed line is one. Bands are 95% base-cluster intervals. The gray region is shown explicitly because it contains no legal exact single-segment target in the held-out saved-frame data.

Horizon (ns)	Support	$N_{\rm base}$	Distribution Score $\uparrow$	Paired angle ( ∘ ) $\downarrow$	Angular energy dist. ( ∘ ) $\downarrow$	Fair energy score $\downarrow$
1.0	seen	33	0.823 [0.783, 0.857]	29.17 [25.28, 33.09]	11.83 [10.27, 13.48]	0.3189 [0.2820, 0.3559]
1.5	interpolation	33	0.793 [0.727, 0.848]	30.86 [26.65, 34.85]	12.48 [10.76, 14.15]	0.3397 [0.3008, 0.3780]
2.0	seen	33	0.777 [0.706, 0.839]	31.38 [27.13, 35.44]	12.74 [11.05, 14.47]	0.3446 [0.3058, 0.3822]
2.5	interpolation	33	0.776 [0.708, 0.838]	32.39 [28.15, 36.84]	13.18 [11.47, 14.95]	0.3594 [0.3207, 0.3986]
3.0	seen	33	0.780 [0.713, 0.843]	32.87 [28.27, 37.21]	13.54 [11.68, 15.46]	0.3622 [0.3225, 0.4020]
3.5	no exact saved target	0	–	–	–	–

Appendix

Table 13: Distribution metrics at every requested integer and half-integer horizon from 1 to 6 ns. Brackets are 95% base-cluster intervals; arrows give the favorable direction. “Seen” means that the duration lies on the training duration grid, not that the held-out physical condition was seen. Unsupported rows are retained rather than extrapolated.

Horizon (ns)	$W_{1}(Q)$	$W_{1}(\overline{m}_{z})$	$W_{1}(E_{\rm ex}\ \mathrm{proxy})$	Model spread ( ∘ )	MuMax 3 spread ( ∘ )	Spread ratio [95% CI]
1.0	0.24 [0.20, 0.29]	0.0215 [0.0171, 0.0261]	0.00294 [0.00261, 0.00329]	49.28	49.64	0.993 [0.982, 1.002]
1.5	0.32 [0.25, 0.40]	0.0266 [0.0205, 0.0335]	0.00356 [0.00315, 0.00402]	48.85	49.28	0.991 [0.979, 1.003]
2.0	0.32 [0.27, 0.38]	0.0261 [0.0204, 0.0331]	0.00386 [0.00336, 0.00444]	48.46	48.90	0.991 [0.978, 1.003]
2.5	0.37 [0.29, 0.44]	0.0312 [0.0238, 0.0406]	0.00424 [0.00374, 0.00480]	48.05	48.30	0.995 [0.980, 1.008]
3.0	0.43 [0.31, 0.55]	0.0316 [0.0234, 0.0428]	0.00439 [0.00391, 0.00494]	47.51	47.58	0.999 [0.984, 1.012]

Appendix

Table 14: Physical-observable distribution errors and ensemble spread over the five supported horizons. Observable columns are empirical five-versus-five Wasserstein-1 errors. Model and MuMax 3 spread are the equal-condition means of within-set angular distances in degrees. Brackets are 95% base-cluster intervals.

Variant	Distribution Score $\uparrow$	Angular energy dist. ( ∘ ) $\downarrow$	Paired angle ( ∘ ) $\downarrow$	Fair energy score $\downarrow$
FLARE (one checkpoint)	0.762 [0.660, 0.844]	15.30 [13.32, 17.26]	36.72 [31.93, 41.38]	0.393 [0.356, 0.430]
Direct U-Net	0.472 [0.188, 0.778]	20.91 [17.89, 24.24]	38.88 [34.61, 43.04]	0.541 [0.499, 0.580]
CNO-FM	0.717 [0.619, 0.799]	20.52 [17.23, 24.15]	35.14 [31.28, 38.94]	0.535 [0.488, 0.582]
DPOT-Ti	0.744 [0.660, 0.818]	25.86 [20.06, 32.43]	40.76 [35.23, 46.34]	0.569 [0.510, 0.628]
Poseidon-T	0.703 [0.610, 0.780]	18.93 [16.06, 22.04]	33.31 [29.12, 37.41]	0.503 [0.456, 0.548]
PDEArena U-Net	0.515 [0.264, 0.760]	24.86 [19.55, 30.91]	38.96 [33.38, 44.56]	0.558 [0.493, 0.623]

Appendix

Table 15: Five-seed ensembles (78–82) on the exact-control complete path, with 95% base-cluster intervals conditional on those members; the frozen FLARE seed-78 reference is repeated for comparison. Each path applies the driven control from its saved drive-start state, uses the predicted boundary for the exact 3.5-ns zero-current relaxation, and scores the final endpoint. Distribution Score, angular energy distance, and paired angle use 33 base conditions with matched five-versus-five endpoints; fair energy uses 165 exact anchors. The three-seed representation/objective controls are reported in Table 2 (a).

Figure 11: Separate single-segment qualitative geometry-ablation predictions. Rows are the 10th, 50th, and 90th percentiles of the seed-78 Cartesian-CFM-minus-core exact-anchor mean-draw angular-error gap over 330 anchors, selected without visual inspection. Each method panel shows the medoid of five draws under inter-draw angular distance; inset values are mean draw-to-MuMax 3 angular errors. Cartesian CFM is competitive in the low-gap case but shows larger domain-wall displacement in the upper tail. All panels show

m_{z}\in[-1,1]

.

Jitter ( ∘ )	Direct U-Net	CNO-FM	DPOT-Ti	Poseidon-T	PDEArena U-Net
0	0.8077	0.7451	0.8140	0.7515	0.7743
0.25	0.7914	0.7437	0.8123	0.7481	0.7721
0.5	0.7819	0.7425	0.8109	0.7453	0.7710
1	0.7633	0.7401	0.8083	0.7402	0.7675
2	0.7338	0.7355	0.8035	0.7312	0.7638
5	0.6778	0.7233	0.7873	0.7062	0.7502

Appendix

Table 16: Fine-grained anchor-jitter sensitivity in fair energy. Each deterministic method uses its seed-78 checkpoint; jitter is the prescribed pre-decoding valid-cell RMS tangent-space rotation magnitude applied only at the initial anchor. Values are point estimates. Lower is better. Bold denotes the best result for each method across all jitter amplitudes. FLARE is not perturbed and has fair energy 0.393 (Table 1 ).

Figure 12: Teacher-forced and autoregressive rollouts. (a) Angular error for the compute-matched Stage–2 controls. (b) rollout degradation, defined as autoregressive minus teacher-forced error. (c) segment-resolved error for Stage 1, mixed Stage 2, and GT-only Stage 2. The first segment shares the real initial state; subsequent rollout steps expose each method to model-generated states. Intervals cluster 736 rollouts by 97 base conditions.

Figure 13: Stage-2 multi-stage evaluation using the same formal results as Table 17 . (a) Exact-boundary and autoregressive angular error. (b) Autoregressive minus exact-boundary error. Whiskers are 95% base-condition bootstrap intervals.

Method	Ground-truth-boundary angle ( ∘ )	Rollout angle ( ∘ )	Rollout degradation ( ∘ )
FLARE-S1	32.82 [30.27, 35.38]	37.74 [34.95, 40.44]	4.92 [4.12, 5.79]
FLARE-S2-GT	31.82 [29.30, 34.33]	36.52 [33.80, 39.22]	4.70 [3.86, 5.62]
FLARE-S2-mixed	32.27 [29.73, 34.81]	36.80 [34.05, 39.55]	4.53 [3.70, 5.42]

Appendix

Table 17: Multi-segment evaluation and compute-matched Stage–2 controls. Brackets are 95% base-condition cluster-bootstrap intervals; lower is better.

Figure 14: Composition consistency and boundary sensitivity. (a) One-shot/composed defect versus the energy distance between the two output ensembles. (b) final error as the second-segment input moves from the exact boundary (

\lambda=0

) to a generated boundary (

\lambda=1

). (c) the same curves parameterized by measured input error. Internal consistency does not imply physical fidelity: a small path defect can coexist with a worse ensemble distance.

Figure 15: Generalization to a geometry family excluded before training. (a) FLARE ID and ring-OOD angular errors. (b) the matched auxiliary FNO architecture control (not a Table 1 comparator). (c) relative OOD changes for single-segment, teacher-forced, autoregressive, and final-rollout-step evaluation. Error bars resample base conditions separately within each stratum.

Evaluation	ID angle ( ∘ ) $\downarrow$	Ring-OOD angle ( ∘ ) $\downarrow$	$\Delta$ ( ∘ )	$\Delta$ (%)
Single segment	29.05 [26.42, 31.73]	29.96 [28.36, 31.59]	+0.91	+3.12
Autoregressive handoff	31.40 [28.64, 34.15]	32.35 [30.64, 34.05]	+0.95	+3.04

Appendix

Table 18: Generalization to the completely held-out ring geometry. Brackets are 95% base-condition cluster-bootstrap intervals.

Evaluation	Non-ring ID [95% CI]	Held-out ring [95% CI]
Single-segment angle ( ∘ )	29.05 [26.42, 31.73]	29.96 [28.36, 31.59]
Single-segment MSE	0.42 [0.37, 0.48]	0.44 [0.41, 0.47]
Single-segment $\|\Delta Q\|$	0.54 [0.477, 0.60]	0.56 [0.52, 0.60]
Teacher-forced rollout angle ( ∘ )	28.85 [26.14, 31.60]	29.92 [28.26, 31.59]
Autoregressive rollout angle ( ∘ )	31.40 [28.64, 34.15]	32.35 [30.64, 34.05]
Final-step angle ( ∘ )	35.42 [32.28, 38.60]	36.21 [34.28, 38.13]

Appendix

Table 19: FLARE under the ring-family leave-one-factor-out split. The ID column uses unseen non-ring bases from the same refitted split; the OOD column contains only ring controls, all excluded before fitting and normalization. Intervals resample base conditions.

$T$	Metric	ID interpolation reference	Unseen temperature	Absolute degradation [95% CI]	Relative degradation [95% CI]
90 K	Single-segment angle ( ∘ )	23.815	24.118	$+0.303$ [ $-0.230$ , $+0.843$ ]	$+1.27\%$ [ $-0.96\%$ , $+3.55\%$ ]
90 K	Autoregressive handoff angle ( ∘ )	31.213	29.917	$-1.296$ [ $-2.573$ , $-0.117$ ]	$-4.15\%$ [ $-8.06\%$ , $-0.39\%$ ]
90 K	Distribution Score	0.6687	0.7380	$-0.0693$ [ $-0.1616$ , $+0.0270$ ]	$-10.36\%$ [ $-24.99\%$ , $+3.99\%$ ]
90 K	Angular energy distance ( ∘ )	12.749	12.385	$-0.365$ [ $-1.388$ , $+0.598$ ]	$-2.86\%$ [ $-10.50\%$ , $+4.89\%$ ]
225 K	Single-segment angle ( ∘ )	37.291	37.903	$+0.612$ [ $-0.102$ , $+1.300$ ]	$+1.64\%$ [ $-0.28\%$ , $+3.45\%$ ]
225 K	Autoregressive handoff angle ( ∘ )	44.320	45.372	$+1.052$ [ $-0.781$ , $+2.777$ ]	$+2.37\%$ [ $-1.77\%$ , $+6.36\%$ ]

Appendix

Table 20: Zero-shot generalization to in-range temperatures absent during training. The frozen Stage 1 checkpoint is trained at 30, 150, and 300 K. For 90 K and 225 K, the ID reference is the paired per-base midpoint of the two adjacent training temperatures. Brackets report 95% paired base-condition cluster-bootstrap intervals. Positive degradation denotes worse performance at the unseen temperature.

Figure 16: Time-resolved autoregressive rollout A: difficult stripe and ring-wall reorganization. MuMax 3 , the FLARE rollout, and the full-vector angular error are aligned at every displayed time. The first relaxation frame has

24.4^{\circ}

mean error; the longer relaxation reaches

14.5^{\circ}

. Error localizes on moving walls and newly resolved branches, while the large-scale magnetic regime remains recognizable.

Figure 17: Time-resolved autoregressive rollout B: stable stripe-wall evolution. The same fully autoregressive protocol is applied to a second fixed case. Wall motion and the disappearance of small domains remain qualitatively similar over the long second segment; its final mean angle is

22.9^{\circ}

, compared with

23.9^{\circ}

in the first relaxation frame.

Figure 18: Time-resolved autoregressive rollout C: many-body domain extinction. The model follows the removal and contraction of several small reversed domains across both controls. Its errors remain concentrated at interfaces and surviving cores instead of spreading through saturated regions; the first-relaxation and final means are

25.6^{\circ}

and

23.3^{\circ}

.

Figure 19: Time-resolved autoregressive rollout D: large-domain contraction. FLARE preserves the two long walls while a lower domain retracts and leaves the field of view. The full-vector map reveals thin registration errors that are hard to see from

m_{z}

alone; mean angular error stays near

12

–

23^{\circ}

throughout the rollout.

Figure 20: Single-segment magnetic evolution across numerical difficulty quantiles. Rows are selected at the 25th percentile, median, and 90th percentile without visual inspection. Columns show the anchor, MuMax 3 target, Stage 1, S2-mixed, and the Stage 2 three-component angular-error map. Stage 2 mean angles for these displayed draws are

17.1^{\circ}

,

31.4^{\circ}

, and

43.0^{\circ}

. Wall displacement and changes in small domains become more pronounced in the difficult case.

Figure 21: Method-level field and error comparison on the fixed median case. The top row shows

m_{z}

; the bottom row shows per-cell full-vector angular error relative to the same MuMax 3 target. The methods retain different amounts of fine-scale wall roughness. The auxiliary internal FNO control’s localized defects are more visible in the angle map than in the scalar state average; FNO is not a Table 1 comparator. This plate is illustrative; aggregate rankings use all held-out cases.

Figure 22: Vector texture and local topology. The formal q25 and median cases from Figure 20 are reused without further selection. MuMax 3 and S2-mixed are shown with sparse in-plane arrows over

m_{z}

. The third column gives full-vector angular error; the fourth gives predicted-minus-reference topological density, computed by open-boundary central differences in grid-cell units on a shared symmetric color scale. The angle maps highlight displaced walls and changes in small domains; local topological residuals retain both positive and negative contributions.

Figure 23: Fully autoregressive magnetic-field rollouts. Rows are the median and 90th-percentile rollouts under the fixed numerical selection rule. Stage 1 and mixed Stage 2 receive their own first-segment prediction as the next anchor. For these displayed draws, the mean final error decreases from

38.5^{\circ}

to

36.4^{\circ}

at the median and from

55.2^{\circ}

to

52.6^{\circ}

in the hard tail. These two examples do not establish a uniform improvement over all cases or over the equal-budget GT-only control. The final two columns show full-vector angular errors relative to MuMax 3 .

Figure 24: One-shot versus composed magnetic transitions. For the fixed median-defect case, each row compares a direct transition over the summed horizon with two adjacent generated transitions. The rightmost map is their local composition defect. The selected case has a

15.15^{\circ}

mean Stage 1 defect over the evaluation draws; the separate displayed draws have defects of

26.2^{\circ}

(Stage 1) and

27.2^{\circ}

(S2-mixed). Both routes retain a stripe pattern but disagree on wall placement and branch structure. Agreement between model routes does not establish agreement with MuMax 3 .

Figure 25: Continuous rollout-state perturbation in magnetization space. The first row shows the true and generated segment boundary, the final MuMax 3 state, and Stage 2 predictions as

\lambda

moves from the exact to the predicted boundary. The second row shows boundary or continuation-state angular error. The continuation sampling seed is held fixed across

\lambda

; the displayed mean output error increases from

28.8^{\circ}

at the true boundary to

36.4^{\circ}

at the predicted boundary. The angle maps expose changes around displaced walls along this unit-sphere interpolation path.

Method class	Full state	Direct	Stoch.	Magn.	Status
Classical solvers	$\checkmark$	$\times$	$\checkmark$	$\checkmark$	reference
Hybrid learned solvers	$\checkmark$	$\times$	$\checkmark$	$\checkmark$	compared
Reduced-order dynamics	$\checkmark$	partial	$\times$	$\checkmark$	not compared
Observable-level neural ODE	$\times$	$\checkmark$	$\times$	$\checkmark$	n/a
Static-field / energy minimization	–	–	$\times$	$\checkmark$	n/a
Texture recognition / inversion	–	–	–	$\checkmark$	n/a

Appendix

Table 21: Positioning by method class. Full state : returns the full spatial state (for magnetization, the full field

\mathbf{m}

) rather than a low-dimensional observable. Direct : generates a state at a requested finite horizon without integrating intermediate LLG steps. Partial : the requested-horizon interface is absent or method-specific within the class (latent-space advancement, fixed-step rollout, or fixed-length trajectory windows) rather than a uniform one-shot interface. Stoch. : natively returns a distribution over outcomes rather than a point prediction. Magn. : developed for magnetization dynamics. n/a marks classes that cannot be scored by the field-based metrics used here. Not compared means that no result under the common protocol is reported. Citations for each class are given in the preceding paragraph.

Method	$B{=}1$ latency (ms)	Distribution Score $\uparrow$	Paired angle ( ∘ ) $\downarrow$	Angular energy dist. ( ∘ ) $\downarrow$	Fair energy score $\downarrow$	$B{=}1$ speedup $\uparrow$	Best-batch speedup $\uparrow$
Reference
MuMax 3 5 vs 5	–	0.862 [0.801, 0.915]	30.25 [23.89, 36.66]	12.20 [9.62, 14.86]	–	–	–
FLARE ODE-10	995.2	0.762 [0.660, 0.844]	36.72 [31.93, 41.38]	15.30 [13.32, 17.26]	0.393 [0.356, 0.430]	$448.3\times$	$3{,}062.4\times$
Poseidon-T	61.91	0.341 [0.077, 0.749]	35.60 [30.60, 41.11]	25.05 [17.86, 34.56]	0.751 [0.672, 0.830]	$7{,}206.3\times$	$686{,}423.6\times$
CNO-FM	72.89	0.695 [0.577, 0.787]	34.43 [30.41, 38.33]	21.82 [17.94, 26.23]	0.745 [0.682, 0.807]	$6{,}120.7\times$	$133{,}012.0\times$
DPOT-Ti	91.13	0.707 [0.599, 0.803]	39.81 [34.46, 45.59]	27.02 [19.98, 36.58]	0.814 [0.735, 0.893]	$4{,}895.3\times$	$48{,}327.7\times$

(a) Representation and objective controls: three training seeds
Variant	Distribution Score $\uparrow$	Angular energy dist. ( ∘ ) $\downarrow$	Paired angle ( ∘ ) $\downarrow$	Fair energy score $\downarrow$
FLARE	$\bm{0.7592\pm 0.0056}$	$\bm{15.31\pm 0.24}$	$\bm{36.70\pm 0.23}$	$\bm{0.3944\pm 0.0028}$
FLARE(Fixed 2D tangent)	$0.7562\pm 0.0221$	$16.70\pm 2.09$	$37.94\pm 1.41$	$0.3991\pm 0.0155$
Cartesian CFM	$0.7176\pm 0.0073$	$17.98\pm 0.12$	$39.96\pm 0.29$	$0.4210\pm 0.0033$
Riemannian FM	$0.7341\pm 0.0071$	$18.00\pm 0.19$	$39.01\pm 0.10$	$0.4330\pm 0.0029$
Direct U-Net	$0.4741\pm 0.0040$	$21.26\pm 0.13$	$38.60\pm 0.07$	$0.8087\pm 0.0013$

FLARE: Flow Matching with Local Axis-Angle Representations for Stochastic Micromagnetic Evolution

Abstract

Figures & tables

Appendix

Explore similar work

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

FlowForge: A Staged Local Rollout Engine for Flow-Field Prediction

Capturing non-Markovian dynamics in non-equilibrium stochastic systems using flow matching