Diffusion and transformer audit¶

Measured on AMD Ryzen 9 5900X 12-Core Processor (CPU) with JAX 0.10.0. Benchmarks synchronize device work, separate export/first-call cost from five warm calls, and use a fixed seed. Runtime numbers are indicative of this machine; these are not GPU measurements. Gaussian and mixture samplers draw 4,096 scalar events. The transformer uses batch 8, width 64, two layers and four heads.

Confirmed bugs fixed¶

• Diffusion sampling now initializes at the distribution view’s overridden t_max.
• Explicit-state sampling uses that state for the base distribution as well as the solver; Python scalar state leaves can be snapshotted.
• DDIM/v solver coefficients use the physical SDE instead of differentiating clipped endpoints.
• Exponential AB2 uses phi2-weighted extrapolation, the current/previous step-size ratio, and midpoint scalar linear coefficients. Analytic forward/reverse, uniform/nonuniform grids verify second-order convergence.
• TimeMLP and DiffusionTransformer accept broadcast batch time vectors; bfloat16 positional encoding casts explicitly.

Gaussian accuracy and runtime (64 time points)¶

Exact zero-network Gaussian denoiser. Percent errors below include Monte Carlo error: the shared ODE initial sample itself is 0.282% above the true standard deviation. VE ODE previously had +11.703% error at 64 points; after the AB2 correction it has +0.282%. The coefficient correction alone did not remove that numerical error.

Nonlinear mixture (64 time points)¶

Exact posterior denoiser for an equal mixture of N(-3, 0.5²) and N(3, 0.5²). The terminal Gaussian is a moment-matched approximation; its error is not isolated from solver error. Target mean absolute value is approximately 3, and target probability of |x|<1 is approximately 0.000032. This is a nonlinear score test, not a learned-model or high-dimensional quality benchmark.

Transformer scaling¶

Remaining weaknesses¶

• Coarse SDE grids can be unstable: VE at 16 points produces Gaussian standard deviation about 5.97 instead of 1. More points are essential; finite outputs alone are not a quality check.
• Oracle SDE runs are much slower than ODE because of per-step random draws; expensive learned denoisers may change that ratio.
• Full self-attention grows quadratically in token count. bfloat16 does not guarantee a CPU speedup. These checks establish finite forward/gradient outputs, not low-precision training equivalence.
• The categorical linear/sigmoid defaults retain about 95% signal at the endpoint; the previous audit documented this without changing defaults. Categorical mixing quality remains unaudited here.
• The exponential AB3 method was regression-tested for shared-state compatibility, but its convergence/order on nonuniform grids was not established by this audit.
• Precompile a distribution with dist.compile() before placing its calls inside an outer jax.jit; lazy export during tracing currently fails.
• No GPU benchmark or end-to-end trained diffusion comparison was run in this audit.

Reproduce¶

JAX_PLATFORMS=cpu python -m docs.scripts.benchmark_diffusion_audit --output /tmp/gaussian.json
JAX_PLATFORMS=cpu python -m docs.scripts.benchmark_diffusion_audit --profile mixture --output /tmp/mixture.json
JAX_PLATFORMS=cpu python -m docs.scripts.benchmark_transformer_audit --output /tmp/transformer.json

Raw measurements are stored alongside this report as JSON, including the Gaussian baseline taken before this audit’s fixes.

Validation: 97 focused audit/diffusion/transformer/sampling tests and 152 selected split-drift ODE/SDE regression tests passed on CPU. The positional-encoding benchmark was rerun after the bfloat16 cast fix.