llada-distilled-24L-v3

LLaDA Distilled 24L Draft Model (v3)

A 24-layer structurally pruned + distilled version of LLaDA-8B-Instruct (32 layers), designed as a draft model for speculative decoding in masked diffusion language models.

Last updated: 2026-04-12 01:59 UTC — Step 13000, Epoch 3 Best argmax agreement: 0.6502

Current Status

Architecture

Layers kept (original indices): [0–11, 20–31] Layers removed: [12, 13, 14, 15, 16, 17, 18, 19]

How It Was Made

Step 1: Structured Pruning (Data-Driven)

12 ablation experiments identified redundancies:

• Layer removal: Layers 12–19 (contiguous middle block) have ~89% leave-one-out agreement — most expendable
• FFN pruning: Bottom 25% of neurons (by gated activation magnitude) removed per layer — different neurons per layer
• Pre-distillation agreement: ~60% argmax agreement with full model

Step 2: Knowledge Distillation (Spec-Dec Optimized)

Loss function — designed to maximize speculative decoding acceptance rate:

• 0.5× Top-1 CE: CE(student_logits, teacher_argmax) — directly optimizes the CtV/VtC acceptance criterion
• 0.3× Top-K KL: KL(student_topK ∥ teacher_topK) over teacher's top-32 tokens only — sharpens decision boundary, ignores 126K irrelevant vocab tokens
• 0.2× Confidence MSE: MSE(max(softmax(student)), max(softmax(teacher))) — calibrates confidence for threshold-based drafting

Training config:

• Data: StarCoderData Python (40%) + OpenMathInstruct-2 (40%) + C4 (20%)
• Mask ratio: Beta(2,2) ∈ [0.15, 0.85]
• Optimizer: AdamW (lr=5e-05, cosine schedule, 5% warmup)
• Effective batch size: 16
• Sequence length: 512
• Precision: bf16

Training History

Intended Use

This is a draft model for speculative decoding with LLaDA-8B-Instruct as the target. NOT for standalone generation.

from transformers import AutoModel, AutoTokenizer
import torch

draft_model = AutoModel.from_pretrained(
    "jaygala223/llada-distilled-24L-v3",
    trust_remote_code=True,
    torch_dtype=torch.bfloat16
).to("cuda").eval()

target_model = AutoModel.from_pretrained(
    "GSAI-ML/LLaDA-8B-Instruct",
    trust_remote_code=True,
    torch_dtype=torch.bfloat16
).to("cuda").eval()

Limitations

• Not standalone — designed for speculative decoding only
• Same tokenizer as LLaDA-8B-Instruct (mask_id=126336)
• Training ongoing — later checkpoints may improve