Research map / Large language models / Instruction and fine-tuning
Learning Distribution-wise Control in Representation Space for Language Models
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.
Replaces deterministic ReFT-style intervention layers with stochastic reparameterized (mean and variance) layers, learning distributions over representation interventions, mainly in early layers, and evaluates on commonsense and arithmetic reasoning with Llama models.
Problem
Pointwise learnable interventions (representation fine-tuning) control LM behavior only at single points in representation space; the paper asks whether learning distributions over interventions (the neighborhood of the concept subspace) improves controllability and robustness.
Approach
Replace the deterministic intervention network with a stochastic one: separate MLPs learn mean and log-variance, and outputs are sampled via the reparameterization trick (no KL term), applied as a drop-in to ReFT, RED, MLP and SwiGLU, with model-specific clamping; best when applied to the first 25% of layers.
Key results
- D-ReFT(25%) on LLaMA-3-8B reaches 89.1 average on eight commonsense benchmarks vs 86.6 for ReFT.
- On arithmetic reasoning with Llama-3-8B, D-ReFT(25%) averages 77.4 vs 76.6 for ReFT, with gains on GSM8K and SVAMP.
- When deleting up to ~8 words from arithmetic test inputs, D-ReFT accuracy stays stable while pointwise ReFT drops about 30%.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Replace the first 25% of ReFT layers with D-ReFT | ReFT · Commonsense reasoning (8 benchmarks), LLaMA-7B | accuracy (Avg.): 80.2 -> 82.2 | Helped |
| Replace the first 25% of ReFT layers with D-ReFT | ReFT · Commonsense reasoning (8 benchmarks), LLaMA-13B | accuracy (Avg.): 83.3 -> 85.1 | Helped |
| Replace the first 25% of ReFT layers with D-ReFT | ReFT · Commonsense reasoning (8 benchmarks), LLaMA-2-7B | accuracy (Avg.): 81.8 -> 83.6 | Helped |
| Replace the first 25% of ReFT layers with D-ReFT | ReFT · Commonsense reasoning (8 benchmarks), Llama-3-8B | accuracy (Avg.): 86.6 -> 89.1 | Helped |
| Replace the first 25% of ReFT layers with D-ReFT | ReFT · Arithmetic reasoning (7 benchmarks), LLaMA-7B | accuracy (Avg.): 59.9 -> 61.2 | Helped |
| Replace the first 25% of ReFT layers with D-ReFT | ReFT · Arithmetic reasoning (7 benchmarks), LLaMA-2-7B | accuracy (Avg.): 62.3 -> 63.7 | Helped |
| Replace the first 25% of ReFT layers with D-ReFT | ReFT · Arithmetic reasoning (7 benchmarks), LLaMA-3-8B | accuracy (Avg.): 76.6 -> 77.4 | Helped |
| Increase the fraction of layers using D-ReFT beyond 25% (ablation) | D-ReFT mixed-layer strategy · Commonsense reasoning (8 benchmarks) | accuracy: 82.2 -> 81.0; 82.2 -> 79.4; 82.2 -> 79.3 (LLaMA-7B) | Hurt |
| Use deterministic inference (τ=0) instead of baseline sampling (ablation) | D-ReFT · Commonsense reasoning | performance: 88.6 -> 89.3 | Helped |
| Use deterministic inference (τ=0) instead of baseline sampling (ablation) | D-ReFT · Arithmetic reasoning | performance: 76.6 -> 79.2 | Helped |
| Use deterministic inference (τ=0) instead of baseline sampling (ablation) | D-ReFT · Instruction following | performance: 82.4 -> 80.1 | Hurt |
| Increase inference temperature to τ=2 (ablation) | D-ReFT · Instruction following | performance: 82.4 -> 85.9 | Helped |
Surprising findings
- Applying distribution-wise intervention to all layers (D-ReFT 100%) underperforms the mixed 25% configuration and makes training hard to converge at high stochastic fractions.
- Tuning the rank to 8-16 gives the best results, and higher ranks do not improve ReFT or D-ReFT, suggesting parameter count alone does not explain gains.
Limitations
- Gains on arithmetic reasoning are small (about +0.8 average over ReFT for Llama-3-8B) and some baselines win on individual tasks.
- Stochastic intervention at high layer fractions causes training instability and numerical issues, requiring model-specific clamping.
- Hyperparameter ablations are not included in the main text and results are only on Llama-family models.
Experimental setup
Datasets: Commonsense170K Math10K UltraFeedback GSM8K development split
Benchmarks: BoolQ PIQA SIQA HellaSwag WinoGrande ARC-e ARC-c OBQA AddSub SingleEQ MultiArith AQuA
Metrics: accuracy win rate
Baselines: ReFT RED LoRA DoRA Prefix-tuning Adapter ChatGPT SwiGLU MLP
Builds on: ReFT variational autoencoder LoFiT RED LoRA
Closest papers by meaning
- An Emulator for Fine-tuning Large Language Models using Small Language Models (ICLR 2024 · similarity 0.86)
- Model Fusion through Bayesian Optimization in Language Model Fine-Tuning (NeurIPS 2024 · similarity 0.86)
- Beyond Linear Activation Steering: Invertible Latent Transformations for Controlling LLM Behavior (NeurIPS 2026 · similarity 0.86)
- Improved Representation Steering for Language Models (NeurIPS 2025 · similarity 0.86)
- Improving Reasoning Performance in Large Language Models via Representation Engineering (ICLR 2025 · similarity 0.86)
- Focus On This, Not That! Steering LLMs with Adaptive Feature Specification (ICML 2025 · similarity 0.86)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.