Efficient large-scale Cox regression with grouped sets and better rates
Scalable Cox Regression via Grouped Risk Sets and Sharper LogSumExp Rates
Machine Learning
Summary
Cox regression is a popular method used to analyze survival data, but it can be very slow when the data is huge. The authors found a new way to speed it up by grouping related events and using a special mathematical trick that reduces errors in calculations. They proved their method is more accurate and faster than older techniques and tested it on both synthetic and real data, showing it works well. This means it can help process large survival datasets more efficiently.
What this means in practice
- •For biomedical data teams: Analyze large patient survival datasets faster by using grouped risk sets to speed up Cox regression computations.
- •For insurance analytics teams: Improve the speed and accuracy of survival time predictions for risk modeling when using large-scale Cox regression.
Authors
Elizaveta Iashchinskaia, Egor Gladin
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 $\widetilde{O}(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.