Combining Allegro with an MLP that takes in molecular information

[varunpatro1]

Hello,

I'm working on an idea by which I might have Allegro train on atomic information as it's originally designed to, and simultaneously, jointly train an MLP on molecular information (in my case obtained from rdkit on the qm9 set). I'd concatenate the last (or possibly a different) layers of both Allegro and my MLP and then send them through a small prediction MLP to estimate energy.

Here is a diagram from the CheMixNet (2018) paper which does a similar task using SMILES and molecular fingerprints:

I'm wondering how/where I might begin this development -- specifically, where I could do the concatenation of the final layers of each model. I'm initially thinking that I could remove the final projection layer of Allegro and then concatenate its output with that of the MLP. I'd project this concatenation through another MLP and then backpropagate through both models. Could you please guide on how I might go about this? Any help would be greatly appreciated. Thank you very much.