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

> Discover how flow matching sampling in TRELLIS.2 generates samples from noise by training a neural network to predict velocity fields and integrating them backward using Euler steps.

- Repository: [Microsoft/TRELLIS.2](https://github.com/microsoft/TRELLIS.2)
- Tags: deep-dive
- Published: 2026-08-04

---

**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`](https://github.com/microsoft/TRELLIS.2/blob/main/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\)

```python
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:

```python
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`](https://github.com/microsoft/TRELLIS.2/blob/main/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:

```python
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.

```python
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`](https://github.com/microsoft/TRELLIS.2/blob/main/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`](https://github.com/microsoft/TRELLIS.2/blob/main/trellis2/pipelines/samplers/flow_euler.py) | Euler ODE solver for sampling | `FlowEulerSampler`, `sample()`, `sample_once()` |
| [`trellis2/pipelines/samplers/base.py`](https://github.com/microsoft/TRELLIS.2/blob/main/trellis2/pipelines/samplers/base.py) | Abstract sampler interface | `Sampler` base class |
| [`trellis2/trainers/flow_matching/mixins/classifier_free_guidance.py`](https://github.com/microsoft/TRELLIS.2/blob/main/trellis2/trainers/flow_matching/mixins/classifier_free_guidance.py) | CFG functionality | `ClassifierFreeGuidanceSamplerMixin` |
| [`trellis2/models/sparse_structure_flow.py`](https://github.com/microsoft/TRELLIS.2/blob/main/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`](https://github.com/microsoft/TRELLIS.2/blob/main/trellis2/pipelines/samplers/base.py) and implement the alternative integration scheme using the same `denoiser` interface.