Scalable Cox Regression via Grouped Risk Sets and Sharper LogSumExp Rates
Abstract
Motivated by the computational challenges of large-scale Cox regression, we study stochastic minimization of LogSumExp objectives over large sets. Mini-batch normalizer estimates generally yield biased gradients. We instead use a softplus surrogate that introduces one auxiliary scalar per normalizer and admits unbiased single-sample gradients. For smooth convex LogSumExp objectives, we prove an O(T−1/2) averaged objective bound, improving the previous T−1/4 analysis. With a strongly convex regularizer on the original variable, we also obtain a last-iterate squared-error rate of Õ(T−1) without strong convexity in the auxiliary variables. For Cox regression, the normalizers are defined over nested risk sets. We exploit this structure by grouping neighboring failures and sharing one auxiliary variable per group. The resulting compressed objective admits uniform score and curvature bounds that control the errors from grouping and softplus approximation. Together with the general optimization result, these bounds give a mean-square rate of T−4/5, up to logarithmic factors, relative to the full Cox solution. The compressed estimator also matches the full estimator's asymptotic distribution. Experiments on synthetic and real survival datasets with slowly decreasing risk sets show a favorable performance relative to stochastic baselines.
Related
- [EDGE] A grouped calibration test for logistic regression that tolerates a few corrupted records
- Recursive moment updates for scalable Bayesian nonparametric regression
- Scalable and Adaptive Distribution Regression via Random Features
- Scalable Nonparametric Change-Plane Regression with Random Features: Estimation and Inference
- Scalable Logistic Gaussian Process Density Regression with Kinetic Langevin Sampling
BibTeX
@misc{iashchinskaia2026scalable,
title = {{Scalable Cox Regression via Grouped Risk Sets and Sharper LogSumExp Rates}},
author = {Elizaveta Iashchinskaia and Egor Gladin},
howpublished = {arXiv},
year = {2026},
month = sep,
doi = {10.48550/arxiv.2609.40120},
eprint = {2609.40120},
archivePrefix = {arXiv},
url = {https://arxiv.org/abs/2609.40120},
}
RIS
TY - UNPB
TI - Scalable Cox Regression via Grouped Risk Sets and Sharper LogSumExp Rates
AU - Elizaveta Iashchinskaia
AU - Egor Gladin
PY - 2026
DA - 2026/09/30
PB - arXiv
DO - 10.48550/arxiv.2609.40120
UR - https://arxiv.org/abs/2609.40120
AB - Motivated by the computational challenges of large-scale Cox regression, we study stochastic minimization of LogSumExp objectives over large sets. Mini-batch normalizer estimates generally yield biased gradients. We instead use a softplus surrogate that introduces one auxiliary scalar per normalizer and admits unbiased single-sample gradients. For smooth convex LogSumExp objectives, we prove an O(T^−1/2) averaged objective bound, improving the previous T^−1/4 analysis. With a strongly convex regularizer on the original variable, we also obtain a last-iterate squared-error rate of Õ(T^−1) without strong convexity in the auxiliary variables. For Cox regression, the normalizers are defined over nested risk sets. We exploit this structure by grouping neighboring failures and sharing one auxiliary variable per group. The resulting compressed objective admits uniform score and curvature bounds that control the errors from grouping and softplus approximation. Together with the general optimization result, these bounds give a mean-square rate of T^−4/5, up to logarithmic factors, relative to the full Cox solution. The compressed estimator also matches the full estimator's asymptotic distribution. Experiments on synthetic and real survival datasets with slowly decreasing risk sets show a favorable performance relative to stochastic baselines.
ER -