Mini-batch Estimation for Deep Cox Models: Statistical Foundations and Practical Guidance

This paper establishes the statistical foundations of the mini-batch maximum partial-likelihood estimator (mb-MPLE) for deep Cox models optimized via stochastic gradient descent, proving its consistency and optimal convergence rates while providing practical guidance on hyperparameter tuning and demonstrating its effectiveness in large-scale applications where standard estimation is intractable.

The Gist

Imagine you are a doctor trying to predict how long a patient will live based on their medical history and genetic data. You have a massive database of millions of patients. To make the best prediction, you need to find the "perfect formula" that connects the data to the outcome. This is what statisticians call a Cox Model.

For a long time, finding this perfect formula was like trying to taste a giant pot of soup by taking a spoonful of the entire pot at once. You had to look at every single patient in your database to calculate the next step of your formula. If the pot was too big (a huge dataset), your kitchen (computer memory) would explode, or the spoon would break (the calculation would take forever).

This paper introduces a smarter way to cook: Mini-Batch Estimation. Instead of tasting the whole pot, you take a small spoonful (a "mini-batch") of patients, taste it, adjust your recipe, and repeat. This is called Stochastic Gradient Descent (SGD). It's fast and efficient, but it raises a big question: If we only taste small spoonfuls, are we still finding the true "perfect recipe," or just a lucky guess?

Here is the breakdown of what the authors discovered, using simple analogies:

1. The "Spoonful" Problem: It's Not Just a Smaller Pot

The authors realized that when you use a mini-batch for survival analysis (predicting time-to-event), you aren't just looking at a smaller version of the whole dataset.

• The Analogy: Imagine a race. To know who is winning, you need to know who is still running at every moment. In a full dataset, you know exactly who is still running. In a mini-batch, you only see a few runners. The "risk" calculation changes because the group of people you are comparing against is different.
• The Discovery: The authors proved that the "perfect recipe" found by tasting small spoonfuls (called mb-MPLE) is actually slightly different from the "perfect recipe" found by tasting the whole pot. However, they showed that as you get more data, this small-batch recipe gets closer and closer to the real truth. It's consistent and reliable.

2. The "Sweet Spot" of the Recipe (Batch Size vs. Learning Rate)

When you are training a neural network (a fancy computer brain), you have two main knobs to turn:

1. Batch Size: How big is your spoonful?
2. Learning Rate: How big of a step do you take when you adjust the recipe?

In normal machine learning, there is a famous rule: "If you double the spoon size, you can double the step size, and the result stays the same." This is called the Linear Scaling Rule.

• The Twist: The authors wondered if this rule works for survival analysis, where the "taste" depends on the group size.
• The Finding: Yes, it works! Even though the math is different, the relationship holds. If you use a bigger batch, you can take bigger steps. This gives doctors and data scientists a huge shortcut: they don't have to guess both knobs; they just need to keep the ratio between them constant.

3. The "Double-Edged Sword" of Batch Size

Here is a surprising finding that applies specifically to survival data (unlike other types of data):

• The Analogy: Imagine trying to find the bottom of a valley.
  • With a small batch, the ground feels a bit bumpy and wobbly. You might wander a bit before finding the bottom.
  • With a large batch, the ground becomes smoother and steeper (more "convex"). It's easier to slide straight to the bottom.
• The Discovery: In survival analysis, using a larger batch size actually makes your final answer more accurate (statistically more efficient). In many other types of AI, a larger batch just makes the training faster, but the final accuracy is the same. Here, bigger batches give you a better "statistical score."

4. The "Guardrails" for the Algorithm

The authors also looked at how the algorithm moves over time. They found that for survival data, the "valley" isn't perfectly shaped everywhere; it can get flat or weird at the edges.

• The Solution: They suggested putting up "guardrails" (a mathematical projection step) to keep the algorithm from wandering off into weird territory. This ensures that even if you run the algorithm for a long time, it will eventually settle on the correct answer.

5. Real-World Proof: The Eye Disease Study

To prove this wasn't just math on paper, they tested it on a real-world dataset involving Age-Related Macular Degeneration (AMD), a disease that causes blindness.

• The Challenge: They had thousands of high-resolution eye images. Trying to process all of them at once would crash a standard computer.
• The Result: Using their mini-batch method, they successfully trained a deep learning model to predict disease progression.
  • They found that using a smaller batch size (32 images) with a specific learning rate worked just as well as larger batches, provided they adjusted the "step size" correctly.
  • They achieved a high prediction accuracy (C-index of 0.85), proving that you don't need a supercomputer to analyze massive medical image datasets; you just need the right "spoon size" and "step size."

Summary

This paper tells us that we can analyze massive medical datasets using small, manageable chunks of data without losing accuracy.

• The Good News: You can use the "Linear Scaling Rule" (keep the ratio of batch size to learning rate constant) to tune your models easily.
• The Bonus: In survival analysis, using larger batches actually makes your predictions more precise, not just faster.
• The Bottom Line: This gives researchers the confidence to use powerful AI on huge medical datasets (like millions of patient records or images) without needing infinite computer memory, paving the way for better personalized medicine.

Technical Summary

Here is a detailed technical summary of the paper "Mini-batch Estimation for Deep Cox Models: Statistical Foundations and Practical Guidance" by Zeng et al.

1. Problem Statement

The Cox proportional hazards model is a cornerstone of survival analysis, but its application to large-scale data (especially with Deep Neural Networks, or Cox-NN) faces significant computational bottlenecks.

• Optimization Challenge: Standard training relies on the Maximum Partial Likelihood Estimator (MPLE), which requires computing gradients over the entire dataset at every iteration (Gradient Descent, GD). This is computationally prohibitive and memory-intensive for large datasets (e.g., high-dimensional medical imaging).
• Limitation of Standard SGD: While Stochastic Gradient Descent (SGD) using mini-batches is the standard solution for Deep Learning, it cannot be directly applied to Cox models in the same way as standard regression (e.g., MSE). In Cox models, the partial likelihood of an event depends on the "at-risk set" (all subjects surviving longer than the event time). A mini-batch gradient does not simply approximate the full gradient because the at-risk set changes based on the specific samples in the batch.
• The Gap: Existing literature applies SGD to Cox models heuristically, but the statistical properties of the resulting estimator—the Mini-batch Maximum Partial Likelihood Estimator (mb-MPLE)—remain unexplored. Specifically, it is unclear if the mb-MPLE is consistent, what its convergence rate is, and how hyperparameters (batch size vs. learning rate) should be tuned.

2. Methodology

The authors propose a rigorous theoretical framework to analyze the mb-MPLE and provide practical guidance for SGD in Cox models.

A. Statistical Framework

• Objective Function: They define the objective function that SGD actually minimizes: the expected value of the mini-batch partial likelihood, $E[L^{(s)}_{Cox}(\theta) | D^{(n)}]$. This differs from the standard full-data partial likelihood $L^{(n)}_{Cox}(\theta)$.
• Two Scenarios:
  1. Cox-NN (Non-parametric): Analyzing deep neural networks where the risk function $f_0(x)$ is unspecified.
  2. Cox Regression (Parametric): Analyzing linear covariate effects where $f_0(x) = \theta^T x$.
• Sampling Strategies: They distinguish between:
  • Stochastic Batch (SB): Sampling $s$ subjects without replacement from the full dataset at each step (standard SGD).
  • Fixed Batch (FB): Partitioning data into $m$ non-overlapping batches and cycling through them.

B. Theoretical Analysis

• Consistency & Convergence (Cox-NN): Using assumptions on smoothness (Hölder classes) and network sparsity, they prove the mb-MPLE is consistent. They establish that it achieves the minimax optimal convergence rate (up to a polylogarithmic factor), similar to the full-data MPLE, effectively circumventing the curse of dimensionality if the intrinsic dimension is low.
• Asymptotic Normality (Cox Regression): They prove the mb-MPLE is $\sqrt{n}$-consistent and asymptotically normal. Crucially, they derive the asymptotic variance, showing it depends on the batch size $s$.
• Convexity & Convergence:
  • They show that while the global objective is not strongly convex, it is locally strongly convex around the true parameter.
  • They introduce Projected SGD (constraining parameters to a ball) to guarantee convergence to the global minimizer in the parametric setting.
• Batch Size Impact: They analyze how the Hessian (local convexity) changes with batch size $s$. They prove that doubling $s$ increases the local convexity, thereby improving statistical efficiency (reducing variance), a phenomenon not observed in standard empirical risk minimization (e.g., MSE).

C. Practical Guidance (Hyperparameter Tuning)

• Linear Scaling Rule: In standard Deep Learning, the ratio of learning rate ($\gamma$) to batch size ($s$) is often kept constant to maintain training dynamics. The authors investigate if this holds for Cox-NN.
• Finding: They demonstrate that despite the batch-size dependence of the objective function, the linear scaling rule ($\gamma/s = \text{const}$) remains approximately valid for Cox-NN, especially when batch sizes are large. This allows practitioners to fix one parameter and tune the other.

3. Key Contributions

1. Statistical Foundations of mb-MPLE: First rigorous proof of consistency and minimax convergence rates for the mini-batch estimator in Deep Cox models.
2. Efficiency of Batch Size: Discovery that increasing batch size improves statistical efficiency (reduces asymptotic variance) in Cox models, unlike in standard MSE-based SGD where efficiency is independent of batch size.
3. Convergence Guarantees: Proof that Projected SGD converges to the mb-MPLE in Cox regression, providing theoretical backing for iterative optimization.
4. Hyperparameter Tuning Strategy: Validation of the linear scaling rule for Cox-NN, offering a practical method to tune learning rates and batch sizes without exhaustive search.

4. Results

• Theoretical Results:
  • Cox-NN: The estimator $\tilde{f}^{(s)}_n$ converges at rate $O_p(\Upsilon_n \log^2 n)$, matching the full-data MPLE.
  • Cox Regression: The estimator is asymptotically normal. The variance of the Fixed Batch (FB) estimator decreases as $s$ increases ($H_{2s}^{-1} \preceq H_s^{-1}$). The Stochastic Batch (SB) estimator is shown to be asymptotically more efficient than FB because FB ignores ranking information between non-overlapping batches.
  • Convergence: Projected SGD converges to the global optimum with a rate of $O(1/t^\alpha)$ depending on the learning rate schedule.
• Simulation Studies:
  • Verified that local convexity increases with batch size but plateaus for large $s$.
  • Confirmed that SGD-SB is more efficient than SGD-FB, particularly for small batch sizes.
  • Demonstrated that the Linear Scaling Rule holds: keeping $\gamma/s$ constant yields similar training trajectories (test error vs. epochs) across different batch sizes.
• Real-World Application (AREDS Data):
  • Applied Cox-NN to predict Age-related Macular Degeneration (AMD) progression using fundus images (7,865 eyes).
  • Feasibility: Full Gradient Descent was infeasible due to memory constraints (requiring ~48GB+). SGD with mini-batches made training possible.
  • Performance: The tuned model achieved a C-index of 0.85 on the test set.
  • Efficiency: Reducing batch size (and proportionally increasing learning rate) accelerated training time without sacrificing final predictive performance, validating the linear scaling rule.

5. Significance

• Enabling Large-Scale Survival Analysis: This work provides the theoretical justification for using SGD in Cox models, enabling the application of Deep Learning to massive survival datasets (e.g., electronic health records, medical imaging) that were previously computationally intractable.
• Bridging Theory and Practice: It moves beyond heuristic application of SGD to survival analysis, offering concrete statistical guarantees (consistency, normality) and actionable tuning guidelines (linear scaling).
• Novel Statistical Insight: The finding that batch size directly influences statistical efficiency in Cox models (unlike in standard regression) offers a new perspective on how to design experiments and training protocols for survival analysis.
• Generalizability: The authors note that their results extend to other ranking-based tasks using Plackett-Luce models (e.g., learning-to-rank, contrastive learning), suggesting broad applicability beyond just survival analysis.