[BUG] Multinomial probabilities are raw logits; Tangos crashes on single-row batches; sparsemax gradients wrong for size-1 dims

[ChrisW09]

Describe the bug

Three unrelated numerical/robustness defects.

---

MultinomialDistribution's softmax is never applied — predict(raw=False) returns raw logits labelled as probabilities

Where: deeptab/distributions/categorical.py (111-115 (transform lookup at deeptab/distributions/base.py:124))

The class stores its softmax as self.probs_transform while its param_names are p_0..p_{K-1}, so BaseDistribution.forward looks up the non-existent p_0_transform etc., silently falls back to identity, and predict(raw=False) returns unnormalised (possibly negative) logits.

Observed: d(torch.tensor([[1.,2.,3.]])) returns [[1., 2., 3.]] (identity) instead of softmax [[0.09, 0.24, 0.67]]. End to end, predict(raw=True) and predict(raw=False) are bit-identical (np.allclose == True); a transformed row is [0.486, 0.075, -0.494] — sums to 0.067 and contains a negative 'probability'. The registered default metric for this family is LogLoss (metrics/registry.py:47), which then receives negative probabilities. The same attribute-name mismatch exists in JohnsonSuDistribution (param_names contains "location" but the attribute is loc_transform, deeptab/distributions/student_t.py:91/96), so a caller-supplied loc_transform is applied inside compute_loss but not by forward.

Expected: predict(raw=False) should return softmax probabilities in [0,1] summing to 1 per row, consistent with the class docstring ('converted to probabilities via softmax') and with compute_loss, which does apply self.probs_transform.

Repro

import os, warnings; warnings.simplefilter('ignore')
import numpy as np, pandas as pd, torch
from deeptab.distributions import get_distribution
from deeptab.models import MLPLSS
from deeptab.configs import MLPConfig, TrainerConfig

d = get_distribution('multinomial', num_classes=3)
print(d.param_names, d(torch.tensor([[1., 2., 3.]])))   # -> ['p_0','p_1','p_2'] tensor([[1., 2., 3.]])

rng = np.random.default_rng(0); N = 48
X = pd.DataFrame({'a': rng.normal(size=N), 'b': rng.normal(size=N)})
oh = np.eye(3)[rng.integers(0, 3, size=N)]
m = MLPLSS(model_config=MLPConfig(layer_sizes=[16]),
           trainer_config=TrainerConfig(max_epochs=1, batch_size=16))
m.fit(X, oh, family='multinomial', distributional_kwargs={'num_classes': 3, 'total_count': 1},
      accelerator='cpu', devices=1, enable_progress_bar=False,
      enable_model_summary=False, logger=False)
print(np.allclose(m.predict(X, raw=True), m.predict(X, raw=False)), m.predict(X, raw=False)[0])

---

Tangos penalty_forward crashes with IndexError whenever a batch contains exactly one row

Where: deeptab/architectures/experimental/tangos.py (190 and 199)

jacobian.squeeze() removes all size-1 dimensions, so a batch of one row collapses the batch axis and neuron_attr.shape[2] no longer exists, aborting training.

Observed: IndexError: tuple index out of range at tangos.py:199 (spec_loss = torch.norm(neuron_attr, p=1) / (batch_size * h_dim * neuron_attr.shape[2])), reached from training_step -> estimator.penalty_forward(*data). With batch size 1 the jacobian (1, 1, h, F) squeezes to (h, F); after swapaxes(0, 1) it is (F, h) — two dimensions — and h_dim is silently taken to be the feature count. Direct unit calls with batch sizes 16, 4 and 2 all work and give finite penalties (0.1112 / 0.0158 / 0.0152) with flowing gradients, so this is purely the trailing-partial-batch case. Any dataset whose training split is ≡ 1 (mod batch_size) fails.

Expected: Training should complete; the batch dimension must be preserved, e.g. jacobian = jacobian.squeeze(1) (the representation's own singleton axis introduced by repr_forward's x.unsqueeze(0)) instead of a bare squeeze().

Repro

import os, warnings; warnings.simplefilter('ignore')
import numpy as np, pandas as pd
from deeptab.models.experimental import TangosRegressor
from deeptab.configs import TangosConfig, TrainerConfig

rng = np.random.default_rng(3)
N = 42                       # 42 rows, val_size 0.2 -> 33 train rows
X = pd.DataFrame({'a': rng.normal(size=N), 'b': rng.normal(size=N)})
y = rng.normal(size=N)
m = TangosRegressor(model_config=TangosConfig(layer_sizes=[16, 8]),
                    trainer_config=TrainerConfig(max_epochs=1, batch_size=16, val_size=0.2))
m.fit(X, y, accelerator='cpu', devices=1, enable_progress_bar=False,
      enable_model_summary=False, logger=False)   # 33 = 16 + 16 + 1  -> IndexError

---

sparsemax backward returns silently wrong gradients when a non-reduced dimension has size 1 (e.g. NODE with depth=1)

Where: deeptab/nn/blocks/common.py (112)

v_hat = grad_input.sum(dim=dim) / supp_size.to(output.dtype).squeeze() uses a bare .squeeze(), which removes every singleton axis rather than only dim. When any other axis of the input has length 1, the divisor loses that axis, the division broadcasts into a wrong-rank tensor, and the returned gradient is silently incorrect — no exception is raised. The forward pass stays correct, so the error is invisible.

Observed: depth=3: fwd=1.1e-16 grad_maxdiff=0.0000
depth=1: fwd=5.6e-17 grad_maxdiff=9.1890
A 2-D-with-singleton case shows the same thing (shape (3,1,5), dim=-1 -> grad_maxdiff 3.34). End-to-end, NODERegressor(model_config=NODEConfig(num_layers=1, layer_dim=4, depth=1)) fits for an epoch with no error at all — the tree feature-selection logits are just trained on garbage gradients.

Expected: grad_maxdiff == 0 for every shape, as it already is for depth=3. The divisor should be supp_size.to(output.dtype).squeeze(dim).

Repro

import torch
from deeptab.nn.blocks.common import sparsemax

def ref(x, dim):                      # differentiable reference simplex projection
    srt, _ = torch.sort(x, descending=True, dim=dim)
    cs = srt.cumsum(dim) - 1
    k = torch.arange(1, x.size(dim)+1, dtype=x.dtype)
    shape = [1]*x.dim(); shape[dim] = -1
    k = k.view(shape)
    ksz = (k*srt > cs).to(x.dtype).sum(dim=dim, keepdim=True)
    tau = cs.gather(dim, ksz.long()-1)/ksz
    return torch.clamp(x - tau, min=0)

torch.manual_seed(0)
# NODE feature_selection_logits shape = [in_features, num_trees, depth], reduced over dim=0
for nf, nt, depth in [(6,4,3), (6,4,1)]:
    x  = torch.randn(nf, nt, depth, dtype=torch.double, requires_grad=True)
    xr = x.detach().clone().requires_grad_(True)
    y, yr = sparsemax(x, 0), ref(xr, 0)
    g = torch.randn_like(y)
    y.backward(g); yr.backward(g)
    print(f"depth={depth}: fwd={float((y-yr).abs().max()):.1e} "
          f"grad_maxdiff={float((x.grad-xr.grad).abs().max()):.4f}")

---

Expected behavior
predict(raw=False) should return actual probabilities for multinomial; a batch of one row is normal
(it is simply the last batch of an odd-sized dataset) and must not crash Tangos; and sparsemax should
return correct gradients regardless of the size of the non-reduced dimensions.

Related but distinct: #426 records that sparsemax mutates its input in place — this is a separate
defect in its backward.

Screenshots
n/a

Desktop (please complete the following information):

• OS: macOS (Darwin 25.5.0, arm64)
• Python version: 3.11.15
• deeptab Version: 2.0.0 (main @ 4e6a359)

Additional context
torch 2.9.1, lightning 2.6.5, scikit-learn 1.9.0, numpy 2.4.6. Found in a second-pass review of v2.0.0
(seven independent lenses, each finding adversarially re-verified by a second reviewer, then re-run by
hand). Distinct from the already-filed #409-#426.