Research map / Large language models / Reasoning and chain of thought

Flexible Realignment of Language Models

NeurIPS 2025 · Large language models › Reasoning and chain of thought · approach: Test-time compute scaling · 0 citations (OpenAlex)

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Proposes TrRa (training-time realignment via distillation toward a logit interpolation of reference and aligned models) and InRa (a zero-initialized duplicated bottom-layer adapter that enables inference-time logit interpolation), giving controllable alignment for reasoning length and 3H dialogue behavior.

Problem

Language models often need realignment after training (e.g., overthinking, sycophancy), but retraining under different reward strengths with RL is costly, and existing inference-time realignment (DeRa) doubles decoding cost and memory.

Approach

Distill the reference model toward a teacher distribution softmax[λ·aligned logits + (1−λ)·reference logits] (TrRa), and add a zero-initialized duplicated bottom layer (layer adapter) trained alone so that interpolating its logits with the original path gives controllable alignment inside one model at inference (InRa).

Key results

• TrRa-iter1 with λ=2.0 on DeepSeek-R1-Distill-Qwen-1.5B reduces tokens by 54.63% on AIME24 with Pass@1 of 29.17 vs 30.00 for the base model, while DeepScaleR-1.5B-Preview reduces tokens by 33.86%.
• On Qwen2.5-7B, the layer adapter (InRa λ=1.0) reaches AlpacaEval2 LC 25.74% and Arena-Hard 45.10%, comparable to full DPO (25.20% and 45.50%) while training 222M params in about 50 min versus 3h35min for full fine-tuning.
• InRa with λ<0 on DeepSeek-R1-Distill-Qwen-7B yields higher accuracy than the base model on AIME-24, AIME-25 and MATH-500 (e.g., 55.00 vs 51.67 on AIME24 at λ=-0.25 vs λ=0).

Measured effects reported in the paper

Surprising findings

• Extrapolating InRa to negative λ (more deliberate thinking) improved reasoning accuracy above the original reasoning model.
• Increasing the number of layer adapters (copying the bottom layer 2-3 times) did not significantly improve alignment under the same hyperparameters.

Limitations

• Inference-time InRa doubles the KV-cache, reducing throughput.
• Second-turn MT-Bench performance declines as λ increases, which the authors attribute to alignment tax from single-turn UltraFeedback data.
• Evaluation relies on LLM-as-judge (Qwen2.5-72B-Instruct) and limited benchmarks; training-time realignment used a reduced data subset due to compute constraints.

Experimental setup

Datasets: OpenR1-Math-220K UltraChat-200k UltraFeedback

Benchmarks: AIME24 AIME25 MATH-500 AlpacaEval 2 Arena-Hard v0.1 MT-Bench MMLU CMMLU TruthfulQA GSM8K

Metrics: Pass@1 token count length-controlled win rate win rate MT-Bench score

Baselines: DeepScaleR-1.5B-Preview DPO full fine-tuning DeRa LoRA SFT

Builds on: DeRa DPO LoRA knowledge distillation vLLM

Closest papers by meaning

• Aligning Large Language Models with Representation Editing: A Control Perspective (NeurIPS 2024 · similarity 0.88)
• Decoding-time Realignment of Language Models (ICML 2024 · similarity 0.88)
• InfAlign: Inference-aware language model alignment (ICML 2025 · similarity 0.88)
• Leveraging Importance Sampling to Detach Alignment Modules from Large Language Models (NeurIPS 2025 · similarity 0.88)
• Aligner: Efficient Alignment by Learning to Correct (NeurIPS 2024 · similarity 0.87)
• Test-Time Alignment for Large Language Models via Textual Model Predictive Control (ICLR 2026 · similarity 0.87)