Tensor Gauge Flow Models

TL;DR

This work introduces Tensor Gauge Flow Models (TGFM), a unified framework that extends Gauge Flow Models and Higher Gauge Flow Models by embedding higher-order Tensor Gauge Fields into the neural ODE governing generative flows. TGFM uses a gauge-corrected velocity term, $v_{\\theta}$, and a tensor gauge contribution driven by a Tensor Gauge Field $\\mathcal{A}_{\\mu_{1}...\\mu_{n}}$ acting on a Tensor Field $\\hat{T}$, projected to the tangent bundle via $\\Pi_M$, within a Riemannian Flow Matching training regime. The authors provide a rigorous mathematical foundation on Tensor Fields on Fiber Bundles, describe concrete model variants, and demonstrate improved generative performance on synthetic Gaussian mixtures compared to baselines, with similar parameter counts. The results suggest that enriching local geometric structure with higher-order gauge data can yield more expressive and robust continuous-time generative models, potentially benefiting structured-data applications and geometry-aware learning.

Abstract

This paper introduces Tensor Gauge Flow Models, a new class of Generative Flow Models that generalize Gauge Flow Models and Higher Gauge Flow Models by incorporating higher-order Tensor Gauge Fields into the Flow Equation. This extension allows the model to encode richer geometric and gauge-theoretic structure in the data, leading to more expressive flow dynamics. Experiments on Gaussian mixture models show that Tensor Gauge Flow Models achieve improved generative performance compared to both standard and gauge flow baselines.

Figures (4)

• Figure 1: A Tensor Gauge Flow Model can be visualized locally via its Tensor Gauge Field $\mathcal{A}^{\mu_{1}\dots\mu_{m}}{}_{\nu_{1}\dots\nu_{n}}{}^{a}$. The black region represents the base manifold $M$ on which the trajectory $x(t)$ evolves. Above each point of $M$, the red regions depict the fiber spaces associated with the indices $\mu_{1},\dots,\mu_{m}$ while the green region represents the internal vector space carrying the index $a$. The yellow regions illustrate the Tensor Field $\hat{T}(x(t),t)$ evaluated along the flow, which is acted on by $\mathcal{A}$ to produce the Tensor Gauge correction in the dynamics.
• Figure 2: Training loss comparison (lower is better). The training loss for each model is normalized by the loss of the PlainVF + Tensor Gauge Field model. Values are shown for several ambient dimensions $N$.
• Figure 3: Test loss comparison (lower is better). The test loss for each model is normalized by the loss of the PlainVF + Tensor Gauge Field model. Values are shown for several ambient dimensions $N$.
• Figure 4: Number of trainable parameters for each model family as a function of ambient dimension $N$.