Sample Efficient Alignment Learning With Episodic Control (2025)

Abstract: Aligning large language models (LLMs) with specific task objectives is challenging, especially when access to feedback signals for guiding the model is limited. While existing parametric methods perform reasonably, they rely heavily on large datasets and frequent feedback, making them impractical in scenarios with limited human feedback. We introduce Alignment Learning with Episodic Control (ALEC), a non-parametric framework that aligns LLM outputs during inference without fine-tuning. ALEC employs a key-value memory to store the associations between generated text and its corresponding values. It leverages a novel confidence-based writing scheme to update these stored values, maximizing the use of available data. During inference, ALEC utilizes a nearest-neighbor mechanism to estimate the values of generated texts, enabling the selection of the optimal text for decoding. Our method outperforms state-of-the-art baselines on harmless, helpful, and summarization tasks, demonstrating improved alignment with minimal interactions with the true reward model.

## introduction large language models (llms) have demonstrated remarkable performance across a wide range of natural language processing tasks, with their effectiveness improving as their parameter count grows @xcite @xcite . however, despite these advancements, pre-trained llms often generate responses that are misaligned with human preferences and values, largely due to the vast and diverse nature of the data used to train them @xcite . this misalignment underscores the critical challenge of ensuring that llms meet performance benchmarks and adhere to human objectives and safety considerations. addressing this misalignment is particularly difficult, as aligning models to human values typically requires fine-tuning based on expert feedback, which is both costly and scarce @xcite . the high expense and limited availability of expert input highlight an urgent need for more efficient alignment learning methods that reduce dependence on expert feedback while ensuring that llm outputs remain consistent with human goals. a key approach to addressing the alignment challenge is the reinforcement learning from human feedback (rlhf) framework @xcite @xcite , where a reward model (rm) evaluates outputs based on human preferences, guiding the llm to adjust accordingly. direct policy optimization (dpo) offers an alternative by aligning models directly with preference data, avoiding rlhf's training instability @xcite . despite their effectiveness, both rlhf and dpo rely heavily on frequent reward model interactions and large amounts of preference data, which may not always be available. moreover, they require fine-tuning llms, which can be expensive and impractical for applications with limited resources. another approach in alignment learning focuses on inference-time methods. for instance, @xcite introduces a lightweight prefix scorer applied during inference, while @xcite and @xcite use rl-based frameworks that combine signals from a scorer and the base model. these methods allow output refinement without fine-tuning the llms, but they heavily rely on training a parametric value model to estimate candidate quality. this approach faces severe limitations when rm calls are restricted. small models often fail to effectively map candidates to value scores @xcite , while large, complex models are prone to overfitting on limited preference data and suffer from slow execution times @xcite , making them impractical in low-resource scenarios where frequent reward queries are not feasible. in this paper, we address the research question: how can alignment learning be effectively achieved under constraints on the number of queries to the expert reward model? to this end, we propose alignment learning with episodic control (alec), a novel and efficient training-free approach designed to perform alignment learning with minimal reliance on expert reward model queries. drawing inspiration from the rapid and instance-based learning observed in the hippocampus region of the human brain @xcite , our method avoids the complexities and overfitting risks associated with training parametric estimators while maintaining reliability through performance-driven optimization. we argue that in scenarios where sample efficiency is critical, our non-parametric, memory-based learning framework can outperform traditional parametric methods. concretely, we leverage episodic control @xcite , treating each training prompt-response sequence as an episode and storing its value in memory. the traditional storing method, which only inserts the episode's states, fails to revisit or refine stored states' values, leading to inaccuracy in estimation. to address this, we propose a confidence-based memory writing that optimizes the sample order for memory writing and updates neighboring states to reduce bias toward familiar states, thereby enhancing overall performance. at test time, the memory functions as a non-parametric nearest-neighbor model, using stored values to approximate the expert reward model. it guides generation by selecting the next token candidates that maximize the estimated values. we show the overall framework in figure 1 . in summary, the main contributions of this paper are: (1) we propose alec, a training-tuning-free framework for determining optimal sequences of tokens at inference time using a nearest-neighbor reading from episodic memory. (2) we introduce a confidence-based writing scheme that enhances value estimations stored in our memory. ## related work generator improvement solutions for alignment learning. fine-tuning lms to reflect human preferences has gained attention recently due to its flexible nature. in terms of generator improvement solutions, reinforcement learning from human feedback (rlhf) offers a direct route @xcite . within the rlhf framework, one notable solution is to use proximal policy optimization and its variants @xcite @xcite . however, due to the instability and data-intensive nature of rl training, researchers have sought alternative approaches. @xcite derives a preference formula and alleviates rl training challenges by directly optimizing lms based on preference data. inference time alignment learning. this approach uses inference-time solutions for alignment learning without fine-tuning the base lm, where the model is sampled multiple times to select the best output. st-bon (2025) improves sampling efficiency with early truncation, but relies on heuristics rather than learned control. @xcite introduces a method that trains a prefix scorer using a deep-q network (dqn) framework @xcite , providing additional signals for sampling. similarly, @xcite employs a scorer to enhance sampling, but uses a standard reward model training approach based on the bradley-terry model (bradley and @xcite . however, training these scorers can be data-intensive, and their performance may degrade when only a limited number of queries to the true rm are available. ## experiments we aim to demonstrate how alec can guide llms to decode in alignment with human preferences while minimizing calls to the expert reward model. the experiments are conducted using 3 base opensource lms: llama-2-7b-chat-hf; vicuna-7b-v1.5; mistral-7b-instruct-v0.2 @xcite @xcite . ## model analysis and ablation study why not using generation llm embeddings? we include results from alec using generation llm embeddings to validate our choice of sentence embeddings. the model we use is llama-2-7b-chat-hf. for embedding type llama-2-7b-chathf, we take the mean pooling of embeddings across all layers across all non-padding token of the last hidden state. as seen in ## conclusion we introduce alignment learning with episodic control (alec), a non-parametric framework that aligns llm outputs during inference without finetuning. alec stores key-value associations and uses a confidence-based writing scheme for data efficiency. evaluations on harmless, helpful, and summarization tasks show its effectiveness, with detailed model analyses highlighting its performance advantages with minimal rm interactions.