Pretraining Curricula Enable Selective Fine-tuning

2026-07-06Machine Learning

Machine LearningArtificial Intelligence
AI summary

The authors studied how the order of training tasks affects what a Transformer model learns and how well it can be adjusted later. They found that training tasks in an unbalanced order, where one task is learned before the other, helps the model keep the tasks separate inside its 'brain.' This separation makes it easier to fine-tune the model to refuse unwanted behaviors. Their experiments with both simple memory tasks and language tasks support this idea, suggesting that choosing the right training order can improve model safety and control.

Transformer modelspretraining curriculumfine-tuningin-context learningdisentangled representationsactivation patchingrule representationmodel robustnessAI safetytask scheduling
Authors
Sebastian A. Bruijns, Jirko Rubruck, Mia H. Whitefield, Kai J. Sandbrink, Fazl Barez, Christopher Summerfield
Abstract
Transformers follow implicit curricula whereby some tasks are learned before others. However, how explicit pretraining curricula influence learning, generalization, and the selectivity of fine-tuning is unclear. This is important for AI safety, where fine-tuning is used to selectively suppress misaligned behaviors. Here, we compare curricula that pretrain tasks in a balanced (sampled uniformly) or an imbalanced (one task early, the other late) fashion. We show that imbalanced learning of two conflicting copy tasks promotes in-context learning and improves the selectivity of refusal fine-tuning. Ablations and activation patching show that this occurs because imbalanced pretraining encourages tasks to be disentangled in separable neural circuits, whereas balanced training routes both tasks through a common pathway. We extend these findings to a synthetic language learning task involving rule-consistent and rule-violating data, where imbalanced curricula similarly lead to more localized, less entangled rule representations, resulting in more robust rule-following behavior. Together, these results suggest that imbalanced pretraining curricula may be an important tool for promoting disentangled representations, with direct consequences for the precision and reliability of safety fine-tuning.