🤖 machine learning

SESaMo: Symmetry-Enforcing Stochastic Modulation for Normalizing Flows

This paper introduces SESaMo, a novel normalizing flow framework that utilizes stochastic modulation to incorporate symmetry inductive biases, enabling the effective learning of both exact and broken symmetries across diverse physical and statistical models.

Original authors: Janik Kreit, Dominic Schuh, Kim A. Nicoli, Lena Funcke

Published 2026-07-24

📖 6 min read🧠 Deep dive

CC BY 4.0

Original authors: Janik Kreit, Dominic Schuh, Kim A. Nicoli, Lena Funcke

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to predict the weather, but instead of looking at clouds, you are trying to figure out how a million tiny particles in a chemical reaction or a piece of metal are going to arrange themselves. In physics and chemistry, these particles don't just sit still; they dance around in a chaotic tangle, trying to find the most comfortable spot. Scientists call this "sampling from a Boltzmann distribution." It's a fancy way of saying, "What are all the possible ways these particles can arrange themselves, and how likely is each way?" The problem is that the number of possibilities is so huge it's like trying to find a single specific grain of sand on every beach on Earth at the same time.

To solve this, scientists use "generative models," which are like super-smart AI artists. These artists learn to draw the most likely arrangements of particles. A popular type of artist is called a "Normalizing Flow." Think of a Normalizing Flow as a magical, stretchy sheet of rubber. You start with a simple, boring pattern (like a perfect circle of dots), and the AI stretches and squishes this sheet until it looks exactly like the complex, messy pattern of the particles you want to study. The beauty of this method is that it can tell you exactly how likely any specific arrangement is, which is a superpower for scientists.

However, there's a catch. Many physical systems have "symmetries." This means if you flip the system upside down, or rotate it, or swap two identical particles, the physics doesn't change. It's like a snowflake: it looks the same if you rotate it by 60 degrees. If your AI artist doesn't know about these rules, it wastes time learning things it shouldn't have to, or worse, it might get confused and only learn one part of the picture, ignoring the rest. This is called "mode collapse," where the AI gets stuck in a rut and forgets that there are other equally valid ways the particles could be arranged.

The New Trick: The Symmetry-Enforcing Stochastic Modulator

In this paper, the authors introduce a new method called SESaMo (Symmetry-Enforcing Stochastic Modulation). They wanted to fix the problem of AI artists getting stuck in a rut when dealing with symmetrical systems, especially when those symmetries are "broken" (meaning the system almost looks the same when flipped, but not quite).

Here is how SESaMo works, using a playful analogy:

Imagine you are trying to paint a mural of a butterfly. The butterfly has two wings that are mirror images of each other.

The Old Way (Canonicalization): You force the artist to only paint the left wing. Then, you take a mirror and hold it up to the left wing to create the right wing. This works great if the butterfly is perfectly symmetrical. But what if the butterfly is sick, and the left wing is huge while the right wing is tiny? If you just mirror the left wing, your painting will be wrong. The old method struggles when the two sides aren't equal.
The SESaMo Way: Instead of forcing the artist to paint only one side, SESaMo gives the artist a "magic dice." First, the artist paints one wing (let's say the left one) perfectly. Then, the magic dice is rolled.
- If the dice lands on "Heads," the artist keeps the left wing as is.
- If the dice lands on "Tails," the artist flips the wing over to make the right wing.
- The Secret Sauce: The dice isn't fair. The AI learns how to load the dice. If the butterfly is perfectly symmetrical, the dice is fair (50/50). But if the butterfly is sick and one wing is bigger, the AI learns to load the dice so that "Tails" comes up more often, or less often, matching the real world.

This "stochastic modulation" (the loaded dice) allows the model to learn the exact shape of the butterfly, even if the wings are different sizes. It forces the AI to respect the symmetry rules (you can't just paint a random wing; it has to be a flip of the other) while still allowing the model to learn that one side might be more common than the other.

What They Found

The authors tested this new trick on several challenges, ranging from simple math puzzles to complex physics simulations.

The Math Puzzle: They started with a "Gaussian Mixture," which is like a picture with several distinct blobs of color. They made the blobs different sizes (breaking the symmetry). The old methods either got confused and only painted the biggest blob, or they tried to paint all of them but did a sloppy job. SESaMo, however, painted all the blobs perfectly, capturing their exact sizes and shapes.
The Physics Simulations: They moved on to real-world physics problems:
- The $\phi^4$ Theory: This models how particles interact in a field. They tested it with a "broken" symmetry where the particles prefer one state over another. SESaMo learned the correct distribution of particles much faster and more accurately than the old methods.
- The Hubbard Model: This is a famous model for how electrons move in materials, crucial for understanding things like superconductors. This is a very hard problem. The authors showed that SESaMo could handle a large grid of electrons (up to $18 \times 100$ sites) and correctly learn the broken symmetries, whereas other methods failed or crashed.

The Bottom Line

The paper suggests that SESaMo is a powerful new tool for scientists. It doesn't just force the AI to follow symmetry rules; it teaches the AI how to handle it when those rules are slightly broken. In their simulations, SESaMo consistently outperformed existing methods, achieving higher accuracy and converging (finishing the learning) faster. It successfully learned to capture all the different "modes" or possibilities of a system, avoiding the trap of getting stuck on just one. While the results are currently based on computer simulations and mathematical models, the authors show that this approach could be a game-changer for understanding complex systems in physics and chemistry, from how drugs fold to how new materials conduct electricity.

Technical Summary: SESAMO: Symmetry-Enforcing Stochastic Modulation for Normalizing Flows

Problem Statement
Sampling from unnormalized Boltzmann-like distributions, $p(x) = \exp(-f[x])/Z$ , is a fundamental challenge in physics, chemistry, and economics. Traditional Markov Chain Monte Carlo (MCMC) methods often suffer from slow convergence due to autocorrelations and high-energy barriers separating metastable states, leading to ergodicity issues. While deep generative models, particularly Normalizing Flows (NFs), offer a variational inference approach to approximate these distributions with closed-form likelihoods, they face significant training challenges. Specifically, standard NFs often struggle with "mode collapse," failing to capture all modes of a target distribution, particularly when the distribution exhibits symmetries or broken symmetries where probability mass is unevenly distributed across symmetric sectors. Furthermore, incorporating prior knowledge such as symmetries into NF architectures to accelerate convergence and ensure physical consistency remains an open challenge, especially for broken symmetries where standard equivariant approaches may be insufficient.

Methodology: Symmetry-Enforcing Stochastic Modulation (SESaMo)
The paper proposes SESaMo, a framework designed to incorporate both exact and broken symmetries (continuous and discrete) into flow-based generative models. Unlike previous approaches that rely on rigid architectural constraints or canonicalization, SESaMo introduces a novel stochastic modulation mechanism.

Core Mechanism: SESaMo operates by learning a flow $g_\theta$ that maps a simple prior distribution $q_0$ (e.g., a Gaussian) to a single mode (or a specific sector) of the target distribution within a defined "canonical cell" $\Omega$ . This alignment is enforced via a penalty term $\Lambda$ in the loss function, which ensures the flow does not map samples outside the canonical cell, thereby preserving bijectivity.
Stochastic Modulation: Once the flow captures a single mode, a stochastic modulation map $S_u$ $S_{u}$ is applied. This map is conditioned on a learnable random variable $u$ $u$ drawn from a distribution $p_{S,b}(u)$ $p_{S, b} (u)$ parameterized by $b$ $b$ . The map $S_u$ $S_{u}$ transforms the sample from the canonical cell to other symmetric regions of the configuration space.
- For exact symmetries, the modulation probability is uniform (e.g., $u \sim \text{Bernoulli}(0.5)$ for $Z_2$ ).
- For broken symmetries, the parameter $b$ is learned during training, allowing the model to infer the correct probability mass distribution across different modes (e.g., $u \sim \text{Bernoulli}(e^b)$ where $b \in \mathbb{R}^-$ ).
Training Objective: The model is trained by minimizing the reverse Kullback-Leibler (KL) divergence. The loss function includes the standard KL term, the unnormalized log-probability of the target, and the bijectivity penalty term. Because the modulation variable $u$ is discrete and stochastic, the REINFORCE estimator is employed to compute gradients with respect to the modulation parameters $b$ .
Bijectivity Enforcement: A key technical contribution is the introduction of a regularization term $\Lambda(\tilde{z}_c)$ (Eq. 10) to the loss function. This term penalizes samples that map outside the canonical cell $\Omega$ , ensuring the transformation remains bijective, a strict requirement for valid NFs.

Key Contributions

Novel Framework: The introduction of SESaMo, which enables the incorporation of arbitrary symmetries (both exact and broken) into normalizing flows via stochastic modulation rather than fixed architectural equivariance.
Handling Broken Symmetries: The ability to learn the probability mass distribution across symmetric modes during training, addressing the limitation of canonicalization methods which often assume equal mass or struggle with asymmetry.
Numerical Bijectivity: The proposal of a specific penalty term in the loss function to numerically enforce the bijectivity of the flow, preventing samples from escaping the canonical cell.
Extensive Benchmarking: Validation of the method on toy problems (Gaussian mixtures) and physically relevant lattice quantum field theories, including the complex $\phi^4$ theory and the Hubbard model.

Results
The authors conducted numerical experiments comparing SESaMo against Flow Annealed Importance Sampling Bootstrap (FAB), Variational Mixture of Normalizing Flows (VMoNF), and RealNVP with canonicalization.

Gaussian Mixture Model: SESaMo achieved an Effective Sample Size (ESS) of 0.9986 for exact $Z_8$ symmetry and 0.9947 for broken $Z_8$ symmetry, significantly outperforming VMoNF (which suffered from mode collapse) and canonicalization.
Complex $\phi^4$ Theory: For both exact and broken $U(1)$ symmetries on an $8 \times 8$ lattice, SESaMo achieved the highest ESS (0.9472 and 0.941, respectively), whereas FAB and VMoNF yielded ESS values around 0.22–0.28. Canonicalization was not applicable to continuous symmetries in this setting.
Hubbard Model: On a small volume ( $2 \times 1$ ), SESaMo achieved an ESS of 0.996 for broken $Z_4$ symmetry, outperforming canonicalization (0.839) and VMoNF (0.37). Crucially, SESaMo demonstrated scalability to large volumes ( $18 \times 100$ ), achieving an ESS of 0.74, while FAB was infeasible and canonicalization dropped to 0.024.
Convergence: In all benchmarks, SESaMo exhibited faster and more stable convergence compared to baseline methods, effectively mitigating mode collapse.

Significance and Claims
The paper claims that SESaMo provides a flexible and powerful tool for incorporating inductive biases (symmetries) into generative models, particularly in scenarios where target distributions exhibit broken symmetries or complex mode structures. By decoupling the learning of the flow (which captures a single mode) from the redistribution of probability mass (handled by the learnable stochastic modulation), the method overcomes fundamental challenges in variational inference such as mode collapse. The authors position SESaMo as a state-of-the-art approach for sampling lattice quantum field theories, capable of handling both discrete and continuous symmetries where previous methods (like canonicalization) face limitations. The work suggests that this approach is particularly valuable for physics and chemistry applications where symmetries are inherent but often broken, enabling more accurate and efficient sampling of complex Boltzmann distributions.

More like this