Joint Flow Matching for Generator-Consistent Classification

2026-07-27Machine Learning

Machine Learning
AI summary

The authors introduce Joint Flow Matching (JFM), a new training method for models called continuous normalising flows that handle several variables together. Unlike usual methods that treat all variables the same over time, JFM assigns each variable a different role at the start and end, allowing both forward and backward predictions from the same model. They prove this method creates a consistent joint distribution useful for tasks like classification and generation, and show it works well on conditional datasets with accurate confidence scores and consistent image generation. Their approach links generation and interpretation in a unified way.

Joint Flow MatchingContinuous Normalising FlowsConditional InferenceJoint DistributionForward and Reverse IntegrationDiscriminative ModelsGenerative ModelsClassifier CalibrationImage GenerationConditional Datasets
Authors
Hayden McAlister, Lech Szymanski
Abstract
We introduce Joint Flow Matching (JFM), a training framework for continuous normalising flows over multiple variables. Standard flow matching transports variables from noise to data simultaneously, offering no natural mechanism for forward and reverse conditional inference from a shared joint model. JFM resolves this by assigning opposite roles to each variable at the temporal endpoints. We prove that JFM produces a consistent joint distribution where that forward or reverse integration are conditionals of the same joint. We explore this consistency in the context of joint classification and generation as the basis for interpretability in discriminative-generative models. We validate JFM on conditional datasets producing competitive accuracy with inherently well-calibrated confidence scores without post-hoc calibration, and classifier-consistent image generation.