Convex-Ordered Count Head

Integrate this into the final scalar-output head for a count-regression network. Let the network output unconstrained values $a,b$; set $\mu=\operatorname{softplus}(a)+\epsilon$ and $\nu=\nu_{\min}+\operatorname{softplus}(b)$, with for example $\nu_{\min}=0.2$. For each example, solve for the CMP rate $\lambda$ whose mean is exactly $\mu$. Truncate the support at $M=\lceil\mu+8\sqrt{\mu+\mu^2}\rceil$, capped at a batch-wide maximum. Compute log-weights $w_m=m\log\lambda-\nu\log(m!)$, normalize with log-sum-exp, and estimate the mean and variance. Use Newton updates on $r=\log\lambda$: $r\leftarrow r-(\widehat{\mathbb{E}}X-\mu)/\widehat{\operatorname{Var}}(X)$, because the derivative of the CMP mean with respect to $\log\lambda$ equals its variance. Stop after 5--10 iterations or when the mean error is below $10^{-4}$. Train with the truncated negative log-likelihood. To impose a desired convex-order relation, choose a reference dispersion $\nu_{\rm ref}$, solve its rate to the same mean $\mu$, and compute convex test expectations for $f_1(m)=m^2$, $f_2(m)=(m-t)_+^2$ at several thresholds, and $f_3(m)=\exp(\alpha m)$ with small $\alpha$. If predictions must not be more dispersed than the reference, add $\rho[\mathbb{E}f(X)-\mathbb{E}f(X_{\rm ref})]_+$ for each test function. The theorem justifies the constraint for all convex functions when the death-rate ratio is monotone; the finite test set is an empirical approximation. First test on synthetic counts generated from CMP laws with means sampled from $[1,20]$ and $\nu\in\{0.5,1,2\}$, comparing Poisson, negative-binomial, unconstrained CMP, and regularized CMP heads on a two-layer MLP. Measure held-out NLL, mean bias, 90% interval coverage, tail-probability error, and convex-risk violations. Success means lower tail-calibration error and fewer catastrophic underpredictions at matched NLL, while preserving mean accuracy. Also numerically verify the ordering by evaluating all selected convex expectations over a grid of matched means before enabling the regularizer.