Reproduce “Train Once, Reuse Everywhere: Generalizable Implicit In-Context Learning by Routing Attention” on OpenResearch

About this paper

Jiaqian Li, Yanshu Li, Ligong Han, Ruixiang Tang, Wenya Wang

Researchers developed In-Context Routing (ICR), a framework that improves Large Language Models' (LLMs) zero-shot performance by implicitly integrating in-context learning (ICL) through direct modulation of attention logits. This approach, which trains a small router while keeping the LLM frozen, demonstrates robust generalization across both in-domain and out-of-domain tasks, achieving a 6.5% accuracy improvement over the best implicit baseline on OOD tasks for Qwen2.5-7B.

Key insights

• Generalizable ICL patterns are found to reside within the attention space of LLMs, and direct modulation of attention logits provides a more principled mechanism for implicit ICL than post-hoc residual stream perturbations.
• Pooled PCA on multi-domain ICL bases effectively recovers domain-stable, general ICL patterns, consistent with the Spiked Covariance Model, which enhances out-of-domain stability.
• ICR's ability to selectively activate 'ICLness' tokens (reasoning-oriented and structural semantics) and induce attention patterns similar to few-shot prompting suggests it effectively internalizes and adapts the underlying structural dynamics of ICL.

Summary by alphaXiv.