Amplify Adjacent Token Differences: Enhancing Long Chain-of-Thought Reasoning with Shift-FFN

TL;DR

This work identifies Cyclical Reasoning as a common failure mode when fine-tuning for long Chain-of-Thought reasoning, correlating it with small differences between adjacent token representations. It introduces Shift-FFN, which inserts an Editor before the FFN to edit the current token using the previous token, thereby amplifying representation differences and reducing cycles. Empirical results show that LoRA+Shift-FFN improves accuracy and lowers Cyclical Reasoning across data sizes and model families, outperforming full fine-tuning and standard LoRA under comparable parameter budgets. The approach provides a practical, architecture-level fix that complements long CoT distillation and PEFT strategies, with code and data available for reproduction.

Abstract

Recently, models such as OpenAI-o1 and DeepSeek-R1 have demonstrated remarkable performance on complex reasoning tasks through Long Chain-of-Thought (Long-CoT) reasoning. Although distilling this capability into student models significantly enhances their performance, this paper finds that fine-tuning LLMs with full parameters or LoRA with a low rank on long CoT data often leads to Cyclical Reasoning, where models repeatedly reiterate previous inference steps until the maximum length limit. Further analysis reveals that smaller differences in representations between adjacent tokens correlates with a higher tendency toward Cyclical Reasoning. To mitigate this issue, this paper proposes Shift Feedforward Networks (Shift-FFN), a novel approach that edits the current token's representation with the previous one before inputting it to FFN. This architecture dynamically amplifies the representation differences between adjacent tokens. Extensive experiments on multiple mathematical reasoning tasks demonstrate that LoRA combined with Shift-FFN achieves higher accuracy and a lower rate of Cyclical Reasoning across various data sizes compared to full fine-tuning and standard LoRA. Our data and code are available at https://anonymous.4open.science/r/Shift-FFN

Figures (10)

• Figure 1: (Left), performance comparison of LoRA and Full Fine-Tuning (Full FT) on Accuracy (%) and Length Exceed (%) metrics for short CoT and long CoT datasets. "Accuracy" represents the average accuracy across four mathematical tasks. "Length Exceed" indicates the percentage of model outputs that exceed the maximum length limit. (Right), an example of Cyclical Reasoning.
• Figure 2: Distribution of the $M(\boldsymbol{X})$ for Exceed and Normal samples, comparing LoRA and Full fine-tuned models. The dashed red line represents the mean value.
• Figure 3: The architecture of Shift-FFN, the left side describes the process of shifting token, the right side demonstrates the detail of the Editor module. $\sigma$ is the ReLU function. $\odot$ and $\oplus$ are element-wise multiplication and addition, respectively.
• Figure 4: Proportion of length-exceeded and repetition samples in different models.
• Figure 5: Pass@K of models with different training setups on AIME24 and OlympiadBench.