Enhancing lattice kinetic schemes for fluid dynamics with Lattice-Equivariant Neural Networks

TL;DR

This work introduces Lattice-Equivariant Neural Networks (LENNs) to enforce local lattice symmetries in neural surrogates for Lattice Boltzmann collision operators. By constraining layer weights according to the symmetry group of the lattice (e.g., $D_4$ in 2D and $O_h$ in 3D), LENNs achieve symmetry-preserving representations that match the accuracy of group-averaged approaches while delivering substantial speedups, especially in 3D. The study demonstrates, through a priori and a posteriori evaluations on laminar and turbulent flows, that LENNs offer improved parameter efficiency, training stability, and scalable performance over non-equivariant networks and competitive accuracy with reduced computational cost. The proposed framework generalizes beyond cubic lattices and holds promise for practical ML-augmented lattice CFD and stencil-based computations in higher dimensions and other domains.

Abstract

We present a new class of equivariant neural networks, hereby dubbed Lattice-Equivariant Neural Networks (LENNs), designed to satisfy local symmetries of a lattice structure. Our approach develops within a recently introduced framework aimed at learning neural network-based surrogate models Lattice Boltzmann collision operators. Whenever neural networks are employed to model physical systems, respecting symmetries and equivariance properties has been shown to be key for accuracy, numerical stability, and performance. Here, hinging on ideas from group representation theory, we define trainable layers whose algebraic structure is equivariant with respect to the symmetries of the lattice cell. Our method naturally allows for efficient implementations, both in terms of memory usage and computational costs, supporting scalable training/testing for lattices in two spatial dimensions and higher, as the size of symmetry group grows. We validate and test our approach considering 2D and 3D flowing dynamics, both in laminar and turbulent regimes. We compare with group averaged-based symmetric networks and with plain, non-symmetric, networks, showing how our approach unlocks the (a-posteriori) accuracy and training stability of the former models, and the train/inference speed of the latter networks (LENNs are about one order of magnitude faster than group-averaged networks in 3D). Our work opens towards practical utilization of machine learning-augmented Lattice Boltzmann CFD in real-world simulations.

Figures (16)

• Figure 1: Velocity stencils $\mathcal{V} = \{ \mathbf{c}_i \}$, with (a) the D2Q9 stencil and (b) the D3Q19 stencil.
• Figure 2: (a,b) Visualization of the action of the dihedral symmetry group $D_4$ and the octahedral symmetry group $O_h$ as actions on the velocity stencils D2Q9 and D3Q19. (a) The two $D_4$ generators: $\mathbf{R}$ (rotation by 90 degrees) and $\mathbf{S}$ (reflection). (b) Representation of all symmetry group actions in $D_4$, that can be obtained as composition of the two generators (8 elements). (c) The three $O_h$ generators: $\mathbf{R_1}$ (rotation by 90 degrees around the $x$ axis), $\mathbf{R_2}$ (rotation by 90 degrees around the $y$ axis), and $\mathbf{S}$ (reflection).
• Figure 3: Visualization of the notion of generalized population features in LENN. (a) the rotation transformation $\mathbf{R}$ acting on a single population vector, $\mathbf{f} = \{ f_i \}$, as permutation of populations (cf. Eq. \ref{['eq:representation']}). (b) the same action on the population features $\mathrm{x} = \{ \mathrm{x} _{i,c} \}$, where the permutation now acts identically on the $C$ different channels (cf. Eq. \ref{['eq:genpop']})
• Figure 4: Visualization of an equivariant layer, as outlined in Eq. \ref{['eq:equi']}. This layer operates with $C_{in}$ input channels and produces $C_{out}$ output channels. The input $\mathrm{x}$ and output $\mathcal{A}(\mathrm{x})$ have dimensions of $q \times C_{in}$ and $q \times C_{out}$, respectively, where $q$ is the number of populations. The linear matrix, with shape $q \times q \times C_{out} \times C_{in}$, can be conceptualized as comprising $C_{out}$ blocks, each contributing to one of the output channels of $\mathcal{A}(\mathrm{x})$. Within each block, $C_{in}$ matrices of size $q \times q$ act on the channels of $\mathrm{x}$. The submatrices of size $q \times q$ are constrained to conform to a specific shape, as depicted in Eq. \ref{['eq:mat-d2q9']} (for 2D) and \ref{['eq:mat-d3q19']} (for 3D).
• Figure 5: Comparison of a-priori error versus relative execution time for the non-equivariant MLP (blue), MLP with group averaging (GAVG) (red) and LENN (green), for 2D (a) and 3D (b). The a-priori error is measured as average Relative Mean Square Error (RMSE) on the post-collision populations on the test set (Eq. \ref{['eq:apriori']}). The execution times are evaluated at testing, and are normalized on the non-equivariant MLP value (the faster time). The barplots are computed considering 20 different architectures trained independently. The equivariant models (GAVG and LENN) show improved accuracy with respect to the non-equivariant MLP. GAVG shows a slight improved accuracy over LENN, especially in the 2D setting, not observed in the a-posteriori error (Fig. \ref{['fig:aposteriori']}). However, LENN is considerably faster than GAVG, especially in the 3D setting.