SLOT: Sample-specific Language Model Optimization at Test-time

TL;DR

SLOT introduces a lightweight test-time adaptation where a per-sample delta is added to the final hidden features before the output head, enabling gradient-based refinement of a single vector for each prompt. By optimizing this delta over a few steps using only the prompt as supervision, SLOT shifts the logits via a logit modulation vector to encourage deeper reasoning without full-model fine-tuning. Empirical results across diverse LLMs and benchmarks show meaningful gains on GSM8K and GPQA Diamond, including state-of-the-art performance for open-source 70B models, while incurring only modest inference-time overhead due to caching and updating just the last-layer representation. The approach highlights the potential of efficient, sample-specific test-time training to better align LLM behavior with individual instructions, with broad implications for reasoning-heavy tasks and practical deployment.

Abstract

We propose SLOT (Sample-specific Language Model Optimization at Test-time), a novel and parameter-efficient test-time inference approach that enhances a language model's ability to more accurately respond to individual prompts. Existing Large Language Models (LLMs) often struggle with complex instructions, leading to poor performances on those not well represented among general samples. To address this, SLOT conducts few optimization steps at test-time to update a light-weight sample-specific parameter vector. It is added to the final hidden layer before the output head, and enables efficient adaptation by caching the last layer features during per-sample optimization. By minimizing the cross-entropy loss on the input prompt only, SLOT helps the model better aligned with and follow each given instruction. In experiments, we demonstrate that our method outperforms the compared models across multiple benchmarks and LLMs. For example, Qwen2.5-7B with SLOT achieves an accuracy gain of 8.6% on GSM8K from 57.54% to 66.19%, while DeepSeek-R1-Distill-Llama-70B with SLOT achieves a SOTA accuracy of 68.69% on GPQA among 70B-level models. Our code is available at https://github.com/maple-research-lab/SLOT.

Figures (4)

• Figure 1: SLOT significantly boosts Qwen-2.5-7B's GSM8K format and answer accuracy at test-time. The left shows the presented prompt and question and compares the output responses from the language model without and with the SLOT used. The right graph shows that both answer accuracy (blue, left axis) and format accuracy (red, right axis) are continuously improved with increasing SLOT optimization iterations from 0 to 5.
• Figure 2: SLOT pipeline during inference. The process consists of two stages: (1) Prompt Stage: Sample-specific parameters $\delta \in \mathbb{R}^{1 \times d}$ are initialized and optimized over $T$ iterations to minimize cross-entropy loss on the input prompt. Each sample in the batch has its own parameters (shared across sequence positions), allowing for efficient adaptation. The original hidden features $H$ are modified by adding $\delta$ to produce $H'$. (2) Generation Stage: During token generation, the optimized $\delta$ parameters are reused without further optimization, modifying the hidden features of newly generated tokens. This approach achieves performance gains with minimal computational overhead, as optimization cost is amortized across the entire generation process. The dashed line illustrates parameter reuse between stages.
• Figure 3: About the most increased and decreased tokens in Qwen-2.5-7B on GSM8K with SLOT(iter=5). We compute the $W_{\text{LM}}(\delta)$ as the degree of change of the logits of the tokens. We find that the most increased tokens are ‘reasoning' related tokens(e.g. reasoning, think, thinking), and the most decreased tokens are numerical tokens(e.g., 0, 1, 2, 3). Also, we find the modal verbs (e.g. should, will, can) is decreased, Figure \ref{['fig:SLOT_iters']}'s left answer example shows, the wrong anser uses many ‘I'll' but don't do actual caculations.
• Figure 4: Samples of the comparison between with SLOT and without SLOT in GSM8K. We mark the mistakes in the original responses with red, and the correct answer with blue.