Technical Summary: Zero-Flow Encoders

Problem Statement

While flow-based generative models (e.g., diffusion models, flow matching) have achieved significant success in capturing complex data distributions for generative tasks, their potential for unsupervised representation learning and structural inference remains underexplored. Existing unsupervised methods, such as contrastive learning (e.g., SimCLR), often rely on maximizing mutual information, which can lead to "shortcut solutions." These models may exploit superficial, instance-specific features (e.g., watermarks or random colors) rather than learning high-level semantic structures, failing to capture the nuanced details required for tasks like identifying Markov blankets or learning robust latent representations.

Methodology

The paper proposes Zero-Flow Encoders, a framework inspired by the properties of rectified flows to enforce conditional independence and extract sufficient information.

The Zero-Flow Criterion

The core theoretical insight is the observation of a "zero-flow" phenomenon in rectified flows trained with independent coupling.

• Theorem 3.1: For a rectified flow transporting a source distribution $p_X$ to a target distribution $p_{X'}$ using independent coupling, the velocity field $v_t$ at the midpoint $t=0.5$ is zero everywhere if and only if $p_X = p_{X'}$.
• Conditional Extension: The authors extend this to conditional distributions. They define a modified rectified flow objective where the velocity field $v_t$ is conditioned on an encoder output $f(Y)$. Theorem 3.4 establishes that $v_{t=0.5} = 0$ if and only if the conditional distributions are identical ($p_{X|Y} = p_{X|f(Y)}$), which is equivalent to the conditional independence condition $X \perp\!\!\!\perp Y | f(Y)$.

The Zero-Flow Loss

To make this criterion tractable for learning, the authors derive a simulation-free least-squares loss function (Equation 9):
$$\min_{u, f} \int_0^1 \omega(t) \mathbb{E}[\|u_t(X_t, f(Y), Y)\|^2] dt + \int_0^1 \mathbb{E}[\|X' - X - u_t(X_t, f(Y'), Y)\|^2] dt$$

• The first term (weighted by $\omega(t)$, peaking at $t=0.5$) acts as the zero-flow criterion, penalizing non-zero velocity at the midpoint to enforce conditional independence.
• The second term is the standard rectified flow loss, ensuring the velocity field learns the transport map.
• The encoder $f$ is trained to minimize this combined loss, effectively learning a representation that renders the target variable conditionally independent of the raw input given the representation.

Applications

The framework is applied to two specific unsupervised tasks:

1. Amortized Markov Blanket Learning: Instead of learning a separate model for each target variable subset, the encoder takes a mask $m$ as input to select a subset of features $Z_m$ (the Markov blanket) from the remaining features $Z_{-m}$. This allows for flexible inference on arbitrary target sets at test time without retraining.
2. Self-Supervised Representation Learning: The method is applied to multi-view data (e.g., augmented images). By enforcing $Z_1 \perp\!\!\!\perp Z_2 | f(Z_2)$, the encoder is forced to extract the shared semantic information $T$ while discarding view-specific noise, avoiding the shortcut problem inherent in mutual information maximization.

Key Contributions

1. Theoretical Discovery: Identification and mathematical proof of the zero-flow condition in rectified flows, establishing it as a necessary and sufficient condition for the equality of source and target distributions (both unconditional and conditional).
2. Non-Parametric Encoder: Development of a tractable, simulation-free loss function that learns sufficient encodings without explicit parametric assumptions about the underlying data distribution.
3. Amortized Inference: Introduction of an amortized strategy for Markov blanket selection, enabling the model to handle arbitrary target feature configurations at inference time, a capability difficult to achieve with traditional parametric graphical model estimators.
4. Robustness to Shortcuts: Demonstration that the zero-flow encoder avoids the "shortcut problem" common in contrastive learning, successfully learning semantic features even when superficial instance-specific cues are present.

Experimental Results

The paper evaluates the method on both simulated and real-world datasets:

• Graphical Model Recovery: On synthetic Gaussian, non-paranormal, and truncated-Gaussian graphical models, the zero-flow encoder (using MLP and LSTM architectures) outperforms Graphical Lasso and the PC algorithm in recovering the true graph structure (measured by AUC). The LSTM variant, leveraging temporal inductive bias, achieved near-perfect scores on Markov chains.
• Markov Blanket Generalization: The amortized encoder successfully identified Markov blankets for target variables seen during training and, crucially, for unseen target variables, demonstrating strong generalization.
• Real-World Time Series: Applied to S&P 500 data (2019–2020), the method detected a structural shift in market dependencies around March 2020 (COVID-19 outbreak), where the Markov blanket shifted from relying on past trading days to future days, indicating a market disruption.
• Self-Supervised Learning: On image datasets (Fashion-MNIST, CIFAR-10, STL-10, TinyImageNet, ImageNet-1K) with artificial "shortcut" watermarks:
  • SimCLR suffered severe performance degradation (loss saturation and accuracy drops) as it learned the watermarks.
  • Zero-Flow Encoders maintained stable loss and high linear probing accuracy, preserving semantic information (e.g., object shapes) despite the presence of shortcuts.
  • Reconstruction tasks showed that Zero-Flow encoders retained rich visual structures, whereas SimCLR reconstructions were dominated by the watermark artifacts.

Significance and Claims

The paper claims that the Zero-Flow framework provides a rigorous, non-parametric alternative to existing unsupervised learning methods. Its primary significance lies in:

• Theoretical Grounding: Offering a mathematically provable criterion for sufficiency and conditional independence derived from flow dynamics.
• Robustness: Solving the "shortcut problem" by directly enforcing conditional independence rather than maximizing mutual information, which is prone to greedily selecting easy, superficial features.
• Flexibility: Enabling amortized inference for complex structural tasks like Markov blanket discovery, allowing models to adapt to new target configurations without retraining.

The authors position this work as a step toward extending flow-based methods beyond generation into robust representation learning and structural inference, validated by competitive performance on both synthetic and real-world benchmarks.