VFM (Variational Flow Map): few-step diffusion via learned noise adapters on bidirectional LLMs

[johndpope]

What is VFM?

Variational Flow Map (VFM) is a method to enable 1–4 step diffusion generation from masked diffusion models, replacing the standard 16–512 step iterative unmasking loop. It works by training a learned noise adapter q_φ(z|prompt) that maps prompt embeddings → an initial noise distribution in embedding space, then a single forward pass through a frozen bidirectional LLM (the "flow map") produces token predictions.

Architecture

Prompt tokens → embed_tokens → VFMNoiseAdapter (2-layer Transformer) → (μ, log σ)
                                                                      ↓
                                                              z = μ + σ·ε (reparameterize)
                                                                      ↓
                                                        z_full = [prompt_embeds; z_gen]
                                                                      ↓
                                                    VFMFlowMapWrapper (frozen LLM, is_causal=False)
                                                                      ↓
                                                              logits → token IDs

Training objective (Eq. 19)

L = (1/2τ²) · L_data + (1/2σ²) · L_obs + L_KL

• L_data: MSE reconstruction loss in embedding space at generation positions
• L_obs: Cross-entropy observation loss at prompt positions (prompt tokens preserved)
• L_KL: KL(q_φ(z|prompt) || N(0,I)) — keeps the adapter's noise distribution close to standard normal

Why bidirectional?

Standard autoregressive LLMs use causal attention, which cannot condition on future tokens. Masked diffusion models require bidirectional attention — each masked token should see all other tokens (both left and right). Using a causal base model for the flow map produces poor reconstruction. Using a pretrained bidirectional model (like fredzzp/open-dcoder-0.5B) gives dramatically better results.

---

Key finding: bidirectional base model is critical

We compared VFM training on two 0.5B base models:

Metric @ 500 steps	Qwen2.5-Coder-0.5B (causal)	open-dcoder-0.5B (bidirectional)
data_loss	8.5	2.26 (3.8× better)
kl_loss	10.6	9.98
obs_loss	7.6	7.09
total_loss	18.7	14.65 (22% better)
grad_norm	200–2000 (unstable)	38 (50× more stable)
NaN steps	5/500 (1%)	1/500

The bidirectional model reconstructs embeddings 3.8× better with dramatically more stable gradients. This validates the core VFM hypothesis: the flow map needs to attend to all positions, not just past ones.

---

Reproduce

Prerequisites

• 1× GPU with ≥16 GB VRAM (tested on RTX PRO 4000 Blackwell, 24 GB)
• Model weights downloaded: fredzzp/open-dcoder-0.5B (1.2 GB)
• Training data: /run/media/johndpope/12TB/open_dllm/ldlm_data/data.jsonl (FineWeb 100K sample)

Step 1: Download open-dcoder-0.5B

python -c "
from huggingface_hub import snapshot_download
snapshot_download('fredzzp/open-dcoder-0.5B', local_dir='/path/to/open-dcoder-0.5B')
"

Step 2: Run VFM smoke test (500 steps, ~2 min)

CUDA_VISIBLE_DEVICES=1 PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True \
torchrun --nproc_per_node=1 tasks/train_vfm.py \
    configs/pretrain/vfm_open_dcoder.yaml

Config: configs/pretrain/vfm_open_dcoder.yaml

Step 3: Compare with causal baseline

CUDA_VISIBLE_DEVICES=1 PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True \
torchrun --nproc_per_node=1 tasks/train_vfm.py \
    configs/pretrain/vfm_06b_smoke.yaml

Step 4: Verify open-dcoder's native MDM generation (16 steps, baseline)

import torch
from torch.nn import functional as F
from transformers import AutoTokenizer, AutoModelForCausalLM

tok = AutoTokenizer.from_pretrained("fredzzp/open-dcoder-0.5B", trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained("fredzzp/open-dcoder-0.5B", torch_dtype=torch.bfloat16, trust_remote_code=True).to("cuda")

prompt = "def fibonacci(n):"
pids = tok.encode(prompt, return_tensors="pt").to("cuda")
x = F.pad(pids, (0, 64), value=151643)
timesteps = torch.linspace(1, 1e-3, 17, device=x.device)
for i in range(16):
    mask_index = (x == 151643)
    if not mask_index.any(): break
    with torch.no_grad():
        outputs = model(input_ids=x, is_causal=False)
    logits = torch.cat([outputs.logits[:, :1], outputs.logits[:, :-1]], dim=1)
    probs = torch.softmax(logits[mask_index].float(), dim=-1)
    log_probs = torch.log(probs.clamp(min=1e-10))
    confidence = (probs * log_probs).sum(dim=-1)
    _, x0 = probs.max(dim=-1)
    t, s = timesteps[i], timesteps[i+1]
    num_masked = mask_index.sum(dim=-1, keepdim=True)
    num_to_unmask = (num_masked * (1 - s / t)).long()
    full_conf = torch.full_like(x, -float("inf"), dtype=confidence.dtype, device=x.device)
    full_conf[mask_index] = confidence
    _, unmask_indices = torch.topk(full_conf, k=num_to_unmask.max(), dim=-1)
    rows = torch.arange(x.size(0), device=x.device).unsqueeze(1)
    sel = torch.zeros_like(x, dtype=torch.bool)
    sel[rows, unmask_indices] = True
    sel = sel & (torch.cumsum(sel.long(), dim=-1) <= num_to_unmask)
    proposals = torch.full_like(x, fill_value=151643)
    proposals[mask_index] = x0
    x[sel] = proposals[sel]
print(tok.decode(x[0][pids.shape[1]:], skip_special_tokens=True))

---

Wandb runs

Run	Base model	Wandb
vfm-06b-smoke	Qwen2.5-Coder-0.5B (causal)	project
vfm-open-dcoder	open-dcoder-0.5B (bidirectional)	project

---

Bugs found and fixed

1. VFMNoiseAdapter position embedding OOB (veomni/models/vfm/model.py:72): When prompt_length + gen_length > max_seq_len, the position indices exceeded the embedding table size, causing a CUDA assert crash at ~step 25. Fixed by clamping positions to max_seq_len - 1.

2. VFMFlowMapWrapper missing is_causal=False (veomni/models/vfm/model.py:135): The flow map was calling the model with default causal attention, defeating the purpose of using a bidirectional base. Fixed by adding is_causal=False to the model call.

3. Embedding resize for larger model vocabs (tasks/train_vfm.py:60): resize_token_embeddings(len(tokenizer)) would shrink the open-dcoder's embedding table from 151936 → 151666. Fixed to only resize when tokenizer is larger.

---

Key files

File	Purpose
veomni/models/vfm/model.py	VFM implementation (VFMNoiseAdapter, VFMFlowMapWrapper, VariationalFlowMap)
veomni/models/vfm/__init__.py	Module exports
tasks/train_vfm.py	VFM training script (joint loss, NaN skip, wandb, eval, generation probe)
configs/pretrain/vfm_open_dcoder.yaml	Config for open-dcoder-0.5B VFM training
configs/pretrain/vfm_06b_smoke.yaml	Config for Qwen2.5-Coder-0.5B causal baseline
veomni/utils/arguments.py	Added vfm and eval_size fields

---

Next steps

• Train VFM with more steps (5K+) and full dataset to evaluate generation quality
• Implement multi-step VFM inference (K=2,4 steps) with iterative refinement
• Benchmark VFM 1-step vs MDM 16-step vs AR generation for throughput (tok/s)
• Scale to 27B: use Repr-Align to convert Qwen3.6-27B to bidirectional, then train VFM adapter
• Evaluate perplexity on standard benchmarks (HumanEval, MBPP) for open-dcoder + VFM

---

Open questions

1. VFM generation quality: 500 steps of training produces good loss but garbled text. Is this a training duration issue, or does the VFM architecture need modification (e.g., more adapter layers, different loss weighting)?
2. Optimal step count: What is the minimum number of VFM steps needed to match 16-step MDM quality? Is 4 steps sufficient?
3. Scaling: Does the 3.8× data_loss improvement from bidirectional attention hold at 27B scale?
4. Mask token: open-dcoder uses mask_token_id=151643 (reusing bos/eos) in generation_config but 151665 (<M>) in the tokenizer. Both work, but which was used during training?