low_var_kl can return negative values near equal log probabilities due to cancellation

[yixiaoer]

I first observed this while running a GRPO update on a TPU, where a matched-reference control produced a small negative K3 value. I then reduced it to the minimal CPU reproducer below.

Expected Behavior

For method="low_var_kl", the per-token K3 estimator K3(d) = exp(d) - d - 1, with d = ref_logp - policy_logp, is nonnegative in real arithmetic. Near equal log probabilities it should return a small nonnegative value (about 5.0e-13 for the example below).

Actual Behavior

The float32 result is -5.9604645e-08. The current implementation evaluates jnp.exp(diff) - diff - 1 directly, so cancellation near diff = 0 can cause the computed value to violate the estimator's nonnegativity property. The float32 cast is not itself the issue.

Steps to Reproduce the Problem

Run with JAX_PLATFORMS=cpu:

import math
import jax.numpy as jnp
from tunix.rl.common import compute_kl_divergence

policy_logps = jnp.asarray([-1e-6], dtype=jnp.float32)
ref_logps = jnp.asarray([0.0], dtype=jnp.float32)
print(compute_kl_divergence(policy_logps, ref_logps, method="low_var_kl"))
print(math.expm1(1e-6) - 1e-6)

Observed output:

[-5.9604645e-08]
5.000001665616551e-13

Environment

• OS: Ubuntu 22.04, JAX CPU backend
• Project Version: google-tunix==0.1.7; the same direct expression is present on main at commit 091403963d42bda79f5e2ca60823e4c568a8039b (source).
• JAX / jaxlib: 0.11.1 / 0.11.1

Relation to #1560

#1560 added an optional symmetric output clamp for large outliers and overflow. Its default is None, and an ordinary positive clamp bound does not change this small negative result. This report concerns cancellation near zero, not the large-positive tail.

Possible direction

Use a numerically stable equivalent formulation near diff = 0 and add a test for near-equal policy and reference log probabilities.

Checklist

• I have searched the existing issues for a similar bug report.
• I have provided the environment details.
• I have provided a minimal, reproducible example.

Would you like to help us fix it?

Yes, I can follow up with a regression test and a PR.