fix(layers): make SoftExponential autograd-correct and NaN-safe (#788) (#809)

SoftExponential.forward branched on the learnable `alpha` parameter with
a Python `if`, which had three problems:

- `if self.alpha < 0.0` forces a host-device sync and drops the branch
  decision from the autograd graph, so `alpha` could only ever learn
  within whichever sign region it was initialised in.
- `self.alpha == 0.0` is an exact-float test that is unreachable after
  the first optimizer step, so the intended near-zero identity never
  fired.
- the `alpha < 0` formula `-log(1 - alpha*(x + alpha))/alpha` produces
  NaN/Inf for sufficiently negative `x`, which then poisons `alpha.grad`.

Select the branches with `torch.where` (keeping both formulas in the
graph), treat `|alpha| < eps` as the identity region, and guard the log
argument and denominators so the activation and its gradient stay finite.
Values in the well-defined region are numerically unchanged.

Adds regression tests for finiteness of the output and of `alpha.grad`
on negative inputs, and that `alpha` receives a usable gradient for both
signs.

Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Author: 0xSoftBoi

docs/changes.md

@@ -7,6 +7,14 @@ nav_order: 3
 # Change Log
 
 ## Unreleased
+- **Fix: `SoftExponential` activation autograd correctness and NaN safety (#788).** `forward` now selects
+  its `alpha < 0` / `alpha > 0` / `alpha ≈ 0` branches with `torch.where` instead of a Python `if` on the
+  learnable `alpha` parameter. The old `if self.alpha < 0.0` forced a host-device sync and dropped the
+  branch from the autograd graph (so `alpha` was effectively trapped in its initial sign region); the
+  `alpha == 0.0` exact-float test was unreachable after the first optimizer step; and the `alpha < 0`
+  formula produced NaN/Inf for sufficiently negative inputs. The log argument and denominators are now
+  guarded so both the activation and `alpha.grad` stay finite. Values in the well-defined region are
+  unchanged.
 - **New: `matgl.utils.MCDropoutWrapper` for uncertainty-aware inference.** Enables Monte Carlo Dropout
   (Gal & Ghahramani, 2016) on any pretrained MatGL model (CHGNet, M3GNet, TensorNet, …) without
   retraining: the backbone stays in `eval()` while only the readout dropout is sampled, and

src/matgl/layers/_activations.py

@@ -56,6 +56,11 @@ class SoftExponential(nn.Module):
     References: https://arxiv.org/pdf/1602.01321.pdf
     """
 
+    # |alpha| below this is treated as the identity (alpha -> 0) region, and it
+    # also floors the log argument / denominator to keep the activation and its
+    # gradient finite.
+    _eps = 1e-6
+
     def __init__(self, alpha: float | None = None):
         """Init SoftExponential with alpha value.
 
@@ -75,17 +80,41 @@ def __init__(self, alpha: float | None = None):
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         """Evaluate activation function given the input tensor x.
 
+        Branch selection uses ``torch.where`` rather than a Python ``if`` on
+        ``self.alpha``: a host-side ``bool(self.alpha < 0)`` test forces a device
+        sync and drops the branch from the autograd graph, so ``alpha`` could
+        only ever learn within whichever sign region it was initialised in.
+        ``torch.where`` keeps both formulas in the graph. Because ``where`` still
+        evaluates both sides, the denominators and the ``log`` argument are
+        guarded so the discarded branch cannot inject NaN/Inf into the gradient
+        (the "double where" trick).
+
         Args:
             x (torch.tensor): Input tensor
 
         Returns:
             out (torch.tensor): Output tensor
         """
-        if self.alpha == 0.0:
-            return x
-        if self.alpha < 0.0:
-            return -torch.log(1.0 - self.alpha * (x + self.alpha)) / self.alpha
-        return (torch.exp(self.alpha * x) - 1.0) / self.alpha + self.alpha
+        alpha = self.alpha
+        # Treat |alpha| < eps as the identity region (the alpha -> 0 limit)
+        # instead of testing exact equality to 0.0, which a learned float never
+        # hits after the first optimizer step.
+        near_zero = alpha.abs() < self._eps
+        # Never divide by a (near) zero alpha, even on the branch that gets
+        # discarded, otherwise NaN/Inf would poison alpha.grad.
+        safe_alpha = torch.where(near_zero, torch.ones_like(alpha), alpha)
+
+        # alpha < 0 branch: -log(1 - alpha*(x + alpha)) / alpha. The log argument
+        # can go <= 0 for sufficiently negative x; clamp it to a small positive
+        # floor on the live branch so it never produces NaN/Inf.
+        neg_log_arg = torch.where(alpha < 0.0, 1.0 - alpha * (x + alpha), torch.ones_like(x))
+        neg = -torch.log(neg_log_arg.clamp_min(self._eps)) / safe_alpha
+
+        # alpha > 0 branch: (exp(alpha*x) - 1)/alpha + alpha; expm1 is accurate near 0.
+        pos = torch.expm1(safe_alpha * x) / safe_alpha + safe_alpha
+
+        out = torch.where(alpha < 0.0, neg, pos)
+        return torch.where(near_zero, x, out)
 
 
 def softplus_inverse(x: torch.Tensor):

tests/layers/test_activations.py

@@ -18,14 +18,53 @@ def test_softplus2(x):
 
 
 def test_soft_exponential(x):
+    # alpha == 0 is the identity, but the output now participates in autograd
+    # (alpha is a learnable parameter), so detach before going to numpy.
     out = SoftExponential()(x)
-    np.testing.assert_allclose(out.numpy(), np.array([1.0, 2.0]))
+    np.testing.assert_allclose(out.detach().numpy(), np.array([1.0, 2.0]))
     out = SoftExponential(1.0)(x)
     np.testing.assert_allclose(out.detach().numpy(), np.array([2.7182817, 7.389056]))
 
     out = SoftExponential(-1.0)(x)
     np.testing.assert_allclose(out.detach().numpy(), np.array([0.0, 0.693147]), atol=1e-5)
 
 
+def test_soft_exponential_negative_x_is_finite():
+    """alpha < 0 must not produce NaN/Inf for large-negative inputs.
+
+    The alpha < 0 branch evaluates ``-log(1 - alpha*(x + alpha)) / alpha``. For
+    sufficiently negative ``x`` the log argument goes <= 0, which yields NaN/Inf.
+    Such inputs are easy to hit early in training when features are large.
+    """
+    act = SoftExponential(-1.0)
+    x = torch.tensor([-10.0, -5.0, -1.0, 0.0, 1.0])
+    out = act(x)
+    assert torch.isfinite(out).all(), f"non-finite output: {out}"
+
+
+def test_soft_exponential_alpha_grad_is_finite():
+    """dL/d(alpha) must stay finite even when inputs hit the undefined region.
+
+    A NaN in the forward pass propagates into ``alpha.grad`` and poisons the
+    optimizer, silently killing training.
+    """
+    act = SoftExponential(-1.0)
+    x = torch.tensor([-10.0, -5.0, 1.0])
+    act(x).sum().backward()
+    assert act.alpha.grad is not None
+    assert torch.isfinite(act.alpha.grad).all(), f"non-finite alpha.grad: {act.alpha.grad}"
+
+
+def test_soft_exponential_alpha_is_learnable_both_signs():
+    """alpha must receive a finite, non-zero gradient in both sign regions so it can train."""
+    for alpha_init in (-0.7, 0.7):
+        act = SoftExponential(alpha_init)
+        x = torch.tensor([0.5, 1.5, 2.0])
+        act(x).sum().backward()
+        assert act.alpha.grad is not None
+        assert torch.isfinite(act.alpha.grad).all()
+        assert act.alpha.grad.abs() > 0
+
+
 def test_softplus_inverse(x):
     assert torch.allclose(softplus_inverse(torch.nn.functional.softplus(x)), x, atol=1e-5)