EMO: A Closed-Loop Test-Time Computation Framework via Alternating Context Optimization and Logits Steering
Authors: Wei Xia, Yehui Tang, Zhi-Hong Deng
Organizations: State Key Laboratory of General Artificial Intelligence, School of Intelligence Science and Technology, Peking University · Samsung Research, Beijing, China
Test-time context optimization adapts frozen language models through textual feedback, but ordinary conditional generation provides no separate control over how strongly to reinforce the distributional changes induced by a context update. We introduce EMO, a training-free framework inspired by expectation--maximization that alternates question-specific context estimation (E-step) and distribution refinement (M-step). The E-step revises a system context using textual gradients from an external judge without access to gold answers. The M-step contrasts next-token distributions under the revised and preceding contexts at a shared prefix to construct a confidence-constrained logit correction. Feedback on the steered response drives the next context update, coupling adaptation across rounds with intervention within each generation. Theoretically, we connect parameter and context optimization through a shared target-KL objective and additive logit representations, and derive the steering vector from context-induced logit displacements. Under context-improvement and faithful-steering assumptions, we prove a one-round decrease in target KL divergence for sufficiently small steering strength. Under greedy decoding across eight models and six benchmarks, EMO improves the zero-shot baseline in 39 of 40 aggregated settings, averaging +9.10 percentage points. With the same judge, it adds +1.53 points on average over context-only optimization (TextGrad), with consistent MMLU-Pro gains across models. At matched target-model generation budgets, accuracy advantages over best-of-N are strongest on Llama, while oracle coverage is higher in most comparisons despite fewer candidates. The coverage--selection gap motivates allocating test-time computation jointly to candidate generation and verification.
Figures & tables
AIME (%)
GPQA-D (%)
GSM8K (%)
MATH-500 (%)
MMLU-Pro (%)
Model
Base
TG
EMO
Base
TG
EMO
Base
TG
EMO
Base
TG
EMO
Base
TG
EMO
Llama-3.2-3B
0.0
6.7
3.3
26.8
40.4
34.3
71.9
86.5
87.1
39.8
61.0
60.4
34.8
45.7
53.0
Qwen2.5-7B
5.0
11.7
16.7
32.8
42.9
45.5
87.8
90.7
91.7
70.0
78.0
79.2
54.8
61.6
65.5
Llama-3.1-8B
0.0
3.3
3.3
29.3
43.9
38.9
82.1
89.6
89.4
43.2
65.4
65.6
45.4
53.9
63.9
Qwen3-8B
20.0
43.3
46.7
49.0
54.5
56.6
91.4
92.7
92.3
75.0
82.0
83.0
63.6
64.8
67.5
DR1-Qwen3-8B
8.3
18.3
16.7
46.0
48.0
48.0
82.2
88.1
87.9
56.2
66.6
64.0
62.0
67.0
68.2
Table 1: Main results. Accuracy (%) across 8 models and 6 benchmarks (AIME-2025/2026 merged into one task group), greedy decoding. Each group reports the zero-shot baseline (Base), context-only optimization (TextGrad), and EMO; the best per group is bolded.
Table 2
Model
iter/ N
SC
BoN
EMO
Pass (oracle) SC/BoN
Pass (oracle) EMO
Qwen3-8B
0/1
65.5 ± 1.2
65.5 ± 1.2
63.9 ± 1.3
65.5 ± 1.2
63.9 ± 1.3
1/3
66.8 ± 1.2
67.1 ± 1.0
67.9 ± 1.1
75.2 ± 0.8
75.8 ± 1.5
2/5
67.6 ± 0.8
68.2 ± 0.7
69.4 ± 0.6
78.7 ± 0.3
80.3 ± 1.7
3/7
67.9 ± 0.3
68.2 ± 0.6
70.0 ± 0.8
80.6 ± 0.3
83.0 ± 2.2
4/9
67.1 ± 1.2
68.4 ± 0.5
70.9 ± 0.7
81.9 ± 0.5
84.6 ± 1.7
Llama-3.1-8B
0/1
45.1 ± 0.6
45.1 ± 0.6
45.1 ± 1.0
45.1 ± 0.6
45.1 ± 1.0
Table 4: Matched generation budgets on MMLU-Pro (accuracy %, mean ± std over 3 seeds, top- p sampling). iter/ N : EMO rounds K / target-model generation equivalents N=2K+1 . Pass : oracle coverage over K+1 candidates for EMO and N for SC/BoN.
Figure 2: Generation-budget scaling under top- p sampling (mean ± std over three seeds). At N=2K+1 , EMO has K+1 candidates; SC/BoN have N independent samples. Dashed curves show oracle coverage.
Appendix figures & tables17 assets
Supplementary material from the paper’s appendix.
Appendix
Benchmark
n
Task type
Answer form
Source
AIME-2025
30
math_boxed
integer
( 2025 )
AIME-2026
30
math_boxed
integer
( 2025 )
GPQA-Diamond
198
MC (A–D)
option letter
( 2023 )
GSM8K
1,319
math
#### number
( 2021 )
MATH-500
500
math_boxed
LaTeX / numeric
( 2023 )
MMLU-Pro (sampled)
560
MC (A–J)
option letter
( 2024 )
Appendix
Table 5: Benchmarks. Task format, size, and answer form. MC = multiple choice. MMLU-Pro is our stratified 560-question subset (14 categories × 40, sampling seed 42).
Model
HF checkpoint
Axis probed
Thinking
Llama-3.2-3B-Instruct
meta-llama/ Llama-3.2-3B-Instruct
scale (low)
✗
Llama-3.1-8B-Instruct
meta-llama/ Llama-3.1-8B-Instruct
family, non-thinking
✗
Qwen2.5-7B-Instruct
Qwen/ Qwen2.5-7B-Instruct
generation (prev), non-thinking
✗
Qwen3-8B
Qwen/Qwen3-8B
anchor
✓
Qwen3.5-9B
Qwen/Qwen3.5-9B
generation (newest)
✓
Qwen3-14B
Qwen/Qwen3-14B
scale (mid)
✓
Appendix
Table 6: Evaluated models. “Thinking” marks models with an explicit long chain-of-thought mode. The selection spans scale (3B–32B), three Qwen generations, two families, and thinking/non-thinking modes.
Experiment
Decoding
temp
top- p / top- k
max_new_tokens
Greedy main (8 models)
greedy
–
– / –
8,192; 12,288 (AIME); 16,384 (DR1-v2)
top- p main (Qwen3-8B)
sampling
0.6
0.95 / 0
24,576
top- p main (Llama-8B/3B)
sampling
0.6
0.95 / 0
12,288
top- k ablation
sampling
0.6
– / 20
as top- p main (seed 0 only)
BoN/SC pool
sampling
0.6
0.95 / 0
as top- p main, n=10 /question
Judge (Qwen3.5-122B-A10B FP8, shared by all experiments)
Appendix
Table 7: Consolidated experimental settings. Generation budget ( max_new_tokens ) is per response; the M-step runs a batch of two contexts per decoding step. The bottom block lists judge scoring and E-step decoding settings.
(a) AIME-2025
Model
Base
TextGrad
EMO
pass@5
Llama-3.2-3B
0.0
6.7
3.3
3.3
Llama-3.1-8B
0.0
3.3
3.3
3.3
Qwen2.5-7B
6.7
13.3
23.3
26.7
Qwen3-8B
16.7
36.7
43.3
60.0
DR1-Qwen3-8B
13.3
20.0
16.7
33.3
Appendix
Table 8: AIME-2025 and AIME-2026 reported separately (accuracy %, greedy; 30 problems per year, so 1 problem =3.3 points). pass@5 is the oracle upper bound over the 5 EMO rounds. The main-paper Table 1 reports the average of the two years.
Figure 3: GSM8K: accuracy vs. EMO iteration for all models (greedy). Solid: judge-selected accuracy; dashed: pass@ n (oracle); grey: base model.
Figure 4: MATH-500: accuracy vs. EMO iteration for all models (greedy).
Figure 5: MMLU-Pro: accuracy vs. EMO iteration for all models (greedy).
Figure 6: GPQA-Diamond: accuracy vs. EMO iteration for all models (greedy).
Figure 7: AIME-2025: accuracy vs. EMO iteration for all models (greedy).
Figure 8: AIME-2026: accuracy vs. EMO iteration for all models (greedy).
Dataset
Base
EMO
w/o M-step
w/o Judge
w/o Both
AIME-2025
10.0
43.3
43.3
30.0
26.7
AIME-2026
26.7
53.3
43.3
46.7
43.3
GPQA-D
59.6
62.1
64.6
61.6
64.1
GSM8K
91.4
92.0
91.9
91.7
91.7
MATH-500
78.6
86.0
83.8
83.2
82.6
MMLU-Pro
71.4
74.3
72.3
75.4
74.8
Appendix
Table 9: Component ablation on Qwen3-32B (accuracy %, greedy), mirroring Table 3 (Qwen3-8B). AIME rows are reported per year (30 problems each).
Model
Metric
iter ( N )
0 (1)
1 (3)
2 (5)
3 (7)
4 (9)
Qwen3-8B
acc GPQA-D
54.4 ± 1.4
55.4 ± 0.2
57.1 ± 1.1
58.4 ± 1.3
60.3 ± 2.1
pass
54.4 ± 1.4
64.0 ± 0.2
70.0 ± 1.7
75.6 ± 1.6
79.1 ± 0.6
acc MATH-500
79.6 ± 0.4
83.5 ± 0.6
84.3 ± 0.7
84.9 ± 0.6
85.0 ± 0.6
pass
79.6 ± 0.4
86.9 ± 0.6
88.9 ± 0.2
90.0 ± 0.2
90.4 ± 0.2
acc MMLU-Pro
63.9 ± 1.3
67.9 ± 1.1
69.4 ± 0.6
70.0 ± 0.8
70.9 ± 0.7
Appendix
Table 10: EMO vs. iteration under top- p sampling (accuracy %, mean ± std over 3 seeds). iter k uses N=2k+1 target-model generation equivalents; Pass is oracle coverage over the k+1 candidate responses.
Model
Task
SC@ N , N=1,…,10
1
2
3
4
5
6
7
8
9
10
Qwen3-8B
GPQA-D
52.4 ± 1.4
57.4 ± 1.0
56.4 ± 1.9
58.1 ± 0.8
57.9 ± 1.7
58.9 ± 2.3
58.4 ± 1.4
58.4 ± 3.2
57.9 ± 2.8
57.6 ± 1.1
MATH-500
79.2 ± 0.2
82.7 ± 0.7
84.2 ± 1.1
84.9 ± 0.7
85.1 ± 0.2
85.6 ± 0.0
85.8 ± 0.4
86.0 ± 0.4
86.1 ± 0.2
86.3 ± 0.5
MMLU-Pro
65.5 ± 1.2
66.4 ± 0.7
66.8 ± 1.2
67.2 ± 0.9
67.6 ± 0.8
67.6 ± 0.8
67.9 ± 0.3
67.7 ± 1.0
67.1 ± 1.2
67.7 ± 0.6
Llama-3.1-8B
GPQA-D
26.4 ± 0.9
27.8 ± 1.5
28.6 ± 1.7
27.8 ± 1.2
27.6 ± 2.1
29.0 ± 1.7
30.0 ± 1.3
28.8 ± 1.1
28.5 ± 0.9
28.8 ± 1.5
MATH-500
45.3 ± 1.9
51.7 ± 1.8
53.4 ± 1.9
54.8 ± 1.9
55.5 ± 1.3
55.7 ± 1.0
56.1 ± 0.4
56.0 ± 0.3
56.9 ± 0.2
57.5 ± 0.9
Appendix
Table 11: Self-consistency (SC@ N ) under top- p sampling (accuracy %, mean ± std over 3 seeds), majority vote over N independent samples.
Model
Task
BoN@ N , N=1,…,10
1
2
3
4
5
6
7
8
9
10
Qwen3-8B
GPQA-D
52.4 ± 1.4
57.4 ± 1.0
58.2 ± 0.6
59.9 ± 0.6
59.8 ± 0.6
59.9 ± 0.5
60.3 ± 0.2
60.6 ± 0.0
60.4 ± 0.9
59.8 ± 0.6
MATH-500
79.2 ± 0.2
82.7 ± 0.7
84.1 ± 0.7
84.6 ± 0.6
84.7 ± 0.5
84.9 ± 0.7
85.0 ± 0.6
85.0 ± 0.7
84.9 ± 0.6
85.1 ± 0.6
MMLU-Pro
65.5 ± 1.2
66.4 ± 0.7
67.1 ± 1.0
67.7 ± 0.4
68.2 ± 0.7
68.0 ± 0.5
68.2 ± 0.6
68.3 ± 0.9
68.4 ± 0.5
68.7 ± 0.4
Llama-3.1-8B
GPQA-D
26.4 ± 0.9
27.8 ± 1.5
28.6 ± 1.3
29.5 ± 1.0
29.5 ± 0.9
30.0 ± 0.2
30.8 ± 0.4
31.6 ± 0.9
32.2 ± 0.9
33.0 ± 0.5
MATH-500
45.3 ± 1.9
51.7 ± 1.8
54.5 ± 1.5
55.6 ± 1.3
56.9 ± 1.2
58.1 ± 0.9
59.1 ± 0.1
60.1 ± 0.2
60.8 ± 0.0
61.3 ± 0.2
Appendix
Table 12: Best-of- N (BoN@ N ) under top- p sampling (accuracy %, mean ± std over 3 seeds), judge-selected best of N independent samples.
Model
Task
N=3
N=5
N=9
SC/BoN
EMO
SC/BoN
EMO
SC/BoN
EMO
Qwen3-8B
GPQA-D
69.4 ± 0.6
64.0 ± 0.2
74.7 ± 1.1
70.0 ± 1.7
80.0 ± 0.2
79.1 ± 0.6
MATH-500
87.8 ± 0.0
86.9 ± 0.6
89.7 ± 0.2
88.9 ± 0.2
90.9 ± 0.2
90.4 ± 0.2
MMLU-Pro
75.2 ± 0.8
75.8 ± 1.5
78.7 ± 0.3
80.3 ± 1.7
81.9 ± 0.5
84.6 ± 1.7
Llama-3.1-8B
GPQA-D
45.1 ± 1.3
49.3 ± 2.7
51.2 ± 0.9
62.1 ± 1.2
60.3 ± 2.7
73.7 ± 0.4
MATH-500
60.7 ± 1.6
59.3 ± 1.7
65.9 ± 1.1
66.1 ± 1.7
71.9 ± 0.4
73.4 ± 2.0
Appendix
Table 13: Oracle coverage at matched generation budgets (accuracy %, mean ± std over 3 seeds): at least one candidate is correct. At budget N=2k+1 , SC/BoN use N independent samples and EMO uses k+1 candidate responses.
Model
Task
Comparison
Δ
n
pHolm
sig.
EMO vs. Base
Qwen3-8B
GPQA-D
EMO − Base
+5.9
594
2.2×10−3
**
Qwen3-8B
MATH-500
EMO − Base
+5.4
1500
1.0×10−16
Qwen3-8B
MMLU-Pro
EMO − Base
+7.0
1680
7.7×10−18
Llama-3.1-8B
GPQA-D
EMO − Base
+7.7
594
2.2×10−3
**
Llama-3.1-8B
MATH-500
EMO − Base
+19.7
1500
5.9×10−62
Appendix
Table 14: Paired significance tests under top- p sampling ( Δ in points; n = question–seed pairs pooled over seeds {0,1,2} ). p -values are Holm-corrected over 18 comparisons; stars mark ∗ p<0.05 , ∗∗ p<0.01 , ∗∗∗ p<0.001 (paired t -test). Wilcoxon significance decisions agree for all 18 comparisons.
(a) Total strength μ ( α=0.1 )
μ
EMO
Δ vs. base
1.6
56.6
+7.6
2.0
54.5
+5.6
2.4
56.6
+7.6
2.8
56.1
+7.1
3.2
54.5
+5.6
Appendix
Table 15: Hyperparameter sweep on Qwen3-8B / GPQA-Diamond (accuracy %, greedy; base =49.0 ). Panel (a) sweeps the total strength μ at α=0.1 ; panel (b) sweeps the confidence threshold α at μ=2.4 . The main-experiment default ( μ=2.4 , α=0.1 ) attains 56.6 ( +7.6 ).
Personalizing language models (LMs) to individual user preferences is essential for aligning responses with diverse goals and backgrounds. Existing methods typically train a separate adapter for each user or learn a reward model whose scores depend on the user. Despite explicitly optimizing for each user, these methods must learn from limited observations and therefore suffer from data sparsity and poor generalization to unseen users and domains. In-context learning (ICL) and Context Steering (CoS) can instead provide more effective personalization by conditioning the base LM directly on user context and leveraging its pretrained capabilities without per-user training. Yet neither adapts the influence of that context across decoding steps: ICL leaves it uncontrolled, whereas CoS applies a fixed steering coefficient and requires two LM forward passes per step. We propose Cautious Context Steering (CCS), which adds a lightweight adapter to a frozen backbone LM to decide at each token whether and how strongly user context should affect generation. The adapter learns this behavior from an oracle context-conditioned LM and preserves the base LM when the context is not helpful. A single CCS adapter trained on only one dataset improves generation quality both in-domain and across four out-of-distribution personalization benchmarks, demonstrating robust generalization to new users and domains. CCS also avoids per-user fine-tuning and the additional context-conditioned forward pass required by CoS, substantially reducing inference cost.
Gihoon Kim, Jeyoung Lee, Suhan Woo +4
Yonsei University · Hyundai Motors Company · Korea Institute of Science and Technology
Feedback signals used to train Large Language Models (LLMs) are the primary driver of their behavior and our main lever for instilling alignment with human values and objectives. However, a key limitation of current post-training methods is the inability of human annotators and automated reward functions to faithfully capture the feedback we would like to give. We introduce Evaluation-Conditioned Training (ECT), a post-training framework that uses natural language to condition each training sample on the fidelity of the feedback we provide and then elicits the desired behavior by conditioning the LLM on a high-fidelity monitor in deployment. ECT is aimed at improving performance under imperfect feedback and works as an add-on to existing algorithms such as SFT and PPO. We first provide a conceptual framework for ECT and discuss its potential to address persistent sources of reward mis-specification. Then we motivate ECT in the context of the eliciting latent knowledge (ELK) problem. Finally, we evaluate ECT on two proof-of-concept experiments: increasing even-handedness in news article generation and reducing sycophancy on an arithmetic task. In each setting, we utilize imperfect feedback, rewarding bias and agreement with the user, respectively. In both settings, ECT improves the targeted behavior relative to direct training.
Alec Harris, Kasey Corra, Archie Chaudhury +1
AI Safety Initiative at Georgia Tech · University of Chicago · Independent
Test-time training (TTT) lets a model store information in its weights during inference. When the model learns from its own output, however, each update also changes the model that generates the next training example. Across 128K-token streams, retaining generated-text updates worsens prediction on independent human-written text with three TTT-E2E model configurations (labeled 125M, 760M, and 3B). The same failure occurs when Adam updates Qwen3-4B's existing weights. The same update mechanisms can improve on real text, so writing itself is not the failure. Three matched comparisons trace the causal pathway. Fixed Generation removes over 98% of the damage at 125M and 760M by using a frozen model to generate training chunks. Recorded Replay separates the loss caused by reading degraded text from the additional loss stored by updating on it. A paired one-update comparison then shows the local conflict: an update predicts its source better but new real text worse. This cost grows after Closed Loop adaptation, with a few trajectories accounting for most large failures. Finally, Settlement evaluates the candidate state on independent real text before commitment. It leaves mean endpoint gaps of 0.07 and -0.02 nats at 125M and 760M while retaining real-text adaptation. These results motivate checking prediction on independent evidence before retaining an update.
Cheng Luo, Bing Li, Bernard Ghanem
King Abdullah University of Science and Technology (KAUST)