not KL/reverse-KL loss for on policy distillation
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 8.5k
- Forks
- 1.3k
- Avg merge
- 5h 36m
- Merged PRs (30d)
- 22
Description
when I read the implementation of on policy distillation based on slime, the loss function is not a KL loss. Concretely, the advantage is defined as:
https://github.com/THUDM/slime/blob/fcf4e6a2e54f2541f1fa62a3ab982115e12ae72f/slime/backends/megatron_utils/loss.py#L296
with no gradient. and then the loss is defined as:
https://github.com/THUDM/slime/blob/fcf4e6a2e54f2541f1fa62a3ab982115e12ae72f/slime/utils/ppo_utils.py#L133
where ratio is -P_s/P_ref. The loss increase P_s when P_t>P_s and decrease P_s when P_t<P_s. It is reasonable. But the supervised information it provides is not as dense as that provided by KL divergence. If we calculate the KL divergence, we need to pass the teacher_logits by rollout_data. Why didn't do that? Are there any concerns?
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start by reading the advantage calculation at slime/backends/megatron_utils/loss.py#L296 and the loss definition at slime/utils/ppo_utils.py#L133, then trace how rollout_data is constructed and consumed. Determine whether teacher_logits can be passed safely and whether a KL or reverse-KL objective is intended; done requires a maintainer decision and corresponding implementation or documentation with tests.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- machine-learning
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Stale
- Clarity
- Needs clarification
- Newbie friendliness
- 25/100