SESaMo: Symmetry-Enforcing Stochastic Modulation for Normalizing Flows
Authors: Janik Kreit
Dominic Schuh
Kim Andrea Nicoli
Lena Funcke
Venue: ICLR 2026 Poster (2026)
Keywords: Generative models
Normalizing Flows
Symmetries
Equivariant Models
Primary area: probabilistic methods (Bayesian methods, variational inference, sampling, UQ, etc.)
TL;DR: SESaMo provides an extension to Normalizing Flows that enforces (broken) symmetries to the output distribution

Abstract: Deep generative models have recently garnered significant attention across various fields, from physics to chemistry, where sampling from unnormalized Boltzmann-like distributions represents a fundamental challenge. In particular, autoregressive models and normalizing flows have become prominent due to their appealing ability to yield closed-form probability densities. Moreover, it is well-established that incorporating prior knowledge—such as symmetries—into deep neural networks can substantially improve training performances. In this context, recent advances have focused on developing symmetry-equivariant generative models, achieving remarkable results.
Building upon these foundations, this paper introduces Symmetry-Enforcing Stochastic Modulation (SESaMo). Similar to equivariant normalizing flows, SESaMo enables the incorporation of inductive biases (e.g., symmetries) into normalizing flows through a novel technique called \textit{stochastic modulation}. This approach enhances the flexibility of the generative model by enforcing exact symmetries while, for the first time, enabling the model to learn broken symmetries during training.
Our numerical experiments benchmark SESaMo in different scenarios, including an 8-Gaussian mixture model and physically relevant field theories, such as the $\phi^4$ theory and the Hubbard model.

Published as a conference paper at ICLR 2026
SESAMo: SYMMETRY-ENFORCING STOCHASTIC MODULATION FOR NORMALIZING FLOWS
Janik Kreit Dominic Schuh 1.2- Kim A. Nicoli!.2.3+
Lena Funcke!
Transdisciplinary Research Area (TRA) Matter, University of Bonn, Germany
HHelmholtz Institute for Radiation and Nuclear Physies (HISKP)) University of Bonn, Germany 3 Oldendorff Carriers GmbH & Co. KG, Lübeek, Germany
ABSTRACT
Deep generative models have recently garnered significant attention across various fields, from physics to chemistry, where sampling from unnormalized Boltzmann- like distributions represents a fundamental challenge. particular, autoregressive models and normalizing flows have become prominent due to their appealing ability to yield closed-form probability densities. Moreover, incorporating prior inowledge-sure such as symmetries into deep neural networks can substantially improve training performances, this context, recent advances have focused developing symmetry -euueetr--quivariann generative models, achiev- ing remarkable results. Building upon these foundations, this paper introduces Symmetry- Symmeety-Enforrcnng Stochastic Modulation (SESaMo). Similar to equivariant normalizing flows, SESaMo enables the incorporation of inductive biases (c.g.. symmetries) into normalizing flows through novel technique called stochastic modulation. This approach enhances the flexibility the generative model by enforcing exact symmetries while, for the first time, enabling the model to learn bro- ken symmetries during training. Our numerical experiments benchmark SESaMo in different scenarios, including an 8- Gaussiian mixture model and physically relevant field theories, such as the o² theory and the Hubbard model.
well-established that
INTRODUCTION
Sampling from unnormalized Boltzmann distributions is an ubiquitous yet challenging task across various fields, including physies 1. chemistry (2), and economics (3). These distributions are typically of the form p(x:)= cxp (-f(=))/Z where f[-] is functional representing, for example, the potential chemical compound or the action of physical system, while Z, the normalization constant (or partition function), is often unknown. While f(-] is usually available in closed form, as it describes the microscopic dynamics of the system under study, computing would require solving a functional or high-dimensional integral, which is generally intractable. In fact, for many systems of interest, sampling from Boltzmann distributions has been proven to be NP-hard (4), making it highly unlikely that a polynomial time algorithm exists for this problem. Due to this complexity, sampling from unnormalized Boltzmann distributions traditionally performed using Markov Chain Monte Carlo (MCMC) methods (5), where randomly initialized Markov chain guaranteed to converge to the target distribution. Despite numerous advanced MCMC techniques, significant challenges remain. In chemical and biological systems, for instance, sampling can hindered by high-energy barriers separating metastable states, posing a major obstacle for tasks such as protein folding (6). In physics, MCMC methods often suffer from slow convergence due autocorrelations between samples, necessitating longer simulations to obtain statistically independent samples and thereby increasing computational costs
Over the past decade, deep generative models (8) have achieved remarkable success in sampling from Boltzmann distributions within the framework variational inference (VT) (9). In particular, Ref. (10)
introduced Boltzmann Generators (BGs), an approach in which a variational (parametrized) probabil-
*Correspondence to Streit@uni-honn.de schuh hinb@hhikpp.ui-boon.de kim.a.nicoli aaaiiool@@mmml.coom and
fuuncke@ @uni-honn.de

---

Published as a conference paper at ICLR 2026
ity density de E Q is learned, using a generative model, to approximate the target distribution I of chemical system, i.e., 90 2 p. Around the same time, concurrent studies proposed similar ideas in the contexts of statistical physics (2 and lattice quantum field theories (33 (14). A distinctive feature of BGs is that they rely on generative models capable of providing the learned variational density in closed form. These include autoregressive neural networks (15) 16) and normalizing flows (NFs) (7718) 18), which are particularly suited for sampling discrete and continuous degrees of freedom, respectively. In the remainder of this work, we primarily focus on NFs, although extensions to other generative models that allow exact likelihood computation are also possible.
Despite their potential overcome some limitations of traditional MCMC sampling, deep generative models present challenges of their own. In particular, to ensure reliable sampling from the target density with suitable asymptotic guarantees (12), these models must first trained to sufficiently high standard. Deep generative models, such as NEs, are parametrized by deep neural networks with numerous trainable parameters, which may require substantial computational effort (training) to converge. To accelerate training, it has been shown that incorporating inductive biases, such as symmetry constraints, into the model architecture can lead to faster and more robust convergence. A seminal example
convolutional neural networks (CNNs), which exhibit built-in translational equivariance (19). This concept has been generalized to arbitrary symmetry groups and manifolds (20) 21 (2).. Similar ideas have been extensively leveraged in scientific applications, where chemical and physical systems are often rich in symmetries (23 24: 25 (6)).
In this paper, we propose general framework for embedding arbitrary symmetries into the training protocol of NFs, which we term Symmetry-Enforcing Stochastic Modulation (SESaMo). Our approach leverages the prior knowledge (symmetries) from the unnormalized log probability train a NF and uses an independent random variable to infer the correct probability mass for each mode of the target distribution. Crucially, this approach holds promise for mitigating and potentially overcoming -ht fundamental challenge of mode collapse in variational inference (87. We focus on setting in which no samples from the target distribution are available, and mode collapse is particularly prominent. However, we emphasize that our method is not limited to the data-free case. In summary, the contributions of this work are threefold:
We propose Cemmmetry-Enforeinnn Stochastic Modulation (SESaMo), novel approach to incorporate continuous and discrete symmetries (broken and exact) into flow-based models.
We numerically enforce bijectivity by introducing a penalty term in the loss function.
We conduct extensive numerical experiments to validate our theory on both toy problems and real-world benchmarks for lattice quantum field theories.
The remainder of this paper is organized as follows. In Sec. we introduce the necessary background on NFs and variational inference. We also discuss how symmetries can be incorporated into flow- based models and establish the notation used throughout the manuscript. In Sec. 3 we present our stochastic modulation approach. Finally, in Sec. 4 we validate our approach, both on standard benchmark and on tasks of practical relevance, such as sampling lattice quantum field theories, including the o² theory and the Hubbard model. In Sec. 5 we discuss the limitations of our algorithm, and we conclude by summarizing our findings and outlining potential directions for future work in Sec. d
1.1 RELATED WORK
The field of geometric deep learning 22 which investigates the mathematical foundations of deep learning on geometric structures particularly group-equivariant and gauge-equivarte neural networks- has advanced significantly in recent years. For comprehensive review of common
methodologies, we refer to Refs. 30 31.
Köhler et al. 32) proposed way to build NFs that are equivariant under the symmetries of the target p. ensuring that the variational distribution de inherently respects these symmetries, thereby improving both accuracy and efficiency. This work laid the foundation for the development of equivariant NFs across various applications. Satorras et al. (33) proposed a generative model equivariant to Euclidean
For notational convenience, we use the same symbol for distribution and its density with respect to the Lehesgue measure.
2

---

Published as a conference paper at ICLR 2026
symmetries, integrating E(n)- Eqqiivaarar Graph Neural Networks (EGNNs) 25 within a continuous NF framework (34), yielding an invertible map that preserves Euclidean invariances. Bose et al. (35) addressed the general problem of constructing equivariant diffeomorphisms with an equivariant finite NF, specifically targeting finite symmetry groups and compact spaces. In high-energy physics, Kanwar et al. (36) and Boyda et al. (37) adapted NFs respect Abelian and non-Abelian gauge symmetries, respectively. In condensed matter physics, Schuh et al. (38) demonstrated the importance of enforcing equivariance NFs for symmetry-rich systems like the Hubbard model, showing that equivariance is crucial for accurately learning the target density and overcoming ergodicity issues. For atomistic systems (39) and atomic solids (40), Wirnsberger al. introduced NFs equipped with permutation-equivariiie diffeomorphisms. More recently. Midgley et al. (41) introduced NF's that inherently respects SE(3) group symmetries-compriane
translations, rotations, and reflections- well as permutation invariance. Furthermore, Klein et al. (42) proposed equivariant flow matching, a training objective based on optimal transport flow matching that leverages inherent symmetrics in physical systems, enabling simulation free training of equivariant continuous normalizing flows (CNFs). In the context of diffusion models (43), Hoogeboom et al. (44) introduced an E(3)-equivariant diffusion model for 3D molecular generation, which, similar (3), enforces Euclidean invariance under translations and rotations. Lastly, Pires et al. (45) showed that a Variational Mixture of NFs (VMoNF) can model the symmetries of labelled data.
2
PRELIMINARIES
2.1 NORMALIZING FLOWS
Normalizing flows (NFs) (8) are class of generative models that provide an effective framework for approximating complicated probability distributions. Commonly employed in the context of vari- ational inference (VI) (46), NFs operate by transforming simple, well-understood, prior distribution (typically Gaussian) into target distribution through a sequence of invertible and differentiable mappings. A key advantage of NFs is their ability to efficiently sample from approximated high- dimensional distributions while retaining the capability to compute exact likelihoods. This exact likelihood computation distinguishes NFs from many other generative models, making them particu- larly well-suited for learning probability distributions in scientific applications, such as chemistry (IO) and physics (14). NFs can be categorized based on how the mappings between the prior density and the target distribution are constructed. These categories include coupling-based NF' (47) 48 49).
autoregressive NF's (50), and continuous NFs (34). For the sake of simplicity, this paper primarily focuses on coupling- based NFs, although extensions to other types of NFs are possible.
At the heart of NFs lies the concept of a bijective transformation that maps samples from prior distribution # (o(2) (such as a multivariate Gaussian) to samples from a variational distribution ~ go. which is meant to approximate target This typically happens by means of a learnable function
(1) where the transformation ge parametrized by a neural network. To increase the flexibility of NFs, multiple transformations (coupling blocks) can be composed, allowing for more expressive mappings between prior and target distributions, g0 = g0, o g0 o... o go, (=). A key feature of NFs is that the transformation must be invertible, allowing the likelihood of the target distribution to be computed exactly using the change of variables formula
90 To + x 2o (xx)
qo () = go (9 (x=) det (gg=
(xx) (x)
dx
For comprehensive overview of NFs, refer to the review papers (1)).. In this work, we focus mainly on affine NF architectures, such as RealNVP (48) and NICE (47)
2.2 THE KULLBACK- LEIBLER DIVERGENCE
In the context of Variational Inference, the parameters 0 of NFs are trained by minimizing the so-called (reverse) Kullback-eibler (KL) divergence (52)
In f(xz)
KL ||) -E
In Z,
(3)
90 (x(x))
3

---

Published as a conference paper at ICLR 2026
Qo
q.
q0
Cr
90
Ti
&
Visualization of the canonicalization approach making a flow-based model equivariant with respect to a 2 symmetry.
Figure 1:
where f(æ=) exp distributions, respectively. The logarithm of the unknown partition function simply appears as an additive term, which vanishes upon taking the gradient. It should also be noted that the KL divergence is not symmetric, i.e., KL( p) * KL(p) ll go). Consequently, training using Eq. (3) differs from the practice of maximum likelihood training. which employs the forward KL-divergence a common approach in, e.g., computer vision applications (50). This distinction is significant: In Variational Inference, access to training data is often unavailable, and models must trained solely using the closed-form unnormalized log probability D. which is the focus of this work.
(-)(())) f|=]) and (x)) are the unnormalized target and the parametrized probability
2.3 EQUIVARIANT NORMALIZING FLOWS
In previous works, several attempts have been made to incorporate prior knowledge into NFs and make them equivariant with respect to certain symmetry groups. The main result stemming from (22 is summarised in the following theorem:
Theorem (Köhler et al., (2020)) Let's assume H a group acting on IR,", qo the base density of a flow-based transformation with qe being the transformed density under the diffeomorphism go TR" - TR... If go is an H H-equivariant diffeomorphism and qo an H Hiimvariant density with respect to the same group H, then qo is also an H --iivariant density on R..
Specifically, this theorem provides a general protocol to build an equivariant NF by choosing an appropriate invertible map So that is H--qquivariant. However, despite the generality of this result, defining equivariant diffeomorphisms that allow for tractable inverses and Jacobians both essential for building Ne-remains open challenge. Indeed, different approaches have been leveraged in recent works to build equivariant flow-based models.
2.3.1 EQUIVARIANT NEURAL NETWORKS
In coupling-bant NFs, the diffeomorphism go often parametrized by neural network (NN). straightforward approach to enforce equivariance (or invariance) (53) (54 is to design NN that explicitly satisfies these symmetry requirements. However, a significant limitation of this method is that constructing such constrained architectures is neither always possible straightforward, One instance where this approach is feasible in the case of a Zo symmetry. Indeed, recent work showed how to build manifestly sign-equivariant architectures (55). For example, simple strategy to achieve sign equivariance in NNs is to use equivariant activation functions, such as tanh, and omit bias terms, ensuring that the resulting NN remains equivariant. Indeed, this approach was successfully applied for training Zg Eeqiivrrai NFs in the context of lattice quantum field theories (14) 56 57))
2.3.2 CANONICALIZATTIN
The idea canonicalization, largely motivated by Theorem context of flow-based sampling for lattice field theories (36). Indeed, physical systems are rich in global (and local) symmetries, and being able to develop equivariant flows fulfilling these constraints very active area of research. The key idea is to a transformation CT map sample from the base density to a so-called canonical cell Ω see (37). The NF then transforms the canonicalized sample, before the inverse C is applied to map the sample back its original space. We refer to Fig. Find App. B for more details. A parametric map go equivariant generic transformation
! has been widely explored in the
90 90 (Tx) Tgo(x:) (0) T-1 T (Tx).
4

---

Published as a conference paper at ICLR 2026
For example, for the sign-flipping Zg transformation mentioned above, the transformation reads
T2, + -
(5)
canonical map Cr, transforms samples a E R", where 2 ~ qu, to the canonical f
Cr R" Ze Ω
with the inverse
CT Ze Ω x
(6)
The two manifolds 52 and 92 are connected by the diffeomorphism go acting in the canonical space, i.e.,
90 : 2e EQ _ s2
and
2 E 5 2c Ω
Note that Cr. depends some specific symmetry transformation T, see Eq. ( which makes the canonicalized flow [=( CT 90 (z,z equivariant. Focusing on the sign-fliping Z2 transformation mentioned above, we have
{ --, else
if E
CzZ
(8)
where the canonical cell in this case is = {z E R s.t. Ei= 4 2 0}. See Fig. for a visual intuition. This approach can be generalised for In and a set of M symmetry transformations (Tt) such that
2,
if e Ω
a. if Ω
Tiz,
elif 2 E A
Ti Ø, elif A
OT, 12
with inverse
T
Tuz
elif = E A&
TM'x, x elif z E AM
where A; denote different regions of the target space, for example corresponding to different modes of the target distribution. Note that the inverse transformation CT. depends on the input z, i.e. the information about the origin of the sample in the base space must be retained in the transformation. proof that the canonicalization approach is equivariant is given in App.
2.3.3 CONSTRAINTS CANONICALIZATION
In order to enforce equivariance via canonicalization, two constraints must be met: first the prior distribution do must be invariant under any symmetry transformation T; (32), i.e., (o(2) )= (T(
Second, Ge should not map samples outside the canonical cell, i.c., ã. = go 9o z) Ω. While the former constraint this latter constraint by introducing a regularization term
be readily verified, the latter may not hold for any general NE. We enforce
A(ž.)
(B X(2.) ((X(E.))
(10)
where A(ž.) is penalty function being zero for a general input ă. at the boundary An of the canonical cell Q negative for 2. E Ω2 and positive for 2 $ Ω. The Heaviside step function O(-) ensures that the penalty term zero for že E Ω2 while the sigmoid function a(-) ensures that the penalty function has gradient pointing toward the canonical cell Ω. The hyperparameters A, B e R are used to scale the amplitude and the gradient the function, respectively. While needs match at least the order of magnitude of the loss function, B must be chosen sufficiently small to ensure non-vanishing gradients for 2. & Ω. This regularization term is added to Eq. (3) during the training of NF, We provide further details about the penalty term in App. D
3 PROPOSED METHOD: SESAMO
Crucially, certain symmetries may be difficult to incorporate through naive canonicalization strategies and are unlikely to be effectively captured by standard flow-based generative models. A representative case is dineennnnio multimodal distribution with modes of unequal probability mass (see Stochastic Modulation (SESaMo),
App. and App. H). Our proposed method, Symmetry-Enfonn introduces novel stochastic modulation mechanism that is described in detail in Sec. 3.1
5

---

Published as a conference paper at ICLR 2026
q0
g0
90
96
S
Figure 2: Visualization of the stochastic modulation approach for enforcing a Z2 symmetry in flow-based model.
3.1 STOCHASTIC MODULATION
Stochastic modulation involves drawing samples 5 from a flow-based sampler with density ()())) Prior samples
go (left panel Fig. [ are shifted by the flow ge into the canonical cell f2 to align with one of the modes, producing a go with æ E S (center panel). This shift is induced by the penalty term A (()) Eq. (10). The samples 7 are then transformed according the stochastic modulation S. (right panel), which is conditioned on random variable u, resulting in output density
(0,6(x) = Σ Ps,(uu) qu (Te, det
(11)
Ox
Here, ps. the modulation probability that depends on learnable parameters b, whose number is fixed by the symmetry type. The diffeomorphic map from the base density go(z) to the final density reads
(z) Su(o)(z)),
(12)
such that
det (0go,u Oz
(8S,)
(0go
det
det
(13)
Age Oz
general stochastic modulation ST,, .u for set of transformations (Ti) reads
2, if
Tix, elif u =
ST. x m
with W ~ PS.b
(14)
TMÆ, elif u M
where the transformations T. map samples space, potentially corresponding different modes of the target distribution p(z). We note that
(z( to distinct regions in the configuration
T. # T, vi * and that Tix & Ω, in order to avoid any overlap between distributions and thus preserve bijectivity.
From Eqs. as
11
13
u (u) is sampled and not marginalized, we can write the log probability
as,
0ge
In qo (ãe. (2)) In (u) In go()) In det
In det
(15)
Ago
Oz
where _o (a:) and (u) is the probability of sampling the random variable To better understand this mechanism, let again consider target density with Z₂ symmetry. In this specific case, we define random variable u E {0, 13 that follows a Bernoulli distribution B(c") and
if x
5, : x
with
u~ Ps.6 = B(e²)
and
b b= 0.5..
(16)
x
if u
For broken Ze symmetry on the other hand, the modulation probability is given in terms of learable parameter E IR and reads
11-e" e if u =
PS.6 =
e if u 1.
(17)

---

Published as a conference paper at ICLR 2026
So
ST.
Target:
Prior: Q0
T go
2 ()())
CT
90
Ta
Target:
Prior: qo
~v aic
Qz.
I ()(x)
Figure Illustaatoon of Symmetry-Enforeee Stochastic Modulation (SESaMo) (top row) and canonicalization (bottom row), shown for example target distribution and corresponding prior.
Unlike canonicalization, SESaMo (illustrated in the upper row of Fig. 3 employs a flow ge that shifts the prior density do so that it captures a single mode of the target density within the canonical cell. This shift is induced by the penalty term (@r). After this alignment step, the probability mass redistributed according to Sy- In contrast, canonicalization (shown the lower row Fig. 3 applies fixed transformation CT, that maps directly into the canonical cell. Consequently, unlike canonicalization, SESaMo does not require the prior go to be invariant under the symmetry transformation, which makes it applicable in broader range of settings.
Similarly to canonicalization, stochastic modulation requires the map S. to be bijective. This property is enforeed by the penalty term introduced in Sec. 2.3.3 An extended comparison between canonicalization and stochastic modulation is provided Appl A pseudocode representation of SESaMo can be found in Sec. A
the probability mass is not evenly distributed among the modes of the target density (e.g.. b b In 0.5 for &2 symmetry), having a learnable parameter b allows SESaMo effectively capture the broken symmetry. This case is further detailed in App. E We also stress that SESaMo extends beyond discrete symmetries. Appendix [ develops formulation for (broken) continuous symmetries, where is continuous variable instead of discrete one.
3.2 LOSS FUNCTION AND REINFORCE ESTIMATOR
For training, the standard reverse KL divergence extended by the penalty term A(-)
kg Ipp) Be ge (ū)..())) + f[)o,.(=)] A(go(z))]
(18)
stemming from Sec. 2.3.3 Since the parameter b of the stochastic modulation is only present in the modulation probability Ps,b the REINFORCE estimator (58) is used to enable gradient computation through the random variable
R Pu (In (ooo. + [[ Euppss (Inou(u.))))) + (z)) In ps, ab
0b
(19)
NUMERICAL EXPERIMENTS
In this section, we present numerical experiments that compare the performance of four approaches: Flow Annealed Importance Sampling Bootstrap (FAB) (59), RealNVP with Variational Mixture of Normalizing Flows (VMoNF) (45). RealNVP with canonicalization (37), and RealNVP with stochastic modulation (SESaMo). We benchmark these approaches both toy problems and

---

Published as a conference paper at ICLR 2026
Ground truth
VMoNF
Canonicalization
SESaMo (ours)
10
-10
-10
10
0 X1
10
-00
-10
0
10
X1
A1
A1
Figure 4: Gaussian mixture target density (broken Zs symmetry). All flow-hased models trained until convergence. From left to right we show: the ground truth, VMoNF, canonicalization, and SESaMo (ours). We refer App. forr more details on the experiments.
physically relevant tasks. To evaluate the effectiveness each approach, we use the the KL divergence (if the normalization Z is known) and the effective sample size (ESS) (12), performance metric. The ESS quantifies the accuracy with which the approximation do matches the target probability distribution Bounded between zero and one, the ESS reaches its optimal value (ESS 1) when the approximation is exact (gg = p). Although the ESS not a good metric if the distribution lacks full mode coverage, the symmetry-enforcitt the ESS a valid metric to quantify SESaMo's performance.
architecture guarantees to capture all modes, making
The code used these experiments builds earlier release of (60). A live version is available at
and static, archived snapshot is provided Ref. (61)
4.1 Toy EXAMPLE: GAUSSIAN MIXTURE
We initially consider probahility distribution in two dimensions whose density is given by
N Σ 1=1
(x
m a(x1 + x22
2T
T
p(x)= = 2
exp
th = R
cos
sin
2
N
N
(20)
where Z 2T Ex=1 (a² aRR2ssinn - sinn(e 2mk + 5)) is the normalization, E N+ is the number
of Gaussians, R R+ is the radius of the circle around which they are located, and a E E breaks the ΣN symmetry of this model. In this study, which results a (broken) Zg symmetry. In Tab. D we report the ESS and KL divergence achieved after convergence, and visualize the corresponding target density in Fig. 4 Even though VMoNF has the ability that different NFs learn different sectors of the target distribution, collapses to the three most likely modes in the lower left. This is not issue of the architecture itself but the mode collapsing behaviour of the reverse KL that is used in this work. Overall, SESaMo achieves the best performance, outperforming the other baselines and yielding higher accuracy. For more details refer to App.
N = 8, R. 12, and = 0 (a 0.05),
4.2 PHYSICS EXAMPLE: LATTICE QUANTUM FIELD THEORY
Sampling using NFs has become ubiquitous across various fields of physics, yielding particularly notable results for sampling lattice quantum chromodynamies (62). scalar lattice quantum field theories (14 63), and condensed matter systems (38). We refer to (64 for comprehensive overview. In what follows, primarily focus on two pertinent benchmarks: the complex ø² theory and the Hubbard model. direct readers secking further technical details regarding the physics to App. G
In lattice quantum field theory, the probability distribution of system is given by Boltzmann-like density p(z)= fxx(=f(x))/ where f[z] is a functional known as the action, Z is an unknown partition function, and denotes the lattice fields. Note that, discussed in Sec. 3 for the following experiments we optimize the symmetry breaking parameter during training, which, as shown in App. H perfectly agrees with the analytical prediction.
8

---

Published as a conference paper at ICLR 2026
The complex 6² scalar field theory in two dimensions The complex 6² theory offers simple yet versatile framework for investigating interacting scalar fields. It plays a crucial role in understanding spontaneous symmetry breaking (including the Higgs mechanism) and critical phenomena (65), while providing a key testbed for the machine learning community to develop theoretical techniques and We consider the action with quartic interactions,
numerical methods (66: 67 68).
f|z]
Σ jEV
2k Σ A=1
(2)0)+H)+ + 22)=² + Aa² + aRe(z,,
(21)
where 2: = 21 + i are the complex scalar fields, the subscript j labels the lattice sites in the two-dimensional lattice volume V N₂ N₂ 8 8, the and X are the couplings of the theory, and fi denotes the interactions between nearest neighbours. The term a Re(æ) introduces an additional component designed to break the continuous (11 symmetry of the theory, thereby increasing the complexity of the learning task P We emphasize that while prior studies have often focused on real scalar fields, physical fields are complex-valued. Therefore, we compare SESaMo with FAB and VMoNF when sampling æ E C Canonicalization could not be applied here. The ESS obtained by each model is detailed in Tab. [ for both broken (a * 0) and unbroken (a 0) (1) symmetry. Across both conditions. SESaMo achieved the highest ESS, indicating its superior ability to incorporate the underlying physical symmetries into the flow model. Additional results, including the density plots, are available in App. H Moreover, App. H also demonstrates how SESaMo outperforms the baselines of ReaINVP and canonicalization in the case of real scalar field theory.
The Hubbard model in two dimensions The Hubbard model a cornerstone of condensed matter physics, providing fundamental description of interacting electrons a lattice and playing pivotal role in studying phenomena such as magnetism, metal-insulator transitions, and high-temperature superconductivity (70). For numerical experiments, we adopt the setup detailed in (38 with the action-feature a broken Z4 symmetry- cerinettry-vvennn by
[|x]
Σ x² jEV
log det M [ax] - log det M[-
(22)
20
where the coupling U describes the interaction strength, M-] is the fermion matrix describing the interacting fermions (particles), x are auxiliary bosonic fields, and the subscript j labels the lattice sites in the lattice volume V N₂ No. While for the GMM and the complex 6² theory we focused on small lattice volumes as proif-lf-principle demonstration, here we include both a small volume of V = x 1 to compare with analytical solutions and large volume of V 18 x 100 to demonstrate and (38) for more details about the model. For learning the Boltzmann distribution, we agann compare FAB. VMoNF, canonicalization, and SESaMo. For the smaller volume of V == 2 x 1, the ESS and KI. divergence are reported in Tab. [ and the resulting probability density after training shown in Fig. [ For the larger volume of V = 18 100, the ESS is also listed in Tab. [ whereas the KL divergence cannot be reported due to the unknown normalization of the target distribution. FAB is not feasible at this lattice volume. As before, SESaMo achieves the highest ESS and exhibits faster and more stable convergence compared to the other baselines. For further results and density plots illustrating that SESaMo mitigates mode collapse 77). we refer the reader to App. H While Schuh et al. (38) first demonstrated the application of NFs to the Hubbard model using canonicalization and small lattice volumes of V 2 x 2, SESaMo with the objective in Eq. (18) crucially achieves higher ESS, perfectly learns the broken Z, symmetry, and scales to lattice volumes of up to V = 18 x 100 (broken (Zz)16 symmetry), thereby establishing new state-of-the-ard
the scalability of our approach. We refer to Apps.
LIMITATIONS
A primary limitation of SESaMo stems from the requirement that the symmetry sectors must be known a priori to apply the stochastic modulation. Nevertheless, for applications in physics and
This system can serve as proxy to describe quantum field theory with two flavors of differing masses (69).
'hh reverse KI. divergence can only be computed if the normalization of the probability distribution is known, which is the case for the GMM and the Hubbard model with V x 1.

---

Published as a conference paper at ICLR 2026
Ground truth
VMoNF
Canonicalization
SESaMo (ours)
20
-20
-20
20
-20
0
20
-20
20
-20
0
20
X1
X1
A1
A1
Figure 5: Density for the Hubbard model (broken 24 symmetry) for V = N₂ N; = 2 All flow- based models are trained until convergence. From left to right we show: the ground truth, VMoNE, canonicalization, and SESaMo (ours). We refer App. more details on the experiments.
Model
Volume
Symmetry
FAB
VMoNF
Canon.
SESaMo
GMM
2 X
exact Zg
0.78(3)
0.61(1)
0.91(8)
0.9986(2)
GMM
2
broken Zg
0.81(1)
0.83(11)
0.747(2)
0.9947(3)
ESS
64 theory
8
exact (1)
0.26(3)
0.22(2)
0.9472(8)
6² theory
8 X 8
broken (11)
0.28(5)
0.23(1)
0.941(2)
Hubbard
2
broken L4
0.946(9)
0.37(12)
0.839(5)
0.996(1)
Hubbard
18 X 100
broken (Z2)
0.06(5)
0.024(1)
0.74(1)
GMM
2
exact Za
1.19(37)
0.79(11)
0.013(2)
0.0008(1)
GMM
2 x
broken Zs
0.84(26)
1.02(14)
0.189(3)
0.0024(2)
Hubbard
2
broken E,
0.28(8)
0.74(9)
0.112(7)
0.0013(8)
Table 1: Effective sample size (ESS, higher is better) and KL divergened (smaller is better) after con- vergence for different benchmarks. Best results (averages over ten different models) are highlighted in bold.
chemistry, this may not pose a significant problem. Indeed, the well-defined symmetries inherent in many physical and chemical systems often allow for the prior determination of the symmetry sectors, thereby enabling the application of stochastic modulation. Another limitation arises from the penalty term in Eq. (IO), which enforces bijectivity of the NF. If the target density assigns non-zero probability at the border of the canonical cell, bijectivity can only be maintained approximately. As result, the ESS may decrease if only samples that strictly preserve bijectivity are accepted. For example, in the Gaussian mixture model discussed in Sec. 4.1 decreasing the radius R causes the modes to move closer together, thereby increasing the density near the border of the canonical cell. However, in many high-dimensional physics applications, the distance between modes typically increases with the dimensionality the system, thereby mitigating the impact of bijectivity violations. Nonetheless, emphasize that these limitations not specific to SESaMo, but is shared with canonicalization.
CONCLUSIONS
This paper introduces Symmetry- -ymmetry-Enforcing Stochastic Modulation (SESaMo)-4 novel and flexible approach for constructing symmetry-enhaneee NFs. Moreover, we propose an additional penalty term to the reverse KL. divergence to enforce numerical bijectivity. Our extensive numerical experiments demonstrate that stochastic modulation outperforms naïve NFs, mixture models and canonicalization methods. We envision SESaMo as a powerful tool for incorporating inductive biases into generative models when learning target probability densities with challenging symmetries an essential feature in fields like physics and chemistry. Future work will explore the broader capabilities of SESaMo and assess its potential to achieve state-of- state-of-the-an the-hrr performance not only against generative neural samplers but also relative to established numerical techniques, such Hamiltonian Monte Carlo.
10

---

Published as a conference paper at ICLR 2026
REPRODUCIBILITY STATEMENT
We provide our source code under the MIT license at ggituub..o//iiirrsscccceeessss A static, archived version of the code available Ref. (6I). The repository contains training scripts and instructions to reproduce all main experiments. We specify all hyperparameters in Appendix and provide default configuration files. Experiments were conducted with Python 3.9 and PyTorch 2.7 on NVIDIA A100 GPUs, but the code runs other CUDA-enabled GPUs as well.
ACKNOWLEDGMENTS
The authors thank Luca Johannes Wagner for inspiring discussions during the development of this method and Simran Singh for suggesting to extend SESaMo to continuous symmetries. The authors also thank Shinichi Nakajima and Jan Gerken for useful discussions on earlier versions of this work. The authors gratefully acknowledge the access to the Marvin cluster of the University of Bonn. This project was supported by the Deutsche Forschungsgemminschaff (DFG, German Research Foundation) as part of the CRC 1639 NuMeriQS project no. 511713970.
REFERENCES
[1] Nicholas Metropolis, Arianna W. Rosenbluth, Marshall N. Rosenbluth, Augusta H. Teller, and Edward Teller. Equation of state calculations by fast computing machines. The Journal of Chemical Physics, 1086 06 1953.
[2] Linus Görlitz, Zhenglei Gao, and Walter Schmitt. Statistical analysis of chemical transformation kinetics using markov-chain monte carlo methods. Environmental Science & TTchhnlooy, 45(10):4429-443 4437, 05 2011.
[3] A. Dragulescu and V. M. Yakovenko. Statistical mechanics of money. The European Physical Journal B Condensed Matter and Complex Systems, 17(4):723-729, 2000.
[4] F. Barahona. On the computational complexity of ising spin glass models. Journal of Physics A: Mathematical and General, 15(10):3241, oct 1982.
[5] Christophe Andrieu, Nando de Freitas, Arnaud Doucet, and Michael I. Jordan. An introduction to meme for machine learning. Machine Learning, 50(1)::-43,, 2003.
[6] Alan R. Fersht and Valerie Daggett. Protein folding and unfolding at atomic resolution. Cell, 17734)5573-582 2002.
[7] Ulli Wolff. Critical slowing down. Nuclear Physics B Proceedings Supplements, 17:93-102, 1990.
[8] Sam Bond Bond-Taylor, Adam Leach, Yang Long, and Chris G. Willcocks. Deep generative modelling: comparative review of vaes, gans, normalizing flows, energy- and autoregressive models. IEEE Transactions on Pattern Analysis and Machine Intelligence, 44(1 (A(11):73277347, 2022.
[9] Denis Blessing, Xiaogang Jia, Johannes Esslinger, Francisco Vargas, and Gerhard Neumann. Beyond ELBOs: A large-scale evaluation of variational methods for sampling. In Forty-first International Conference on Machine Learning. 2024.
[10] Frank Noé, Simon Olsson, Jonas Köhler, and Hao Wu. Boltzmann generators: Sampling equilibrium states of many-body systems with deep learning. Science, 365(645)) 147, 2019.
[11] Dian Wu, Lei Wang, and Pan Zhang. Solving statistical mechanics using variational autoregres- sive networks. Phys. Rev. Lett. 122:080602, 2019.
[12] Kim A. Nicoli, Shinichi Nakajima, Nils Strodthoff, Wojciech Samek, al. Asymptotically unbiased estimation of physical observables with neural samplers. Phys. Rev. E, 101:023304, 2020.
11

---

Published as a conference paper at ICLR 2026
[13] M. S. Albergo, G. Kanwar, and P. E. Shanahan. Flow-based generative models for Markov chain Monte Carlo in lattice field theory. Phys. Rev. D, 100:034515, aug 2019.
[14] Kim A. Nicoli, Christopher J. Anders, Lena Funcke, Tobias Hartung, et al. Estimation of thermodynamic observables in lattice field theories with deep generative models. Phys. Rev. Lett., 126:032001, 2021.
[15] Aaron van den Oord, Nal Kalchbrenner, Lasse Espeholt, koray kavukcuoglu, et al. Conditional image generation with pixelenn decoders. In Advances in Neural Information Processing Systems, volume 29, page 4797-4805, 2016.
[16] Aäron van den Oord, Nal Kalchbrenner, and Koray Kavukcuoglu. Pixel recurrent neural net- works. In Proceedings of The 33rd International Conference on Machine Learning, volume 48 of Proceedings of Machine Learning Research, pages 1747-1756, 2016.
[17] Danilo Rezende and Shakir Mohamed. Variational inference with normalizing flows. In Proceedings of the 32nd International Conference on Machine Learning, volume 37, pages 1530- 1538, 2015.
[18] Ivan Kobyzev, Simon JD Prince, and Marcus A Brubaker. Normalizing flows: An introduction and review of current methods. IEEE transactions on pattern analysis and machine intelligence, 13(11):3964-397) 2020.
[19] Y. LeCun, B. Boser, J. S. Denker, D. Henderson, R. E. Howard, al. Backpropagation applied to handwritten zip code recognition. Neural Computation, 1(4):541-551, 1989.
[20] Taco Cohen and Max Welling. Group equivariant convolutional networks. In Proceedings of The 33rd International Conference on Machine Learning, volume 48 of Proceedings of Machine Learning Research, pages 2990-2999, PMLR, 20-22 Jun 2016.
[21] Taco S. Cohen, Mario Geiger, Jonas Köhler, and Max Welling, Spherical CNNs. In International Conference on Learning Representations, 2018
[22] Taco Cohen, Maurice Weiler, Berkay Kicanaoglu, and Max Welling. Gauge equivariant convolutional networks and the icosahedral CNN. In Proceedings of the 36th International Conference on Machine Learning, volume 97 of Proceedings of Machine Learning Research, pages 1321-1330. PMLR, 09-15 Jun 2019.
[23] Kristof Schütt, Oliver Unke, and Michael Gastegger. Equivariant message passing for the prediction of tensorial properties and molecular spectra. In Proceedings of the 38th International Conference Machine Learning, volume 139 of Proceedings of Machine Learning Research, pages 9377-9388. PMLR, 18-24 Jul 2021.
[24] Alexander Bogatskiy, Brandon Anderson, Jan Offermann, Marwah Roussi, David Miller, al. Lorentz group equivariant neural network for particle physics. In Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pages 992-1002. PMLR, 13-18 Jul 2020.
[25] Victor Garcia Satorras, Emiel Hoogeboom, and Max Welling. E(n) equivariant graph neural networks. In ICML, pages 9323-9332 2021.
[26] Di Luo, Giuseppe Carleo, Bryan K. Clark, and James Stokes. Gauge equivariant neural networks for quantum lattice gauge theories. Phys. Rev. Letr., 127:276402, dec 2021.
[27] Roman Soletskyi, Marylou Gabrié, and Bruno Loureiro. theoretical perspective on mode collapse in variational inference. arXiv preprint orXiv:2410.13300, 2024.
[28] Kim A. Nicoli, Christopher J. Anders, Tobias Hartung, Karl Jansen, Pan Kessel, et al. Detecting and mitigating mode-collapse for flow-based sampling of lattice field thcories. Phys. Rev. D. 108:114501, dec 2023.
[29] Michael M. Bronstein, Joan Bruna, Yann LeCun, Arthur Szlam, and Pierre Vandergheynst. Geometric deep learning: Going beyond euclidean data. IEEE Signal Processing Magazine, 34(4):18-43 2017.
12

---

Published as a conference paper at ICLR 2026
[30] Jan E. Gerken, Jimmy Aronsson, Oscar Carlsson, Hampus Linander, Fredrik Ohlsson, et al. Geometric deep learning and equivariant neural networks. Artificial Intelligence Review, 56(12)11400-116622 2023.
[31] Michael M Bronstein, Joan Bruna, Taco Cohen, and Petar Veličković. Geometric deep learning: Grids, groups, graphs, geodesics, and gauges. arXiv preprint arXiv: 204. 13478, 2021.
[32] Jonas Köhler, Leon Klein, and Frank Noe. Equivariant flows: Exact likelihood generative learning for symmetric densities. In Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pages 5361- 5370. PMLR, 13-18 Jul 2020.
[33] Victor Garcia Satorras, Emiel Hoogeboom, Fabian Bernd Fuchs, Ingmar Posner, and Max Welling. E(n) equivariant normalizing flows. In Advances in Neural Information Processing Systems, 2021.
[34] Ricky T. Q. Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud. Neural ordinary differential equations. In Advances in Neural Information Processing Systems, volume 31. Curran Associates, Inc., 2018.
[35] Avishek Joey Bose, Marcus Brubaker, and Ivan Kobyzev. Equivariant finite normalizing flows. arXiv preprini arXiv:21 2021.
[36] Gurtej Kanwar, Michael S. Albergo, Denis Boyda, Kyle Cranmer, Daniel C. Hackett, et al. Equivariant flow-based sampling for lattice gauge theory. Phys. Rev. Lett., 125:121601, sep 2020.
[37] Denis Boyda, Gurtej Kanwar, Sébastien Racanière, Danilo Jimenez Rezende, Michael S. Albergo, Kyle Cranmer, Daniel C. Hackett, and Phiala Shanahan. Sampling using SU(N gauge equivariant flows. Phys. Rev. D, 103(7):074504, 2021.
[38] Dominic Schuh, Janik Kreit, Evan Berkowitz, Lena Funcke, Thomas Luu, Kim A Nicoli, and Marcel Rodekamp. Simulating the hubbard model with equivariant normalizing flows. arXiv:2501.07371, 2025.
[39] Peter Wirnsberger, Andrew J. Ballard, George Papamakarios, Stuart Abercrombie, Sébastien Racanière, et al. Targeted free energy estimation via learned mappings. The Journal of Chemical Physics, 153(14):144112, oct 2020.
[40] Peter Wirnsberger, George Papamakarios, Borja Ibarz, Sébastien Racanière, Andrew J Ballard, al. Normalizing flows for atomic solids. Machine Learning: Science and Technology, 3(2):025009, may 2022.
[41] Laurence Illing Midgley, Vincent Stimper, Javier Antoran, Emile Mathieu, Bernhard Schölkopf, et al. SE(3) equivariant augmented coupling flows. In Thirty-seventh Conference on Neural Information Processing Systems, 2023.
[42] Leon Klein, Andreas Krämer, and Frank Noe. Equivariant flow matching. In Thirty-seventh Conference on Neural Information Processing Systems, 2023.
[43] Ling Yang, Zhilong Zhang, Yang Song, Shenda Hong, Runsheng Xu, al. Diffusion models: A comprchensive survey of methods and applications. ACM Comput. Surv., 56(4), 2023.
[44] Emiel Hoogeboom, Víctor Garcia Satorras, Clément Vignac, and Max Welling. Equivariant diffusion for molccule generation in 3D. In Proceedings of the 39th International Conference on Machine Learning, volume 162 of Proceedings of Machine Learning Research, pages 8867-8887. PMLR, 17-23 Jul 2022.
[45] Guilherme G. P. Freitas Pires and Mário A. T. Figueiredo. Variational mixture of normalizing flows. European Symposium on Artificial Neural Networks, 2020.
[46] David M Blei, Alp Kucukelbir, and Jon D McAuliffe, Variational inference: A review for statisticians. Journal of the American statistical Association, 112(518):559887, 2017.
13

---

Published as a conference paper at ICLR 2026
[47] Laurent Dinh, David Krueger, and Yoshua Bengio. NICE: non-linear independent components estimation. In 3rd International Conference on Learning Representations, ICLR 2015, Workshop Track Proceedings, 2015.
[48]
Laurent Dinh, Jascha Sohl-Dickstein, and Samy Bengio. Density estimation using real NVP. In International Conference on Learning Representations, 2017.
[49] Conor Durkan, Artur Bekasov, lain Murray, and George Papamakarios. Neural spline flows. In Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019.
[50] Durk Kingma and Prafulla Dhariwal. Glow: Generative flow with invertible 1x1 convolutions. In Advances in Neural Information Processing Systems, volume 31. Curran Associates, Inc.,
2018.
[51] George Papamakarios, Eric Nalisnick, Danilo Jimenez Rezende, Shakir Mohamed, and Balaji Lakshminarayanan, Normalizing flows for probabilistic modeling and inference. Mach. Learn. Res., 22(1), jan 2021.
[52] S. Kullback and R. A. Leibler. On information and sufficiency. The Annals of Mathematical Statistics, 22(1):79 86, 1951.
[53] Haggai Maron, Heli Ben-Hamu, Nadav Shamir, and Yaron Lipman. Invariant and equivariant graph networks. International Conference on Learning Representations, 2019.
[54] Haggai Maron, Ethan Fetaya, Nimrod Segol, and Yaron Lipman. On the universality invariant networks. In Proceedings of the 36th International Conference on Machine Learning, volume 97 of Proceedings of Machine Learning Research, pages 4363-4371. PMLR, 09-15 Jun 2019.
[55] Derek Lim, Joshua Robinson, Stefanie Jegelka, Yaron Lipman, and Haggai Maron. Expressive sign equivariant networks for spectral geometric learning. In ICLR 2023 Workshop on Physics for Machine Learning, 2023.
[56] Michele Caselle, Elia Cellini, Alessandro Nada, and Marco Panero. Stochastic normalizing flows as non- non-equilibrium transformations. JHEP, 07:015, 2022.
[57] Mathis Gerdes, Pim de Haan, Corrado Rainone, Roberto Bondesan, and Miranda C. N. Cheng. Learning lattice quantum field theories with equivariant continuous flows. SciPost Phys., 15:238, 2023.
[58] John Schulman, Nicolas Heess, Theophane Weber, and Pieter Abbeel. Gradient estimation using stochastic computation graphs. Advances in Neural Information Processing Systems 28, abs/1506.00254, 2015.
[59] Laurence Illing Midgley, Vincent Stimper, Gregor N. C. Simm, Bernhard Schölkopf, and José Miguel Hernández-Lobate Flow annealed importance sampling bootstrap, 2023.
[60] Kim A. Nicoli, Christopher J. Anders, Lena Funcke, Karl Jansen, Shinichi Nakajima, et al. NeuLat: toolbox for neural samplers in lattice field theories. PoS, LATTICE200232266, 2024.
[61]
Janik Kreit. Sesamo: Official implementation. Zenodo, 2026.
[62] Ryan Abbott, Denis Boyda, Daniel C Hackett, Gurtej Kanwar, Fernando Romero-López, Phiala E Shanahan, and Julian M Urban. Progress in normalizing flows for 4d gauge theories. arXiv preprint xrxv: 00263, 2025.
[63] Andrea Bulgarelli, Elia Cellini, Karl Jansen, Stefan Kühn,, Alessandro Nada, Shinichi Nakajima, Kim A. Nicoli, and Marco Panero. Flow-based sampling for entanglement entropy and the machine learning of defects. Phys. Rev. Lett., 134:151601, Apr 2025.
[64] Miranda CN Cheng and Niki Stratikopoulon. Lecture notes normalizing flows for lattice quantum field theories. arXiv preprint arXiv:2504.18126, 2025.
[65] Hagen Kleinert and Verena Froolte-. Critical Properties of Thi4-hheories. WORLD SCIENTIFIC, 2001.
14

---

Published as a conference paper at ICLR 2026
[66] Lorenz Vaitl, Kim Andrea Nicoli, Shinichi Nakajima, and Pan Kessel. Path-gradient estimators for continuous normalizing flows. In Proceedings of the 39th International Conference on Machine Learning, volume 162 of Proceedings of Machine Learning Research, pages 21945- 21959. PMLR, 17-23 Jul 2022.
[67] Alex Matthews, Michael Arbel, Danilo Jimenez Rezende, and Arnau Doucet. Continual repeated annealed flow transport Monte Carlo. In Proceedings of the 39th International Conference on Machine Learning, volume 162 of Proceedings of Machine Learning Research, pages 15196-15219 PMLR, 17-23 Jul 2022.
[68] Lorenz Vaitl, Ludwig Winkler, Lorenz Richter, and Pan Kessel. Fast and unified path gradi- estimators for normalizing flows. In The Twelfth International Conference Learning Representations, 2024.
[69] Edward Witten. Phases of = 2 theories in two dimensions. Nuclear Physics B, 2,493()):19--222, 1993.
[70] Daniel P. Arovas, Erez Berg, Steven A. Kivelson, and Srinivas Raghu. The hubbard model. Annual Review of Condensed Matter Physics, (3(11:239-277) March 2022.
[71] Thomas Luu and Timo A. Lähde. Quantum monte carlo calculations for carbon nanotubes. Phys. Rev. B, 93:155106, Apr 2016.
[72] Kim A. Nicoli, Christopher J. Anders, Tobias Hartung, Karl Jansen, Pan Kessel, et al. Detecting and mitigating mode-collapse for flow-based sampling of lattice field theories. Phys. Rev. D, 108: I 14501, dec 2023.
[73] Conor Durkan, Artur Bekasov, lain Murray, and George Papamakarios. Neural spline flows. Advances neural information processing systems, 32, 2019.
[74] Jan-Lukas Wynen, Evan Berkowitz, Christopher Körber, Timo A. Lähde, and Thomas Luu. Avoiding ergodicity problems in lattice discretizations of the hubbard model. Phys. Rev. B, 100:075141, Aug 2019.
[75] Michael E. Peskin and Daniel V. Schroeder. An introduction to quantum field theory. Frontiers in Physics, 1995.
[76] Daniel Naegels. An introduction to goldstone boson physics and the coset construction. arXiv preprint arXiv:21 10.14444 2021.
[77] Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization, 2017.
15

---

Published as a conference paper at ICLR 2026
A PSEUDOCODE
To offer an intuitive overview of proposed method, we include pseudocode example in Alg.
Algorithm 1 Training loop of SESaMo
1: Initialize flow ge with parameters 0
2: Initialize stochastic modulation S,, with parameters b
3: for iteration = to N do
4: Sample (11B from qo and evaluate In 4o (z (2(::))
5: Obtain (1:B) and In det day
from flow (1:B))
dz
6:
Compute regularization A(z)) Sample (I:B) from ps.6 and evaluate In (s.s(u(I.)))
7:
8:
Obtain (I:B) and In det 8S..
from stochastic modulation Sym-a (x(1:.1)
0ge
Compute unnormalized target probability f (I:1))
9:
10:
Compute log probability In qo (a(1..))
Compute loss (g. || p) (1z(1.B))
12:
Update parameters 0, according to the REINFORCE estimator
13:
end for
B INTUITIVE COMPARISON OF CANONICALIZATION AND STOCHASTIC MODULATION
In the main text, two approaches for effectively incorporating symmetries into generative models such as NFs were introduced: canonicalization in Sec. Sec.2.3.2 and Symmetry-Enforeet Stochastic Modulation (SESaMo) in Sec. 3.1 In this section, we summarize the differences between these approaches a more intuitive level. To help the reader familiarize with the underlying ideas, we provide an illustration for both SESaMo (top row) and canonicalization (bottom row) in Fig. 3 showing example target distribution and corresponding prior. In Fig. 3 the goal is to sample from a toy target density that exhibits three modes, visually represented by the three red triangles on the left of Fig. Both approaches start from Gaussian prior density qo. represented by circle.
In the case of SESaMo (top row), a random sample 2 ~ qu is transformed by an NF, i.e., parametric map Ge. such that the probability mass of the prior density is shifted and transformed cover one of the modes of the target density (depicted as the triangle with a solid black line), which lies within the canonical cell Ω (dashed black line), while the other symmetric modes (triangles with dotted black line) remain uncovered. The transformed density is denoted as qe. At this stage, the model has captured only one mode of the target density. Subsequently, SESaMo employs the stochastic modulation Sr,, to redistribute the probability mass towards the other modes of the target density, resulting the final variational probability distribution go- This distribution (visualized by the three triangles) approximates the target density
The canonicalization approach, depicted in the bottom row of Fig. 3 also starts with prior Gaussian distribution Qo- Samples drawn from the prior distribution are transformed such that any sample z ~ Yo is mapped to the canonical cell (dashed black line), resulting in the density 4x. (solid black line), while the other symmetric modes (dotted black line) remain uncovered. Subsequently, an NF ge learns a bijective map to transform these samples in the canonical space. Canonicalized samples, denoted as E... are then drawn from the resulting distribution go, which is illustrated in Fig. 3 as a triangle with a solid black line. Given that the resulting parametrized distribution Go in the canonical space, it needs be transformed back the input space. This is achieved by applying the inverse of the initial transformation C to the samples T., resulting in the final parametrized probability distribution qo- The support of go is visualized in the right-most plot of the bottom row by three triangles that approximate the target density p.
16

---

Published as a conference paper at ICLR 2026
C EQUIVARIANCE OF THE CANONICALIZATTIN METHOD
Let us consider a general symmetry transformation under which some function E(-) x: E R - {(x) invariant, i.e., {(x) <TTx). A concrete example of such function can be the action of equivariant
a physical system, such as Eqs. 21 and (2 learnable map Ge z E -
under T, and is thus denoted go. if it satisfies the following condition:
(T() = go
(23)
The canonicalization approach, introduced in Sec. 2.3.2 leverages a so-called canonical transforma- tion Cr. : R² Ω to map samples from the input space into the canonical cell Ω2 thereby making the map ge equivariant with respect to I", The equivariant map ge thus reads
ge (x) = CT, (o(,r,
(24)
where Cr,s maps sample into the canonical cell 22 ge denotes a specific NF, and Cr.s maps the canonicalized (and transformed) sample = ==, back to the original input space.
In this section, we restrict ourselves involutory symmetry transformations, i.e., T = 1. Our goal is thus to show that canonicalization fulfils the equivariant condition in Eq. (23). We define the canonical transformation
if EΩ Tz,
2,
CT 2
(25)
if E Ω
with the inverse transformation
x if TTx, if f2 E Q
CT +
(26)
is crucial to note that the inverse transformation CT.2 still depends on the sample to which the canonical transformation Cr. was initially applied, i.e., the information about the initial sample is implicitly stored in the transformation. One way to check if the map Eq. (24) is really equivariant under the transformation to sequentially apply the transformation and then Cr, to the input z,
(T2, TTz, if ET
if EQ
[Tz,
if EΩ
CT,T: Tz
(27)
2
e
Note that the involutory property TT = 1l has been used here CT,T= and Cr, are equivalent,
It follows that the transformations
CT.T:Tz = Cr 2;
(28)
while the inverse transformation
C reads
if E (Tx, i
x
CT 1 +
(29)
Additionally, compute TC-1
Tx, if EΩ
(Tx, e
TCT +
(30)
TTx,
if EΩ x if
and verify that indeed
T.r = TCT.
(31)
Leveraging the identities in Eqs. 28 and B1). one can finally show that the overall map ğo is
cquivariant with respect to the transtormation T,
Jo (Tx) = CT go (CT.T= Tx) = TCT1 (CT, z) = [z)
(32)
which proves the initial equivariance condition in Eq.
'Noe that while the subscript T, means that the forward canonical transformation is applied to the input the subscript T, T2 means that the transformation is applied to the transformed input Tz.
17

---

Published as a conference paper at ICLR 2026
A
B 1
B 2
B 10
Sample
Figure 6: Example of penalty term with M(x) = |a| - т. The penalty term is zero for x E [-TT,] m]
and approaches A as + 300 The parameter B controls the scaling of the penalty gradient.
An essential part of the canonicalization is that the map go must not move the canonicalized sample Cr, outside the canonical cell, i.e., into Ra \ 12. This requirement arises because if the map go maps a sample outside of the canonical cell Ω2 that is, if go (CT,a z) f2 then it is possible for two distinct imputs 2 jf 22 with 21, 22 IR" be mapped to the same output via canonicalization and transformation: (=
(z2). This leads to a loss of injectivity and, consequently, the transformation ge is no longer bijective. This poses a problem, as NFs require the map Ge to be bijective in order to perform density estimation via Eq. (2) As described in Sec. 2.3.3] this constraint can be numerically enforced using a penalty term a E Ih + ((x) E R which is zero for æ E Ω and greater than towards the canonical cell Ω. This ensures that if the NF pushes sample Σ ge (Cr.., 2) outside of Ω the gradient of A acts to pull it back into the cell. Further details on the penalty term and the enforcement of bijectivity are provided in Sec. 2.3.3 and further elaborated in App. D
for æs Ω. Furthermore, it is essential that the gradient 2_A (go())) points
D PENALTY TERM FOR THE KL DIVERGENCE
we introduced penalty term that is necessary to numerically enforce the bijectivity required for the NF to serve valid transport map between probability densities. In this section, we further elaborate on this penalty term and provide an example in Fig.
In Eq. [T0) from Sec. 2.3.3
Crucially, the penalty term A(x) and the associated penalty function X(x) are necessary for ensuring that the NF ge does not map samples outside of the canonical cell Ω. For convenience, we recall the penalty term,
A(w) = A A. o(B >(x)) ((((z)),
(33)
where the set {A, B} denotes all hyperparameters, while (-) and (--) refer to the sigmoid and Heaviside theta functions, respectively.
Fig. shows an example for penalty term for the canonical cel p = (a E IR |a| S =). The function A(x) |z| is chosen so that it becomes
the boundary |z]| and positive outside the canonical cell, i.c., for |a] > æ. Correspondingly, the penalty term A(x) is zero for all |- л and smoothly approaches the value A too. The parameter B controls the scaling of the gradient of the penalty term.
E GENERALIZATION OF SESAMO
E.
EM STOCHASTIC MODULATION
The stochastic modulation for the Z, symmetry introduced in Sec. 3.1 can be generalized to a ZM symmetry. The transformation S, randomly rotates two-dimensional vector æ (x1, R
Note that the example is in one-dimensional space R hut be straightforwardly generalized.
18

---

Published as a conference paper at ICLR 2026
4s
q=
Probability
Probability
0.5
0.5
0
4
2
4
-2
2
4
Sample x
Sample x
Figure 7: Prior Gaussian distribution with mean 2 and standard deviation transformation Su, implementing the Z2 symmetry, randomly flips the sign of a sample r; ~ go with probability determined by the breaking parameter b. When b = In 0.5 (left), the resulting distribution is (yellow) is symmetric around zero, with both modes carrying equal probability mass. When b In 0.25 (right), according to Eq. (38), the sign flip occurs with probability Ps, 0.25. leading to asymmetric modes at f = +2 that carry 25% and 75% the total probability mass, respectively.
1. The
b
ps, = 0)
Ps 1)
O In 0.5
0
1/2
1/2
0
Table 2: Probability ps. of not flipping (u 0) and flipping (u = 1) the sign of the input for examples of the breaking parameter including the even case and the edge cases.
about the origin by angle of 2Tu/M, i.c.,
cos
2m M
sin 2
S, x
M
with
~~ Usise (0, M),
(34)
sin
STT
M
COS
2nd
where u ~ Mdiso (0, M) is discrete uniform random variable taking values in the set {0, 1, 2.... M 1). The modulation probability is therefore given by Ps 1/M. To ensure the bijectivity of the transformation S., the penalty term A is added to the KL divergence in Eq. (10). where
(ax) ( + (x)),
(35)
and the bijectivity function is expressed
A (x: tan (T/M) +
(36)
(1 (an(=/MM)
The canonical cell defined by this penalty corresponds to a sector of angular width 2 2 centered around the Ij -axis, with boundaries at angles T /M. The bijectivity function then measures the distance of sample to the border the canonical cell. For more details on the penalty term, we refer back to Sec. 2.33
E.2 BROKEN Z, STOCHASTIC MODULATION
In the main text, the exact Eg symmetry was considered to illustrate how canonicalization and SESaMo transform the base density. A Z₂ symmetry called exact when both modes (as shown in Fig. and Fig. carry equal probability mass. the following, we extend this to a more general case where the probability mass unevenly distributed across the modes.
is important to note that under these conditions, the canonicalization approach faces challenges. Specifically, it is longer sufficient learn single mode and evenly distribute the probability mass
19

---

Published as a conference paper at ICLR 2026
among the others. In contrast, SESaMo, owing to its greater flexibility, effectively handle this asymmetry. To accommodate such cases, a learnable breaking parameter b E R is introduced to account for the imbalance in probability mass between the modes. When b - 0, the sign of æ is always flipped, whereas in the limit + ~2, the sign is never flipped. The transformation S,, for a broken Z2 symmetry therefore yields
if 0 with
:
B(E")
and
E R
(37)
=
where B() denotes a Bernoulli distribution. Note that when = In 0.5, the transformation reduces to the symmetric Z2 case, where each mode is selected with equal probability. Tab. shows the modulation probability ps,6 for the even case and the edge cases of the breaking parameter b discussed above. For an arbitrary breaking parameter b, the modulation probability ps,b is given by
Ps { e²
if =0
(38)
1,
where u. B(e"). The corresponding bijectivity constraint, used in the penalty term A introduced in Eq. 10 reads
N
(xx)
Σ
(39)
where the sum is taken over of all components of the vector E RN The breaking parameter b is used in the exponential to ensure numerically stable simulations, which becomes particularly important in the limits ps,6 0 and ps,6
Fig. 7 (left) shows a one one-dimensional Gaussian distribution qo (blue), centered at x: = 2 with standard deviation 1. Applying the stochastic modulation S,, corresponding the Z, symmetry, with the breaking parameter = In 0.5, yields a new distribution (yellow) that is symmetric around zero. In this case, the probability mass is equally distributed across both modes. When the breaking parameter b * In 0.5, the stochastic modulation accounts for the imbalance between the modes, resulting unequal probability masses the transformed density qs- Fig. (right) shows an example for b = In 0.25, where the mode at x < O carries less mass than the one at æ > O.
Numerically, so-called breaking ratio can be estimated by counting the number of samples in each mode of the distribution:
N4 N N_ = 2e", N +
R
(40)
where N₂, and N_ denote the number of samples in the positive and negative modes of qa. respectively. As an example, the experiments for the Hubbard model presented in the main text feature a broken 2 symmetry, composed of an exact Z, and broken Σ, symmetry. SESaMo is able to learn this broken E, symmetry by combining an exact and a broken Z, transformation, i.c., effectively modulating the sign of one of two field components.
E.3 BROKEN (Z STOCHASTIC MODULATION
In the main text, we presented results for the Hubbard model for both small and large spatial lattice extents, i.e., N₂ = 2 and N₂ 18. The Hubbard models exhibits combination of N₂ distinet Ze symmetries, forming a (Z2) symmetry. In terms of the stochastic modulation, we introduce breaking parameters b R=" with n 2N= The stochastic variable
drawn from categorical distribution (bt(b) with log-odds 6 and u = 0, I,..., N2. - 1. The canonical cell is defined as Ω {æ E R. XM, E, F. 0 Vi 0, ..... N, 1} and the stochastic modulation is given by
1 L,.., 1, 1, Dooo =fu=0
L.... 1, 1.-1)" ox iu=1
L,.., 1,-1, 1)a 0x fu=2
Sy : +
I..., 1,-1,-17'oz
fu=3
with u Cat(b), (41)
(-1 -1,
ox
if =
2
20

---

Published as a conference paper at ICLR 2026
where @ denotes component cimponent-wise multiplication. This formulation ensures that can be transformed to one the 2% orthans.
We combine this broken symmetry with the exact 2g symmetry present in the system, resulting in an effective broken (Zz) N=2 exact Eg symmetry. For N₂ = 2, this reduces to a broken Σ symmetry.
STOCHASTIC MODULATION FOR CONTINUOUS SYMMETRIES
In Sec. vcc.3.1 the stochastic modulation S, was introduced for discrete symmetries, where S. has a finite number of possible outcomes, each selected according to the modulation probability ps.b. This approach is well-suited for discrete symmetries such as sign-flip or Zas symmetries. However, is not applicable continuous symmetries such as rotational translational symmetries- where the transformation space is uncountably infinite. In these cases, a modified formulation of stochastic modulation is required account for the continuous nature of the symmetry group.
The continuous stochastic modulation proceeds as follows: first, draw a sample u from a distribution (u), which can be the uniform distribution (0, 1). Then, apply a trainable map hy [0, 1) [0, 1) with parameters ó to obtain hy (2). This output parametrizes a continuous transformation (M), such
as a rotation matrix where the rotation angle is determined by ha(u). The stochastic transformation is thus given by
Su (n)
(42)
The modulation probability, which enters the density transformation in Eq. [15), follows from the change-of -aa------vvrralll formula of the transformation RAG (u) and can be expressed as
(oRputus)
(s,((u) qu qu(u)=
det
(43)
Ou
where (u) is the probability density of u and the determinant captures the local volume change under the inverse transformation R (n)"
F.1
BROKEN AND EXACT U(1) STOCHASTIC MODULAATION
In Sec.
4.2
the complex of scalar field theory is introduced, in which the action [(x:] (as defined in remains invariant under (1) transformation of the form
Eq. 21
R = 2m
(44)
where the angle yp lies in the interval [0. 1). If term aRe[æ] is added to the action / [(z) this (())
symmetry is broken, meaning that the Boltzmann-like density p(æ) exp (=f(==)) /2 becomes dependent on the angle sp. This angular dependence can be captured within the stochastic modulation
framework by introducing a trainable map o ho(u). In particular, spline flow (73) is used for this purpose. The modulation probability in Eq. (43) then simplifies to
det (Dholu) )
ps,((u) = 2π
(45)
Ou
where the chain rule used compute d R Du in Eq. (43). as well the fact that
(RR
ha(a)
det
(46)
ahb
2x
This is given because the rotation R, 2mrp in Eq. 44 corresponds a full angular cycle over the interval [0, 1), scaling the Jacobian by the full rotation angle 2n. Meanwhile, we used qu = since u is sampled from uniform distribution on [0, 1), which has constant density of one.
The sample æ must be completely real before applying the stochastic modulation. This means that prior sample z 21 iiz, where 21, 22 IR. must satisfy 22 0, i.e., it lies the real axis,
21

---

Published as a conference paper at ICLR 2026
and is transformed by an NF go R - R After applying the stochastic modulation Rng(")) the sample æ becomes complex-valuse valued, given by
e 2mthss()) go g0()))
(47)
Note that by omitting the spline flow h2 and using h an exact (1) symmetry be enforced instead of a broken one. Furthermore, this approach can similarly be used to enforce a broken or exact rotational (2) symmetry.
G TECHNICAL DETAILS OF THE PHYSICAL THEORIES
In this section, we discuss some fundamental aspects of the complex o² theory and the Hubbard model that are relevant to our study.
G.1 THE COMPLEX 6² SCALAR FIELD THEORY IN TWO DIMENSIONS
In recent years, the & theory has become popular benchmark for generative models in the machine learning community (66) 67 68). Originally developed as physical model, it describes interacting particles with integer spin. On finite lattice with points j E V, the theory is specified by the action
fbl
Σ JEV
a
(4))tp ;
m P4
9u 4
(48)
a²
where s denotes the field value site The first term inside the brackets corresponds to the kinetic term, the second the mass term governed by the bare mass mo, and the quartic y term describes the interaction, weighted by the bare coupling strength gu- Using the more standard redefinitions (similarly adopted by Nicoli et al. (4))
(amo)² = 2A 4,
6Λ
& = (2x) 1/2 (2x)1/2>) x,
d'
(49)
R²
we rewrite the action in the form presented in the main text:
f|z]
Σ
2K
a [4) + (1 21)a² + Xa + aRe(z;]
(50)
JEV
A=1
Here, À is known as the coupling parameter, while is the hopping parameter. Additionally, we added term aRe(ar, to progressively break the ((1) symmetry of the のi theory as the parameter a increases. Such symmetry- breaaknkttybbbbkiggg.... term also arises in quantum field theories with non-degenerate particle flavor masses, providing a physically motivated example.
G.2 THE HUBBARD MODEL IN TWO DIMENSIONS
The Hubbard model fundamental model in condensed matter physics that describes how electrons interact on a fixed lattice of ions 77).. By neglecting lattice vibrations and other atomic excitations, captures the essential physics of electrons hopping between valence orbitals and interacting through their electric charge. This is further illustrated in Fig. g We describe the system in the so-called spin basis, where the degrees of freedom correspond to spin-up and spin-down electrons. Other basis choices exist but are not considered here.
The action of the system is given by
71
f
Σ
25 log det M [x] log det M [[=x),
(51)
20
J.KEV
where δ denotes the on-site Coulomb-like interaction strength, æ are auxiliary bosonic fields, and the subscripts j. k label the spatial and temporal lattice sites in the lattice volume V, respectively. Since we do not consider temporal extent throughout this manuscript, i.e. N₂ = 1, we have dropped the index & in Sec. 4.2 for brevity, Lastly, the fermion matrix M is defined
M [x] y .i kgr & - Ie"ly Me By dy A+1
(52)
22

---

Published as a conference paper at ICLR 2026
F
Figure 8: Illustration of lattice described by the Hubbard model. Blue and red circles represent spin-up and spin-down electrons, respectively. The hopping term A allows electrons to move between neighbouring lattice sites, while the on-site Coulomb interaction U penalizes the presence of two electrons with opposite spins at the same site.
hopping matrix, where is the hopping amplitude and o(y' enforces hopping only between nearest neighbours j', on the lattice, and B, is factor implementing periodic (anti-periodic boundary conditions in the temporal direction for N, = (N, > 1). The action in Eq. (51) consists of two main contributions: the Gaussian term, which encodes the on-site interaction, and the fermionic term, represented by the product of fermion matrices, which captures the electron hopping dynamics across the lattice.
Here, = -
The Boltzmann-like density of the Hubbard model features widely separated modes, which can lead to ergodicity problems and biased estimates of observables when using Monte Carlo-based sampling methods such Hybrid Monte Carlo (HMC) (74). NFs have demonstrated the ability overcome these challenges, particularly when they incorporate prior knowledge of the system's 's symmetries (38).
H ADDITIONAL NUMERICAL EXPERIMENTS
In this section, we present additional experiments for the Gaussian mixture model, the Hubbard model, and the 6² theory.
H.1 GAUSSIAN MIXTURE
The Gaussian mixture model introduced in Sec. exhibbts multi-modas density, where locating all modes is poses a significant challenge for RealNVP. This issue is mitigated by applying canonical- ization and further improved with SESaMo, which achieves higher accuracy. Fig. 9 (left) shows the ESS as a function GPU training time in minutes. The solid lines and shaded regions indicate the mean and standard deviation
ten models trained with different seeds. Both canonicalization and SESaMo lead to faster convergence compared to RealNVP, which suffers from strong fluctuations due to frequent mode collapse.
H.2 THE HUBBARD MODEL IN TWO DIMENSIONS
In Fig. 9 (right), the ESS is shown as function of the GPU training time for the Hubbard model. The solid lines and shaded regions indicate the mean and standard deviation over ten models trained with different seeds. SESaMo not only achieves higher accuracy than both canonicalization and RealNVP, but also converges faster than RealNVP. The canonicalization method fails to capture the unequal probability masses across the modes, illustrated in Fig. []] while RealNVP suffers from droppin.... contrast, SESaMo successfully identifies all four modes and accurately predicts their relative probabilities.
23

---

Published as a conference paper at ICLR 2026
1.00
1.00
0.75
0.75
ESS
ESS
0.50
0.50
0.25
0.25
SESaMo (ours)
RealN VP
SESaMo (ours)
RealN VP
0.00
Canonicalization
0.00
Canonicalization
3
3 GPU minutes
GPU minutes
(a) Gaussian mixture
(b) Hubbard model
Figure 9: ESS function of the GPU training time (minutes) for the Gaussian mixture (left) and the Hubbard model (right). Solid lines represent the mean and shaded areas indicate the standard deviation across ten models trained with different seeds. The results show that SESaMo achieves higher ESS compared to both canonicalization and ReaINVP.
The effect the broken Zn symmetry becomes more pronounced as the inverse temperature B increases. To investigate this behaviour, train SESaMo and canonicalization models for values of E [1, 4), as shown in Fig. 10 (left). SESaMo consistently achieves high accuracy across all values of B8 while the canonicalization method exhibits significantly lower accuracy. This demonstrates that SESaMo successfully learns the broken Z₂ symmetry.
To further verify whether the probability is predicted correctly, we compare against the ground truth. In Fig. (right), the breaking ratio At from Eq. (40) is shown, where N can be computed analytically by integrating the probability distribution p(z) for volume V 2 1, i.e., x (21,22) E R The probability distribution is known up to constant factor and given by
*f+28
p(æ) x h(x)/((-)
(53)
where
Ø1 + D2
[1 - x2
h(x) cosh
+ cosh
cosh (R)
(54)
2
2
The theoretical prediction of the breaking ratio matches perfectly with the expression R obtained from the learned breaking parameter b.
2
H.3 THE REAL of SCALAR FIELD THEORY IN TWO DIMENSIONS
In Sec. 4.2 and G.1 we introduced the complex 6f scalar field theory in two dimensions. In its general form, this theory consists of complex-valuee fields.
Most recent works in the context of generative models (see, e.g., 1 14)). however, have focused on real scalar fields. Under this assumption, the &²t theory belongs to the same universality class the Ising model and serves as an instructive toy model for exploring spontaneous symmetry breaking and the Higgs mechanism (65). Assuming real scalar fields, the action Eq. (2I) simplifies
í[a] Σ
2x (maaz+a)) 22)ar Xay + ony
(55)
with R". This form of the action corresponds the studied in Ref. (14 72). up to the addition of symmetry-breen breaking factor aæ. The coefficient a introduces an exponential suppression of the probability with respect to the field thereby explicitly breaking the Za symmetry when a > 0. this context, the symmetry-breakinn parameter introduced App. can be learned such
'Note that this distribution is exact for V 1. For larger volumes, it becomes exact only in the strong-coupling limit U while keeping fixed.
24

---

Published as a conference paper at ICLR 2026
Amiytical
SES (ours)
-0.3
08
-0
0.6
SES (surs) Capoticalization
Breaking
0.4
-0.3
0.2
-0.9
2.0
3.0
Inverse tompersture
g
Inverse temperature B
(a) ESS function of inverse temperature
(b) function of inverse temperature
Figure 10: Left: ESS for different values of the inverse temperature B. The blue and yellow markers correspond to canonicalization and SESaMo, respectively. Means and standard deviations computed by averaging over three independently trained models (for each method) using three different random seeds. Right: Breaking ratio R as function of [3. The analytical curve (yellow) obtained by integrating the analytically derived probability weight (see Eq. (79) in Ref. (74)). The numerical estimate from Eq. (40), computed using a trained SESaMo model, agrees with the analytical result within error bars. The uncertainties- often too small to be visible at the scale of the plot-are are estimated by averaging over three independently trained models with different seeds.
0.4
Samples
Fit
Samples
Counts (normalized)
Counts (nommllzzed)
0.3
0.2
0.2
0.1
0.1
0.0
0.0
-4
-2
2
4
(x)
(x)
(a) Magnetization for broken Za (a 0.001).
(b) Magnetization for unbroken E2 (a 0).
Figure 11: Histograms the magnetization for real s² scalar field theory for a broken Z, symmetry (left, 0.001) and an exact Ze symmetry (right, 0). Samples for the histograms are drawn from two SESaMo models trained for the corresponding values of the breaking factor a.
that SESaMo redistributes the probability mass of the learned probability in accordance with the asymmetry of the target distribution. We train SESaMo using the REINFORCE estimator the KL divergence discussed in Sec. 3.2 Additional hyperparameters and experimental details are provided in App. [ Unless stated otherwise, all experiments in this setting are conducted on lattices of size 16 x 8, with action parameters fixed to к 0.3 and X = 0.022.
Since the theory now consists of scalar real fields two dimensions, it enters the so-called broken phase for couplings {k, A} = {2 0.3, 0.022]. This phase is characterized by bimodal probability density with the centers the modes located at the vacuum expectation values (VEVs) 72) of the theory. When a 0, both modes are identical, and the resulting distribution is symmetric. In this case, both SESaMo and canonicalization are able to accurately learn the target distribution, achieving high ESS without mode collapse (72). In the following, we compare the performance of SESaMo and canonicalization in the case a > O, where the Zz symmetry the double well potential is explicitly broken. To study this scenario. increasing values of a. The results are shown in Fig. [2 (left), which displays the ESS obtained from models trained for theory defined lattice of size V 16 8 for various values of
trained different models using both approaches for
25

---

Published as a conference paper at ICLR 2026
0.8
0.0
Analytical
SESaMo (ours)
0.6
Breaking catio R
-0.4
ESS
SESaMo (ours)
0.4
Canonicalization
0.6
0.2
-0.8
0.0
1.0
0.00
0.05
0.10
0.15
0.20
0,000
0.002
0.004
0.006
0.008
0.010
Breaking factor
Breaking factor
(a) ESS as function of the breaking factor a.
(b) R(aa) as function of the breaking factor a.
Figure 12: Left: ESS for different values of the breaking factor a. The blue and yellow markers refer canonicalization and SESaMo, respectively. Mean and standard deviations are computed by averaging three models (for both approaches) trained with three different seeds. Right: Breaking ratio R for different values The analytical (yellow) curve is obtained by plotting Eq. (58) of a. The numerical estimate in Eq. (40), obtained with a trained SESaMo model, is compatible with the analytical result within errors. The uncertainties (sometimes too small to be visible in the scale of the plot) are estimated by averaging three models trained with three different seeds.
function
Yellow and blue markers indicate results from SESaMo and canonicalization, respectively. Error bars represent standard deviations computed from three independently trained models with different seeds. Crucially, while the performance of SESaMo and canonicalization comparable of canonicalization drops to zero increases, and the potential becomes increasingly asymmetric. In contrast, SESaMo maintains a stable ESS across the entire range of a, thanks to the stochastic modulation enabled by the learned symmetry-breakee parameter.
0, the ESS
Interestingly, this analysis can be made fully quantitative. The distribution of the magnetization for the o² theory (see Fig. I1] yields Gaussian distribution with two modes located at the VEVs= =u, and is modulated by the gymmetty--beakinng factor a,
f(x) A A(c
( 2
(2tp)
Re
-Vr
(56)
where V = 16 is the volume of the lattice. The parameters {A, o, μ²) can be inferred from numerical of the histogram at a = 0 (see Fig. [1b). yielding
A 0.499(2),
p 2.126(3),
0.629
These parameters fully characterize the distribution defined in in Eq. 55 To quantify the effect of symmetry breaking, we define (a) and (ax) as the integrated probability mass over the right and left modes, respectively:
f dx fx
N (a)=
and
Na - dx fa(
(57)
We then define the breaking ratio A as the relative imbalance between the two modes N4 and N_. Using standard Gaussian integrals, this ratio can be computed analytically, resulting
(a) (N
Von [1 erf (r- (a))] + Yar [1 + erf (T+(a))]
R(a)= N+ (a) + (a)
(58)
2 cosh( Vap)
where (a) are defined by
Ta 22
(Va= 11 o2)
(59)
The analytical result from Eq. 588 can be compared to the numerical estimate from Eq. 40
Fig. 10
(right) shows both the analytical prediction and the numerical estimate for the ratio R(a). ) for breaking
factors a e [0, 0.01]. The theoretical value in Eq. (58) and the numerical estimate in Eq. 440). for
different a, represented with solid (black) line and (blue) markers, respectively. For 0. the
26

---

Published as a conference paper at ICLR 2026
Ve
SESaMo (ours)
Ve
SESaMo (ours)
Im
Im
Re[A]
Re[x]
Re[i] (b) 0.005
Re[x]
(a)
Figure 13: Continuous Symmetries: Density plot for real and imaginary components the complex- valued fields of complex o² scalar field theory, introduced in Sec. 4 and sampled from trained generative models, i.e., RealNVP and SESaMo. The models have been trained to sample from the target density in Eq. (1) for lattices of volume
8 8 and coupling values {k, }} {0.3, 0.022). The models are trained until convergence and the density plots are made by drawing M samples. The left and right plots refer to continuos (11 symmetries in the unbroken (a 0) and broken (a = 0.005) case, respectively. Note that canonicalization is not shown that approach is not capable handling continuos symmetries.
1.0
1.0
0.6
ESS
SESaMo (ours)
ESS
SESaMo (ours)
ReaINVP
RealNVP
0.4
0.2
0.2
0.0
0.0
150
200
0
50
150
200
GPU
GPU minutes
(a) 0
(b) 0.005
Figure 14: ESS as a function of the GPU training time (minutes) for the o² theory experiments (see Fig. 13. The solid line and the shadows represent the mean and the standard deviations of models trained with different seeds. The curve shows the substantially faster convergence of SESaMo compared naive ReaINVP. Again, canonicalization is not shown as cannot straightforwardly incorporate continuos symmetries into the model.
estimated ration from the model is zero, suggesting that the fact the E₂ symmetry been correctly learned by SESaMo. By increasing the ratio R converges to -1, which corresponds fully broken Zg symmetry i.e., that is, the probability for x > is zero. Crucially, SESaMo is able to always learn the correct breaking parameter, hence estimating the correct breaking rutio R for broken Zz-symmetric action.
broken has
With this simple example, we conclude that SESaMo is capable of incorporating symmetries inside a flow-based generative model even when those are broken. One could foresee the power of this approach in incorporating other types of broken symmetries, such as the chiral symmetry breaking in quantum chromodynamics (QCD) (75). The QCD Lagrangian with two flavors, i.e., up and down quarks, has broken chiral symmetry due the different masses of the up and down quarks. Results for toy model of such scenario were presented in the main text (see Tab. I and we further elaborate on them in App. H.4 below.
27

---

Published as a conference paper at ICLR 2026
Experiment
Ne
Ng
NN
Activation
NB
LR
Steps
p
Var
GMM
6
40
ReLU (TanlL)
8k
5 x 10-4
10k
I
(20)
Complex 64
6
4
100
ReLU
8k
5 x 10 4
400
0
1
Hubbard 2 x 1
6
4
40
ReLU (Tanh")
8k
5 x 10 4
6k
18
Hubbard 18 100
10
100
ReLU (Tanh*)
8k
x 10
20k
0.18
Table 3: Hyperparameters for the Gaussian mixture model (GMM), the complex 6² theory and the Hubbard model. Shown are the number of couplings No. number of layers Nr. number of neurons per layer Nv. actívation function, batch size NB. learning rate (LR), training steps epochs, and the mean 1l and variance of the prior Gaussian distribution.
H.4 THE COMPLEX o$ SCALAR THEORY IN TWO DIMENSIONS
In light of these considerations, in the main text (see Tab. tested how SESaMo is capable of deal- ing with continuous (broken and unbroken) symmetries, and we showed a remarkable outperformance compared naive RealNVP model. Furthermore, in App. we discussed the details of SESaMo when dealing with continuos symmetries. In this section we complement the results from the main text with some further insights. First, Fig. shows the density of the real and imaginary components of the complex fields x E C summed across the lattice volume, i.c., Re[ã] = Evev Re( and Im[2] = Eiev Im [. Fig. 13a shows a ring-shaped potential projected on the complex plane stemming from the spontaneous symmetry breaking of an exact (1) symmetry in the of theory, which leads to the emergence of Goldstone Bosons (see (76) and Fig. therein). When the U (1) symmetry itself is broken (a 0.005), the probability density around the ring is no more evenly distributed, as it is visualized in the density learned by SESaMo in Fig. [3b] The reader should note that crucially, in the setting of continuous symmetries, only naive RealNVP and SESaMo can be applied. Indeed, the canonicalization approach could not straightforwardly be applied.
Fig. 13 demonstrates the greater capability of SESaMo to incorporate exact and broken continuous symmetries to enhance the model training and convergence. Moreover, this is further confirmed by the speed of convergence to a relatively high ESS function of training time, as shown in Fig. [4 After only seven minutes of training (one a single A100 NVIDIA GPU), the ESS achieved by SESaMo already surpasses 60% for both bias induced by stochastic modulation, struggles to learn meaningful of the target density. The low ESS reflects this failure in learning the target probability density, also shown in the RealNVP plots from Fig. 13
0 and 0.005. In contrast, ReaINVP, lacking the inductive
I DETAILS OF NUMERICAL EXPERIMENTS
In this section, we present details of the numerical simulations and the hyperparameters used in the main paper. All NFs are trained single A100 NVIDIA GPU, using floating precision. For the Hubbard model, however, double precision is used to ensure numerically stable estimation of the fermion determinant. The Adam optimizer (77) is employed with a learning rate of 5 x 10-4 Additionally, a learning rate scheduler is used: if the standard deviation of the loss has not changed over the last 2000 epochs, the learning rate is multiplied by a factor of 0.92. The learning rate is bounded from below at 1 x 10-6 As discussed in Sec. 3.2] the gradients are estimated with the REINFORCE algorithm for all experiments with SESaMo. For the VMoNF method, the number of NFs is chosen accordingly to the number of symmetry sectors, i.c., four in the Hubbard case and eight for the Gaussian Mixture Model. For the o² theory we chose eight NFs there are no clear symmetry sectors for the (1) symmetry. The remaining experiment-specific hyperparameters are summarized in Tab. 3 For FAB, the hyperparameters are chosen similarly, while for the Hubbard 18 x 100 system they are reduced to enable training on single GPU. The number of training steps is chosen to match the training time of the other architectures.
For VMoNF it was used Tanh, due to numerical instabilities.
Thii variance was only used for the RealNVP and VMoNF models to alleviate mode-dropping.
28

---

Published as a conference paper at ICLR 2026
LARGE LANGUAGE MODEL USAGE
This work was partly supported by large language models. The usage contains polishing previously written text, discovering related work and assisting coding.
29


Reviews:
[{"review_id": "3cfq6vyrIr", "reviewer": "Reviewer_HyQJ", "summary": "The paper presents a sampling method that handles symmetries in unnormalized distributions in a principled way. It maps the random tensor of interest $x$ to a reduced space $\\tilde x$. It learns a Normalizing Flow model to sample $\\tilde x$ and then map it, along with a latent variable $u$, back to $x$.\n\nExperiments are conducted on several distributions to validate the effectiveness of the proposed method.", "questions": "- Please answer my concerns raised in the Weaknesses section.\n- It seems the paper considers only discrete symmetries as $u$ is discrete, and not continuous ones. But the experiments are performed for U(1) symmetry too. What is $u$ in that case?\n- Although bijectivity has been enforced by a penalty term, its effectiveness has not been assessed. Experiments could be added to address this.\n- Why do we need to learn $C_{T,z}$? Can't it be a heuristic map, or can't we sample $\\tilde z$ directly? I would appreciate it if another example(s) could help me understand better.\n\n- Minor typos:\n  - Line 228: It think it should be $g_\\theta(z)$ instead of just $g_\\theta$.\n  - Line 283: a -> an", "limitations": "", "rating": 2, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "- The paper introduces a systematic way to incorporate symmetries in NF.\n- Experimental testbed is good. The complex $\\phi^4$ and Hubbard models are known to be challenging distributions.\n- The idea of using a penalty term to enforce bijectivity is interesting.\n- The paper is well written and is mostly comprehensible.", "weaknesses": "I think several essential questions/issues are overlooked both in theory and experiments.\n- Experimental results report only ESS and RKL (reverse KL divergence). However, the performance on two metrics can be good (i.e., high ESS and low RKL) even in the case of mode collapse (see [AdvNF 2025]), i.e., when the model learns and explores only a part of the target distribution. A standard metric robust to this problem is the negative log-likelihood (NLL).\n- The baselines are weak. The state of the art baselines such as [FAB 2023] and [IDEM 2024] should be used.\n- Are all $S_{T,u}$ independently trained for each $u$? If no, this could be a problem because the symmetry breaking bias may cause a probability mass drift even within a mode (of symmetry). If yes, the number of parameters involved will shoot up, and this should be discussed in the paper (comparing model sizes and training times of different methods/models). Or is the model conditioned on $u$ (i.e., $u$ is an input)?\n\nIn the core, the proposed approach factorizes $q(x)$ as $\\sum_u q(x|u)p(u)$, where $u$ is a random variable corresponding to a symmetry mode. Theoretically, the contribution seems modest. More rigorous experiments could make the paper more appealing.\n\nThe presentation could be simplified to highlight the core contribution. There are too many variables and equations, and many are superfluous.\n\n\n[AdvNF 2024] AdvNF: Reducing mode collapse in conditional normalising flows using adversarial learning, SciPost Phy. 2024\n[FAB 2023] FLOW ANNEALED IMPORTANCE SAMPLING BOOTSTRAP, ICLR 2023\n[IDEM 2024] Iterated Denoising Energy Matching for Sampling from Boltzmann Densities, NeurIPS 2024", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "30WyfhpQde", "reviewer": "Reviewer_goWZ", "summary": "This paper introduces SESaMo (Symmetry-Enforcing Stochastic Modulation), a simple and general framework for incorporating both exact and broken symmetries into normalizing flows (NFs). The method augments a standard flow with a symmetry transformation whose selection probabilities are learnable, allowing the model to capture multiple symmetric modes and symmetry-breaking effects without having to enforce symmetries as hard constraints within the architecture. Experiments on Gaussian mixtures and lattice field theory benchmarks ($\\phi^4$ and Hubbard models) demonstrate improved sample efficiency and effective sample size compared to normalizing flow baselines, including with canonicalization.", "questions": "* If I understand correctly, the stochastic variable + REINFORCE estimator used in the modulation map is likely to increase variance. Furthermore, since the approach still minimizes reverse KL, the approach remains mode-seeking and it seems possible (with poor initialization or high-variance training) that a mode may be entirely missed. Did the authors observe this potential issue?\n* Did the authors evaluate the method on datasets beyond toy or physical models where symmetry-enforcement is important, e.g. rotation/reflections with natural images?", "limitations": "", "rating": 4, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "* The proposed method is a simple and flexible way to impose both exact and broken symmetries on NFs without requiring explicitly equivariant architectures or canonicalization.\n* The work is technically sound and supported with solid empirical validation.\n* Furthermore, the proposed method integrates naturally with standard NFs (e.g. RealNVP) and requires relatively minimal overhead (a REINFORCE estimator to enable gradients through a stochastic variable).", "weaknesses": "* While the proposed method is elegant in principle, the exposition is overly abstract which makes the main idea underlying the method difficult to understand. For example, the authors provide several intuitive examples in increasing generality, but relegate them to the appendix (E, F). The paper could benefit from a clear algorithm box/pseudocode implementation describing the key components (flow, modulation map, and probability weights) and, in increasing generality, how each can be set or augmented with learned components. I believe this would make the work much more accessible.\n* The paper presents the method and empirical results, but could benefit from a formal analysis of conditions for convergence, expressivity, or relations to existing equivariant NFs. For example, as far as I understand, the REINFORCE estimator and stochastic variable introduced for the modulation map is likely to introduce additional variance, but it is unclear to what extent this affects the training stability or the consistency of the results (see questions.) Additional experiments that quantify the variability of the results with SESaMo may further strengthen the paper.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "cTeSdU65mZ", "reviewer": "Reviewer_jKsj", "summary": "This paper introduces SESaMo (Symmetry-Enforcing Stochastic Modulation), a framework that integrates group symmetry constraints into normalizing flows via a stochastic modulation mechanism. Specifically, the method augments a base flow g_\\theta with a family of symmetry transformations S_u, where u is sampled from a learnable distribution p_{S,b}(u). This allows the model to both enforce exact symmetries and capture symmetry breaking. The approach is evaluated on several discrete-symmetry benchmarks, showing improved sampling performance and better mode coverage compared to RealNVP, canonicalization-based flows, and VMoNF.", "questions": "1. Could the authors clarify whether recent equivariant generative models cannot be directly applied to the symmetry-enforcing scenarios considered here? If so, please explain the technical limitations preventing their use. Otherwise, their inclusion as baselines would significantly strengthen the experimental evidence.\n\n2. Could the authors discuss the potential of extending SESaMo to more complex or continuous symmetry groups?\n\n3. The bijectivity penalty appears heuristic. How sensitive is the training process to the choice of its parameters (A and B)? Have the authors tried alternative formulations?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper proposes a creative mechanism to integrate group symmetries into flow-based generative models using stochastic modulation. This differs from previous deterministic equivariant designs (e.g., Equivariant Flows) by introducing a learnable mixture structure that explicitly models symmetry breaking. Experiments demonstrate that SESaMo achieves competitive or superior ESS on discrete-symmetry datasets while maintaining interpretability of the learned symmetry-breaking parameter b.", "weaknesses": "1. The bijectivity penalty used to enforce separability is heuristic and not theoretically justified. The experimental comparisons only include RealNVP, VMoNF, and canonicalization methods. These are several years old and do not represent the current state of equivariant or symmetry-aware generative modeling. Modern approaches, including equivariant flow mathcing and equivariant diffusion models, should at least be discussed or justified as not directly applicable. Without this clarification, it is unclear whether SESaMo’s advantage is due to the stochastic modulation itself or simply the choice of older baselines.\n\n2. The method assumes that each symmetry sector maps to a distinct, non-overlapping region of configuration space, which ensures bijectivity but restricts applicability to simple discrete symmetries.\n\n3. The bijectivity penalty used to enforce separability is heuristic and not theoretically or experimentally justified.\n\n4. The experiments focus on low-dimensional or lattice-based toy systems. There is no demonstration on realistic and high-dimensional systems, where the claimed symmetry advantages would be most meaningful.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "bjCRaDyyyy", "reviewer": "Reviewer_Kqpf", "summary": "The paper introduces a new way to incorporate symmetries into discrete normalizing flows, especially suited for symmetry breaking. In this scheme the NF always maps to the same symmetry sector (canonical mode or cell). Samples are then possibly transformed to one of the other symmetry sectors (picked stochastically according to learned probability). To allow gradient flow to the parametrization of this “stochastic modulation”, the authors introduce the “self-reparametrized KL divergence” (bijectivity penalty; REINFORCE estimator to differentiate through the random variable). Since the learned probabilities can differ between the modes, the method allows for symmetry breaking, i.e. different mass per mode. The authors test on a Gaussian mixture with exact and broken Z8 symmetry, 2D complex $\\phi^4$ scalar field with U(1) symmetry, and Hubbard model with Z4 symmetry.", "questions": "- Appendix A and B (especially figure 5) are quite crucial to understanding the method. In my opinion they should be in the main text, possibly at the cost of moving parts of the background into the appendix\n- How does accuracy/training complexity scale with dimension and number of symmetry sectors?\n- Can stochastic modulation be used with continuous NF (like flow matching)?", "limitations": "", "rating": 6, "confidence": 2, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "- The idea is technically sound and novel\n- The paper shows clear benefits in terms of effective sample size (ESS) and lower KL on the tasks considered, especially for symmetry breaking", "weaknesses": "- The proposed stochastic modulation seems to only work for discrete symmetries\n- The symmetry sectors (and the corresponding transformations from the canonical cell) must be known a priori", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)