Qwen 3 4B Base Fine‑Tuned on 100 Zebra Puzzles Boosts MATH‑500 by 31%

Resources & Quick Links:

• Read this Post: HuggingFace (Rendered Blog) | GitHub (Markdown Source)
• GitHub Repository: tamewild/pcss
• Reproduction Notebooks (Single H100/H200): Qwen 3 4B (6.5 mins) | Granite 4.1 3B (23 mins) | Qwen 3.5 9B (~40 mins)
• Models: PCSS-Qwen3-4B | PCSS-Granite-4.1-3B | PCSS-Qwen3.5-9B
• Datasets: zebra_100 | instruct5

Encountering the research on the Hyperfitting Phenomenon sparked an interest in exploring how pre-trained base models behave when micro-datasets are fine-tuned to near-zero training loss. We focused on training small pre-trained base models on 100 to 500 pure logic deduction traces (specifically 5x5 zebra puzzles) with zero mathematical data in the fine-tuning dataset. This synthetic dataset was generated by Qwen3-235B-A22B-Instruct-2507 in September 2025 using rejection sampling.

Across multiple model families and architectures:

• Qwen 3 4B Base fine-tuned on100 zebra puzzles (~6.5 minutes on 1 GPU) reaches85.26% on the full 5,000-problem MATH benchmark (+31.16% delta over base) and21.67% on AIME 2025 , rivaling the official post-trained model in non-thinking mode on MATH-500 while outperforming it on AIME 2025.
• Granite 4.1 3B Base fine-tuned on500 zebra puzzles achieves77.73% on MATH-500 and19.44% on AIME 2025 , outperforming the official IBM Granite 4.1 3B Instruct model (+11.13 and +12.77 percentage points, respectively).
• Qwen 3.5 9B Base fine-tuned on500 zebra puzzles jumps from 17.00% to80.50% on 5x5 zebra puzzles (Pass@3:97.05% ), while jumping from 5.42% to22.92% on the uncontaminated ArXivMath 05/26 benchmark (Pass@3:34.88% ) and reaching96.60% on MATH-500 and60.67% on AIME 2025 .

While this out-of-domain generalization appears primarily driven by the pure logic dataset, we chose to manage the optimization process through adaptive loss scaling. For our primary results, we utilized PCSS (Per-Example Calibrated Sigmoid Scaler): an adaptive, data-aware loss scaler derived directly from KTO.

PCSS originates from KTO. But why adapt an alignment algorithm for Supervised Fine-Tuning?

The inspiration came from an experiment in the original KTO paper (Table 4). The authors demonstrated that running KTO using only desirable/positive examples (the exact same unpaired data used for SFT) yielded results highly competitive with standard SFT. While the paper conservatively notes this setup "yields similar results to SFT when is small," the results with raised the overall average benchmark score from 29.7% to 31.1%, driven by a substantial relative increase on the GSM8K mathematical reasoning benchmark (from 1.0% to 12.5%).

Why does an alignment algorithm behave this way? We can answer this by performing a gradient analysis on standard KTO when applied to purely desirable data.

Setting the gain sensitivity to 1 , the KTO loss for a desirable ground-truth sequence is: where is the unnormalized log ratio . The loss for the ground-truth sequence is evaluated at , and is the KL divergence estimate. Crucially, KTO explicitly does not backpropagate through .

Taking the gradient with respect to : Because the reference model is frozen, . Since the standard SFT loss is the negative mean log-probability over tokens, . Therefore, .

Substituting this back in gives the exact gradient update for standard KTO on positive data:

When you apply KTO to purely desirable examples, it is no longer performing preference alignment against negative data. Instead, it mathematically reduces to the standard SFT gradient multiplied by a dynamic scaler.

However, this same equation highlights why standard KTO would struggle on logic chains spanning up to 11,000+ tokens:

1. Length-Induced Gradient Inflation: The scaler is multiplied by the sequence length . For long traces, the gradient magnitude increases sharply by a factor of hundreds.
2. Length Bias & Gradient Saturation: Because is an unnormalized sum, its magnitude also scales linearly with , quickly saturating the sigmoid. The derivative drops toward zero, which would cause the optimizer to prematurely stop learning on long sequences.

To convert this mechanism into the stable PCSS scaler, we make two fundamental mathematical transformations:

PCSS explicitly normalizes the implied reward by sequence length using standard mean cross-entropy losses (where ):

• Current model SFT loss :
• Frozen base model reference loss :

The normalized reward becomes the difference: .

The canonical KTO value function includes as a baseline expectation .

While KTO's batch-dependent estimate breaks down at , mathematical analysis indicates that if we want the scaler to natively decay upon mastery in SFT, dropping the dynamic KL anchor is a straightforward architectural adjustment, making a natural design choice.

During fine-tuning on a ground-truth logic trace , the policy's probability mass concentrates entirely on the correct answer . If we plug this mastered state back into the formal sequence-normalized definition of :

Upon convergence , the reward is . If we kept the baseline anchor , the advantage would evaluate to:

Because the derivative of the sigmoid reaches its peak value at an input of exactly , keeping the baseline anchor would cause the scaler to trend back toward its maximum value at the end of training, failing to suppress residual gradients and increasing the risk of late-stage overfitting.

By explicitly setting , the advantage becomes an absolute measure of progress against the base model . Upon mastery , the advantage becomes strictly positive . The scaler structurally decays toward zero as mastery increases, naturally suppressing residual gradients to help mitigate late-stage overfitting.

The PCSS loss for a chosen example becomes:

Because this gradient mathematically reduces to the standard SFT gradient scaled by a dynamic multiplier, we can implement it directly by scaling the standard SFT loss output by the detached multiplier:

To standardize the dynamic multiplier across model families and decouple its magnitude from its shape, PCSS applies two scaling adjustments:

1. Scaled Advantage : , where is the pre-computed standard deviation of reference losses across the training split.
2. Decoupled Peak Scaling :

Since , peak_scale directly defines the maximum multiplier applied to the SFT gradient.

Applying these scaling adjustments, the programmatic PCSS loss becomes:

While PCSS is implemented by scaling the standard SFT loss with a detached multiplier, this mechanism is not an arbitrary gradient heuristic. Since the gradient of our implementation with respect to is , we can find the antiderivative of with respect to to recover the exact integrated objective being optimized (choosing the constant of integration to match the form of the original KTO objective):

By anchoring the constant of integration to KTO's baseline, this explicit form demonstrates that PCSS preserves the exact functional shape of the original objective (the core ). More importantly, it reveals how standard deviation normalization and decoupled peak scaling scale the overall magnitude of the loss surface by a factor of .

To accommodate training with an effective batch size of 1, we apply the constant token half-life scaling heuristic from recent small-batch training research to mitigate the rapid decay of AdamW's second moment estimate. Scaling a baseline from batch size 16 down to batch size 1 yields:

Across all PCSS experiments, we fix:

• Effective Batch Size : 1
• Optimizer : 8-bit AdamW
• Epochs : 5
• Scheduler : Constant learning rate with 1-epoch linear warmup
• Weight Decay :0.01
• Parameter :0.99994
• PCSS Parameters : ,
• MiSS Parameter :

For modern foundation architectures that incorporate native Multi-Token Prediction (such as dense Qwen 3.5 models, which iteratively unroll an MTP decoder layer across speculative steps), PCSS can be extended to jointly govern both the primary next-token prediction and subsequent look-ahead steps.

For the Qwen 3.5 9B experiments in this post, we utilized a preliminary formulation with a single speculative step , computing independent scalers for each prediction step:

While treating the single speculative step independently performed well for Qwen 3.5 9B, it remains an initial exploration. Scaling to deeper speculative horizons introduces optimization challenges for independent scalers.

We plan to address these limitations and propose a more generalized approach in a future post.

All training runs were conducted on vast.ai instances with a single H200 NVL GPU. Unless otherwise noted, models were evaluated on their zero-shot Chain-of-Thought (CoT) capabilities with 16,384 max completion tokens via a pinned version of vLLM (0.19.1).

Evaluation Notes:

Exploratory Runs: Results represent exploratory single-seed training runs; score differences between configurations should not be taken as definitive conclusions. Multiple Generations: To account for natural sampling stochasticity as well as vLLM's runtime batching non-determinism (which affects even greedy decoding), most benchmarks were evaluated by drawing multiple generations per problem (e.g., on GPQA Diamond, and on AIME). Single generations were used only when the evaluation set was large enough to provide natural stability (e.g., full MATH, MMLU Redux, and MATH-500 for certain ablations). Pass@3: Reported only for benchmarks with at least generations per question.

Our fine-tuned models were evaluated using greedy decoding with 16,384 max completion tokens. Scores for the official post-trained and base models are self-reported figures taken directly from the Qwen 3 Technical Report.

The configurations evaluated below share the following default hyperparameters:

• Shared Settings: 5 epochs, 1-epoch linear warmup, 8-bit AdamW optimizer.
• PEFT Defaults: LoRA uses , , learning rate . MiSS uses .
• Optimizer : for , and for .

(Note: Training times were recorded on individual vast.ai H200 instances across TRL and custom trainer setups; minor variations within 6–8 minutes reflect framework callbacks and host instance differences. For , a micro-batch size of 16 required gradient checkpointing to fit in VRAM, whereas reaching the same effective batch size via gradient accumulation would roughly equalize speed).

Ablation Takeaway: Across all tested configurations—including standard batch size 16 with cosine decay—MATH-500 performance consistently approached or exceeded 80%, suggesting that reasoning generalization from pure logic traces is robust across different training setups rather than dependent on a single hyperparameter configuration.

Evaluated with greedy decoding:

Granite Delimiters & Template Ablation: During exploratory evaluations of Granite 4.1 3B Base, we observed that performance appeared sensitive to the string <extra_id_1>. We compared two fine-tuning configurations: one using the generic chat template from our Qwen runs, and another adapted to use <extra_id_1>User and <extra_id_1>Assistant. Despite <extra_id_1> being processed as standard text tokens rather than a special control token, the run formatted with these delimiters demonstrated slightly better loss convergence and higher downstream accuracy. This is the primary checkpoint reported in the table above.

Chat Template Pre-training Note: Qwen 3.5 9B Base was released with its chat template control tokens (e.g., <|im_start|>, <|im_end|>) explicitly pre-trained "to allow efficient LoRA-style PEFT with the official chat template." Applying this template allowed us to evaluate the untrained base, our fine-tuned checkpoint, and the official post-trained model directly in non-thinking mode for strict comparability.

Evaluated with sampled decoding ( , with presence_penalty=1.5 for the official post-trained model, as recommended by Qwen):

To better understand the structural generation dynamics accompanying these performance gains—specifically regarding inference length and repetition failure modes—we measured completion lengths and looping rates.

Token counts were calculated using the Qwen 3.5 tokenizer on MATH-500 generation outputs:

Despite training exclusively on long-form logic traces (with many spanning up to 11,000+ tokens), the PCSS fine-tuned Qwen 3.5 9B model exhibits a shorter median completion length on MATH-500 than both its untrained base model and the official post-trained model.

To evaluate repetition failure modes without stochastic sampling variance, models were evaluated under greedy decoding on GPQA Diamond. Repetition loops were detected using Liquid AI's antidoom inner-repetition detection heuristic:

The structural shifts and repetition reductions observed here should not be viewed as universal features of this training process. Exploratory runs indicate that different base model families internalize and generalize from these logic traces in substantially different ways—some models contract their output, while others expand it. Furthermore, although longer sequences inherently provide more opportunities for repetition, a model's underlying susceptibility to looping can shift in ways that are not strictly tied to generation length. Because this analysis is not exhaustive, even the specific dynamics of Qwen 3.5 9B may vary across different reasoning tasks. While these initial results suggest its performance leap does not strictly rely on pure inference scaling and it reduces repetitive loops on GPQA Diamond, generation behaviors remain complex and highly model-dependent.

Our exploratory runs highlight two distinct takeaways for post-training base models:

1. Impact of Pure Logic Micro-Datasets: In these exploratory runs, the primary driver of performance appears to be the data itself. Fine-tuning on as few as 100 pure logic deduction traces (zebra puzzles) can elicit substantial out-of-domain reasoning generalization, enabling small base models to match or outperform multi-stage instruct models on competitive mathematics benchmarks.
2. Data-Aware Optimization (PCSS): While standard SFT is capable of eliciting this generalization, shifting from static time-based learning rate decay to data-aware loss calibration viaPCSS provides a principled optimization mechanism.

Ultimately, fine-tuning on these micro-datasets—whether through standard SFT or an adaptive loss scaler like PCSS—can provide a highly efficient way to elicit a base model's latent reasoning capabilities. Since training on these logic traces with PCSS appears to rapidly approximate the reasoning performance of a fully post-trained model, it may offer practical utility beyond standard alignment. For researchers evaluating pre-trained models, this methodology could potentially serve as an efficient evaluation proxy—providing a fast, low-cost way to gauge downstream reasoning performance without the need to run extensive post-training pipelines.