CanonFlow: Canonical Generative Models

• CanonFlow is a family of generative modeling methods that impose canonical structure on latent representations through local orthogonality, sparsity, and symmetry invariance.
• It integrates canonical regularization into manifold flows to yield efficient statistical representations, improved sample quality, and accelerated learning.
• Applied to molecular graph generation, CanonFlow achieves state-of-the-art performance by reducing learning variance and enhancing stability via structured canonicalization.

CanonFlow is a family of generative modeling methods that systematically impose canonical structure on latent representations or data, in either manifold learning flows or symmetry-invariant diffusion and flow models. The CanonFlow paradigm enforces either local orthogonality and sparsity (for latent manifolds) or global orbit representativeness (for group-invariant structured data such as molecular graphs), thus yielding more efficient statistical representations, improved sample quality, and tractable or accelerated learning objectives. CanonFlow frameworks have been demonstrated for both general manifold flows (Flouris et al., 2023) and molecular graph generation under $S_n \times SE(3)$ symmetries (Zhou et al., 16 Feb 2026).

1. CanonFlow in Manifold Flows: Canonical Intrinsic Bases

Manifold learning flows (MLFs) generalize normalizing flows (NFs) by learning injective (non-bijective) mappings $g_\phi: \mathbb{R}^d \to \mathbb{R}^D$ ($d < D$), thus restricting the data support to a low-dimensional manifold $\mathcal{M} = g_\phi(\mathbb{R}^d)$. The induced density on $\mathcal{M}$ is tractable via the Riemannian metric $G(z) = J_{g_\phi}(z)^\top J_{g_\phi}(z)$, where $J_{g_\phi}$ is the Jacobian, and the likelihood is

$$p_X(x) = p_Z(g_\phi^{-1}(x)) \cdot |\det G(g_\phi^{-1}(x))|^{-1/2}.$$

A reconstruction loss $\|x - g_\phi(g_\phi^{-1}(x))\|_2^2$ ensures alignment of the learned manifold to the observed data. In unconstrained settings, the latent dimensions of MLFs tend to form degenerate, entangled bases in which information is redundantly spread, reducing latent efficiency and interpretability (Flouris et al., 2023).

CanonFlow resolves this by introducing a canonicality-promoting regularization: $$R_\mathrm{can}(z) = \| G(z) - \operatorname{diag}(G(z)) \|_1 = \sum_{i=1}^d \sum_{j \ne i} |G_{ij}(z)|.$$ Minimizing $g_\phi: \mathbb{R}^d \to \mathbb{R}^D$0 simultaneously enhances local orthogonality (forcing $g_\phi: \mathbb{R}^d \to \mathbb{R}^D$1, $g_\phi: \mathbb{R}^d \to \mathbb{R}^D$2) and sparsity (shrinking small tangent directions), yielding a "canonical intrinsic basis" that organizes information into prominent, non-redundant coordinates.

2. CanonFlow Loss Function and Training Protocol

For manifold learning, the CanonFlow objective combines likelihood, reconstruction, and canonical basis regularization: $g_\phi: \mathbb{R}^d \to \mathbb{R}^D$3 where $g_\phi: \mathbb{R}^d \to \mathbb{R}^D$4 and $g_\phi: \mathbb{R}^d \to \mathbb{R}^D$5 are positive hyperparameters balancing these terms. Log-determinants are efficiently approximated using the Hutchinson trace estimator with conjugate-gradient solvers. Differentiation through $g_\phi: \mathbb{R}^d \to \mathbb{R}^D$6 proceeds via standard autodiff. Training uses Adam or AMSGrad optimizers on minibatches (Flouris et al., 2023).

3. CanonFlow for Symmetry-Invariant Generative Modeling

In the context of group-invariant distributions, such as molecules under $g_\phi: \mathbb{R}^d \to \mathbb{R}^D$7, CanonFlow employs a canonicalization map $g_\phi: \mathbb{R}^d \to \mathbb{R}^D$8 selecting a unique orbit representative per input. For structured data $g_\phi: \mathbb{R}^d \to \mathbb{R}^D$9 (coordinates, atomic features, adjacency tensors), $d < D$0 consists of a spectral canonical permutation (via the signed Fiedler vector of a random-walk Laplacian on a graph built from pairwise distances and bond structure) and a canonical $d < D$1 frame defined by optimized landmarks in the sorted coordinate array.

Generative modeling (e.g., score-based diffusion or flow-matching) is trained on the canonical "slice" $d < D$2, and symmetry invariance is restored at generation time by randomizing the output over the group using Haar sampling. Canonicalization removes the mixture complexity in the learning signal, resulting in score and flow fields with reduced variance and improved smoothness. The approach is universally expressive for all group-invariant distributions (Zhou et al., 16 Feb 2026).

4. CanonFlow Architectures and Algorithmic Details

Manifold Learning

• Embedding parametrization: $d < D$3, with RealNVP couplers applied on both latent and ambient segments.
• No batch-norm: Avoided to prevent interference with vector-Jacobian-product operations.
• Hyperparameters: For images, learning rate $d < D$4, batch size 120, $d < D$5 and $d < D$6 (images), up to $d < D$7 (tabular data). Ten layers in $d < D$8-flow, $d < D$9 hidden units per coupler net. Tabular experiments use $\mathcal{M} = g_\phi(\mathbb{R}^d)$0, between 4–10 layers, and fitted batch sizes per dataset.

Symmetry-Invariant Models

• Backbone: SemlaFlow with an added "rank stream" incorporates canonical information directly.
• Inputs: Canonicalized coordinates, features, and normalized permutation ranks (with optional sinusoidal positional encodings).
• Operator blocks: Three parallel attention streams update coordinates, features, and rank; edge messages constructed using projected coordinate sets and node features.
• Output: Updated coordinates, atom/bond logits, and optional rank prediction.
• Training: Canonicalize each input, select appropriate priors (aligned Gaussian or isotropic), pair initial/final states by product or approximate optimal-transport coupling, sample interpolation time, optimize flow-matching loss, and (if required) cross-entropy losses.
• Sampling: Reverse ODE/SDE steps followed by optional projection back to the slice and Haar-randomization for invariance.

5. Empirical Performance and Applications

Manifold Flows

Experiments on synthetic geometries, image datasets (MNIST, FMNIST, Omniglot, CIFAR-10, SVHN, CelebA), and tabular UCI benchmarks (POWER, GAS, HEMPMASS, MINIBOONE) yield the following:

Dataset/Metric	MFlow	RNF	CMF (CanonFlow)
MNIST FID (d=10)	77.2	68.8	63.4
FMNIST FID (d=10)	663.6	545.3	529.8
Omniglot FID (d=10)	174.9	150.1	147.2
CIFAR-10 FID (d=30)	541.2	544.0	532.6
SVHN FID (d=30)	102.3	95.6	88.5
CelebA FID (d=30)	—	9064	9060

CMF consistently outperforms rectangular normalizing flow (RNF) and unconstrained manifold flow (MFlow), especially with a limited latent budget. CanonFlow identifies and orders the few active latent dimensions with minimal degeneracy, while pruning redundant coordinates without quality loss (Flouris et al., 2023).

Molecular Graph Generation

On QM9 and GEOM-DRUG, CanonFlow with the "Canon" architecture and flow-matching achieves state-of-the-art values in atom/molecule stability and validity, and enables high-quality, few-step sampling:

Method	Mol Stab↑	Valid↑	Opt-RMSD↓	NFE
CanonFlow (GEOM-DRUG, OT-anneal)	98.4	95.9	—	100
CanonFlow (QM9, full Canon)	99.8	99.0+	0.15	100

Few-step CanonFlow retains 97.5% molecular stability and 94.9% validity in 50 denoising steps—nearly matching full runs, but with an order-of-magnitude reduction in compute (Zhou et al., 16 Feb 2026).

6. Theoretical Insights and Practical Impact

CanonFlow achieves universality for all invariant distributions via canonicalization plus Haar-randomization. Canonical basis penalties in latent manifolds yield "automatic relevance determination," using only as many coordinates as required by the intrinsic data geometry. In symmetry-invariant modeling, canonicalization removes group-induced mixture complexity from the learning objective, yielding smoother scores, reducing variance in flow-matching, and accelerating convergence. Alignment of priors (e.g., Gaussian matched to data mean/covariance) and near-Monge OT couplings further improve training robustness and sample quality.

A measurable cost is a single eigen-decomposition per sample (for the graph Laplacian in molecular data), which is outweighed by gains in learning and sample quality (Zhou et al., 16 Feb 2026).

7. Relation to Prior Work and Conclusion

CanonFlow synthesizes advances in manifold learning flows and group-invariant generative models. In manifold flows, it extends the approach of rectangular normalizing flows (RNF) by imposing local canonicality regularization. In symmetry-invariant flows and diffusion models, CanonFlow departs from equivariance-by-design, instead leveraging canonicalization and post-hoc group randomization.

The framework demonstrates that canonicalization can simplify learning, reduce sample inefficiency, and unlock expressive model architectures without sacrificing invariance. CanonFlow represents a unified method for inducing efficient latent structure or group-invariance in flows and diffusion models, with empirical state-of-the-art performance across diverse structured data regimes (Flouris et al., 2023, Zhou et al., 16 Feb 2026).