Fixing Reasoning Loops in Small Language Models by Editing Just 0.016% of the Weights

1. Why Reasoning Loops Overwhelm Small Models

In one mathematical reasoning evaluation of a small language model, 85 out of 96 test samples failed to complete because they reached the output limit of 32,768 tokens. The model was not simply reasoning slowly. It had entered a reasoning loop, repeatedly generating highly similar passages.

For example, the model might continue asking itself, “Is 20 the total? No, it is the remaining time,” without advancing beyond the same step. This repetition continues until the available context window is exhausted.

The behavior can be understood at three levels:

1. Error accumulation during autoregressive decoding: Each generated token is conditioned on the model’s previous output. Once a passage begins to repeat, that repeated context increases the probability of producing more of the same, creating a positive feedback loop.
2. Mismatch between model capacity and task complexity: Small models often struggle to maintain a global solution plan. When uncertain about the current state, they may assign higher probability to restatement or self-questioning than to meaningful progress.
3. Amplification by the decoding strategy: Low-temperature or greedy decoding further narrows the exploration space, making the model more likely to become locked into a locally repetitive pattern.

The consequences are substantial: reasoning latency and token costs increase, many samples never produce an answer, and repetitive text consumes context that should have been available for solving the problem.

Common black-box mitigations—such as rewriting prompts, increasing repetition penalties, or adding more fine-tuning data—often address the symptoms rather than the internal cause. The approach described here instead examines the model’s internal representations, locates the source of the loop, and applies a minimal intervention.

2. Turning a Behavioral Failure into a Representation Problem

The central hypothesis is that a reasoning loop is not random noise. Instead, it results from the repeated amplification of a specific direction in the model’s internal representation space. If such a direction exists, it should be observable, localizable, and open to intervention.

The first step is to construct a controlled comparison. For the same math problem, researchers prepare two reasoning traces:

• Productive reasoning: The model computes a step such as 10 + 8 + 20 = 38 and continues toward the answer.
• Repetitive reasoning: The model becomes stuck on one step and repeatedly produces similar questions and answers.

Both traces use the same problem and the same model. The only intended difference is whether the reasoning enters a loop. This controls for confounding factors such as problem difficulty and overall model capability.

Paired samples can be created by taking a real looping trajectory and attaching a productive continuation, or by manually rewriting one of the traces. The experiment described here used four such pairs.

The next step is to collect internal states. With teacher forcing, the productive and repetitive traces are each passed through the original model as complete sequences. A single forward pass is sufficient to retrieve the hidden states, or residual stream representations, and MLP activations at every layer and token position.

Compared with running full token-by-token decoding, teacher forcing is inexpensive and ensures that activations under the two conditions can be compared directly.

3. Extracting the Repetition Direction

For residual stream layer (l), take the mean hidden state over the last 36 tokens of each sample. Let the resulting representations be (\bar{h}{normal}^{(l)}) and (\bar{h}{loop}^{(l)}). The contrastive direction is defined as the average difference across the four sample pairs:

d^(l) = (1/4) · Σ(i=1..4) ( h̄_loop,i^(l) − h̄_normal,i^(l) )

The direction is then normalized:

u^(l) = d^(l) / ‖d^(l)‖₂

Three design choices are important.

Why use the last 36 tokens?

The most useful evidence for detecting a loop is concentrated in what happened recently. Earlier positions mostly represent the original question and already established reasoning steps, which are less sensitive to whether the model later becomes repetitive.

A window that is too long dilutes the repetition signal, while an extremely short window produces higher statistical variance. In this experiment, the final 36 tokens provide the comparison window.

Why average across four sample pairs?

A difference computed from a single pair contains both a repetition-related signal and problem-specific noise. Averaging across several problems helps cancel out task-specific components and preserve the direction shared by repetitive traces.

This design assumes that reasoning loops across different problems share an approximately common representation direction. The experimental results suggest that this assumption can hold for small models.

Why normalize the direction?

Activation norms can vary substantially between layers. Normalization makes the relative strength of the direction easier to compare across layers and supports subsequent projection-based analysis.

A simplified implementation looks like this:

# Collect activations with one teacher-forced forward pass per condition
out_n = model.forward(texts_productive, output_hidden_states=True)
out_l = model.forward(texts_repetitive, output_hidden_states=True)
d = []
for l in range(num_layers):
    diff = (out_l.hidden[l][:, -36:, :].mean(1)     # (4, d_model)
          - out_n.hidden[l][:, -36:, :].mean(1))
    d.append(diff.mean(dim=0))                      # Average across four pairs
u = [d[l] / d[l].norm() for l in range(num_layers)] # Unit direction per layer

4. Locating the Layers That Amplify Repetition

After obtaining (d^{(l)}) for every layer, examine the change between adjacent layers:

Δd^(l) = d^(l+1) − d^(l)

If (\lVert \Delta d^{(l)} \rVert) increases sharply after a particular layer, the difference between productive and repetitive activations is being amplified there. This suggests that the layer is actively pushing the model toward repetitive behavior rather than merely carrying a signal introduced elsewhere.

In practice, the first step is usually to plot (\lVert \Delta d^{(l)} \rVert) across all layers. A small number of prominent spikes can then be selected as candidates for further analysis, avoiding an exhaustive layer-by-layer investigation.

5. Attributing the Effect to Attention or the MLP

Each Transformer layer writes two main components into the residual stream through residual connections: the attention output and the MLP output. The change in the contrastive signal can therefore be decomposed approximately as:

Δd^(l) ≈ Δ_attention^(l) + Δ_MLP^(l)

The MLP contribution can be calculated directly:

Δ_MLP^(l) = W^(l) · ( ā_loop^(l) − ā_normal^(l) )

Here, (\bar{a}) is the mean MLP activation over the final 36 tokens, measured after the activation function and before the down-projection. (W^{(l)}) is the MLP down-projection matrix for layer (l).

The attention contribution can then be estimated by subtracting the MLP contribution from the total layer-to-layer difference.

In the reported case, most of the growth in the repetition signal came from MLP channels in a small number of layers. Under looping conditions, these MLPs repeatedly wrote components aligned with the repetition direction into the residual stream. Functionally, this acted like an internal instruction telling the model to continue repeating itself.

6. How Editing 0.016% of the Weights Can Work

Once the relevant layers and sublayers have been identified, there are two common intervention paths.

Path A: Direction Ablation at Inference Time

Without changing the model’s weights, the system can remove the component aligned with (u) from the target layer’s residual stream or MLP output:

def ablate(x, u, alpha=1.0):
    coeff = (x @ u).unsqueeze(-1)   # Projection magnitude along the repetition direction
    return x - alpha * coeff * u

This approach requires no training and can be rolled back immediately. However, it introduces a small amount of additional computation during inference and only affects services on which the ablation logic has been installed.

Path B: Low-Rank Weight Editing

The intervention can instead be converted into a rank-one, or otherwise low-rank, weight update that reduces the model’s tendency to write along direction (u):

W' = W − α · u vᵀ

Only a small portion of the parameters in a few selected layers must be changed. In the reported case, the edited parameters represented approximately 0.016% of the full model.

Such a small change can be effective because the looping behavior depends heavily on a few concentrated representation directions rather than being uniformly distributed throughout parameter space. Precisely suppressing those directions can reduce repetition while leaving most other model capabilities untouched.

After the edit, the small model’s reported math benchmark score increased to several times its original level, with a reported improvement of sixfold. The model did not necessarily become fundamentally better at mathematics. Much of the gain came from allowing samples that previously exhausted their token budgets to complete their reasoning chains. In effect, the intervention returned the context budget to the problem-solving process.

7. Validating Effectiveness and Avoiding Regressions

The model should be evaluated before and after the intervention across four dimensions:

1. Loop rate: The proportion of new samples that reach the token limit without terminating normally should decrease substantially.
2. Task accuracy: Improvements should be measured on a held-out mathematical reasoning test set.
3. Token efficiency: The average number of generated tokens should return to a healthy range.
4. General capability regression: Non-mathematical tasks such as summarization, translation, and dialogue should be tested to detect catastrophic forgetting or unintended style changes.

Ablation studies are also important. Useful comparisons include:

• Editing only attention components
• Editing only MLP components
• Modifying one layer
• Modifying multiple layers

These experiments help confirm that the gains come from the components identified by the representation analysis rather than from an unrelated effect or chance.

8. Engineering Considerations and Limitations

White-box access is required

This method depends on reading internal activations and, for weight editing, modifying model parameters. It is therefore applicable to internally developed models or models with accessible weights, but not to closed-source API-only systems.

The direction may not generalize universally

A repetition direction may depend on the task type, model architecture, or model scale. Before applying the method to a different task or model, new contrastive samples should be collected and the consistency of the extracted direction should be validated.

Excessive suppression can remove useful behavior

In mathematical reasoning, useful restatement—such as rereading a condition from the problem—differs from a pathological loop mainly in degree. The coefficient (\alpha) should therefore be tuned in small increments to avoid suppressing legitimate reasoning behavior.

Production monitoring remains necessary

Whether inference runs on an internal cluster or through a managed cloud service, including model inference products offered by providers such as Baidu AI Cloud, loop rate and token-limit truncation rate should be included in production monitoring. These metrics can reveal regressions after model, prompt, decoding, or infrastructure changes.

The method also remains a targeted treatment rather than a complete cure. It addresses a systematic failure mode in a particular capability of a small model, but it cannot replace larger-scale pretraining or higher-quality post-training. Whether the same directional assumption holds for larger models must be evaluated case by case.

9. Conclusion

The main value of this approach is not limited to the headline result of achieving a severalfold improvement by editing approximately 0.016% of the model’s weights. It also demonstrates a broader model-debugging workflow:

1. Find a geometric representation of the failure in activation space.
2. Trace that representation through the model to identify the responsible layers and sublayers.
3. Apply a minimal, targeted intervention.
4. Validate both the intended benefit and possible regressions.

This workflow moves model debugging away from intuition-driven trial and error and toward an engineering discipline based on observability, attribution, and controlled intervention. For teams deploying small language models under limited compute and context budgets, it provides a practical technique worth adding to the model optimization toolkit.