$C_4$-Equivariant Flow Matching on Anisotropic Power-Diagram Graphs for Microstructure Generation

AI summary

Overview

Research area: Generative machine learning for computational materials science, specifically the synthesis of polycrystalline microstructures using flow matching, graph neural networks, and symmetry-equivariant architectures.

Technical level: Advanced. The paper assumes familiarity with flow matching, continuous normalising flows, optimal transport, graph message passing, and group-equivariance.

Scope: The paper presents a resolution-independent generative model that treats polycrystalline microstructures as anisotropic power diagrams, trains a $C_4$-equivariant graph network to sample their parameters, and adds a training-free guidance scheme for user-directed morphology control.

What This Paper Is About

Polycrystalline materials derive their properties from the geometry and arrangement of their grains, so realistic synthetic microstructures are needed to model material behaviour. Experimentally acquiring such data via Electron Backscatter Diffraction (EBSD) is costly and time-consuming. The paper's goal is to generate realistic microstructures cheaply by learning a compact geometric parametrisation, rather than modelling raw pixel images directly, and to allow users to steer generation toward specific structures without retraining.

Key Contributions

1. Flow matching over anisotropic power diagram parameters. Instead of generating images, the model generates the compact geometric parameters of an anisotropic power diagram (APD). Each grain is represented by its seed position in $[0,1]^2$ plus an ellipse-based anisotropy parametrisation. This representation is resolution-independent, supports variable grain counts, and allows rendering at arbitrary pixel resolution.

2. A graph-based, $C_4$-equivariant architecture. A graph neural network models dependencies between grains and guarantees that a rotation of the input noise by a multiple of $90^\circ$ produces a correspondingly rotated microstructure.

3. Training-free guidance via differentiable objectives. The model parameters stay fixed while user-defined differentiable objectives bias sampling, enabling complex, spatially varying global structure without retraining.

4. Qualitative generation of four motivated morphologies. Using the guidance framework, the authors generate microstructures qualitatively inspired by a copper weld, a cast metal slab, 3D-printed stainless steel, and heterogeneous lamella titanium.

Code is released at https://github.com/DawidLipin/Flow_APDs.

Main Findings

• Compact representation works for generation: The model learns a distribution over complete grain configurations parametrised as APDs, in contrast to existing APD methods that prescribe selected grain statistics a priori and do not directly model spatial dependencies.

• Equivariance holds by construction: The architecture separates the update into invariant scalar coefficients predicted by message passing and fixed geometric basis vectors with known transformation properties, so that $v_\theta(R\mathbf{z}, t) = R,v_\theta(\mathbf{z}, t)$ for any $90^\circ$ rotation $R$. Under a $90^\circ$ rotation, the anisotropy parameter transforms as $a_i \mapsto -a_i$; a $180^\circ$ rotation leaves $a_i$ unchanged.

• Guidance reshapes global structure without retraining: From the same initial noise, unguided and guided samples differ substantially. For the 3D-printed stainless steel example, guidance encourages four horizontal bands with alternating grain orientations of $+45^\circ$ and $-45^\circ$. The paper shows guided results alongside representative experimental microstructures taken from five sources (Buze et al., Balit et al., Reed-Hill, Savolainen et al., Wu et al.). These reference images are used only for qualitative comparison and do not enter any objective.

• Four morphology families reproduced qualitatively: The layer objective imposes elongated grains with alternating layer-wise directions; the cast-slab objective combines spatially varying shape, relative grain weight, and generator-density targets for chill, columnar, and central zones; the weld objective prescribes a smoothly varying inward growth field; the heterogeneous-lamella objective places a high-weight coarse population in wavy bands while surrounding smaller grains stay approximately equiaxed.

• Fidelity to experimental statistics is not validated: The model is trained only on synthetic APDs, so agreement with higher-order experimental statistics remains unvalidated. No quantitative accuracy metrics or benchmark comparisons are reported.

• Strong guidance can leave the learned distribution: Because objectives constrain only selected APD statistics, the method provides approximate control rather than exact sampling from a conditional distribution. Sufficiently strong or conflicting objectives may move samples away from the learned data manifold.

Methodology in Plain English

Representing a microstructure as geometry rather than pixels. A polycrystalline microstructure is described as a tessellation of the unit square into regions, where each region is defined by a seed point, an anisotropy matrix (which makes grains elongated rather than round), and a scalar weight (which controls region size). This is called an anisotropic power diagram. Because the description is parametric, it does not depend on raster resolution and it regularises away small pixel-level segmentation artefacts. For fixed seeds, anisotropies, and target areas, the weights are obtained by solving an optimal transport problem so that regions match prescribed areas.

Making the anisotropy matrices valid by construction. The anisotropy matrix must be symmetric positive definite. Rather than having the network predict entries directly, the authors write $A_i = \exp(S_i)$ where $S_i = \log A_i$ is real and symmetric, and additionally impose $\det A_i = 1$ so that the matrix encodes shape only, not size (size is handled by the power weights). This constraint implies $\operatorname{tr}(S_i) = 0$, so $S_i$ is a symmetric trace-free $2\times 2$ matrix determined by just two unconstrained scalars $a_i \in \mathbb{R}^2$. Any network output therefore yields a valid anisotropy matrix.

Learning to sample with flow matching. The model learns a continuous normalising flow that transports Gaussian noise to the distribution of APD parameters, trained with the standard simulation-free flow matching objective: predict the velocity $\mathbf{z}_1 - \mathbf{z}_0$ along a linear interpolation between noise $\mathbf{z}_0$ and data $\mathbf{z}_1$. Sampling integrates the learned vector field from $t=0$ to $t=1$. The vector field is parametrised by a graph neural network.

Graph structure. Each APD becomes a graph: nodes are the grains $(\mathbf{x}_i, a_i, w_i)$ and directed edges connect each seed to its $k$ nearest neighbours. Edge features (15 dimensional) are built only from quantities invariant under $90^\circ$ rotations, including pairwise distances, weights, anisotropy norms and inner/cross products, projections of anisotropies onto a quadratic map $\rho$ of the displacement, and a function $b$ that encodes position relative to the square boundary. Message-passing layers are conditioned on flow time $t$, and aggregation preserves permutation equivariance. The output head predicts invariant scalar coefficients (four $\alpha$, six $\beta$, one $\omega$, plus an attention score) that multiply geometric basis vectors with known transformation properties.

Steering generation without retraining. Following gradient-based guidance for diffusion and flow models, the authors bias sampling toward a differentiable objective $\mathcal{J}$ by sampling a tilted distribution proportional to $p_1(\mathbf{z}_1)\exp[-\beta\mathcal{J}(\mathbf{z}_1)]$. At each time $t$ they form a first-order estimate of the terminal state $\widehat{\mathbf{z}}_1 = \mathbf{z}t + (1-t)v\theta(\mathbf{z}t,t)$, differentiate the objective with respect to the current state, and subtract a scaled gradient from the velocity. Each gradient is divided by its root mean square (with $\epsilon = 10^{-8}$ for stability) so that numerical scale does not determine influence. Separate strengths $\lambda{k,c}$ control the effect of objective $k$ on position, anisotropy, and weight channels independently. Position strengths can differ horizontally and vertically. A cosine ramp schedule $s(t)$ applies no guidance over the first 10% of the trajectory.

Training and sampling configuration. The model has three message-passing layers, latent dimension 256, and $k=16$ nearest neighbours. Training used 12,000 synthetic APDs with between 100 and 500 grains, generated with PyAPD, trained for 100 epochs with the AdamW optimiser. Sampling uses $N=300$ generators, $k=16$ neighbours, and $K=60$ Euler steps. Each normalised channel of the base state is drawn independently from a standard Gaussian. The four generated columns in the main figure use random seeds 117, 237, 377, and 527; within a column every row starts from the same base state. The nearest-neighbour graph is recomputed after every step as positions evolve, and the updated state is detached before the next step, so gradients flow through the current terminal estimate but not through the full integration history. Neighbour selection itself is discrete and not differentiated.

Objective design details. The layer objective uses a signal $\ell_i = \tanh(12\sin(6\pi y_i))$ alternating sign across six horizontal layers, a target axis ratio of 5.5, and a stop-gradient confidence term that removes ambiguous orientation at layer interfaces and prevents grains from reducing the loss by drifting toward interfaces. The cast-slab objective uses three independently guided terms — shape, relative weight via centred mean-squared error, and generator density — with zone boundaries $b_{\mathrm{ch}} = 0.11$ and $b_{\mathrm{co}} = 0.58$, zone fractions $f_{\mathrm{ch}} = 0.34$, $f_{\mathrm{col}} = 0.48$, $f_{\mathrm{core}} = 0.18$, and a columnar axis ratio increasing from 2.2 to 9.2. The weld objective derives a target direction from the field $(1.9(0.5-x_i), 0.62)$ with axis ratio 5.5. The heterogeneous-lamella objective splits the population into a coarse fraction of $\mathrm{round}(0.22N)$ grains and the rest, uses four wavy bands from $h(x) = 0.035\sin(2\pi x)$, a relative-weight target of 1.55, a coarse axis ratio of 1.6, and a sorted mean-squared error that removes dependence on grain labels. Guidance strengths in the paper's Table 1 range from 0 to 18 depending on objective and channel.

Why This Matters

Impact on research. The work offers an alternative to pixel-space generative modelling of microstructures, where individual samples may contain millions of pixels and available EBSD data are often limited and noisy. By combining a compact geometric prior, discrete symmetry equivariance, and a plug-in guidance mechanism, it suggests that domain-specific structure can be injected at inference time rather than baked into training. The paper also notes that many existing approaches offer limited control over global characteristics such as collective grain alignment or microstructure-spanning spatial patterns.

Real-world applications (as presented or implied by the paper's objectives):

• Weld microstructures, where grain orientation changes smoothly across a joint.
• Cast metal slabs, which have distinct chill, columnar, and central zones with different grain shapes, sizes, and densities.
• 3D-printed metals, where layer-wise processing produces elongated grains with alternating directions.
• Heterogeneous lamella metals, where coarse and fine grain populations coexist in wavy bands.

Industry relevance. Since microstructures strongly influence material properties, cheap generation of representative samples supports modelling material behaviour and studying structure–property relationships. Resolution independence means a single generated sample can be rasterised at whatever pixel resolution downstream simulation requires, and the training-free guidance means practitioners can target a desired morphology without curating a dataset of that morphology or retraining a model.

Future Directions

• Validate and fine-tune on real EBSD data. The model is currently trained only on synthetic APDs; the authors state that fidelity to higher-order experimental statistics remains unvalidated and identify evaluation and fine-tuning on EBSD data as future work.

• Incorporate crystallographic orientation. The current representation captures grain geometry but not crystallographic orientation, which the authors list as an extension.

• Extend to three dimensions. The representation is two-dimensional; a 3D extension is proposed.

• Physics-informed guidance objectives. The authors highlight exploring alternative guidance objectives, particularly those based on physics-informed loss functions, as a promising direction. The open question is how to specify guidance strengths principledly — currently they were selected manually for qualitative agreement, and the authors note they need not be optimal and could instead come from domain knowledge, tuning against an application-specific validation criterion, or estimation from real data.

Target Audience

This paper is best suited to researchers working at the intersection of generative machine learning and materials science: those developing equivariant or geometric deep learning architectures, computational materials scientists who need synthetic microstructure data, and practitioners interested in training-free, inference-time control of generative models. Readers need a working understanding of flow matching, graph neural networks, and group-equivariance; the appendix-level derivations of the anisotropy parametrisation and equivariant update rule make it most rewarding for readers comfortable with differential geometry and symmetry arguments. Someone seeking benchmark numbers or head-to-head comparisons with other microstructure generators will not find them here.

Authors’ abstract

Acquiring realistic microstructure data through Electron Backscatter Diffraction (EBSD) is costly and time consuming, often relying on specialised equipment. As microstructures strongly influence material properties, generating realistic samples is essential for modelling the behaviour of polycrystalline materials. We introduce a generative model for synthesising realistic polycrystalline microstructures using flow matching and graph neural networks. By representing microstructures as anisotropic power diagrams, our model learns a compact geometric parametrisation and can render generated samples at arbitrary pixel resolution. A $C_4$-equivariant architecture incorporates rotational symmetry directly into the model, ensuring that rotations of the input noise produce corresponding rotations of the generated microstructure. We also demonstrate how training-free guidance can be used to generate complex microstructures, based on user defined objective function. In particular, we generate microstructures resembling a copper weld, cast metal slab, 3D-printed stainless steel and heterogeneous lamella titanium.