Multi-Teacher On-Policy Distillation selects one teacher for every batch item. Student logits have shape (B,S,V) for batch size B, sequence length S, and vocabulary size V. Teacher logits have shape (3,3,B,S,V): the first axis selects the domain and the second selects the reasoning-effort level. Domain and effort index tensors each have shape (B), and sampled token identifiers have shape (B,S).
For sampled token yb,t, use the two index tensors to select the teacher logits for batch item b, compute student and selected-teacher log probabilities with log-softmax, and gather the probability of yb,t. For positive clipping threshold Rmax, compute
rb,t=clip(stopgrad[logπteacher(db,eb)(yb,t)−logπθ(yb,t)],−Rmax,Rmax).The reward must have no gradient connection to either policy. Return a dictionary with rewards and teacher_token_log_probs. Both tensors have shape (B,S) and preserve the logits dtype and device.
Examples
Input: student_logits = [[[0,0]]], teacher_logits = [[[[[0,-0.02]]],[[[-0.24,-0.14]]],[[[-0.48,-0.26]]]],[[[[0.15,0.28]]],[[[-0.09,0.16]]],[[[-0.33,0.04]]]],[[[[0.3,0.58]]],[[[0.06,0.46]]],[[[-0.18,0.34]]]]], domain_indices = [0], effort_indices = [0], sampled_tokens = [[0]], clip_threshold = 2
Output: {"rewards": tensor([[0.009950000833311101]]), "teacher_token_log_probs": tensor([[-0.6831971797266342]])}
Explanation: The selected teacher gives token 0 a slightly higher log probability than the student, producing a small positive unclipped reward.
Input: student_logits.shape = (2, 2, 3), teacher_logits.shape = (3, 3, 2, 2, 3), domain_indices.shape = effort_indices.shape = (2,), sampled_tokens.shape = (2, 2), clip_threshold = 1.5
Output: {"rewards": tensor of shape (2, 2), "teacher_token_log_probs": tensor of shape (2, 2)}
Input: student_logits.shape = (1, 2, 2), teacher_logits.shape = (3, 3, 1, 2, 2), domain_indices.shape = effort_indices.shape = (1,), sampled_tokens.shape = (1, 2), clip_threshold = 0.75
Output: {"rewards": tensor of shape (1, 2), "teacher_token_log_probs": tensor of shape (1, 2)}
Hint 1
Use batch indices together with domain and effort indices to select one teacher tensor per item.
Hint 2
Apply log-softmax before gathering the sampled-token positions.
Hint 3
Detach the log-probability difference before clipping it to the reward interval.
Requirements
- Return a dictionary with exactly rewards and teacher_token_log_probs. Both tensors have shape (batch, sequence).
- Reward values lie within the inclusive clipping interval.
- The reward tensor has no gradient connection to either policy.
- Preserve floating-point dtype and device.
Constraints
- Teacher logits have leading dimensions (3, 3, batch, sequence, vocabulary).
- Domain and effort indices lie in [0, 2].
- Sampled token indices are valid vocabulary positions.
- Clipping threshold is positive.
All steps in Kimi K3
Accepts: array
Accepts: array
Accepts: array
Accepts: array
Accepts: array
Accepts: number
Multi-Teacher OPD
Multi-Teacher On-Policy Distillation selects one teacher for every batch item. Student logits have shape (B,S,V) for batch size B, sequence length S, and vocabulary size V. Teacher logits have shape (3,3,B,S,V): the first axis selects the domain and the second selects the reasoning-effort level. Domain and effort index tensors each have shape (B), and sampled token identifiers have shape (B,S).
For sampled token yb,t, use the two index tensors to select the teacher logits for batch item b, compute student and selected-teacher log probabilities with log-softmax, and gather the probability of yb,t. For positive clipping threshold Rmax, compute
rb,t=clip(stopgrad[logπteacher(db,eb)(yb,t)−logπθ(yb,t)],−Rmax,Rmax).The reward must have no gradient connection to either policy. Return a dictionary with rewards and teacher_token_log_probs. Both tensors have shape (B,S) and preserve the logits dtype and device.
Examples
Input: student_logits = [[[0,0]]], teacher_logits = [[[[[0,-0.02]]],[[[-0.24,-0.14]]],[[[-0.48,-0.26]]]],[[[[0.15,0.28]]],[[[-0.09,0.16]]],[[[-0.33,0.04]]]],[[[[0.3,0.58]]],[[[0.06,0.46]]],[[[-0.18,0.34]]]]], domain_indices = [0], effort_indices = [0], sampled_tokens = [[0]], clip_threshold = 2
Output: {"rewards": tensor([[0.009950000833311101]]), "teacher_token_log_probs": tensor([[-0.6831971797266342]])}
Explanation: The selected teacher gives token 0 a slightly higher log probability than the student, producing a small positive unclipped reward.
Input: student_logits.shape = (2, 2, 3), teacher_logits.shape = (3, 3, 2, 2, 3), domain_indices.shape = effort_indices.shape = (2,), sampled_tokens.shape = (2, 2), clip_threshold = 1.5
Output: {"rewards": tensor of shape (2, 2), "teacher_token_log_probs": tensor of shape (2, 2)}
Input: student_logits.shape = (1, 2, 2), teacher_logits.shape = (3, 3, 1, 2, 2), domain_indices.shape = effort_indices.shape = (1,), sampled_tokens.shape = (1, 2), clip_threshold = 0.75
Output: {"rewards": tensor of shape (1, 2), "teacher_token_log_probs": tensor of shape (1, 2)}
Hint 1
Use batch indices together with domain and effort indices to select one teacher tensor per item.
Hint 2
Apply log-softmax before gathering the sampled-token positions.
Hint 3
Detach the log-probability difference before clipping it to the reward interval.
Requirements
- Return a dictionary with exactly rewards and teacher_token_log_probs. Both tensors have shape (batch, sequence).
- Reward values lie within the inclusive clipping interval.
- The reward tensor has no gradient connection to either policy.
- Preserve floating-point dtype and device.
Constraints
- Teacher logits have leading dimensions (3, 3, batch, sequence, vocabulary).
- Domain and effort indices lie in [0, 2].
- Sampled token indices are valid vocabulary positions.
- Clipping threshold is positive.
All steps in Kimi K3
Accepts: array
Accepts: array
Accepts: array
Accepts: array
Accepts: array
Accepts: number