# Scalable Cox Regression via Grouped Risk Sets and Sharper LogSumExp Rates

arXiv: https://arxiv.org/abs/2609.40120 · HTML: https://mlsanity.com/paper/2609.40120/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 12 of its 25 references that are in the map.
- First posted on arXiv 30 Sep 2026.

## Builds on (most cited first, 12 of 12)

- [Large-Scale Methods for Distributionally Robust Optimization](https://mlsanity.com/paper/2010.05893/) (2020)
- [Time-to-Event Prediction with Neural Networks and Cox Regression](https://mlsanity.com/paper/1907.00825/) (2019)
- [Finite-Sum Coupled Compositional Stochastic Optimization: Theory and Applications](https://mlsanity.com/paper/2202.12396/) (2022)
- [Unbiased Simulation for Optimizing Stochastic Function Compositions](https://mlsanity.com/paper/1711.07564/) (2017)
- [Improved Stochastic Optimization of LogSumExp](https://mlsanity.com/paper/2509.24894/) (2025)
- [A Near-Optimal Single-Loop Stochastic Algorithm for Convex Finite-Sum Coupled Compositional Optimization](https://mlsanity.com/paper/2312.02277/) (2023)
- [NeuCLIP: Efficient Large-Scale CLIP Training with Neural Normalizer Optimization](https://mlsanity.com/paper/2511.08417/) (2025)
- [SGD with Variance Reduction beyond Empirical Risk Minimization](https://mlsanity.com/paper/1510.04822/) (2015)
- [Unbiased scalable softmax optimization](https://mlsanity.com/paper/1803.08577/) (2018)
- [Mini-batch Estimation for Deep Cox Models: Statistical Foundations and Practical Guidance](https://mlsanity.com/paper/2408.02839/) (2024)
- [Approximating Partial Likelihood Estimators via Optimal Subsampling](https://mlsanity.com/paper/2210.04581/) (2022)
- [A Geometry-Aware Efficient Algorithm for Compositional Entropic Risk Minimization](https://mlsanity.com/paper/2602.02877/) (2026)

## Abstract (arXiv)

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 widetildeO(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.