Temporal Graph Networks (TGN) | GNN Paper Explained

Key Takeaways

i want to make a shout out to shantanu who recommended i create a video on tgn uh temporal uh graph networks and it was about time since dynamic networks is such an uh interesting like uh concept uh interesting research uh direction and many real world applications such as social media networks rely on uh having the modeling these uh dynamic graphs so what's a dynamic graph so basically whenever you have so so in contrast with steady graphs basically here you have uh the nodes can be added deleted modified same goes with edges so we basically have two types of events we have uh node wise events and we have interaction or adjuvants so let's say we have some simple network here like graph and uh it can be multi-graph in in the general case and basically uh i'll be using twitter as one example since the authors of this paper emmanuel rossi ben fabrizio david federico and michael bronstein are all from twitter let's say anna decides to join twitter so basically she signs up she fills up this signup form and now she's on new noting in twitter uh social graph and that's basically the node edition event now let's say after a couple of days she was just kind of retweeting people and every single retweet is basically an edge in that graph now uh ben joins twitter and because those two broke up anna's pissed off and she just deletes the twitter account altogether and that's the note deletion event now ben is also pissed off and so he decides to update his uh user profile updates his description and by updating the description he's basically changing his feature vector and you can think of it maybe we take this text and we pass it in a birth uh encoder and you basically have cls here token uh hopefully you're familiar with bird if you're not just treat it as a black box we input the text here we get out the uh condensed representation here and we use exactly this thing as the feature vector so that's the node modification event now after a couple of days of activity he decides to retweet andrew ang uh and school post and basically now whatever the retweet text is that's the edge feature vector and now maybe he made a typo so he decides to to modify uh the text even though that's not possible on twitter but let's say it is and for the sake of argument and basically that's the edge modification event where now the new text is again fed through bert and we get the new feature vector and finally he decides to just delete the tweet altogether and that's the node that's the edge deletion event so that's how the dynamic graph works in a nutshell it's basically a multi-graph because every single node can have multiple edges between so the nodes nodes are connected by multiple edges and that's that's what's known as a multi graph okay that out of the way let's jump into the paper and see what they say so they say here a few approaches have been proposed for dealing with graphs that are dynamic in nature so basically most of the research in graph ml was focusing on static graphs uh and not on dynamic graphs and uh while it is possible to apply steady graph deploying models to dynamic graphs by ignoring the temporal evolution this has shown been shown to be sub-optimal so what that means is the following so you have this graph here and basically every single edge has a timestamp attached to it so let's say uh ben here has got maybe three retweets of androin and this one happened at t1 this one happened at t2 this one happened just recently like a couple of seconds ago at t3 and you can treat this as a static graph just a static multi-graph uh or you can um uh kind of acknowledge the the like the existence of these timestamps and use that to when you sub sample your neighborhood later and we'll see what we do that a bit later but when we subsample the neighborhood we can just take the most recent edges and uh by doing that we we have a much more optimal performance than just by treating this as a steady graph and doing unifor a uniform sampling of the neighborhood so yeah basically that's suboptimal you can treat it as a as a steady graph but it's sub-optimal okay and yeah so uh learning on dynamic graphs is relatively recent and uh most of the previous work was uh focused on these discrete time dynamic graphs uh which is just a sequence of snapshots of the graph so basically uh you you you have your graph that's evolving and what you do in the in this discrete uh time dynamic graph is the following you just take a snapshot at equidistant moments so maybe you you had some graph at t1 that looked like this and then uh after some period and the important part here is that this is pretty much the same constant so you have equidistant moments when you take the snapshot and maybe a couple of more edges and notes appeared etc so like this i don't know whatever and you do the same thing after the same time period t this is basically the improvement over the steady graph modeling but this discrete time the dynamic graphs are still limited compared to the continuous ones that we're going to see in this paper okay a small recap basically when you want to compute the embeddings for the static graph you have you just have basically accumulation over the neighborhood of these messages and h is just a trainable uh update function and that's how you get your embeddings that's what this equation here stands for but it's not so important now so i just want to kind of walk you through some terminology that we'll be using and that's as i already mentioned we have this discrete time dynamic graphs uh or dt dtdgs for sure and we have continuous time dynamic graphs on the other hand and those are the ones we'll be treating in this video and basically uh the the the point here is we we treat it as a as a sequence of time stamped events so every single event be no wise even be the adjuvant has a timestamp associated with it uh conscious that with the steady graphs where we had those equidistant snapshots and i already explained those two and now the the the interesting concept new concept that's not encountered in steady graphs is this concept of a temporal neighborhood so what it means basically is the following so again you have some nodes in the graph and let's say there are multiple edges here and let's say this one happened at t5 and we're currently looking for some reason we're looking at at a moment that happened prior to t5 like at t4 that basically means that the neighborhood of this graph here is will ignore this edge so everything that comes after this particular target moment of time is ignored and that's what you're left with so all of the other edges here is what's known as the temporal neighborhood of this particular node at this point of time okay so that's a new concept and with that out of the way let's let's jump and see uh what what are the main modules in in this tgn network okay let me give you a a quick uh high level overview of the system and then we'll slowly start building on and adding new details so high level how we train this uh graph neural network this whole system actually because it also has this thing called memory uh is the following so we do a self-supervised uh learning and we're learning how to predict link to do link prediction so basically because we have have all those all of those retweets and we have time stamps associated with them we we can just do chronological sort over all of those interactions and we can just uh predict the the interactions that have that haven't yet happened and that's that's that's basically the you just patch that chronologically sorted uh array of of those uh edge interactions and you just predict the follow-up uh edges it's a similar thing as transformers there you're trying to predict the follow-up word and here you're trying to predict what the next interaction will be so that's the high level uh overview of how the thing works so now let's see on this chart what happens basically so you have you have and i'll i'll just create a simple example here we have node one we have nodes two and we have node three so one as you can see here at t1 one retweeted node two and at t2 node two retreated node three so like this so now what we do is the following we have so because node one also had some other interactions maybe that happened even prior to these interactions so let's say these happen that whatever t zero maybe okay so now what happens is that we have this embedding method and that's basically get you if you if you haven't watched my video uh on graph attention networks you can check it out i'll link it somewhere here and you basically do uh one layer get over this node 1 over its temporal neighborhood and we already mentioned what that is we do get and we calculate the embeddings for nodes 1 2 and node 3 right and those are these z's once we have those we can just calculate the probability using a simple decoder so they used mlp here so you just concatenate these two embeddings you just pass them through mmp mlp and you just add the nonlinear function like sigmoid non-linearity and that means this is basically probability and now in order to train you do the following so you want to make sure that p so given t1 what's the probability of interaction between one and two happening and we know it's 100 it's one because it already happened we have it in the data set right uh we did the chronological sword so you want to push you want to push the loss uh to go to one to p to go to one uh i loss will be be zero because we're just using simple binary cross entropy here and we'll also have negative edges which we'll want to make sure that the probability goes to zero and uh yeah that's it so if you have maybe some notes x and nodes y and this interaction hasn't happened we want to make sure that uh we we push the the probability to zero for those for those uh note pairs okay and we use a simple as i mentioned two losses just a binary uh cross entropy uh that means that's a fancy name for you you just put a log here so basically if we have positive examples we'll just the loss will have a format like this minus log of p right and um if i just draw it here uh they will look like this and basically we want to push the probability to go to one because then the loss goes to zero for when we have negative examples we'll just have minus log of one minus p and that will just kind of um mirror this chart here and it will look like this so we have p on the x-axis we have loss on the y-axis and the loss will look like this basically at zero we have loss uh that's uh that's equal to zero and at one it goes to infinity so uh that's how we train it simple bc we have positives we have negatives uh we have the data set with uh and we just batch it and we use those batches to to to to to predict uh those those those links because we know they're there okay so how the embedder part works so get leverages something called uh memory or states in this tgn network so what it is is basically every single node or in our graph in twitter graph has has a state associated with it so we have a huge table and every single node in this table has some representation and they said here basically the purpose is to represent the node's history in a compressed format and that's exactly what it does so basically over the uh when the interactions are happening in twitter we're basically using all of those to just update this memory state and then we can use those to later do some recommendations but i'll get to that in a moment okay so we have these states and they're somehow calculated so we have some messages we do some aggregation we somehow use them to update the the memory state and then the embedding method the get we'll just use those to calculate embeddings and we can predict the links so that's that's the high level overview and now we'll slowly start digging into details hopefully that was clear enough um maybe one more thing to worth mentioning um when we are using get in this in this model uh the way we we subsample the neighborhood is using the most recent edges so that's why we have timestamps so basically uh if we have a bunch of edges we'll use will only subsample the most recent k1 so we we sort the neighborhood edges by timestamps and we just take the k most recent ones that may be like 10 most recent neighbors and then you just do the simple get aggregation and we'll see what exact features go inside here in a sec okay now the problem with this uh chart is if you maybe if you noticed is that we are using these interactions to calculate the messages to update the states and then we're using that in those states to predict the links so that means we have information leakage here we're basically using this information to update states and then we're using the states to predict those links and that doesn't make any sense right but because you somehow have to uh you somehow have to calculate and use these message functions to get the states because otherwise you won't have any gradients flowing towards these trainable functions the aggregated messages and messages and this mem function they are all trainable so you somehow have to have uh this forward pass here but you can use dispatch you'll just use the previous batches and that's the solution so basically they have a graph here okay um you have to introduce this concept of a raw memory storage and now what you do here is from the previous batch you just stored all the necessary information inside here so now when you try to predict the the the existence of a link between one and two the the existence of of retweet in the case of twitter uh what you do is the following you use the the last patch and i mean not just the last patch actually the whole history and you calculate the messages and we'll see how we do that so we have some data here that's necessary to calculate these messages then we just aggregate the messages so you see here that messages for the node 2 are kind of accumulated and then we update the states and they just use group here and we'll see those details so what happens is basically you you have the previous state so this is group g r u you just use s2 and this is the old state and you use the newest message this m2 aggregated and this will just spit out new state and we just write it down here so that's how how we update the states now the states have uh now we have this thing like included in the computation graph and we just uh use get to calculate embeddings and to we get to the loss so now gradient gradients can now flow through the whole system and all of these functions which we'll see what they exactly are so these egg these mem uh functions will be uh updated in the backpropagation step okay so that's the important solution so introduction of raw message store is important because that allows us to train this part of the pipeline this part here okay okay hopefully that was clear enough um let me let me slowly start digging into into details of how the messages are are computed and then once we have all of those details i'll again zoom out and try and explain everything okay so messages are computed the following once the interaction happens between nodes i and j so we have node i we have j one of them retreated other so we have interaction that happened and it has feature vector associated with it so that's e i j right so how we calculate the message is basically uh we input the previous states and this t minus just means the state before this interaction happened for those nodes uh so that's basically whatever is currently in the state table so we have si here we have sj so whatever is currently in there we use those and we input them in the message function and we additionally use the uh this feature vector i just mentioned so that's like the maybe birth encoding of retweet and we finally use this delta t which is the time that has elapsed between the previous interaction and this current interaction so that may be like 30 seconds or whatever and that's how we compute the message for this for this node i and then we do analogously to node j so that's the destination so that's the source and destination nodes and the only difference is these are permuted if you can notice here so just permutation simple permutation and although this can be trainable what it did is they in practice use just identity function which means they are just doing simple concatenations uh between all of these features so you have some vector here vector you concatenate all of them and that's your message and this here is an interesting thing and i'll later explain uh how we convert time to a vector uh to a vector representation and just stay tuned for that okay so that's the uh so that's the part with with the messages okay and they said here a more complex message function that involves additional aggregation from the neighborhood from the neighbors of nodes i and j is also possible and is left for future study so the thing with this paper is there is there is a it's a work in progress there's a lot of things they still want to try and they haven't tried it so for example they have this um node-wise uh memory uh but they also so they say here while a global graph-wise memory can also be added to the model to track the evolution of the entire network we'll leave this as a future work so that's that's my point uh there is uh still a lot of things uh that evolving in this work so yeah just just keep that in mind um okay so once we have the message function so that's this part right so that's the messages how do we aggregate them and here we just have a message aggregator and again it can be a trainable function but what they did is they just do a simple most recent message and mean message heuristic and you can even see it here so what it did is you have m2 at t1 you have m2 at t2 they just stick the last message and that's the aggregation okay um emoji general case they can also just do a mean or something else just kind of do make it trainable maybe pass it through an rnn or whatever okay but it's stuck with the most recent message and let's use it as the running example uh finally we have the memory updater i already explained that they used some form you can use whatever like rnnd or lstm they used gru which is a specific type of of lstm basically and i already explained how it works but let me just repeat okay let's this is group and we have we have the state state i t minus we just get the new message and we just run this and there are like a if you know how if you don't know how group works i'll just link uh chris ola's blog he has a really nice blog intuitive to understand how this works basically a bunch of forget and update gates and this spits out we can treat it as a black box for now it just spits out the new state s i at point t so this was a t minus and here we have the current state so that's the memory updating part okay finally we have the embedding module and so i haven't mentioned this but the main goal of the embedding module is to avoid the so-called memory stillness problem so what that means is the following so let's say we we have a graph and um like we have we had been here and maybe ben stopped using twitter for maybe a month or two and now what happens is that his if you take a look at the memory so and we we take ben's state uh because he's not interacting with anybody anymore um we this this representation is not updated and so um we we we need to have uh some sort of of of embedding to to to make this relevant for ben when he returns so if you want to recommend to ben which tweets should he uh see so that's like your your homepage on twitter you see a bunch of tweets so twitter somehow needs to figure out which tweets are relevant for you and also uh twitter should recommend you uh what's the what's what what other persons you should probably follow so how they do that is the following so ben has uh his state b and in the simplest case you could just do the following you could just take sp and maybe andrew ang has his own state as a and you just do a dot product or mlp and uh whatever and that spits out some probability so if the probability is close to one uh twitter may recommend that you follow uh uh and you anc okay but now the problem is if this remains stale these predictions will not be as relevant anymore for ben because he he has changed in the meanwhile okay so the solution is to use what i already explained and that's using some kind of like a gnm module like get and uh you you get the updated version uh of the of the embedding because his friends presumably uh presumably his friends haven't stopped using twitter so their states are are being constantly being updated so we can use the the blood so so we can use that information to make a more relevant predictions for ben okay hopefully that makes sense and they had a couple of baselines so the one i mentioned the simplest one is just you take the you just take the state uh of of the user so that's this uh identity uh embedding so the final embedding is just whatever the state is and they show that this has a really bad performance on the later benchmarks okay and uh then we have this time projections simple equation you can check it out yourself uh i'll focus on on on the temporal graph attention so this is the important part this is the strongest baseline basically they had uh and the way it works it's it's the same thing as get just instead of using the the attention that get used originally they just use the original vasani attention so that means they have queries keys and values and everything else remains the same basically the new information here is this five function and the fact they're using multiple features concatenated here to get the final uh embedding representation okay so let's see how it exactly works so this h so for a particular node at a particular time t how they calculate the h is like this so they they take the state of that node at a particular time t and they take the notes features so that's basically maybe a description on your profile and then birth embedding blah blah blah blah and you basically just add them up and that's your zeroth series that's the ro those are the raw features okay so we concatenate those with the uh basically edge information so let me let me draw an example here it'll be easier we have a node we have a bunch of nodes here again we'll do the most recent uh neighbors we'll just kind of disregard all of the others so let's say we only take these two into account okay so we'll these are these will have some um these will have some features so if this is not one this is maybe node i so this is uh node i1 and this is node i j sorry j1 and uh so you you just combine those information and you combine with phi so what's what's phi exactly here and uh it's a there is a this nice paper called time to back time to wek and you can check it out uh but basically so that's one thing you could do and basically what you do is in order to get the vectorized representation of time you just map it into the following thing so you map it into a vector that has features like this so maybe w one t plus ah uh n1 and then the second feature will be sine of w2 t plus n2 etc so you just continue using science here and you increment the arguments by one so we have w3 etc so these are learnable so you learn w2 2 w1 and 1 etc and if if you take a closer look and you're familiar with the transformer this kind of resembles the positional encodings except that these are time aware okay and basically so you learn these and uh they showed in the time to back paper that you basically these are capturing the periodicity in the signal whereas this one is uh capturing something more uh constant in the signal so this one if you take a look this is just a simple uh basically a line and these here are sinusoids so you basically are learning sinusoids like this and the the higher the w the the higher the frequency will be and if you have some smaller w then you'll have slower smaller frequency etc and these are just uh the offset uh these are just uh encoding the offset of your sinusoids so that's the simple uh heuristic they use to encode a time and once you have that you just do simple multi-head attention and you you get the aggregated feature vector you concatenate it with the current feature vector and you pass it through an mlp so this thing here h i corresponds to this node and the second term here is the aggregated combination of these two because they are the most recent neighbors right so we kind of we associate some alphas alpha one alpha two here and uh we combine them like a simple weighted uh sum and we get the features so that's how we calculate embedding and they show that graph attention network is the best baseline we'll we'll see that in a minute so that was that was all the nitty-gritty details um we saw how to train these um uh how to how to pass how to pass some graded information to these modules we saw how they exactly work and now we'll see a couple more details and then i'll zoom out i already mentioned this part about the information leakage and so i'll skip it and so that's the reason we had to introduce the raw memory storage and this part is interesting so while from the perspective of the first interaction the batch the memory is up to date since it contains information about all previous interactions in the graph from the perspective of the last interaction in the batch the same memory is out of date since it lacks information about previous interactions in the same batch this incentivizes the use of a big batch size so let me break it down for you um that means the following so we have the previous batch information we update the states and now this first because this is a batch that's chronologically sorted we have t5 d6 so this first interaction will be using states which are totally up to date but once this one is done we should basically have s1 and s2 updated because once an interaction happens we calculate the messages for the nodes that that are associated with that interaction and we update the states but we can do that while we are in the batch so that means this to this one here two to three will use the same state as this one here and that means it's slightly out of date and now imagine the bigger batch and the used 200 as a nice trade-off and the the the deeper you are in the batch let's call it that way uh the the more the more stale this current uh state representation is so that's the reason you don't want to have a too big of a batch but then again you don't want to keep it too slow too small because otherwise it will be really compute inefficient to have a small batch size okay so that's that part and now let me just show you the results before before i do a high level overview so um they had some uh continuous uh this dynamic graph uh baselines here like jody tgad which is basically the same thing as t like this this paper tgn it just doesn't have memory and dissociated modules so it's a simplification of this of this paper basically it's a specific case specific instance and they also experimented with different uh uh basically different embedding methods using memory or not using memory module etc so they have a bunch of baselines and the results are the following so the tgn with attention achieves the best results overall on all the three data sets so i was continuously using twitter as a running example but i also have wikipedia and reddit where this one is just users or notes and pages or notes and user editing a page is an interaction event and basically the edit text is what the feature for the edges etc okay similarly for edit so those are bipartite graphs so we have pages here we have users here and let me see what sounds interesting they also did uh node wise classification but that's not so so important for this paper and uh here are some ablation studies basically you can see that let me just zoom in a little bit uh basically the the tgn with the tension is the best tradeoff overall because it takes uh less time than these two um whereas the performance is really similar so this one just uses two get layers and the reason one get layer is enough is because they are using this memory thing and the memory thing already implicitly contains the information from the neighboring nodes so using one layer you're basically uh you're basically accessing the features from the two hop neighborhood okay and uh using mean i think the mean just means um uh they're using for the for the memory aggregation they're using the the mean instead of the last uh last heuristic and you can see tget and jody and dyrap these other bass lines are way lower uh performance wise than than this method uh some ablations again um with their own model uh when they don't use the memory you can see the the worst curve is here uh when they uh add two layers again no memory it's a bit better and then as we add the memory and we add the the the last uh message um aggregation heuristic we get the best results overall so that's pretty much it let me see if there is something else left that's it um now let me just now with all the information you have in your head let me just briefly go ahead and explain it again um okay now that we have all of the information um explained let me just kind of give you a holistic overview of the whole pipeline so here we have information from the previous batches of these uh edge interactions okay and that means we have um so rm1 is basically s i contains si j contains the delta t so that the time that elapsed between the previous interaction and this interaction and we have e i j which is the no the interaction features and we have the same tuples for all of the others the thing is we only focus on the nodes that we need so this thing may be huge and contain information for for many other nodes in our graph we only care about nodes one two and three and that's why we only calculate those messages so we have we we just basically do a concatenation that's what they did and you get m1 you do a concatenation you you get m2 etc at different timestamps okay now the aggregation they used either the mean or the last message so basically you just take let's say we we use the last message we just pass this message into the next step and that's the aggregated messages array once we have those we just have group here and for every one of those we pass the state we have feed in the message we spit out the new state when we update the states so now we have the states and the the compute this part is included into the computational graph now we use those states uh and we use get to create the embeddings now interesting here thing here to note is so we have notes one two and three but because uh this node one maybe has some other uh neighbors uh it will make sense it would make sense to have other states here as well so whatever the neighbors are i'll call them sn so those are these guys here um it would be useful to have uh the states for them updated as well because those will be included by a get because get accumulates the neighborhood information right so those will be accumulated by the get so i guess they just missed you to place sn here like the neighbors but whatever and finally you calculate the embeddings and now you just uh depending on the ground truth so these exist but the negative negative ones don't exist so you just use binary cross entry plus to train this whole pipeline end to end and now you have a system that can successfully predict uh given two nodes given to yeah given two nodes you can calculate their embeddings and you can predict whether they are likely to interact in the future and that's useful for twitter because you can use that to recommend for node a which node should that node follow in the future like maybe andrew ang or which which uh tweets should that user see etc a couple of strange things i've noticed in the paper maybe maybe a typo or something uh basically uh if you take a look at the query uh it doesn't have the edge information here so that means if you're doing a scale.product between the the the query and the keys you have a like a differently dimension vectors so you can do it so it's either a typo or they have some additional projection layer here that they didn't explicitly include here or i've missed something the second thing is i mentioned the memory stillness problem but basically if it's like they are implicitly stating that you need to retweet somebody else in order for your memory state to get updated but what if somebody retweets my post so i assume because that interaction involves my node and that person's node so because the way how the messages are computed my nodes state should be updated as well so i'm not sure um whether they're just simplifying it here in the paper but anyways uh just just keep that in mind that that's not super clear because for me it looks like that my stage should be updated even though i'm i'm not active on twitter because other people are interacting with me and thus my memory state is being updated constantly uh they're probably using a direct edge assumption here but it's not super clear from the paper uh one more thing to keep in mind is that this memory uh table is not uh let me open a pen it's not trainable so what is trainable is the group here the lstm and potentially the message and aggregation functions even though they just use the concatenation so there's no learner learnable parameters here neither there is uh learnable parameters in their aggregation function so they're left up with learning mem so that means once the the model is fully trained uh you're basically once an interaction happens you just calculate the messages and you basically uh aggregate them and you update the states and that's it and once you once you want to predict do some recommendation then you use the the get module and create those embeddings and just uh make the recommendation so that was all i had to say for this video uh you know the drill subscribe hit the bell icon and until next time keep learning deep