Accelerating Conjugate Gradient Solvers for Homogenization Problems with Unitary Neural Operators

Rapid and reliable solvers for parametric partial differential equations (PDEs) are needed in many scientific and engineering disciplines. For example, there is a growing demand for composites and architected materials with heterogeneous microstructures. Designing such materials and predicting their behavior in practical applications requires solving homogenization problems for a wide range of material parameters and microstructures. While classical numerical solvers offer reliable and accurate solutions supported by a solid theoretical foundation, their high computational costs and slow convergence remain limiting factors. As a result, scientific machine learning is emerging as a promising alternative. However, such approaches often lack guaranteed accuracy and physical consistency. This raises the question of whether it is possible to develop hybrid approaches that combine the advantages of both data-driven methods and classical solvers. To address this, we introduce UNO-CG, a hybrid solver that accelerates conjugate gradient (CG) solvers using specially designed machine-learned preconditioners, while ensuring convergence by construction. As a preconditioner, we propose Unitary Neural Operators as a modification of Fourier Neural Operators. Our method can be interpreted as a data-driven discovery of Green’s functions, which are then used to accelerate iterative solvers. We evaluate UNO-CG on various homogenization problems involving heterogeneous microstructures and millions of degrees of freedom. Our results demonstrate that UNO-CG enables a substantial reduction in the number of iterations and is competitive with handcrafted preconditioners for homogenization problems that involve expert knowledge. Moreover, UNO-CG maintains strong performance across a variety of boundary conditions, where many specialized solvers are not applicable, highlighting its versatility and robustness.

💡 Research Summary

The paper introduces UNO‑CG, a hybrid solver that accelerates the conjugate gradient (CG) method for parametric partial differential equations (PDEs) arising in homogenization of heterogeneous microstructures. Classical CG offers guaranteed convergence for symmetric positive‑definite (SPD) systems but can suffer from slow convergence when material contrast is high or the discretization is fine. Preconditioners are essential to mitigate this, yet traditional preconditioners either rely on expert‑crafted Fourier‑based schemes (e.g., FANS) or algebraic multigrid, both of which require substantial domain knowledge and may not generalize across boundary conditions.

UNO‑CG addresses these limitations by learning a data‑driven preconditioner that is provably SPD, thus preserving CG’s convergence guarantees. The authors modify the Fourier Neural Operator (FNO) to enforce a unitary constraint in the spectral domain, creating the Unitary Neural Operator (UNO). This constraint ensures that the learned operator is symmetric and positive‑definite for any set of parameters, effectively making it a learned approximation of the Green’s function of the underlying PDE.

Training data are generated by solving a wide variety of microstructure‑parameter combinations with high‑fidelity finite‑element or existing FFT‑based solvers. Each sample consists of material tensors, microstructure images, and loading conditions as inputs, and the corresponding optimal preconditioning matrix in Fourier space as the target. The loss function combines a reconstruction term that forces the product of the learned preconditioner and its transpose to approximate the identity, a regularization term that enforces symmetry and positive‑definiteness, and a unitary penalty that maintains the spectral constraint.

Extensive numerical experiments cover two‑ and three‑dimensional homogenization problems with millions of degrees of freedom, including high‑contrast metal‑matrix composites and a range of boundary conditions (periodic, Dirichlet, Neumann). UNO‑CG consistently reduces the number of CG iterations by 30 %–70 % compared with unpreconditioned CG and matches or outperforms state‑of‑the‑art FFT‑based preconditioners. Notably, the learned preconditioner generalizes well to unseen microstructures and parameter regimes, eliminating the need for retraining when new designs are introduced.

Key contributions are: (1) a principled way to embed SPD requirements into a neural operator, guaranteeing CG convergence; (2) the introduction of the Unitary Neural Operator as a physically interpretable, Green’s‑function‑like preconditioner; (3) demonstration that the hybrid approach rivals handcrafted preconditioners without requiring expert insight; and (4) validation on large‑scale 3‑D problems, showing practical applicability to industrial multiscale simulations.

The work establishes a new paradigm that blends scientific machine learning with classical numerical analysis, offering a scalable, robust, and theoretically sound tool for rapid many‑query simulations in materials design, topology optimization, and other fields that involve parametric PDEs.