Optimizing Pretraining for Safety Posttraining

Research Question: How can we optimize pretraining data curation to maximize the effectiveness of safety posttraining?

Motivation: The dominant paradigm for aligning LLMs relies almost entirely on posttraining to instill safe and desirable behavior. But posttraining pipelines are very fragile: they suffer from instability, reward hacking, and alignment taxes where safety gains come at the cost of capability. The field has largely treated the pretraining corpus as a neutral backdrop and focused all its optimization effort on the posttraining process itself. Emerging work (Mizrahi et al.) suggests this isn't optimal, and that the composition of the pretraining data materially affects how well posttraining can align a model. Separately, work on posttraining data curation (Zhang et al.'s DORM, Deng et al.'s BeeS) shows that curating for stability rather than pure quality produces more robust preference-tuned models, while Moskovitz et al. show that optimizing purely for quality often creates brittle models due to conflicting objectives.

This project sits at the intersection of these findings: instead of only curating posttraining data for stability, we ask whether we can identify and remove the pretraining documents that are the root cause of downstream posttraining instability. The goal is to build a pipeline that produces a stabilizing pretraining corpus, making safety posttraining cheaper, more robust, and less prone to reward hacking.

Methodology:

1. Data Curation: Use a standard pretraining mix (e.g. FineWeb) and a standard safety preference dataset augmented with synthetically generated paraphrases per prompt, creating "prompt families" that are semantically identical but lexically diverse. This lets us later measure variance in model behavior across meaning-equivalent inputs.
2. Baseline Pretraining: Train a policy model (e.g. Qwen 3.5) on the full pretraining corpus as a baseline.
3. Variance Auditing: Generate the policy model's responses across each prompt family, then score them with an off-the-shelf reward model (e.g. Skywork Reward V2) to find prompt families where paraphrases of the same underlying request receive very different rewards (meaning the model's behavior is unstable).
4. Influence Attribution: For high-variance prompt families, compute a "variance vector" as the difference between the parameter gradient of the high-reward response and the low-reward paraphrase. Then use TRAK (Park et al.), a data attribution method, to compute the dot product between this variance vector and representations of individual pretraining documents. Documents with high positive scores are mathematically implicated as sources of the instability.
5. Curation & Retraining: Filter out the top destabilizing documents identified by TRAK and retrain the policy model on the curated corpus, optionally iterating this cycle.
6. Posttraining & Evaluation: Align the retrained model via GRPO, then evaluate (a) stability, measured as reduction in reward variance across held-out perturbed prompt families, and (b) capability, measured on standard safety benchmarks, to check whether stability gains come at a capability cost.

Starting points for exploration: Mentees could begin by prototyping the "prompt family" generation and variance-scoring pipeline on a small scale (e.g, a few hundred prompts and a small open-source model) to validate that reward variance across paraphrases is a meaningful and measurable signal before scaling up. From there, the TRAK attribution step and retraining loop can be built incrementally, with several fallback options (e.g. using existing high-diversity prompt datasets instead of synthetic paraphrases, computing gradients only over final layers if full-parameter TRAK proves too noisy, or using simpler heuristics like leave-batch-out ablation if TRAK doesn't scale)

Risk Mitigation: Each stage of the pipeline (data curation, baseline training, variance auditing, influence attribution, curation/retraining, posttraining, evaluation) has documented smaller risks (e.g. paraphrases lacking lexical diversity, noisy TRAK attributions) & larger risks (e.g. compute exhaustion, GRPO non-convergence, filtered data removing real knowledge rather than noise) along with contingency plans, simplified fallback approaches, and small validation tests for each (e.g. injecting a synthetic "poisoned" document to confirm TRAK can detect it, training tiny 100M-parameter models on filtered vs. unfiltered data to check for catastrophic forgetting).

Relevant References:

• Mizrahi et al: Importance of pretraining for posttraining outcomes
• Moskovitz et al: Quality-optimization tradeoffs and instability from conflicting objectives
• Rafailov et al: Overoptimization/reward hacking in offline posttraining
• Zhang et al. (DORM) & Deng et al. (BeeS): Noise-reduced preference data curation for PPO/DPO
• Park et al: TRAK, the data attribution method central to this project

Timeline: I've created a rough draft for a 12-week plan covering tooling/setup, evaluation framework and reward model selection, prompt-family & variance-auditing pipeline construction, TRAK & GRPO pipeline setup, an initial pretraining-filtering-retraining cycle, adjustment based on results, ablation studies, and documentation, including buffer time

Note: this plan is a first draft, and mentees are encouraged to help reshape the methodology, timeline, or scope based on their own interests and expertise.

As models approach transformative capability, posttraining becomes the primary safeguard against catastrophic misuse or misalignment. But posttraining alone is fragile: it patches representations already shaped during pretraining, and current methods are known to suffer instability and reward hacking, which is exactly where safety matters most. This means transformative-model safety is currently bottlenecked on techniques we already know are unreliable at the margins.

This project's theory of change is that making pretraining itself safety-aware is a more scalable foundation for alignment than treating posttraining as the sole safety mechanism. If we can show specific pretraining documents causally drive posttraining instability (via TRAK-based influence attribution), we establish a proof of concept for curating pretraining data for alignment properties, not just capability. This matters in two ways:

1. It reduces reliance on brittle, late-stage safety patches. Removing destabilizing pretraining data before posttraining begins makes safety less dependent on getting posttraining exactly right under time and compute pressure, which is a realistic failure mode as capabilities advance.
2. It shrinks the reward-hacking surface at its source. Reward hacking often exploits inconsistencies learned during pretraining. A model trained on a stability-curated corpus should be intrinsically harder to push into hackable, inconsistent behavior, which matters more as models become capable enough to find and exploit subtle reward model weaknesses.

Other safety work I've completed:

• BiasGRPO: https://aclanthology.org/2026.findings-acl.2052/
• A Geometric Perspective (SPAR): https://arxiv.org/abs/2607.17946
• We Finally Have a Blueprint for Safe AGI: https://medium.com/@sreddyj2023/we-finally-have-a-blueprint-for-safe-agi-ab8fa512aa55
• Designing a Social Movement for Responsible AI (AISC, in progress)

Mentees will be active co-researchers, not just implementers of a fixed plan. The current methodology is a first draft, and mentees will have real input into shaping the direction (e.g. how to implement TRAK in a scalable way, how to measure a variance vector & whether a variance vector is the correct vector to compute, etc).

More specifially, mentees will:

• Own specific pipeline components. Depending on team size and interests, mentees may take primary responsibility for one or more stages (e.g. prompt-family generation & variance auditing, the TRAK/gradient-attribution pipeline, or the GRPO posttraining and evaluation setup)
• Make technical decisions. Mentees will decide things like which reward model to use after empirically testing candidates, how to cluster and interpret TRAK attributions, and how to design ablations
• Contribute to write-up and interpretation. Mentees will help analyze final results and author the documentation of findings in weeks 11-12

I expect to meet regularly (likely weekly) to review progress, unblock issues, and help decide on methodology adjustments.

• Proficient in Python and software engineering practices
• Proficient in Hugging Face, including experience fine-tuning 1-2B language models using Trainer
• Familiarity with core concepts in LLM alignment & posttraining (RLHF, DPO, GRPO)
• Comfortable reading & applying ideas from ML research papers