PaperlayerView exact revision on arXiv

Source-linked AI summary

Distributionally Robust Survival Models under Subpopulation Shift and Outlier Contamination

Seonghwi Kim, Sung Ho Jo, Minwoo Chae

arXiv:2610.02868v1cs.LGcs.AI

TL;DR

Survival models can perform poorly on shifted subpopulations, and outlier contamination compounds the challenge; existing DRO methods do not directly handle non-decomposable survival losses with sample interactions. The paper couples nominal-distribution refinement with adversarial reweighting and applies it to survival losses. Experiments report improved worst-group performance and competitive overall performance under shift and contamination.

  • Problem

    Subpopulation shifts can leave survival models performing poorly on some groups, while outlier contamination and cross-sample-dependent survival losses complicate robust learning.

    Source: §1 · pp. 1–2, §2 · p. 3

  • Method

    The framework refines a nominal distribution to reduce contaminated samples’ influence, then emphasizes challenging subpopulations while preserving non-decomposable survival-loss interactions.

    Source: §1 · p. 2

  • Results

    Experiments show robust performance under simultaneous shift and contamination, with improved worst-group performance and competitive overall performance across survival benchmarks.

    Source: §1 · p. 2, §7 · p. 9

  • Takeaways & Limitations

    The formulation preserves cross-sample interactions in survival losses, including the Cox risk-set structure, rather than relying only on additive losses.

    Source: §1 · p. 2

    The explicit two-stage trimming and reweighting characterization applies only to additive losses, not losses depending on multiple samples simultaneously.

    Source: §4.2 · p. 5

Abstract

from arXiv · show

Learning robust survival models under distribution shift is an important but challenging problem in many applications. In heterogeneous populations, a model that performs well on average may still perform poorly on certain subpopulations, and this issue becomes even more severe when the training data are contaminated by outliers. In this paper, we propose a novel distributionally robust framework for survival analysis that jointly addresses latent subpopulation shift and outlier contamination. The proposed method combines an outer minimization that selects a refined nominal distribution by reducing the influence of contaminated samples and an inner maximization that focuses on the most challenging subpopulation. This formulation directly accommodates non-decomposable survival losses while preserving interactions across samples, including the risk-set structure of the Cox negative partial log-likelihood. We develop an alternating gradient-based algorithm with outer updates derived from the KKT conditions of the inner maximization. Experiments on simulated data and two survival benchmarks demonstrate that the proposed method remains robust when subpopulation shift and outlier contamination occur simultaneously. It stabilizes training in contaminated settings and substantially improves worst-group performance across both linear and nonlinear survival models, while maintaining competitive and sometimes superior overall performance.

1 INTRODUCTION

The paper addresses survival models that can fail on shifted subpopulations and the difficulty of applying existing DRO to losses with cross-sample interactions. It couples nominal-distribution refinement with adversarial reweighting and reports robust performance under simultaneous shift and contamination.

Motivation

  • Existing additive-loss DRO methods do not directly accommodate survival losses with cross-sample interactions, while f-DRO can be sensitive to outliers.Cox negative partial log-likelihood depends on risk sets, and ranking losses use pairwise comparisons.Source: §1 · p. 2

Contribution

  • The coupled framework refines the nominal distribution to reduce contaminated samples’ influence, then emphasizes the most challenging subpopulation.The two stages jointly address latent subpopulation shift and outlier contamination.Source: §1 · p. 2
  • An alternating gradient-based algorithm optimizes non-decomposable survival losses while preserving sample interactions, with outer updates derived from inner-maximization KKT conditions.The framework is instantiated with the Cox negative partial log-likelihood and its risk-set interactions.Source: §1 · p. 2

Results

  • Across simulations and METABRIC and FLC, the method remains robust under simultaneous subpopulation shift and outlier contamination, improves worst-group performance, and maintains competitive overall performance.Simulations show stable training under contamination where f-DRO trajectories are unstable; on real benchmarks, worst-group performance improves across linear and nonlinear models.Source: §1 · p. 2

2 RELATED WORK

Prior survival models include classical statistical, machine-learning, and deep-learning approaches, while robustness to heterogeneity has been studied through group fairness, federated methods, and DRO. The paper identifies jointly handling latent subpopulation shift, outliers, and non-decomposable survival losses as an unresolved challenge.

Survival Analysis under Subpopulation Shift

  • Survival models span classical statistical methods, nonlinear machine learning, and deep learning, but do not explicitly account for subpopulation shift and can perform poorly on particular groups.The cited approaches include proportional hazards models and flexible risk-prediction methods.Source: §2 · p. 2
  • Fairness methods target parity across observed groups, whereas federated survival methods address institutional heterogeneity; both generally rely on prespecified group or institution structure.Source: §2 · p. 3

Distributionally Robust Optimization

  • DRO for latent shift emphasizes high-loss regions, including through f-DRO and CVaR-type objectives, and has also been extended to survival analysis.Source: §2 · p. 3
  • Existing robust survival methods address outliers but not latent shift, while outlier-robust f-DRO variants do not cover non-decomposable survival losses with cross-sample dependence.The paper frames jointly handling contamination and latent shift in survival analysis as an important problem.Source: §2 · p. 3

3 PROBLEM SETUP

The setup considers a latent-group target distribution whose mixture proportions may differ from clean training data, with observed training data further subject to arbitrary outlier contamination. The goal is robustness to both shifts, which ERM and standard f-DRO do not provide in this setting.

Assumptions

  • The target distribution is a mixture of latent groups, and clean training data may have different group proportions from the target.Observed training data follow Huber’s ϵ-contamination model with an arbitrary outlier distribution.Source: §3 · pp. 3–4

Learning Objective

  • The learning goal is robustness to training contamination and latent subpopulation shift, whereas ERM accounts for neither distributional shift nor contamination.Source: §3 · p. 4

Existing Alternatives

  • Group DRO requires group labels unavailable in this setting, while f-DRO reweights high-loss samples but remains vulnerable to outliers.Thus, f-DRO may fail when subpopulation shift and contamination occur together.Source: §3 · p. 4

4 PROPOSED METHOD

The method couples two distribution choices to refine nominal weights and emphasize difficult regions, then characterizes the structure for additive losses and extends the framework to non-decomposable survival losses. The survival-loss extension preserves cross-sample dependence, unlike the additive-loss closed-form characterization.

4.1 FORMULATION

  • The coupled formulation uses outer minimization to refine the nominal distribution and inner maximization to reweight challenging regions within that refined distribution.The inner ambiguity set depends on the refined distribution, linking the two optimization stages.Source: §4.1 · p. 4

4.2 STRUCTURAL ANALYSIS UNDER ADDITIVE LOSSES

  • For additive losses, the outer minimization performs ϵ-trimming and the inner maximization reweights the highest-loss α-fraction of remaining probability mass.This characterization depends on additivity and does not extend to losses depending on multiple samples simultaneously.Source: §4.2 · p. 5

4.3 NON-DECOMPOSABLE SURVIVAL LOSSES

  • The framework is extended to right-censored survival data, with the Cox weighted negative partial log-likelihood serving as a representative non-decomposable loss.The observed time is the minimum of survival and censoring times, and the event indicator records whether the event occurs by censoring.Source: §4.3 · p. 5
  • Because sample weights enter the Cox risk-set logarithms, the survival risk is nonlinear in weights and additive-loss trimming and reweighting do not carry over directly.The method therefore uses a gradient-based scheme that preserves cross-sample dependence.Source: §4.3 · p. 5

5 ALGORITHM

The algorithm alternates updates to sample weights and model parameters for the coupled objective. KKT multipliers from the inner maximization provide the outer update while preserving interactions in the full survival loss.

5 ALGORITHM

  • The gradient-based procedure assumes the loss is continuously differentiable in q and θ.The paper also notes that mini-batch implementations are possible.Source: §5 · pp. 5–6
  • The procedure alternates updates over q, q′, and θ, repeating weight updates before a gradient-descent step on model parameters.Full-batch updates are described, and stochastic mini-batch implementations are also possible.Source: §5 · pp. 6–7
  • KKT multipliers from the inner maximization provide a supergradient for updating q′, whose outer objective is otherwise implicit.The characterization relies on regularity and local stability assumptions.Source: §5 · p. 6
  • The q update ascends the full-loss objective while the q′ update descends using a supergradient estimate and projection.Both updates account for how sample-weight changes affect the full loss, including other samples’ Cox risk-set terms.Source: §5 · pp. 6–7

6 EXPERIMENTS

Across the contaminated survival benchmarks, the proposed method improves worst-group performance, with a trade-off in some average-performance measures. Simulations also show that setting ε to 0.05 stabilizes training and improves worst-group C-index.

6.1 SIMULATION

  • The simulation shifts the mixture proportions of two survival groups between training and a clean test distribution, while including a 2% training outlier component.f-DRO uses χ2-divergence and sample splitting.Source: §6.1 · p. 7
  • With ε = 0.05, the proposed method’s overall and worst-group C-index trajectories stabilize, and worst-group C-index is near 0.70.Both f-DRO and the proposed method at ε = 0 show large fluctuations; the standard Cox model is relatively stable but has worst-group C-index approximately 0.55.Source: §6.1 · pp. 7–8

6.2 REAL DATA ANALYSIS

  • On METABRIC, the method leads all baselines in overall and worst-group C-index in every setting; on FLC, it leads worst-group C-index and test tail loss, with weaker IBS and test log-likelihood than Cox models and Robust-Cox.METABRIC test tail loss is second-lowest.Source: Table 1 · p. 8, §6.2 · pp. 8–9, Table 2 · p. 9
  • METABRIC and FLC each use four clinically defined subgroups whose labels are unavailable during training and validation, with linear and neural Cox models evaluated under risk-inverted outliers.Outliers are constructed from the highest-risk 2% of training samples for linear models and 5% for nonlinear models.Source: §6.2 · p. 8
  • For METABRIC networks with hidden-layer sizes h = 4, 8, and 16, the method leads overall and worst-group C-index, IBS, and test log-likelihood, while f-DRO leads test tail loss.The method’s overall and worst-group discrimination gains persist across the larger networks evaluated.Source: Table 3 · p. 9, §6.2 · p. 9

7 CONCLUSION

The paper couples nominal-distribution refinement with adversarial reweighting to address subpopulation shift and outlier contamination in survival models. Experiments show improved worst-group C-index and low test tail loss under contamination, with competitive overall C-index.

7 CONCLUSION

  • The framework jointly addresses subpopulation shift and outlier contamination in survival analysis.It couples nominal refinement with adversarial reweighting while retaining cross-sample interactions in non-decomposable losses.Source: §7 · p. 9
  • Outer updates derived from the inner problem’s KKT multipliers yield an alternating algorithm that directly optimizes the coupled objective.Source: §7 · p. 9
  • Under contamination, experiments demonstrate improved worst-group C-index and consistently low test tail loss, alongside competitive overall C-index.Source: §7 · p. 9
  • Future work could develop ambiguity sets that capture within-group covariate shifts while retaining outlier robustness.Source: §7 · p. 9

A.1 FORMULATION AND ALGORITHM DETAILS

This section details the coupled DRO formulation, including trimming and worst-case reweighting for additive losses and KKT-based updates for non-decomposable losses, with conditions for weighted Cox loss. It also describes an alternating weight-and-parameter algorithm, its computational cost, and local convergence under exact inner maximization and stated assumptions.

A.1.1 ADDITIVE LOSSES

  • When α = 1, an optimal solution assigns mass to smaller-loss samples first, yielding exact ϵ-trimming.The allocation fills samples in increasing loss order up to their bounds, with a possible fractional boundary sample.Source: §A.1.1 · pp. 16–17
  • The inner maximization admits a weighted CVaR dual representation, which underlies the two-stage solution for general α.The passage identifies this representation as the basis for the two-stage structure.Source: §A.1.1 · p. 18
  • The general additive-loss solution trims toward smaller losses in the outer minimization, then concentrates inner mass on larger losses among those retained.Both stages have threshold structures with at most one fractional boundary index.Source: §A.1.1 · pp. 18–19
  • With uniform nominal weights, the solution exactly ϵ-trims the empirical distribution, then assigns uniform mass to the largest α(1 − ϵ)n retained losses.When boundary counts are nonintegral, the outer and inner allocations can be fractional at their respective cutoffs.Source: §A.1.1 · pp. 20–21

A.1.2 NON-DECOMPOSABLE LOSSES

  • For non-decomposable losses, the method alternates ascent in q, a q′ update based on an estimated supergradient, and descent in θ.The outer supergradient is characterized through upper-bound KKT multipliers under local stability.Source: §A.1.2 · pp. 21–22
  • For the weighted Cox loss, the supergradient characterization applies when the inner maximizer is unique, event risk-set sums are positive, and gradient coordinates are pairwise distinct.Under these conditions, the objective is continuous, continuously differentiable near the maximizer, and the inner optimizer is locally stable.Source: §A.1.2 · pp. 22–23
  • KKT stationarity and complementary slackness identify the upper-bound multipliers for an optimal inner solution.The proof applies these conditions to the inner maximizer and its constraints.Source: §A.1.2 · p. 24
  • When q′_i = b_i = 0, the lower- and upper-bound constraints are both active, so their multipliers are not separately identifiable from stationarity alone.This case is excluded from the multiplier characterization.Source: §A.1.2 · p. 24
  • At an exact inner maximizer, the gradient on interior coordinates equals ν⋆; approximate maximizers estimate this value for the q′ update.The estimate substitutes the approximate weights, gradient, and ν into the multiplier formula; exact characterization requires the stated assumptions at an inner optimum.Source: §A.1.2 · p. 25
  • The outer variables (q′, θ) converge locally under exact inner maximization and the stated assumptions; implementation approximates the inner maximizer with finite gradient steps.KKT multipliers estimated from the resulting iterate are used for the outer update.Source: §A.1.2 · p. 25

A.1.3 COMPUTATIONAL COMPLEXITY

  • O(KTn(d + log n)) is the linear Cox model’s total computational complexity, which is nearly linear in n for fixed K, T, and d.Risk scores and parameter derivatives cost O(nd), and the capped-simplex projections are computed by sorting.Source: §A.1.3 · p. 25
  • The algorithm requires O(n) additional memory for sample weights and cumulative risk-set quantities, beyond storage for data and model parameters.Empirical model-fitting times are reported in Appendix A.3.11.Source: §A.1.3 · p. 25

A.1.5 ALTERNATING GRADIENT ALGORITHM

  • Algorithm 3 alternates updates of adversarial weights q and nominal weights q′, followed by a model-parameter update.The algorithm is presented for the coupled DRO problem.Source: §A.1.5 · p. 27

A.1.4 LOCAL CONVERGENCE ANALYSIS

Under stated regularity conditions, the alternating scheme converges locally and R-linearly to a strict local minimizer. The analysis assumes exact global inner maximizers and sufficiently close feasible initialization.

A.1.4 LOCAL CONVERGENCE ANALYSIS

  • With exact global inner solves and sufficiently close feasible initialization, the outer iterates converge R-linearly to a strict local minimizer.The guarantee holds for sufficiently small step sizes under the theorem’s assumptions.Source: §A.1.4 · pp. 25–26
  • The analyzed cycle takes T KL-projected weight steps with θ fixed, then a parameter step using the last weight step’s inner solution.The analysis sets the weight and parameter step sizes equal.Source: §A.1.4 · p. 25
  • The guarantee requires a unique, smoothly varying inner maximizer, outer stationarity with positive active-bound multipliers, and positive curvature along feasible directions.The inner loss is twice continuously differentiable near the solution.Source: §A.1.4 · p. 25
  • Nearby KL updates retain the solution’s active upper bounds; fixing those coordinates lets the full update cycle contract, yielding geometric error decay.The proof also uses continuity of the inner maximizer to obtain convergence of the exact inner solutions.Source: §A.1.4 · p. 26

A.2.1 ROBUSTNESS PARAMETERS

The appendix tests robustness across parameter choices, simulated conditions, and real-data settings. Across these experiments, the proposed method often improves worst-group performance, with trade-offs in overall performance that vary by setting.

A.2.1 ROBUSTNESS PARAMETERS

  • Smaller α improves worst-group C-index under subpopulation shift, with a slight overall C-index decrease at α = 0.05; positive ϵ smooths training trajectories under contamination.The sensitivity experiment varies α with ϵ fixed at 0 and varies ϵ with α fixed at 0.1.Source: §A.2.1 · p. 28, Figure 2 · p. 28

A.3.1 SIMULATION: SUBPOPULATION SHIFT

  • Under simulated subpopulation shift, both robust methods can improve worst-group C-index at an overall C-index cost; the proposed method is faster than f-DRO but slower than standard Cox and Exact-Cox.The shift experiment uses 20 independent runs; the runtime comparison uses 100 iterations.Source: §A.3.1 · p. 29, Figure 3 · p. 29, §A.3.11 · p. 36

A.3.2 SIMULATION: NONLINEAR RISK FUNCTIONS

  • With ϵ = 0.05, the proposed method achieves higher late-stage worst-group C-index than Cox-NN and f-DRO in nonlinear simulations, with lower overall C-index than Cox-NN.At epoch 70, it achieves the lowest worst-group IBS, while Cox and Cox-NN have the lowest overall IBS in their respective settings.Source: §A.3.2 · p. 29, Figure 4 · p. 30, §A.3.4 · p. 31
  • With ϵ = 0.05, the proposed method outperforms f-DRO on all four metrics and improves worst-group C-index and IBS over Cox-NN, at a cost in overall performance.Table 4 reports results averaged over five independent repetitions.Source: Table 4 · p. 30, §A.3.2 · p. 30

A.3.3 SIMULATION: HIGHER DIMENSION, SMALLER SAMPLES, AND HEAVIER CENSORING

  • For d = 10, n = 500 or 1,000, and 50% censoring, robust methods improve worst-group C-index and IBS without contamination, while increasing ϵ improves overall C-index and IBS under contamination.Intermediate ϵ values achieve the best worst-group C-index and IBS among evaluated settings.Source: §A.3.3 · p. 30, Table 5 · p. 31, Table 6 · p. 31

A.3.8 REAL DATA ANALYSIS: SUBPOPULATION SHIFT

  • On METABRIC, the proposed method attains the highest overall and worst-group C-index with linear models; on FLC, its worst-group C-index is within 0.003 of standard Cox models.Across settings, it reduces test tail loss relative to standard Cox models, while those baselines achieve better IBS and test log-likelihood.Source: §A.3.8 · pp. 34–35, Table 7 · p. 34

A.3.9 ADDITIONAL OUTLIER SETTING

  • With appended outliers on FLC, the proposed method achieves the best worst-group C-index and test tail loss under both validation criteria, while its overall C-index is close to Elastic-Cox.Elastic-Cox achieves the best IBS, test log-likelihood, and overall C-index.Source: §A.3.9 · p. 35, Table 8 · p. 35

A.3.10 CONTAMINATION SENSITIVITY

  • At 5% contamination, the proposed method has the highest worst-group C-index in all four model–dataset settings; at 2%, it leads in linear settings but trails the best nonlinear baselines.Its test tail loss is lowest or second lowest in every setting.Source: §A.3.10 · p. 36, Table 9 · p. 36

Related papers

Loading 2610.02868v1…