Technical Summary: AURORA – Contextual Orthogonalization for Geometric Representation Learning in Healthcare Foundation Models

1. Problem Statement

Recent healthcare foundation models have achieved strong predictive performance through large-scale self-supervised learning (SSL) on diverse biomedical modalities, including structured Electronic Health Records (EHR), imaging, and physiological signals. However, these models typically optimize for a single dominant objective: compressing patient information into a unified latent representation that maximizes downstream predictive accuracy.

The paper identifies a critical limitation in this paradigm: latent entanglement. In current models, distinct clinical concepts—such as physiological severity, intervention intensity, observational structure (e.g., measurement frequency), and institutional workflow—are frequently collapsed into shared embedding directions. While these entangled representations may yield high AUROC scores, they suffer from:

• Semantic Opacity: The latent space is difficult to interpret geometrically.
• Instability: Representations are unstable under contextual shifts (e.g., changes in hospital workflow or measurement protocols).
• Lack of Disentanglement: A single vector may simultaneously encode disease severity and operational intensity, making it difficult to distinguish physiological states from observational artifacts.

Existing SSL objectives (reconstruction, contrastive learning, self-distillation) provide no explicit mechanism to preserve the relational organization of contextual evidence or to enforce geometric separability between these factors.

2. Methodology: AURORA Framework

The authors propose AURORA (Adaptive Uncertainty-aware Representations through Orthogonalized Relational Alignment), a framework that reframes healthcare representation learning as a problem of latent semantic geometry rather than purely predictive compression.

Core Hypothesis

Clinically meaningful representations should not only preserve predictive information but also maintain the relational organization of contextual evidence. Distinct semantic factors should occupy geometrically separable (orthogonal) directions in the latent space.

Representation Decomposition

Instead of learning a single unified embedding $z$, AURORA decomposes the latent representation into semantically structured, orthogonal subspaces:
$$z = z^{(phys)} + z^{(int)} + z^{(obs)} + z^{(ctx)}$$
Where:

• $z^{(phys)}$: Captures physiological state.
• $z^{(int)}$: Captures intervention structure.
• $z^{(obs)}$: Captures observation and measurement behavior.
• $z^{(ctx)}$: Captures residual contextual information.

Learning Objectives

The framework is fully self-supervised, deriving relational structure directly from patterns within the EHR data without requiring expert annotations. It optimizes two primary objectives:

1. Contextual Alignment ($L_{align}$):
   For each semantic factor $k$, the model constructs patient neighborhoods based on contextual similarity (e.g., shared intervention patterns or similar laboratory acquisition structures). It minimizes the distance between representations of contextually similar patients within the corresponding subspace:
   $$L^{(k)}_{align} = \mathbb{E}_{i,j} \left[ R_k(x_i, x_j) \| z^{(k)}_i - z^{(k)}_j \|^2_2 \right]$$
   Here, $R_k(x_i, x_j)$ is a relational alignment operator measuring contextual similarity under factor $k$.

2. Orthogonalization Regularization ($L_{orth}$):
   To prevent the collapse of distinct factors into the same latent directions, the model penalizes the overlap between different semantic subspaces:
   $$L_{orth} = \sum_{k \neq l} \| (z^{(k)})^\top z^{(l)} \|_F^2$$
   This encourages the latent vectors for different factors to be geometrically orthogonal.

Total Loss:
$$L_{AURORA} = \sum_{k} L^{(k)}_{align} + \lambda L_{orth}$$
Where $\lambda > 0$ controls the trade-off between semantic alignment and contextual disentanglement.

3. Key Contributions

• Geometric Reinterpretation: The paper proposes treating the organization of the latent manifold itself as a primary modeling target, rather than a byproduct of predictive compression.
• Contextual Orthogonality Principle: It introduces a mechanism to explicitly separate clinically distinct factors (severity, intervention, observation) into orthogonal subspaces, addressing the entanglement inherent in standard SSL.
• Self-Supervised Disentanglement: The framework achieves semantic disentanglement without manual labels, relying solely on relational structures inherent in EHR data.
• Scalability: The approach is compatible with transformer-based encoders, multimodal architectures, and large-scale distributed pretraining systems.

4. Experimental Results

The authors evaluated AURORA against three classes of baselines: Masked Autoencoding (MAE), Contrastive SSL, and DINO-style Self-Distillation. All models used the same transformer backbone and parameter budget, pretrained on a large-scale EHR corpus.

Downstream Clinical Prediction

• Performance: AURORA consistently outperformed all baselines across mortality prediction, sepsis detection, readmission prediction, and patient retrieval.
• Retrieval: The most significant gains were observed in retrieval tasks (Recall@10), suggesting that contextual orthogonalization produces semantically cleaner neighborhoods in latent space.
• Robustness: Under institutional distribution shifts (changes in intervention frequency, lab ordering, and demographics), AURORA demonstrated substantially higher robustness (AUROC 0.879) compared to baselines (e.g., MAE at 0.781).

Disentanglement and Geometry

• Disentanglement Quality: AURORA achieved the lowest Mutual Information (MI) overlap (0.181) and highest Orthogonality Score (0.854), indicating successful separation of semantic factors.
• Neighborhood Purity: The model produced significantly higher neighborhood purity (0.843) and lower contextual entropy (0.61) compared to baselines, confirming that local manifold regions are organized by clinically meaningful semantics rather than mixed correlations.

5. Significance and Claims

The paper argues that the primary limitation of current healthcare foundation models is not insufficient scale (parameters or data), but insufficient structure within the latent representation space.

• Beyond Predictive Compression: The authors claim that optimizing solely for predictive performance leads to representations that are semantically opaque and unstable. By explicitly structuring the latent space according to contextual semantics, AURORA provides a complementary direction for foundation model design.
• Geometric Capacity as a Scaling Axis: The work suggests that "geometric capacity"—the semantic organization of the latent space itself—constitutes an additional axis of scaling for healthcare models, potentially improving robustness and interpretability without requiring larger architectures.
• Clinical Interpretability: The results suggest that disentangled representations are not merely more interpretable but are fundamentally more robust and predictive, as they allow the model to distinguish physiological severity from operational artifacts (e.g., measurement density).

The authors conclude that latent geometry is not merely an intermediate representation but a clinically meaningful object in its own right, and that explicitly structuring this space offers a path toward more reliable healthcare foundation models.