GNN Project #4.2 - GVAE Training and Adjacency reconstruction

Key Takeaways

hello everyone today i want to talk a bit about my progress with the graph variational autoencoder and want to share some learnings regarding the generation of new molecules all of the code is available in a new github repository which is linked in the video description molecule generation is a pretty tough topic as first of all the space of new molecules is huge and secondly several validity constraints need to be considered when generating new molecules i read a couple of papers to get an overview of where the current research is and i quickly want to share two valuable resources here first of all peta villichkovic has published some slides about gnns for molecules on his website in this presentation there is also a section about generative graph models this section provides a pretty structured overview of different papers prominent examples are for instance the junction tree variational auto encoder or the graph variation auto encoder then i found this survey on deep generative models for graph generation it's a bit more general and doesn't focus on graph neural networks but summarizes the whole research field the methods for generating graphs can be divided into one shot and iterative approaches one-shot methods output the whole molecule representation in one run as these plots show we can see that these models for instance using a fully connected network directly output the matrices such as the adjacency matrix iterative models on the other hand sequentially build a graph for instance the first node is added and then the information such as the node embeddings are updated and the next node is added and so on and the unit of each element that is added can also be larger for instance predefined substructures of a molecule such as ch2 groups this is typically called a motif in that context both kinds of methods so iterative and one shot have their advantages and disadvantages iterative models for instance can be pretty slow and more complex to implement but they can generate much larger graphs as they keep updating the context one-shot methods are pretty simple to implement but can usually only generate smaller molecules for this video series i decided to try out one shot generation as it's much easier to implement as mentioned in the last video the first attempt is to reconstruct the adjacency matrix so let's have a look at the architecture of my model the encoder are simply three graph transformer layers which create embeddings of size 32. then i split the network to output mu and sigma separately with an additional gene n layer i chose a latent vector size of 8. it would also be possible to do that last step with a fully connected layer in between these layers i apply batch normalization then for each node in the graph i sample from the latent distribution which is defined by mu and sigma that's currently the special part of this architecture because later the plan is to have only one latent embedding and here we have one per node however the layers that output mu and sigma are shared anyways so each node passes through the same network later we could simply apply some pooling before passing through these layers and then we have only one latent representation for the decoder i originally planned to use the inner product decoder which i mentioned in the last video however i thought it's better to have another learnable component in this model that's why i built all possible combinations of the adjacency matrix and pass them through another fully connected network in my example it has three layers with a size of 128. the output is simply if there is a connection between two nodes or nots so zero or one for the future it will be easy to change that to also output different edge types so not only this binary value so one more remark here so these combinations of vectors they have a source node and a target node and the prediction of the model is if there is a connection between these two nodes so this these combinations are supposed to stand for the adjacency matrix so all combinations of potential connections of course i faced some challenges when training this and i had to do further refinements the first problem was that the predicted full adjacency matrix led to some inconsistencies because it is supposed to be symmetric for a molecule we have undirected edges and therefore we know that the adjacency matrix will be mirrored along the diagonal that on the other hand means that we can also consider this in the model that's done by predicting only the upper triangular part of the matrix we can also ignore the diagonal itself because we don't have self loops in the molecules and this significantly reduces the number of values we need to predict so the number of combinations from source to target nodes still the model had a tendency to always predict zero because typically there are more zeros than ones in the upper triangular part to account for this i added a weight term to the binary cross-entropy loss to consider this class imbalance what happened is that the model now either always predicted one or always predicted zero depending on the weight value finally i calculated this weighted loss separately per graph which means i had an individual weight factor depending on how many connections a graph had with this the model learnt much much better by the way another approach that i've also tried is negative sampling that simply means we don't take all of the zero values from the adjacency matrix but only instead as many as edges we have so the same number of ones with the same number of zeros this way the loss term is also sort of weighted so with these adjustments the training worked much better but there was still one more thing because at the moment i had only a batch size of one and batching is usually used to speed up the training and smoothen the loss and when i had this batch size one so i processed one graph after the other and then back propagated the loss then it led to a very unstable training and therefore i implemented a simple batching strategy which simply loops over all graphs and then back propagates the loss for the whole batch the last thing that didn't work out of the box for my model is the balancing of the loss terms in the variational autoencoder loss function as you know the loss consists of a reconstruction term and a regularizer for the latent space however when training with this loss function my model didn't learn anything because the kl divergence was somehow too restrictive for the model to resolve this i added a weight term to the kl divergence part which is called beta here actually there is also a paper about this called beta vae and in the paper they find a correlation between the size of the latent space and the optimal beta values as my latent vector is also relatively small i also chose a small weight for the beta and the model suddenly started to reproduce the inputs so that's the theory part behind everything i did and now let's have a look at the code so that is the project and we see we have a train function for training the model then this gvae file which is the model itself the data set is the same as in the previous parts so it's this hiv data set with those 40 000 molecules i think and i use this featurizer which i also talked about in the last videos to generate the same features as i used for the hiv prediction task then i also added this config file in which i specify the device and at the moment it's cpu only because i have some issues with my gpu however there's no big speed difference between using gpu or cpu okay now let's have a look at the strain function which is pretty simple we'll load the train data set the test data set and the train and test loader with a batch size of 32 at the moment i've limited the data set to only use the first 10 000 molecules and on the test set i only test on the first 4 thousand then i specify the model and also the loss function the optimizer and this kl divergence better weight which i just talked about and i have this function run one epoch which is called for the train and test section so here is the train part and every five epochs i call this function for the test data so i pass in the test loader and here i pass in the train loader and all this function now does is simply calculate or store some metrics which i lock later with ml flow again and inside of this function i just iterate over the specified loader and then i pass in everything to the model and i calculate the loss here it is and i back propagate it here we can see it and calculate some metrics so nothing fancy it looks like in every other pytorch implementation so next let's have a look at the model file here in the init function i specify all of the layers we need so we see a couple of encoder layers and then this additional split layer which outputs the mu and sigma which is another gnn layer for all of these layers i chose the transformer conf layer and i pass in the edge dimension so that i can also use the edge features i chose four attention heads and some other modifications and then here are the decoder layers are simply some linear layers and in between i have batch normalization again and the output is one value which will be used to specify if we have a connection or not as you know every pytorch model calls the forward function when you pass in data so if we have a look at the trend function one more time in this step where we pass in the node features the edge features the adjacency information and the batch info when that happens all that happens is that we call this forward function and here we have now three steps so first we encode everything and the output of this is the mu and sigma so we have the information about the distribution then i call this reparametrization function which does nothing else but sample using this mu and sigma and then third of all we have the decode function and the output of this are the triangular upper part of the matrix and the logic values for that part so let's quickly have a look at the encode function it's pretty simple all we do is we call these message passing layers and in the last layer we separately get mu and sigma and actually this is called log variance because the idea is to get more numeric stability by applying the log here and later when we use it so in this reprimatization function we apply the exponential function to get back the original value for mu and sigma so this is just a trick that is commonly done so that's all about the encode function we get those two as the output of the last gnn layer and then i call this function and this torch rent and like does nothing else but sample so this is this epsilon if you remember from the last video and now we have the mu the epsilon and the standard deviation and all together we can calculate the sampled values so the sampling takes place outside of the network so that we can back propagate and don't have to handle this random node and then we have the decoder and i thought for the decoder let's use the debugger to see what's going on because it will be much easier and then we can step through the the process in the decoder so here we are at the breakpoints and now let's have a look at this decoder function so if we go here we can see first of all that i implemented this for a batch of graphs and we have this graph index which looks like this so it's basically a map that tells us for each node so each value represents one node in the batch it tells us to which graph this node belongs so these nodes all belong to the first graph and these nodes all belong to the second graph and so on and using this batch index i can now specify a graph mask let's quickly go in here and this mask now basically masks out everything except for the graph in the batch i'm interested in for the first graph id i only get the first graph and then the second one and so on and in order to get the predictions i first need to construct the input features and i said i have all combinations for the triangular part of the adjacency matrix of source and target nodes so i get using this trio indices which is a predefined function from torch i get the indices for the triangular upper parts and so essentially in this graph we have so if we check this shape we have 46 notes and this trio indices now gives me the edge indices for these nodes and what i do now is i calculate or i fetch the source features here and the target features for each of these source and target indices that i get out of this edge indices so it sounds a little bit complicated but actually all it does is if we look at this it fetches so these are the features this is the feature of node 0 and then we have again node 0 node 0 so several times node 0 and here in the end we have node 43 43 those are these two and i have 44 that's the last one and then i do the same for the target features and so this is the second row the first the second the third and the the corresponding node features and now all i do is i concatenate them so i append them and now the graph inputs are simply let's check the shape 1035 notes or pairs of notes and each of them have a dimension of 32 and that's exactly what's going into my model and here we simply pass through these decoded layers and what we see in the end here is the edge logics and we get for each pair a value and when we later apply sigmoids we get probabilities and we can use them to say if there's a edge or not so i hope that was not too confusing there's quite some data wrangling or tensor wrangling you need to do and now we can have a look at this utils function we have also some quite some helper functions here but the main important part here is that we have the loss function and now all the loss function does is with those logits we got from the decoder we can now calculate for each graph the individual loss and all i do here is i get the target values so the target triangle the upper part of the adjacency matrix and i get the predictions for this specific graph so remember that these trial logics here are for the whole batch that's why i need to slice them from the batch for a specific graph id and that's what's happening in those two slicing functions which i have specified somewhere here it's nothing fancy it's just also calculating this graph mask and then using it to get the corresponding values in the batch and i do that for the targets for the predictions and then i have the binary cross entropy loss here and i calculated between those two and additionally i specify a weight here which says how many positive edges do we have in this graph and that's all and this gives me the graph reconstruction loss for one graph in the batch so i append all of them to the batch reconstruction loss which is then essentially nothing else but the average of each of the graph losses and then i use the k or i calculate the kl divergence loss this is happening somewhere here and i got this formula from pytorch geometric they've also implemented the kl divergence and essentially this is the closed form for the kl divergence for normal distributions and using that i can then calculate the full loss which is the kl divergence multiplied with this beta term and then the batch reconstruction loss now there's one more thing i do in the strain function and this is so we have seen the loss function we have seen the model the only thing that is missing is this reconstruction accuracy function and this is supposed to give me an accuracy value that tells me how many of the or which percentage of the edges are correctly predicted and for that i have a function in this uts file which i call reconstruction accuracy it gets the triangular upper logits and the workflow in this function is pretty similar as in the loss so i slice everything for each of the graphs in the batch and now what i do is i calculate the accuracy and this is basically happening here so between the batch targets and the actual predictions and what i also wanted to know is if we are able to reconstruct with 100 percent which molecules are those and this is what's happening in this check trio graph reconstruction function so i hope it's not too confusing uh but essentially so i loop over each of the graphs and now i want to check if they are able to reconstruct with 100 percent so i get the predictions here for the specific graph i get the labels for the specific graph and now i compare them with this torch equal and if they are equal then this means uh we have 100 accuracy and then what i do is i generate the original molecule out of this so basically i use the predictions to get back the molecule and this is done in this graph representation to molecule function and what i do here is i create an empty molecule and then i add all of the atoms in the node types and then i add all of the all of the edges or bonds with this at bond function which are in the predicted adjacency triangular upper and what this gives me is a full molecule object and what i do during training if a molecule was successfully reconstructed i print out the smart string and at the moment i use bond type unspecified and this will create these tilts in the smart string because this they are basically a placeholder that tell us there is some connection but at the moment it's not sure which one and for the next video i plan to replace this with the actual bond type and at the moment i support these atomic numbers and in the data set so in this featurizer every other atom type than those is specified with other and i've mapped this other to a minus one and here in this function i check if -1 is in those atom types and if that is the case i say okay it was successfully reconstructed but there are some unsupported nodes because with this other i cannot map it back to a molecule so i will only generate the molecules which have these atoms and by the way hydrogen is implicit so it's not specified here and if that is the case i can reconstruct the molecule in this function okay and now i would say we can run the training so that we can see what's going on and while this is running we will see that it starts to reconstruct molecules and the sizes of these molecules would get bigger and bigger so it starts with reconstructing very small molecules and then it should start to reconstruct larger ones and those 300 steps here are 300 times the batch size 32 so that's my full data set and at the moment we see no prints because nothing is reconstructed in the meantime we can have a look at the logs of previous runs so let's have a look at ml flow here we can see that i log several things such as the accuracy the loss the kl divergence and i think this run was not very stable but for instance the number of reconstructed molecules is going further up so the maximum number in this training was 971 so out of these 10 000 roughly 10 percent were reconstructed and we can also have a look at the kl divergence it's also decreasing but i still have the feeling that some weird things are still happening here and also overall that the kl divergence is too high but this is something we will see in the next videos and uh on the test data so on new data that is unknown to the model we get also an increasing number of reconstructed molecules so here we have 4 000 molecules and 500 of them are successfully reconstructed and now let's go back to the training and as you can see it started to reconstruct some molecules so if we see if we go back here it starts with the smaller molecules we have only seven nodes here and six connections and then it starts to get bigger and bigger and as you can see here it already reconstructed molecules with 18 nodes and 18 edges and if we keep that training it will reconstruct more and more molecules and here i print the smart string so this is what happened in this utils function where i call this reconstruct molecule or whatever it's called here i print the small string and that's exactly what we see here yeah and as you can see it starts to reconstruct everything which is already a good sign already bigger molecules with 27 nodes i think the largest one i've seen so far it was i don't know 42 notes or so so this first training gives me the feeling that this is going in the right direction because the gnn learns to compress the information into node embeddings and then is able to pick out two of these node embeddings and tell me if there is a connection or not between those and it also works for larger graphs with up to 40 nodes with this basis i've implemented it's getting series now because my next personal goal for this video series is to generate entirely new molecules until the next video to make this happen there are a couple of points on the agenda first of all i need to change the decoder to output different edge types this should be fairly easy to implement then i need to make sure that the latent space is only one vector now for this i need to sample the atom types as well as the number of atoms for the generated molecule so after we've implemented that we should be able to sample new molecules and that will only work if the latent space is working properly so if the kl divergence is small enough when this works we can then start to have a look at conditional generation so that we don't produce arbitrary molecules but instead molecules with desired properties in our case this means potential hiv inhibitors and finally i didn't optimize any hyper parameters yet so that's also something we can try out all in all i plan to do two more videos on the generation part of this gnn series and after that i will move on with the next topic so that's all for today thanks for watching and have a great week [Music] you