Yang Jing's  Blog

Efficient Test-Time Memory Learning in Speculative Decoding

TL;DR

We propose integrating trainable memory modules (RNN or Diffusion-based) into the draft model of speculative decoding pipelines to enable efficient test-time adaptation. Rather than updating memory in the main model (heavy cost), we leverage the lightweight draft model to learn task-specific context representations in parallel with token verification.

This approach achieves the speed benefits of speculative inference while progressively optimizing memory representations, outperforming LoRA-based alternatives in flexibility and computational efficiency.

1. Introduction

image.png

Large language models face significant limitations in context memorization, dynamic adaptation, and handling long-sequence generation. Solving these problems typically requires expensive fine-tuning. Existing solutions either introduce training overhead (like LoRA, Prefix-Tuning) or are too rigid to learn new knowledge (like standard speculative decoding).

Here's the core insight: The draft stage of speculative decoding already performs computations. What if we could use those same computations to update memory as well? Instead of discarding gradient signals during verification, we can use them to evolve lightweight memory in the draft model. The main model stays frozen, the computational overhead gets amortized across work we're already doing, and the approach works for both RNN and Diffusion architectures. It's not quite a free lunch, but it's remarkably close.

Recent research has made meaningful progress on individual subproblems, though most approaches remain siloed.

Titans demonstrated that memory doesn't have to be ephemeral. By combining long-term memory modules with standard attention, Titans can compress and retrieve historical context across windows of 2 million+ tokens. The design is elegant: attention acts as working memory while learned modules serve as persistent storage. The cost? Inference speed suffers due to additional computation.

Speculative Decoding solves a completely different problem. The insight is almost deceptively simple: have a small draft model propose multiple tokens at once, then have the main model verify them in parallel. Invalid tokens are discarded. This achieves 2–4× speedup with minimal quality loss. However, the draft model is frozen—typically just a pruned or quantized version of the main model—so it can't adapt to the current task or remember the context it's already processed.

LoRA and variants enable adaptation without full fine-tuning through low-rank updates. These updates are surprisingly expressive and can even be swapped at test time. The limitation: gradients still flow through the main model, introducing unavoidable inference overhead.

These three research directions converge on a gap: a method that adapts and memorizes without touching the main model or requiring a separate training phase. That's exactly where our approach fits in.

3. Method: RNN-Based Draft Model Memory

Architect ，ureHere's how the pipeline works:

image.png

Speculative Decoding Pipeline:

Draft Phase (RNN Memory Update):
  ├─ h_{t-1} = Previous memory state
  ├─ x_t = Current tokens + prompt
  ├─ h_t = RNN(h_{t-1}, x_t, [context_embeddings])  ← TRAINABLE
  ├─ memory_t = Memory_Encoder(h_t)
  └─ draft_tokens ~ LLM_draft(prompt + memory_t)

Verification Phase:
  ├─ verified_tokens = check(draft_tokens)
  ├─ loss = CrossEntropy(verified_tokens)
  └─ grad = ∂loss / ∂h_t  ← Backprop to memory RNN

Update:
  └─ h_t ← h_t - α * grad  ← Lightweight SGD step

During the draft stage, the RNN memory module receives the previous timestep's hidden state and current token, updating to a new memory representation that gets injected into draft token generation. During verification, the main model validates draft tokens and produces a cross-entropy loss that backpropagates to the RNN's memory parameters, completing one lightweight SGD update. The main model never participates in parameter updates.

We can choose from several proven RNN architectures. LSTMs offer robust gating mechanisms that excel at capturing long-range dependencies—a critical capability when working with extended context windows. GRUs provide a lighter-weight alternative with simpler structure, faster computation, and fewer parameters, making them ideal when latency matters. For cutting-edge applications, Mamba/S6 represent modern linear RNNs with O(n) complexity, specifically designed for long sequences without the quadratic scaling of traditional attention.

Draft models are designed to be lightweight, and RNNs fit this profile perfectly. RNNs naturally capture sequence evolution token-by-token, aligning with speculative decoding's autoregressive nature. Gradient signals flow cleanly from verification back to the memory module, providing complete learning signals. Critically, during the draft stage, memory updates happen in parallel with main model inference—so the additional computational cost is essentially zero.

4. Method: Diffusion-Based Draft Model Memory

4.1 Architecture

The Diffusion approach models memory state as a random variable evolving in latent space, using score matching for gradient-based optimization:

image.png

Diffusion-Based Memory Evolution:

Initialization:
  z_0 ~ N(0, I)  or z_0 = learned_prior

Forward Diffusion (Inference):
  z_t = SQRT(1 - β_t) * z_{t-1} + SQRT(β_t) * ε_t

Score Guidance (Gradient-Based):
  ∇_z loss = backward(LLM_draft(memory_decoder(z_t)))

Reverse Step (Optimization):
  z_{t+1} ← z_t - η * ∇_z loss  ← Score matching for memory

Decoded Memory:
  memory_t = Memory_Decoder(z_t)

By performing diffusion in latent space, the optimization becomes smoother and the loss landscape easier to navigate. Each token can take multiple gradient steps, further improving memory quality. Memory uncertainty can be naturally quantified through a probabilistic framework, while reverse diffusion enables "rollback" when bad updates are detected.

4.2 RNN vs. Diffusion

Dimension RNN Diffusion
Speed Fast, low latency Slower, multi-step overhead
Optimization Stability Moderate, gradient-clipping dependent Smooth, more stable latent optimization
Uncertainty Modeling Unsupported Naturally supported
Reversibility Irreversible Reversible via reverse diffusion
Implementation Complexity Low Medium to high
Best For Latency-sensitive tasks High-quality applications

5. Comparison: Memory Update vs. LoRA

5.1 Why Memory Update > LoRA?

image.png

Aspect Memory Module LoRA
Location Draft model (lightweight) Main model (expensive)
Gradient Flow Only memory params All model weights
Overhead O(hidden_dim²) O(d × r) for each layer
What it learns Context representation Weight corrections
Interpretability Memory state (checkpointable) Implicit in weight updates
Scalability Constant with model size Linear with model layers
Test-time speed Negligible (amortized in draft) Noticeable (gradient computation)

Memory modules work in the draft model, never touching the main model, so their computational cost is independent of main model scale. LoRA requires gradients to flow through all main model weights—overhead that grows linearly with layer count on large models.

LoRA shines when tasks require changing deeper model behaviors. For fine-grained weight adaptation to specific tasks or when transfer across tasks matters, LoRA has clear advantages.

Memory modules and LoRA aren't competitors—they're complements. Memory modules handle context compression and immediate adaptation; LoRA handles task-level weight refinement. Combined, they enable more comprehensive multi-level adaptation.

6. Practical Implementation Considerations

Memory State Management: Initialize memory via zero initialization, learned priors, or embeddings from pre-training. Reset strategy depends on task: for in-context learning, reset per prompt; for long document generation, maintain continuity within sequences. In multi-turn dialogue, consider preserving memory across turns to capture longer conversation history.

Gradient Stability: Test-time gradient updates lack batch norm's protective effects, so gradients can become unstable. We recommend several practical mitigations: apply Layer Norm to memory states for normalization stability, set a gradient clip threshold around 1.0 to prevent explosions, use exponential moving average (EMA) for memory smoothing to reduce jitter, and keep learning rates conservatively small—typically in the 1e-4 to 1e-3 range.

Computational Overhead : For a typical 7B main model, the overhead remains remarkably small. You'd typically see hidden dimensions around 4096, with memory state dimensions of 4096 or 8192. The RNN parameters add up to roughly 50M—about 1% of the main model. A single gradient update runs in under 5ms on GPU, and the amortized cost per token is approximately 0.1ms. Critically, since memory updates happen in parallel with draft token generation, the actual amortized latency impact is minimal—you're essentially getting this for "free" by overlapping computation.

Draft token acceptance rate directly impacts memory update quality. Above 90% acceptance suggests reliable, low-noise signals; below 70% suggests noisy updates needing filtering or regularization. An adaptive strategy: dynamically decay learning rate based on acceptance rate, reducing it when acceptance drops to suppress noisy updates.

7. Experimental Design & Key Questions | TODO

This approach raises several key empirical questions worth investigating systematically.

Convergence and Stability: On in-context learning tasks, we should compare frozen memory (standard speculative decoding) against trainable modules, plotting loss curves over the first 100 tokens and tracking accuracy across generation length. This tells us whether the learning signal is actually productive.

Speed-Quality Tradeoff: Under fixed inference budgets, we'd compare RNN, Diffusion, and LoRA approaches on diverse tasks—in-context learning, long-document QA, and reasoning—measuring BLEU/ROUGE scores, latency, and memory footprint. This reveals which method delivers the best bang for your computational buck.

Generalization: Train memory on one task, evaluate on another. This measures whether learned memory representations transfer across problem domains or if they overfit to specific tasks.

Ablations: Test different RNN variants (LSTM vs. GRU vs. Mamba), various update frequencies (every token vs. every 5 tokens), different learning rate schedules, and initialization strategies.

Based on theoretical analysis, we expect to see roughly 2–8% accuracy gains on in-context learning tasks, 1–3% improvements on generation quality metrics, with memory typically converging within 50–200 tokens depending on task complexity.

8. Discussion and Future Directions

8.1 Open Questions

Does improved memory enhance draft token quality or introduce more noise? Should the main model also have memory? Can memory states be shared between them? How should memory dimension and update frequency scale to 70B+ models? In batch inference, can multi-sequence memory updates be efficiently parallelized?

8.2 Potential Extensions

Multi-task memory represents one natural extension. Rather than a single memory module, you could maintain separate modules for different task types with task-aware routing that learns specialization. This trades a bit of parameter overhead for the ability to excel at diverse problems simultaneously.

Persistent memory is another direction: save learned memory across prompts to build genuine long-term user or task context. This opens possibilities for truly personalized model behavior that evolves over extended interactions.

Hierarchical memory combines coarse-grained and fine-grained levels, allowing the system to capture both global structure and local detail. Think of it as a two-level filing system—important themes at one scale, specifics at another.

Integration with RAG completes the picture. Memory modules handle context compression and immediate adaptation; retrieval-augmented generation handles external knowledge retrieval. These aren't competitors—they're complementary components of a robust system.

9. Conclusion

We've proposed integrating trainable memory modules into speculative decoding's draft model. By reusing computation that already happens, we achieve continuous test-time memory updates with negligible overhead. The approach leaves the main model untouched, scales independently of model size, and plays well with existing techniques like LoRA. RNN variants suit latency-critical scenarios; Diffusion variants suit quality-critical ones.

As LLM inference optimization evolves, lightweight adaptation mechanisms in inference pipelines will become increasingly important. Draft models' dual utility represents a promising entry point.

References