JAX PathwaysCheck your starting point
Should response-only loss remove prompt tokens from the model context?
Compare your reasoning
No. Prompt inputs remain visible under the causal attention contract; the shifted response mask controls which targets receive loss.
If this is unfamiliar, work through Supervised fine-tuning with response-only token loss before attempting the project.
Build in stages
1. Supervised fine-tuning with response-only token loss
Write an input/next-target/role table, audit ragged likelihood sums, identify supervised positions in the gradient map, and demonstrate a context distinction the transition model cannot represent.
2. LoRA: adapt, save and merge low-rank updates
Derive and check both adapter-factor gradients, explain asymmetric initialization, verify nonunit merge scaling, and calculate a rank-one capacity limit before interpreting the adaptation curve.
3. Learn a reward model from pairwise preferences
Derive preference likelihood and its linear gradient, inspect contradictory labels, distinguish offsets from scale, and explain the measured comparison margins before handing off a reward model.
4. RLHF mechanics: a frozen reward and PPO policy updates
Identify current, behavior, reference and reward states; work both clipping signs; verify baseline cancellation by enumeration; and compare the trained bandit with its regularized optimum and probability plot.
5. Direct preference optimization and reference-corrected margins
Work through reference correction and its gradient, reproduce a falling chosen-probability counterexample, and explain why positive margins do not by themselves prove improved response behavior.
Try a changed condition
Make a prediction before opening the reasoning. Keep your answer and the evidence you would collect.
Reward rises, but independently reviewed outputs get worse. What has been established?
Compare an approach
Only improvement against the learned reward. Investigate reward exploitation, preference coverage and independent task metrics; do not label the result human alignment.
Diagnose a likely failure
Symptom: A PPO ratio is always one during multiple optimizer epochs.
Check next: Check whether old log probabilities are incorrectly recomputed from the changing policy.
Decide what is ready
Complete stages 4–8 of training-methods. Retain shifted token-mask checks, frozen-base/merge evidence, preference oracles, signed PPO clipping and fixed-reference DPO comparisons. Add the new lesson evidence: ragged likelihood oracle, rank-capacity counterexample, conflicting preferences, exact bandit optimum and a DPO ratio improvement with lower chosen probability. State which model integrations remain absent.
Where to put your project code
Write your code in: projects/posttraining-audit/starter/model.py (copy projects/posttraining-audit/starter/model.py first: cp projects/posttraining-audit/starter/model.py projects/posttraining-audit/starter/model.py).
Verify your implementation: python3 projects/posttraining-audit/tests/check.py --stage all --impl projects/posttraining-audit/starter/model.py
Run the staged project checks on your implementation in projects/posttraining-audit/starter/model.py and record your transfer experiment results in your portfolio.
Open Audit pretraining and post-training objectives
What still needs further work
The PPO lab is a terminal synthetic preference bandit, not full language-model RLHF. Human preference collection, sequence PPO/value models, large-model PEFT and held-out human evaluation remain separate work. For the learned-reward PPO lesson, complete rl-01 through rl-04 first. Response-only SFT and dense LoRA do not require the RL phase.