Scalable Cox Regression via Grouped Risk Sets and Sharper LogSumExp Rates

Elizaveta Iashchinskaia; Egor Gladin

cs.LGSep 30, 2026

Scalable Cox Regression via Grouped Risk Sets and Sharper LogSumExp Rates

Authors: Elizaveta Iashchinskaia, Egor Gladin

Organizations: HSE University

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.

Figures & tables

Dataset	$N$	$d$	$m_{\rm tr}$	$K$	Tuning cap	Final cap
SUPPORT2	8,873	22	3,621	61	3	4.5
NWTCO	4,028	11	342	15	3	4.5
Correlated 100k	100,000	10	6,005	133	10	15
Correlated 1M	1,000,000	20	60,136	178	100	150
Independent 100k	100,000	20	41,921	166	10	15

Table 1: Datasets, training event counts

m_{\rm tr}

, group counts

K

\delta=0.05

, and per-fit vector budgets (millions). Dimension

d

is after preprocessing.

Figure 1: Regularized full-Cox training gap against vector work. Each dataset has the same cap for all five methods. Thin lines show ten optimizer seeds, thick lines their medians, and bands their interquartile ranges. Ours and BigSurv use running arithmetic averages of post-update coefficient iterates. The other methods use current coefficient iterates.

Figure 2: Coefficient-gradient approximation errors at

\beta_{\rm ref}

. The left panel isolates risk-set grouping using exact group normalizers; the right isolates softplus profiling at

\delta=0.05

. Filled markers satisfy the measured pointwise conditions of Theorem 4.1 ; hollow markers violate at least one applicable condition. Dashed lines are proportionality guides with arbitrary vertical offsets.

Appendix figures & tables10 assets

Supplementary material from the paper’s appendix.

Appendix

Method	Candidates
Ours	$(b,\eta_{0})\in\{128,256\}\times\{0.003,0.01,0.03,0.1,0.3,1\}$
Batch LSE	$(Q,c)=(8,8)$ : $\eta\in\{0.0009,0.003,0.009,0.03,0.09,0.3\}$ ; $(Q,c)=(16,64)$ : $\eta\in\{0.0003,0.001,0.003,0.01,0.03,0.1\}$
Minibatch Cox	$(b,\eta)\in\{32,256\}\times\{0.001,0.003,0.01,0.03,0.1,0.3\}$
BigSurv	Stratum size 20; $(\ell,\eta_{0})\in\{1,16\}\times\{0.006,0.02,0.06,0.2,0.6,2\}$
Cox-CC	$Q=32$ ; $(c,\eta)\in\{1,8\}\times\{0.0003,0.001,0.003,0.01,0.03,0.1\}$

Appendix

Table 2: Twelve candidate configurations per method: two structural settings and six learning rates per setting, shared across datasets. Here

b

is batch size,

Q

is the number of sampled cases,

c

the controls per case, and

\ell

the strata per update.

Method	SUPPORT2	NWTCO	Correlated 100k	Correlated 1M	Independent 100k
Ours	$(256,0.03)$	$(128,0.03)$	$(128,0.1)$	$(256,0.1)$	$(128,0.03)$
Batch LSE	$(16,64,0.001)$	$(16,64,0.01)$	$(8,8,0.009)$	$(16,64,0.0003)$	$(16,64,0.0003)$
Minibatch Cox	$(256,0.01)$	$(256,0.03)$	$(32,0.003)$	$(256,0.001)$	$(256,0.001)$
BigSurv	$(1,0.02)$	$(16,0.06)$	$(1,2)$	$(1,2)$	$(16,0.06)$
Cox-CC	$(8,0.001)$	$(8,0.03)$	$(8,0.03)$	$(1,0.0003)$	$(8,0.0003)$

Appendix

Table 3: Configurations minimizing the mean, over two tuning seeds, of the minimum saved validation Cox loss within the tuning budget. Tuples are

(b,\eta_{0})

for Ours,

(Q,c,\eta)

for Batch LSE,

(b,\eta)

for Minibatch Cox,

(\ell,\eta_{0})

for BigSurv, and

(c,\eta)

for Cox-CC (with

Q=32

). Candidate menus are given in Table 2 .

Dataset	Ours	Batch LSE	Minibatch Cox	BigSurv	Cox-CC
SUPPORT2	$\mathbf{3.93\cdot 10^{-5}}$	$1.47\cdot 10^{-3}$	$2.13\cdot 10^{-3}$	$1.47\cdot 10^{-4}$	$1.16\cdot 10^{-3}$
NWTCO	$\mathbf{4.96\cdot 10^{-5}}$	$7.77\cdot 10^{-3}$	$2.98\cdot 10^{-3}$	$1.18\cdot 10^{-4}$	$1.19\cdot 10^{-2}$
Correlated 100k	$\mathbf{7.92\cdot 10^{-6}}$	$1.76\cdot 10^{-2}$	$1.61\cdot 10^{-3}$	$1.28\cdot 10^{-4}$	$1.24\cdot 10^{-2}$
Correlated 1M	$\mathbf{1.39\cdot 10^{-6}}$	$4.79\cdot 10^{-4}$	$7.11\cdot 10^{-4}$	$3.64\cdot 10^{-5}$	$6.46\cdot 10^{-4}$
Independent 100k	$\mathbf{1.36\cdot 10^{-5}}$	$4.49\cdot 10^{-4}$	$2.27\cdot 10^{-4}$	$1.65\cdot 10^{-5}$	$2.68\cdot 10^{-4}$

Appendix

Table 4: Median terminal full-Cox training gap over ten optimizer seeds

Dataset	Method	$10^{-3}$	Cost (M)	$10^{-5}$	Cost (M)
SUPPORT2	Ours	10/10	0.3115	0/10	–
SUPPORT2	Batch LSE	1/10	4.4980	0/10	–
SUPPORT2	Minibatch Cox	0/10	–	0/10	–
SUPPORT2	BigSurv	10/10	0.0654	0/10	–
SUPPORT2	Cox-CC	3/10	4.4999	0/10	–
NWTCO	Ours	10/10	0.3108	0/10	–

Appendix

Table 5: Sustained training-gap attainment at saved post-update checkpoints through the final cap. Threshold columns give the number of attainers out of ten runs. Costs are median work in millions of vector units among attainers. A first crossing at the terminal checkpoint qualifies, a dash indicates no attainment.

Dataset	Method	Test Cox loss	Test C-index
SUPPORT2	Ours	$6.866840\pm 0.000700$	$0.602019\pm 0.000693$
SUPPORT2	Batch LSE	$6.867794\pm 0.001330$	$0.600758\pm 0.001080$
SUPPORT2	Minibatch Cox	$6.866488\pm 0.002054$	$0.601296\pm 0.001637$
SUPPORT2	BigSurv	$6.867411\pm 0.000304$	$0.602970\pm 0.000435$
SUPPORT2	Cox-CC	$6.867097\pm 0.001045$	$0.601770\pm 0.001126$
NWTCO	Ours	$6.227649\pm 0.001645$	$0.731458\pm 0.000142$

Appendix

Table 6: Test metrics at validation-selected checkpoints: mean

\pm

sample SD over ten optimizer seeds. One fixed split; unpenalized Cox loss per event (

\downarrow

), Harrell C-index (

\uparrow

Dataset	Method	Validation Cox loss	Validation C-index
SUPPORT2	Ours	$6.871160\pm 0.000429$	$0.594114\pm 0.000440$
SUPPORT2	Batch LSE	$6.871841\pm 0.000956$	$0.593927\pm 0.001286$
SUPPORT2	Minibatch Cox	$6.870807\pm 0.000519$	$0.594091\pm 0.000511$
SUPPORT2	BigSurv	$6.871334\pm 0.000152$	$0.594797\pm 0.000370$
SUPPORT2	Cox-CC	$6.870968\pm 0.000646$	$0.594495\pm 0.000652$
NWTCO	Ours	$6.153779\pm 0.000821$	$0.723944\pm 0.001528$

Appendix

Table 7: Validation metrics at checkpoints selected by validation Cox loss: mean

\pm

sample SD over ten optimizer seeds. One fixed split; unpenalized Cox loss per event (

\downarrow

), Harrell C-index (

\uparrow

Dataset	$K$	$\delta$ slope	$\rho$ slope
SUPPORT2	61	1.783	1.000
NWTCO	15	1.601	0.997
Correlated 100k	133	1.372	0.999
Correlated 1M	178	1.716	0.999
Independent 100k	166	1.866	0.999

Appendix

Table 8: Group count

K

\delta=0.05

and ordinary least-squares log–log slopes of coefficient-gradient error against

\delta

(grouping) and

\rho

(softplus), using every point in the respective grid at

\beta_{\rm ref}

Figure 3: Regularized full-Cox training gap against vector work and setup plus optimizer time. Thin lines show individual optimizer seeds, thick lines and bands give the median and interquartile range. Ours and BigSurv use running arithmetic averages of post-update coefficient iterates; Ours averages after projection. The other methods use current coefficient iterates.

Figure 4: Unpenalized validation Cox loss per event against vector work (left) and setup plus optimizer time (right); lower is better. Matched panels share Y-axis limits.

Figure 5: Validation Harrell C-index against vector work (left) and setup plus optimizer time (right); higher is better. Matched panels share Y-axis limits.

Explore similar work

Oct 22, 2025stat.ML

Survival of the fittest Cox model: Pivotal variable selection for time-to-event data

We revisit Cox's proportional hazards model to improve variable selection in survival analysis. A square-root transformation of the partial likelihood renders the selection of the regularization parameter pivotal, free of the unknown baseline hazard and censoring mechanism. The resulting criterion borrows from information criteria such as BIC and from penalized regression methods such as the lasso, taking the best of both. On simulated and real data, our method substantially improves upon state-of-the-art approaches used daily in support recovery.

Maxime van Cutsem, Sylvain Sardy

Department of Mathematics, University of Geneva

Time-To-Event Data Survival Analysis Model Selection Lasso

Ask PDF

May 15, 2026stat.ML

Isotonic Survival Regression: Calibrated Survival Distributions from Deep Cox Models

Time-to-event data is widespread across the life sciences and engineering, but it is typically encountered together with censoring, which complicates the application of standard machine learning methods. Deep Cox models have emerged as a popular method for analyzing time-to-event data because they gracefully handle censoring and can be used with unstructured data such as clinical text reports, genomic sequences, and pathology images. However, their predicted survival probabilities are often poorly calibrated, thus limiting their practical utility. In this paper, we propose a novel post hoc calibration method for Deep Cox models that uses isotonic regression to refine predicted survival probabilities without affecting discriminative power. We establish favorable theoretical guarantees, including a double-robustness property and asymptotic calibration. Experiments on synthetic and real-world clinical data demonstrate the empirical effectiveness of our method.

Anchit Jain, Kevin Zhang, Stephen Bates

EECS, MIT, Cambridge MA, USA.

Time-To-Event Data Survival Analysis Exponential Family Post-Hoc

Ask PDF

Jun 6, 2024cs.LG

On Regularization via Early Stopping for Least Squares Regression

A fundamental problem in machine learning is understanding the effect of early stopping on the parameters obtained and the generalization capabilities of the model. Even for linear models, the effect is not fully understood for arbitrary learning rates and data. In this paper, we analyze the dynamics of discrete full batch gradient descent for linear regression. With minimal distributional assumptions, we characterize the trajectory of the parameters and the expected excess risk. Using this characterization, we show that when training with any learning rate schedule and finite time horizon, the early stopped solution is equivalent to the minimum norm solution for a generalized ridge regression problem. We also prove that early stopping is beneficial for generic data with arbitrary spectrum and for a wide variety of learning rate schedules. We provide an estimate for the optimal stopping time and empirically demonstrate the accuracy of our estimate.

Rishi Sonthalia, Jackie Lok, Elizaveta Rebrova

Department of Mathematics, Boston College · ORFE Department, Princeton University

Early Stopping Kernel Ridge Regression Least Squares Regularization Gradient Descent Batch

Ask PDF