Research projects

Current work on making large language models cheaper to train and run: two projects on transformer training and two on long-context inference. All are in progress; write-ups are in preparation.

BLT: Bilinear Attention with a Shared Interaction Matrix

A drop-in replacement for multi-head attention that swaps every layer's query and key projections for a single D×D matrix shared across all layers and heads. That cuts attention parameters by up to 50%, removes a matrix multiply per layer, and for smaller models lets the shared matrix stay resident in GPU cache instead of streaming per-layer weights from HBM.

Trained GPT-2 Small with BLT, standard MHA and grouped-query attention: BLT and GQA pay a similar overall quality cost for different reasons.
A hybrid of half MHA and half BLT layers recovers most of the lost quality while keeping half the memory benefit.
A low-rank variant (M ≈ UV^T) beats full-rank BLT across three seeds and nearly matches MHA on zero-shot benchmarks, with about half the attention parameters and 32% faster training.

EMA Per-Token Loss Weighting

Standard pretraining weights every token equally, so easy, frequent tokens dominate the gradient. This method keeps a running loss estimate for each vocabulary token and reweights the training loss by it. Unlike prior hard-token methods it needs no reference model and no extra forward pass: it uses only what the model already computes.

Improves LAMBADA accuracy on GPT-2 Small from 0.225 to 0.268 (cumulative running mean) at a bounded cost to general next-token perplexity.
The trade-off holds across architectures (standard attention and BLT) and model scale (GPT-2 Medium, 355M parameters).
Blending the weighted objective with standard cross-entropy from the start of training beats fine-tuning afterward, which hits a “forgetting cliff.”

Faithfulness over Accuracy: Rethinking KV Cache Compression

KV cache compression methods are usually judged by whether the compressed model still gets the right answer, and by that measure the leading methods look nearly free. But the real goal is to produce the same output the full-context model would have, right or wrong. This project introduces two faithfulness metrics: Weighted TV Faithfulness (F_WTV), an entropy-weighted total-variation distance between the compressed and full model's token distributions at every generation step, and Output Faithfulness (F_out), word-level overlap between the two models' generated text. Methods are evaluated on Llama 3.1 8B and Mistral 7B across 16 LongBench tasks.

F_WTV is bounded between 0 and 1, discounts near-deterministic steps that any method matches, and isn't distorted by rare catastrophic divergences the way KL divergence is.
SnapKV and PyramidKV, the best methods by ground-truth accuracy, score far worse on distributional faithfulness than naive truncation.
The cause is structural: tokens kept at their original RoPE positions after eviction corrupt every decode step. Re-rotating the retained keys to compact positions, a single tensor operation, recovers most of the loss.
KL divergence, the usual distributional metric, gives misleading method rankings; F_WTV and F_out agree with each other while KL disagrees with both.

Phrase-Based Prompt Compression

KV cache eviction happens after the full prompt has been processed, so it can't reduce time-to-first-token. Phrase-based compression instead builds a shorter prompt before the model sees it: the context is split into fixed-size phrases, ranked by word overlap with the query, and the best are kept in their original order along with a recency tail. The model gets a structurally intact prompt, with no cache patching and no model internals needed.

Reduces time-to-first-token by 61–74% while staying competitive on KL faithfulness.
Uses only string matching: no forward pass, attention scores or second model.