# SAD: Soft Ancestor Diffusion: Optional Intermediate States for Masked Discrete Diffusion Language Models

SAD is a block-wise hierarchical diffusion model for unconditional text generation. It supports three variants:

- **SAD** (`sad`) — block-wise diffusion with learnable ancestor hierarchy
- **Block Diffusion** (`block_diffusion`) — block-wise binary mask diffusion (no ancestors)
- **AR** (`ar`) — standard autoregressive GPT-2 baseline

---

## Directory Structure

```
sad/
├── configs/              # YAML config files
├── scripts/              # Entry-point scripts (training / inference / eval)
├── src/                  # Core library
│   ├── models/           # Model definitions
│   ├── diffusion/        # Ancestor table & noisy state builder
│   ├── losses/           # Loss functions
│   ├── data/             # Data loaders
│   ├── eval/             # Sampling & metric utilities
│   └── utils/            # Helpers
├── data/                 # Preprocessed data & caches (see below)
├── models/               # Downloaded pretrained models (see below)
├── tokenizers/           # Local GPT-2 tokenizer
├── outputs/              # Training checkpoints & logs
└── eval/                 # Evaluation results
```

---

## What to Put in `data/` and `models/`

### `data/`

| File / Subdir | What it is | How to get it |
|---|---|---|
| `data/owt_cache/` | HuggingFace `openwebtext` cache | Auto-created on first training run |
| `data/hierarchy_prototypes_*.pt` | Hierarchy prototype tensors (k-means centroids) | Run `scripts/build_hierarchy.py learn` |
| `data/ancestor_lut_*.pt` | Sparse soft top-k ancestor lookup tables | Run `scripts/build_hierarchy.py build_lut` |

**You do NOT need to commit `data/` to git.** Add it to `.gitignore`.

### `models/`

| File / Subdir | What it is | How to get it |
|---|---|---|
| `models/gpt2/` | GPT-2 base weights (for leaf embeddings & eval LM) | `huggingface-cli download gpt2 --local-dir models/gpt2` |
| `models/gpt2-large/` | GPT-2 Large (for MAUVE featurization & gen-PPL eval) | `huggingface-cli download gpt2-large --local-dir models/gpt2-large` |

**You do NOT need to commit `models/` to git.** Add it to `.gitignore`.

### `tokenizers/`

| File / Subdir | What it is | How to get it |
|---|---|---|
| `tokenizers/gpt2/` | GPT-2 tokenizer files | `huggingface-cli download gpt2 --local-dir tokenizers/gpt2` |

---

## Scripts

### Training

| Script | Purpose | Quick Start |
|---|---|---|
| `scripts/train_sad.py` | Train SAD (with ancestors) | `torchrun --nproc_per_node=8 scripts/train_sad.py --config configs/sad_owt_b32_top3.yaml` |
| `scripts/train_block_diffusion.py` | Train block-mask diffusion (no ancestors) | `torchrun --nproc_per_node=8 scripts/train_block_diffusion.py --config configs/block_diffusion_owt_b32.yaml` |
| `scripts/train_ar.py` | Train AR baseline | `torchrun --nproc_per_node=8 scripts/train_ar.py --config configs/ar_owt.yaml` |

All training scripts support:
- `--config <path>` — YAML config (required)
- `--resume <path>` — resume from checkpoint
- `--num_steps <N>` — override training steps
- `--batch_size <N>` — override per-GPU batch size

### Inference / Sampling

| Script | Purpose | Quick Start |
|---|---|---|
| `scripts/inference_sad.py` | Sample from SAD checkpoint | `python scripts/inference_sad.py --config configs/sad_owt_b32_top3.yaml --checkpoint outputs/sad/latest.pt --num_samples 4` |
| `scripts/inference_block_diffusion.py` | Sample from block-diffusion checkpoint | `python scripts/inference_block_diffusion.py --config configs/block_diffusion_owt_b32.yaml --checkpoint outputs/block_diffusion/latest.pt --num_samples 4` |
| `scripts/inference_ar.py` | Sample from AR checkpoint | `python scripts/inference_ar.py --config configs/ar_owt.yaml --checkpoint outputs/ar_baseline/latest.pt --num_samples 4` |

### Evaluation

| Script | Purpose | Quick Start |
|---|---|---|
| `scripts/eval_gen_ppl.py` | Generative perplexity of SAD / block-diffusion samples under GPT-2 Large | `python scripts/eval_gen_ppl.py --checkpoint outputs/sad/latest.pt --model_type sad --num_samples 256` |
| `scripts/eval_ar_gen_ppl.py` | Generative perplexity of AR samples under GPT-2 Large | `python scripts/eval_ar_gen_ppl.py --checkpoint outputs/ar_baseline/latest.pt --num_samples 256` |
| `scripts/compute_mauve.py` | Compute MAUVE score against OpenWebText | `python scripts/compute_mauve.py --checkpoint outputs/sad/latest.pt --model_type sad --num_samples 5000` |
| `scripts/compute_diversity.py` | Compute distinct-n / repetition-n diversity metrics | `python scripts/compute_diversity.py --input samples.json --output diversity.json` |

### Hierarchy Construction (SAD only)

| Script | Purpose | Quick Start |
|---|---|---|
| `scripts/build_hierarchy.py learn` | Learn hierarchy prototypes via cosine k-means++ | `python scripts/build_hierarchy.py learn --config configs/sad_owt_b32_top3.yaml` |
| `scripts/build_hierarchy.py build_lut` | Build sparse soft top-k ancestor LUT | `python scripts/build_hierarchy.py build_lut --config configs/sad_owt_b32_top3.yaml` |
| `scripts/build_hierarchy.py extend` | Extend existing prototypes with deeper levels | `python scripts/build_hierarchy.py extend --input data/hierarchy_50257-128.pt --levels 128,32,8 --output data/hierarchy_50257-128-32-8.pt` |
| `scripts/build_hierarchy.py merge` | Merge per-level LUT files into one | `python scripts/build_hierarchy.py merge --source 1:data/lut_top3.pt --source 2:data/lut_top2.pt --output data/lut_mixed.pt` |

---

## Configs

| Config | Model | Description |
|---|---|---|
| `configs/sad_owt_b32_top3.yaml` | SAD | Base SAD on OpenWebText (512 tokens, 2 levels) |
| `configs/sad_owt_b32_h2_mixed.yaml` | SAD | batch_size=32, 2 ancestor levels, mixed top-k |
| `configs/sad_owt_b32_h3_mixed.yaml` | SAD | batch_size=32, 3 ancestor levels, mixed top-k |
| `configs/sad_owt_b32_top1.yaml` | SAD | top-1 ancestors |
| `configs/sad_owt_b32_top2.yaml` | SAD | top-2 ancestors |
| `configs/sad_owt_b32_top3.yaml` | SAD | top-3 ancestors |
| `configs/block_diffusion_owt_b32.yaml` | Block Diffusion | Mask-only block diffusion (no ancestors) |
| `configs/ar_owt.yaml` | AR | Standard autoregressive GPT-2 baseline |

---

## Core Source Modules

| File | Role |
|---|---|
| `src/models/sad_model.py` | `SADModel` — backbone with block-diff attention mask & FlexAttention |
| `src/models/ar_model.py` | `ARModel` — standard causal decoder baseline |
| `src/models/dit_components.py` | DiT blocks, Rotary embeddings, EmbeddingLayer |
| `src/diffusion/ancestor_table.py` | `AncestorTable` — fixed LUT + learnable ancestor embeddings |
| `src/diffusion/noisy_state.py` | `NoisyStateBuilder` — sample level per position & build noisy embeddings |
| `src/losses/sad_loss.py` | Leaf CE (ancestor CE loss was not experimented) |
| `src/data/__init__.py` | `build_owt_dataloader`, `build_debug_dataloader` |
| `src/eval/metrics.py` | Evaluation metric helpers |

---

# 1. SAD (`scripts/train_sad.py`)

## Quick Start

```bash
# Single GPU
python scripts/train_sad.py --config configs/sad_owt_b32_top3.yaml

# Multi-GPU DDP (e.g. 8 GPUs)
torchrun --nproc_per_node=8 scripts/train_sad.py \
    --config configs/sad_owt_b32_top3.yaml

# Resume
torchrun --nproc_per_node=8 scripts/train_sad.py \
    --config configs/sad_owt_b32_top3.yaml \
    --resume outputs/sad/latest.pt

# Override steps / batch size
torchrun --nproc_per_node=8 scripts/train_sad.py \
    --config configs/sad_owt_b32_top3.yaml \
    --num_steps 500000 --batch_size 64
```

## Training Paradigm

- Sequence is divided into `block_size` token blocks; **intra-block bidirectional**, **inter-block causal**.
- Vectorized training: `x_full = [noisy | clean]`, single forward with block-diff mask:

| q→k           | Allowed Condition |
|---------------|-------------------|
| noisy → noisy | Same block only |
| noisy → clean | Strictly earlier clean block |
| clean → clean | Same or earlier clean block |
| clean → noisy | Never |

- No timestep conditioning; only a single learnable `cond_bias`.
- Input is continuous embeddings (leaf / ancestor / mask), not discrete token ids.
- Loss computed only on the noisy half: leaf cross-entropy. (Note: `lambda_ancestor > 0` was not experimented.)

## Positional Encoding

Three-way additive + RoPE:

- `block_idx_embed` — cross-block position (AR)
- `intra_pos_embed` — intra-block position
- `segment_embed` — noisy(0) / clean(1)
- RoPE applied to both halves with `0..L-1`; noisy[i] and clean[i] have relative offset 0

## Config Skeleton (`configs/sad_owt_b32_top3.yaml`)

```yaml
model:
  vocab_size: 50257
  hidden_size: 768
  n_blocks: 12
  n_heads: 12
  max_seq_len: 512
  block_size: 8
  num_levels: 2
  level_sizes: [50257, 128]

ancestor:
  lut_path: data/ancestor_lut_50257-128_top3_t1.0.pt    # fallback: random debug LUT
  proto_path: data/hierarchy_prototypes_50257-128.pt

loss:
  lambda_ancestor: 0.0      # ancestor CE loss was not experimented
  mask_only: false

training:
  seed: 0
  batch_size: 64          # per-GPU
  num_steps: 1_000_000
  lr: 3.0e-4
  weight_decay: 0.01
  grad_clip: 1.0
  dtype: bf16
  log_interval: 100
  eval_interval: 5000
  save_interval: 10000
  compile: default        # "off" to disable torch.compile

data:
  dataset: openwebtext    # "debug" uses TinyDebugDataset
  seq_len: 512
  cache_dir: data/owt_cache
  mode: subsample         # subsample | pack

logging:
  use_wandb: true
  project: sad
  save_dir: outputs/sad
```

## DDP Invariants (do not change without care)

1. **AncestorTable is outside DDP**, but its parameters are in the optimizer. After `loss.backward()` and before clip, manually `dist.all_reduce(grad, op=AVG)` to sync its gradients.
2. **AncestorTable initial value broadcast**: after `build_ancestor_table`, broadcast parameters and buffers from rank 0 via `dist.broadcast` to all ranks, to counteract the drift from `set_seed(seed + local_rank)`.
3. **Streaming data sharding**: `build_owt_dataloader(shard_across_ranks=True)` for train, `False` for val (val only runs on rank 0; sharding would bias metrics to a single shard).
4. **FlexAttention ignores padding mask**, protected by an assert that `pad_token_id == eos_token_id`. Future multi-doc packing would need doc ids in `mask_mod`.
5. **Debug LUT cross-rank consistency**: `build_ancestor_table` uses an independent `torch.Generator` for random LUT, not the global RNG.

## Output

- `outputs/sad/ckpt_{step}.pt` — periodic checkpoint
- `outputs/sad/latest.pt` — latest checkpoint (overwritten)
- `outputs/sad/config.yaml` — runtime config snapshot
- wandb: `project=sad` (when `logging.use_wandb: true`)

Checkpoint contains `step` / `model` / `ancestor_table` / `optimizer` / `config` / `metrics`; resume restores all of them (data iterator position is not restored; streaming data restarts from the beginning).

## Evaluate

Runs on rank 0 only, every `eval_interval` steps, over `num_batches=50` batches (see `evaluate()`). Non-main ranks wait for rank 0 at the next backward all-reduce, so long eval will drop throughput but will not deadlock.