Technical Summary: Enhancing Spectral Embedding through Robust and Flexible Knowledge Transfer in Electronic Health Records

Problem Statement

The paper addresses the challenge of unsupervised representation learning for rare disease cohorts using Electronic Health Records (EHRs). In these settings, data matrices are high-dimensional (many clinical concepts, $p$) but sample sizes are limited (few patients, $n$), often with $n \ll p$. Standard unsupervised methods like PCA require large datasets to learn robust low-dimensional embeddings for patients and clinical concepts.

While external knowledge (e.g., general biomedical knowledge graphs or embeddings from broader populations) offers a potential solution, direct application is problematic. Existing methods often assume a restrictive one-to-one signal alignment, where each principal component (PC) of the target data aligns perfectly with a single PC of the knowledge matrix. In reality, clinical data exhibits mixed alignment, where a target PC may align with multiple directions in both the knowledge subspace and its orthogonal complement. Furthermore, existing methods lack robustness against negative transfer, where incorporating weak or misaligned external signals degrades performance compared to using the target data alone.

Methodology: SENT (Spectral Embedding with kNowledge Transfer)

The authors propose SENT, a two-step spectral framework designed to handle mixed alignment and mitigate negative transfer.

1. Model Setup

The observed data $\hat{X} \in \mathbb{R}^{p \times n}$ is modeled as a noisy low-rank matrix: $\hat{X} = X + Z$, where $X = U\Lambda^{1/2}V^T$.

• $U$: Latent concept embeddings.
• $F = V\Lambda^{1/2}$: Patient embeddings.
• $W$: A transferable knowledge matrix derived from external sources, sharing a partially overlapping subspace with $U$.

The authors define the relationship between the latent subspace $\text{span}(U)$ and the knowledge subspace $\text{span}(W)$ using principal angles ($\theta_j$). They introduce a nonseparability parameter $\delta = \|Q_1^T \Lambda Q_2\|$ to quantify the mixing between shared components (aligned with $W$) and heterogeneous components (orthogonal to $W$).

2. Step I: Knowledge Preprocessing

To avoid negative transfer from irrelevant directions in a raw knowledge matrix $W_0$, the method first estimates a transferable rank ($r_W$) and a transferable knowledge matrix ($\hat{W}$).

• The algorithm performs SVD on $W_0$ to obtain a basis.
• It constructs nested candidate subspaces and selects the largest rank $k$ such that the subspace distance between the candidate and the baseline data estimator $\hat{U}^{(0)}$ (from PCA) remains within a tolerance bound ($c_W \epsilon_0$).
• This step effectively filters out non-transferable directions that do not align sufficiently with the target data structure.

3. Step II: Embedding Estimation

Using the filtered matrix $\hat{W}$, the method estimates embeddings via a projection-based approach:

• Shared Component: Projects the data covariance onto $\text{span}(\hat{W})$ to recover directions shared with the knowledge.
• Heterogeneous Component: Projects the data covariance onto the orthogonal complement $\text{span}(\hat{W})^\perp$ to recover directions unique to the rare disease cohort.
• The final estimator $\hat{U}$ combines these blocks. If $r_W = 0$, it reverts to standard PCA; if $r_W = r$, it relies entirely on knowledge-guided directions.

Key Contributions

1. Flexible Transfer Framework: The authors relax the one-to-one alignment assumption common in methods like JIVE and AJIVE. They formalize mixed alignment using subspace distances and a nonseparability parameter, allowing for more realistic modeling of clinical data where signals are cross-component.
2. Two-Step Design for Negative Transfer: The preprocessing step (Step I) explicitly identifies and removes non-transferable directions. This design provides robustness in scenarios where shared signals are weak relative to heterogeneous variation (e.g., demographics), preventing the degradation of performance seen in direct transfer methods.
3. Non-Asymptotic Theoretical Guarantees: The paper provides rigorous error bounds for both concept and patient embeddings.
   • The error bound decomposes into three terms: subspace distance ($\gamma_F$), mixing bias ($\tilde{\delta}$), and finite-sample noise.
   • Theoretical results show that informative concept-level transfer improves patient embedding estimation, provided the subspace distance is small.
   • The bounds are sharper than prior work by analyzing individual signal components rather than just aggregate spans.

Results

Simulations

• Performance: SENT consistently outperforms PCA, JIVE, AJIVE, and the method by Li et al. (2024), particularly in regimes with small sample sizes ($n$) and weak shared signals.
• Negative Transfer: In scenarios where shared signals are weak, competing methods (including Li et al., 2024) often underperform PCA due to negative transfer. SENT avoids this by filtering irrelevant knowledge.
• Ablation: Removing the preprocessing step leads to performance degradation, confirming the necessity of filtering non-transferable directions.

Real-World Application: Multiple Sclerosis (MS)

The method was applied to EHR data from the CLIMB registry (Mass General Brigham) for MS patients.

• Setup: External knowledge was derived from ONCE embeddings (large-population EHRs). The goal was to learn representations for a limited MS cohort ($n=100$ or $500$).
• Evaluation:
  • Concept Relations: SENT achieved the highest concordance with GPT-4o generated reference rankings of MS-relevant concepts.
  • Patient Prediction: SENT yielded the best AUC for predicting disability outcomes (EDSS and PDDS) at 6, 12, and 24 months.
• Interpretability:
  • SENT: Identified factors heavily loaded on neurologically specific concepts (e.g., "brain," "cerebral hemisphere structure").
  • Competitors: Methods like JIVE, AJIVE, and PCA tended to load heavily on generic utilization concepts (e.g., "pregnancy," "body weight," "edema") or mobility/rehabilitation terms, failing to capture the specific neuro-degenerative structure of MS as effectively as SENT.

Significance and Claims

The paper claims that SENT offers a statistically tractable and interpretable solution for unsupervised representation learning in rare disease settings. Its primary significance lies in:

• Robustness: It successfully leverages external knowledge without suffering from negative transfer, a critical issue when target cohorts are small and heterogeneous.
• Flexibility: By accommodating mixed alignment, it better reflects the complex reality of clinical data where disease-specific signals are often entangled with general population structures.
• Practical Utility: The method improves both the quality of learned concept embeddings (aligning with clinical relevance) and downstream patient-level predictions (disability outcomes), demonstrating that careful knowledge transfer can overcome the limitations of small sample sizes in EHR analysis.

The authors note that while they focus on a single knowledge source, the framework can be extended to multiple sources, and future work could explore nonlinear extensions or longitudinal settings.