Source-linked AI summary

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

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

arXiv:2609.34070v1cs.LGcond-mat.mes-hallcs.AI

TL;DR

Long-horizon stochastic micromagnetic simulation is costly, while learned methods have largely retained sequential propagation or deterministic prediction. FLARE uses conditional flow matching over anchor-relative axis-angle rotations to generate horizon-conditioned stochastic endpoints. It reports improved distributional metrics against external baselines and a large speedup over MuMax3 on a composed 5-ns workload.

  • Problem

    Long-horizon and ensemble micromagnetic simulations are costly, and direct generation of full-field stochastic endpoints across physical horizons remains largely unexplored.

    Source: §I · p. 1
  • Method

    FLARE uses conditional flow matching over anchor-relative axis-angle rotations, with physical-horizon conditioning and Rodrigues decoding for stochastic endpoint generation.

    Source: §I · p. 1
  • Results

    FLARE achieves 29.9% lower angular energy distance (15.30° versus 21.82° for CNO-FM) and a 3,062× best-batch speedup over MuMax3 on a 5-ns workload.

    Source: B. Main Results · p. 6
  • Takeaways & Limitations

    Anchor-relative rotational coordinates provide an approach to modeling finite-temperature micromagnetic dynamics as conditional future-state distributions.

    Source: §V · p. 7

    Reported observable-distribution evidence uses a matched five-sample budget and covers several observables, including topological charge and mean out-of-plane magnetization.

    Source: §V · p. 7

Abstract

from arXiv · show

Long-horizon micromagnetic simulation remains expensive because conventional and learned solvers typically propagate Landau--Lifshitz--Gilbert (LLG) dynamics step by step. Existing learned approaches generally retain stepwise integration or model deterministic evolution, leaving full-field, direct-horizon stochastic prediction largely unexplored. We propose FLARE, a flow-matching framework that recasts stochastic finite-time magnetization prediction as conditional transport over anchor-relative local axis-angle rotations. This rotation-space formulation respects the intrinsic geometry of magnetization dynamics and preserves pointwise unit norm by construction. By explicitly conditioning on the physical prediction horizon, FLARE directly generates full-field stochastic endpoints across multiple target times without stepwise integration. Against the strongest single-checkpoint external baseline on each metric, FLARE achieves 29.9% lower angular energy distance ($15.30^\circ$), and a 37.3% lower fair energy score (0.393). On a representative composed 5-ns two-segment protocol, FLARE achieves a $3{,}062\times$ best-batch speedup over the widely used GPU micromagnetic solver MuMax$^3$ on a single GPU.

I. INTRODUCTION

FLARE addresses the cost of long-horizon stochastic micromagnetic simulation by generating full-field endpoints directly across physical horizons. It represents changes as anchor-relative rotations, preserving unit magnetization through Rodrigues decoding.

I. INTRODUCTION

  • Conventional micromagnetic solvers integrate LLG dynamics through many small steps, making long-horizon and ensemble simulation expensive.Source: §I · p. 1
  • Existing learned approaches largely target specialized tasks, sequential simulation, or deterministic prediction, leaving direct full-field stochastic endpoints across horizons largely unexplored.Source: §I · p. 1
  • FLARE models finite-time magnetization changes as anchor-relative axis-angle rotations and uses Rodrigues decoding to preserve pointwise unit norm.Explicit physical-horizon conditioning supports stochastic endpoint generation across target times without stepwise LLG integration.Source: §I · p. 1
  • Contributions include matched-control evaluation of rotational coordinates, generalization to held-out horizons and self-generated anchors, and a framework spanning field error, distribution fidelity, observables, and latency.Source: §I · p. 1

II. BACKGROUND AND RELATED WORK

Micromagnetic simulation and learned surrogates have reduced costs in some settings, but sequential propagation and deterministic or reduced-order modeling remain common. FLARE applies generative modeling to full-field stochastic micromagnetics using local anchor-relative rotations.

II. BACKGROUND AND RELATED WORK

  • Micromagnetic solvers numerically integrate stochastic or deterministic LLG dynamics, while high-resolution, long-horizon, and ensemble runs remain costly because they propagate full fields through successive time steps.Source: §II · p. 1
  • Learned micromagnetic methods span surrogate fields, inverse problems, reduced-order dynamics, and neural operators, but many retain sequential integration or focus on deterministic or reduced-order prediction.Source: §II · p. 2
  • Generative models and flow matching have been used for stochastic physical systems, including fluids, weather, and molecular dynamics, but these applications do not themselves address micromagnetic geometry.Source: §II · p. 2
  • FLARE models finite-time changes as local rotations relative to each anchor magnetization, using Euclidean flow matching and rotation-based decoding to enforce the spherical constraint.Source: §II · p. 2

A. Problem setup

The setup treats magnetization as a unit-length spatial field evolving under stochastic LLG dynamics and asks for its conditional future-state distribution. A query specifies the anchor, physical conditions, and prediction horizon; multi-stage protocols recursively use predicted endpoints as new anchors.

A. Problem setup

  • The state is a three-component unit magnetization field on a two-dimensional domain, governed by stochastic LLG dynamics.Source: A. Problem setup · p. 2
  • Given an anchor, physical condition, and target horizon, FLARE models the conditional distribution of the future state; conditions combine global controls with spatially resolved features.Source: A. Problem setup · p. 2
  • A multi-stage protocol assigns conditions and durations to segments, then uses each generated endpoint as the next segment's anchor to form a generated state sequence.Source: A. Problem setup · p. 2

B. Model formulation

FLARE encodes each transition as an axis-angle rotation field relative to the anchor and learns a conditional flow in rotation space. At inference it decodes sampled rotations to magnetization, while conditioning on the horizon, controls, and spatial features.

B. Model formulation

  • At inference, ODE sampling transports an initial Gaussian rotation field along the learned velocity field to a terminal rotation, which is decoded back to magnetization by Rodrigues rotation.Generated endpoints can serve as anchors for subsequent segments under their corresponding conditions.Source: Figure 1 · p. 3, B. Model formulation · p. 4
  • Each anchor-to-target transition is represented by a principal local axis-angle rotation field, whose direction gives the rotation axis and magnitude the angle.Rodrigues reconstruction preserves pointwise unit norm; special approximations handle nearly identical and near-antipodal moments.Source: B. Model formulation · p. 3
  • Conditional flow matching interpolates between a Gaussian source rotation field and the target rotation field, learning a velocity conditioned on the anchor, horizon, controls, and spatial features.The flow time is distinct from the physical prediction horizon, with uniform weighting across sampled flow times, horizons, and driving categories.Source: B. Model formulation · p. 4
  • The conditional vector field uses a native-resolution U-Net with anchor magnetization, rotation state, and spatial features as input, while horizon, flow time, and global controls modulate features through AdaGroupNorm.An eight-head bottleneck self-attention block captures long-range spatial dependencies.Source: B. Model formulation · p. 4
  • Training pairs sample an anchor and valid horizon within a constant-control segment, using the nearest saved frame that remains within that segment as the target.The actual timestamp difference defines the physical horizon; validation uses the earliest valid anchor in each admissible range.Source: B. Model formulation · p. 4

C. Distributional Evaluation Metric.

The metric evaluates sets of stochastic magnetization outcomes while tolerating small spatial displacements of similar structures. Its local translation-matching distance feeds a density-based Distribution Score for agreement with MuMax3 endpoint distributions.

Local translation-matching distance

  • Local translation matching compares magnetization patches after bounded shifts, reducing penalties for small displacements of domain walls, vortices, and other localized structures.The field is divided into 4 × 4 patches, each matched within ±16 grid cells along each spatial direction.Source: C. Distributional Evaluation Metric · p. 4, Figure 3 · p. 5

Distribution Score

  • For each MuMax3 outcome, the Distribution Score compares a model-generated local density with a leave-one-out density from the remaining MuMax3 repeats.Its kernel bandwidth is set only from nearest-neighbor distances among MuMax3 samples, independently of model outputs.Source: C. Distributional Evaluation Metric · p. 4
  • The Distribution Score ranges from 0 to 1, with values closer to 1 indicating better agreement between generated and MuMax3 endpoint distributions.The score is based on a symmetric distance designed to remain tolerant to small spatial translations.Source: C. Distributional Evaluation Metric · p. 4

A. Experimental Setup

FLARE is evaluated on finite-temperature MuMax3 trajectories with condition-level splits, external and matched internal comparators, and separate quality and runtime protocols. Distribution fidelity is assessed across held-out complete paths with clustered uncertainty estimates.

Dataset and splits

  • The dataset contains 7,770 completed trajectories from 1,000 base conditions on a 256 × 256 × 1 grid, with splits made by base condition to avoid leakage.The corpus varies temperature, field, current amplitude, initial texture, and current-mask family; material parameters are fixed.Source: A. Experimental Setup · p. 5

Comparators and execution protocol

  • Seven author-released external models and two matched internal controls are compared after adaptation to the same training split and 50,000 optimization updates.External comparators are deterministic per checkpoint except NeuralMAG-x5, which retains the thermal field.Source: A. Experimental Setup · p. 5
  • Quality is scored at the final endpoint after each method's predicted driven-segment boundary state enters the exact 3.5-ns relaxation, while runtime uses a standardized two-segment 5-ns workload.Runtime reports batch-one latency and the minimum per-output latency across a power-of-two batch sweep.Source: A. Experimental Setup · p. 6

Metrics and uncertainty

  • Distribution fidelity is evaluated on 33 held-out complete-path conditions with five repeats using Distribution Score, angular energy distance, and fair energy score.Uncertainty uses 95% base-condition-clustered bootstrap intervals; physical forecast spans range from 3.75 to 4.42 ns, averaging 4.10 ns.Source: A. Experimental Setup · p. 6

B. Main Results

FLARE leads the external baselines on all three distributional metrics, while CNO-FM has a lower repeat-paired angle. FLARE also provides substantial speedup on the composed 5-ns workload, and multi-segment autoregressive robustness remains below an equal-budget ground-truth-only continuation.

Main benchmark results

  • 15.30° angular energy distance is 29.9% lower than CNO-FM's 21.82°, and FLARE's fair energy score of 0.393 is 37.3% lower than NeuralMAG-x5's 0.627.FLARE also records a Distribution Score of 0.762 versus 0.707 for DPOT-Ti (+7.9%); it leads all seven author-released baselines on the three distributional metrics.Source: B. Main Results · p. 6
  • 3,062× best-batch speedup over MuMax3 is achieved on the composed 5-ns workload, with a 448× speedup at batch size 1.On the driven segment, FLARE's fair energy score is 0.322 [0.290, 0.353] versus 0.283 [0.230, 0.335] for the MuMax3 reference.Source: B. Main Results · p. 6

Multi-segment composition

  • Across 736 rollouts from 97 held-out conditions, predicted-state exposure improves autoregressive robustness over Stage 1 but does not beat an equal-budget ground-truth-only continuation.The study evaluates both ground-truth-boundary and autoregressive handoffs; predicted-state exposure is treated as an ablation rather than a separation contribution.Source: B. Main Results · p. 6

C. Out-of-Distribution Generalization

FLARE is tested for out-of-distribution generalization to an unseen current-mask family and unseen temperatures. Reported errors indicate stable transfer to the held-out mask family and robust interpolation at the evaluated in-range temperatures.

Unseen current-mask family

  • On the held-out ring-mask family, single-segment and autoregressive angular errors rise by 0.91° (3.12%) and 0.95° (3.04%) relative to the in-distribution reference.The model was trained after excluding ring-current-mask conditions from training and validation.Source: C. Out-of-Distribution Generalization · p. 6

Unseen temperature

  • At 90 K, single-segment and autoregressive angular errors change by +1.27% and −4.15%; at 225 K, they change by +1.64% and +2.37%.No evaluated metric shows statistically resolved adverse degradation at either unseen in-range temperature.Source: C. Out-of-Distribution Generalization · p. 6
  • Temperature generalization is tested zero-shot at 90 and 225 K using a frozen Stage 1 checkpoint trained only at 30, 150, and 300 K.The reported comparisons use paired interpolation references from adjacent training temperatures.Source: C. Out-of-Distribution Generalization · p. 6

D. Ablation Studies

Matched ablations show that rotational coordinates and target-horizon conditioning improve predictions, and FLARE uses the 3D axis-angle form to avoid per-spin tangent frames and reduce seed variability.

D. Ablation Studies

  • The matched controls isolate learned endpoint dynamics, anchor-relative rotational coordinates, physical-horizon conditioning, and the flow-matching objective.Source: D. Ablation Studies · p. 6
  • Table I reports exact-control complete-path quality and standardized 5-ns timing; learned rows use 50,000 x5 updates and brackets are 95% base-group cluster-bootstrap intervals.Its simulator-only top row is a descriptive scale, not a ranked entry.Source: Table I · p. 7
  • Timing uses a fixed two-segment 5-ns workload, with speedups referenced to MuMax3's 446.125-s runtime.The workload comprises a 2-ns drive and 3-ns relaxation with one predicted-state handoff; best batch is the fastest supported power of two.Source: D. Ablation Studies · p. 7
  • Both 3D axis-angle and fixed 2D tangent outperform absolute-spin Cartesian CFM and Riemannian FM across all four metrics.The comparison associates the gain with rotational coordinates rather than dimensionality.Source: D. Ablation Studies · p. 7
  • Target-horizon conditioning improves complete-path predictions and interpolation at held-out 1.5- and 2.5-ns horizons.Source: D. Ablation Studies · p. 7
  • FLARE uses 3D axis-angle coordinates to avoid per-spin tangent frames and reduce seed variability.Source: D. Ablation Studies · p. 7

V. DISCUSSION AND LIMITATIONS

The discussion examines FLARE’s distributional performance and practical scope while qualifying the results through metric limitations, supported-horizon and efficiency diagnostics, and controls on representation, ensemble size, and anchor perturbation. It distinguishes distributional quality from realization-level accuracy and identifies limits on broader claims from the available tests.

1. Rotation-space transition parameterization

  • Rodrigues decoding preserves unit norm, while the deterministic near-antipodal fallback applies to 0.1590% of valid cells but occurs in 327/330 anchor–target pairs.The fallback approximates near-antipodal targets with angular error below 2.563°; the rotation map’s orthogonality, rather than Cartesian-prediction repair, preserves the norm.Source: §1 · p. 12

4. Comparator and timing protocol

  • Timing uses a standardized two-segment 5-ns workload, while confidence intervals resample held-out base conditions and exclude simulator-parameter and model-family selection uncertainty.FLARE uses ten Heun steps per segment, or 40 velocity evaluations for the two-segment workload; training-seed variation is reported separately.Source: §4 · p. 13, §2 · p. 15

1. Field, topology, and ensemble metrics

  • Field error, topology and energy observables, distribution fidelity, and repeat-paired angular error measure distinct aspects of prediction and are not interchangeable.Local-translation matching tolerates small physically equivalent shifts, while global observables can reveal topology-changing defects concealed by field averages.Source: §1 · p. 15

Appendix C: Distribution fidelity

  • 15.30° angular energy distance and 0.762 Distribution Score make FLARE the best reported method on both complete-path distribution metrics.The comparison uses five outcomes per method across 33 complete-path conditions; Cartesian CFM, the matched representation control, scores 0.711 on Distribution Score.Source: Appendix C · p. 16, Figure 4 · p. 17

2. Sensitivity, simulator self-consistency, and physical observables

  • FLARE ranks first under every tested bandwidth and without translation matching; its driven fair energy score is 0.322 [0.290, 0.353], versus the MuMax3 reference 0.283 [0.230, 0.335].The fair-energy comparison uses a separate 17-base x30 audit, not the 33 x5 bases in the main comparison, so it is descriptive rather than a paired test or universal bound.Source: §2 · p. 16, Table VI · p. 17, Table VII · p. 17

1. Conditional outcome variability under identical physical inputs

  • Across fixed-input cases, FLARE generates diverse outcomes that overlap substantially with MuMax3 in field and observable space, unlike deterministic Poseidon-T’s single prediction.The audit compares 30 MuMax3 realizations and 30 fresh FLARE draws per case under identical anchors, controls, and horizons.Source: §1 · p. 16, Figure 5 · p. 18

3. Metric and finite-ensemble pressure tests

  • Distribution Score changes from 0.691 to 0.694 and fair energy score from 0.3095 to 0.3092 as model outcomes rise from 5 to 20, while angular energy distance drifts downward through estimator bias.The same angular-energy drift appears between two MuMax3 ensembles, so cross-method comparisons retain equal five-versus-five budgets; the single-segment subset ranking remains stable at one outcome and above.Source: §3 · p. 16, Table VIII · p. 17, Figure 6 · p. 19
  • Local matching can attenuate topology-changing defects, so translation-tolerant field scores require complementary global observables; the nominal 4 × 4, ±16 setting is not uniquely selected by the examples.FLARE is best for mean mz, while NeuralMAG-x5 is best for charge and exchange texture after matched 50,000-component adaptation.Source: §3 · p. 20

Appendix D: Solver, efficiency, and target-time diagnostics

  • At 1.5 and 2.5 ns, target-time conditioning improves the matched no-time control by 8.91° [7.32, 10.69] and 8.31° [6.97, 9.73], respectively.Across 1,493 single-segment test transitions, masking target time increases angular error by 7.625° [6.64, 8.71]; these are interpolation queries within realized duration support.Source: Appendix D · p. 20, Figure 9 · p. 22

1. Complete supported-horizon distribution audit

  • Across the legal 1–3 ns range, paired angle rises from 29.17° to 32.87°, angular energy distance from 11.83° to 13.54°, and fair energy score from 0.3189 to 0.3622.Distribution Score remains between 0.776 and 0.823; the 3.5–6 ns entries cannot support one-shot long-horizon fidelity claims, and spread ratios of 0.991–0.999 show no statistically resolved mode collapse or over-diffusion over evaluated support.Source: §1 · pp. 21, 24, Figure 10 · p. 22, Table XIII · p. 23

Appendix E: Multi-seed representation and objective controls

  • The core’s mean beats Cartesian CFM, Riemannian FM, and Direct U-Net on all four metrics, while the fixed 2D basis is close in mean but has substantially greater seed variation.The controls share the split, architecture, global batch 128, 50,000 updates, and seeds 78–80; the results support the coordinate family while separating that claim from prior- or basis-specific claims.Source: Appendix E · p. 24
  • 19.2% lower angular energy distance (15.30° versus 18.93°) and 21.8% lower fair energy (0.393 versus 0.503) than the best external five-model ensemble, Poseidon-T.This matched five-output comparison is not equal in total training compute; Poseidon-T has a lower paired angle (33.31° versus 36.72°).Source: Appendix E · p. 24
  • Anchor jitter produces strongly non-monotone performance, but even the best perturbed deterministic result remains above FLARE’s unperturbed fair energy of 0.393.Direct U-Net reaches 0.5022 at 40°; the sweep perturbs only the initial anchor and scores against the original MuMax3 endpoints.Source: Appendix E · p. 24

Appendix F: Multi-segment composition

The appendices examine multi-segment composition, transfer to unseen geometries and temperatures, rollout behavior, failure modes, and the scope of the comparisons. Results show qualified transfer and composition consistency, alongside trade-offs and boundaries on the supported conclusions.

Appendix F: Multi-segment composition

  • S2-mixed reduces autoregressive rollout error by 2.48% [1.49%, 3.45%] versus Stage 1, but has 0.78% [0.09%, 1.51%] higher error than equal-budget S2-GT.Its 3.61% rollout-degradation reduction versus S2-GT has an interval crossing zero, so mixed rollout-state exposure is not superior to GT-only continuation on absolute rollout error.Source: Appendix F · pp. 24–25
  • FLARE Stage 1 has a 14.88° [13.58, 16.17] one-shot/composed defect and 13.62° [12.51, 14.70] ensemble energy distance; rotation transport lowers these by 2.79° and 1.52° versus matched Cartesian CFM.These consistency measures do not replace comparison with MuMax3; mixed Stage 2 final error grows from 32.32° at the exact boundary to 36.67° at the predicted boundary, without abrupt failure on the sampled path.Source: Appendix F · pp. 25–26

Appendix G: Leave-one-mask-family-out: held-out ring controls

  • With all ring conditions excluded before fitting, FLARE’s error shifts by +0.91° (3.12%) in single-segment evaluation, +0.95° (3.04%) over complete rollouts, and +0.80° (2.25%) at the final segment.The result supports stable transfer under one unseen mask family, not invariant performance; resolution, simulator, representation, and broad physical-parameter support remain fixed.Source: Appendix G · p. 29

Appendix H: Zero-shot generalization to unseen temperatures

  • At unseen in-range temperatures, the frozen model follows the temperature-dependent ID trend without temperature-specific adaptation; at 400 K, angular errors and angular energy distance rise relative to the 300 K endpoint.At 400 K, single-segment and autoregressive angular errors increase by 17.54% [15.00%, 20.18%] and 16.75% [12.58%, 21.21%], respectively, while angular energy distance increases by 44.24% [35.47%, 53.20%].Source: Appendix H · p. 29

1. Selection and display conventions

  • Displayed cases are selected by fixed numerical rules before visualization, and printed errors describe individual draws rather than selection quantiles.Figures 20 and 23–25 use a frozen case manifest and fixed visualization seeds; time-resolved examples follow each prediction into the next hop.Source: §1 · p. 29

2. Time-resolved autoregressive rollouts

  • Autoregressive examples retain recognizable magnetic regimes, while errors concentrate around moving walls, interfaces, and small-domain changes.The displayed rollouts include stripe-wall reorganization, domain extinction, and large-domain contraction; the examples do not establish uniform improvement across cases or over equal-budget GT-only control.Source: Figure 16 · p. 31, Figure 17 · p. 32, Figure 18 · p. 33

Appendix J: Failure modes, negative results, and reproducibility

  • The analysis identifies wall displacement, uncertainty in small structures, and imperfect segment boundaries as recurring error sources; more ODE steps slightly improve charge error but worsen angular error and cost.Rollout-state-only continuation reduces the exposure gap further than GT-only but worsens clean-boundary accuracy, while a 1:1 mixture is a compromise.Source: Appendix J · p. 38

Appendix K: Positioning: scope of the comparison

  • The comparison includes methods only when they return a full magnetization field at a requested time, and the paper distinguishes direct stochastic endpoint generation from stepwise or deterministic alternatives.The authors report finding no prior learned model matching the combination of requested-horizon full-field prediction without stepwise LLG integration and conditional future-state distributions in their literature search through September 2, 2026.Source: Appendix K · p. 38
Loading 2609.34070v1…