Contrastive Neuron Regulation for Chain-of-Thought Reasoning

Discover expansion-associated neurons. Regulate them sparsely. Keep the backbone frozen.

Highlights · Results · Method · Quick start · Reproduction · Citation

English | 简体中文

---

This repository is the official implementation of Contrastive Neuron Regulation (CNR), accepted to the EMNLP 2026 Main Conference. CNR is a sparse neuron-level framework for controlling how chain-of-thought (CoT) trajectories unfold: it builds correctness-controlled expanded/concise trajectory pairs, identifies MLP neurons associated with reasoning expansion, and suppresses only those neurons during generation.

Tip

Accepted to the EMNLP 2026 Main Conference.

Note

The release contains the core end-to-end pipeline: pair construction, activation collection, neuron scoring, fixed-coefficient intervention, adaptive coefficient learning, and evaluation.

Highlights

	What CNR provides
Controlled discovery	Compares expanded and concise correct trajectories for the same problem, reducing correctness and difficulty confounds.
Sparse intervention	Regulates fewer than 1% of MLP neurons in the reported experiments.
Frozen backbone	Never updates the language-model backbone; adaptive CNR learns only neuron-level suppression coefficients.
Two operating modes	Supports diagnostic fixed-alpha sweeps and learned neuron-specific coefficients.
Reproducible pipeline	Includes configurable launchers for all six stages, structured outputs, citation metadata, and automated static checks.

Results at a glance

Macro-average accuracy across the seven MATH subjects reported in the paper:

Model	BASE	CNR	Gain	Selected MLP neurons
DeepSeek-R1-Distill-Qwen-1.5B	65.6	68.2	+2.6	0.56%
Qwen2.5-7B	58.8	62.0	+3.2	0.42%

_{Results are macro averages, not uniform gains on every subject. See the paper for subject-wise results, controls, diagnostic sweeps, and limitations.}

How CNR works

_{Figure 2. CNR constructs correctness-controlled expanded/concise pairs, identifies expansion-associated neurons, and regulates them with fixed or learned suppression coefficients.}

Stage	Operation	Entry point	Main artifact
01	Construct same-question expanded/concise correct pairs	collect_paired_data.py	paired_data.jsonl
02	Collect reasoning-span MLP mean/max/std activations	collect_activations.py	activation arrays + metadata
03	Rank expansion-associated neurons with CAS	analyze_neurons.py	top_neurons_positive.json
04	Sweep fixed suppressive coefficients	intervene_and_eval_vllm.py	fixed-alpha evaluations
05	Learn neuron-specific coefficients with a frozen backbone	train_neuron_factors.py	best_alpha.json
06	Evaluate adaptive CNR on the test split	eval_learned_cnr.py	learned-CNR evaluations

The selected neurons are not assumed to be pure “length neurons” or complete reasoning mechanisms. They are neurons whose activations are associated with expansion under a correctness-controlled contrast; their functional role is tested through intervention.

Quick start

1. Install

The release targets Linux, Python 3.10+, CUDA-capable GPUs, and recent Hugging Face/vLLM environments.

git clone https://github.com/Blancacacher/Contrastive-Neuron-Regulation.git
cd Contrastive-Neuron-Regulation

python -m venv .venv
source .venv/bin/activate
python -m pip install --upgrade pip
pip install -r requirements.txt

Important

Full reproduction is GPU intensive. The paper reports experiments on NVIDIA A800 80GB GPUs. Other Hugging Face-compatible causal LMs may require changes to chat templates or activation hooks.

2. Configure paths

export MODEL_KEY=qwen25_7b
export MODEL_PATH=/path/to/Qwen2.5-7B
export DATA_ROOT=/path/to/MATH
export SUBJECT=algebra
export GPU_ID=0

3. Run the six-stage pipeline

bash scripts/collect_pairs.sh
bash scripts/collect_activations.sh

# The paper uses 80 neurons/layer for Qwen2.5-7B.
TOP_K_PER_LAYER=80 bash scripts/analyze_neurons.sh

bash scripts/eval_fixed_alpha.sh
bash scripts/train_learned_coefficients.sh
bash scripts/eval_learned_cnr.sh

The launchers run jobs in the background with nohup and print the PID and log path. Follow the corresponding log file to monitor progress.

Data and model setup

The code expects local Hugging Face-format model weights and MATH-style Parquet files. Supported models in the reported experiments are:

• DeepSeek-R1-Distill-Qwen-1.5B
• Qwen2.5-7B

Expected MATH directory layout

/path/to/MATH/
├── algebra/
│   ├── train-00000-of-00001.parquet
│   └── test-00000-of-00001.parquet
├── counting_and_probability/
├── geometry/
├── intermediate_algebra/
├── number_theory/
├── prealgebra/
└── precalculus/

Each example must contain at least:

{
  "problem": "...",
  "solution": "..."
}

The column names default to QUESTION_COL=problem and ANSWER_COL=solution.

Reproducing the paper configuration

Neuron-selection budget

Setting	DeepSeek-R1-Distill-Qwen-1.5B	Qwen2.5-7B
Transformer layers	28	28
MLP intermediate dimension	8,960	18,944
Selected neurons per layer	50	80
Selected-neuron ratio	0.56%	0.42%

Adaptive CNR

Hyperparameter	Value
Development ratio	0.1
Epochs	3
Training / evaluation batch size	2 / 4
Learning rate	1e-2
Regularization weight	1e-3
alpha_min, alpha_max, alpha_init	0, 1, 0.5
Maximum sequence length	8192
Precision	bfloat16
Random seed	42

The main-paper Contrastive Activation Score uses uniform component weights; those are the defaults in cnr/analyze_neurons.py.

Configuration reference

Common environment variables

Variable	Description	Default
MODEL_KEY	Output namespace	deepseek15b
MODEL_PATH	Local Hugging Face model path	/path/to/your/model
DATA_ROOT	Local MATH dataset path	/path/to/MATH
SUBJECT	MATH subject	algebra
GPU_ID	GPU index	0
OUTPUT_ROOT	Root for generated artifacts	./outputs/${MODEL_KEY}
USE_CHAT_TEMPLATE	Apply the tokenizer chat template	1
MAX_MODEL_LEN	vLLM maximum model length	8192
MAX_NEW_TOKENS	Maximum generation tokens	8192
BATCH_SIZE	Per-stage batch size	stage-dependent

Valid subjects are algebra, counting_and_probability, geometry, intermediate_algebra, number_theory, prealgebra, and precalculus.

Useful stage-specific overrides

# Pair construction
N_SAMPLES=8 TEMPERATURE=0.7 TOP_P=0.95 \
MIN_EXPANSION_RATIO=1.5 bash scripts/collect_pairs.sh

# Activation collection
DTYPE=bfloat16 MAX_SEQ_LEN=8192 BATCH_SIZE=4 \
bash scripts/collect_activations.sh

# Neuron analysis
TOP_K_PER_LAYER=80 ACT_STAT=fused TEST_NAME=ttest \
bash scripts/analyze_neurons.sh

# Fixed-coefficient sweep
ALPHA_VALUES=1.0,0.9,0.8,0.7,0.6,0.5,0.4,0.3,0.2,0.1,0.0 \
bash scripts/eval_fixed_alpha.sh

Output map

Artifact	Default path
Paired trajectories	outputs/{MODEL_KEY}/paired_data_{SUBJECT}_train/paired_data.jsonl
MLP activations	outputs/{MODEL_KEY}/activations_{SUBJECT}_train/activations/
Selected neurons	outputs/{MODEL_KEY}/analyze_{SUBJECT}_train/top_neurons_positive.json
Fixed-alpha evaluation	outputs/{MODEL_KEY}/eval_{SUBJECT}_test_fixed_alpha_scale/
Learned coefficients	outputs/{MODEL_KEY}/tune_factors_{SUBJECT}_perneuron/best_alpha.json
Adaptive-CNR evaluation	outputs/{MODEL_KEY}/eval_{SUBJECT}_test_learned_cnr/

Repository map

Show repository structure

Contrastive-Neuron-Regulation/
├── assets/
│   ├── cnr-hero-v2.png
│   ├── cnr-hero.svg
│   └── paper-method-overview.png
├── cnr/
│   ├── model_utils.py
│   ├── collect_paired_data.py
│   ├── collect_activations.py
│   ├── analyze_neurons.py
│   ├── intervene_and_eval_vllm.py
│   ├── train_neuron_factors.py
│   └── eval_learned_cnr.py
├── scripts/
│   ├── collect_pairs.sh
│   ├── collect_activations.sh
│   ├── analyze_neurons.sh
│   ├── eval_fixed_alpha.sh
│   ├── train_learned_coefficients.sh
│   └── eval_learned_cnr.sh
├── .github/workflows/quality.yml
├── CITATION.cff
├── CONTRIBUTING.md
├── LICENSE
├── README.md
├── README_zh.md
└── requirements.txt

Troubleshooting

MODEL_PATH or DATA_ROOT is not set

Pass the paths before the launcher:

MODEL_PATH=/path/to/model DATA_ROOT=/path/to/MATH \
bash scripts/collect_pairs.sh

The model does not use a chat template

USE_CHAT_TEMPLATE=0 bash scripts/collect_pairs.sh

CUDA out of memory

Reduce BATCH_SIZE, MAX_MODEL_LEN, MAX_NEW_TOKENS, or GPU_MEMORY_UTILIZATION:

BATCH_SIZE=32 MAX_NEW_TOKENS=4096 \
bash scripts/eval_fixed_alpha.sh

Paper

Contrastive Neuron Regulation for Chain-of-Thought Reasoning
Zitao Su, Xinyu Tang, and Xin Zhao
Gaoling School of Artificial Intelligence, Renmin University of China
EMNLP 2026 Main Conference

Citation

If you use this code, please cite the paper:

@misc{su2026contrastive,
  title  = {Contrastive Neuron Regulation for Chain-of-Thought Reasoning},
  author = {Su, Zitao and Tang, Xinyu and Zhao, Xin},
  year   = {2026},
  note   = {EMNLP 2026 Main Conference}
}

GitHub's Cite this repository menu can also export the metadata in CITATION.cff.

Contributing

Focused bug reports and pull requests are welcome. Please read CONTRIBUTING.md and avoid committing datasets, model weights, checkpoints, activations, or credentials.

License

This project is released under the MIT License.