Preferences and scaling

Implement the reward-model loss used in RLHF, plus two scaling calculations.

• preference_prob(r_a, r_b): the Bradley–Terry probability that A is preferred, $\sigma (r_a-r_b)$. Works element-wise on arrays.
• reward_loss(r_chosen, r_rejected): the mean of $-\log \sigma (r_{\text{chosen}}-r_{\text{rejected}})$ over arrays of rewards. Compute it stably: $-\log \sigma (m)=\log (1+e^{-m})$, which np.logaddexp(0, -m) gives without overflow.
• training_flops(n_params, n_tokens): the $6ND$ estimate.
• chinchilla_loss(N, D): the fitted formula from the lesson with $E=1.69$, $A=406.4$, $B=410.7$, $\alpha =0.34$, $\beta =0.28$.