Best way to scale GNN transformer down?

[inwaves]

Hi!

I'm trying to scale the GNN transformer model down so I can run it with less compute. I can change the GNN embedding dimension, number of layers, transformer feedforward dimension etc., but my network remains relatively large in number of parameters.

I think this is mostly from the last set of layers in the model, the Linears that I think correspond to the readout function. These layers' output dimension seems tied to the dataset; how can I scale these down?

Thanks!
Andrei