Technical Summary: AI-Augmented Statistical Network Estimation with Proxy Gene Embeddings

Problem Formulation

The paper addresses the statistical challenge of recovering latent gene–gene interaction networks when the target network is observed only on a restricted set of genes ($Q$), while modern biomedical foundation models provide high-dimensional "proxy" embeddings for a substantially larger universe of genes ($P \supseteq Q$).

Existing methods for covariate-assisted network estimation typically assume covariates are observed on the same nodes as the network or treat embeddings as noisy proxies on the exact same node set. This work considers a distinct regime where:

1. Target Network ($A_Q$): An adjacency matrix observed only on a subset of genes $Q$.
2. Proxy Embeddings ($W$): A matrix of embeddings observed on a larger set $P$, which includes $Q$ and an additional set of genes $Q^c = P \setminus Q$.
3. Goal: To leverage the information in the embeddings for genes in $Q^c$ (which have no observed network data) to improve the estimation of the latent positions $U_Q$ for the target genes, thereby enhancing network denoising, community detection, and link imputation.

The authors formalize this via a joint latent-factor model where both the network and the embeddings are noisy observations of a shared low-dimensional latent geometry $U \in \mathbb{R}^{n_P \times k}$:
$$\mathbb{E}[A_Q | U_Q] = U_Q U_Q^\top, \quad \mathbb{E}[W | U] = U B^\top$$
where $B$ maps latent factors to the embedding space.

Methodology: PLANE

The authors propose Proxy-Latent Assisted Network Estimation (PLANE), an adaptively weighted joint estimator.

1. Objective Function

PLANE minimizes a weighted least-squares criterion combining the reconstruction errors of the network and the embeddings:
$$\ell(U, B) = \lambda_1 \|A_Q - U_Q U_Q^\top\|_F^2 + \lambda_2 \|W - U B^\top\|_F^2$$
where $\lambda_1 + \lambda_2 = 1$. The weights $(\lambda_1, \lambda_2)$ balance the information from the network channel and the embedding channel.

2. Optimization Algorithm

To solve this non-convex problem, the authors employ a blockwise normalized gradient descent scheme (Algorithm 1). Unlike standard gradient descent, the update steps are scaled by block-specific Gram matrices (e.g., $G_{U_Q} = 4\lambda_1 U_Q^\top U_Q + 2\lambda_2 B^\top B$). This normalization:

• Accounts for the local curvature of the loss landscape.
• Reduces sensitivity to heterogeneous singular values and scale imbalances.
• Facilitates a cleaner recursive analysis of the estimation error.

The algorithm iteratively updates $U_Q$ (target latent positions), $U_{Q^c}$ (embedding-only latent positions), and $B$ (loadings), aligning iterates via Procrustes rotation at each step to handle rotational non-identifiability.

Key Contributions and Theoretical Results

1. Identification and Zero-Order Limits

• Identifiability: Under mild rank conditions (full column rank of $U_Q$ and $B$), the pair $(U, B)$ is identifiable up to an orthogonal transformation. The network channel is shown to be essential for removing the linear indeterminacy inherent in factor models based solely on embeddings.
• Zero-Order Analysis: The authors demonstrate that a zero-order optimality analysis (comparing the objective value of the estimator to the truth) yields a sharp bound on the weighted reconstruction loss. However, this approach is insufficient to characterize the trade-off between the two data sources or to identify the optimal weighting $(\lambda_1, \lambda_2)$ for recovering the latent factors themselves.

2. Gradient Dynamics and Error Contraction

The paper's primary theoretical contribution is a deterministic contraction analysis of the normalized gradient descent algorithm.

• Curvature-Weighted Error: The authors define a specific error metric $e_t$ that couples the errors of $U_Q$, $U_{Q^c}$, and $B$, weighted by the local Gram matrices.
• Contraction Bound: They prove that this error contracts geometrically up to an explicit statistical tolerance $\Delta_S$. This tolerance depends on the operator norms of the noise matrices ($D$ for the network, $E$ for embeddings) and the tuning weights.
• Implication: This analysis extends gradient-dynamics theory from single-network objectives to joint network-factor models, explicitly characterizing how the weights $(\lambda_1, \lambda_2)$ and noise levels affect convergence.

3. Optimal Weighting and Trade-off Characterization

By specializing the error bound to the target block $U_Q$, the authors derive an explicit error function $E(\lambda_1, \lambda_2)$ that separates the contributions of network noise and embedding noise.

• Sharpness: They prove that the dependence of the error on the weights is sharp (up to constants) by comparing it to a partial-oracle estimator.
• Optimal Weight: The analysis yields an optimal, data-adaptive choice for the weights. The optimal weight $\lambda^*$ follows a harmonic-mean identity, effectively balancing the signal-to-noise ratios of the two channels.
• Guarantee: The theory guarantees that the optimal weighted estimator performs no worse than the network-only benchmark and, in unconstrained oracle comparisons, no worse than the better of the single-channel estimators.

Empirical Results

1. Simulations

• Adaptive Weighting: Simulations confirm that PLANE adaptively upweights the embedding channel when the proxy embeddings are informative (low noise) and downweights them when the network is more reliable.
• Benefit of Extra Nodes: The inclusion of embedding-only nodes ($Q^c$) improves the recovery of target latent positions ($U_Q$), provided the embeddings are informative. Cross-validation successfully downweights "null" proxies that offer no additional signal.

2. Real-Data Application: Single-Cell CRISPR Screens

The method is applied to the Norman et al. (2019) CRISPRa Perturb-seq dataset.

• Setup: The target network ($Q$) consists of 99 perturbed genes with observed differential expression signatures. The proxy embeddings ($P$) are constructed from a concatenation of multiple biological sources (GO pathways, STRING, GenePT text/protein embeddings, ESM) covering 1,099 genes.
• Performance: PLANE with cross-validated weights ($\lambda_2 \approx 0.6$) significantly reduced the held-out target-network prediction error compared to a network-only baseline.
• Biological Interpretability: The estimated latent positions were clustered into modules. The largest module centered around SLC6A9 (glycine transporter), while other modules captured immune, hematopoietic, and developmental regulatory programs (e.g., ZBTB1, FOXL2).
• Ablation Study: Removing individual embedding sources (e.g., GenePT text, ESM protein) showed that the full concatenated embedding provided the most robust improvement, indicating that the benefit arises from integrating complementary information rather than a single source.

Significance and Claims

The paper claims to be the first work to characterize the overall effect of joint estimation in a factor-assisted network estimation problem where embeddings are available on a superset of the network nodes.

• Statistical Insight: It resolves the "how to weight" question for integrating network and embedding data, proving that an optimal, data-adaptive weighting exists and can be approximated via cross-validation.
• Methodological Advance: It introduces a normalized gradient descent framework that handles the rotational non-identifiability and heterogeneous noise scales inherent in joint network-embedding models.
• Practical Utility: The results demonstrate that externally learned foundation model embeddings can serve as effective proxies to denoise sparse biological networks and impute missing links, provided the embeddings capture relevant biological signals. The method is robust, automatically downweighting uninformative proxies.

The authors maintain modesty regarding future directions, suggesting that extending the model to generalized link functions (e.g., for binary or signed networks) is a natural next step, but they do not propose specific new applications beyond the demonstrated utility in single-cell perturbation analysis.