Graph Neural Network based Hierarchy-Aware Embeddings of Knowledge Graphs: Applications to Yeast Phenotype Prediction

Abstract: We present a method for finding hierarchy-aware embeddings of knowledge graphs (KGs) using graph neural networks (GNNs) enriched with a semantic loss derived from underlying ontologies. This method yields embeddings that better reflect domain knowledge. To demonstrate their utility, we predict and interpret the effects of gene deletions in the yeast Saccharomyces cerevisiae and learn box embeddings for KGs in the absence of a prediction task. We further show how box embeddings can serve as the basis for evaluating KG revisions. Our yeast KG is constructed from community databases and ontology terms. Low-dimensional box embeddings combined with GNNs are used to predict cell growth for double gene knockouts. Over 10-fold cross validation, these predictions have a mean $R<sup>2$~score~of~0.360, significantly higher than baseline comparisons, demonstrating that high-level qualitative knowledge is informative about experimental outcomes. Incorporating semantic loss terms in the training of the models improves their predictive performance ($R<sup>2$=0.377) by aligning embeddings with ontology structure. This shows that class hierarchies from ontologies can be exploited for quantitative prediction. We also test the trained models on triple gene knockouts, showing they generalise to data beyond those seen in training. Additionally, by identifying co-occurring relations in the yeast KG important for the cell-growth predictions, we construct hypotheses about interacting traits in yeast. A biological experiment validates one such finding, revealing an association between inositol utilisation and osmotic stress resistance, highlighting the model's potential to guide biological discovery.

• The paper introduces hierarchy-aware GNNs that combine heterogeneous knowledge-graph message passing with box embeddings and semantic losses for modeling ontological containment and disjointness.
• The distance-based semantic loss achieved the strongest digenic fitness prediction, raising mean R² from 0.348 for a standard GNN to 0.377 and reaching 0.415 on held-out trigenic data.
• The model generated a testable link between inositol utilization and NaCl stress resistance, which laboratory experiments validated through significant interaction effects at medium and high inositol levels.

Overview

This paper introduces Hierarchy-aware GNNs, a framework that combines graph neural network (GNN) representations of knowledge graphs (KGs) with box-based volumetric embeddings of ontological class hierarchies, enforced through semantic loss terms during training. The method is evaluated on a newly constructed heterogeneous KG for Saccharomyces cerevisiae, built from SGD, GO, APO, ChEBI, INO, MI, RO, and BioCyc. The authors demonstrate three applications: (i) end-to-end prediction of digenic gene deletion fitness, (ii) unsupervised box embedding learning over the KG, and (iii) a link-evaluation scheme for ranking candidate KG revisions. Notably, an interpretability-driven hypothesis generated by the model — an interaction between inositol utilisation and NaCl stress resistance — was validated experimentally in an automated laboratory.

Method: hierarchy-aware GNNs

The core idea is to treat the entire KG as TBox-style axioms: class assertions $C(a)$ are rewritten as $\{a\} \sqsubseteq C$ and role assertions $r(a,b)$ as $\{a\} \sqsubseteq \exists r.\{b\}$, yielding a computational graph whose nodes are classes and nominals and whose edges are existential restrictions. Each message-passing layer's output is interpreted as latent variables $\theta_l$, converted into axis-aligned boxes via a MinDeltaBox construction ($Z_i = z_i + \text{softplus}(\theta_i^Z)$), and regularised by losses enforcing subsumption containment and disjointness separation.

Two families of hierarchy losses are considered: a distance-based loss penalising violation of box containment, and an overlap-based loss on intersection-to-volume ratios, using Gumbel-box smoothing for non-zero gradients. Disjointness penalties prevent collapse; volume and small-box regularisers stabilise training. The total loss combines a task loss (MSE) with weighted semantic and negative-semantic terms plus L2 weight decay. The framework is architecture-agnostic; here it uses heterogeneous GraphSAGE layers with per-edge-type SAGEConv modules, domain-specific embedding dimensionalities, and separate embedding spaces per ontology domain.

A practical limitation acknowledged by the authors is that the source ontologies (GO, ChEBI) contain no disjointness axioms, so negatives must be synthesised from random sampling or manually introduced disjointness statements. Additionally, representing individuals as boxes permits instance intersections without clear semantic interpretation — a simplification relative to point-based KGE conventions.

The yeast knowledge graph

The KG integrates eight domains spanning roughly 296,000 classes, including ChEBI material entities (217,381 classes, depth 24), GO biological processes and molecular functions, APO phenotypes, BioCyc reactions, and a custom gene taxonomy. Phenotypes are modelled as subclasses of observable–qualifier intersections, linked to genes via has phenotype and, when chemical-specific, directly to ChEBI classes via custom relations such as hasChemNutrientUtilization_Decreased. Regulation and interactions are encoded as gene-specific INO/MI subclasses conditioned on GO processes. After adding reverse edges and dropping edge types with fewer than 1,000 instances (which improved performance, likely by reducing overfitting), 72 relation types remain.

Digenic fitness prediction results

The prediction task regresses fitness scores for 10,085,183 digenic deletions (Costanzo et al., 30 °C subset), with genetic-interaction edges removed to avoid leakage. Cross-validation splits are gene-based, discarding cross-set pairs. Gene-pair embeddings are combined by element-wise product (the best of four tested combination schemes) before a fully connected regression head.

Model	Mean $R^2$	SD
ComplEx + LightGBM	0.191	0.039
Phenotype instantiations + LightGBM	0.211	0.022
GNN without box embeddings	0.348	0.050
GNN with subClassOf-links in KG	0.350	0.049
GNN with prior box embeddings	0.360	0.043
GNN + overlap semantic loss	0.368	0.038
GNN + distance semantic loss	0.377	0.046

The best model reaches $R^2 = 0.377$, significantly outperforming all baselines and ablations under paired one-sided t-tests ($p < 0.05$). Two findings stand out. First, box prior embeddings outperform explicitly encoding hierarchies as graph links, suggesting volumetric containment acts as a useful inductive bias rather than mere additional connectivity. Second, the distance-based loss beats the overlap-based loss when combined with a task loss; the authors attribute this to its informative gradients whenever constraints are violated, whereas Gumbel-smoothed overlap losses remain nonzero even when constraints hold. Loss diagnostics show the pretrained priors have substantial inter-class overlap while GNN outputs discriminate classes well almost immediately, indicating the semantic loss primarily shapes the deeper layers.

On held-out trigenic deletion data (15,095 examples, Kuzmin et al.), models generalise with $R^2 = 0.380$ (priors only) and $\{a\} \sqsubseteq C$0 (with distance loss) — higher than digenic cross-validation, though partly because constituent genes were seen in double-deletion training pairs. Parity plots reveal a shrinkage effect: extreme low-fitness deletions are systematically overestimated.

The authors concede that no sequence information was used despite being fully available; this was deliberate to isolate the contribution of qualitative KG knowledge, but it bounds achievable performance.

Hypothesis generation and experimental validation

Using input×gradient attribution (Captum), importance scores are assigned to nodes adjacent to each deleted gene, and pairwise products aggregated across deletions yield globally ranked interacting trait pairs. Filtering to nutrient-utilisation-related, safely testable pairs selected an association between myo-inositol utilisation and NaCl stress resistance (the top-ranked pair involving bleomycin was discarded on safety grounds).

An experiment in the automated laboratory cell Eve used a $\{a\} \sqsubseteq C$1ino1 mutant across a factorial grid of inositol (0–0.25 mM) and NaCl (0–0.6 M) concentrations. A Gaussian GLM on growth-curve AUC found significant interaction terms for high-NaCl conditions with both medium ($\{a\} \sqsubseteq C$2) and high ($\{a\} \sqsubseteq C$3) inositol supplementation: inositol rescued cells from salt-induced stress, consistent with known roles of inositol in membrane integrity. This constitutes an experimental validation of a model-derived hypothesis, though the authors stress it is a single proof-of-concept; broader evaluation of reliability, novelty, and expert-perceived interest remains open.

Embeddings without a prediction task

Training the same architecture purely on hierarchy losses (with synthetic disjointness between molecular-function branches and random negatives) produces two-dimensional box embeddings that partially capture ontological semantics: disjointness between structural molecule activity, molecular function regulator activity, and small molecule sensor activity is respected. However, after message passing, the distance-loss embeddings show clear hierarchy violations and near-degenerate subclass placements, attributed to the flat gene-domain hierarchy and the absence of hierarchy information in message passing itself. Overlap-loss embeddings preserve superclass positions better but produce elongated boxes. The authors note that two dimensions are likely insufficient for this domain, and that the demonstration is illustrative rather than a benchmark claim.

Link evaluation for KG revisions

A proposed revision procedure adds candidate edges one at a time, re-runs the frozen GNN, and measures the distance between original and revised box embeddings. On an 80:20 stratified split, rank distributions of real test edges differ significantly from class-constrained random edges for most relation types (two-sided Mann–Whitney U, $\{a\} \sqsubseteq C$4), but not uniformly — e.g., hasChemStressResistance variants show no appreciable embedding change. Median distances are small with very large variance, so embedding distance alone does not discriminate true from random revisions. The authors suggest combining distance with semantic-loss change as a "surprise" signal for active learning, but this remains untested.

Limitations and open questions

Several limitations are stated plainly. Generalisation beyond this specific KG and task is unassessed; effectiveness plausibly depends on informative class hierarchies. The TBox reformulation simplifies semantics (box-valued individuals, existential restrictions treated identically to instance relations). Domain splits were manual. Hypothesis filtering introduces selection biases toward lab-feasible nutrient experiments. The link-evaluation method is rudimentary and fails for some relation types. Whether box-level message aggregation (as in BoxGNN-style geometric composition) would outperform auxiliary-loss enforcement is left as an open question, as is integration of sequence-encoded initial gene embeddings.

Conclusion

The paper contributes a flexible mechanism for injecting ontological hierarchy into GNN-based KG embeddings via box-containment semantic losses, showing measurable gains on quantitative genotype–phenotype regression ($\{a\} \sqsubseteq C$5 0.348 → 0.377) over strong baselines, successful transfer to trigenic data, and one experimentally validated biological hypothesis. The results support the claim that high-level qualitative knowledge carries quantitative predictive signal, while leaving open the breadth of applicability, faithful handling of DL semantics beyond shallow hierarchies, and systematic use of embedding perturbation for scientific hypothesis evaluation.