Dynamic Steering With Episodic Memory For Large Language Models (2025)

Abstract: Large Language Models (LLMs) exhibit emergent in-context learning (ICL) capabilities, allowing them to adapt to unseen tasks based on example demonstrations. Traditional ICL embeds examples within the prompt, while activation steering, uses a vector derived from examples to guide the latent states of LLMs toward desired behaviors. However, traditional ICL is difficult to control quantitatively and consumes valuable context space. Existing activation steering methods apply a single sentence-level steering vector uniformly across all tokens, ignoring LLMs’ token-wise, auto-regressive nature. This coarse control can lead to inconsistencies and suboptimal adjustments during generation. To address this problem, we introduce Dynamic Steering with Episodic Memory (DSEM), a novel training-free framework that aligns LLMs to given demonstrations by steering at the token level conditioned on the input query. DSEM employs a key-value memory to store associations between generated tokens and steering vectors. During inference, it uses a nearest-neighbor mechanism to dynamically compute steering vectors for each token chunk, enabling more precise and adaptive guidance. Our method surpasses strong baselines across diverse alignment tasks - including safety, style transfer, and role-playing - demonstrating improved alignment as demonstration size scales.

## introduction large language models (llms) have demonstrated remarkable performance in a wide range of natural language processing tasks, with their effectiveness improving as the number of parameters increases @xcite @xcite . one of the most intriguing abilities of larger models is in-context learning (icl), where llms can perform tasks they have not encountered during training, given only a few demonstrations @xcite . this enables an alternative approach to achieve strong downstream performance by modifying the input prompt in text space, incorporating instructions and in-context examples @xcite . despite the intriguing capabilities of icl, it is highly sensitive to nearly every aspect of the prompt. the selection of demonstrations, the phrasing of instructions, or even subtle variations in example formatting can lead to significant fluctuations in final performance @xcite . this sensitivity raises concerns about the robustness and reliability of icl in realworld applications. furthermore, a fundamental limitation of traditional icl is the restricted context length of most open-source llms, which constrains the number of demonstrations that can be included in the prompt. as a result, effectively utilizing icl requires careful prompt engineering and selection of examples to mitigate these challenges. one recent approach to address the limitations of traditional icl is activation steering, which involves modifying the activation layers of llms. for instance, @xcite have proposed computing the principal direction of the difference between source and target representations and using this as a single steering vector for all test data. more recently, @xcite introduced a binary masking method to identify the most influential elements in the steering process. a fundamental limitation of these methods is that they are either query-agnostic or lack fine-grained control, hampering final alignment performance. this raises a research question: how can activation steering be adaptively applied in a query-dependent manner and at the token level to effectively guide llm generation toward a desired goal? to this end, we propose dynamic steering with episodic memory (dsem), a novel and efficient steering method designed to steer llm generations flexibly, enabling quick adaptation to a given input and fine-grained control over the genera-tion process. drawing inspiration from the rapid and instance-based learning mechanisms observed in the hippocampus region of the human brain @xcite , our method utilizes an external episodic memory that stores input representations, partially generated outputs, and their stored steering vectors. we use the input-output demonstrations to define the partial outputs and compute the stored steering vectors. during inference, dsem utilizes the memory to generate steering vectors for each input query. it retrieves stored steering vectors and estimates the chunk-level steering vector using a nearest-neighbor mechanism, ensuring that steering vectors adapt across queries instead of remaining static. another significant benefit of dsem is its ability to compute steering vectors at the token level. instead of using a fixed steering vector for an entire sentence, dsem dynamically determines steering vectors for generating the next few tokens. by leveraging token-level information saved in memory (i.e., the partially generated outputs), dsem retrieves steering vectors that are better aligned with the desired outputs. this approach is motivated by the fact that next-token prediction in transformers @xcite is inherently a token-wise process, suggesting that steering should also be applied at the token rather than the sentence level. finally, to preserve the global semantics of the demonstration, dsem retrieves the most relevant steering vector at inference time and interpolates it with a global steering vector during each token generation step. this interpolation ensures a smooth transition between context-sensitive adjustments and overarching guidance, preventing abrupt shifts in model behavior. the interpolation weight is determined adaptively based on the uncertainty of the memory prediction, allowing dsem to balance local specificity with global coherence. in summary, our main contributions are: (1) we propose dsem, a fine-grained activation steering method that dynamically interpolates between memory-based and global steering vectors, enabling precise control over generation. (2) we propose an uncertainty-driven interpolation mechanism that dynamically balances local specificity with global coherence. (3) we empirically demonstrate that dsem outperforms strong baselines, including fine-tuning and state-of-the-art activation steering methods, across multiple tasks such as detoxification, formality adjustment, and roleplaying, as well as across different model architec-tures and demonstration sizes. ## experiments our experiments focus on open-source llms: llama-2-7b-chat, gemma-2-9b-it, falcon-7b, and mistral-7b-instruct-v0.2 @xcite @xcite . the same top_p and temperature values are used across all methods within each model. ## related work traditional icl improvements. several methods have been proposed to improve icl, including instruction tuning, demonstration selection, and ordering @xcite @xcite . instruction tuning fine-tunes models on diverse tasks to enhance generalization, while demonstration selection identifies the most relevant examples to optimize performance. ordering strategies explore how the arrangement of examples impacts output. while these approaches enhance icl by editing instructions or selecting and arranging examples, they remain constrained by the fixed context window, limiting scalability. activation steering refers to techniques that modify the latent space of models to produce desired outputs. burns et al. (2022); park et al. (2024) demonstrate that latent knowledge encoded in the activation space is linearly separable. liu et al. (2024) proposes finding the principal direction of the steering vectors in a demonstration set and applying it to the entire test set. wang et al. (2025) introduces a binary mask for the latent space to enhance the values of the most important elements. however, while liu et al. (2024) identifies a queryagnostic steering vector, wang et al. (2025) relies heavily on brute-force search to determine the optimal number of important elements and the strength of the intervention, a limitation acknowledged in their work. in contrast, dsem introduces a dynamic, query-dependent approach to identify steering vectors at the token level, allowing the steering direction to adapt smoothly during generation. ## conclusion we introduce dynamic steering with episodic memory (dsem), a training-free framework that aligns llms with given demonstrations by steering at the token level, conditioned on the input query. dsem stores representation-vector associations and uses a nearest-neighbor mechanism to extract the steering vector during generation. it leverages prediction uncertainty to determine the interpolation weight with a global steering vector. evaluations on safety, style transfer, and role-playing tasks demonstrate its effectiveness, with detailed model analyses highlighting its advantages.