Score a batch of sampled responses with a KL penalty. For response i and token t, mit marks a selected response token, ℓitπ and ℓitref are its policy and reference log probabilities, ri is its reward, and β is the penalty coefficient. Compute
Di=t∑mit(ℓitπ−ℓitref) ri=ri−βDi J=B1i∑ri.Here, Di is the sampled sequence divergence, ri its penalized reward, B the batch size, and J the mean objective. This sampled divergence may be negative; do not clamp it or divide by response length. Masked positions contribute zero. An empty batch has objective zero.
Return a Python dictionary of float32 PyTorch tensors: objective is scalar; penalized_rewards and sampled_divergences each have shape (B,). All outputs remain on the input device.
Input: policy_log_probs = [[-1,-2]], reference_log_probs = [[-1.5,-2.5]], rewards = [3], response_mask = [[true,true]], beta = 0.5
Output: {"objective":2.5,"penalized_rewards":[2.5],"sampled_divergences":[1]}
Explanation: The two selected differences sum to one, so the reward receives a penalty of 0.5.
Input: policy_log_probs = [[-1,-2,-3],[-0.5,-0.5,-0.5]], reference_log_probs = [[-1.2,-2.1,-3.4],[-0.7,-0.4,-0.2]], rewards = [1,-1], response_mask = [[true,true,false],[true,true,true]], beta = 0.2
Output: {"objective":-0.01,"penalized_rewards":[0.94,-0.96],"sampled_divergences":[0.3,-0.2]}
Input: policy_log_probs = [[-1,-2]], reference_log_probs = [[-3,-4]], rewards = [-2], response_mask = [[false,false]], beta = 3
Output: {"objective":-2,"penalized_rewards":[-2],"sampled_divergences":[0]}
torch.where can exclude masked entries before reduction.
Tensor.sum reduces a chosen axis.
A PyTorch scalar tensor has shape ().
Accepts: array
Accepts: array
Accepts: array
Accepts: array
Accepts: number
Score a batch of sampled responses with a KL penalty. For response i and token t, mit marks a selected response token, ℓitπ and ℓitref are its policy and reference log probabilities, ri is its reward, and β is the penalty coefficient. Compute
Di=t∑mit(ℓitπ−ℓitref) ri=ri−βDi J=B1i∑ri.Here, Di is the sampled sequence divergence, ri its penalized reward, B the batch size, and J the mean objective. This sampled divergence may be negative; do not clamp it or divide by response length. Masked positions contribute zero. An empty batch has objective zero.
Return a Python dictionary of float32 PyTorch tensors: objective is scalar; penalized_rewards and sampled_divergences each have shape (B,). All outputs remain on the input device.
Input: policy_log_probs = [[-1,-2]], reference_log_probs = [[-1.5,-2.5]], rewards = [3], response_mask = [[true,true]], beta = 0.5
Output: {"objective":2.5,"penalized_rewards":[2.5],"sampled_divergences":[1]}
Explanation: The two selected differences sum to one, so the reward receives a penalty of 0.5.
Input: policy_log_probs = [[-1,-2,-3],[-0.5,-0.5,-0.5]], reference_log_probs = [[-1.2,-2.1,-3.4],[-0.7,-0.4,-0.2]], rewards = [1,-1], response_mask = [[true,true,false],[true,true,true]], beta = 0.2
Output: {"objective":-0.01,"penalized_rewards":[0.94,-0.96],"sampled_divergences":[0.3,-0.2]}
Input: policy_log_probs = [[-1,-2]], reference_log_probs = [[-3,-4]], rewards = [-2], response_mask = [[false,false]], beta = 3
Output: {"objective":-2,"penalized_rewards":[-2],"sampled_divergences":[0]}
torch.where can exclude masked entries before reduction.
Tensor.sum reduces a chosen axis.
A PyTorch scalar tensor has shape ().
Accepts: array
Accepts: array
Accepts: array
Accepts: array
Accepts: number
