fix(dpmodel): gate the phantom denominator floor to the negative-count regime

Two review findings on the exponential-tail floor (OutisLi):

1. The floor fired for phantom_count >= 0, where no pole exists (the
   denominator is a plain positive softmax sum), breaking the documented
   term-for-term dense parity by orders of magnitude when every logit
   sits below phantom_logit ([0.5, 0.5] became [0.0009, 0.0009]).

2. Float32 autodiff produced NaN gradients for ordinary finite logits:
   with a tiny delta the backward of -denom/delta forms denom/delta**2,
   which overflows before being multiplied by the underflowed-to-zero
   exponential (0 * inf -> NaN), in both torch and jax.

Fix: apply the floor only where (a) the count is negative and (b)
denom < 40*delta -- beyond that the correction is < 1e-19 relative,
sub-ulp in float32 and float64, so the gate boundary is bit-invisible
and every in-design segment (and ALL count >= 0 segments) is now
BIT-exact with the floor-free formula. The inactive branch substitutes
(0, 1) for (denom, delta) BEFORE the division (safe-where), so the
pathological ratio is never constructed under autodiff. The exponent
caps move to 80 (finite exp in float32 too), including a new cap on
ph_term itself that kills a latent inf - inf NaN for logits > 80 below
phantom_logit.

Tests: phantom_count == 0 below-phantom exact-dense (bitwise vs the
phantom-free primitive), positive-count below-phantom exact-dense,
torch + jax float32 backward finite-and-correct in both floor regimes,
and LocalAtten / Atten2Map below-phantom graph-vs-dense parity at the
reviewer's construction.

Author: Han Wang

deepmd/dpmodel/utils/neighbor_graph/segment.py

@@ -113,11 +113,14 @@ def segment_softmax(
     term-for-term equal to the dense softmax.  Real logits are NOT bounded
     below by ``phantom_logit`` (the smooth envelope maps pre-shift logits
     ``raw < -attnw_shift`` to ``l < phantom_logit`` at ``sw > 0``), so for
-    negative counts the raw signed denominator can cross zero; a smooth
+    NEGATIVE counts the raw signed denominator can cross zero; a smooth
     strictly-positive floor (see the inline comment at the denominator)
-    keeps the normalization finite, positive and C-infinity for arbitrary
-    finite logits while perturbing the in-design regime by only
-    ~``exp(-2 * attnw_shift)`` relative.
+    keeps the normalization finite and positive there.  The floor is gated
+    to the negative-count, small-denominator regime, so for
+    ``phantom_count >= 0`` (a plain positive softmax sum -- no pole exists)
+    and for every in-design negative-count segment the output is
+    BIT-IDENTICAL to the floor-free formula: the dense term-for-term parity
+    is exact, not merely approximate.
     """
     xp = array_api_compat.array_namespace(data)
     dev = array_api_compat.device(data)
@@ -158,9 +161,18 @@ def segment_softmax(
         ex = ex * xp.astype(mask_b, ex.dtype)
     denom = segment_sum(ex, segment_ids, num_segments)
     if ph_b is not None:
-        ph_term = xp.exp(xp.full_like(seg_max, phantom_logit) - seg_max)
-        denom = denom + xp.astype(ph_b, denom.dtype) * ph_term
-        # Smooth strictly-positive floor for the SIGNED compensation: real
+        # exponent clamped at 80 (exp(80) ~ 5.5e34 is finite in BOTH float32
+        # and float64): pl - seg_max > 80 means every real logit sits > 80
+        # below the phantom logit; an unclamped exp would overflow (inf in
+        # float32 already at ~88) and poison the signed sum with inf - inf.
+        ph_arg = xp.minimum(
+            xp.full_like(seg_max, phantom_logit) - seg_max,
+            xp.full_like(seg_max, 80.0),
+        )
+        ph_term = xp.exp(ph_arg)
+        ph_f = xp.astype(ph_b, denom.dtype)
+        denom = denom + ph_f * ph_term
+        # Smooth strictly-positive floor for NEGATIVE counts only: real
         # logits are not bounded below by ``phantom_logit`` (the smooth
         # envelope maps ``raw < -shift`` to ``l < -shift`` at ``sw > 0``),
         # so with a negative count the signed denominator D can cross zero
@@ -169,19 +181,25 @@ def segment_softmax(
         #
         #     D~ = D + delta * exp(-D / delta),   delta = sel*ph_term / 40
         #
-        # Properties: (i) C-infinity in the logits and monotone-decreasing
-        # towards larger deficits; (ii) D~ >= delta > 0 for ANY finite
-        # logits (minimum at D == 0), so the weights stay finite and are
-        # bounded by ``40 * n_real / sel``; (iii) in the in-design regime
-        # every real logit is >= phantom_logit, hence D >= sel * ph_term
-        # == 40 * delta and the correction is <= exp(-40) * delta -- more
-        # than 1e-17 RELATIVE below D, i.e. it underflows to EXACT identity
-        # in float64 and the dense term-for-term parity is bit-preserved;
-        # (iv) the boundary-edge cancellation survives (a smooth function
-        # of the already-continuous D).  The exponent is clamped at 700 to
-        # avoid inf*0 pathologies for astronomically negative D (there the
-        # floor is astronomically large and the weights are ~0, as a
-        # maximal phantom deficit should give).
+        # applied ONLY where (a) the count is negative (for count >= 0 the
+        # denominator is a plain positive softmax sum -- no pole -- and the
+        # dense term-for-term parity must stay BIT-exact) and (b)
+        # D < 40 * delta (beyond that the correction is <= exp(-40) * delta
+        # < 1e-19 RELATIVE to D -- sub-ulp in float32 AND float64, so the
+        # gate boundary is bit-invisible; in-design every real logit is
+        # >= phantom_logit, hence D >= sel * ph_term == 40 * delta and the
+        # floor never fires).  Properties inside the active region: C-inf
+        # in the logits, D~ >= delta > 0 for any finite logits, weights
+        # bounded by ~40 * n_real / sel, and the boundary-edge cancellation
+        # survives (a smooth function of the already-continuous D).
+        #
+        # The inactive branch substitutes (0, 1) for (D, delta) BEFORE the
+        # division: even a zero cotangent cannot rescue exp(-D/delta) under
+        # autodiff when delta underflows (backward forms D/delta**2 -> inf,
+        # and 0 * inf = NaN in torch/jax float32), so the pathological
+        # division must never be constructed.  The exponent cap 80 bounds
+        # the active branch's exp for float32 while keeping D~ positive for
+        # any physical edge count (needs n_real/sel < exp(80)/40 ~ 1e33).
         if mask is not None:
             n_real_seg = segment_sum(
                 xp.astype(mask, denom.dtype), segment_ids, num_segments
@@ -192,12 +210,15 @@ def segment_softmax(
                 segment_ids,
                 num_segments,
             )
-        sel_b = xp.reshape(
-            n_real_seg, n_real_seg.shape + (1,) * (denom.ndim - 1)
-        ) + xp.astype(ph_b, denom.dtype)
+        sel_b = (
+            xp.reshape(n_real_seg, n_real_seg.shape + (1,) * (denom.ndim - 1)) + ph_f
+        )
         delta = xp.maximum(sel_b, xp.ones_like(sel_b)) * ph_term / 40.0
-        e_arg = xp.minimum(-denom / delta, xp.full_like(denom, 700.0))
-        denom = denom + delta * xp.exp(e_arg)
+        active = xp.logical_and(ph_f < 0.0, denom < 40.0 * delta)
+        denom_a = xp.where(active, denom, xp.zeros_like(denom))
+        delta_a = xp.where(active, delta, xp.ones_like(delta))
+        e_arg = xp.minimum(-denom_a / delta_a, xp.full_like(denom, 80.0))
+        denom = xp.where(active, denom + delta_a * xp.exp(e_arg), denom)
     denom_e = xp.take(denom, segment_ids, axis=0)
     safe = xp.where(denom_e > 0, denom_e, xp.ones_like(denom_e))
     return ex / safe

source/tests/common/dpmodel/test_repformer_graph_ops.py

@@ -328,6 +328,79 @@ def test_local_atten_parity(smooth):
     )
 
 
+def test_local_atten_below_phantom_dense_parity():
+    """OutisLi review: finite valid projection weights can push every smooth
+    logit below ``-attnw_shift``.  With ``n_real == sel`` the signed phantom
+    count is 0 -- the denominator is a plain positive softmax sum -- and the
+    graph route must match dense EXACTLY (the always-on floor used to return
+    ~0.0018 where dense gives 1.0)."""
+    la = LocalAtten(1, 1, 1, smooth=True, precision="float64", seed=1)
+    la.mapq.w = np.array([[1.0]])
+    la.mapkv.w = np.array([[-30.0, 1.0]])  # key -30, value 1
+    la.head_map.w = np.array([[1.0]])  # identity head
+    la.head_map.b = np.array([0.0])
+    nf, nloc, nnei = 1, 1, 2
+    g1 = np.ones((nf * nloc, 1))
+    gg1 = np.ones((nf, nloc, nnei, 1))
+    mask = np.ones((nf, nloc, nnei), dtype=bool)
+    sw = np.ones((nf, nloc, nnei))
+    ref = la.call(g1.reshape(nf, nloc, 1), gg1, mask, sw)  # == [[[1.0]]]
+    dst = np.zeros(nnei, dtype=np.int64)
+    got = la.call_graph(
+        g1,
+        gg1.reshape(-1, 1),
+        mask.reshape(-1),
+        sw.reshape(-1),
+        dst,
+        nf * nloc,
+        nnei,  # sel == n_real: phantom count 0
+    )
+    np.testing.assert_allclose(
+        np.asarray(got), ref.reshape(1, 1), rtol=1e-12, atol=0.0
+    )
+    # anti-vacuity: the logits really are below -attnw_shift and the dense
+    # result is the nontrivial value from the review
+    np.testing.assert_allclose(np.asarray(got), [[1.0]], rtol=1e-12)
+
+
+def test_atten2map_below_phantom_dense_parity():
+    """Same below-``-attnw_shift`` regime for the pair attention map: with
+    all key slots real (phantom count 0) graph must equal dense exactly."""
+    a2m = Atten2Map(1, 1, 1, has_gate=False, smooth=True, precision="float64", seed=2)
+    a2m.mapqk.w = np.array([[1.0, -30.0]])  # query 1, key -30 -> logits -30
+    nf, nloc, nnei = 1, 1, 2
+    g2 = np.ones((nf, nloc, nnei, 1))
+    h2 = np.zeros((nf, nloc, nnei, 3))
+    h2[..., 0] = 1.0
+    mask = np.ones((nf, nloc, nnei), dtype=bool)
+    sw = np.ones((nf, nloc, nnei))
+    ref = a2m.call(g2, h2, mask, sw)  # (nf, nloc, nnei, nnei, nh)
+    n_total = nf * nloc
+    dst = np.repeat(np.arange(n_total, dtype=np.int64), nnei)
+    q_e, k_e, pm = center_edge_pairs(
+        dst,
+        mask.reshape(-1),
+        n_total,
+        include_self=True,
+        ordered=True,
+        static_nnei=nnei,
+    )
+    got = a2m.call_graph(
+        g2.reshape(-1, 1),
+        h2.reshape(-1, 3),
+        sw.reshape(-1),
+        q_e,
+        k_e,
+        pm,
+        nf * nloc * nnei,
+        nnei,
+    )
+    ref_pairs = ref[0, 0, np.asarray(q_e) % nnei, np.asarray(k_e) % nnei, :]
+    np.testing.assert_allclose(np.asarray(got), ref_pairs, rtol=1e-12, atol=0.0)
+    # anti-vacuity: weight 0.5 per key slot times h2h2t = 1/sqrt(3)
+    np.testing.assert_allclose(np.asarray(got), 0.5 / np.sqrt(3.0), rtol=1e-12)
+
+
 def test_atten2map_graph_torch():
     import torch
 

source/tests/common/dpmodel/test_segment_softmax.py

@@ -2,6 +2,7 @@
 """segment_max / segment_softmax (NeighborGraph PR-D segment toolkit)."""
 
 import numpy as np
+import pytest
 
 from deepmd.dpmodel.utils.neighbor_graph import (
     segment_max,
@@ -122,9 +123,14 @@ class TestSignedPhantomStrictPositivity:
     raw signed denominator ``sum exp(l) + (sel - n) exp(ph)`` then crosses
     zero (e.g. ``sel=1``, logits ``[-21, -21]``, ``ph=-20``), which used to
     hit the ``denom > 0`` where-guard and return weights ``[1, 1]`` (sum 2)
-    on one side and a pole on the other.  The smooth floor
-    ``(D + sqrt(D^2 + delta^2)) / 2`` keeps the normalization finite,
-    positive and smooth for any finite logits.
+    on one side and a pole on the other.  The exponential-tail floor
+    ``D + delta * exp(-D / delta)`` keeps the normalization finite,
+    positive and smooth there; it is GATED to the negative-count,
+    small-denominator regime so that ``phantom_count >= 0`` (no pole
+    possible) and every in-design negative-count segment stay BIT-exact
+    with the floor-free dense formula, and its inactive branch is
+    constructed without the ``-D / delta`` division so float32 autodiff
+    stays NaN-free.
     """
 
     def test_reviewer_repro_finite_positive(self) -> None:
@@ -199,10 +205,121 @@ def _sweep(lo: float, hi: float, num: int) -> np.ndarray:
             f"-> fine {step_fine:.4f}): discontinuity or pole at the crossing"
         )
 
+    def test_zero_phantom_count_below_phantom_exact_dense(self) -> None:
+        """OutisLi review: ``phantom_count == 0`` means ``n_real == sel`` --
+        no phantom term exists and the denominator is a plain positive
+        softmax sum, so the output must be the EXACT dense softmax even
+        when every logit sits far below ``phantom_logit`` (the always-on
+        floor used to return ~0.0009 instead of 0.5 here).
+        """
+        data = np.array([-30.0, -30.0])
+        seg = np.array([0, 0], dtype=np.int64)
+        w = segment_softmax(
+            data, seg, 1, phantom_count=np.array([0.0]), phantom_logit=-20.0
+        )
+        np.testing.assert_array_equal(w, [0.5, 0.5])
+        # and BIT-equal to the phantom-free primitive on generic data
+        rng = np.random.default_rng(11)
+        data = rng.normal(size=9) * 30.0  # spans far below AND above -20
+        seg = np.array([0, 0, 0, 1, 1, 1, 1, 2, 2], dtype=np.int64)
+        w_zero = segment_softmax(
+            data, seg, 3, phantom_count=np.zeros(3), phantom_logit=-20.0
+        )
+        w_none = segment_softmax(data, seg, 3)
+        np.testing.assert_array_equal(w_zero, w_none)
+
+    def test_positive_count_below_phantom_exact_dense(self) -> None:
+        """count > 0 with all real logits below ``phantom_logit``: the
+        denominator is positive (phantom terms only ADD), so the fixed-width
+        dense softmax must be reproduced exactly -- the floor must not fire.
+        """
+        data = np.array([-30.0, -30.0])
+        seg = np.array([0, 0], dtype=np.int64)
+        w = segment_softmax(
+            data, seg, 1, phantom_count=np.array([1.0]), phantom_logit=-20.0
+        )
+        full = np.array([-30.0, -30.0, -20.0])
+        ref = np.exp(full - full.max())
+        ref = ref / ref.sum()
+        np.testing.assert_allclose(w, ref[:2], rtol=1e-15, atol=0.0)
+
+    def test_float32_torch_backward_finite(self) -> None:
+        """OutisLi review: with a tiny ``delta`` the floor's backward forms
+        ``D / delta**2`` -> inf, and ``0 * inf = NaN`` poisons the gradients
+        in float32 even though the forward is exact.  The gated safe-where
+        construction must keep float32 autodiff finite and correct, in both
+        the inactive (reviewer repro) and active (deep-deficit) regimes.
+        """
+        import torch
+
+        # inactive-floor regime: logits far above phantom_logit
+        data = torch.tensor([20.0, 21.0], dtype=torch.float32, requires_grad=True)
+        ids = torch.tensor([0, 0], dtype=torch.int64)
+        w = segment_softmax(
+            data,
+            ids,
+            1,
+            phantom_count=torch.tensor([-1.0], dtype=torch.float32),
+            phantom_logit=-20.0,
+        )
+        v = torch.tensor([1.0, 2.0])
+        (w * v).sum().backward()
+        assert torch.all(torch.isfinite(data.grad))
+        # analytic softmax-jacobian reference: g_i = w_i * (v_i - sum_j w_j v_j)
+        # (the phantom term exp(-41) is negligible at float32 resolution)
+        w64 = np.exp([20.0, 21.0] - np.float64(21.0))
+        w64 = w64 / w64.sum()
+        ref = w64 * (np.array([1.0, 2.0]) - (w64 * [1.0, 2.0]).sum())
+        np.testing.assert_allclose(data.grad.numpy(), ref, rtol=1e-4)
+        assert np.abs(ref).max() > 0.1  # nontrivial gradient, not all-zero
+
+        # active-floor regime: signed denominator below the floor threshold
+        data = torch.tensor([-21.0, -21.0], dtype=torch.float32, requires_grad=True)
+        w = segment_softmax(
+            data,
+            ids,
+            1,
+            phantom_count=torch.tensor([-1.0], dtype=torch.float32),
+            phantom_logit=-20.0,
+        )
+        (w * v).sum().backward()
+        assert torch.all(torch.isfinite(data.grad))
+
+    def test_float32_jax_backward_finite(self) -> None:
+        """Same float32 autodiff guarantee through JAX (the reviewer
+        reproduced the NaN gradient there as well).
+        """
+        jax = pytest.importorskip("jax")
+        jnp = jax.numpy
+
+        ids = np.array([0, 0], dtype=np.int64)
+        v = np.array([1.0, 2.0], dtype=np.float32)
+
+        def loss(x):  # noqa: ANN001, ANN202
+            w = segment_softmax(
+                x,
+                jnp.asarray(ids),
+                1,
+                phantom_count=jnp.asarray([-1.0], dtype=jnp.float32),
+                phantom_logit=-20.0,
+            )
+            return (w * jnp.asarray(v)).sum()
+
+        # inactive-floor regime (reviewer repro)
+        g = jax.grad(loss)(jnp.asarray([20.0, 21.0], dtype=jnp.float32))
+        assert np.all(np.isfinite(np.asarray(g)))
+        w64 = np.exp([20.0, 21.0] - np.float64(21.0))
+        w64 = w64 / w64.sum()
+        ref = w64 * (np.array([1.0, 2.0]) - (w64 * [1.0, 2.0]).sum())
+        np.testing.assert_allclose(np.asarray(g), ref, rtol=1e-4)
+        # active-floor regime
+        g = jax.grad(loss)(jnp.asarray([-21.0, -21.0], dtype=jnp.float32))
+        assert np.all(np.isfinite(np.asarray(g)))
+
     def test_floor_does_not_disturb_dense_parity_regime(self) -> None:
         """In the in-design regime (all logits >= phantom_logit, count >= 0)
-        the floor's correction is ~exp(-2*shift) relative -- invisible at
-        the 1e-12 level used by the dense-parity suites.
+        the floor never fires (its gate excludes non-negative counts), so
+        the fixed-width dense softmax is reproduced exactly.
         """
         rng = np.random.default_rng(3)
         data = rng.normal(size=7)  # normal-scale logits, shift 20 below