GNN for Trees

[ilyalasy]

Hey, what algorithms should we use when we exactly know that our dataset consists of not arbitrary graphs but trees?
I have an example:
My dataset contains of DOM trees, and I'm solving node classification problem. I tried to use GATConv which works fairly well, with one exception: it can classify one element correctly but also it will classify all (most) its children with the same class. This behaviour is not ideal and it leads to a lot of false positives as only parent elements are labeled.
What layer/model should one use to classify only parent nodes?
Thank you!

[rusty1s]

If I understand you correctly, this question is more related on how to evaluate/train the model rather than which layer to use inside a model. For example, you can integrate a node-level parent_mask into your data that denotes which node corresponds to a parent and not to a child, and then only evaluate/train on parent nodes:

out = model(data.x, data.edge_index)[data.parent_mask]
F.cross_entropy(out, data.y[data.parent_mask]

[ilyalasy]

Yes, but to some extent we can say that most of such nodes will be a parents, because its a nested structure. E.g.

-N1 <-- node i want to classify, mask should be True
--N2 <-- mask should be False, but at the same time its a parent to N3
---N3

I also interested, aren't there any models that works better on such structures than on arbitrary graphs?

[rusty1s]

I don't have any particular pointers for you regarding GNN design, but I suggest to look into GNNs that operate on ASTs for understanding code. This should give you a pretty good overview on alternative GNN formulations on tree-structured data.

[ilyalasy]

Thanks for the info, this is definitely helpful