Structured Treatment Modeling in Deep Survival Analysis via Hazard Factorization

Deep learning models trained on electronic health records are increasingly used for clinical risk prediction, yet modeling heterogeneous treatment effects remains challenging. Most approaches treat treatment as an undifferentiated covariate (S-Learner), conflating treatment effects with baseline risk, while training separate models for treated and untreated patients (T-Learner) suffers from treatment imbalance and sparsity. We propose a structured hazard factorization that decomposes the hazard into a shared baseline component and a treatment-specific hazard ratio network, enabling direct estimation of time-varying, covariate-dependent hazard ratios without post-hoc computation. By sharing a baseline while isolating treatment effects, the framework acts as a hybrid between S- and T-Learners, improving efficiency and reducing majority-group dominance under imbalance. We further extend the model with differentiable subgroup assignment for regularized treatment effect estimation and inverse propensity weighting to adjust for confounding. In simulations with known ground truth, our approach improves hazard ratio recovery while maintaining competitive survival prediction, and the subgroup extension recovers latent heterogeneity when assumptions hold. On two real-world clinical cohorts from the UK Clinical Practice Research Datalink, the framework produces time-varying hazard ratios and identifies subgroups characterized by established risk factors. Our results demonstrate that explicit hazard factorization provides useful inductive bias for incorporating treatment into deep survival models, bridging flexible neural architectures with hazard ratio estimation familiar to clinical practice.