PixelDiT — conditional velocity model for MNIST

Pixel-space conditional DiT that predicts the flow-matching velocity v(z_t, t | c) on MNIST. Trained as the amortized posterior estimator for CoDiff (Bayesian Experimental Design via Contrastive Diffusions, arXiv:2410.11826).

Conditioning c = [A_Ξ θ, mask_Ξ] (2 channels) is a masked image plus its mask, where Ξ is a random union of 1–6 soft 7×7 squares. Sampling with c bound draws p(θ | measurements) directly in a few steps — no guidance needed. Conditioning dropout (10%) keeps the unconditional model available via zero conditioning.

Files

• model.msgpack — flax serialized parameters (EMA, 6.49M)
• config.json — architecture and training metadata

Architecture

Shared trunk of 5 blocks + 3-block velocity head, hidden size 256, 4 heads, patch size 4. In-context conditioning (learnable time tokens, no AdaLN), 2D RoPE on patch tokens, QK RMSNorm, SwiGLU MLP, bias-free linears. The network predicts x and the velocity follows from v = (z - x_pred) / max(t, 0.05). Architecture after pMF, single-head.

Training

Rectified flow on [0, 1] with target v* = e - x on z_t = (1-t)x + t e. 28k steps, batch 256, AdamW 1e-4, EMA decay 0.999. Data scaled to [-1, 1].

Usage

from flax import nnx, serialization
from huggingface_hub import hf_hub_download
from models.pixel_dit import PixelDiT   # from the CoDiff repo

net = PixelDiT(rngs=nnx.Rngs(0))
graphdef, state = nnx.split(net)
with open(hf_hub_download("jcopo/mnist", "model.msgpack"), "rb") as f:
    state.replace_by_pure_dict(serialization.from_bytes(state.to_pure_dict(), f.read()))
nnx.update(net, state)

v = net(x, cond=c, t=t)      # x: (B,28,28,1) in [-1,1], c: (B,28,28,2)

Citation

@article{iollo2024bayesian,
  title={Bayesian Experimental Design via Contrastive Diffusions},
  author={Iollo, Jacopo and Heinkel{\'e}, Christophe and Alliez, Pierre and Forbes, Florence},
  journal={arXiv preprint arXiv:2410.11826},
  year={2024}
}