Flow Matching Sampling in TRELLIS.2: How the Velocity-Based Diffusion Architecture Works

Flow matching sampling in TRELLIS.2 trains a neural network to predict the velocity field of a forward diffusion process, then integrates that field backward via Euler steps to generate samples from pure noise.

TRELLIS.2 implements a velocity-based diffusion model that replaces traditional denoising objectives with flow matching. According to the microsoft/TRELLIS.2 source code, this architecture learns to predict how data points move through time rather than how to remove noise, enabling more direct and efficient sampling.

What Is Flow Matching in TRELLIS.2?

Flow matching formulates generative modeling as learning a velocity field (v(x, t)) that describes the instantaneous direction of data transformation. In trellis2/trainers/flow_matching/flow_matching.py, the FlowMatchingTrainer class implements this by:

• Defining a forward diffusion process that linearly interpolates between clean data (x_0) and Gaussian noise (z)
• Computing the analytical velocity target (v = (1-\sigma_{\text{min}})z - x_0)
• Training the denoiser to predict this velocity through mean-squared error

This differs from score-based or DDPM approaches that predict noise or score functions. The velocity formulation provides a direct ODE description of the generative process.

Forward Diffusion: Creating Noisy Training Data

The forward diffusion in TRELLIS.2 follows a linear interpolation schedule controlled by sigma_min. In FlowMatchingTrainer.diffuse() (lines 69-88), the noisy sample (x_t) at time (t \in [0,1]) is computed as:

[ x_t = (1-t)x_0 + \bigl(\sigma_{\text{min}} + (1-\sigma_{\text{min}})t\bigr)z ]

Where:

• (x_0) is the clean latent tensor
• (z \sim \mathcal{N}(0, I)) is standard Gaussian noise
• (\sigma_{\text{min}}) (default 1e-5) prevents numerical issues at (t=0)

from trellis2.trainers.flow_matching.flow_matching import FlowMatchingTrainer

# Example: forward diffusion during training

trainer = FlowMatchingTrainer(
    models={"denoiser": model},
    dataset=dataset,
    sigma_min=1e-5,
    t_schedule={"name": "logitNormal", "args": {"mean": 0.0, "std": 1.0}}
)

# The trainer automatically:

# 1. Samples timestep t ~ logit-normal distribution

# 2. Creates x_t via diffuse()

# 3. Computes velocity target v via get_v()

The time schedule is configurable. FlowMatchingTrainer.sample_t() (lines 134-140) implements logit-normal sampling by default, which concentrates training effort near the middle of the trajectory where the velocity field changes most rapidly.

Velocity Target Computation

The velocity target is derived analytically from the forward diffusion equation. In FlowMatchingTrainer.get_v() (lines 100-105):

[ v = \frac{dx_t}{dt} = (1-\sigma_{\text{min}})z - x_0 ]

Notice this does not depend on (t)—the velocity is constant in time for each sample, making the ODE straight-line in (x)-space. This simplification is specific to the linear interpolation schedule and enables efficient training.

The loss function in FlowMatchingTrainer.training_losses() (lines 62-73) is straightforward MSE between predicted and target velocity:

def training_losses(self, x_0, cond=None):
    t = self.sample_t(x_0.shape[0])           # Sample timesteps

    x_t = self.diffuse(x_0, t)                 # Forward diffusion

    v_target = self.get_v(x_0, x_t, t)         # Analytical velocity

    
    v_pred = self.models["denoiser"](x_t, t, cond)  # Network prediction

    
    return {"loss": F.mse_loss(v_pred, v_target)}

Euler Integration for Sampling

Sampling reverses the flow by integrating (\frac{dx}{dt} = -v(x,t)) from (t=1) (pure noise) to (t=0) (clean data). The FlowEulerSampler in trellis2/pipelines/samplers/flow_euler.py implements first-order Euler integration:

Single Euler Step

FlowEulerSampler.sample_once() (lines 79-82) performs one update:

[ x_{t_{\text{prev}}} = x_t - (t - t_{\text{prev}}) \cdot \hat v(x_t, t)

```python
from trellis2.pipelines.samplers.flow_euler import FlowEulerSampler

# Initialize from Gaussian noise

noise = torch.randn(batch_size, channels, height, width)
sampler = FlowEulerSampler(sigma_min=1e-5)

# The sampler automatically:

# 1. Creates linear time grid: t = [1.0, 0.98, ..., 0.0]

# 2. At each step, calls model to predict velocity

# 3. Updates state via Euler rule

result = sampler.sample(model, noise, steps=50)
generated = result.samples  # Final x_0

Full Sampling Loop

FlowEulerSampler.sample() (lines 94-126) implements the complete reverse integration:

def sample(self, model, noise, steps=50, guidance_strength=0.0, **kwargs):
    sample = noise
    t_seq = np.linspace(1, 0, steps + 1)  # Time from noise to data

    
    for t, t_prev in zip(t_seq[:-1], t_seq[1:]):
        # Predict x_0, epsilon, and velocity from current state

        pred_x_0, pred_eps, pred_v = self._get_model_prediction(
            model, sample, t, cond
        )
        
        # Euler integration step

        sample = sample - (t - t_prev) * pred_v
    
    return SampleResult(samples=sample, pred_x_0=pred_x_0, pred_eps=pred_eps)

The method _get_model_prediction() recovers (x_0) and (\epsilon) from the predicted velocity using the algebraic relationships defined by the forward diffusion, which is useful for intermediate visualization and guidance.

Classifier-Free Guidance Extensions

TRELLIS.2 provides two enhanced samplers for conditional generation:

FlowEulerCfgSampler

Inherits from ClassifierFreeGuidanceSamplerMixin (trellis2/trainers/flow_matching/mixins/classifier_free_guidance.py). This sampler computes conditional and unconditional velocities, then combines them:

[ \hat v_{\text{guided}} = \hat v_{\text{uncond}} + g \cdot (\hat v_{\text{cond}} - \hat v_{\text{uncond}}) ]

Where (g) is the guidance_strength parameter.

from trellis2.trainers.flow_matching.flow_matching import FlowMatchingCFGTrainer

cfg_trainer = FlowMatchingCFGTrainer(
    models=trained_models,
    p_uncond=0.1  # 10% chance to drop conditioning during training

)

sampler = cfg_trainer.get_sampler()  # Returns FlowEulerCfgSampler

result = sampler.sample(
    cfg_trainer.models["denoiser"],
    noise,
    cond=positive_condition,
    neg_cond=negative_condition,  # Unconditional or negative prompt

    steps=50,
    guidance_strength=3.0
)

FlowEulerGuidanceIntervalSampler

Restricts classifier-free guidance to a specific timestep interval, applying standard sampling outside that range. This balances sample quality with computational efficiency.

Key Implementation Files

File	Purpose	Key Classes/Functions
trellis2/trainers/flow_matching/flow_matching.py	Core trainer with diffusion, loss, and sampler factory	FlowMatchingTrainer, diffuse(), get_v(), training_losses()
trellis2/pipelines/samplers/flow_euler.py	Euler ODE solver for sampling	FlowEulerSampler, sample(), sample_once()
trellis2/pipelines/samplers/base.py	Abstract sampler interface	Sampler base class
trellis2/trainers/flow_matching/mixins/classifier_free_guidance.py	CFG functionality	ClassifierFreeGuidanceSamplerMixin
trellis2/models/sparse_structure_flow.py	Example velocity-predicting architecture	SparseStructureFlow

Why Flow Matching for 3D Generation?

TRELLIS.2's choice of flow matching over traditional diffusion reflects several advantages for structured latent generation:

• Straight-line trajectories: The constant-velocity ODE enables fewer sampling steps (25-50) compared to curved diffusion paths
• Deterministic sampling: No stochasticity in the base sampler, enabling reproducible generation
• Direct velocity supervision: The network learns a geometrically interpretable quantity rather than abstract scores

The linear interpolation schedule specifically suits TRELLIS.2's sparse structured latents, where maintaining geometric consistency across the trajectory is critical for coherent 3D output.

Summary

• Flow matching in TRELLIS.2 trains a velocity-predicting network rather than a denoiser, with the objective defined in FlowMatchingTrainer.training_losses().
• Forward diffusion mixes clean data and noise via linear interpolation in diffuse(), with velocity targets computed analytically in get_v().
• Sampling integrates the learned velocity field backward using first-order Euler steps in FlowEulerSampler.sample().
• Classifier-free guidance is available through FlowEulerCfgSampler, mixing conditional and unconditional predictions.
• The complete pipeline from training to inference is contained in trellis2/trainers/flow_matching/ and trellis2/pipelines/samplers/.

Frequently Asked Questions

What is the difference between flow matching and standard diffusion in TRELLIS.2?

Standard diffusion predicts noise or score functions and requires stochastic sampling with many steps. TRELLIS.2's flow matching predicts velocity along straight-line trajectories, enabling deterministic sampling with 25-50 Euler steps. The training objective is MSE on velocity rather than noise, implemented in FlowMatchingTrainer.training_losses().

Why does the velocity target not depend on timestep t?

The velocity (v = (1-\sigma_{\text{min}})z - x_0) is constant because TRELLIS.2 uses a linear interpolation schedule for forward diffusion. This design choice simplifies both training and sampling—trajectories are straight lines in data space, making Euler integration exact for the ODE and reducing numerical error.

How many sampling steps does TRELLIS.2 require?

The default configuration uses 50 Euler steps, though the sampler accepts any steps parameter. Flow matching's straight-line trajectories allow quality generation with as few as 25 steps, compared to 1000+ steps often required by DDPM. The FlowEulerSampler linearly spaces timesteps from 1.0 to 0.0 regardless of step count.

Can I use other ODE solvers with TRELLIS.2's flow matching?

The base FlowEulerSampler implements first-order Euler integration. While the repository does not include higher-order solvers, the velocity field formulation is compatible with Runge-Kutta or Heun methods—you would subclass Sampler from trellis2/pipelines/samplers/base.py and implement the alternative integration scheme using the same denoiser interface.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}