Local-Curvature-Aware Knowledge Graph Embedding: An Extended Ricci Flow Approach

TL;DR

This work tackles the mismatch between fixed geometric priors and local curvature heterogeneity in knowledge graphs by introducing RicciKGE, a framework that jointly evolves manifold geometry and entity embeddings through an extended Ricci flow that couples curvature with the KGE loss. The method updates edge weights via a curvature term and a gradient-driven term, and updates embeddings under a constrained optimization that respects the flow, forming a bidirectional feedback loop. The authors prove exponential curvature decay toward Euclidean geometry and linear convergence of the embedding objective under mild conditions, and they validate the approach with consistent improvements on link prediction and node classification benchmarks. The work advances curvature-aware KG embedding by enabling local-geometry adaptation and provides theoretical and empirical foundations for stable, co-evolving representations.

Abstract

Knowledge graph embedding (KGE) relies on the geometry of the embedding space to encode semantic and structural relations. Existing methods place all entities on one homogeneous manifold, Euclidean, spherical, hyperbolic, or their product/multi-curvature variants, to model linear, symmetric, or hierarchical patterns. Yet a predefined, homogeneous manifold cannot accommodate the sharply varying curvature that real-world graphs exhibit across local regions. Since this geometry is imposed a priori, any mismatch with the knowledge graph's local curvatures will distort distances between entities and hurt the expressiveness of the resulting KGE. To rectify this, we propose RicciKGE to have the KGE loss gradient coupled with local curvatures in an extended Ricci flow such that entity embeddings co-evolve dynamically with the underlying manifold geometry towards mutual adaptation. Theoretically, when the coupling coefficient is bounded and properly selected, we rigorously prove that i) all the edge-wise curvatures decay exponentially, meaning that the manifold is driven toward the Euclidean flatness; and ii) the KGE distances strictly converge to a global optimum, which indicates that geometric flattening and embedding optimization are promoting each other. Experimental improvements on link prediction and node classification benchmarks demonstrate RicciKGE's effectiveness in adapting to heterogeneous knowledge graph structures.

Figures (3)

• Figure 1: The framework of RicciKGE, which couples KGE gradients and local graph curvature in an extended Ricci flow, driving the latent manifold toward Euclidean flatness. The evolving geometry, in turn, guides curvature and gradient convergence and induces a closed-loop mechanism for bidirectional fitting between space and embedding.
• Figure 2: (a) Normalized loss and curvature variance curves on WN18RR (embedding dimension = 32). Curvature variance converges faster than loss, with both eventually stabilizing. (b) DistMult performance on WN18RR under varying coupling coefficients, showing that overly small or large values degrade effectiveness.
• Figure 3: Performance comparison of baseline KGE models and RicciKGE-enhanced variants on WN18RR over training steps. Blue curves represent the vanilla baseline models, while yellow curves denote models trained with our RicciKGE extension.

Theorems & Definitions (5)

• Definition 1: Wasserstein–gradient Lipschitz constant
• Theorem 1: Curvature-flattening convergence
• proof : Proof sketch
• Corollary 1: Linear convergence of distances
• Definition 2: Wasserstein–gradient Lipschitz constant