Table of Contents
Fetching ...

Knowledge-Embedded Latent Projection for Robust Representation Learning

Weijing Tang, Ming Yuan, Zongqi Xia, Tianxi Cai

TL;DR

Knowledge-Embedded Latent Projection (KELP) tackles the challenge of estimating latent representations from high-dimensional, imbalanced binary matrices by tying column embeddings to semantic side information through a reproducing kernel Hilbert space (RKHS) mapping, so $\mathbf{V}=\mathbf{K}_c\mathbf{A}$ with $\mathbf{v}_j=\boldsymbol{\varphi}(\mathbf{e}_j)$. The estimation combines kernel PCA-based subspace selection with scalable projected gradient descent, and a data-driven kernel selection procedure adapts to the alignment between semantic information and observed data. The authors derive non-asymptotic error bounds that decompose into estimation and approximation components, establish local linear convergence of the non-convex optimizer, and demonstrate improved performance over standard GLFM in simulations and a real-world electronic health records application for multiple sclerosis. The results show that incorporating semantic embeddings yields more efficient latent representations, enabling reliable downstream tasks such as MS knowledge reconstruction and disability phenotyping even when $n\ll p$, and offer practical extensions to unseen columns via Nyström methods and to covariate-augmented models.

Abstract

Latent space models are widely used for analyzing high-dimensional discrete data matrices, such as patient-feature matrices in electronic health records (EHRs), by capturing complex dependence structures through low-dimensional embeddings. However, estimation becomes challenging in the imbalanced regime, where one matrix dimension is much larger than the other. In EHR applications, cohort sizes are often limited by disease prevalence or data availability, whereas the feature space remains extremely large due to the breadth of medical coding system. Motivated by the increasing availability of external semantic embeddings, such as pre-trained embeddings of clinical concepts in EHRs, we propose a knowledge-embedded latent projection model that leverages semantic side information to regularize representation learning. Specifically, we model column embeddings as smooth functions of semantic embeddings via a mapping in a reproducing kernel Hilbert space. We develop a computationally efficient two-step estimation procedure that combines semantically guided subspace construction via kernel principal component analysis with scalable projected gradient descent. We establish estimation error bounds that characterize the trade-off between statistical error and approximation error induced by the kernel projection. Furthermore, we provide local convergence guarantees for our non-convex optimization procedure. Extensive simulation studies and a real-world EHR application demonstrate the effectiveness of the proposed method.

Knowledge-Embedded Latent Projection for Robust Representation Learning

TL;DR

Knowledge-Embedded Latent Projection (KELP) tackles the challenge of estimating latent representations from high-dimensional, imbalanced binary matrices by tying column embeddings to semantic side information through a reproducing kernel Hilbert space (RKHS) mapping, so with . The estimation combines kernel PCA-based subspace selection with scalable projected gradient descent, and a data-driven kernel selection procedure adapts to the alignment between semantic information and observed data. The authors derive non-asymptotic error bounds that decompose into estimation and approximation components, establish local linear convergence of the non-convex optimizer, and demonstrate improved performance over standard GLFM in simulations and a real-world electronic health records application for multiple sclerosis. The results show that incorporating semantic embeddings yields more efficient latent representations, enabling reliable downstream tasks such as MS knowledge reconstruction and disability phenotyping even when , and offer practical extensions to unseen columns via Nyström methods and to covariate-augmented models.

Abstract

Latent space models are widely used for analyzing high-dimensional discrete data matrices, such as patient-feature matrices in electronic health records (EHRs), by capturing complex dependence structures through low-dimensional embeddings. However, estimation becomes challenging in the imbalanced regime, where one matrix dimension is much larger than the other. In EHR applications, cohort sizes are often limited by disease prevalence or data availability, whereas the feature space remains extremely large due to the breadth of medical coding system. Motivated by the increasing availability of external semantic embeddings, such as pre-trained embeddings of clinical concepts in EHRs, we propose a knowledge-embedded latent projection model that leverages semantic side information to regularize representation learning. Specifically, we model column embeddings as smooth functions of semantic embeddings via a mapping in a reproducing kernel Hilbert space. We develop a computationally efficient two-step estimation procedure that combines semantically guided subspace construction via kernel principal component analysis with scalable projected gradient descent. We establish estimation error bounds that characterize the trade-off between statistical error and approximation error induced by the kernel projection. Furthermore, we provide local convergence guarantees for our non-convex optimization procedure. Extensive simulation studies and a real-world EHR application demonstrate the effectiveness of the proposed method.
Paper Structure (24 sections, 4 theorems, 16 equations, 2 figures, 1 algorithm)

This paper contains 24 sections, 4 theorems, 16 equations, 2 figures, 1 algorithm.

Key Result

Proposition 1

Let ${\bm{\Theta}} = \{\rho, {\bm{\alpha}}, {\bm{U}}, {\bm{V}}\}$ and $\overline{{\bm{\Theta}}} = \{\overline{\rho}, \overline{{\bm{\alpha}}}, \overline{{\bm{U}}}, \overline{{\bm{V}}}\}$ be two sets of parameters that specify the same conditional distribution in the model (eq: glfm-kf). Assume that

Figures (2)

  • Figure 1: Comparison between the proposed KELP estimator with data-driven kernel selection and the standard GLFM. Panels (a) and (b) show log-log plots of the relative estimation error for ${\bm{\Theta}}$ versus the number of rows $n$ and the the number of columns $p$, respectively. The upper and lower rows correspond to the linear (1) and nonlinear (2) mapping function settings. Panel (c) shows the relative errors for ${\bm{U}}$ and ${\bm{V}}$ versus the proportion of zero entries (sparsity level) under the linear mapping setting.
  • Figure 2: Evaluation of learned feature and patients embeddings on clinical downstream tasks. Panel (a) reports AuROC for recovering known MS-related clinical relationships, comparing three sets of feature embeddings: knowledge-fused LSM, GLFM trained solely on the MS cohort, and the original VA embeddings. Panel (b) reports AuROC for severe disability phenotyping at enrollment (left) and one-year follow-up (right), comparing two feature sets: the patient embeddings in an $r$-dimensional latent space and the original high-dimensional binary EHR features.

Theorems & Definitions (10)

  • Proposition 1: Identifiability
  • Remark 1: Comparison with generalized linear factor models
  • Remark 2: Connection to low-rank regression models
  • Remark 3: Complexity Comparison with Dual Formulation
  • Remark 4: Generalization to New Unseen Column Entities
  • Theorem 1
  • Remark 5: Impact of Sparsity
  • Remark 6
  • Theorem 2
  • Corollary 2.1