praful-goel/speculative_decoding_models

Model Architecture & Training

Both models are Decoder-only Transformers (GPT-style) implemented in pure PyTorch.

Key Components

1. 1.Token + Position Embeddings
2. 2.Learned token embedding table: (vocab_size, n_embd)
3. 3.Rotary Position Embedding: parameter-free sinusoidal rotations on attention queries and keys

1. 1.Decoder Blocks Each block consists of:
2. 2.Masked Self-Attention (MHA/MQA/GQA)
3. 3.Key-Value (KV) Caching for O(1) generation latency
4. 4.Rotary Positional Embedding to understand relative positions
5. 5.Feed-Forward Network (Linear → GELU → Linear → Dropout)
6. 6.Residual connection
7. 7.Pre-norm RMSNorm

1. 1.Final RMSNorm + Language Modeling Head
2. 2.Linear projection from n_embd → vocab_size to produce logits

Training Strategy

Main Model

The main model was trained in 4 distinct phases to hande hardware constraints and optimize convergence

Draft Model (Small)

The draft model (small) was trained only once

max_iters = 40_000
warmup_steps = 1_000
eval_iter = 20
eval_interval = 1_000
accumulation_steps = 16
base_lr = 3e-4
weight_decay=0.1
batch_size = 16

Draft Model (Medium)

The draft model (medium) was trained only once

max_iters = 40_000
warmup_steps = 2_000
eval_iters = 20
eval_interval = 2_000
accumulation_steps = 16
base_lr = 3e-4
weight_decay = 0.1