---
title: Score-Regularized Continuous-Time Consistency Models (rCM)
url: https://www.emergentmind.com/topics/score-regularized-continuous-time-consistency-models-rcm
type: topic
---

# Score-Regularized Continuous-Time Consistency Models (rCM)

Score-Regularized Continuous-Time Consistency Models (rCM) are a framework for scalable, high-quality distillation of large-scale diffusion models—especially in the application domains of text-to-image (T2I) and text-to-video (T2V) generation. rCM unifies forward-KL (mode-covering) and reverse-KL (mode-seeking) objectives into a practical distillation methodology that preserves sample diversity and high-fidelity detail, while enabling few-step sampling for models exceeding 10 billion parameters [2510.08431].

## 1. Mathematical Foundations

### Continuous-Time Consistency Models (sCM)

Given a pre-trained diffusion "teacher" model defined by the probability-flow ODE
\[
\frac{dx_t}{dt} = f_{\rm teacher}(x_t, t)
\]
which maps data $x_{t=0} \sim p_{\rm data}$ to terminal noise at $t=T$, a student consistency model $f_\theta(x_t, t)$ is trained to predict the original $x_0$ directly from $x_t$ at arbitrary time $t$. Discrete-time consistency minimizes prediction error over intermediate states, while the continuous-time limit yields the sCM objective:
\[
L_{\rm cons}(\theta) = \mathbb{E}_{x_0, t} \big\| f_\theta(x_t, t) - f_{\theta^-}(x_t, t) - w(t) \frac{d\,f_{\theta^-}(x_t, t)}{dt} \big\|_2^2
\]
with $w(t) = \cos t$ and $x_t = \cos t\,x_0 + \sin t\,\varepsilon$. The differential
\[
\frac{d\,f_{\theta^-}(x_t, t)}{dt} = \nabla_{x_t} f_{\theta^-}(x_t, t) \cdot f_{\rm teacher}(x_t, t) + \partial_t f_{\theta^-}(x_t, t)
\]
necessitates Jacobian-vector products (JVPs) for efficient computation.

### Limitations of sCM

sCM objectives are fundamentally forward-KL-divergence minimizing (teacher $\rightarrow$ student), which enforces mode coverage but can accumulate errors and cause smoothing in reconstructed details. Specifically:
- **Error Accumulation:** Small inaccuracies at earlier diffusion times are amplified by self-feedback, resulting in blurred fine structures and reduced temporal coherence in video synthesis.
- **Mode-Covering:** The forward-KL focus penalizes missed data modes, broadening density at the cost of sharpness and high-frequency features.

## 2. Score Regularization and The rCM Loss

rCM augments sCM with a mode-seeking reverse-KL regularization, inspired by Distribution-Matching Distillation (DMD). Defining $p_\theta^t$ and $p_{\rm teacher}^t$ as the time-$t$ marginals for student and teacher,
\[
x_0^\theta \sim p_\theta,\quad x_t^\theta = \cos t\,x_0^\theta + \sin t\,\varepsilon
\]
the reverse-KL component is approximated via a “fake-score” network $\psi$ trained using flow matching, leading to the DMD loss:
\[
L_{\rm DMD}(\theta) = \mathbb{E}_{x_0^\theta, t}\left\| x_0^\theta - \operatorname{sg}\left[x_0^\theta - \frac{\psi(x_t^\theta, t) - f_{\rm teacher}(x_t^\theta, t)}{\mathbb{E}\|x_0^\theta - f_{\rm teacher}(x_t^\theta, t)\|} \right] \right\|_2^2
\]
where “sg” denotes stop-gradient.

The combined rCM objective is:
\[
L_{rCM}(\theta) = L_{\rm cons}(\theta) + \lambda L_{\rm DMD}(\theta), \qquad \lambda \approx 0.01
\]
$L_{\rm cons}$ preserves diversity, while $L_{\rm DMD}$ restores sharp detail via mode-seeking correction.

## 3. Implementation Techniques

### FlashAttention-2 JVP Kernel

A custom Triton-based kernel fuses the JVP operation into the block-wise, tiled FlashAttention-2 forward pass. The implementation:
- Accumulates both standard attention $O = P V$ and JVP outputs $tO = P\,tV + (P \odot (P \odot tS))\,V - \mathrm{diag}(\mathrm{rowsum}(P \odot tS))\,O$, where $tS = tQ\,K^T + Q\,tK^T$.
- Extends to both self- and cross-attention, preserving parallelism efficiency.

### Parallelism Strategies

- **FSDP:** Each layer exposes a “JVP-aware” interface, enabling tangent and forward computations without recomputing gradients.
- **Ulysses Context-Parallelism:** Query, key, and value (QKV) tensors are sharded and communicated by all-to-all protocols, with tangent vectors following the same pattern.

### Training and Inference Protocol

A generator/critic step alternates:
- **Generator:** Samples $(x_0, t)$, computes $x_t$ and JVP, forms $L_{\rm cons}$, and—after bootstrap—forms $L_{\rm DMD}$ using backward simulation of the student.
- **Critic:** Updates $\psi$ by matching student rollouts via $\mathbb{E}\|\psi(x_t, t) - x_0^\theta\|^2$.
- **Inference:** Few-step sampling ($1 \sim 4$ steps), iterating $x_0 \leftarrow f_\theta(x_{t_k}, t_k)$.

The stack employs PyTorch 2.0 (torch.func.jvp), custom Triton kernels, FSDP, Ulysses context-parallelism, BF16 mixed precision, and A100/H100 GPUs.

## 4. Experimental Results

Experiments span the Cosmos-Predict2 and Wan2.1 model families in T2I and T2V, with up to 14B parameters and 5-second video outputs.

| Model                    | Params | NFE | GenEval (T2I) | VBench (T2V) |
|--------------------------|-------:|----:|:-------------:|:------------:|
| Cosmos-Predict2 (teacher)| 14 B   | 70  | 0.84          | —            |
| + DMD2                   | 14 B   | 4   | 0.80          | 84.6         |
| + rCM                    | 14 B   | 4   | **0.83**      | **84.9**     |

- **Speedup:** rCM enables $15\times\sim50\times$ acceleration (number of function evaluations, NFE) relative to teacher models.
- **Diversity:** Unlike DMD2, rCM maintains mode coverage and avoids mode collapse in T2V (e.g., varied object poses).
- **Qualitative Quality:** rCM sharply renders fine details (e.g., text in T2I, frame-to-frame coherence in T2V).

## 5. Theoretical and Practical Implications

rCM leverages the complementary strengths of forward-KL (mode-covering) and reverse-KL (mode-seeking) training signals. This dual-objective strategy counters errors and smoothing artifacts observed with sCM alone, while preserving the ability to generate high-diversity samples in few steps. JVP-based self-feedback remains a computational bottleneck and source of numerical instability, especially in low precision; proposed remedies include higher-order flow formulations and Hutchinson-trace approximations.

Implementation complexity increases due to custom kernel requirements and parallelism adaptation, yet rCM achieves state-of-the-art quality without requiring adversarial finetuning or extensive hyperparameter search. A residual performance gap is noted for single-step T2V sampling.

## 6. Extensions and Future Work

Potential avenues for improving rCM include:
- Adaptive regularization scheduling (time-dependent $ \lambda(t) $).
- Integration with multi-step consistency trajectory approaches (e.g., sCTM).
- Incorporating adversarial objectives to further improve sample realism.
- Addressing the efficiency and error limitations intrinsic to JVP computation.

A plausible implication is that rCM’s unification of consistency and score distillation principles may generalize to other generative model distillation regimes, especially where balancing sample sharpness with output diversity remains critical [2510.08431].

Source: https://www.emergentmind.com/topics/score-regularized-continuous-time-consistency-models-rcm