Train Once, Reuse Everywhere: Generalizable Implicit ICL by Routing Attention
Abstract
Implicit in-context learning (ICL) has newly emerged as a promising paradigm that simulates ICL behaviors in the representation space of large language models (LLMs), aiming to attain few-shot performance at zero-shot cost. However, existing approaches largely rely on injecting shift vectors into residual flows, which are typically constructed from labeled demonstrations or task-specific alignment. Such designs fall short of utilizing the structural mechanisms underlying ICL and suffer from limited generalizability. To address this, we propose In-Context Routing (ICR), a novel implicit ICL method that internalizes generalizable ICL patterns at the attention logits level. It extracts reusable structural directions that emerge during ICL and employs a learnable input-conditioned router to modulate attention logits accordingly, enabling an efficient train-once-and-reuse framework. We evaluate ICR on 12 real-world datasets spanning diverse domains and multiple LLMs. The results show that ICR consistently outperforms existing implicit ICL methods that require task-specific retrieval or training, while demonstrating robust generalization to out-of-domain tasks where they struggle. These findings position ICR to push the boundary of the practical value of ICL.
Lay Summary
Large language models can often solve new tasks better when they are given a few examples in the prompt. However, including these examples every time can make the prompt longer, slower, and more expensive, and the model’s performance can depend heavily on which examples are chosen. In this work, we ask whether a model can reuse the benefits of examples without actually seeing them at test time. We propose In-Context Routing (ICR), a method that learns how examples usually guide the model’s attention during in-context learning. Instead of storing or retrieving demonstrations for every new query, ICR learns a compact routing pattern from training tasks and uses it to adjust how the model processes new inputs. Our results show that this learned routing can transfer across tasks, including tasks that are different from those used during training. This suggests that part of what makes in-context learning useful can be captured and reused more efficiently. ICR may help make language models faster and easier to deploy when examples are costly, unavailable, or difficult to choose reliably.