Rule Learning for Knowledge Graph Reasoning under Agnostic Distribution Shift

TL;DR

This paper introduces StableRule, a framework for logical rule learning in knowledge graphs that remains effective under distribution shifts, addressing a key limitation of existing methods by enhancing robustness in out-of-distribution scenarios.

Contribution

The paper proposes StableRule, a novel end-to-end framework combining feature decorrelation with rule learning to improve out-of-distribution generalization in knowledge graph reasoning.

Findings

StableRule outperforms existing methods on seven benchmark KGs.

It demonstrates superior robustness across diverse environments.

The approach effectively mitigates covariate shift effects.

Abstract

Logical rule learning, a prominent category of knowledge graph (KG) reasoning methods, constitutes a critical research area aimed at learning explicit rules from observed facts to infer missing knowledge. However, like all KG reasoning methods, rule learning suffers from a critical weakness-its dependence on the I.I.D. assumption. This assumption can easily be violated due to selection bias during training or agnostic distribution shifts during testing (e.g., as in query shift scenarios), ultimately undermining model performance and reliability. To enable robust KG reasoning in wild environments, this study investigates logical rule learning in the presence of agnostic test-time distribution shifts. We formally define this challenge as out-of-distribution (OOD) KG reasoning-a previously underexplored problem, and propose the Stable Rule Learning (StableRule) framework as a solution. StableRule is an end-to-end framework that combines feature decorrelation with rule learning network, to enhance OOD generalization in KG reasoning. By leveraging feature decorrelation, StableRule mitigates the adverse effects of covariate shifts arising in OOD scenarios, improving the robustness of the rule learning network. Extensive experiments on seven benchmark KGs demonstrate the framework's superior effectiveness and stability across diverse heterogeneous environments, highlighting its practical significance for real-world applications.