DiffusionGemma-Lite

From-scratch PyTorch block-AR discrete diffusion LM at ~343.5M params — 256-token canvases denoised in parallel from uniform noise, block-causal attention (causal across canvases, bidirectional within), entropy-bounded adaptive stopping, and a chunked-CE loss that never materializes full-vocab logits. Pure PyTorch. Read it like a field notebook: a name, a wiring sketch, then the measurements.

PARALLEL DECODE
256 TOKENS / STEP
One forward pass refines a whole canvas; KV cache grows once per finalized canvas, not per token
ADAPTIVE STOPPING
~2× FLOP CUT
Entropy-bounded early stop · closed rule, no learned halting policy · fixed-schedule fallback is bit-identical
CE MEMORY
6.6→1.1 GB @ ms8
Chunked x0-CE over 8192-token vocab chunks · never materializes the (B,T,V) logits tensor
ZERO-INIT SELF-COND
Δ=0 AT INIT
f(xt, t, sc) == f(xt, t) bit-for-bit at init (fp64, atol=0) · conditioning is learned, not imposed
FIG. A0 BLOCK-AR DECODE PIPELINE STATIC · SEE DIFFUSION.MD FOR THE FULL DERIVATION
• PREFILL PROMPT → [ DENOISE CANVAS (T steps, 256 tokens in parallel) → COMMIT unchanged-or-stronger → RE-ENCODE into KV ] × n_canvases
§ core

Core Documents

5 entries
README

Project Overview

343.5M-param block-AR discrete-diffusion LM — uniform-state canvases of 256 tokens, decoded in parallel.

 DIFFUSION

Authoritative Technical Doc

Diffusion math, mask derivation, self-conditioning, sampler semantics, training recipe, recipe deltas vs upstream.

 AGENTS

Agent Contract

Hard rules (pure PyTorch, chunked-CE, commit monotonicity), numerical-stability rules, known caveats.

 SKILLS

Skills Map

Day-to-day workflows: test suite, chunked-CE probe, sampler smoke, eval CLI, docs gates.

 DOCS

Documentation Index

Map of the concepts, guides, and symbol-anchored references in this portal.
§ concepts

Architecture & Concepts

3 entries
C1

Diffusion Core

D3PM-style uniform corruption, the cosine schedule, x0-parameterization, and the chunked-CE memory story.

 C2

Block-Causal Attention

Causal across canvases, bidirectional within: the mask that makes one transformer a denoiser and an LM.

 C3

Self-Conditioning

Zero-init posterior feedback: bit-exact at init, detached pre-pass, cross-step conditioning at eval.
§ guides

Guides & Playbooks

2 entries
G0

Quickstart

Data → train → sample → evaluate, end to end, with the shard-path contract spelled out.

 G1

Debugging Playbook

NaN rollback, resume divergence, anchor failures, shape errors — symptom-first recipes.
§ refs

API References

2 entries
R1

Config Reference

Every key of configs/pretrain_a100_380m.yaml and its TrainingConfig default.

 R2

API Reference

The public surface of models/, training/, data/, inference/, utils/ — symbol-anchored.