Gauge Flow Models

TL;DR

Gauge Flow Models are introduced, a novel class of Generative Flow Models that incorporate a learnable Gauge Field within the Flow Ordinary Differential Equation (ODE) and have potential for enhanced performance across a broader range of generative tasks.

Abstract

This paper introduces Gauge Flow Models, a novel class of Generative Flow Models. These models incorporate a learnable Gauge Field within the Flow Ordinary Differential Equation (ODE). A comprehensive mathematical framework for these models, detailing their construction and properties, is provided. Experiments using Flow Matching on Gaussian Mixture Models demonstrate that Gauge Flow Models yields significantly better performance than traditional Flow Models of comparable or even larger size. Additionally, unpublished research indicates a potential for enhanced performance across a broader range of generative tasks.

Figures (5)

• Figure 1: A fiber bundle can be visualized as a base space with a fiber attached to each point. As shown, the base space corresponds to the plane and the blue lines represent the fibers.
• Figure 2: A connection on a fiber bundle defines a mapping that relates each fiber to its infinitesimally near neighboring fibers. This enables local ‘parallel transport’ within the bundle.
• Figure 3: Train Loss Comparison (lower is better): The normalized train loss is plotted against the dimension $N$ to compare the performance of the Gauge Flow Model using directional vector fields $v_{\theta}(x(t), t)$ and $v_{s}(x(t), t)$ with the standard plain Flow Model.
• Figure 4: Test Loss Comparison (lower is better): The normalized test loss is plotted against the dimension $N$ to compare the performance of the Gauge Flow Model using directional vector fields $v_{\theta}(x(t), t)$ and $v_{s}(x(t), t)$ with the standard plain Flow Model.
• Figure 5: Number of Parameters: A comparison of the number of parameters for the plain Flow Model and the Gauge Flow Models across various dimensions $N$.