DLoFT: Gradient-Decoupled Fine-Tuning for Generalizable Long Chain-of-Thought Reasoning

NeurIPS 2025

Sitong Wu, Haoru Tan, Jingyao Li, Shaofeng Zhang, Xiaojuan Qi, Bei Yu, Jiaya Jia

Introduction

DLoFT (Decoupled LongCoT Fine-Tuning) is a supervised fine-tuning algorithm that teaches an LLM the generalizable Long Chain-of-Thought (LongCoT) reasoning paradigm while preventing it from overfitting to problem-specific knowledge in the training data.

Standard SFT on high-quality LongCoT traces couples two learning signals in a single gradient: (1) the reasoning paradigm — how to structure thoughts, explore ideas, reflect, backtrack, and correct errors; and (2) the problem-specific content — the facts, formulas, and heuristics needed for a particular problem. Because internalizing the LongCoT paradigm takes many iterations, the model tends to memorize content along the way, which degrades out-of-distribution generalization.

DLoFT exploits the intrinsic two-phase structure of a LongCoT response — a thinking phase T followed by a solution phase S — to disentangle these signals at the gradient level. In each step it computes two gradients and decouples them via projection + orthogonalization:

Step 1: Full-response gradient over the complete response T ⊕ S:

$$ g_{\text{full}} = \nabla_\theta , \mathcal{L}_{\text{full}}, \qquad \mathcal{L}_{\text{full}} = -\frac{1}{B}\sum_{i=1}^{B}\sum_{k=1}^{|T_i \oplus S_i|} \log p_\theta\big((T_i \oplus S_i)_k \mid (T_i \oplus S_i)_{<k}, P_i\big) $$

Step 2: Reference gradient over the solution only (conditioned on the problem P, with the thinking phase removed):

$$ g_{\text{ref}} = \nabla_\theta , \mathcal{L}_{\text{ref}}, \qquad \mathcal{L}_{\text{ref}} = -\frac{1}{B}\sum_{i=1}^{B}\sum_{k=1}^{|S_i|} \log p_\theta\big(S_i^{k} \mid S_i^{<k}, P_i\big) $$

Since g_ref captures only problem-specific content, we project g_full onto it to isolate the content-relevant component, then orthogonalize to obtain the paradigm-relevant component:

$$ g_{\text{con}} = \frac{\langle g_{\text{full}}, g_{\text{ref}}\rangle}{\lVert g_{\text{ref}}\rVert^2}, g_{\text{ref}}, \qquad g_{\text{par}} = g_{\text{full}} - g_{\text{con}} $$

Step 3: Update with the paradigm-relevant gradient only:

$$ \theta \leftarrow \theta - \eta \cdot g_{\text{par}} $$

By dropping g_con, the model concentrates on the transferable reasoning paradigm rather than memorizing content, yielding large out-of-domain gains while preserving in-domain performance.

Highlights

• 🧠 Gradient decoupling for LongCoT: a single global projection/orthogonalization over the full gradient vector separates paradigm-relevant from content-relevant signals — exactly Algorithm 1 of the paper.
• 🔌 Drop-in on LLaMA-Factory: enable with one flag (use_dloft: true); reuses the standard SFT data/model/template stack.
• 🧩 Configurable thinking/solution delimiter: dloft_delimiter (default </think>) marks where the solution phase begins, so it works for both native reasoning models and short-CoT models that are learning to emit <think>...</think>.
• 🖥️ Distributed-aware: supports single-GPU, DDP (with gradient accumulation) and FSDP2 gradient decoupling out of the box.
• 🎯 Full fine-tuning & LoRA: the decoupling is applied over whatever parameters are trainable.

This repository implements DLoFT on top of LLaMA-Factory.

Installation

DLoFT follows the standard LLaMA-Factory setup.

# 1. Create the environment
conda create -n llamafactory python=3.10 -y
conda activate llamafactory

# 2. Enter the repository
cd DLoFT

# 3. Install LLaMA-Factory (editable) with training extras
pip install -e ".[torch,metrics]"

Training

1. Prepare the data

We use the mixed-domain s1K dataset (1,000 high-quality LongCoT traces from Gemini Flash Thinking). The provided script downloads it from the HuggingFace Hub and formats each sample as a LongCoT response — <think>\n{thinking}\n</think>\n\n{solution} — in the Alpaca format, then writes data/s1k_dloft.json (already registered in data/dataset_info.json as s1k_dloft):

conda activate llamafactory

# Download & build data/s1k_dloft.json
python scripts/prepare_s1k_dloft.py

# If the HuggingFace Hub is blocked, use a mirror:
HF_ENDPOINT=https://hf-mirror.com python scripts/prepare_s1k_dloft.py

# Options: choose the solution field / cap the number of samples
python scripts/prepare_s1k_dloft.py --solution_field attempt --max_samples 1000

Any LongCoT dataset works, as long as each response contains the dloft_delimiter (default </think>) that separates the thinking phase from the solution phase. Everything after the last delimiter is treated as the deterministic solution S used for the reference gradient.

2. Prepare the model

Any HuggingFace causal LM. The example configs teach a short-CoT model (Qwen2.5-7B-Instruct) to perform LongCoT reasoning. Edit model_name_or_path in the YAML to point to your local checkpoint or a Hub id.

3. Launch training

Enable DLoFT with a single flag in the config, then launch with llamafactory-cli (it starts torchrun automatically, so a single GPU, DDP, or FSDP2 all work):

conda activate llamafactory
llamafactory-cli train examples/train_lora/qwen2.5_lora_dloft.yaml

Citation

If you find this work useful, please cite:

@inproceedings{
wu2025dloft,
title={{DL}o{FT}: Gradient-Decoupled Fine-Tuning for Generalizable Long Chain-of-Thought Reasoning},
author={Sitong Wu and Haoru Tan and Jingyao Li and Shaofeng Zhang and Xiaojuan QI and Bei Yu and Jiaya Jia},
booktitle={The Thirty-ninth Annual Conference on Neural Information Processing Systems},
year={2025},
}

Acknowledgement

This codebase is built on top of the excellent LLaMA-Factory training framework, and uses the s1K dataset. We thank both communities for their open-source contributions.

Contact

For questions about the paper, please contact Sitong Wu at stonewst@163.com, or open an issue on GitHub.