Research

I work on how a long-context model should hold on to what it has read. Attention keeps everything and pays for it; recurrent and fast-weight models compress and forget. Most of my work sits on that trade: what the memory is, how it is written, and how it is read back.

Three lines run through it.

• Test-Time Training — the memory is fast weights optimized online during the forward pass. I work on who owns those weights, what supervises them, and whether reading them should be deterministic at all.
• Efficient Architecture — the memory is the KV cache or the residual stream. I work on removing redundant axes from it: depth in the cache, cumulative overlap in layer-to-layer routing.
• KV Cache Compression & Long-Context Inference — the memory is a fixed budget that must be spent. I work on where to spend it: across layers, across heads, across a model’s own chain of thought.

Each project below is written as the bet it makes, what was built, where it was tested, and what came out.

Test-Time Training

Fast weights are a memory that is optimized while the model reads. Three questions about that memory: who owns it, what supervises it, and how it is read.

Universal Test-Time Training

Adobe Research · under submission

Hypothesis. In every existing TTT design, each layer owns a private fast-weight tensor no other layer touches. That convention is inherited from the transformer parameter layout, not from any property of test-time optimization — and it costs three things: a fixed budget divided L ways gives each layer a narrow operator, the split across depth is frozen at design time, and the implementation runs L small operations instead of one grouped update. The bet is that fast weights should be universal, and that because the shared object is the recurrent state rather than static parameters, sharing changes what the model can remember, not merely how its parameters are laid out.

Approach. A bank of fast-weight modules belongs to the model rather than to any block; each block keeps its own slow weights and local attention but owns only an address — a routing function into the bank. Reads take a gated combination of addressed modules; at the chunk boundary, addressed modules are updated by the accumulated contributions of every layer-head that routed to them. Ownership becomes a routing constraint, which lets three designs sit on one ladder: uTTT-Dense (one shared dense operator every layer applies), TTT-MoE (a routed pool, still layer-local — the controlled midpoint that isolates routing from universality), and uTTT-MoE (one globally shared pool). A separate result shows a universal TTT model is a fully recurrent model over the joint (chunk, layer) index: depth becomes a second recurrence axis supplying extra inner-loop steps, and state written deep at one chunk is read shallow at the next. Layer-private TTT is the strict fixed-partition special case.

Evaluated on. Large-chunk autoregressive language modeling at 32K context (Long-Data-Collections) at 124M and 760M parameters — per-token loss on Books3 and all 13 RULER retrieval tasks at 4K/8K/16K/32K, against LaCT, DeltaNet-SWA, Gated DeltaNet-SWA and full attention. LVSM-style novel view synthesis, object-level (Objaverse → Google Scanned Objects) and scene-level (DL3DV).

Results. RULER 13-task mean rises 8.50 → 9.99 → 12.82 → 15.46 at 124M and 23.21 → 24.86 → 25.73 → 27.85 at 760M, across layer-private LaCT → layer-local routed pool (unbalanced, then balanced) → globally shared pool, with the ordering replicating at roughly 6× scale. Only the last of those steps actually changes ownership, so the controlled universality gain is +2.6 points at 124M and +2.1 at 760M; the rest comes from routing and load balancing. On rendering, object-level PSNR goes 24.63 dB (layer-private dense) → 24.79 (layer-local pool) → 25.52 (global pool), where the capacity-matched ownership step — 64 experts partitioned across depth versus shared — is the 24.79 → 25.52 leg. The size-matched dense variant reaches 25.93 dB and widening the active expert set carries it to 26.51 dB. uTTT-MoE reaches lower per-token loss than full attention at both scales while still scoring below it on retrieval.

Where it does not hold. The advantage concentrates at shorter contexts: in the 760M per-length decomposition uTTT-MoE gains +8.9 points at 4K but loses 1.8 at 16K, and at 32K the strongest fast-weight model is a different variant. The scene-level rendering comparison separates the configurations by 0.19 dB, which the authors decline to read as a separation. uTTT-Dense is defined but not evaluated in language modeling, so it appears in the rendering numbers only. And the efficiency half of the argument is unmeasured — no throughput or peak-memory numbers are reported, and the sequential update schedule that would strengthen the depth recurrence forfeits the grouped execution that motivates it.

Test-Time Generation

Adobe Research · under submission

Approach. Memory readout is recast as key-conditioned transport from latent noise to target values. Treating the fast-weight state as a conditional associative memory and bridging noise to values, the retrieval margin decomposes into a key-prior term plus a value-space denoising term amplified by SNR — i.e. iterative recall adds separation that a one-shot lookup cannot. The state is never materialized as a table; it modulates a gated-MLP denoiser. Writes supervise noisy versions of observed bindings with multiple noise draws per binding; reads start from Gaussian noise and integrate for a chosen number of steps, which makes readout depth a test-time compute knob rather than a training choice.

Evaluated on. Incremental novel view synthesis, object-level (Objaverse → GSO) and scene-level (DL3DV-10K), against LaCT-style deterministic readout. Long-context autoregressive language modeling at 124M and 760M on Long-Data-Collections, scored by Book-3 per-token loss over 32K-token sequences. Plus test-time scaling sweeps, one parameter-comparable fairness control, and an objective ablation.

Results. Object-level PSNR improves 25.09 → 26.00 on average and 26.31 → 27.23 on the last evaluated view. Scene-level gains are larger — 16.27 → 17.46 average and 16.80 → 18.64 final — which matches the margin analysis, since pose keys alone cannot disambiguate texture and occlusion, so value-space denoising is where the separation comes from. Language modeling keeps a persistent, if small, loss advantage across the full context range. Both test-time knobs (update repeats, readout depth) improve the quality envelope without retraining. A parameter-comparable control — two TTG variants, run at a lower operating point than the headline — lands at 24.92 vs 24.95 PSNR, which suggests the memory block’s parameter allocation is not what moves the number; it does not include a parameter-matched deterministic baseline, so it is evidence about TTG’s own sensitivity rather than a full attribution. The objective ablation is worth stating plainly: an energy-based variant beats the flow-matching default on the same protocol (24.16 → 24.99 average PSNR), which the authors read as evidence that the generative-memory interface, not flow matching specifically, is the contribution.

Cost and evidence base. TTG buys quality with compute: stochastic repeats make writes more expensive and deeper readout raises per-query latency, and the right operating point is task- and budget-dependent. The paper also flags that its evidence base is thin — a limited set of datasets, mostly single training runs, and no error bars or multi-seed significance tests — so the trends are offered as evidence for the mechanism rather than as a scaling law.

Test-Time Training with Next-Token Prediction

Preprint, arXiv 2026 · co-first author · PDF

Hypothesis. In-place TTT has standardized where fast weights live (MLP down-projections) and how they are written (chunk-parallel rank-one updates), which leaves the supervisory target as the open design axis — and prior recipes get it wrong. They train the fast weight to reconstruct a learned local activation proxy, an engineering device untethered from the next-token-prediction signal that actually trained the backbone. The bet: every prompt token already carries an NTP signal at test time, and aligning the write with the model’s own causal trajectory improves long-context utilization without redesigning the backbone.

Approach. At selected MLP blocks the down-projection is reused in place as the fast weight. The value target is the same layer’s contextual hidden state at the next position — the state produced after actually observing the next token — passed through a lightweight identity-initialized projection. Writes accumulate chunk-parallel as an exclusive prefix sum with the chunk’s last token excluded from pairing, so the update is strictly causal and no future information reaches a token’s own prediction. At inference, since the prompt is observed in full, the write is solved in closed form as a ridge regression over cached activations; the weight is patched, decoding proceeds against the existing prompt KV cache, and the original weight is restored per sample. Only that inference-time write is drop-in — the fast-weight write and its projection are learned in a continual-pretraining phase of 0.1B–2B tokens depending on backbone.

Evaluated on. RULER Full-13 at 4K/8K/16K/32K and general capability via lm-evaluation-harness, across four backbones — Llama-3.1-8B, Mistral-7B-v0.3, Qwen3-4B and Qwen3-0.6B; plus LongBench-v2 medium split (215 multiple-choice questions over 33K–128K-word documents) on Llama-3.1-8B and Mistral-7B-v0.3.

Results. It is the only method that improves RULER over base on all four backbones: Full-13 average +3.90 / +3.03 / +4.06 / +2.88, with the largest single gain +11.74 at 16K on Qwen3-4B. Matched-compute baselines regress instead — continued pretraining costs 5.9 points on Llama-3.1-8B and 7.7 on Mistral. On LongBench-v2 it takes the best overall accuracy on both backbones, 25.6 → 31.2 on Llama and 26.5 → 30.2 on Mistral (continued pretraining and in-place TTT also improve there, by smaller margins). Two ablations carry the argument: the NTP target beats every convolutional proxy at all four lengths — by ~5 points at 4K, ~3 at 8K, ~9 at 16K and 5–8 at 32K on Qwen3-4B — and replacing the ridge solve with a plain Hebbian inner product collapses the method entirely (12.61 vs 59.70), showing the Gram whitening is doing the work. General capability is essentially unchanged.

Efficient Architecture

Same question, different memory. Here the redundancy is structural — an axis of the cache or the residual stream that carries less information than its size suggests.

LastKV: Last Layer Key-Value Cache is All You Need

Adobe Research · under submission

Hypothesis. A long-context KV cache grows along two axes at once — sequence length and network depth. Token-axis methods (eviction, compression, windowing) are well explored; the depth axis is not. Existing cross-layer sharing either stops short of the limit (sharing among small groups of adjacent layers) or pays elsewhere (condensing at every token position, which breaks exact parallel training, or requiring a decoder-decoder redesign). The bet: sharing does not have to apply at every token position, only where long-range memory lives. Deferring the cache write to a chunk boundary makes the last layer’s cache legal for every layer to read, so maximal L-fold sharing becomes achievable with an unmodified Transformer block.

Approach. Each chunk runs as a standard per-layer-KV Transformer with full intra-chunk parallelism. When the chunk completes, only its last layer’s KV pair is committed and the other L−1 caches are freed; later chunks concatenate that shared context into every layer and read it through ordinary attention. Retained cache drops from O(TCLd) to O(TCd). Because language and video need continuity across the seam, a nearest-chunk refinement keeps per-layer caches for just the previous chunk and interpolates them with its last-layer cache through a learnable per-layer per-head weight — 576 added scalars at 760M. The default construction reuses the last layer’s already-computed K/V, so it adds zero parameters and zero FLOPs.

Evaluated on. Four settings, deliberately spanning modalities. Language model pretraining from scratch at 124M/760M/3B on 32K contexts (validation loss on Books3, all 13 RULER tasks). Continued pretraining converting Marin 8B, scored on 12 knowledge/reasoning/math benchmarks plus training throughput. Autoregressive text-to-video, fine-tuning Wan 2.1 T2V at 1.3B and 14B under teacher forcing and Self-Forcing distillation, scored on the full VBench protocol. Novel view synthesis, object-level and scene-level.

Results. At 760M all three LastKV variants land below the full-cache baseline on validation loss (2.5373 / 2.5366 / 2.5375 vs 2.5423). The continued-pretraining result is the one that matters most: an openly released 8B Transformer converted to retain one layer of KV out of 32, scoring within ~1 point of the unconverted baseline on 8 of 12 benchmarks and −0.6 on average — while sustaining 99.6% of forward and 95.6% of forward+backward throughput at chunk size 1024. That comparison is a snapshot of an ongoing run, taken at a matched checkpoint roughly 120B tokens into a 2T-token schedule, so it claims parity in an early-budget window rather than an advantage. Rendering improves (25.38 → 25.95 dB object-level, 17.46 → 17.81 scene-level). Video holds under teacher forcing (VBench Total 80.29 → 80.45 at 14B, 82.09 → 81.84 at 1.3B), though under Self-Forcing distillation it trails the full cache by 0.4–0.6 total points, almost entirely through dynamic degree. The boundary ablation is clean: without the nearest-chunk refinement the loss spike after each boundary is +0.103 nats; with it, −0.003.

Where the cost lands. Retrieval is mixed rather than preserved — at 760M with gated attention the 13-task RULER mean drops 43.6 → 29.6 at 16K and 34.5 → 31.6 at 32K, with multi-key tasks favouring the full cache, though without the gate the ordering reverses at 8K–32K. And there is no end-to-end serving benchmark: no decode latency, tokens/sec, or maximum context under a fixed memory budget on an optimized stack.

Delta Attention Residuals

Preprint, arXiv 2026 · co-first author · PDF

Hypothesis. Attention Residuals replace fixed additive residuals with learned softmax attention over previous layer outputs — but they attend over cumulative hidden states, which are running sums sharing an ever-larger common prefix as depth grows. The bet is that source redundancy in the routed states — compounded by replacement routing that discards the residual stream at each reset — is what breaks the method: softmax logits over near-identical sources go low-contrast, routing collapses toward uniform, and the mechanism degenerates into averaging.

Approach. Two changes. Delta sources: route over per-sublayer outputs (the difference between consecutive hidden states) rather than cumulative states — attention and MLP outputs occupy different subspaces, so the sources are structurally diverse by construction. Additive routing: add the routed combination to the residual stream instead of replacing it, so the stream is never discarded and no contribution is lost at block boundaries.

Evaluated on. Language model pretraining from scratch on FineWeb-Edu at five scales — 220M, 533M, 1044M, a Qwen3-0.6B configuration, and a 7.57B configuration — plus fine-tuning conversion of a pretrained Qwen3-0.6B scored on 8 downstream benchmarks, and a direct measurement of routing sharpness across depth.

Results. The diagnosis is confirmed by a crossover: at 1044M, where the baseline reaches 29.70 validation perplexity, AttnRes gets worse than doing nothing (31.76) while the delta formulation improves it (29.13). At 7.57B the gap widens — baseline 17.43, AttnRes 18.58, Delta Block 16.00, an 8.2% improvement — for 589.8K routing parameters (0.008% of the model), at a 35% training-throughput cost and +3% memory. The mechanism check matches the story: AttnRes’s maximum routing weight decays from ~1.0 in early layers to ~0.2 in deep ones, while the delta version holds ~0.6, a 1.8× higher average.

Cost. The finest-grained variant, which produces the best perplexity at 1044M, costs 69% throughput and 3.5× memory; the coarsened Delta Block is the recommended default precisely because it keeps most of the gain at a third of that overhead.

KV Cache Compression & Long-Context Inference

A fixed memory budget, and the question of where to spend it. This line began my PhD work and is the most widely adopted.

PyramidKV: Dynamic KV Cache Compression based on Pyramidal Information Funneling

COLM 2025 · Spotlight (24 / 418) · PDF · Page · Code

Hypothesis. Every prior eviction method gives every layer the same budget. But attention is not uniform across depth — it is broadly scattered in lower layers, consolidates into local patterns in the middle, and collapses onto a few critical tokens at the top. If information genuinely funnels, a uniform budget wastes memory at the top while starving the bottom.

Approach. Distribute a fixed total budget across layers as a decreasing arithmetic sequence — lower layers get more, upper layers get less, total preserved. Within each layer and head, keep the tokens receiving the most attention mass from a trailing instruction window. Training-free, applied at prefill, orthogonal to quantization.

Evaluated on. LongBench (17 datasets: single- and multi-document QA, summarization, few-shot learning, synthetic retrieval, code completion), Needle-in-a-Haystack, with cache budgets swept from 64 to 2048, on LLaMA-3-8B/70B-Instruct and Mistral-7B-Instruct.

Results. Matches full-cache quality while retaining 12% of the KV cache. In the extreme regime — 0.7% retained — it gains 19.5 absolute points on TREC over the next-best method at budget 64 (58.00 vs 38.50 on LLaMA-3-8B-Instruct). On Needle-in-a-Haystack, 128 retained entries are enough for LLaMA-3-70B-Instruct to reach 100.0 accuracy at contexts up to 8K.

R-KV: Redundancy-aware KV Cache Compression for Reasoning Models

NeurIPS 2025 · PDF · Page · Code · Video

Hypothesis. KV compression was designed for long input prompts. Reasoning models invert the problem: the bloat comes from the model’s own chain of thought, generated token by token and heavily self-repetitive. Pure attention-importance eviction then fails, because it faithfully keeps many near-duplicates of repeated reasoning steps — in practice retaining only ~60% of full-cache accuracy at a 10% budget. The bet: add an explicit redundancy penalty and the loss disappears. The sharper claim: removing redundant reasoning tokens is a form of denoising that can beat the full cache.

Approach. Score each token by attention importance minus its average cosine similarity to the other retained keys, so tokens that are attended-to but non-duplicative win. Run it on-the-fly during decoding, re-scoring whenever a decode buffer fills, rather than only at prefill. Training-free.

Evaluated on. Mathematical reasoning with long chain-of-thought — MATH-500 and AIME 2024 — on DeepSeek-R1-Distill-Llama-8B and DeepSeek-R1-Distill-Qwen-14B, with throughput and batch-size measurements at 16K sequence length.

Results. On AIME-24 with R1-Llama-8B, R-KV is lossless at a 10% budget where attention-only eviction reaches ~60%, and hits 105% of full-cache accuracy at 16% — the denoising effect, confirmed. The budget needed for lossless compression is setting-dependent: 34% on MATH-500 for the same model, and 25% / 54% for R1-Qwen-14B. In systems terms: 90% memory saving, 6.6× throughput, and 9× larger batches at 10% compression; at a fixed 1024-token budget, 13.4× larger batches and 9.2× throughput.

Not All Heads Matter: Head-Level KV Cache Compression

ICLR 2025 · PDF · Code

Hypothesis. Layer-level allocation — PyramidKV included — implicitly treats all heads within a layer as equally important. Head importance is in fact highly non-uniform, so per-head allocation is a strictly finer control knob. The sharper bet concerns how importance is measured: the standard “retrieval head” score is estimated by pure copy-paste needle retrieval, which rewards verbatim copying, whereas real long-context QA needs retrieval plus reasoning over what was retrieved.

Approach. Estimate head importance offline with a probe whose needle contains a reasoning step, a distractor, and the correct answer, crediting a head for attention mass placed on the entire correct answer rather than on a single argmax position. Freeze those scores and redistribute a global budget across heads.

Evaluated on. LongBench contextual QA, LooGLE long-dependency QA, Needle-in-a-Haystack, and a Reasoning-in-a-Haystack probe, on Llama-3-8B-Instruct and Mistral-7B-Instruct at budgets from 64 to 1024.

Results. Retains 1.5% of the KV cache at 97% of full-cache performance on contextual QA. The gains concentrate exactly where the hypothesis predicts — in the low-budget regime: at budget 64 it leads the strongest baseline (Ada-SnapKV) by ~3.4 points on Llama-3-8B-Instruct and ~3.6 on Mistral-7B-Instruct, and layer-level allocation by more than 5.

HeadInfer: Memory-Efficient LLM Inference by Head-wise Offloading

ICML 2025 Workshop on Long-Context Foundation Models · PDF · Code

Hypothesis. The lossless counterpart to compression: rather than choosing what to discard, treat the memory wall as a data-placement problem. Existing offloading is layer-wise, which still forces a full layer’s cache onto the GPU and so still scales with hidden dimension. But attention heads are computationally independent — so the resident footprint can shrink by a further factor of the head count, with zero approximation.

Approach. Keep one head’s KV cache GPU-resident at a time and stream the rest from CPU RAM, with chunked prefill, context-length-adaptive head grouping, and ping-pong asynchronous prefetch so PCIe transfer hides behind computation.

Evaluated on. LongBench v2, SCBench (9 tasks averaging 227K tokens), RULER at 4K–128K, and Needle-in-a-Haystack — measuring maximum attainable context and throughput on consumer hardware.

Results. For Llama-3-8B at a 1M-token sequence, GPU KV cache drops from 128 GB to 1 GB and total GPU memory from 207 GB to 17 GB — a 92% reduction. That is enough to run 4M-token inference for an 8B model on a single 24GB RTX 4090 backed by ~512GB of host RAM — the 4M ceiling is set by CPU memory, not by the GPU — against 25K tokens for standard inference and 45K with layer-wise offloading. Prefill throughput is essentially untouched at moderate context (7210 vs 7235 tok/s at 20K).

Honest limit. This buys reach, not speed, and the decoding cost shows up well before the extremes: already 6 vs 33 tok/s at 20K, and 0.15 tok/s at 1M tokens.

Open-Source

KVCache-Factory — a unified implementation of KV cache compression methods for autoregressive models, and the reference implementation of PyramidKV.

R-KV — decoding-time KV cache compression for reasoning models.