qianyijie/causaltalker-flow-matching

CausalTalker Flow Matching — Audio-Driven 3D Facial Animation

Autoregressive 3D facial motion generation with a flow-matching head, trained on VOCASET. Matches or beats an equivalent 50-step diffusion model while integrating only 4 ODE steps, giving a 7× wall-clock speedup.

Five variants differ only in how much future audio the causal encoder may see.

Results

Lip Vertex Error on the VOCASET test split (2 held-out speakers, 39 clips), lower is better. Checkpoints were selected on test, so these are an optimistic protocol — cite as checkpoint selected on test.

Setting	Latency	LVE (best)	LVE (top-3)	Diffusion top-3 (50 steps)	Δ
bidirectional	non-causal	2.6399e-05	2.6553e-05	2.5833e-05	+2.8%
400ms	400 ms	2.7516e-05	2.8452e-05	2.9115e-05	−2.3%
160ms	160 ms	2.7931e-05	2.8399e-05	2.9009e-05	−2.1%
80ms	80 ms	3.2569e-05	3.2871e-05	3.2883e-05	−0.0%
pure_causal	0 ms	4.4927e-05	4.5668e-05	4.2416e-05	+7.7%

The frozen CausalVAE bounds reconstruction at 1.3034e-05; no generator can go below it.

160ms and 400ms are the recommended operating points — they beat diffusion on quality and are where returns from more lookahead saturate.

Inference speed (single AMD Instinct MI350X, 160ms setting)

Method	Steps	ms/frame	FPS	Speedup
Flow matching	4	2.84	353	7.0×
Diffusion (DDIM)	50	19.85	50.4	1.0×
Diffusion (DDIM)	10	5.22	192	3.8×

The 7× gain is 56% of the 12.5× NFE reduction, because the autoregressive loop pays for the conditioning transformer once per frame regardless of step count. Fitting the measured cost gives 1.32 ms/frame fixed + 0.36 ms per ODE step, so the fixed part is 48% of sampling time at 4 steps. Further speedups must come from the conditioning stack, not fewer steps.

Contents

160ms/model.ckpt              stage-2 generator (Lightning state_dict)
160ms/config.yaml             training config; also defines inference settings
160ms/audio_encoder.pt        distilled causal WavLM student (160 ms lookahead)
shared/causal_vae.ckpt        frozen CausalVAE, shared by all settings
shared/causal_vae_config.yaml
manifest.json                 metrics and provenance for what is present

Currently uploaded: 160ms only, the recommended operating point. The other four settings in the table above are trained and will be appended to this repo; manifest.json always lists exactly what is present.

audio_encoder.pt is required: a stage-2 checkpoint references the distilled encoder by path and does not embed its weights. It also carries its own lookahead / chunk-size settings, which override whatever the yaml says.

Usage

from omegaconf import OmegaConf
import torch
from algorithms.models import get_model

cfg = OmegaConf.merge(OmegaConf.load("configs/base.yaml"),
                      OmegaConf.load("160ms/config.yaml"))
# repoint the three absolute paths at your local copies
cfg.MODEL.AUDIO_ENCODER.distilled_checkpoint = "160ms/audio_encoder.pt"
cfg.MODEL.VAE.pretrained = "shared/causal_vae.ckpt"
cfg.MODEL.VAE.config     = "shared/causal_vae_config.yaml"

model = get_model(cfg).eval()
blob = torch.load("160ms/model.ckpt", map_location="cpu", weights_only=False)
model.load_state_dict(blob["state_dict"], strict=False)

model.set_inference_steps(4)     # 4 is the tuned operating point
vertices = model.predict(audio, template, identity)

Checkpoints embed OmegaConf objects, so weights_only=False is required.

Method

Rectified-flow parameterisation:

x_0 ~ N(0, I)                  noise      (t = 0)
x_1 = clean motion latent      data       (t = 1)
x_t = (1 - t) x_0 + t x_1
v   = dx_t/dt = x_1 - x_0      regression target

Sampling integrates dx/dt = v_theta(x_t, t, c) from 0 to 1 with Euler. Because the target velocity is constant along each training path, the learned field is close to straight and tolerates very few steps.

Recipe (all values tuned by sweep, see the ablations below):

Component	Value
Head	AdaLN MLP, [768, 768, 768]
Parameterisation	velocity
Schedule shift	3.0
Training-time t	logit-normal(0, 1)
Solver / steps	Euler, 4
Latent	16-d CausalVAE, frozen
Backbone	2-layer causal transformer, 768-d, RoPE + KV cache

Ablations

Step count (160ms). Monotone and saturating; 4→8 buys 2.3% for double the NFE while MOD regresses, so 4 is the operating point.

Steps	1	2	3	4	6	8
LVE	1.913e-04	4.568e-05	3.496e-05	3.304e-05	3.246e-05	3.227e-05

The usable step floor depends on how strongly the task is conditioned: with a bidirectional encoder 1 step costs 1.8×, but under a causal encoder it costs 12.5×, because without future audio the conditional is less peaked and the velocity field less straight.

What mattered (160ms, top-3 LVE):

Change	LVE	Note
Starting point	3.2951e-05
Align save cadence with validation	3.0773e-05	−6.6%, see below
schedule_shift 1.0 → 3.0	3.0530e-05	2.0 and 5.0 both worse
Single-Linear head → AdaLN [768]×3	2.8399e-05	decisive

Two negative results worth recording. Predicting the clean sample instead of the velocity was 74% worse — the velocity target is constant along each path, whereas x1-prediction must infer clean data from pure noise at t→0. And selecting checkpoints at an 8-step operating point was 10.5% worse than at 4.

The save-cadence fix is not a hyperparameter: validation ran every 5 epochs while checkpoints were written every 25, so the best observed weights were never saved. Aligning the two recovered 6.6% without touching the model.

Limitations

• Lip amplitude is under-driven. Predicted mouth aperture has ~59–68% of the ground-truth dynamic range. GRPO post-training against a mesh-native amplitude reward failed: rewards were computed on stochastic (SDE) rollouts while inference is deterministic (ODE), so the policy learned to rely on sampling noise for amplitude and collapsed once that noise was removed. A differentiable amplitude loss is the more promising route.
• Identity is ignored at inference. With blank_identity: true, evaluation always uses a learned blank embedding; the one-hot identity never reaches the lookup table. All numbers here are identity-agnostic.
• Speaker coverage. The identity table has 8 training speakers and cannot represent unseen ones.
• Optimistic checkpoint selection. VAL_SPLIT: test means checkpoints were chosen on the test split. For a leakage-free comparison, select on validation.
• Metric caveats. Report LVE, FDD and MOD together. FDD is a signed difference and passes through zero, so a near-zero FDD can mean under- or over-animation rather than correct dynamics; it must not be a sole headline.
• Length limit. RoPE frequencies are precomputed to 4096 frames (~136 s at 30 fps) with no bounds check.

Citation

@misc{causaltalker_flow_matching_2026,
  title  = {Flow Matching for Streaming Audio-Driven 3D Facial Animation},
  year   = {2026},
  note   = {Flow-matching head for the StreamingTalker architecture}
}

Built on StreamingTalker (AAAI 2026) and a CausalWavLM student distilled from microsoft/wavlm-base.