LLM Pretraining: KV-Structural Reduction

Research Question

Design a more KV-efficient causal attention structure for GPT-style pretraining, with the primary focus on the tradeoff between KV head sharing and latent KV compression.

• how much quality can be preserved by reducing the realized KV state
• whether grouped/shared KV heads or latent KV bottlenecks give the better quality-memory tradeoff under a fixed small-scale pretraining budget

What You Can Modify

One editable region in custom_pretrain.py:

1. Attention-structure region (lines 35-154), including:
   • build_kv_heads(...): how many KV heads are materialized relative to query heads
   • cross_layer_share(...): optional structural sharing hook inside the attention stack
   • latent_kv_project(...): whether K/V are compressed into a lower-rank latent space
   • CausalSelfAttention: how the above choices are instantiated inside the attention block, including the internal query/KV projection and attention mixing path

Intended Task Boundary

• This task studies KV-state reduction inside the attention block.
• The main comparison axes are:
  • dense MHA vs grouped/shared KV heads
  • grouped/shared KV heads vs latent KV compression
• cross_layer_share(...) remains available as an auxiliary structural hook inside the same block.
• The evaluator enforces the top-level boundary of this region with an AST validator: only the allowed helper functions plus CausalSelfAttention may appear in the editable span. That keeps edits inside the attention block, even though the internal contents of CausalSelfAttention remain flexible.

Evaluation

• Primary metric: validation loss (cross-entropy, lower is better)
• Secondary metrics:
  • kv_bytes_per_token (lower is better; evaluator-derived analytic KV footprint from the realized attention structure)
  • head_sharing_ratio (higher means more aggressive KV head sharing; evaluator-derived from realized head topology)
  • latent_rank_ratio (lower is better only for methods that truly use latent KV compression; evaluator-derived from latent projection dimensions)
  • heldout_loss (lower is better; perplexity-style held-out evaluation on the packaged eval corpus, reported consistently for every visible benchmark)
  • generation_toks_per_s (higher is better; lightweight autoregressive generation throughput on held-out tokens, reported consistently for every visible benchmark)
• Visible benchmark regimes:
  • train-ctx1024: 124M pretraining at block_size=1024
  • train-ctx256: 124M short-context pretraining at block_size=256
  • eval-wikitext2: 124M pretraining followed by held-out eval and generation-side diagnostics on WikiText-2
  • eval-wikitext103: 124M pretraining followed by held-out eval and generation-side diagnostics on WikiText-103
  • eval-lambada: 124M pretraining followed by held-out eval and generation-side diagnostics on LAMBADA
• Training data: ClimbMix tokenized training split
• Held-out eval data: packaged eval dependency (WikiText-2, WikiText-103, LAMBADA)
• Training schedule: same budget family as the existing nanoGPT MLS-Bench pretraining tasks