From Probability Flows to Reusable Representations — ICML 2024, Part 4 of 6

Learning a whole distribution from sample-specific paths

Suppose you know how to transport samples along a probability path ending at one training example. How can those sample-specific paths produce a generator for the entire dataset? Part 4 of the ICML 2024 recording begins with this flow-matching problem. The answer is to average conditional velocities given the current state, weighting each possible data endpoint by its conditional probability. This average is the marginal velocity.

The continuity equation connects velocity to probability: density falls where probability flux has net outward divergence. To understand divergence, imagine a tiny neighborhood around a point. Count the mass leaving, subtract the mass entering, and shrink the neighborhood toward an infinitesimal region. Net outflow means a corresponding decrease in density.

The marginal probability path is a mixture of conditional paths. Applying the continuity equation to each path, then exchanging integration with the linear divergence operator, yields the continuity equation for the mixture. Averaging conditional velocities preserves probability transport, provided that exchange is valid.

This also gives a tractable training objective. A neural velocity field regresses onto the available conditional velocities. That objective has the same expected gradients as directly regressing onto the inaccessible marginal velocity. At the optimum, conditional regression recovers the marginal velocity needed for generation.

Crystals expose the limits of continuous embeddings

The next extension handles manifolds: spaces with local Euclidean approximations called tangent spaces. Velocities live in those tangent spaces, and Riemannian divergence replaces Euclidean divergence. The averaging argument still applies when divergence and integration can be exchanged. Constructing suitable conditional velocities, choosing volume elements, and handling boundaries remain separate mathematical work.

A materials-generation project at Meta makes this geometry concrete. A crystal is represented by a unit cell that repeats in every direction, giving periodic boundary conditions. The representation contains the cell, fractional coordinates locating atoms inside it, and atom types. Training on stable materials aims to produce novel stable candidates. The presenter informally connects stability with real-world synthesizability, while emphasizing the difficulty of checking it; these results are not experimental synthesis claims.

Periodicity complicates translation handling. For an ordinary point cloud, subtracting the mean and using zero-mean noise keeps the generative path centered. Periodic coordinates lack that same notion of an origin and zero mean. The project instead projects the velocity field to prevent collective mean motion. The geometry and the treatment of translations must therefore be designed together.

Results differed between two tasks: - Fixed atom types: generating continuous cell and coordinate variables reportedly modestly improved on diffusion baselines for finding stable conformations. - De novo generation: generating atom types as well performed only on par with some language-model baselines, which the presenter found disappointing.

The atom-type approach embedded discrete identities in continuous space, learned their generation with flow matching, and converted generated embeddings back into types. Its weak performance motivated a different approach: apply flow matching directly to discrete states, without assuming an embedding space or a distance metric.

Discrete flow moves probability, then samples a state

In discrete space, velocity changes probabilities rather than coordinates. Start with all probability at the current state, redistribute some mass, and sample the next state. Moving 20% of the mass means each particle moves with probability 20% and stays with probability 80%. This connects the construction to a continuous-time Markov chain, whose rate matrix describes transitions between states. The presenter credits overlapping prior work by Campbell and colleagues.

A discrete velocity must obey two constraints: - Nonnegative contributions elsewhere: states other than the current state can only gain probability from this update. - Zero total change: probability added elsewhere must be balanced by probability removed from the current state.

Discrete divergence measures outflow minus inflow using probability fluxes on graph edges. No geometric ordering of token values is required. To avoid modeling a complete graph over every possible sequence, the derivation assumes a single-token-change graph: flux is zero between states differing in more than one coordinate.

Despite that restriction, a finite sampling step can update coordinates in parallel. Under the discrete continuity equation, independently sampling each coordinate gives an Euler update that follows the probability path to first order. The omitted terms shrink faster than the step size h. The same continuity equation supports averaging conditional velocities into a marginal velocity, carrying the flow-matching recipe into discrete space.

One conditional path lets each coordinate independently choose between its source and target values, with a time-dependent mixture probability kappa(t). Making the source an entirely masked sequence produces gradual unmasking. But a probability path does not uniquely determine its velocity field, leaving room for different sampling behavior.

The presentation distinguishes two endpoint predictions: - Predict the target: transport forward toward completion; in the masked setting, this unmasks tokens. - Predict the source: transport backward; in the masked setting, this permits re-masking. Both predictors output full distributions. Combining a large forward step with a small backward step creates a predictor-corrector-style sampler: a revealed token can be reconsidered, re-masked, and generated again.

A 1.7-billion-parameter text-completion experiment did not meet the presenter’s goal of outperforming language models. Generating code from docstrings offered executable checks. The non-autoregressive design—generation without a required left-to-right order—also supported infilling conditioned on arbitrary portions of code, although the presenter identified a lack of a good benchmark for that flexibility.

On OpenWebText, a path mixing masked, uniform, and target distributions reportedly beat a purely masked path mainly with a small sampling budget; the purely masked approach remained competitive with more sampling computation. The proposed explanation is that parallel unmasking can produce mutually inconsistent choices. Corrector steps or intermediate states containing alternative tokens let the model repair them later.

A masked discrete image-generation experiment was reported as slightly worse than continuous flow matching. The talk closes with the shared recipe: construct conditional velocities along paths ending at target samples, learn the expected or marginal velocity, and insert it into the transport equations to recover the fitted distribution.

Rectified flow: allocate training effort along the trajectory

Patrick Esser turns to the cost of scaling image generation. Larger models, more examples, and more compute improve performance, he argues, but increase both training and sampling costs. His study examines three decisions: the generative formulation, the architecture connecting text and images, and whether validation loss predicts downstream quality.

Rectified flow uses linear interpolation between sample pairs as its forward process. Conditional regression learns a vector field connecting distributions such as Gaussian noise and images; sampling integrates that field. Straight trajectories would permit one-step integration, but straight conditional paths do not guarantee straight learned sampling paths. Esser reports that the learned fields empirically tend to have less curvature than diffusion-derived fields, making sampling more efficient.

Classical rectified flow samples training timesteps uniformly. Esser argues that optimal solutions at the endpoints involve comparatively simple estimates of distribution means. The team therefore explored emphasizing the trajectory’s interior, where perceptually relevant image properties emerge. Because training estimates the objective from sampled timesteps, their distribution determines where optimization effort goes.

Across 61 diffusion and flow-matching variants, uniform-timestep rectified flow performed strongly with few sampling steps but fell behind a strong baseline with more steps. Esser reports that rectified flow with a logit-normal timestep distribution led in both regimes, outperforming the existing variants evaluated in the study. Numerical scores and exact step counts are not given in the talk.

For text-image interaction, MM-DiT uses separate modality weights and joint attention. Text and images have separate weights while exchanging information through full joint attention. Esser reports that this transformer design performed best among comparisons that included directly concatenating the modalities and using cross-attention for text conditioning.

The scaling study measured validation loss at fixed timesteps. Loss comparisons were meaningful only within a single formulation. Within that scope, predicted improvements corresponded to better GenEval image-evaluation results and human preference ratings. Similar observations for video synthesis supported further scaling, though no numerical correlations are supplied.

The closing development lessons concern targeted fixes: - Training stability: the technique identified as QK normalization helped stabilize scaling runs; its mechanics are not explained here. - Resolution-specific timestep distributions: failing to retune them lost substantial performance. Recovering that performance through a larger model would be much more expensive. - Human preference alignment: alignment provided a reportedly quick, inexpensive improvement in preference scores. The resulting model was described as working well across resolutions and aspect ratios, with strong prompt understanding and spelling, supported by human comparisons but without spoken scores.

NaturalSpeech 3 separates voice attributes before generation

Zhentian Zhu introduces NaturalSpeech 3, a text-to-speech system that uses reference audio to synthesize an unseen speaker’s voice. In his comparison, traditional multi-speaker text-to-speech requires the target speaker in training. Zero-shot synthesis instead conditions on a short recording, drawing on training across diverse speakers and acoustic environments.

Zhu connects this change to larger acoustic models and learned codec representations of audio. Conventional residual vector quantization builds representations in a coarse-to-fine hierarchy, but he argues that those levels do not adequately separate speech’s timbre, content, prosody, and recording conditions. Modeling that mixture leaves problems in speaker similarity, quality, and delivery.

FA-Codec separates four speech attributes. It extracts a global timbre vector and uses three factorized vector quantizers to represent prosody, content, and acoustic details in separate subspaces. An information bottleneck limits each token’s capacity, while supervision encourages each representation to retain its intended attributes.

Generation proceeds through duration, prosody, content, and acoustic details in that order. Timbre comes directly from the reference audio. Each discrete diffusion process randomly masks sequence tokens during forward corruption and learns to recover them using context and conditioning.

Attribute prompts remain fixed during diffusion. They are prepended as conditioning prefixes. Taking all prompts from one recording supports zero-shot synthesis; selecting attributes from different recordings provides separate controls over the output.

Zhu reports stronger similarity, robustness, and overall quality than baselines on LibriSpeech and an emotional-speech dataset, including a claim of human-level naturalness without numerical scores or the naturalness evaluation protocol. He also reports high-fidelity FA-Codec reconstruction on LibriSpeech. Demonstrations use three-second reference clips to condition new text, then show sad, calm, and disgusted delivery.

The final demonstrations isolate duration control through the prompt. Slowing or speeding the duration prompt changes the timing of the same generated sentence. Another example substitutes a duration prompt from a different recording. The presenter claims that keeping the other prompts fixed preserves the other attributes; the demonstration does not quantify their independence.

DiffS4L: imperfect speech can still improve pretraining

Yang Zhang begins with a data shortage. Speech self-supervised learning, or SSL, pretrains representations on unannotated audio, then fine-tunes them on a smaller labeled corpus. Zhang argues that successful pretraining generally requires at least a thousand hours, which many languages cannot readily supply. DiffS4L augments a small corpus with diffusion-generated speech before applying standard representation pretraining.

The augmentation procedure has three stages: 1. Train initial representations on the small unannotated corpus. Zhang describes their quality as poor but sufficient for conditioning. 2. Train a diffusion reconstruction model using those representations and speaker embeddings extracted from the original recordings. This model also learns only from the small corpus. 3. Modify the conditioning and generate speech, seeking variation in content, speaker identity, and prosody. Use the expanded corpus for standard SSL pretraining.

The training mixture contains four levels of variation: - Level one: original recordings. Retain the real audio. - Level two: unchanged conditioning. Reconstruct the utterance, allowing small variations, especially in prosody. - Level three: changed speaker. Substitute the speaker embedding while retaining content conditioning. - Level four: new content. Change the speaker embedding and partially mask the speech representations, forcing the generator to fill in missing content.

In the level-two demonstration, Zhang points out a word whose intonation changes from falling to rising despite unchanged conditioning. Level four can sound like nearly nonsensical babble, which he interprets as incomplete learning of grammar and word structure. Yet linguistic plausibility is not required for every useful training sample: he reports that this output can still improve pretraining.

The experiments use two data budgets: - Low resource: 100 real hours become 960 total hours. Add 430 hours from levels two and three combined, plus 430 hours of level-four speech. - High resource: 960 real hours become 2,400 total hours. Start from the full LibriSpeech corpus and add synthetic speech. Both initial and final representation training use the standard wav2vec 2.0 or HuBERT methods.

With ten labeled hours for English automatic speech recognition fine-tuning, Zhang reports substantial error reductions for both methods and a small gap between pretraining on 960 augmented hours and 960 real hours. He also reports the best performance on almost all eight evaluated non-recognition tasks and consistent error reductions across thirteen additional languages. No numerical error rates are spoken.

The closing experiment varies mixture composition at a fixed data budget. With real speech held at 100 hours and total speech at 960 hours, an intermediate balance of levels two/three and level-four babble performs best on English recognition. All four levels are present in that mixture. Zhang concludes that diffusion captures speech information complementary to SSL, with benefits even from seemingly nonsensical output.

DeCAF: a better representation changes the baseline

Trevor Darrell accepts the 2024 ICML Test of Time award for DeCAF—A Deep Convolutional Activation Feature for Generic Visual Recognition—ten years after its original presentation. He retrospectively describes it as an early vision foundation model: a pretrained model whose representations proved useful across many tasks.

The recipe was simple: freeze AlexNet, extract activations, train a linear classifier. The activations—the representations computed inside the pretrained network—became features for a new task. Darrell reports state-of-the-art performance across a wide range of tasks without retraining the feature extractor.

Darrell considers open-source release and broad dissemination DeCAF’s most important contribution. He sees its recognition as a shift toward valuing accessible tools alongside new algorithms and models. The related offerings served different purposes: - DeCAF supplied frozen pretrained features that users could apply to new tasks. - Caffe let users train their own models, with CPU and GPU execution and a codebase usable in academia and industry.

He describes Caffe as prominent from 2013 to 2018 and points to three 2024 Test of Time honors across DeCAF, R-CNN, and the Caffe system paper, with the last still to be presented later that year. In retrospect, he suggests DeCAF may be the most consequential of the three, despite receiving less recognition initially.

Returning to the earlier research context, Darrell contrasts handcrafted image features with learned convolutional representations. He traces the latter through Fukushima’s neocognitron, work by Rumelhart, Hinton and Williams, and Yann LeCun’s models, with AlexNet’s 2012 result demonstrating that the approach could scale. Even then, he recalls skepticism that object-recognition success would extend to fine-grained recognition, segmentation, or difficult transfer problems.

DeCAF challenged that skepticism. Darrell interprets its transfer performance as evidence of semantic structure beyond the training labels. Different layers transferred differently. Higher representations captured broader categories, called “super labels” in the presentation, without explicit supervision for those categories. This helped explain why frozen features could serve tasks beyond the one used for pretraining.

The closing comparison concerns domain adaptation on the Office dataset. Replacing handcrafted SURF features with learned representations reportedly doubled performance figures for some sophisticated adaptation methods and raised simpler baselines by factors of three or four. These retrospective ratios come without named metrics, absolute scores, or precise transfer settings. The consequential observation is that representation transfer could outweigh elaborate adaptation machinery.

Part 4 ends as Darrell introduces individual bird-species recognition as another example of fine-grained recognition where learned representations proved useful. The specific result lies beyond this part.