EDEN Distills a Directed Graph’s Own Hierarchy Into Better GNN Predictions

The gap between data-centric and model-centric graph learning

Graph neural networks have posted strong results on node classification, link prediction and graph level tasks for years now, but the field’s research energy has overwhelmingly gone into model architecture. New attention mechanisms, new convolution operators, new ways of combining incoming and outgoing edge information. What has received far less attention is a more basic question. Is the data itself, the raw topology and node features, actually being used to its full potential before a model ever sees it.

That framing matters more for directed graphs than for undirected ones. In an undirected graph, an edge between two nodes is symmetric, a friendship is a friendship in both directions. In a directed graph, a citation, a trust relationship, or a hyperlink points one way, and that asymmetry carries information a symmetric model architecture cannot easily recover after the fact. Directed graph neural networks, DiGNNs, were built specifically to preserve that asymmetry, using separate learnable weights for incoming and outgoing neighbors. The EDEN authors argue that even the best DiGNNs are still only scratching the surface, because they focus on how to process directed edges rather than on what deeper structural knowledge those edges are hiding.

Turning a graph into a teaching hierarchy

EDEN’s core object is something the paper calls a Hierarchical Knowledge Tree, or HKT. The intuition is closer to organizational structure than to a typical graph neural network layer. Picture the raw nodes of the graph as employees at the bottom of a company. Above them sit team leads, who each represent a cluster of employees who work closely together. Above the team leads sit department heads, and so on up to a single root representing the whole organization. Each level up the tree is a coarser summary of the level below it, and information should be able to flow both ways, specific knowledge climbing up to inform the summary, and general knowledge flowing back down to correct individual nodes.

EDEN builds exactly this kind of structure out of a directed graph, using two separate signals in sequence.

Step one, measure structural entropy from topology alone

The first pass at building the tree uses only the graph’s connectivity pattern, no node features involved yet. The authors borrow the concept of structural entropy, a way of quantifying how much uncertainty or disorder exists in a network’s connection pattern, and adapt it specifically for directed graphs by tracking in degree and out degree separately.

That one dimensional measurement gets extended to a hierarchical, multi level version that scores how well a given partition of nodes into groups captures the graph’s true underlying organization, with lower entropy meaning a cleaner, more informative grouping. A greedy algorithm then searches for the tree structure that minimizes this entropy, essentially asking, at every step, which two nodes or groups should merge next to reduce disorder the most.

One practical wrinkle worth flagging. Directed graphs are usually not strongly connected, meaning a random walk starting from one node frequently cannot reach large portions of the graph by following edge directions alone. The authors found that the proportion of walks able to continue drops sharply after only about five steps in most of their test graphs. Their fix is to let the structural entropy calculation incorporate reverse walking probability alongside forward walking probability, rather than sticking strictly to edge direction, which keeps the entropy measurement informative even in graphs where directed paths dead end quickly.

Step two, refine the tree using node profile similarity

Topology alone produces what the paper calls a coarse grained tree. It groups nodes by connection pattern, but two nodes can be topologically similar while having completely different features and labels, or vice versa. To fix this, EDEN adds a second pass that uses mutual information between node profiles, features and labels together, to reassign nodes that were placed in the wrong branch by the topology only pass.

This is where the paper’s more involved theoretical work comes in. Directly computing mutual information between high dimensional node features is intractable, so the authors build on the MINE approach to mutual information neural estimation, first framing it through a KL divergence lower bound, then showing an equivalent GAN style formulation is more practical to train. The end result is a neural network that scores how well a candidate node represents its assigned partition, based on the dependency between that node’s features and the aggregated features of its neighborhood inside the tree.

Concretely, the method looks at two kinds of node moves. It identifies nodes that already sit in a partition but have unusually high mutual information with that partition, marking them as strong representatives worth trusting more. It also scans nodes sitting in other partitions to see if any of them would actually be a better fit here, effectively correcting mistakes the topology only pass made. The result is a fine grained tree that reflects both connection pattern and node content.

Parent and child as teacher and student

Once the tree exists, EDEN treats it as a live training structure rather than a static preprocessing artifact. Parent nodes in the tree, which represent coarse summaries of a group, act as teachers. Child nodes, closer to the original graph, act as students. Crucially, this is not the standard offline distillation setup where a frozen large model hands down fixed soft labels. Both parent and child representations keep updating throughout training, an arrangement the paper calls online knowledge distillation, so teacher and student evolve together rather than one being fixed before the other starts learning.

Two operators drive this exchange. A knowledge generation operator computes each parent’s representation as a weighted combination of its sampled children, using affinity scores that favor nodes which best represent the current partition while still injecting some diversity from other partitions to avoid overfitting to an overly narrow view. A knowledge transfer operator then sends a personalized version of that parent representation back down to each child, normalized so that the transfer accounts for how confidently the parent representation is expressed in the first place.

For the final prediction step, EDEN does not simply read off the leaf node representations directly. It runs a tree structured random walk that samples from a node’s parents, siblings and children, with the sampling bias tuned differently depending on the task. Node classification benefits from leaning on parent representations for a richer, higher level view of the node’s likely class, while link prediction tasks benefit more from sampling siblings, since two nodes at the same tree level carry more directly comparable contextual information for deciding whether an edge should exist between them.

Two ways to use the same idea

A detail that gives EDEN more practical reach than a typical new architecture paper is that it was designed from the outset to work two different ways. Used on its own, EDEN is a complete data centric alternative to a directed graph neural network, built entirely around the hierarchical tree and its digraph learning function. Used differently, EDEN becomes what the authors call a hot and plug online distillation module, meaning an existing DiGNN, whatever a team already has in production, can supply the digraph learning function at each level of the tree, and EDEN’s hierarchical knowledge transfer simply layers on top to improve that existing model’s predictions without requiring a redesign.

What the experiments actually show

The evaluation is unusually broad for a single paper, spanning 14 datasets, a mix of homophilous graphs where connected nodes tend to share labels and heterophilous graphs where they do not, across four distinct tasks, node classification and three flavors of link prediction, existence, direction and a three way link classification.

As a standalone data centric method, EDEN posts the best result on every dataset in this comparison, with the authors reporting average gains of 2.78 percent on node level tasks and 2.24 percent on link level tasks over the strongest baseline available for each case. As a plug in module added to five existing graph neural networks across eight datasets, including undirected ones like Photo, Computer, PPI and Flickr, the improvement ranges from 2.54 percent for OptBasisGNN up to 4.68 percent for Dir-GNN, with the largest gains generally showing up on directed graphs rather than undirected ones, consistent with the paper’s argument that digraphs simply contain more exploitable structural knowledge to begin with.

What the ablation study rules out

Stripping out individual pieces of EDEN one at a time confirms none of them is dead weight. Removing the diverse knowledge sampling that pulls in nodes from other partitions increases overfitting and lowers accuracy on Tolokers and Slashdot alike. Removing the personalized, node adaptive transfer step and replacing it with a uniform transfer also costs accuracy, confirming that treating every child node identically loses information. Removing the tree based random walk for leaf prediction and falling back to a simpler prediction scheme hurts performance across every dataset tested, and removing the knowledge distillation loss entirely produces the largest single drop of the four ablations, underscoring that the distillation mechanism itself, not just the tree structure, is doing real work.

Where robustness held up and where it did not

A separate set of experiments stress tests EDEN and its baselines under three kinds of data scarcity, missing node features, missing edges, and fewer labeled examples per class. The results are a useful reality check rather than a uniform win. Under feature sparsity, methods that lean heavily on the sheer quantity of node representations, D-HYPR and NAGphormer specifically, degrade noticeably, while DiGCN and MGC hold up better because their high order feature propagation can partially compensate for missing values elsewhere in the neighborhood. Under edge sparsity, every baseline suffers since they all depend on decent topology to power their architectures, but EDEN’s data centric knowledge mining keeps it ahead of the pack even as edges disappear. Label sparsity follows a broadly similar pattern to feature sparsity. The overall picture is that EDEN improves robustness on average without claiming immunity to every kind of missing data, and the paper is specific enough about which baselines struggle where that a reader can judge the comparison rather than just take the summary claim at face value.

The efficiency argument, and its limits

Because EDEN’s tree construction step is independent of the model training loop, the authors point out that its computational cost can be paid once and amortized, rather than repeated every epoch the way ordinary model training cost is. The paper also describes a deliberately lightweight implementation path, Monte Carlo based approximate tree construction instead of running the full greedy algorithm exactly, incremental training with prototype representations for the fine grained mutual information step, and weight free feature propagation for the digraph learning function used at each tree layer. Reported efficiency figures on the Empire dataset show the lightweight preprocessing step cutting tens of seconds off tree construction time, and the lightweight training path running roughly 40 percent faster on average than the compared DiGNN and GNN baselines, while still beating their accuracy.

That efficiency story comes with an honest asterisk the authors attach themselves in their own conclusion, discussed below. A framework built around a hierarchical tree, a neural mutual information estimator, and a tree structured random walk carries more moving parts than a single graph convolution layer, and the paper is explicit that this complexity has not yet been fully tamed for the largest graph scales.

Why this matters beyond one leaderboard

Strip away the specific tree construction algorithm and EDEN is making an argument that applies well past directed graphs specifically. Model centric research keeps asking how to build a more powerful function that maps inputs to outputs. Data centric research asks a different, often neglected question, whether the representation being fed into any model already contains structure that a generic architecture is failing to exploit on its own. For directed graphs specifically, that structure includes the asymmetry between incoming and outgoing connections, which a naive undirected treatment throws away entirely, and the deeper hierarchical organization that both topology and node content jointly encode but that neither one captures completely alone.

The hot and plug design decision is arguably the paper’s most practically important contribution, separate from the raw accuracy numbers. A method that only works as a full architectural replacement asks a team to rip out whatever DiGNN they already have in production. A method that layers cleanly on top of an existing model, improving it without requiring a redesign, is a fundamentally easier sell inside an organization that already has infrastructure and monitoring built around a specific model. Whether EDEN specifically becomes that adopted layer depends on the complexity concerns the authors themselves raise, but the general pattern, distill a model against structure mined from its own data rather than against a separate teacher network, seems like a genuinely reusable idea for other structured data domains beyond graphs.

PyTorch implementation

The following code reconstructs the core mechanics of EDEN at a manageable scale, a simple entropy guided greedy tree builder, a GAN style mutual information estimator following Theorem 3.3, and the parent to child knowledge distillation loss from Equation 10, combined with ordinary cross entropy for node classification. It runs as a self contained smoke test on a small synthetic directed graph.

# eden_digraph_kd.py
# Reconstructed implementation of EDEN, Entropy-driven Digraph knowlEdge
# distillatioN, a data-centric hierarchical knowledge distillation
# framework for directed graph learning.
# Reference: Li, Wu, Yu, Qin, Zeng, Li and Wang, "Toward Data-centric
# Directed Graph Learning: An Entropy-driven Approach", ICML 2025.

import torch
import torch.nn as nn
import torch.nn.functional as F


def directed_structural_entropy(in_degree: torch.Tensor, out_degree: torch.Tensor, num_edges: int):
    """Equation 1. One dimensional directed structural entropy from
    the stationary distribution of in and out degrees."""
    d_in = in_degree.clamp_min(1e-8) / num_edges
    d_out = out_degree.clamp_min(1e-8) / num_edges
    term_in = -(d_in * torch.log(d_in)).sum()
    term_out = -(d_out * torch.log(d_out)).sum()
    return (term_in + term_out).item()


def greedy_hkt_merge(in_degree: torch.Tensor, out_degree: torch.Tensor, num_edges: int, target_groups: int):
    """Simplified stand-in for Algorithm 2's greedy partition tree
    construction. Starts with every node as its own group, then
    repeatedly merges the pair of groups whose combined degree
    profile most reduces overall structural entropy, until only
    target_groups remain. This captures the entropy-minimizing
    spirit of the paper's HKT construction at a scale small enough
    to run without the full topology-aware machinery."""
    n = in_degree.shape[0]
    groups = [[i] for i in range(n)]

    def group_entropy(members):
        gi = in_degree[members].sum()
        go = out_degree[members].sum()
        return directed_structural_entropy(gi.unsqueeze(0), go.unsqueeze(0), num_edges)

    while len(groups) > target_groups:
        best_pair, best_entropy = None, float("inf")
        for i in range(len(groups)):
            for j in range(i + 1, len(groups)):
                combined = groups[i] + groups[j]
                candidate_entropy = group_entropy(combined)
                if candidate_entropy < best_entropy:
                    best_entropy = candidate_entropy
                    best_pair = (i, j)
        i, j = best_pair
        merged = groups[i] + groups[j]
        groups = [g for k, g in enumerate(groups) if k not in (i, j)]
        groups.append(merged)
    return groups  # each inner list is a coarse-grained partition (parent)


class MutualInfoEstimator(nn.Module):
    """Equation 6. GAN-style mutual information neural estimator
    between a node and its generalized neighborhood, used to refine
    the coarse-grained HKT into a fine-grained one."""

    def __init__(self, feature_dim: int, hidden_dim: int = 64):
        super().__init__()
        self.w1 = nn.Linear(feature_dim, hidden_dim)
        self.w2 = nn.Linear(feature_dim, hidden_dim)
        self.score = nn.Sequential(
            nn.Linear(hidden_dim * 2, hidden_dim),
            nn.ReLU(inplace=True),
            nn.Linear(hidden_dim, 1),
        )

    def forward(self, node_features: torch.Tensor, neighborhood_features: torch.Tensor):
        node_emb = self.w1(node_features)
        neighborhood_emb = self.w2(neighborhood_features)
        joined = torch.cat([node_emb, neighborhood_emb], dim=-1)
        return self.score(joined).squeeze(-1)  # raw logits, sigmoid applied by the loss

    def mi_loss(self, node_features: torch.Tensor, true_neighborhood: torch.Tensor, shuffled_neighborhood: torch.Tensor):
        """Positive pairs: node with its real generalized neighborhood.
        Negative pairs: node with a mismatched neighborhood, built by
        shuffling true_neighborhood across the batch."""
        positive_logits = self.forward(node_features, true_neighborhood)
        negative_logits = self.forward(node_features, shuffled_neighborhood)
        positive_loss = F.logsigmoid(positive_logits).mean()
        negative_loss = torch.log(1 - torch.sigmoid(negative_logits) + 1e-8).mean()
        return -(positive_loss + negative_loss)  # maximize MI lower bound -> minimize this


class DigraphLearningFunction(nn.Module):
    """A minimal weight-free-style embedding function, standing in
    for the lightweight layer-wise digraph learning function M(.)
    described in Section 3.4. Any real DiGNN can be substituted
    here without changing the distillation logic below."""

    def __init__(self, in_dim: int, hidden_dim: int):
        super().__init__()
        self.proj = nn.Linear(in_dim, hidden_dim)

    def forward(self, x: torch.Tensor):
        return F.relu(self.proj(x))


def knowledge_distillation_loss(parent_repr: torch.Tensor, parent_confidence: torch.Tensor, child_repr: torch.Tensor, child_projector: nn.Module):
    """Equation 10. Parent representation, normalized by its own
    confidence, is matched against a personalized projection of
    the child representation using a Frobenius norm."""
    normalized_parent = parent_repr / parent_confidence.clamp_min(1e-6)
    projected_child = child_projector(child_repr)
    return torch.norm(normalized_parent - projected_child, p="fro")


def eden_train_step(
    model: DigraphLearningFunction,
    mi_estimator: MutualInfoEstimator,
    child_projector: nn.Module,
    classifier: nn.Module,
    optimizer: torch.optim.Optimizer,
    node_features: torch.Tensor,
    labels: torch.Tensor,
    parent_group_indices: list,
    alpha: float = 1.0,
):
    """One training step. L = L_cross_entropy + alpha * L_kd, per
    Equation 12, plus the mutual information loss used to keep
    refining which nodes belong to which partition."""
    model.train()
    optimizer.zero_grad()

    node_embeddings = model(node_features)
    logits = classifier(node_embeddings)
    ce_loss = F.cross_entropy(logits, labels)

    total_kd_loss = torch.tensor(0.0)
    total_mi_loss = torch.tensor(0.0)
    for members in parent_group_indices:
        members_t = torch.tensor(members, dtype=torch.long)
        child_repr = node_embeddings[members_t]
        parent_repr = child_repr.mean(dim=0, keepdim=True).expand_as(child_repr)
        parent_confidence = torch.sigmoid(parent_repr.abs().mean(dim=-1, keepdim=True)) + 0.5

        total_kd_loss = total_kd_loss + knowledge_distillation_loss(
            parent_repr, parent_confidence, child_repr, child_projector
        )

        shuffled = child_repr[torch.randperm(child_repr.shape[0])]
        total_mi_loss = total_mi_loss + mi_estimator.mi_loss(child_repr, parent_repr, shuffled)

    total_loss = ce_loss + alpha * (total_kd_loss + total_mi_loss)
    total_loss.backward()
    optimizer.step()
    return total_loss.item(), {"ce": ce_loss.item(), "kd": total_kd_loss.item(), "mi": total_mi_loss.item()}


def smoke_test():
    """Builds a tiny synthetic digraph, constructs a coarse HKT with
    the greedy entropy merger, then runs one EDEN-style training
    step combining classification, distillation and MI losses."""
    torch.manual_seed(0)
    num_nodes, feature_dim, hidden_dim, num_classes = 16, 8, 32, 3

    node_features = torch.randn(num_nodes, feature_dim)
    labels = torch.randint(0, num_classes, (num_nodes,))

    # synthetic in/out degree sequence and a rough edge count
    in_degree = torch.randint(1, 5, (num_nodes,)).float()
    out_degree = torch.randint(1, 5, (num_nodes,)).float()
    num_edges = int(in_degree.sum().item())

    parent_groups = greedy_hkt_merge(in_degree, out_degree, num_edges, target_groups=4)
    print(f"built {len(parent_groups)} coarse-grained partitions")
    for idx, group in enumerate(parent_groups):
        print(f"  partition {idx}: {len(group)} nodes")

    model = DigraphLearningFunction(feature_dim, hidden_dim)
    mi_estimator = MutualInfoEstimator(hidden_dim)
    child_projector = nn.Linear(hidden_dim, hidden_dim)
    classifier = nn.Linear(hidden_dim, num_classes)

    optimizer = torch.optim.Adam(
        list(model.parameters())
        + list(mi_estimator.parameters())
        + list(child_projector.parameters())
        + list(classifier.parameters()),
        lr=1e-3,
    )

    loss, parts = eden_train_step(
        model, mi_estimator, child_projector, classifier, optimizer,
        node_features, labels, parent_groups, alpha=0.5,
    )
    print(f"total loss {loss:.4f}")
    for name, value in parts.items():
        print(f"  {name}: {value:.4f}")

    assert torch.isfinite(torch.tensor(loss)), "loss is not finite"
    print("smoke test passed")


if __name__ == "__main__":
    smoke_test()

Conclusion

What makes EDEN worth paying attention to is less any single accuracy number and more the reframing it insists on. Knowledge distillation has almost always meant compressing a large model into a small one. EDEN keeps the machinery of teacher and student, soft targets and a distillation loss, but points it at a different source entirely, the hierarchical structure that a directed graph’s own topology and node profiles already encode, whether or not anyone ever trains a large model on that graph at all. Parent and child in the hierarchical tree take on the teacher and student roles usually reserved for a big network and a small one, and both sides keep learning together rather than one being frozen ahead of the other.

The theoretical scaffolding behind this, structural entropy adapted for directed edges, and mutual information estimated through a GAN style neural network rather than computed directly, does real work rather than existing for its own sake. The entropy measurement is what lets the tree respect the asymmetry that makes directed graphs different from undirected ones in the first place, and the mutual information refinement is what stops the tree from being fooled by topology that looks clean but groups together nodes with genuinely different underlying content.

The dual design as both a standalone method and a hot and plug module is the detail most likely to matter for anyone actually building on this work. It means the paper’s contribution is not tied to convincing a team to abandon whatever directed graph neural network they already trust in production. A model already in use can, at least in principle, keep its architecture and gain EDEN’s hierarchical distillation on top, which is a meaningfully lower barrier to trying the idea than a full architectural replacement would be.

None of that erases the complexity the authors themselves flag. A tree construction step, a neural mutual information estimator, and a random walk based prediction stage is a lot of machinery compared to a single graph convolution, and the paper’s own conclusion names scalability to billion node graphs as unsolved even with the lightweight variant. The honest reading of this work is a genuinely useful new tool for datasets in the range tested here, citation networks, social networks, web graphs up to a few million nodes, alongside an open engineering problem for anyone hoping to run the same idea at web scale.

The broader lesson likely outlasts the specific tree algorithm. Before reaching for a bigger model or a more exotic architecture, it is worth asking whether the data already in hand contains structure a simpler model is failing to use. For directed graphs, that structure turned out to be a hierarchy hiding in the entanglement of topology and node content. For other structured data domains, the equivalent hiding place will look different, but the instinct to look for it before scaling up the model is the part of this paper worth carrying elsewhere.

Frequently asked questions

Related reading