distribution on early exits as different KL scaling factor.

[marichkakorolyuk]

1. Plotting the distribution on early exits as different KL scaling factor.
2. Investigate coherence of final answers for different scaling factors. [not urgent, for final paper]

[puria-radmard]

1. Do normal generation, i.e. with teacher generate mode, for one new token
2. Collect the KL divergences, like teacher generate mode, for one new token
3. Then calculate early exit probabilities, and sample the prescription exit
4. Do prescribed student "SFT" mode with sampled early exit, i.e. freeze at sampled layer, get NTP distribution
5. "Manually" sample next token as if you had exitted at that layer
6. Repeat token by token - autoregressive generation

[karthikviswanathn]

The code and the visualizations can be found here showing the teacher’s behavior with freezing both enabled and disabled.

[puria-radmard]

1. Try with no early exit at prompt
2. Take advantage of huggingface past_key_values (KV cache) mechanism - keys cache list of tensors, each of shape shape [previous tokens, key dimensions] and same with values cache
3. Make sure past_key_values cache after early exits at each token are identical

[karthikviswanathn]

I have updated the code to reflect these changes. The generation looks way better if we use the KV cache.