Discovering Lie Groups with Flow Matching

TL;DR

This work introduces LieFlow, a flow matching framework on Lie groups to automatically discover data symmetries by learning a distribution over a large hypothesis group $G$ and extracting the stabilizer subgroup $H$ that preserves the data distribution. By operating directly on group manifolds and using exponential curves along orbits, LieFlow captures both continuous and discrete symmetries, including reflections via complex-domain flows. The authors identify a 'last-minute mode convergence' phenomenon and propose a time-scheduling strategy to uncover more complex groups, achieving state-of-the-art performance on 2D and 3D symmetry discovery tasks (e.g., discovering $\mathrm{C}_4$, $\mathrm{D}_4$, $\mathrm{Tet}$, and $\mathrm{SO}(2)$ elements) and outperforming the LieGAN baseline in Wasserstein-1 distance. Limitations include challenges in scaling to higher-order or more symmetric groups and handling real-world noisy data, with future work pointing toward larger Lie groups, non-compact cases, and diffusion-based Riemannian formulations to better integrate discovered symmetries into downstream equivariant models.

Abstract

Symmetry is fundamental to understanding physical systems, and at the same time, can improve performance and sample efficiency in machine learning. Both pursuits require knowledge of the underlying symmetries in data. To address this, we propose learning symmetries directly from data via flow matching on Lie groups. We formulate symmetry discovery as learning a distribution over a larger hypothesis group, such that the learned distribution matches the symmetries observed in data. Relative to previous works, our method, \lieflow, is more flexible in terms of the types of groups it can discover and requires fewer assumptions. Experiments on 2D and 3D point clouds demonstrate the successful discovery of discrete groups, including reflections by flow matching over the complex domain. We identify a key challenge where the symmetric arrangement of the target modes causes ``last-minute convergence,'' where samples remain stationary until relatively late in the flow, and introduce a novel interpolation scheme for flow matching for symmetry discovery.

Figures (33)

• Figure 1: Flow matching for symmetry discovery
• Figure 2: 2D Datasets: Generated data samples. The samples match closely to the original dataset symmetries of $C_4$ or $D_4$.
• Figure 3: 2D Datasets: Visualization of trajectories from $t=0$ (left) to $t=1$ (right) of the centroids of the transformed objects over $100$ samples, and a histogram of learned angles for $\mathrm{SO}(2)$ to $C_4$.
• Figure 4: Time progression of $x_t$ when generating $5$ samples over $20$ steps. The gray arrow shows the original $x_1$ and the color represents the determinant of the generated transformation matrix.
• Figure 5: 3D Datasets: Visualization of $5{,}000$ generated elements of $\mathrm{SO}(3)$ by converting them to Euler angles. The first two angles are represented spatially on the sphere using Mollweide projection and the color represents the third angle. The elements are canonicalized by the original random transformation and the ground truth elements of the target group are shown in circles with black borders. The points are jittered with uniformly random noise to prevent overlapping.