Flock: A Knowledge Graph Foundation Model via Learning on Random Walks
Authors: Jinwoo Kim
Xingyue Huang
Krzysztof Olejniczak
Kyungbin Min
Michael M. Bronstein
Seunghoon Hong
Ismail Ilkan Ceylan
Venue: ICLR 2026 Poster (2026)
Keywords: knowledge graphs
link prediction
knowledge graph foundation models
invariance
equivariance
random walks
Primary area: learning on graphs and other geometries & topologies
TL;DR: We present Flock, a knowledge graph foundation model (KGFM) that uses random walks to achieve probabilistic node-relation equivariance and overcome the limitations of previous KGFMs.

Abstract: We study the problem of zero-shot link prediction on knowledge graphs (KGs), which requires models to generalize to novel entities and novel relations. Knowledge graph foundation models (KGFMs) address this task by enforcing equivariance over both nodes and relations, which enables them to learn structural properties of nodes and relations that transfer to novel KGs with similar structure. However, the conventional notion of deterministic equivariance inherently limits the expressive power of KGFMs, as it prevents them from distinguishing relations that are structurally similar but semantically distinct. To overcome this limitation, we propose to leverage probabilistic node-relation equivariance, which preserves equivariance in distribution while using structured randomness to break symmetries at inference time. Building on this principle, we present Flock, a KGFM that iteratively samples random walks, encodes them into sequences, embeds them with a sequence model, and aggregates node and relation representations through learned pooling. Flock respects probabilistic node-relation equivariance and, crucially, is a universal approximator for isomorphism-invariant link-level functions over KGs. Empirically, Flock perfectly solves our new diagnostic dataset Petals on which current KGFMs fail, and achieves state-of-the-art performance on entity and relation prediction tasks across 54 KGs from diverse domains. Code is available at https://github.com/jw9730/flock.

Published as a conference paper at ICLR 2026
FLOCK: A KNOWLEDGE GRAPH FOUNDATION MODEL VIA LEARNING ON RANDOM WALKS
Jinwoo Kim Xingyue Huang?
Krzysztof Olejniczak" Kyungbin Min
Michael Bronstein Seunghoon Hong'
ismail ilkan Ceylan 34
KAIST Universsiy of Oxford
'TU Wien
*AITHYRA
ABSTRACT
We study the problem of zero-shot link prediction on knowledge graphs (KGs), which requires models generalize novel entities and novel relations. Knowl- edge graph foundation models (KGFMs) address this task by enforcing equivari- ance over both nodes and relations, which enables them to learn structural prop- erties of nodes and relations that transfer to novel KGs with similar structure. However, the conventional notion of deterministic equivariance inherently limits the expressive power of KGFMs, it prevents them from distinguishing relations that are structurally similar but semantically distinct. To overcome this limitation, we propose to leverage probabilistic node-relation equivariance, which preserves equivariance in distribution while using structured randomness to break symme- tries at inference time. Building this principle, present FLOCK, KGFM that iteratively samples random walks, encodes them into sequences, embeds them with sequence model, and aggregates node and relation representations through learned pooling. FLOCK respects probabilistic node-relation equivariance and, crucially, is a universal approximator for commrppism-invariaae link-level func- tions over KGs. Empirically, FLOCK perfectly solves our new diagnostic dataset PETALS on which current KGFMs fail, and achieves state-of state-of-the-an the-arr performance entity and relation prediction tasks across 54 KGs from diverse domains.
INTRODUCTION
Knowledge graph foundation models (KGFMs) (Galkin et al. 2024 Zhang et al. 2024 Huang et al., 2025) aim to infer missing links over novel knowledge graphs (KGs) that are not part of the training data or domains. This task requires generalization to both unseen nodes and unseen relation types. To achieve this, KGFMs learn node and relation invariants: structural properties of nodes and relations that are transferable across KGs even when their relational vocabularies differ. This inductive bias is formalized double equivariance (Gao al., 2023 equivariance under permutations of both entities and relations-ani used as a core design principle of current KGFMs.
Problem statement. this work, we challenge a funda- mental assumption dictated by strict equivariance in exist- ing KGFMs: structural isomorphism of relations implies semantic equivalence. Consider, for example, the KG in Figure 1 where the relations like and dislike are struc- turally isomorphic yet semantically opposite. Any KGFM that computes relation invariants is forced to assign the same representation to both like and dislike-losing the ability to distinguish entities with opposite relationships. This expressiveness limitation an architectural one and cannot be resolved through finetuning, further limiting the downstream use of existing KGFMs. This raises central question: how can we design KGFMs that are both expres- sive and have the right inductive bias for generalization?
Yoda
HanSolo
DarthVader
Jabba
Luke
Leia
Emperor
Chewbacca
Figure 1: A KG representing charac- ters' relationships in Star Wars movies. Blue arrows indicate like, red- dislike, and green arrows indicate friendWith.
"qqual contribution.

---

Published as a conference paper at ICLR 2026
Approach. We propose new approach for KGFMs that relies on probabilistic node-relation equivariance as inductive bias, instead of enforcing deterministic equivariance over nodes and rela- tions. Probabilistic equivariance relaxes the hard constraint that "structurally isomorphic relations must have identical representations", and requires only that the representations of structurally iso- morphic relations need to be equivalent in distribution over a model's ' stochastic predictive processes (Srinivasan & Ribeiro, 20200 Abboud et al. (021). This way, the model retains the inductive bias needed for generalizing across different KGs, while the stochasticity of each forward pass ensures that structurally identical but semantically distinct relations are assigned different representations.
Inspired by the success of models that learn probabillittic invariants via random walks (Perozzi et al, 2014 Grover & Leskoved 2016 Nikolentzos & Vazirgiannis, 2020 Kim et al. (2025), we introduce FLOOCK, a KGFM that inherently computes probabilistic node-relation invariants. Given poten- tially unseen KG and query. FLOCKK iteratively samples random walks over the KG based on the query, encoding both encountered nodes and relations with a recording protocol. To ensure the model can generalize to unseen entities and relation types, we anonymize all nodes and relations, enforcing that FLOCK only learn from their structural roles. These anonymized sequences are then fed into sequence processor, and the representations for each node and relation are aggregated via consensus protocol. Finally, we construct per-query features from the aggregated entity and relation embeddings and input them into binary classifier for link prediction.
Key findings and contributions. The design of FLOCK offers several key advantages over exist- ing KGFMs. First, it entirely abandons the conventional two-stage process of encoding relations and node representations via two separate networks, and does not rely on message passing at all, thereby avoiding the well-known expressivity limitations of MPNNs on KGs (Barceló et al. 2022 Huang et al. 2023 2025). Second, FLOCK is universal approximator (Proposition/4.] capable of approximating every link-level function on KGs of any bounded size. Finally, F1,CKK's architecture inherently respects probabilistic node-relation equivariance, enabling strong generalization. Our ex- periments on both entity and relation prediction validate this approach, demonstrating that FLOCK consistently outperforms ofte-of-theearr KGFMs on existing benchmarks. Our contributions are:
We highlight a limitation of existing KGFMs: their reliance on deterministic node-relation equivariance prevents them from distinguishing between structurally similar but semantically different relations, limiting their expressivity.
We propose to leverage probabilistic node-relation equivariance, a property for KGFMs that ensures equivariance only in distribution, an effective solution balancing the model expres- sivity and generalization.
We propose FLOCK,, KGFM that respects probabilistic node-relation equivariance. FLOCK replaces the conventional two-stage, message- -asssing paradigm with direct sequence encod- ing approach based on random walks, and is universal approximator of link-level functions.
We validate our approach on entity and relation prediction tasks across 54 diverse KGs, where FLOCK consistently achieves state-of-the-art performance. We further construct a synthetic dataset PETALS to confirm our theoretical results empirically.
All proofs are in B The code is available at
2 RELATED WORK
Link prediction and KGFMs. Early methods for inferring missing links in KGs relied on learned embeddings (Bordes et al., 2013: Sun al.) 2019 Balazevic al., 2019: Abboud al.) 2020]]
Schlichtkrull et al. 2018; Vashishth et al. 2020). operating in the transductive setting and incapable
of generalizing to unseen entities or relation types. Later MPNN-based methods based on the la-
beling trick (Teru et al. 2020) Zhang et al., 2021) or conditional message passing (Zhu et al.) [202]] 2023 Zhang & Yao 2022 Zhang et al. 2023b Huang et al., 2023) unlocked the node-inductive sce- nario, while remaining restricted to fixed relational vocabulary. KGFMs eliminate this restriction and enable node-relation inductive link prediction over both unseen nodes and relations (Geng et al, 2023), typically by first encoding relations and then nodes. Early examples are InGram (Lee et al. 2023) and ULTRA (Galkin et al. (2224), later extended by TRIX (Zhang et al. (024) into more ex-
pressive framework. KG-ICL. (Cui et al. 2024 achieved full inductiveness by combining in-context
2

---

Published as a conference paper at ICLR 2026
learning with node-relation tokenization. ISDEA (Gao al. 2023) and MTDEA (hou al. 2023
highlighted the benefits of node-relation equivariance, while MOTIF Huang et al. 2025) proposed a general KGFM framework with theoretical analysis of their expressive power. Our work advances the field with a stochastic KGFM that achieves invariance in probability to ensure generalization while being provably more expressive than the existing methods. Notably, FL.CK achieves univer- sality without any form of message passing, instead relying on random walks and sequence models encode both nodes and relations anonymously. This is distinct from prior stochastic KGFMs that rely on random initialization of message passing (Lee et al. 2023 Gao et al. 2023).
Random walks for graph representations. Random walks have been widely adopted in graph learning due to their simplicity and ability to gather context from neighborhoods. DeepWalk (Per-) ozzi et al.) 2014] and node2vec (Grover & Leskovec 2016) were among the first approaches, treating walks as analogues of natural language sentences and processing them with skip- gram models. Sub- sequent work has developed diverse neural architectures based random walks: Nikolentzzos & Vazirgiannis 2020) generates graph-level predictions using joint walks on direct products of graphs and their subgraphs, CRaWI. (Tönshoff al.] (02T) represents graph as collection of walks and processes them with a convolutional network, WalkLM (Tan et al. (023) samples walks from graphs with textual features and embeds them using a language model, RWNN (Kim et al.) 2025) and RUM (Wang & Cho 2024) anonymize walks and process them with sequence networks, and Walker (Chen et al. 2025 integrates walk encodings into message passing.
Probabilistic invariance. Neural architectures that enforce invariance to specific transformations often exhibit more stable training and improved performance (Bronstein al. 2021). but this induc- tive bias can reduce their expressivity by preventing the model from distingunshing non-equivalent inputs. In graph learning, this trade-off is exemplified by MPNNs, whose power is limited by the 1-WI. test (Xu et al. 2019 Morris et al. 2019). Randomization has emerged as a solution, en- hancing expressivity through techniques such noise injection (Abboud et al.] (022), vertex drop- ping (Papp et al. 2021), subgraph sampling (Bevilacqua et al., | /Zhang et al.] (023a). dynamic rewiring (Finkelshtein et al.] (2024), and random walks (Kim et al. [2025] Wang & Cho (024)) Despite their stochasticity, such methods can remain probabilisticall invariant, ensuring that equiv- alent inputs yield identical expected outputs, or even identical output distributions. In line with a prior work (Gao et al., 2023). we extend the notion of probabilistic invariance to KGs and prove that FLOCK satisfies invariance in distribution, and further clarify its usefulness in KG learning.
3 PRELIMINARY
Knowledge graphs. A knowledge graph (KG) is tuple G (V,E, R), where V denotes the set of entities (nodes), R the set of relation types, and E V x R x V the set of labeled edges (facts). A fact is written as (h, r,,) r) (or h 4 interchangeeabl) with E R and h, t E V. A (potential) link in G is any triple (h, r, 1) V x R XxVV regardless of whether it is present in F. We denote by R-1 the set of inverses of relations R. defined as {r-1 E R), and mean when writing (r 1r-1)-1 Further, let K,,,m be the space of knowledge graphs with n vertices and m relation types.
ssmorphism. An isomorphism between two knowledge graphs G E, and G R', such that
(V', E',R') is a pair bijections j = =( (T, where π + V' and & : R + fact (h,r, () belongs to E if and only if the fact j((h,rr.) r,))) = (T(h), ((r),((() belongs to E'.
Two KGs are isomorphic if such a mapping exists, in which case we write G ~ G¹.
Link invariance. In this work, focus on link- invariant functions. Let be a function assigning to each KG G = (V, E, R) E Knnm a map (G) : V x V + Rdd We say that w is link invariant if for every pair of isomorphic KGs G, G' E Kn every isomorphism (T. Φ) from G to G', and every link (h, r, t) in G, we have (G)) r, f)) =
(t)))
Probabilistic invariance. A stochastic KG model o can be viewed as function that takes a KG and returns a random variable (G)) Following Kim et al. (2025). we call yp invariant in probability
VG, G' E Kn,m:
G G'
((G) d ((GG),
i.e., the distributions of (GG) and ((G') are equal. In particular, this implies El(((G)) = El(((G)))
3

---

Figure 2: An overview. In each updating step, FLOCK (1) samples random walks on the KG (two walks indicated by red and teal, respectively), (2) anonymizes the encountered nodes and relations via a recording protocol (for each walk, nodes are anonymized as 1,2, ... and relations as a, 3, ...,, and (3) feeds the sequences in sequence processor compute node and relation representations. Published as a conference paper at ICLR 2026
1 ar' 2. B. 3.
[1 2
Random walks
Recording protocol
Sequence processor
Consensus protocol
(4) A consensus protocol then pools them back to the original KG's nodes and relations.
4 METHODOLOGY
We present FLOCK, a KGFM respecting probabilistic node-relation invariance. FLOCK random- ized function Xs(-) which takes as input KG G = (V, E, Rt) and a link prediction query We consider two types of queries: entity prediction q = (h, ?) and relation prediction q (h, ?,t). FLOCK outputs random variable ý ~ (o(G, q) which is suited for the task at hand. For entity pre- diction, it outputs y For relation prediction, it outputs y R [0, 1] such that link (h, r, t) can be evaluated by (r)).
[0. 1] such that a potential link (h, r, t) can evaluated by ÿ(t) e [0, 1].
At test time, we average multiple (P) independent stochastic predictions to produce the final output. This improves performance and reduces variance through an ensembling effect.
We describe the architecture of FLOCK in $4.1 focused on four main components, and then analyze its theoretical properties in $4.2 showing universality and probabilistic equivariance. An expansion on the model details can be found in SA
4.1 FLOCK
Internally, FLOCK has two lookup tables of hidden states, Rd for entities and Rd for relations, respectively. At each forward pass, it starts from trained initializations of these states (.)) = Vo and (-) = ro. and updates them iteratively v(0), p())) for S J. Each update is done residually using randomized function Ug.:
v(++)) == + Av,
(f+1) (6) + Ar,
(Av, update,, (v( (())).
The final hidden states v() V - Rd and r(), : R + Rd are then processed by a binary classifier head Rd 1] to produce the output ý which V [0, 1] R [0, 1] depending task.
We now describe the randomized updateg. We drop for brevity. It consists of four components:
1. Random walk algorithm produces random walks 7/,..., of length & on the input KG.
2. Recording protocol Ml ++ zj transforms each walk into a graph- aggnstic sequence.
3. Sequence processor So 25 + hy processes each sequence independently, outputting features.
4. Consensus protocol (hI:N; m::) r> (Av, Ar) collects features of all walks and decides state updates for each entity and relation type.
An overview is presented in Figure We note that w, fer and all deterministic, and the random walk the only source of stochasticity. We now discuss the design choice for each. For the ease of exposition, explain for entity prediction tasks q (h, ?), but relation prediction is similar.
Random walks. In FLOCK, random walks are central in two ways: they rewrite the connectivity of nodes and relations as sequences, and support generalization via probabilistic equivariance.
4

---

Published as a conference paper at ICLR 2026
Formally, the random walk algorithm produces random walks ",,,, 7n of length é on KG G. Each random walk 17 is a chain of random variables, written asın as:
E4 UE, E V,ro E R, (v.-1, Ts.
E,
where the underlying transition mechanism and are hyperparameters.
To support probabilistic equivariance, we ask the walk algorithm to be invariant in probability. We say 17 is invariant in probability if for any G in kn,mm with isomorphism (T,, Φ) from G H:
x(2) s(ra), TVV)) o(rz),
s(re), (T2) uoo 64
ug:
where to "i 2 uy and uo respectively. In such case, the isomorphism (π, 6) yields natural translation from walks in G to H.
ue follow the distributions of n(G, E) and mHH, E).
In FLOCK, we use simple random walk algorithm which we show be invariant in probability. Specifically, we use uniform walks with non-backtracking, with minor modifications to handle di- rected multi-edges of KGs. Despite the simplicity, we find that this choice works well in practice, consistent with findings of prior works (Tönshoff et al. [221 Kim et al. (025).
Under this choice, we diversify the starting locations of walks such that local context around the query and broad regions of the nodes and relations in a KG are both well-captured. Our diversi- fication strategy is as follows: given a base walk count n, for entity prediction queries (h, r, ?), we use 3n walks with three types of start locations. The first n walks start at query node h, capturing local context around the query; the second n walks start by traversing a random edge (v, 8, u) where is uniformly chosen relation, broadly capturing the relations of the KG including r; the last walks start at random nodes, broadly capturing various regions of the KG. For relation prediction queries (h, ?, t). we additionally start n walks at the tail node t, sampling total of 4n walks.
We lastly discuss how to choose the base walk count n. While this is a fixed hyperparameter Perain pretraining, we find that scaling it adaptively to input KG at test-time benefits size generalization. We thus propose test-fime adaptation of walk counts, and use:
IVI VIrainn
π Ttrain x harmonic mean
(I)
Eltrain
where |V Virainn E Elrsin are average numbers of nodes and edges in pretraining KGs, respectively. Intuitively, this scales proportionally to the sizes of test KGs relative to pretraining. In practice, we clamp n to the nearest power of 2 and limit its value an interval to avoid out-of-memory errors.
Recording protocol. While random walks provide basis for invariant sequence representations of KGs, two issues remain: (1) They reveal nodes U, and relations т, specific to each KG which obstructs transferability to unseen KGs. (2) They do not offer way to condition current states of entities v, relations r, and the query q = (h, r, ?) as often done in KGFMs via the labeling trick.
ny + Z, resolves this by transforming each walk into graph-agnostic sequence that only leaves structural information. Motivated by prior works on node anonymization for invariance (Kim et al. 2025 Wang & Cho. (024), we propose extension called node-relation anonymization: reserve separate namespaces for nodes and relations, respectively, and assign unique names in the order of discovery. For example, with 1, 2, for nodes and a, B,... for relations:
The recording protocol
1 2 8
Up
12
Up
1,
where (-)-1 marks direction of a relation. The protocol additionally employs simple conditioning on current states (v, r) and query q = (h,r,?3 completing the record z as follows:
1 - z = (1, (vo). (Mo)) a.(())))..))) (2, v(11), Ix(1)) (.r(ra))..
with indicator functions (-). (-) at h and r, respectively. As we will show, the recording protocol keeps node-relation invariance by hiding nodes and relations while leaving their structural roles.
5

---

Published as a conference paper at ICLR 2026
Sequence processor. Now that the recordings only encode structural information of KG, can safely process them with an arbitrary neural network fo : 2 + h without the risk of losing invariance. Since z are sequences, we choose sequence networks leverage their inductive bias. Specifically, we use bidirectional GRU (Cho et al. 2014) equipped with RMSNorm (Zhang & Sen- nnicc) 2019) and SwiGL.U feedforward network (Shazeer 2020), which provided robust results. To convert anonymizations into input feature vectors to the GRU, we use trainable embedding tables.
Given that fo a sequence network, it is convenient to interpret its output h as positionally aligned with each step of the walk 1 or record z. Specifically, for the example in Equation 2 we obtain:
fo z + h (Avo, an) Ari.bi (1v1, Ar2.b2
where Av,, Ar, E IR Axdx and a,,b, b. E R are the decoded outputs at each position using linear projections. Intuitively, , Ar, encode proposals of state updates for entities and relations by So- and ag, by encode respective confidences of /e for the proposed updates. This separation is useful due to the localized, pure-structure nature the recordings z. a random walk 17 densely visited cycle-like region and then terminated in a dangling manner, it is natural to assign more confidence the cycle-like region of the structural encodings h, and less confidence the dangling region.
Consensus protocol. After sequence processing, we are left with a handful of state update pro- posals b from So. that are positionally aligned with random walks 7/1y on KG. The consensus protocol e uses the information to decide final state updates Av V + Rd and Ar R + RR..
Since e can access how each Av, within hy is associated associated to relation E Rt) through the random walk 7n5 a simple way to form consensus is by finding all proposals (Av.) associated to each node v, and all (Ar.) associated to each relation and take averages of these proposals. The drawback is that uninformative proposals from e.g..
node v. E V (and how each år, is
dangling regions of walks are not directly suppressed, and can affect the state updates.
We can leverage the confidences ag, b, from fe to alleviate this issue. For each node v V or relation r E R, we first find all respective associated pairs {(Av,, a.)} or {(Ar,, b,]] of proposals and confidences, and compute multi-head boftmax-normalizae weighted average:
Av(u)= Σ exp(as) @ Av.
Exxp(a.)
Ar(r)
Σ exp(bb) Ar. Σ exp((b),,
where 3 and Ø are rowewise multiplication and division, respectively. Intuitively, this normalization induces competition between state update proposals, naturally leading to uninformative proposals being suppressed. Similar ideas are presented by Locatello et al. (2020).
Again, we can show that the consensus protocol does not operate in way specific to particular KGs, and hence retains node-relation equivariance.
4.2 THEORETICAL ANALYSIS
Expressivity. Following the notion of probabilistic expressivity introduced by) Abboud et al. (2021 we say that FLOCK model Xç is a universal approximator of link invariant functions over K,,,, for any link invariant s Kn,m + (V R x - [0, 1]) and any e, ô > 0, there exists a choice of the network parameters 0 and the length of the sampled random walks E, such that:
P(Ip(G)((h,rr)) r,,,, Xo(G, (h, r, ()) < 1-66 for all graphs G = (V, E, R) E K.,... and all links rh, E V x R.
Proposition 4.1. With a powerful enough sequence processor fe, the FLOCK framework described above is a universal appreximator of link invariant functions over Kmn,m all pairs (n, m)..
All proofs are in B To offer intuition behind the result, provide proof sketch.
Proof sketch. A sufficiently long random walk will cover all edges of the graph with high proba- bility. Then, from its anonymized version, assigning unique positional identifiers to every node and relation, one can reconstruct the input graph, up to isomorphism. Thus, with a sufficiently expressive sequence processor, FLOCK approximate any fin--invarian function.
Invariance. Despite the stochastic nature of our framework, beyond randomized node embeddings Abboud et al. 202T), FLOCK can be provably guaranteed satisfy probabilistic invariance:

---

Published as a conference paper at ICLR 2026
Proposition 4.2. Suppose that the walk sampling protocol rj is invariant in probability and both the recording protocol w and the consensus protocol с are invariant. Then, regardless of the choice of the deterministic sequence processor fo. the corresponding FLOCK model is invariant in probability.
Proof sketch. Since each of these components is invariant (in probability). and invariance of individ- ual component is preserved under composition, we have that FLOCK is invariant.
Moreover, the designs of FLOCK'S 's components provided earlier in this section satisfy the conditions of Proposition 4.2 Therefore, the suggested pipeline is indeed invariant in probability:
Proposition 4.3. Any FLOCK model with components as oullined in this section, and detailed in is invariant in probability.
5 EXPERIMENTS
We test FLOCK over a wide range of KGs for both entity and relation predictions, aiming to answer:
Q1. Can FLOCK approximate functions that existing KGFMs cannot?
Q2. How does FLOCK generalize to unseen entities and relations compared to existing KGFMs?
Q3. How does performance scale with the sizes of pretraining graph mix and test-time ensemble?
Q4. What is the impact of choices of each component on the behavior and performance of FLOCK?
and E and comparisons against current KGFMs augmented with noise injection in $ Experiment details and hyperparameters are in $
We further provide detailed scalability analysis in D
5.1 SYNTIETIC DATASET
Setup. To validate the limitations of KGFMs with node- relation equivariance (Q1), we construct a synthetic bench- mark PETALS. It contains 220 instances, each including: (1) a KG G = (V, F,, R) consisting of 'central' node 8, a 'stem' T V with query relation To, and multiple cyclic "petals', each 'colored' with different pair of relations R. {ro). (2) an entity prediction query (h, ro, ?) with h E {s} T, and (3) two candidate targets t and to from the same 'petal', located at the same distance from s. An example is in Figure 3 See for more details.
1/
Figure 3: A KG in PETALS. KGFMs with relational invariants must equate Fi and r2, thus predicting the same scores for dashed queries with ro.
PETALS is designed such that each instance always admits non-trivial automorphisms, meaning that swapping rela- tions occurring in the same 'petal" results in an isomorphic KG. Consequently, any model computing relation invari- ants will not be able to distinguish between potential links (s. ro. t1) and (s, Fo, tz). However, the samples constructed that these links not isomor- phic from the graph perspective, making them distinguishable for general link-invariant functions.
We say model solves instance if it can classify (s.ro_L) (1) as TRUE and (s.mn. (2) FALSE.
For empirical validation, we train ULTRA (Galkin et al. 2024), MOTIF(JF Path) (Huang et al. 2025
TRIX (Zhang al. 2024). and FLOCK from seratch and validate them the training instances.
Table 1: PETALS accuracies.
Results.
The results are provided in Table [ As expected, all existing KGFMs relying on learning deterministic relational in- variants fail to distinguish between the candidate target triplets completely, achieving 50% accuracy due to random guesses. In contrast, FLOCK succeeds all considered instances, displaying that, while remaining invariant probability, can differentiate between non-isomorphic links, even with isomorphic relations.
Model
PETALS
ULTRA
50%
MOTIF(FR
50%
TRIX
50%
FLOCK
100%

---

Published as a conference paper at ICLR 2026
Table 2: Average entity prediction MRR and Hits@ 10 over 54 KGs from distinct domains.
Inductive e, (23 graphs)
Inductive e (I8 graphs)
Transductive
Total Avg
Pretrained
Model
(13 graphs)
(54 graphs)
(3 graphs)
MRR H@10
MRR
H@10
MRR
H@10 ||
MRR
H@10
MRR H@10
ULTRA (zero-shot)
0.345
0.513
0.431
0,566
0.312
0.458
0,366
0.518
TRIX (zero-shot)
0.368
0,540
0,455
0.592
0.339
0.500
0.390
0.548
FLOCK (zero-shot)
0.369
0.554
0.456
0.604
0.340
0.509
0.391
0.560
ULTRA (finetuned)
0.397
0.556
0,440
0,582
0.379
0.543
0,408
0.562
0,407
0.568
TRIX (finetuned)
0.401
0.556
0,459
0.595
0.390
0.558
0.418
0.569
0.415
0.564
FLOCK (finetuned)
0.417
0.576
0.473
0.619
0.383
0.544
0.427
0.582
0.415
0.561
Table 3: Average relation prediction MRR and Hits@I over 54 KGs from distinct domains.
Inductive e, r
Inductive e
Transductive
Total Avg (54 graphs)
Pretrained
Model
(23 graphs)
(18 graphs)
(13 graphs)
(3 graphs)
MRR H@1 1
MRR
H@1
MRR
H@1
MRR
H@1
MRR H@l
ULTRA (zero-shot)
0.785
0.691
0.714
0.590
0.629
0.507
0.724
0.613
TRIX (zero-shot)
0.842
0.770
0.756
0.611
0.752
0.647
0.792
0.687
FLOCK (zero-shot)
0.898
0.846
0.864
0.782
0.873
0.813
0.881
0.817
ULTRA (finetuned)
0.823
0.741
0.716
0.591
0.707
0.608
0.759
0.659
0.876
0.817
TRIX (finetuned)
0.850
0.785
0.759
0.615
0.785
0.693
0.804
0.706
0.879
0.797
FLOCK (finetuned)
0.929
0.889
0.887
0.808
0.897
0.844
0.907
0.851
0.977
0.959
5.2 ENTITY AND RELATION PREDICTION OVER KNOWLEDGE GRAPIIS
Setup. To answer Q2, we follow the protocol of Galkin et al. 2024); Zhang et al. 2024 and pretrain FLOCK FB15k-25 237 (Toutanova & Chen 2015). WNI (Dettmers al. 2018). and CoDEx Medium (Safavi & Koutra, 2020). We then evaluate its zero-shot and finetuned inference performance with the test set of 54 KGs (see for details). These KGs are extracted from diverse domains across three settings: inductive on nodes and relations (Inductive e, r), inductive on nodes (Inductive e), and transductive. Note that these settings differ only during finetuning setup: in zero-shot setup, all entities and relations are unseen. We choose otate-of-the-art KGFMs ULTRA (Galkin et al. 2024) and TRIX (Zhang et al., 2024) as baselines, as they are pretrained on the same KGs, to ensure a fair comparison. For evaluation, we use the filtered ranking protocol (Bordes et al.) 2013 reporting mean reciprocal rank (MRR) and Hits 10 for entity prediction, and Hits I for relation prediction, as some KGs have fewer than 10 relations. Per-dataset results are in |
Entity prediction. Table 2 shows the entity prediction results. In the zero-shot setting, FLOCK consistently outperforms ULTRA and TRIX, demonstrating strong generalization across diverse do- mains. Notably, on Metafam (Zhou et al., 2023), dataset designed to challenge models with conflicting and compositional relational patterns, FLOCK roughly doubles MRR over ULTRA and achieves x 40% gain in MRR over TRIX in the zero-shot setting. We find that FLOCK distinguishes structurally similar but semantically confficting relations while ULTRA fails ($ explaining the gain. These findings align with our hypothesis that probabilistic equivariance improves expressivity without sacrificing generalization. In the finetuning setting, we observe a similar pattern: FLOCK maintains consistent improvement over all datasets except transductive ones, where KGs are gen- erally larger. We hypothesize that the gap stems from walk coverages. Unlike ULTRA and TRIX, whose message passing guarantees full neighborhood coverage around the queried node, FLOCK relies on sampled walks that may not fully cover target nodes of interest. We find that FLOCK favors sparse KGs (sH), consistent with this hypothesis as random walks cover sparse graphs faster.
Relation prediction. Table 3 presents the relation prediction results. FLOCK substantially outper- forms all existing KGFMs across all categories in the zero-shot setting, achieving an 11.2% relative improvement in MRR compared to the best baseline TRIX. FLOCK shows further performance boost of 12.8% in the finetuned setting. We attribute this huge gain to FLOCK's joint encoding of entities and relations during the updating step via the sequence encoder, while existing KGFMs, ULTRA and TRIX, have separate update steps for entities and relations. This joint update mechanism yields more holistic representations of both entities and relations with minimal information loss.
8

---

Published as a conference paper at ICLR 2026
Averaged MAR
0.405
0.400
Averaged
0.400
0.395
0.395
0.390
0.390
10
22
14
16
Training Graphs
Test Ensemblee
(a) Zero-shot MRR vs. Hipretraining graphs.
(b) Zero-shot MRR vs. #ensembled predictions.
Figure Pretraining and test-time scaling of FLOCK on 41 inductive KG datasets.
Table Ablation study adaptive test-time walks with zero-shot entity prediction task. We show the average number of entities [V, triples EE, base walks n, MRR, and Hits@ 10.
Statistics
FLOCK
FLOCK w/o Adap.
Dataset split
[V]
E
71
MRR
Hits 10
Y
MRR
Hits 10
Inductive e,
5,303
10,838
28.40
0.369
0.554
128
0.357
0.551
Inductive e
7,578
29,090
18.08
0,456
0.604
128
0.441
0.596
Transductive
47,810
387,491
214.15
0.340
0.509
128
0.334
0.493
5.3 SCALING ANALYSIS
Size of pretraining graph mix. To assess whether FLOCK benefits from more pretraining graph and data (Q3), we follow the setup of Galkin et al. (202)). and pretrain FLOCK on increasing number of KGs. We then evaluate them on all 41 inductive benchmarks for a fair comparison. We present the detailed pretraining graph mix in Table 16 As shown in Figure 4a FLOKK''s 's general- ization improves consistently as the number of pretraining KGs increases, exhibiting clear scaling behavior, which is core characteristic of being foundation model.
Number of ensembled predictions. To assess how test-time ensemble size P affects performance (Q3), we take the pretrained FLOCK and run zero-shot entity prediction on 41 inductive KGs by increasing the number of ensembled passes. As shown in Figure 4b performance improves from 1 to 8 passes and then begins to plateau beyond 12. This indicates clear scaling behavior: larger ensembles provide a more accurate estimate of the underlying node and relation distributions.
5.4 ABLATION STUDIES
Setup. To understand the impact of each design choice on the performance and behavior of FLOCK (Q4), we conduct series of ablation studies spanning random walks, sequence processor and the consensus protocol, in the entity prediction task in the zero-shot setting.
Adaptive test-time walks. Recall that we employ test-time adaptation of walk counts, which adap- tively selects the base walk count n based on the graph size, computed via the harmonic-mens mean rule shown in (T) during inference. Table 4 compares this adaptive setting with a fixed setting that uses 128 base walks per sample for all datasets, matching the pretraining setup (Naiin 128). As expected, the average selected base count n is smaller on both inductive splits and larger on the transductive split, yet the adaptive mechanism improves performance across all settings. This is consistent with the intuition that adaptive n scales up walks on larger KGs to improve coverage while allocating fewer walks smaller KGs to reduce redundant visits; FLOCK maintains compa- rable visiting rates and coverage to those seen during pretraining, thereby producing representations closer to the pretraining distribution and resulting in consistent performance gains.

---

Published as a conference paper at ICLR 2026
Table Detailed ablation study with zero-shot entity prediction task. For the transductive split, con- sidering resource limits, we test NELL995, NELL23K, WDsinger, ConceptNet1008 and YAGO310.
Inductive e, (23 graphs)
Inductive (18 graphs)
Transductive
Total Avg (46 graphs)
Model
(5 graphs)
MRR
Hits@ 10
MRR
Hits@ 10
MRR
Hits@ 10
MRR
Hits@ 10
FLOCK (& 128)
0.369
0.554
0.456
0.604
0,360
0,542
0.395
0.567
w/o non-backtracking
0.370
0.549
0.456
0.605
0.334
0.499
0.386
0.551
f = 64
0.372
0.556
0.459
0.606
0.351
0,534
0.394
0,565
t 256
0.360
0.548
0.458
0.605
0.338
0,508
0.385
0,553
w/o diverse starts
0.360
0,539
0.448
0.596
0.319
0,488
0.385
0.553
transformer fo
0.356
0.542
0.410
0.591
0.312
0.477
0.359
0.537
w/o weighted consensus
0.351
0.526
0.448
0.593
0.361
0.515
0.387
0.545
Non-backtracking walks. FLOCK employs non-backtracking uniform random walk, which has an effect of faster exploration and coverage of distant regions (Alon et al.) 2007). In Table[5] we com- pare this with uniform walk that may backtrack and hence slower in global exploration. While non-backtracking does not alter results much on inductive splits, it significantly improves perfor- mance on the transductive split. This is consistent with the idea that improving coverage especially benefits performances on large KGs, which FLOCK achieves via non-backtracking.
Walk lengths. FLOCK uses random walks of length & 128, a choice made by finding the lowest E reliably visiting target node and relation on various KGs. Table 5 compares this with shorter and longer walks by a factor of two. As expected from coverage, shorter walks show degraded results on the transductive split with large KGs. Longer walks overall worse, which is explained by higher learning complexity of the sequence processor that has small hidden dimension (64) for scalability. FLOCK finds balance of coverage and lcarnability, achieving robust results on diverse KGs.
Diverse starting locations of walks, We recall that FLOCK uses a diversification strategy of starting locations of walks, with walks from the query node, walks from random relation, and n walks from random node, adding up to 3n walks capturing both local context near query and global infor- mation of KG. Table compares this against all 3n walks starting at the query node. As expected, this causes degradations on all splits, showing the benefit of using both local and global information.
Sequence processor. FLOCK uses sequence processor fa with bidirectional GRU. In Table 5 compare this against a transformer fe with a similar SwiGLU-RMSNon architecture (Dubey et al.) 2024) and parameter count. This alternative does not deliver good results, which explained by the restrictions on model scales that are enforced to scale to large KGs. FLOCK benefits from reasoning efficiency of GRU in limited parameter regime, gaining good performance and scalability together.
Consensus protocol. FLOCK uses softmax-weighter averaging to pool sequence processor outputs into state updates for nodes and relations, under the intuition that this can suppress uninformative proposals from the sequence processor better than simple, unweighted averaging. Table [3 provides a comparison, showing that weighted consensus outperforms the unweighted counterpart. This ver- ifies our intuition on how the design of the consensus protocol strengthens FLOCK.
CONCLUSIONS
We introduced FLOCK, a knowledge graph foundation model that respects probabilistic node- relation equivariance. FLOCK iteratively samples query-conditioned random walks, records encoun- tered nodes and relations via recording protocol, and relies on a sequence processor and consensus protocol to obtain node and relation representations. We evaluate FLOCK on 54 KGs across different domains for both entity and relation prediction, demonstrating superior zero-shot and finetuned per- formance. We further construct synthetic dataset PETALS to validate our theoretical findings. One limitation is scalability ($E): ensuring coverage of the sampled random walk in large KGs requires many long walks, which can quickly become computationally infeasible. A future direction is to develop approximation strategies (Chamberlain et al. 2023 Lącki et al. 2020) that reduce the cost of random walk sampling while retaining FLOCK'S downstream performance. Another avenue for future work is studying the families of approximable functions when the walk lengths are restricted, for example based connections subgraph-based reconstructions (Cotta et al. 2021).
10

---

Published as a conference paper at ICLR 2026
ACKNOWLEDGMENT
This work was in part supported by the National Research Foundation of Korea BRS-2024-00351212 and RS-2024-000331665) and the Institute Information & Communications Technology Planning & Evaluation (ITTP) (RS-2024-00509279, RS AS-2202-11220926, RS-2022-I1220955, and RS-2019- 11190075) funded by the Korean government (MSIT). Computational resources were part pro- vided by the "HPC support" project funded by MSIT and NIPA. MB is supported by EPSRC Turing AI World-Leading Research Fellowship No. EP/X040062/1 and EPSRC AI Hub on Mathematical Foundations of Intelligence: An "Erlangen Programme" for AI No. EP/Y028872/1.
ETHICS STATEMENT
This work introduces a probabilistic framework for knowledge graph foundation models, aiming to improve the generalization in zero-shot link prediction. Our contributions are methodological and theoretical, with evaluations performed on publicly available benchmarks and a synthetic dataset designed to validate theoretical results. We do not anticipate any direct ethical risks associated with this approach. We acknowledge and adhere to the ICLR Code of Ethics
REPRODUCIBILITY STATEMENT
We make every effort to ensure the reproducibility of the experiments in our paper. We release a codebase at https:: /ttpt/ggttb...t//e wttp://wom/ww//////// Flcck with training and evaluation scripts for FLOCK, including pretraining scripts and checkpoints on FB15k-237, WN18RR, and CoDEx Medium, evaluation over 54 KGs, and the synthetic dataset PETALS generator. All architectural details needed reimplement the method, including the random-walk sampler, recording proto- col, sequence processor, and consensus protocol, are specified $A and our theoretical claims are supported with complete proofs in B We additionally include further experimental details in $
REFERENCES
Ralph Abboud, ismail Ilkan Ceylan, Thomas Lukasiewicz, and Tommaso Salvatori. Boxe: box embedding model for knowledge base completion. In NeurIPS, 2020. (page
Ralph Abboud, ismail ilkan Ceylan, Martin Grohe, and Thomas Lukasiewicz. The surprising power of graph neural networks with random node initialization. IJCAI, 2021. (pages
6
Noga Alon, Itai Benjamini, Eyal Lubetzky, and Sasha Sodin. Non-backtracking random walks mix faster. Communications Contemporary Mathematics, 9(4)):5856633, 2007. (page (0)
Ivana Balazevic, Carl Allen, and Timothy Hospedales. Tucker: Tensor factorization for knowledge
graph completion. In EMNLP-IJCNLP.. 2019. (page
Pablo Barceló, Mikhail Galkin, Christopher Morris, and Miguel Romero. Weisfeiler and leman go relational. In LoG, 2022. (page 2
Beatrice Bevilacqua, Fabrizio Frasca, Derek Lim, Balasubramaniam Srinivasan, Chen Cai, Gopinath Balamurugan, Michael M Bronstein, and Haggai Maron. Equivariant subgraph aggregation net- works. In ICLR. 2022. (page
Antoine Bordes, Nicolas Usunier, Alberto Translating embeddings for modeling multi-relational data. NIPS, 2013. (pages g
Jason Weston, and Oksana Yakhnenko.
Michael M Bronstein, Joan Bruna, Taco Cohen, and Petar Veličković. Geometric deep learning: Grids, groups, graphs, geodesics, and gauges. arXiv preprint arXiv:2104.13475 2021. (page
Benjamin Paul Chamberlain, Sergey Shirobokov, Emanuele Rossi, Fabrizio Frasca, Thomas Markovich, Nils Yannick Hammerla, Michael M Bronstein, and Max Hansmire. Graph neural networks for link prediction with subgraph sketching. In ICLR, 2023. (page 10)

---

Published as a conference paper at ICLR 2026
Dexiong Chen, Till Hendrik Schulz, and Karsten Borgwardt. Learning long range dependencies graphs via random walks. In ICLR, 2025. (page B)
Yihong Chen, Pasquale Minervini, Sebastian Riedel, and Pontus Stenetorp. Relation prediction as an auxiliary training objective for improving multi-relational graph representations. In 3rd Conference on Automated Knowledge Base Construction, 2021. (page 32
Kyunghyun Cho, Bart Van Merriënboer, Dzmitry Bahdanau, and Yoshua Bengio. On the properties of neural machine translation: Encoder-decoder approaches. arXiv preprint 14X9.
2014. (pages 6 16)
Leonardo Cotta, Christopher Morris, and Bruno Ribeiro. Reconstruction for powerful graph rep-
resentations. Advances in Neural Information Processing Systems, 34:1713-1726, 2021. (page
10
Yuanning Cui, Zequn Sun, and Wei Hu. A prompt-based knowledge graph foundation model for universal in-context reasoning. In NearlPS, 2024. (page 2
Tim Dettmers, Minervini Pasquale, Stenetorp Pontus, and Sebastian Riedel. Convolutional 2D
knowledge graph embeddings. AAAI, 2018. (pages
32
Boyang Ding, Quan Wang, Bin Wang, and Li Guo. Improving knowledge graph embedding using simple constraints. In ACL, 2018. (page 32
Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Amy Yang, Angela Fan, Anirudh Goyal, Anthony Hartshorn, Aobo Yang, Archi Mitra, Archie Sravankumar, Artem Korenev, Arthur Hinsvark, Arun Rao, Aston Zhang, Aurélien Rodriguez, Austen Gregerson, Ava Spataru, Baptiste Roz- ière, Bethany Biron, Binh Tang, Bobbie Chern, and et al. The llama 3 herd of models. arXiv, abs/2407.21173, 2024.
(page TOp
Ben Finkelshtein, Xingyue Huang, Michael M Bronstein, and Ismail Ilkan Ceylan. Cooperative graph neural networks. In ICML, 2024. (page B)
Luis Antonio Galárraga, Christina Teflioudi, Katja Hose, and Fabian Suchanek. AMIE: Association rule mining under incomplete evidence in ontological knowledge bases. In WWW, 2013. (page
Mikhail Galkin, Etienne Denis, Jiapeng Wu, and William L. Hamilton. Nodepiece: Compositional and parameter-efffcient representations of large knowledge graphs. In ICLR, 2022. (page 32
Mikhail Galkin, Xinyu Yuan, Hesham Mostafa. Jian Tang, and Zhaocheng Zhu. Towards foundation
models for knowledge graph reasoning. In ICLR, 2024. (pages 2
32
Jianfei Gao, Yangze Zhou, Jincheng Zhou, and Bruno Ribeiro. Double equivariance for inductive link prediction for both new nodes and new relation types. In arXiv, 2023. (pages 30)
Yuxia Geng, Jiaoyan Chen, Jeff Z. Pan, Mingyang Chen, Song Jiang, Wen Zhang, and Huajun Chen. Relational message passing for fully inductive knowledge graph completion. In ICDE, 2023. (page()
Aditya Grover and Jure Leskovec. node2vec: Scalable feature learning for networks. In KDD, 2016.
(pages 3
Daniel Scott Himmelstein, Antoine Lizce, Christine Hessler, Leo Brueggeman, Sabrina Chen, Dexter Hadley, Ari Green, Pouya Khankhanian, and Sergio E Baranzini. Systematic integration of biomedical knowledge prioritizes drugs for repurposing. Elife, 2017. (page
Xingyue Huang, Miguel Romero Orth, ismail ilkan Ceylan, and Pablo Barceló. A theory of link prediction via relational weisfeiler-leman on knowledge graphs. NeurlPS, 2023. (pages 32
Xingyue Huang, Pablo Barcelo. Michael M. Bronstein, ismail ilkan Ceylan, Mikhail Galkin, Juan L Reutter, and Miguel Romero Orth. How expressive are knowledge graph foundation models? In ICMI,, 2025. (pages
12

---

Published as a conference paper at ICLR 2026
Jinwoo Kim, Olga Zaghen, Ayhan Suleymanzade, Youngmin Ryou, and Seunghoon Hong. Revisit- ing random walks for learning on graphs. In ICLR, 2025. (pages
32
Jaejun Lee, Chanyoung Chung, and Joyce Jiyoung Whang. Ingram: Inductive knowledge graph embedding via relation graphs. In ICML, 2023. (pages 30 32
Shuwen Liu, Bernardo Grau, Ian Horrocks, and Egor Kostylev. Indigo: Gnn-based inductive knowl- edge graph completion using pair-wise encoding. In NeurIPS, 2021. (page 32
Francesco Locatello, Dirk Weissenborn, Thomas Unterthiner, Aravindh Mahendran, Georg Heigold, Jakob Uszkoreit, Alexey Dosovitskiy, and Thomas Kipf. Object-centric learning with slot atten- tion. In NeurIPS, 2020. (page ④
Xin Lv, Lei Hou Xu Han, Juanzi Li, Zhiyuan Liu, Wei Zhang, Yichi Zhang, Hao Kong, and Suhui Wu. Dynamic anticipation and completion for multi-hop reasoning over sparse knowledge graph. In EMNLP, 2020. (page 32
Farzanch Mahdisoltami, Joanna Asia Biega, and Fabian M. Suchanek. Yago3: A knowledge base from multilingual wikipedias. CIDR, 2015, (page 32)
Chaitanya Malaviya, Chandra Bhagavatula, Antoine Bosselut, and Yejin Choi. Commonsense knowledge base completion with structural and semantic context. In AAAI, 2020. (page 32
Christopher Morris, Martin Ritzert, Matthias Fey, William L.. Hamilton, Jan Eric Lenssen, Gaurav Rattan, and Martin Grohe. Weisfeiler and Leman go neural: Higher-order graph neural networks In AAAI, 2019. (page
Giannis Nikolentzos and Michalis Vazirgiannis. Random walk graph neural networks. In NeurIPS, 2020. (pages
Pál András Papp. Karolis Martinkus, Lukas Faber, and Roger Wattenhofer. Dropgnn: Random dropouts increase the expressiveness of graph neural networks. In NewrlPS, 2021. (page 3
Bryan Perozzi, Rami Al-Rfou, and Steven Skiena. Deepwalk: Online learning of social representa- tions. In KDD, 2014. (pages
Tara Safavi and Danai Koutra. CoDEx: A Comprehensive Knowledge Graph Completion Bench-
mark. In EMNLP, 2020. (pages 8 32
Michael Sejr Schlichtkrull, Thomas N. Kipf, Peter Bloem, Rianne van den Berg, Ivan Titov, and Max Welling. Modeling relational data with graph convolutional networks. In ESWC, 2018.
(page 2
Noam Shazeer. Glu variants improve transformer. arXiv preprint arXiv:2002.05222 2020. (pages
16
Balasubramaniam Srinivasan and Bruno Ribeiro. On the equivalence between positional node em- beddings and structural graph representations. In ICLR, 2020. (page
Zhiqing Sun, Zhi-Hong Deng, Man--Ynndu Nie, and Jian Tang. Rotate: Knowledge graph embedding by relational rotation in complex space. In ICLR, 2019. (pages 33)
Yanchao Tan, Zihao Zhou, Hang Lv, Weiming Liu, and Carl Yang. Walklm: uniform language model fine-tuning framework for attributed graph embedding. In NeurlPS, 2023. (page 3
Komal K. Teru, Etienne G. Denis, and William L. Hamilton. Inductive relation prediction by sub- graph reasoning. In ICML, 2020. (pages 32)
Jan Tönshoff, Martin Ritzert, Hinrikus Wolf, and Martin Grohe. Walking out of the weisfeiler leman hierarchy: Graph learning beyond message passing. TMLR, 2021. (pages
Kristina Toutanova and Danqi Chen. Observed versus latent features for knowledge base and text inference. In Workshop on Continuous Vector Space Models and their Compositionality, 2015.
(pages 8 32
13

---

Published as a conference paper at ICLR 2026
Shikhar Vashishth, Soumya Sanyal, Vikram Nitin, and Partha Talukdar. Composition-baser multi- relational graph convolutional networks. In ICLR, 2020. (page
Yuanqing Wang and Kyunghyun Cho. Non-convolutional graph neural networks. NeurIPS, 2024.
(pages 5 32
Wenhan Xiong, Thien Hoang, and William Yang Wang. Deeppath: A reinforcement learning method for knowledge graph reasoning. In EMNLP, 2017. (page 32
Keyulu Xu, Weihua Hu, Jure Leskovee, and Stefanic Jegelka. How powerful graph neural networks? ICLR, 2019. (page 3
Biao Zhang and Rico Sennrich. Root mean square layer normalization. Advances in neural infor- mation processing systems, 32, 2019. (pages 6 16)
Bohang Zhang, Guhao Feng, Yiheng Du, Di He, and Liwei Wang. A complete expressiveness hierarchy for subgraph gnns via subgraph weisfeiler-lehman tests. In ICML, 2023a. (page 3}
Muhan Zhang, Pan Li, Yinglong Xia, Kai Wang, and Long Jin. Labeling trick: A theory of using graph neural networks for multi-node representation learning. NeurIPS, 2021. (page 2
Yongqi Zhang and Quanming Yao. Knowledge graph reasoning with relational digraph. In WebConf, 2022. (page
Yongqi Zhang, Zhanke Zhou, Quanming Yao. Xiaowen Chu, and Bo Han. Adaprop: Learning adaptive propagation for graph neural network based knowledge graph reasoning, In KDD, 2023b.
(page 2
Yucheng Zhang, Beatrice Bevilacqua, Mikhail Galkin, and Bruno Ribeiro. TRIX: A more expressive model for zero-shot domain transfer in knowledge graphs. In LoG, 2024. (pages
Jincheng Zhou, Beatrice Bevilacqua, and Bruno Ribeiro. A multi-task perspective for link prediction with new relation types and nodes. In NeurlPS GLFrontiers, 2023. (pages g 30 32
Zhaocheng Zhu, Zuobai Zhang, Louis-Pascal Xhonneux, and Jian Tang. Neural bellman-ford net-
works: general graph neural network framework for link prediction. In NeurlPS, 2021. (pages
32
Zhaocheng Zhu, Xinyu Yuan, Mikhail Galkin, Sophie Xhonneux, Ming Zhang, Maxime Gazeau, and Jian Tang. * scalable path-based reasoning approach for knowledge graphs. In NeurlPS, 2023. (page 2
David Zuckerman. On the time to traverse all edges of graph. Inf. Process. Lett., 1991. (page 32)
Jakub Lęcki, Slobodan Mitrović, Krzysztof Onak, and Piotr Sankowski. Walking randomly, mas- sively. and efficiently. In Proceedings of the 52nd Annual ACM SIGACT Symposium on Theory of Computing, pp. 364-377, 2020. (page 10
14

---

Published as a conference paper at ICLR 2026
A
METHODOLOGY DETAILS
In this section, we expand on the deseriptions individual components FLOOK summarized in the random walk algorithm, the recording protocol, the sequence processor, and the consensus protocol.
$
A.1
UNIFORM RANDOM WALK
Let G (V, E, R) be knowledge graph, and let é be the length of random walks. For each node v E V, we will denote by (v) the set neighbors v:
((v) (w V 3r E R.(w,r, w)) E EV r,,, E E}
and by E(v, w), the set relational edges from v to or (allowing for the inverse direction):
F(u, m) ((v,r, E R (n) x [w) (1, F}
((ee,r") E R x {u) x {w) (u,r,e) E E}
where R-1 is the set symbolizing the inverses of relation types in R. The uniform random walk with backtracking n(G,f) of length & over G. represented as:
Vo _h Vi Re, Re Ve
is a second-order Markov process that follows the rules:
0
if v and W > 1 if and No {u}
1
P(V+2 = V4+1 w, V, =)
if + and E No.. if $ No if fr(( u) E E(w,a) otherwise
(3)
0
P(R = VJ+1 w, V, u)
E
for all i 2 0,j 2 1. Intuitively, at each step of the walk, we first select neighbor (except for the vertex chosen step ago) of the current node uniformly random (disregarding multi-edges and edge directions), and then sample an edge between these two nodes uniformly at random. If the current node has only one neighbor, we are forced to return it.
The initial conditions depend on the selected scenario. Given query g (h,r,?) over G, we can deseribe the process selecting the first step Vo Ri Vi as setting either (each with probability }):
Vo h and selecting the first step uniformly at random described above, meaning:
if E A,
P(V1 = v Vo = h)
NW.
i Wr.
ifr h) E(m,h)
P(R = Vi w)
((w.))
otherwise
setting R1 r and selecting V, Rg Vi uniformly at random from edges with type choosing Vo uniformly at random, and then sampling the first step at random as well:
P(Vo w)
if E N
P(V Vo v)
if $ No
if if v) E E(w.e)
P(R = V 8, Vo w)
E
otherwise
For the relation prediction objective, we add one more scenario. similar to the first one described above, but substituting Vo instead. For that problem, each scenario is chosen with probability 1
15

---

Published as a conference paper at ICLR 2026
A.2 RECORDING FUNCTION
Given a KG G = (V, E, Rt), query q (hg, Tay ?), a walk fj - vo _ e1 2 "4 of length È over G, and set of embeddings of nodes V and of relations R, recording function first splits the walk into a sequence of L +-1 steps:
(ro, vo), (7i, v1).... (re,
with ru = rer being a special marker for no relation. Each step (ra, 0/) is transformed into a 7-tuple:
S: (idy (iy f),, (r n), dir,, = &,, v(e)).((r))
where:
idv( and (ds(r;; fj) are the anonymized id's of the node us and relation 7i, evaluated as:
idy (Ux; r7) = arg min (+ = vz]
(ds( = arg min [rx = r2 V r. = V [
dir; denotes the direction in which we follow the edge. We set dir; = if vi E R (the edge is traversed from head to tail) and dir, if E R- (the edge is taken in the reverse direction).
be,=h, and br,-r, By and if the relation is either the queried relation rg its inverse r,
binary flags representing whether the current node Us is the query head
v(v;), r(r:) are the embeddings of o and ri respectively.
The output of w for given G, g, v, is then:
w(7; G. q, v,r) = (So. .,)
A.3 SEQUENCE PROCESSOR
Once the sampled walks anonymized by the recording protocol w, the output for each walk ij:
w(i;; G, (,G,v,vr) = (So. S
passed through the sequence processor fo. parametrized by the following modules:
Ap, A, E (f++)xxd embedding tables for anonymized vertices and relations, respectively,
D R2xd; look-up table for the direction embedding.
Q Q E 2xd, embedding tables for the binary query labels.
V, R. : Rd + Rd: linear maps applied to the passed embeddings of vertices and relations,
Ω2 a bi-directional GRU (Cho et al. 2014) cell equipped with RMSNorm (Zhang & Sennrich
2019) and SwiGLU (Shazeer 2020 activation function.
For each step, encoding S; of the form:
(idy (1,; n)), id (re ri.), dir,, d Sn.=,, v(ui),r(rs)
we evaluate the processed embedding c; of S; as sum of the corresponding encoded components:
C; = A,, (vv; m):)) + A, (idn(ra))) + D(dir;)
+ +n (6m, + Q (5m.=r,) + V(v(e:) + R(r(r:))
These are then passed to the GRU cell Q which fuses the features across the whole walk and produces multi-hean head embeddings of vertices and relations, well the associated weights:
, a (T) = 22([co. C1. ,I....,c.))
where s (2)
S E ((++))x and a (V) a (R) ()+ Stacking all N them gives us the final output of the sequence processor.
16

---

Published as a conference paper at ICLR 2026
A.4 CONSENSUS PROTOCOL
Given walks TANN over = (V, E, R) and the outputs s(V,s((R) a(V), a of the sequence pro- cessor, the consensus protocol aggregates the signal for each node by evaluating weighted over the appearances of this node across the walks. More precisely, for each node v E V, we find all pairs of indices (i, such that the jth node visited in fJ was v, and concatenate the weighted sums of embeddings produced by each head, with weights exponentially proportional to the scores (V);
(V)
(V) Sij.k
exp
(1y) =
Av(u) =
(V) (..k)
1=1
exp
(1yy m
Similarly, aggregate the encodings for relations, considering their occurrences in both directions:
Σ
exp
a (R)
(
(R) 1
Ar(r) = (r
)
(R) (a..1)
exp
((rr))(
In both formulas above, + denotes concatenation.
Additionally, we say that a consensus protocol invariant if for any pair of isomorphic KGs G = (V, E. R) and H = (V','', EE, R'). any isomorphism
(T, Φ) from G to H. any list of embeddings I with h, E Rdd and any sequence of sampled walks f).. over G, the outputs
(Av, Ar) c(hI.N m.:)
(Av', Ar') c(hI:N. ((TT:N))
satisfy:
Av(u) Av((T(V))
V
Ar(r) Ar(()(r)
Y
17

---

Published as a conference paper at ICLR 2026
B PROOFS
B. EXPRESSIVITY
The main proposition of this section formalizes the fact that FLOCK can approximate any link- invariant function over fixed-size knowledge graphs in probability. Intuitively, when the length of the sampled walks becomes higher, the probability of single walk witnessing all the edges grows to 1. Once walk visits all the edges, sufficiently powerful sequence processor can derive the whole graph structure from its anonymized representation, recreating the graph in its entirety, up to isomorphism. Then, the processor
return the value of the approximated function for that graph.
We start by showing that the edge cover time Cg(-) of graphs in Kn,m is bounded:
Lemma B.1. Let G E Ken,forr some n. m. The edge cover time Cg(G) of G. using the algorithm from J 1il is finite.
Proof. Let = (V. E, R) E Kknm be graph. For any edge e E E and any vertex E let H (e) denote the expected number of steps of the random walk algorithm r described in $A.1 Then, the edge cover time CE(G) of with y), i.c. the expected number of steps that needs to take before visiting every edge in G, is bounded above by:
Cg(G) S Σ max He(e) S m max (e)
eEV
BE
CE
Indeed, consider the event of visiting all these edges in order 61,..., C.:
Cg(G) E #steps to visit all el...., Cm]
< E #steps to visit e1, then C2,..., then eml
E[4steps to visit ei] +
Σ E [l#steps to visit 6 starting from h; or 4.]
m
max H (e)e)
max
H (ei+1), H4, (e(ta)
vEV
m
max (e)) + eVV
max (o(e:)1)
eeV
Σ max He:)
VEV
where hú; and 4, are the head and tail of the edge e,, respectively. Therefore, show that CF(()) finite, it suffices to prove that He)e) is bounded for all e E V.e E E.
Fix E and e E. Consider an infinite random walk generated with 7 over G, starting at ot
v0 24
22 01
We want to bound the expected first index such that e is the edge traversed in step 0t-1 In et- Denote by A maximum degree of a vertex in G (counted as the number of connected vertices (v)),, by p the maximum number of edges between any single pair of nodes and by d - the diameter of the graph, i.c. the length of the longest shortest path between two vertices (in the undirected version of G). Consider the series of events Ag, Ai.... where A is characterized as:
A; = the event that starting from Vi(42) the walk will follow shortest path one the endpoints of e and then go through
Let e = (he, Fe.tel te). For all values of i, by definition, the length of the shortest path from Us(d42) to he or to is at most d. Therefore, the whole part of the walk described in A, is at most + 1 steps long. By the definition of the used random walk algorithm, which only looks at the previously taken edge, we can deduce that the events A; are all mutually independent.
18

---

Published as a conference paper at ICLR 2026
Moreover, let (i(d+2) uo - one of he, te. Note that by minimality, there cannot be any backtracking while following this path. Therefore, the probability of the next visited node is dependent only on the value of the previous one, and we can bound the probability P(A:) of A, from below by:
U1 - {he:te] } be a shortest path from (i(d+2) to
{-1
P(A.) 2 P(pass through after reaching he te)
III j=0
P((ea(+2))++ ug+1 vi(d+2)+3 =u))
The first term on the right hand side is the probability of selecting e while being at he or te, which is the probability of first selecting the other endpoint (out of at most A neighbors) and then picking e over other edges between he and te (of which there is at most p). Hence:
1
1
Pppaass through e after reaching be or te)
A.P
As we never reach backtracking situation by minimality of the shortest path, we can also write:
P(B((42)+5+1 = Uj+1 U((d=2)+ = uy) = [(( A
Combining these observations, we can derive bound for P(A,)) in terms of A, and d:
8-1
(P(A,) 2 P(aass through e after reaching lt or Le)
II
P(o(d-2)+5+11 = 4z+1 Us(d+2)+5 = =u))
j=0
II
P
A
A
A A.P.
A
A
Finally, note that if A; is true, then the first index such that 1/-1 I4 ii traverses e at most (i+ l) + 2). We can therefore bound the expectation of such f, being (e = H. by:
He(e) & (d 2) 2 P(An) + 2(d + P(-Aa A A1) 3(d++) P(-Ap A -A A A2) +...
= (d (d++2 - P(Au) + 2(d P(-Aa) P(A1) + 3(d P(-Ao) P(-A1) P((A2) +...
(d + P(-Ao) (d+2+ ((((-A))((+++-))))) 2 (...))
(d++2) + [ PAd+1
(d+2+
1 PATTT) (1--44)
(d+2+ (1 1 p (...))
=p(d+2)A d+1
Since p S m. d+2 S n. and A I n. we have He(e) & m(n+ + 2)n", which completes the proof.
1
Remark B.2. The bound obtained in the proof of Lemma B. is very crude. In fact, we could transform the given knowledge graph into a simple graph (undirected, with no multi- edges) by sub- stituting each edge u 4 v with two undirected edges u ++ U(ur,,)) + v. The augmented graph will then have + vertices, and our random walk algorithm naturally translates to a weighted random walk on the transformed graph. This hints at an assumption that in practice, the edge cover time of the used random walk algorithm is of the magnitude O( + m))3 = O(n² + m²)
Let us now prove a fact about the number of distinct, up to isomorphism, graphs in Kn,m- Lemma B.3. For any n, m. the number of isomorphism classes in Kn,m is finite.
Proof. Since the number of distinct relation types that a graph in Kn,m is at most m, suffices to show that the number of isomorphism classes of graphs in Kn,m with exactly relation types is bounded, for all & E m...
Fix the number A 2,, m. and consider G (V, E Kn.m with [R] k. We will show that, up to isomorphism, there are finitely many such choices of G. Firstly, as renaming does not change the graph structure, without loss of generality, we can assume that:
V {11, U2 ... }n and
R = fri, r2....,.. rK)
Then, there exactly 2 possible relational edges set of m elements. Hence, there are (n²*) possible choices of E, and hence, at most (n²k) non-isomorphic choices of G. Since l was chosen arbitrarily, this completes the proof.
(V x RX and E CVVXRV
19

---

Published as a conference paper at ICLR 2026
Lemma B.4. For each pair (n, m)) there exists number Cu.m such that the edge cover time, using the algorithm from SA.1 of any knowledge graph in Kn,m is at most Cr.m-
As two isomorphic graphs have identical cover time, we can set C, to be the maximum of cover times of representatives of all isomorphic classes, which, by finiteness of both, is well-defined.
Proof. The result follows from Lemmas B.I and B.3
Lemma B.5. Let G E kn,m be graph, q (Ng.Ta;?) be link query over G. and ij be walk over G. if traverses all edges of G, then using only the output w(ū;; G.q.-.) of the recording func- tion detailed in $A.2 we can construct a graph-query pair (H, q') isomorphic to (G.9).
Proof. Suppose that .ī = uo 01 23 I De visits all edges of G = (V, E, R) and let & be its length. Recall the anonymization functions idy (:: and id n) as defined in $A.2 The out- put wf;; G. q, :) (the embedding functions provided as the last two arguments are irrelevant) is a sequence of tuples So, St,...,,S with each S, equal to:
S: = (idv i)), id 7),, dir, do )
Consider graph H (V', E", R²) constructed as follows:
the vertices V correspond to the anonymized node ids idy (1: i):
V {id (1; v E V}
Since each vertex must have been visited by 77. this is well-defined. the relation types R' are the anonymized relation ids id i)::
R' fidy (r: n) R}
Again, this is well-defined, as each relation must have been noticed by ij.
the edges E' are reconstructed from the consecutive step encodings using the anonymized vertex and relation indices and the direction dir;:
{(idy idnn)))),d(())).)). idy (x))) dir, = 0,1 0, E'
U [(idy (e:), (e:), (ri),,) idy (m--1) dir, = 1,,
and a query q' = (idy (dy 1). (r, m). for i such that 8,,=h, and 8,,=r,
Then by the definition of w ( AA2.)) p. it is straightforward to check that the pair (-:)() id n)) defines an isomorphism from (G. 9) to (H, Indeed, both these functions are injective by con- struction, and as witnesses all nodes and relations, they are well-defined bijections. For each unique edge traversed by 1j. there exists a unique edge in E" translated to the anonymized space, which implies an isomorphism between E and E", Finally, by utilizing the flags du, -h and 8r =ry we can identify the query head node and relation in the new graph. All things considered, we can reconstruct the pair (G, 9), up to isomorphism, from the output of w(;;; G, ".).
We are now ready to prove the main result regarding the universality of FLOCK as an approximation of link invariant functions. The outline of the proof is as follows: 1) Using the upper-bound on the edge cover time of graphs in Kn,,. derived in Lemma B.4, we can bound the probability of sampling a walk that visits all edges, 2) Once such walk is sampled, we can recover the graph and query, up to isomorphism, from its anonymized form (Lemma B.5) 3) Lastly, we can return the value of the approximated function for the derived isomorphic instance. Since the approximated function link invariant, if the reconstructed graph matches the original one, we return precisely the correct value. Proposition 4.1 With a powerful enough sequence processor fe, the FLOCK framework described in is universal approximator of link invariant functions over K, for all pairs (n.m)
Proof. Let s Kn.m + (V x R + 10, 1l) be a link invariant function over Kn,m returning values from the interval [0, 1]. Let G = (V, E, R) e Kn,m- 9 = (h, r, ?) be a link prediction query over G and L E V be target node. Pick some e, 5 >> O. Our goal is to show that:
((,,()))(h,r,)) Xo (h,, ()) < <)) 1-8
(4)
20

---

Published as a conference paper at ICLR 2026
For simplicity, let consider situation where only single walk fj of length é is sampled by the model (otherwise, omit additional walks). We will also restrict the argument to a single refine- ment case the result can be extended to multiple refinement steps by returning Av, Ar 0 during all additional iterations. Consider a sequence processor fo that given the output w G,g,'.) of the recording protocol, creates graph-query pair (H, with o' ( rig, ?) using the strat- egy described in the proof of Lemma B.5 and returns vector
R whose ith entry is equal h; = (( T''', (d; 7)) where o; is the 7th node visited by 17. The consensus protocol e, provided was visited by 7). can then identify L as one of the vy and pull the corresponding em- bedding h, (( ,yg, (t: 7)), returning it as the output v(t) = h, (note that no matter which specific value of j chosen, this value will be the same). Finally, the classification head can work as an identity operation, returning Xo(G, (t) = v(t) = +p( HT) (Hh)((ha "a'' idy ()))
We claim that if the sampled walk ij traverses all edges of G. then the output of the FLOCK model described above satisfies:
(()) r, = Xo(G, (h, (.,r,,) ?))(()
By Lemma B.5 in such case, the reconstructed pair (H, q') is isomorphic to (G, q) by the isomor- phism id (idv (s idæ(= i)). Since is link invariant, we can write:
((G) r,)) = ((id (1),),, id i)), (lv
((()(((g Fag Tay, (lv
X2(G, (h, r))
Therefore, whenever the walk witnesses all edges of G, the output of the FLOCK model satisfies:
((G))((,,) = Xo(G, (h, (t)
Hence, to show 4. suffices prove that we can uniformly choose the length & of the random walk so that the probability of îj covering all the edges is greater than 1 - 6. By Markov's 's inequality:
Pi does not cover all edges) = Pit takes
steps for 1) to cover edges of G)
E fittts such that 1j covers all edges of G]
CE(G)
But by Lemma B.4 CE(G) S Cr,m n for some constant Cr.m. Hence, taking > Symm we get:
Pijj does not cover all edges) S CF(())
Cn.m
<8
l
This means that for such a choice of E:
PP(îi witnesses all edges of G) > 1- 8
which leads to the conclusion that for C.im the proposed FLOCK framework satisfies: 8"
P(Ip(()(((h, ()) Xo(G, (h,, ?)) > for any choice of G == (V, F, R) E Kn,, and (,, E V RXV.
B.2 INVARIANCE
First, let us recall the definition of invariance for the context of knowledge graphs and the associated notion of invariance in probability. We say that function &p taking KGs as input is invariant if for any pair of isomorphic KGs ~ produces the same input, i.e.
= ((G = (H).
We extend the notion of invariance for further types of inputs, not limited to full knowledge graphs, particularly to random walks and link prediction queries. Let G (V, E, R) E Kkn, and let
H = (V', E', Rt' = G be a KG isomorphic to G' via the isomorphism μ (T,, Φ). For any h E V and r R, we identify the link prediction query q = (h,r,?) ?) in H using the isomorphism μ as:
m(q) p((h, r, ?)) (T(h) (()),))
21

---

Published as a conference paper at ICLR 2026
t0 -4 ... "4 be walk of length e in G. The view of n with μ is defined as: Similarly, let 17
po " V1 3 3 vx) =(0)) d(m)), x(1)) d(rz), o(re), x
Let be function taking inputs drawn from KGs. We call ʃ invariant if for any pair of isomorphic graphs G ~ H and an associated isomorphism (I,, p), satisfies
f(z) = f(((x))
where r can be, e.g., walk link prediction query. In words, invariance means that the function preserves output under the re-identifications of the input graph and the induced transformations of queries and walks.
This notion extends to functions with multiple inputs, where we enforce the transformation on each graph-related input. For example, function yo taking KG, query and a d-dimensional vector is invariant if it satisfies:
VG t H ()G.ggv) ((u(C),
Following the definition of invariance in probability, provided in $ extend all the definitions above to the stochastic case, replacing equality (-) with equality in distribution (-).
We prove the main propositions stated in $4.2 Let's begin with the more general:
Proposition Suppose that the walk sampling protocol is invariant in probability and both the recording protocol and the consensus protocol are invariant. Then, regardless of the choice of the deterministic sequence processor fa. the corresponding FLOCK model is invariant in probability.
(V', E",E", R²) be isomorphic knowledge graphs with isomorphism (x, Φ) transforming G into H. Our goal is show that when the statement conditions met for a FLOCK model Xe with refinement steps, then for any link prediction query q = (h, ?) and any target node / V, the prediction of FLOCK for over (G, q) is identical random variable the prediction for (() over (H, p(g)), i.e.
Proof. Let
X)( - ((Q)
where p(g) (T(h), 6()).). ?). Recall that these predictions are defined as:
(o(G,g)(t) head (v(d((()(()) + (r)) Xa x(() = head(v (0) (T(()) - r 1h ((())
As head is deterministic map, it suffices to show that the final embeddings v(), r()) for (G,g)
and (1)) (I) for (H, (H))) satisfy:
(x) _ (() (m()) and ()) ( (((r) Ve VVEV R
We will prove this result by induction on the number of layers The base case 0 is trivial, as we initialize the embeddings all nodes with a pretrained vector Vo. and all relations with ro-
For the induction step, suppose the claim holds for We drop the superscript (i) for readability. The result for i + becomes apparent by unfolding the definitions invariance of the considered components. Since is invariant in probability, we have
p(n(G)) 17(H)
(5)
Let 17..... 7h the random walks over G using 7 and v ...... vh, be random walks over H. Now, mh..... 77 are independent and identically distributed random variables, each following the distri- bution n2 ~ (G). Similarly, using (5):
hi ( & [(n(G)) = n/s p(n))
(6)
As the recording protocol w invariant, w - w for all j. which with 6 yields:
== w(n)) w()(n))) - w(n))) zj
22

---

Published as a conference paper at ICLR 2026
Then, fo is deterministic map, so 7 implies:
hy == f fo(z)) == hi
Let (Av, Ar) c(bIN,,LN)) (Av', Ar') c(h/:N: be the outputs of the consen- sus protocol. We will denote by cv and cr. the restrictions to the first and second output, e.g. Av (hi:,, n. ). Let x R be vector, and denote by W(G) the space of walks over G. For any vertex U E V, the probability that Av(v) = x equals:
F(Av(b)= -x
E P(Av(U) = x|mi-N = ()) P(T::N = 1))
nEW(G)"
Σ
P(cy(hn,, m:N) x|m:N 7)) P(mI:N 17)
"EW(G)"
= Σ
P(c-(ooloo(m)))))o)m = x|maN =i) P(mnN = 7)
(EW(G)"
Σ
(( x/miNN - - PPNINN- 7))
nEW(G)"
Σ
P(NN= =
REWIG)" (x(f
Similarly, we can derive:
((Av(((())) = =
Σ d'EWHH)" (o(oo(((y))) (),r (xx)))x
P( 57))
Using the invariance of the consensus protocol and the invariance of so w, we can write:
(((o((((1)) ((66(
= (((f(((())) (r(oo((())) ()))
The graph isomorphism / defines bijection between walks W(G) in G and walks W(H) in H, so we can use this correspondence to deduce:
P(Av((T(U)) = xx)
Σ N'EW(HY" a (oll N1. (=(o))=
P(T = =7)
Σ
P(TLNN ((Fi)
(8)
Σ
P( p(T))
REWIG)" cv ()))
Since 17 is invariant in probability, P(m:N P( x(())) Applying this to 8 yields:
((AV((F(V) =x)
Σ
P(TL.N = =n()))
FEW(G)"
Σ
P(NIN= = -1) = P(Av(v) = x)
REW(G)"
23

---

Published as a conference paper at ICLR 2026
As x was chosen arbitrarily, we can conclude that Av(u) - (v(a())).
The proof for relations
follows analogously, considering or instead of Cy. This allows us to write:
Av(u) (T((T)
VOEV
(9)
Ar(r) (o(r))
Vre
By the induction hypothesis, (u) & v² (π)()) for all v E V and r ()))))) _ (r) for all E R. Therefore, by (9. combined with properties of sums of random variables:
vl) = (v) + Av( _ v (x(1)) + (v(()()) = (1tt (z(1))
VeE
(r) = (r) + ar(r) = / (d(r)) Ar'()r)) == ((41) (b(r))
VreR
which completes the induction step, and hence the proof.
We can use the conclusion from Proposition 4.2 to prove the probabilistic invariance of the archi- tecture proposed in $ To be able apply we first need to verify the invariance of all used components, which we formalize in the following lemmas.
Lemma B.6. The choice of the first step " the uniform random walk algorithm described in A.7 invariant.
Proof. Let = (V, E, It) be graph and let H G be an isomorphic graph, with the isomorphism = (म. taking G to H. Consider link prediction query q = (h, r, ?) over G. and its identifi- cation q' plq) (m)) (,)), ?). The goal is to show that when using 17 described in A for (G, and (H,d), ), the first steps:
Vo R V1 and Uo S U1
of the execution of over G and H, respectively, satisfy the following property:
* (Vo) o(Ra), +(Vi) Uo - U1
By definition y, there are three scenarios of choosing the first step, each with probability J Hence, it suffices show that within each scenario, the selection process is invariant in probability:
Scenario 1:
selecting the query head the first node, then proceeding by random. First, æ takes the head node of q to the head node of q'. Secondly, as isomorphisms preserve the number of neighboring nodes and number of edges between pair nodes, we have:
if N if & No.
P(V1 = Vo = h)
NAT
if T() E Na TV) & Nz(s))
P(U; = 7(1) T(h))
and
if rf h) E E(w,,)
P(R: =r Vi =))
B B
otherwise
(())) E B((())))))))
B ((((())
otherwise
P(S Φ(r) -U==
Scenario 2: selecting an edge with query relation type at random. Here, we use the fact that isomorphisms preserve the number of edges of given type. Hence, ! defines a bijection between the sets of edges with type r in G and type ((r) in H, which allows us conclude that this scenario is also invariant in probability.
Scenario 3: selecting the first step completely at random. This case is similar to Scenario
using the invariance the number of neighboring nodes under isomorphism, we can repeat similar calculations in straightforward manner show probabilistic invariance.
24

---

Published as a conference paper at ICLR 2026
Either way, we find that the selection process of the first step of 17 over G translates naturally via to the choice of the first step over H, proving the desired statement.
Lemma B.7. Suppose that the first step 2 is chosen in an invariant manner. Then, the uniform random walk with no bucktrucking algorithm v is invariant in probability.
Proof. Let G (V, E, R) be knowledge graph, and let & be the length of random walks. Let H be a KG isomorphic to G via the isomorphism j (T, Φ). We aim to show that:
()) = =T(V) o(Ra), (V1) ((Ra),
d(Re), π TVv))
Up S U1
-
54 Up (H,
2n E W(G) be a walk of length over G. It suffices to show that the probability of sampling v from G identical the probability of sampling μ(Π) from H:
Let uo 24 -.
P(n(G. E) = 7) = P() () = v(())
To see this, let us expand the definitions of P()(GG, () 1)):
P(n(G, = =Vo vn)
P(V Vo = vo)
II
P(V42 tit2 V441 U:41, V, v1)
(10)
1=
f-1
P(R++ = V4+1 = Ui+1, V; = =w)
and P( =(r))): P(n(H, f) = ((n)) -P = π T(Vo))
PPU ( Uo =(vo))
II PPU1+2 x(1+22) Up+1 = T(U441), =(()))
(1I)
8-1
II IS = M(++1) U++1 U;
Given that the graph isomorphism preserves the number of neighbors for each node and is a bijection,
can easily verify using the definitions from 3 that the following indeed hold:
PVV++2 = oit2 Vi+1 vi+1, V-2 v4) P(U ( U1+1 #(++1), U, =(=)) (12)
P(R r1+1 Vi++ = E.4+1 = =P( *(15+1) U3+1 = ((g+1)) U, = T(;)))
for all i E 10.1...... . 2}.j E {l....,.- - 1}. Moreover, by the assumption that the first step
Vo Ri V1 is invariant, we have:
P((V,, R1 Vn) = (vo, rı, v1)) S1, = (1(())
(13)
But by the laws of conditional probability:
P((V.. Ri, Vi) (uo.rie 01)) P(R Vo en. -.). .(VV eo. Vi vi)
=P(R, = r 0u v1) -(V, Vo =0) PPVo vo)
and analogously:
P(UU,,S; Uz) (())))))(rr)))). (v:)))
= P(S = o(ri) Uo = (vo)), U1 = m(x1)) P(U = ((o) | Uo = =(po)) P(U = (
Substituting these equalities into 2).3 {I..... 13, we get precisely the equality of the right sides
13 and multiplying both sides by the equalities from (12) for all
choices i E {0, 1,..., of (10) and (II). Hence,
P(n(G. E) = 77) P (())
and conclude that (n 0)) _ =, (H,,)), and the algorithm 1) invariant in probability. n
25

---

Published as a conference paper at ICLR 2026
Corollary B.8. The random walk algorithm presented in $A. invariant in probability.
Lemma B.9. The recording protocol w, as described in invariant, provided that the embed- ding functions and invariant.
Proof. Let G = (V, E, R) and H = EV", R') be isomorphic knowledge graphs with the iso- morphism μ = (T,, Φ) taking G to H. Let q (hg. rg. ?) be link prediction query over G, and (()) = (Ι (h.)), (rrg), ?) be its identification in H. Let 1j = vo _ - " os E W(G) be walk over G, and r/ (a( (1n)) ((rz), ((1) 6(r2), p(rx), (we) be the analogous walk over H. To prove that the recording protocol w outlined AA.2 is invariant, suffices to show that the encoding of each step:
5 (idy (14: f), (d i),, dir, b,,=h,, S v(e)),
is identical for and r/'. We will show this for each component:
since defines a bijection between nodes in G and H, for any we have:
idy (wg n)) = arg min(e; = v.] = arg min(m (Tx) = +(x))] idy (T())):
similarly to the point above, ø a bijection between relations of G and H, so we can write:
idæ(r;;)) arg arg min (r+ rs V =r
1
arg arg min [Ø(r.) = <(r,) v m(r:) = (rr;)
id)()(r):))))
dir; is clearly preserved, as the isomorphism μ preserves the directions of edges,
as á are bijections the masks dwy=hy, 8, Orer representing whether the i'th node and relation match the types in the query, satisfy:
= hg - π (vx))= = π (hhg)
o = =ox(v))=((h)) 6r, = =6(tr)=o(ra)
= g <(r;) = x((r))
v and r are invariant by assumption, so:
v(e,) v(T(E:)
and
r(r.) =(((r:))
Combining all these observations, can conclude that wo[i; G,g.v. w(y(T): v) and w is indeed invariant.
Lemma B.10. The consensus protocol e, as described A.4 invariant.
Proof. Let G (V, F, be knowledge graph and H be isomorphic to G via an isomorphism (T,, Φ). Let TT E W(G') be sequence of walks in G. To show that the output of the consensus protocol invariant, we need to prove that for each v V and r E R, the following holds:
Av(u) (T(()) and Ar(r) = Ar (
(14)
where (Av, Ar) c(h, M.N) and (Av', Ar') - c(h ((TIN)) for h (s(V), (M))) (
(R).
The result follows from the fact that т and & are bijections whenever w is the j vertex visited in the walk 7) the jth node of ((7)) must be T()) (and vice versa). An analogous result holds for the relations. Hence, the aggregation performed by e for v (resp. r) over 7/.N is equivalent to the aggregation for π (t) (resp. Φ(r)) over p(ñh:N). ), and [4] is indeed satisfied.
Proposition 4.3
FLOCK with components described in $4 is invariant in probability.
Proof. The result follows naturally from aggregating the results of Corollary B.8 and Lemmas B.9 and 10 followed by applying Proposition 4.2
26

---

Published as a conference paper at ICLR 2026
(2)
05
(2)
65
by
02
04
Figure 5: An example of a graph from PETALS with e 4, 1 = 2 and / == 3, and the associated link prediction instances (dashed). The relation types 'red', 'blue', 'pink" and "yellow" are structurally isomorphic, hence become equated in the eyes of the existing KGFMs.
C DETAILS OF THE PETALS BENCHMARK
State-of-the-art knowledge graph foundation models (KGFMs) typically impose relational invari- ance. Formally, given two knowledge graphs G (V, E, R) and И -V'' E', Rt), if there exists an isomorphism (T, σ) from G to H. then for any r E R. the model enforces identical representa- tions for r and its image ((r) E Rt. This design promotes generalization across different graphs, as it aligns analogous relations, but reduces expressivity within a single graph (G = H), where relations related by automorphisms are forced to be indistinguishable. Concretely, if an automor- phism (T, Φ) of G maps to r2, then the model must treat r1 and as identical during inference. While some approaches mitigate this limitation via the labeling trick, assigning distinct embeddings to query- nodes and relations, this only isolates the queried relation type and does not resolve the underlying issue in general.
Motivated by this limitation, we introduce the PETALS benchmark. PETALS comprises 220 graphs. each paired with a link prediction query (h, r, ?) and a target set {t1. f2]. While t1 and t2 are non-isomorphic, KGFMs enforcing relational invariance
unable to distinguish them, producing identical predictions. We empirically validate this property by evaluating the classification accuracy of marking TRUE and f2 FAI.S,, reported in Table
C.I STRUCTURE OF THE STUDIED KGS
Knowledge graphs in PETALS follow a flower-like structure, parametrized by the number c of 'petals', their length I and the length t of the 'stem' (see Figure 5 for visualization).
Vertices. Each 'petal' is a set (0) of 2/ + vertices A
{a { a ..... 22+1 }. while the stem
B consists of t + 1 nodes = {bo, b1,...,brl! b,}.
The full set of entities is then:
V = B A(i) = (b bi..,,....bbb b.) u (a |
i=1
We call by the 'central' node, it is connected to every petal, as described below.
Edges. The nodes of the stem are connected a consecutive manner by the same relation type ro- Precisely, for each E 1,... I, there exists an edge (b,-1, bz). Each petal ( is associated (8) (6) with two edge types ri .22 and is connected the central node by with links (bo.r(" ri ai (t)
and
(bo. 1g 0,2 (1) The rest of the petal connected with edges of type y(0 only, going from a (8)
(4)
(8) (6) 025+1 and from az 02j+2- Finally, there are also edges linking d2--1 and az to 22
(5) ( (8)
27

---

Published as a conference paper at ICLR 2026
Therefore, the full set of edges can be characterized as:
(ba, r1 aa (8)) (bo. (bo. (0) r2)
F. =
21 ( aa..t) (ai . (9tt))))))
==1==1
1f ( (al (
We select each of the types ri and r2 from the set of considered relations R = fri...., | } so that any glaatin-invaarian model will equate all petals (i.e. so that for each pair of petals, there is an automorphism taking one to another). For instance, Figure
cyclic pattern, in which r2 ri ((++1) Such symmetry causes all petals be isomorphic, and leads to the inability of KGFMs to distinguish between the relations inside them.
(e)
Link prediction instances. Although the petals are isomorphic to each other, given the asymmetry
of edge types from by to ai and (+) the nodes within a single petal generally can be distinguished. Therefore, for each graph with the structure as described above, stem nodes b,. and ask the link prediction query (b,, To, ?). For the target nodes, we randomly select (t)
randomly sample one of the
petal index and distance j from the central node bo. and consider the predictions for o and azj. For example, Figure 5 shows the case when b, bu, 1 and j 1, where the query is (bu, ro. ?) (1)
and we are interested in the scores for aj and 02
C.2 PARAMETERS AND GENERATION
We construct PETALS by manually designing 11 relation-assipie relation-assignment schemes that guarantee isomor- phism across all petals. For each such selection, which already determines the number of petals, we generate 20 graphs corresponding to all combinations of t E 2,,2,3, 4) and E 2,,2,3, 3, 4, 5). Each graph is paired with link prediction query and two target nodes, sampled as described above. This yields 11 20 220 instances that constitute the PETALS benchmark.
28

---

Published as a conference paper at ICLR 2026
Table Training scalability analysis on single NVIDIA RTX A6000 (48 GB) with batch size 8. FLOCK using 16 number of base walks and ensemble.
Model
Parameters
Time batch (s)
GPU memory (GB)
ULTRA
168,705
0.117
2.110
TRIX
87,138
0.690
3,442
FLOCK
801,969
1.30
27.89
Table 7: Inference scalability on a single NVIDIA RTX A6000 (48 GB) with batch size = 8. Left columns specify base walks n and ensembled passes Dashes indicate not applicable.
Model
# Base Walks
Ensemble P
Time /batch (8)
GPU memory (GB)
ULTRA
0.073
0.848
TRIX
1
0.500
1.382
16
1
1.26
2,868
16
2
1,99
2.864
16
4
3,24
3,938
16
8
5.45
5.172
FLOCK
16
16
9.45
8.892
128
1
1.77
5,000
128
2.80
7,880
128
4
5.00
14.42
128
8
10.05
43.68
D COMPUTATIONAL COMPLEXITY
Recall that I is the iterations in each forward pass of FLOCK; is the base walk count; is the walk length; L is the number of linear sequence-model layers (such as GRU); and is the hidden dimension for the sequence processor. Note that in practice, perform P forward passes and ensemble their outputs to reduce variance. For a single pass (P=1), walk sampling and recording cost O while the sequence processor with L layers of hidden dimension costs LdF). The consensus protocol costs ((ndd). In total, the time complexity is O OPInnLdd*) which scales linearly with the number (base) walks n, the length of walks L, and the number of ensembled predictions P. We empirically verified this in {
Compared with pesssage-passing KGFMs like ULTRA and TRIX, FLOCK'S complexity is indepen- dent of the graph size and average degree; empirically, however, using more walks (increasing n) and longer walks (increasing E) improves coverage and yields more fine-grained representation.
The space complexity of FLOCK per forward pass is ((ndd) plus model parameters ((Ld²). Note that running ensembles sequentially keeps peak memory near this bound, while parallel ensembling increases it by factor of P.
E SCALABILITY ANALYSIS
To investigate the scalability of the proposed method FLOCK,, we report the training and inference time per batch and peak GPU memory for UL.TRA, TRIX, and variants of FLOCK on a single RTX A6000 (48 GB) 6
Training. During training, we fix FLOCK to n = 16 base walks and with an ensemble size of P 1, which yields higher cost than ULTRA/TRIX but remains feasible on single GPU. addition, unlike ULTRA/TRIX. FLOCK does not rely GNN message passing where highly optimized fused sparse kernels (e.g., RSPMM kernel developed in Zhu et al. (202T) accelerate computation; instead, FLOCK's runtime is dominated by walk sampling and sequence encoding, making time per batch the main bottleneck. As a result, pretraining typically takes about three days. One avenue for future work is to develop similarly highly optimized kernels for random-wall sampling.
29

---

Published as a conference paper at ICLR 2026
Table 8: Noise injection over the best performing KGFM baseline TRIX.
(a) Zero-shot entity prediction.
(b) Zero-shot relation prediction.
(c) Accuracy on PETALS.
MRR
Hits@10
MRR
Hits@
Accuracy
TRIX
0.366
0.518
TRIX
0.792
0.687
TRIX 50%
+ noise
0.385
0.545
noise
0.739
0.643
+ noise
52%
FLOCK
0.391
0.560
FLOCK
0.881
0.817
FLOCK 100%
Additionally, we report the inference results in Table 7 where we vary the number of walks n and ensembled passes P. We observe near-linear growth of latency and VRAM with n, reflecting the dominant costs of walk sampling and sequence processing. Note that during infer- ence, ensembled predictions are parallelizable, meaning that with sufficient GPU memory, these P stochastic passes can be executed concurrently, so the latency grows sublinearly
Inference.
P. while peak VRAM scales roughly linearly with P. In practice, reducing n (walks) or P (ensembled passes) lowers both memory and latency, while larger n/ settings trade extra cost for better coverage and stability on harder KGs.
F NOISE INJECTION OVER EXISTING KGFMS
Setup. Since noise injection is a possible way to build a probabilistic equivariant KGFM in different way from our approach (Gao et al. 2023). it is natural ask how such KGFMs would perform compared to FL.CK. To answer this question, we apply noise injection over the best per- forming KGFMs baselines TRIX. Specifically, sampled from a uniform distribution ~ 0[50.5, 0] to all relation and entity embeddings after the initialization stage. Note that the addition of noise technically breaks deterministic node-relation equivariance, but the resulting model (TRIX noise) still respects probabilistic node-relation equiv- ariance. We then pretrain TRIX using the same experimental setup shown in $5.2 and compare with TRIX without noise injection and FLOCK. To minimize the variance induced by injected noise and to ensure fair comparison, we report ensembled prediction results with 16 samples for both TRIX noise and FLOCK. This is strong baseline implementing the ideas of prior work on noise injection and test-time ensembling for message passing networks on KGs (Lee et al. 2023 Gao et al. 2023).
each forward pass, we add element-wise noise
Results. We report the average zero-shot performance for entity prediction and relation prediction over 54 KGs in Tables Sal and 8b respectively, as well as trained performance for PETALS in Ta- ble Sc Across all tasks, TRIX with naive noise injection fails close the gap between FLOCK. In particular, TRIX + noise degrades compared with vanilla TRIX without noise injection in relation prediction, while boosting the performance in the entity prediction task. We hypothesize that such a difference lies in the added randomization breaks symmetry among entity embeddings more than among relation embeddings, and entity prediction depends more on having distinguishable entity representations than relation prediction does. Additionally, we attribute this performance gap be- tween FLOCK and TRIX noise to the source of randomization. FLOCK introduces stochasticity through random walks, which induces structure- structure-informed perturbations that respect the underlying topology. In contrast, TRIX with naive noise injection attempts to break deterministic node-relation equivariance by introducing grructure-agonisti noise, which might, in turn, hurt the model's general- ization. Together, these findings suggest that simply adding structure-agonistie noise is insufficient; performance gains only arise when stochasticity topology-aware and induced from the graph structure in principled way.
G CASE STUDY: RELATION EMBEDDING ON METAFAM
Setup. To further showcase why expressivity matters for zero-shot generalization, we present a case study the METAFAM dataset (Zhou et al. 2023). METAFAM is built from a fixed family- relations ontology: during training, models observe edges with relations motherof, fatherOf, daughterof, sonOf, while the test queries only involve motherOf and ferher Up to gender symmetries, this reduces to two effective predictive patterns (parent_op of vs. child _hild_of) and the test set focuses a single one (parent_of).
30

---

Published as a conference paper at ICLR 2026
USTRA PARAROA Dwc.dding Mesifars
Pack Pelltic Diccodrrgs Messfar
Figure 6: PCA of relation embeddings on METAFAM. ULTRA (Left) maps several inverse pairs (e.g., fatherOf sonOf) to almost similar embeddings, where FLOCK (Right) yields clearly separated embeddings, indicating that its probabilistic equivariance allow FLOCK to distinguish between these semantically different relations, explaining its strong zero-shot performances.
p-value 5.0e-05
Teen
2.5
Flock MRR /Ulrr MAR
2.0
1.5
1.0
Density
Figure 7: The zero-shot entity prediction performance of FLOCK relative to ULTRA, plotted against the densities of the 53 KGs. Performance of FLOCK and log log-density of KGs have a Pearson correla- tion coefficient of -0.53 with p-value 5.0e-5, showing statistically significant negative correlation.
In the zero-shot setting, KGFMs cannot adapt to this ontology and must rely their pretrained relation representations. Notice that here, METAFAM challenging: many relations are structurally similar (e.g., fatherOf, sonOf, sisterOf, nieceOf) yet encode opposite predictive patterns.
Result. Figure 6 shows that ULTRA'S relation embeddings largely collapse these families, plac- ing fatherOf and sonOf in almost identical positions the PCA plane. This collapse makes it difficult to distinguish who is the parent and who is the child, leading to poor zero-shot perfor- mance. FLOCK, contrast, can distinguish between these relations even if they are structurally similar, thanks to its andom-walk sampling which introduce probabilistic equivariance on nodes and relations. As a results, FLOCK can produce distinct embeddings to fatherOf and sonOf f and achieves much stronger zero-shot performance METAFAM.
H ANALYSIS KG SPARSITY AND PERFORMANCE
Setup. While explaiin FLOCK'S high performance for METAFAM, understanding its perfor- mances for other general KGs would be beneficial. Thus, we present an additional analysis on the 53 remaining KGs by identifying structural property that is indicative of the performance of
31

---

Published as a conference paper at ICLR 2026
FLOCK. For the performance measure, we use the relative gain of FLOCK's 's zero-shot entity predic- tion MRR compared to ULTRA. For the structural property, we focus on density of KGs defined by [ which affects the speed of random walks traversing all edges of KG (intuitively because more edges means more time needed to traverse all of them), and hence is relevant in the context of our theory in $4.2 We make the argument more grounded below.
|B]
Result. Figure shows that FLOCK tends to perform well sparse KGs, while less dense KGs, and the tendency is statistically significant. Interestingly, this agrees with our theoretical analysis in $4.2 in which necessary condition for universality that the random walk covers all edges of a KG with high probability (Proposition 4.1). The time taken until covering all edges is called the edge cover time, and known be e.g.. O(IV) E)) for uniform walks (Zuckerman, 1991), which is proportional to the density of a graph. This suggests the performance of F1.OCK is associated with the easiness to visit as many edges as possible rapidly, which is more challenging for dense KGs. This analysis is consistent with the observations of recent works on graph learning based on random walks, e.g., Wann & Cho (2024, Section 6) and Kim et al. 2025).
I FURTHER EXPERIMENTAL DETAILS
This section provides the full details for all experiments described in the main text. For pretraining. we fit the FLOCK model on three standard transductive knowledge graph com- pletion benchmarks, following Galkin et al. (2024): FB15K-237 (Toutanova & Chen 2015), WN18RR (Dettmers et al. 2018), and CoDEx Medium (Safavi & Koutra 2020). Then, we eval- uate zero-shot transfer of both entity prediction and relation prediction, as well as the finetuning performance on multiple datasets grouped as follows:
Datasets.
Inductive r. Link prediction tasks involving previously unseen nodes and relation types. This includes the 13 datasets from INGRAM (Lee et al. 2023): FB-25, FB-50, FB-75, FB- 100, WK-25, WK-50, WK-75, WK-100, NL-O, NL-25, NL-50, NL-75, NL-100, as well as 10 datasets from MTDEA (Zhou et al. 2023): : MT1 tax, MTI health, MT2 org. MT2 sei, MT3 art, MT3 infra, MT4 sei, MT4 health, Metafam, and FBNELL.
Inductive e. Link prediction tasks involving novel nodes but fixed relation types. This category comprises 12 GralL datasets (Teru et al.) 2020) (WN-v1 through WN-v4, FB-vl through FB- v4, NL-v1 through NL-v4), 4 INDIGO benchmarks (Liu et al. 2021) (HM 1k, HM 3k, HM 5k, HM Indigo), and 2 NodePicce datasets (Galkin et al. 2022): ILPC Small and ILPC Large.
Transductive. Link prediction tasks where both entities and relations are observed dur- ing training. These include CoDEx Small, CoDEx Large (Safavi & Koutral 2020). NELL- 995 (Xiong al.) 2017), YAGO 310 (Mahdisoltani et al. al.], 2015). WDsinger. NELL23K. FB15k- 237(10), 1B133-237(20 FB15K-237(50) (Ly et al. 2020). AristoV4 (Chen et al. (021)) DB- pedia100k (Ding et al. 2018). ConceptNet10Ok (Malaviya et al. 2020). and Hetionet (Him-
melstein et al. 2017).
Full results of 55.2 Full tables zero-shot entity prediction results are presented in Table 9 and
full tables of finetuned performance are given in Table 10 We further provide the complete zero- shot and finetuned relation prediction results in Table 11 and Table 12 The dataset statistics are in
Table [3 Table 14n Table 15] Table [6 the pretraining graph mix shown in 5.3) Finally,
detailed hyperparameter settings can be found in Table 17
Table [8 and Table 19
Training. Following conventions in the literature (Zhu et al. 2021 Huang et al. 2023 for each triple (h,r, r, t), we add the corresponding inverse triple (h, .t), where r- is a fresh relation symbol. All FLOCK instances and its variants are optimized to minimize the negative -ikelioiddood over positive and negative facts under the partial completeness assumption (Galárraga et al. (2013), where negatives are generated by randomly corrupting either the head or the tail entity (for entity prediction) or by corrupting the relation (for relation prediction). To reduce overfitting, we remove edges that directly connect the queried endpoints. The best checkpoint is selected by validation performance. For entity prediction, obtain the score p(h, r, t) by passing into 2-layer MLP. For relation prediction, we concatenate the embedding for source h, target I, and potential relation obtain the score p(h,
take the embedding for potential target t and relations r, and
32

---

Published as a conference paper at ICLR 2026
Let (h, r, t) be positive triple and let & denote the number of negatives sampled per positive, where (he, 7, fx) is the i-th negative samples for entity prediction, and h, rs. t2 is the i-th negative samples for relation prediction. Following Sun et al. (2019), also consider a aelf-adversarial variant where negatives are reweighted according to their current difficulty. With adversarial temperature or > O, the weights for entity and relation prediction, respectively, are
log(1 10 (1))))))
p(h,,
log(1 p(h,, (b1))
wh Softmax
W Softmax
The corresponding losses become
Cadt p(,
Σ
fi log (1 - p(h, t()),
==1
badv rl,
Σ
w log(1 - p(h, r t)).
33

---

Published as a conference paper at ICLR 2026
Table 9: Zero-shot entity prediction results. Bold indicates the best score row.
ULTRA
TRIX
FLOCK
Dataset
MRR
Hits @
MRR
@its@10
MRR
Hits @ 10
Inductive e,
FB-25
0.388
0,640
0.393
0.650
0,404
0,664
FB FB-50
0.338
0.543
0.334
0.547
0.352
0.566
FB-75
0.403
0.604
0.401
0.611
0.418
0.622
FB-100
0.449
0.642
0.436
0.635
0.452
0.663
WK-25
0.316
0.532
0.305
0.496
0.280
0.491
WK-50
0.166
0.324
0.166
0.313
0.136
0.278
WK-75
0.365
0,537
0.368
0.513
0,382
0.538
WK-100
0.164
0.286
0.188
0.299
0.187
0.304
NL-0
0.342
0.523
0.385
0.549
0.381
0.606
NL-25
0.395
0.569
0.377
0.589
0.345
0.590
NI.-50
0.407
0.570
0.404
0.548
0.366
0.565
NL-75
0.368
0.547
0.351
0.525
0.311
0.524
NI.-100
0.471
0.651
0.486
0.676
0.452
0.692
MTI lax
0.224
0.305
0.358
0,452
0,282
0.383
MTI health
0.298
0.374
0.376
0.457
0.385
0.481
MT2 org
0.095
0.159
0.091
0.156
0.100
0.163
MT2
0.258
0.354
0.323
0.465
0.318
0.458
0.259
0.402
0.284
0.441
0.301
0.466
MT3 infra
0.619
0.755
0.655
0.797
0.684
0.821
MT4 sci
0,274
0.449
0,290
0,460
0,301
0,463
MT4 health
0.624
0.737
0.677
0.775
0,680
0.780
Metafam
0.238
0.644
0.341
0,815
0.476
0.935
FBNELL
0.485
0.652
0.473
0.660
0.502
0.700
Inductive
FB-v1
0.498
0.656
0.515
0,682
0.500
0.697
FB-v2
0.512
0.700
0.525
0.730
0.535
0.737
FB-v3
0.491
0.654
0.501
0.669
0.511
0.685
FB-v4
0.486
0.677
0.493
0.687
0,505
0.702
WN-v1
0.648
0.768
0.699
0.791
0.698
0.803
WN-v2
0.663
0.765
0.678
0.781
0.696
0.790
WN-v3
0,376
0.476
0.418
0.541
0.467
0.608
WN-v4
0.611
0.705
0.648
0.723
0.653
0.729
NL-v1
0.785
0.913
0.806
0.898
0.658
0.863
NL-v2
0.526
0.707
0.569
0.768
0.588
0.797
NL-x3
0.515
0.702
0.558
0.743
0.590
0.783
NL-x4
0.479
0.712
0.538
0.765
0.555
0,786
HM 1k
0.059
0.092
0.072
0.128
0.069
0.119
HM 3k
0.037
0.077
0.069
0.119
0,067
0.118
HM 5k
0.034
0.071
0.062
0.110
0.064
0.116
HM Indigo
0.440
0.648
0.436
0.645
0.423
0.638
ILPC Small
0.302
0.443
0.303
0.455
0.309
0.459
ILPC Large
0.290
0.424
0.307
0.428
0.318
0.438
Transductive
NELL995
0.406
0.543
0.472
0.629
0.494
0.655
NEI.I.23
0.239
0.408
0.290
0.497
0.233
0.398
WDsinger
0.382
0.498
0,511
0.609
0.410
0.528
ConceptNet
0.082
0.162
0.193
0.345
0.248
0.453
CoDEx Small
0.472
0.667
0.472
0.670
0.441
0.644
CoDEx Large
0.338
0.469
0.335
0,469
0.342
0.464
YAGO310
0.451
0.615
0.409
0.627
0.414
0.674
AristoV4
0.182
0.282
0.181
0.286
0.308
0.443
DBpedia 100k
0.398
0,576
0.426
0.603
0.450
0.627
Hetionet
0.257
0.379
0.279
0.420
0.246
0.371
FB 237((1))
0.248
0.398
0.246
0.393
0.246
0.402
FB 1 237(20)
0.272
0.436
0.269
0,430
0.273
0.444
FB15k- -37(55))
0.324
0.526
0.321
0.521
0.319
0.518
34

---

Published as a conference paper at ICLR 2026
Table 10: Finetuned entity prediction results. Bold indicates the best score row.
ULTRA
TRIX
FLOCK
Dataset
MRR
@its@10
MRR
@
MRR
Hits@ 10
Inductive e,
FB-25
0.383
0.635
0.393
0.650
0.405
0.666
FB-50
0.334
0.538
0.334
0.547
0.357
0.570
FB-75
0.400
0.425
0.630
FB-100
0.444
0.643
0.436
0.633
0.460
0.668
WK-25
0.321
0.535
0.300
0.493
0.298
0.506
WK-50
0.140
0.280
0.166
0.313
0.127
0.260
WK-75
0.380
0.530
0.368
0.513
0.405
0.556
WK-100
0.168
0.286
0.188
0.299
0.187
0.306
NL-O
0.329
0.551
0.385
0.549
0,418
0.619
NL-25
0,407
0.596
0.377
0.589
0,405
0.626
NI S0-
0,418
0.595
0.405
0.555
0.39
0.562
NL-75
0.374
0,570
0.351
0.525
0.344
0.544
NL-100
0.458
0.684
0.482
0.691
0.486
0.714
MTI tax
0,330
0.459
0.397
0.508
0.413
0.497
MT health
0.380
0.467
0.376
0.457
0.394
0.493
MT2 org
0.104
0.170
0.162
0.107
0.174
sci
0.311
0,451
0.526
0.366
0.525
MT3
0.306
0.473
0.289
0.441
0.330
0.483
MI3 infra
0.657
0.807
0.672
0.810
0.709
0.838
MT4 sei
0.303
0.478
0.305
0.482
0.324
0.509
MT4 health
0.704
0.785
0.702
0.785
0.711
0.790
Metafam
0.997
1.090
0.997
1,000
0.992
1,000
FBNELL
0.481
0.661
0.478
0.655
0.531
0.714
Inductive
FB-v1
0.509
0.670
0.515
0.682
0.549
0.721
FB-v2
0,524
0.710
0.525
0.730
0.553
0.754
FB-v3
0.504
0.663
0.501
0.669
0.528
0.696
FB-v4
0.496
0.684
0.493
0.687
0.510
0.702
WN-v1
0.685
0.793
0.705
0.798
0.715
0.811
WN-v2
0.679
0.779
0.780
0.702
0.795
WN.v3
0.411
0.546
0.425
0.543
0.494
0.627
WN-v4
0.614
0.720
0.650
0.722
0.665
0.741
NL-vI
0.757
0.878
0.804
0.899
0.762
0.928
NL.v2
0.575
0.761
0.571
0.764
0.612
0.806
NL-v3
0.563
0.755
0,571
0.759
0.606
0,803
NL-v4
0.469
0.733
0.551
0.772
0.572
0.801
HM 1k
0.042
0.100
0.072
0.128
0,071
0.153
HM 3k
0.030
0.090
0.069
0.1
0.067
0.153
HM 5k
0.025
0.068
0.074
0.118
0.061
0.130
IIM Indigo
0.432
0.639
0.436
0.645
0.418
0.633
ILPC Small
0.303
0.453
0.303
0.455
0.305
0.454
ILPC Large
0.308
0.431
0.310
0.431
0.320
0.441
Transductive
NELL995
0.509
0.660
0.648
0,531
0.665
NEL.L.3K
0.268
0.450
0.306
0.536
0.280
0.465
WDsinger
0.417
0.526
0,502
0.620
0.435
0.543
ConcepiNett00k
0.310
0.529
0.340
0.564
0.352
0,580
CoDEx Small
0.490
0.686
0.484
0.676
0.463
0.648
CoDEx Large
0.343
0.478
0.348
0.481
0.342
0.467
YAGO310
0.557
0.710
0.541
0.702
0.552
0.700
AristoV4
0.343
0.496
0.345
0.499
0.383
0.523
DBpedial 00k
0.436
0.603
0.457
0.619
0.470
0.623
Hetionet
0.399
0,538
0.394
0.534
0.314
0.465
(5k-23((10)
0.254
0.411
0.253
0.408
0.260
0.420
151-22237
0.274
0.445
0.441
0.284
0.459
FB15K-237(50)
0.325
0.522
0.317
0.517
Pretrained
FB15k-237
0.368
0.564
0.366
0.559
0.343
0.532
WNISRR
0.480
0.614
0.514
0.611
0.550
0.656
CoDEx Medium
0.372
0.525
0.365
0.521
0.351
0.496
35

---

Published as a conference paper at ICLR 2026
Table 11: Zero-shot relation prediction results. Bold indicates the best score row.
ULTRA
TRIX
FLOCK
Datasel
MRR
@its@I
MRR
Hits@1
MRR
Hits
Inductive
FB-25
0,687
0.565
0.805
0.724
0.895
0.839
FB-50
0.696
0.575
0.780
0.699
0.880
0.820
FB-75
0.698
0.555
0.822
0.747
0.903
0.844
FB-100
0.830
0.728
0,921
0.880
0.962
0,938
WK-25
0.857
0.760
0.881
0.823
0.952
0.929
WK-50
0.865
0.793
0.868
0.818
0.921
0.882
WK-75
0.911
0.875
0.916
0.883
0.962
0.944
WK-100
0.887
0.812
0.907
0.869
0.963
0,937
NL-0
0.632
0.502
0.658
0,519
0.714
0,574
NL-25
0.688
0.562
0.742
0.614
0.729
0.632
NL-50
0.680
0.569
0.755
0.636
0.813
0.728
NL-75
0.795
0.692
0.788
0.699
0.833
0.756
NL-100
0.743
0.564
0.884
0.796
0.939
0.889
MTI Lax
0.98S
0.976
0.975
0.958
0.998
0.997
MTI health
0.721
0.561
0.973
0.949
0.991
0.983
MT2 org
0.974
0.951
0.986
0.973
0.991
0.984
MT2
0.976
0.961
0.964
0,941
0.995
0.992
0.881
0.798
0.885
0,825
0.944
0.907
MT3 infra
0.962
0.935
0.940
0,905
0.989
0,980
MT4 sci
0.933
0.891
0.966
0.944
0.974
0.957
M1T4 health
0.826
0.719
0.937
0.898
0.990
0.983
Metafam
0.124
0.000
0.291
0.011
0.490
0.223
FBNELL
0.700
0.564
0.726
0.605
0.833
0.737
Inductive
FB-v1
0.646
0.523
0.705
0.599
0.814
0.723
FB-v2
0.695
0.570
0.713
0.590
0.847
0.761
FB-v3
0.679
0.553
0.742
0.644
0.860
0.780
FB-v4
0.638
0.488
0.766
0.665
0.873
0.799
WN-v1
0.836
0.740
0.792
0.613
0.924
0.858
WN-v2
0.853
0.790
0.764
0.572
0.924
0.863
WN-v3
0.707
0.577
0.741
0.568
0.937
0.888
WN-v4
0.860
0.803
0.764
0.570
0.937
0.886
NL-v1
0.636
0.358
0.657
0.453
0.862
0.731
NL-v2
0.742
0.652
0.780
0.696
0.893
0.855
NL-v3
0.669
0.544
0.725
0.612
0.815
0.731
NL-v4
0.606
0.489
0.794
0.691
0.868
0.807
ILPC Small
0.905
0.843
0.919
0.872
0.955
0.921
ILPC Large
0.875
0.799
0.894
0.829
0.948
0,908
HM 1K
0.626
0.447
0.663
0.414
0.687
0.500
HM 3k
0.592
0.439
0.664
0.418
0.714
0,549
HM
0.605
0.452
0.672
0.428
0.746
0,593
HM Indigo
0.681
0.559
0.852
0.765
0.956
0.921
Transductive
NELL995
0.583
0.437
0.578
0.457
0.684
0.555
NELL.23K
0.669
0.548
0.756
0.657
0.831
0.762
WDsinger
0.668
0.546
0.720
0.621
0.823
0.738
ConceptNet100k
0.181
0.083
0.650
0.469
0.795
0.658
Coll
0.900
0.820
0.961
0.935
0.982
0.970
CoDExLarge
0.892
0.824
0.902
0.837
0.973
0,950
YAGO310
0.646
0.403
0.783
0.598
0.971
0.943
Aristo
0.254
0.201
0.389
0.265
0.597
0.496
DBpedia 100k
0.650
0.509
0.717
0,582
0,919
0.861
Hetionet
0.634
0.524
0.809
0.707
0.940
0.890
FB155-237((0)
0.688
0,550
0.795
0.711
0.918
0.876
FB 15k 237(20)
0.695
0.558
0.834
0.758
0.952
0,923
FB1 (315k)237((0)
0.717
0.591
0.876
0.812
0.968
0.946
36

---

Published as a conference paper at ICLR 2026
Table 12: Finetuned relation prediction results. Bold indicates the best score per row.
ULTRA
TRIX
FLOCK
Dataset
MRR
Hits@1
MRR
@its@1
MRR
Hits@1
Inductive e,
FB-25
0.684
0.563
0.805
0.724
0.909
0.857
FB-50
0.696
0.575
0.780
0.699
0.881
0.820
FB-75
0.754
0.638
0.699
FB-100
0.851
0.769
0.921
0.880
0.965
0.939
WK-25
0.897
0.834
0.905
0.860
0.968
0.954
WK-50
0.865
0.793
0.881
0.840
0.925
0.876
WK-75
0.911
0.875
0.937
0.910
0.965
0.948
WK-
0.924
0.879
0.916
0.885
0.970
0.946
NL-O
0.632
0.502
0.655
0.518
0.731
0,602
NL-25
0.737
0.622
0.709
0.606
0.757
0.634
NL-50
0.808
0.704
0.774
0.683
0.814
0.721
NL-75
0.795
0.678
0.790
0.671
0.848
0.774
NL-100
0.803
0.678
0.885
0.793
0.937
0.887
MTI tax
0.990
0.984
0.995
0.990
0,999
0.998
MTI health
0.929
0.867
0.973
0.949
0.994
0.988
MT2 org
0.981
0.963
0.987
0.978
0,994
sci
0,977
0.961
0,990
0.992
MT3
0.907
0.851
0.887
0.828
0.950
0.916
MI3 infra
0,966
0.947
0.970
0.952
0.996
0.993
MT4 sei
0.954
0.929
0.972
0.952
0,983
0.968
MT4 health
0.951
0.919
0.986
0.979
0.995
0.991
Metafam
0.368
0.112
0.265
0.024
0.997
0.995
FBNELL
0.720
0.576
0.766
0.639
0.879
0.801
Inductive
FB-vI
0.650
0.513
0.705
0.599
0.855
0.766
FB-v2
0.675
0.547
0.713
0.590
0.887
0.812
FB-v3
0.677
0.556
0.742
0.644
0.879
0.810
FB-v4
0.690
0.560
0.766
0.665
0.884
0.807
WN-vl
0.844
0.754
0.776
0.591
0.926
0.879
WN-v2
0.834
0.766
0.765
0.574
WN.v3
0.707
0.577
0.756
0.594
WN-v4
0.861
0.795
0.804
0.651
0.943
0.898
NL-v1
0.719
0.504
0.590
0.341
0,883
0.766
NT.v2
0.668
0.549
0.811
0.740
0.911
0.870
NL-v3
0.646
0.484
0.757
0.643
0.868
0.795
NL-v4
0.570
0.412
0.822
0.735
0.906
0,849
ILPC Small
0.922
0.876
0.919
0.872
0.953
0.918
I.PC Large
0.875
0.799
0.894
0.829
0.953
0.915
HM 1k
0.626
0.447
0.663
0.414
0.756
0.561
HIM 3k
0.592
0.439
0.664
0.418
0.790
0.623
HM 5k
0.605
0.452
0.672
0.428
0.744
0.591
HM Indigo
0.726
0.614
0.835
0.746
0.946
0.903
Transductive
NELL995
0.630
0.513
0.578
0.457
0.713
0.584
NEI.L.33
0.688
0.571
0.755
0.658
0.869
0.805
WDsinger
0.730
0.603
0.721
0.627
0.885
0.815
ConceptNett00k
0.612
0.488
0.712
0.551
0.885
0.813
CoDExSmall
0.942
0.900
0.964
0.943
0.981
0.967
CoDExt Lare
0.907
0.850
0.908
0.845
0.973
0.950
YAGO310
0.930
0.891
0.826
0.666
0.970
0.942
AristoV4
0.254
0.201
0.498
0.381
0.651
0.547
DBpedial 00k
0.650
0.509
0.780
0.665
0.923
0.869
Hetionet
0.737
0.646
0.922
0.862
0,942
0.897
15k 232-1)
0.688
0.550
0.795
0.711
0.940
0.905
FB15k-237((0)
0.695
0.558
0.846
0.778
0.958
FB15K-237(50)
0.728
0.618
0.903
0.858
0.948
Pretrained
FB15k-237
0.795
0.709
0.924
0.870
0.976
0.957
WN18RR
0.914
0.871
0.783
0.634
0.982
0.968
CoDExMedium
0.919
0.870
0.931
0.886
0.974
0.952
37

---

Published as a conference paper at ICLR 2026
Table 13: Dataset statistics for inductive-e, link prediction datasets. Triples are the number of edges given training, validation, test graphs, respectively, whereas Valid and Test denote triples to be predicted in the validation and test graphs.
Training Graph
Validation Graph
Test Graph
Dataset
Entities
Rels
Triples
Entities
Rels
Triples
Valid
Entities
Rels
Triples
Test
FB-25
5190
163
91571
4097
216
17147
S716
4097
216
17147
5716
FB.50
FB-75
4659
134
62809
2792
186
9316
3106
2792
186
9316
3106
FB-100
4659
134
62809
2624
77
6987
2329
2624
6987
2329
WK-25
12659
47
41873
3228
74
3391
1130
3228
3391
1131
WK-50
12022
72
82481
9328
93
9672
3224
9328
93
9672
3225
WK-75
3430
WK 100
9784
67
49875
12136
37
13487
4496
12136
37
13487
4496
112
2287
2287
NL-25
4396
106
17578
2146
120
2230
743
2146
120
2230
744
NL-50
4396
106
17578
2335
119
2576
859
2335
119
2576
859
NL-75
2607
96
11058
1578
116
1818
606
1.78
1818
607
NL-100
1258
55
7832
1709
53
2378
793
1709
53
2378
793
FANELL
4636
100
10275
4636
100
10275
1055
4752
183
1068S
597
Wiki MTI
10000
10
17178
10000
10
17178
1908
10000
16526
1834
Wiki MTI health
10000
7
14371
10000
7
14371
1596
10000
14110
1566
org
Wiki MT2 sci
10000
16
16471
10000
16
16471
1830
10000
16
14852
1650
27262
27262
3026
Wiki MT3 infra
10000
24
21990
10000
24
21990
2443
10000
21646
2405
Wiki MI4
10000
42
12576
10000
42
12576
1397
10000
12516
1388
Wiki MT4 health
10000
15539
10000
15539
1725
10000
15337
Table 14: Dataset statistics for inductive-e link prediction datasets. Triples are the number of edges given training, validation, test graphs, respectively, whereas Valid and Test denote triples to be predicted the validation and test graphs.
Training Graph
Validation Graph
Test Graph
Dataset
Rels
Entities
Triples
Entities
Triples
Valid
Entities
Triples
Test
FB-v1
180
1594
4245
1594
4245
489
1093
1993
411
FB-v2
2608
9739
2608
9739
1166
4145
947
FB-v3
215
3668
17986
3668
17986
2194
2501
7406
1731
FB-v4
219
4707
27203
4707
27203
3352
3051
11714
2840
WN-v1
9
2746
5410
2746
5410
630
922
1618
373
WN-v2
6954
15262
6954
15262
1838
2757
4011
852
WN-v3
12078
25901
12078
25901
3097
5084
6327
1143
WN-v4
9
3861
7940
3861
7940
934
7084
12334
2823
NL-vI
14
3103
4687
3103
4687
414
225
833
201
NL-v2
88
2564
8219
2564
8219
922
2086
4586
935
NL-v3
4647
16393
4647
16393
1851
3566
8048
1620
NL-v4
76
2092
7546
2092
7546
876
2795
7073
1447
ILPC Small
48
10230
78616
6653
20960
2908
6653
20960
2902
ILPC Large
65
46626
202446
29246
77044
10179
29246
77044
10184
HM Ik
36237
93364
36311
93364
1771
9899
18638
476
HM
32118
71097
32250
71097
1201
19218
38285
1349
HM 5k
11
28601
57601
28744
57601
900
23792
48425
2124
HM Indigo
229
12721
121601
12797
121601
14121
14775
250195
14904
38

---

Published as a conference paper at ICLR 2026
Table 15: Dataset statistics for transductive link prediction datasets. Entity task denotes the entity- prediction task: h/t is predicting both heads and tails, and is predicting only tails.
Dataset
Entities
Rels
Train
Valid
Test
Entity Task
FB15k-237
14541
237
272115
17535
20466
hft
WNI8RR
40943
86835
3034
3134
h/f
CoDEx Small
2034
42
32888
1827
1828
CoDEx Medium
17050
51
185584
10310
10311
h/t
CoDEx Large
77951
69
551193
30622
30622
hift
NELL995
74536
200
149678
543
2818
h/f
YAGO310
123182
37
1079040
5000
5000
h/t
WDsinger
10282
135
16142
2163
2203
b/t
NELL2233
22925
200
25445
4961
4952
h/t
Aristo
44949
1605
242567
20000
20000
hft
DBpedia
99604
470
597572
50000
50000
ConceptNet100k
78334
34
100000
1200
1200
h/
FB (B15k-237(10)
11512
237
27211
15624
18150
FB 1B15k-237(20
13166
237
54423
16963
19776
FB FB15k-237(50)
14149
237
136057
17449
20324
Hetionet
45158
24
2025177
112510
112510
h//
Table 16: Different graph pretraining mix shown
$5.3
8
FB15k-237 WNI8RR
V
V
/
/
4
CoDEx Medium
NELL9995
/
YAGO 310
ConceptNec110K
DBpedial00k
AristoV4
39

---

Published as a conference paper at ICLR 2026
Table 17: Hyperparameter for FLOCK in pretraining and finetuning setups.
Hyperparameter
Entity prediction
Relation prediction
Walk length &
128
128
Random walk
Pretraining hase walk Remin
128
128
Test-time or finetuning base walk nt
16-512
16-512
# Layers
Sequence processor
Hidden dimension
64
Heads
Consensus protocol
Head dimension da
16
16
Update
# Update step
6
Ensemble
# Maximum ensembled passes
16
Optimizer
Adam
AdamW
Learning rate
0.0005
0.0005
Training steps
400,000
40,000 1
Pretraining
Adversarial temperature
# Negatives Batch size
512
512
Weight decay
0.01
0.00
Optimizer Learning rate
AdamW 0,0005
AdamW 0.0005
Finetuning
Adversarial temperature
1
Negatives Batch size
256
256
4-32
4-8
40

---

Published as a conference paper at ICLR 2026
Table 18: Detailed finetuning and inference hyperparameters for FLOCK in entity prediction. For each dataset, report the finetuning epochs, batches per epoch, batch size, and the inference set- tings for both zero-shot and finetuned modes: test-time ensemble size P. base walk count n. For Hetionet finetuning we used (P, n) (1, 1024), instead of (2, 512) as zero-shot.
Dataset
Epoch
# Batch/Epoch
Batch Size
# Ensembled Passes 1'
# Base Walk
FB15k-237
full
8
16
128
WNISRR
full full full
8
16
128
CoDEx Small
16
CoDEx Medium
8
16
128
CoDEx Large
2000
4
2
512
NELL-995
full
8
16
128
YAGO310
2000
512
WDsinger
full
16
NELL233
full
8
16
32
FB1 FB1.5k-23((10)
full
8
16
32
FB1.5k-237(20)
full
16
64
(3755-237(50)
16
Hetionet
4000
512
DBpedial00k
1000
2
512
AristoV4
full
4
256
ConceptNet100k
full
16
128
v1-v4
full
16
WN v1-v4
full
16
NL v1-v4
full
16
16
ILPC Small
full
16
16
ILPC Large
full
16
FB 25 100
full
16
WK 25-100
full
16
16
NI. 0-100
full
16
16
Wiki MT1
full
16
health
16
Wiki MT2 org
full
16
16
32
Wiki MT2 sci
full
8
16
16
Wiki MT3 art
full
16
16
32
Wiki MT3 infra
full
16
16
32
Wiki MT4 sci
full
16
Wiki MT4 health
full
8
16
16
Metafam
full
8
16
16
FBNELL
full
8
16
HM
16
HM 3k
full
16
16
32
HM 5k
full
8
16
64
HM Indigo
full
8
16
128
41

---

Published as a conference paper at ICLR 2026
Table Detailed finetuning and inference hyperparameters for FL.CK in relation prediction. For each dataset, report the finetuning epochs, batches per epoch, batch size, and the inference set- tings for both zero-shot and finetuned modes: test-time ensemble size P and base walk count π.
Dataset
Epoch
# Batch/Epoch
Batch Size
# Ensembled Passes
# Base Walk
FB15k-237
1000
16
128
WNI8RR
16
128
CoDEx Small
1000
8
16
16
CoDEx Medium
1000
8
16
128
CoDEx Large
1000
2
512
NELL-995
1000
16
128
YAGO310
1000
512
WDsinger
1000
8
16
16
NELL223
1000
8
16
32
FB1.5k-33((0)
1000
16
32
FBI5k-255 (37(22))
1000
16
7B15k-237(50
1000
8
16
64
Hetionet
1000
2
512
DBpedia 100k
1000
2
512
V4
1000
256
ConceptNet10Ok
1000
16
128
FB vl-v4
1000
8
16
16
WN v1-v4
1000
8
16
16
NL v1-v4
1000
16
ILPC Small
1000
16
ILPC Large
1000
16
64
FB -100
1000
8
16
16
WK 25-100
1000
8
16
16
NI. 0-0--00
1000
16
MTI tax
16
Wiki MTI health
1000
8
16
Wiki MT2
1000
8
16
32
Wiki MT2 sci
1000
16
16
Wiki MT3 art
1000
16
32
Wiki MT3 infra
1000
16
32
Wiki MT4 sci
1000
16
16
Wiki MT4 health
1000
16
16
Metafam
1000
16
16
FBNELL
16
HM 1k
1000
16
HM 3k
1000
8
16
32
HM 5k
1000
16
64
HM Indigo
1000
16
128
42


Reviews:
[{"review_id": "lGlSq7IcHw", "reviewer": "Reviewer_LNUH", "summary": "An architecture for zero-shot entity and relation prediction. FLOCK adopts a probabilistic approach to the equivariance assumption, as opposed to the conventional deterministic interpretation. The method relies on random walks sampling and direct sequence encoding.", "questions": "- What is the recommended interval of $l$ values for random walks, keeping in mind impact on predictive power and training time? \n- Have you considered using a transformer architecture for the sequence processor?\n- line 236: Can you clarify what do you mean by \"to keep GPU memory usage in a range\"?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "- Research problem/limitation of SoTA KGFMs is well explained and clearly justified.\n- Novelty: FLOCK is the first to adopt a stochastic take on node-relation equivariance. Extending the notion of probabilistic invariance to KGs is more than reasonable.\n- Comprehensive evaluation campaign. Up-to-date baselines choice from literature. Protocol shows fair comparison w.r.t. to baselines.\n- Relation prediction results clearly stronger than conventional KGFMs.\n- Writing is clear and narrative flows well and is coherent.\n- Related work is comprehensive and up to date.", "weaknesses": "- Scalability of random walks on large KGs (the authors are well aware and they have discussed this in the conclusions)\n- I could not find an overall training time and inference time comparison against conventional KGFMs such as ULTRA, TRIX, FLOCK. It is not entirely clear to which extent the adoption of random walks is slower than message passing? (apologies if I have missed)\n- Results: entity prediction results only marginally better than prior art.\n- (minor) Figure 2 could help a better caption, to clarify color coding and actions taken at each step.\n- Some examples would help clarify, e.g. \"They reveal nodes $v_s$ and relations $r_s$ specific to each KG which obstructs transferability to unseen KGs.\" (line 239)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "HApMj8EW6i", "reviewer": "Reviewer_qE2N", "summary": "The work proposes a novel knowledge graph foundation model (KGFM) architecture, FLOCK, with stochastic equivariant node and relation representations. The stochastic equivariance property endows the model with better expressivity without sacrificing zero-shot generalizability to novel KGs with unseen node and relation types. The authors achieve this via an innovative adoption of random walks that anonymize the identities of nodes and relation types, thus keeping the model informed of solely the structural roles that the visited nodes and relation types play in the random walks. In the experimental part, the paper then demonstrates the performance of FLOCK on a wide range of both synthetic and real-world KGs, showing superior performance to the baselines and the capability to distinguish isomorphic relation types for the link prediction tasks on the synthetic PETALS dataset.", "questions": "Please see the Weaknesses section for my questions and concerns.", "limitations": "", "rating": 4, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "1. The paper correctly identifies that most of the existing KGFMs employs deterministic node-relation equivariance, which creates an expressivity bottleneck. The paper then follows the well-established, theoretically sound notion of stochastic/probabilistic equivariance. The introduction of random-walk based methodology to tackle the KG representation is novel, and the proposed model design is both intuitive and well-motivated.\n\n2. The paper creates the new PETALS synthetic dataset to cleverly demonstrate the unique advantage of FLOCK at distinguishing isomorphic relation types. This group of experiments is well-designed, providing targeted empirical validation. The rest of the experiments are also diverse and comprehensive, with FLOCK showing consistently superior performance.", "weaknesses": "W1. **My major concern is the paper's novelty claim of \"the first stochastic KGFM\"**. As identified in one of the previous work (Gao et al. 2023), the InGram model (Lee et al., 2023) can be viewed as a stochastic node-relation equivariant KG model, because it injects random initial node and relation type embeddings, thus creating node and relation embeddings that are equivariant in distribution. Based on this observation, Gao et al. then proposed an improvement, named DEq-InGram, that further stabilizes the stochastic equivariant embeddings of the original InGram, showing significant improvement in performance. Thus, I do not believe that this work may claim that FLOCK is the \"first stochastic KGFM.\" Could the author then clarify this point? If DEq-InGram was indeed a predecessor to FLOCK on the line of stochastic KGFMs, could the author further elaborate how FLOCK could be uniquely advantageous compared to DEq-InGram?\n\n\nW2. **Missing error bars on the empirical results**. In all experiment tables in the main text, only the average numbers are reported. The error bars (e.g. standard deviation numbers) are missing. It would be informative to see the comparison of variance of model performance, particularly for stochastic equivariant models, because the stochasticity nature of the model usually results in higher variance of model prediction accuracy. \n\nW3. **Clarity of writing and lack of critical technical details in methodology**. The writing on the random walks section (line 219 - 229) is a bit confusing to read and could be improved. For instance, the paragraph on line 226 seems to repeat the same message with those discussed in the preceding paragraph. Furthermore, in this preceding paragraph that starts on line 219, it is not clearly explained why $3n$ walks are needed. Could the authors clarify why are there $3n$ walks?\n\nLater in the sequence processor section (line 257 - 269), it is critical for the readers to understand how does the \"structural indices\" such as 1, 2, 3, for nodes and $\\alpha$, $\\beta$ from equation (5) and (6) translates to which neural embeddings in the GRU model from equation (7). This part is critical because later, the analysis of the equivariance-ness of the entire pipeline hinges on it. It would be great if more writing effort is devoted here to clarify and elaborate how the structural indices interacts with the GRU model.\n\n\nW4. **Ablation on the Consensus Protocol**. In the consensus protocol section, the authors made the following claim regarding the approach of taking averages of the proposals:\n\n> The drawback is that uninformative proposals from e.g. dangling regions of walks are not directly suppressed, and can affect the state updates.\n\nThis claim seems to be a conjecture at this stage, without theoretical analysis or empirical evidence. Would it be possible to provide further evidence to support this claim, such as some sort of ablation studies?\n\n\nW5. **Lack of proof insight or proof sketch**. Although this is not necessarily a critical concern or weakness, it would be great if the theoretical section 4.2 have some discussions on the insights or ideas or proof sketch to illustrate intuitively, why does the claims and propositions are mathematically correct. Currently, it reads like the section 4.2 is simply an information dump of related theoretical results without much insights into their correctness.\n\n\nW6. **Notations clarity**: On line 154\n\n> Let $\\omega$ be a function assigning to each KG $G = (V, E, R) \\in \\mathbb{K}_{n, m}$ ...\n\nThe notation $\\mathbb{K}_{n, m}$ is first used before being properly defined and introduced (it was later introduced in Line 158).\n\n\nReferences:\n\n[1] Gao, Jianfei, Yangze Zhou, Jincheng Zhou, and Bruno Ribeiro. \"Double equivariance for inductive link prediction for both new nodes and new relation types.\" arXiv preprint arXiv:2302.01313 (2023).\n\n[2] Lee, Jaejun, Chanyoung Chung, and Joyce Jiyoung Whang. \"InGram: Inductive knowledge graph embedding via relation graphs.\" In International conference on machine learning, pp. 18796-18809. PMLR, 2023.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "l3PdAeaGUj", "reviewer": "Reviewer_w7Ej", "summary": "This work argues that existing Knowledge Graph Foundation Models (KGFMs) are limited by deterministic node-relation equivariance, preventing them from distinguishing structurally isomorphic yet semantically distinct relations. They propose probabilistic node-relation equivariance as a relaxation and introduce FLOCK, a KGFM based on sampling and encoding anonymized random walks using sequence models, without using traditional message-passing. FLOCK is claimed to respect probabilistic equivariance and be a universal approximator for link-invariant functions. They also design a diagnostic dataset (PETALS) to showcase the limitations of prior work, and showcase empirical results suggesting state-of-the-art performance on the standard benchmark of the domain with 54 KGs. Overall, I think this work makes a valuable contribution, and I would be happy to increase my score if my concerns are addressed.", "questions": "1. Beyond PETALS, can you provide examples from the benchmark KGs where the ability to distinguish structurally isomorphic relations proved crucial for FLOCK's performance advantage? How common is this structural ambiguity in real KGs?\n2. The universality proof sketch relies on walks covering all edges. How does the approximation quality degrade in practice when walks inevitably provide only partial coverage on large graphs? Does this partial coverage limit the types of functions FLOCK can effectively learn?\n3. Given the computational overhead of sampling and ensembling, under what specific conditions (graph size, structure, task requirements) does FLOCK offer a clear practical advantage over potentially faster, deterministic KGFMs like TRIX or ULTRA, especially if these were to be augmented with techniques to increase expressivity (subgraph features or textual semantics like [2])? SEMMA [2] approaches the same problem of “losing the ability to distinguish between two entities with opposite relationships”, but by directly encoding the textual features into ULTRA. This at least warrants a discussion.\n4. How was the sequence model architecture (biGRU + RMSNorm + SwiGLU) chosen? Did you experiment with other paradigms (like Transformers or SSMs), and how did they perform in terms of accuracy and efficiency?\n5. Could the proposed probabilistic equivariance be achieved through methods other than random walks (like stochastic message passing, randomized node features)? How does the choice of random walks specifically contribute? Appendix F is great, but more than naive noise injection would make it stronger.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "1. **Problem:** Identifies and demonstrates a clear limitation (expressivity with respect to structurally isomorphic relations) of KGFMs based on strict equivariance using a new synthetic dataset (PETALS).\n2. **Elegant solution:** Proposes probabilistic node-relation equivariance as a potentially more expressive alternative inductive bias and introduces FLOCK, a novel non-message-passing KGFM architecture based on random walks and sequence models.\n3. **Theoretical backing:** Provides theoretical analysis regarding universality and probabilistic invariance (though practical relevance is debatable).\n4. **Empirical Backing:** Shows strong empirical performance on a wide range of KGs, particularly in relation prediction and on the PETALS set.", "weaknesses": "1. **Scalability:** The reliance on random walks raises concerns for large graphs, both in terms of sampling time and ensuring adequate graph coverage. The proposed test-time adaptation is heuristic. Efficiency comparisons (Table 6) show FLOCK can be much slower than baselines.\n2. **Practicality:** The theoretical benefits (universality, distinguishing specific isomorphic cases) might not translate into significant practical advantages across all common KG structures and tasks, especially given the variance introduced by stochasticity. Although the results across the standard benchmark suggested improvement, prior work [1, 2] has shown that this benchmark has potential issues. FLOCK's pipeline (sampling, recording, sequence processing, consensus) is also more complex than standard message passing approaches, potentially hindering adoption.\n3. **Ablation Gaps:** More detailed ablation (like Figure 4b) on the random walk strategy (like length, different samplers) and sequence model architecture would be beneficial.\n\n---\n\n*[1] Harry Shomer, Jay Revolinsky, & Jiliang Tang (2025). Towards Better Benchmark Datasets for Inductive Knowledge Graph Completion. In Proceedings of the 31st ACM SIGKDD Conference on Knowledge Discovery and Data Mining (KDD).*\n\n*[2] Arvindh Arun, Sumit Kumar, Mojtaba Nayyeri, Bo Xiong, Ponnurangam Kumaraguru, Antonio Vergari, & Steffen Staab. (2025). SEMMA: A Semantic Aware Knowledge Graph Foundation Model. In The Thirtieth Conference on Empirical Methods in Natural Language Processing (EMNLP).*", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "672xaNwc8T", "reviewer": "Reviewer_sNAr", "summary": "This paper introduces FLOCK, a novel knowledge graph foundation model for zero-shot link prediction. The authors identify that existing knowledge graph foundation models (KGFMs) rely on deterministic node-relation equivariance, which limits their ability to distinguish between structurally similar but semantically different relations. To address this issue, the authors introduce the concept of probabilistic node-relation equivariance and design FLOCK based on this principle. FLOCK iteratively samples random walks from KGs, encodes them into sequences, embeds them with a sequence model, and aggregates node and relation representations through learned pooling.", "questions": "1. Given the random walk coverage issues, have you explored more targeted random walk strategies rather than uniform random walks? For example, could biased random walks based on graph structure or query information improve model performance?\n\n2. In the transductive setting, FLOCK underperforms baseline models, which seems to contradict the model's theoretical advantages. Could you provide a deeper analysis of this phenomenon and how this issue might be resolved while maintaining the model's expressiveness?", "limitations": "", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The paper's concept of probabilistic node-relation equivariance extends the theoretical framework of knowledge graph foundation models and provides detailed proofs showing FLOCK is a universal approximator for link-level functions.\n\n2. The authors clearly identify limitations of existing KGFMs, particularly in handling structurally similar but semantically different relations, and provide an intuitive Star Wars character relationship example (Figure 1).\n\n3. The paper presents extensive evaluations on 54 KGs across different domains, including both entity prediction and relation prediction tasks, in both zero-shot and fine-tuned settings, with FLOCK outperforming state-of-the-art models in most cases.", "weaknesses": "1. While the authors highlight FLOCK's significant advantage over ULTRA and TRIX on the Metafam dataset in lines 288-290, Table 10's fine-tuning results show all three models achieving nearly 1.000 MRR, contrasting with the significant differences in the zero-shot setting. This inconsistency is not explained.\n\n2. While Section 4.2 provides a thorough theoretical analysis, the paper doesn't adequately discuss how these theoretical advantages directly explain specific performance differences observed in experiments, particularly across different types of knowledge graphs.\n\n3. On certain datasets (e.g., WK-50, Table 8 line 1570+), FLOCK performs worse than ULTRA and TRIX, contradicting the overall conclusions, but these anomalies aren't analyzed or explained.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)