Transformers On Large-Scale Graphs with Bayan Bruss - #641

0:07All right, everyone. Welcome to another episode of the TwiML AI podcast. I am your host, Sam Charrington. And today I'm joined by Bayan Bruce. Bayan is Vice President of Applied Machine Learning Research at Capital One. Before we get into today's conversation, be sure to take a moment to head over to Apple Podcasts or your listening platform of choice. And if you enjoy the show, please leave us a five-star rating and review. Bayan, welcome back to the podcast. Thanks, Sam. It's great to be here again. It's great to have you back on the show and congratulations are in order. The last time you were here, you were Senior Director of Applied ML Research.

0:42That's right. It's very exciting and it's a sign of the continued growth and investment that Capital One has made in this area. So it's exciting for me personally and then also for the team. Awesome. Since you are a repeat guest, we'll keep introductions brief, but I would love for you to share briefly a little bit about your background and then we will jump right into our conversation, which is going to focus on some of the research that you and your team are presenting at ICML Conference. Great. Yeah, brief introduction for those of us who didn't hear the last podcast. So I lead our Applied Machine Learning and, yes, AI as the new addition to that title team at Capital One.

1:24And our main focus is how do we take all of the amazing academic and industry research and shorten the time until it can be used in production systems across the company, across our applications. I think, as I mentioned last time, Capital One has a very long history of using machine learning in many of our most important applications. And as part of that history and the foundation that we built, we know that we have to stay on top of the advancements in the technology to be competitive and to provide the best possible experiences for our customers. And so we use our applied research function as a way to shorten the time from when there's been some great breakthrough outside the company to when we can use it for the various things that we do.

2:10And your team has a couple of papers at the conference we'll be talking about. The first that we'll dig into is called Identifying Interpretable Subspaces in Image Representations. Give us a high-level overview of that paper and what it's hoping to accomplish. Sure. I guess let me back up and provide a little bit of additional context. One of the things that we do as a team in order to, as I say, accelerate that adoption of research and technology is we partner very closely with a wide range of academic partners across a number of universities, University of Maryland, Carnegie Mellon, MIT, Columbia, NYU, and others.

2:52And we work side by side with researchers, grad students, and we work on publications with them so that we can understand where the state of the art lies and co-generate that knowledge and then take that back and quickly apply it to the things that we're seeing inside the company. So in this particular case, this was a paper that we worked on with some colleagues at the University of Maryland, Neha Khalibat, Shweta Bharadwaj, Hamid Firouz, Maziar Sanjabi, and Sohail Faizi. And as I mentioned, stepping back again at Capital One, and I actually think this is more broadly across industry, while the big focus in the headlines and other places is generative AI, the workhorses of industrial machine learning applications are still discriminative.

3:42There are models that are taking in inputs, making classification decisions, making regression decisions. They are taking those decisions and then applying business logic to them and then determining what happens in a given application or a given system that we use. And given that we have so many of these models that are discriminative in nature and they're running in all of these different applications that we have, one of the best and easiest ways for us to take deep learning and AI technology and quickly integrate it into the various systems that we have is through embeddings, is through representation learning.

4:22And you actually see this in a number of industries isn't just Capital One. Representation learning and high quality embeddings has become the workhorse of many, many applications of machine learning today. And the way that it typically works is you have some complex high dimensional data set, you have some pre-trained embedding generation model that produces embeddings to some centralized repository, and then downstream applications can come to those embedding repository and do a wide range of tasks on top of that, either they can take those embeddings and they can do similarity search, nearest neighbor functions on top of the embeddings.

5:01This is why we're seeing such a huge proliferation of things like vector databases. Or another common path, especially when you have a lot of machine learning models, is that the downstream models can use those embeddings as inputs in addition to other features that they might be using to make predictions. and they can thereby take advantage of all the representation learning power of a deep learning model without having to be fully integrated with that deep learning system. And so we actually see oftentimes like a mix and match where you'll have a traditional tree-based model using embeddings as an input.

5:33And it's kind of a way to accelerate deep learning into our various systems. But it leaves this gap, which is that traditional features and traditional models were hand-engineered by the data scientists who built the models, they had an intuitive sense of what those features were. When they put those features in the models and deploy the models, when they do feature importance or any kind of explanation technique, post hoc explanation technique, like if the SHAP values, what they are learning is this feature was important to this prediction. And because I know that this feature was and where it came from, I understand kind of implicitly what was going on that led that prediction to happen or why the model made that prediction.

6:17When you take those traditionally hand-engineered features where there was some clear sense of what they mean, then you append a 64-dimensional or 128-dimensional feature vector, 256 dimensions, embedding, and pass that into the model. And then you ask the same question of it. They might say, oh, well, these two features that you understand were important. And then embedding dimension six and embedding dimension 53 were important. And as a data scientist or as somebody who's in the business trying to understand model behavior, that becomes very challenging to interpret. And so having an agnostic way of understanding representations and what goes into those representations is a big area of interest to us.

7:03And so that is the setup for identifying interpretable subspaces in image representations. Yeah, that's a great setup and context for the paper. One of the immediate things that jumped out at me just kind of looking through it is it throws around the term feature, but everything's a feature in these systems to some degree or another. In this context, you're talking about dimension 53 in some long embedding as opposed to full embedding or something else. You've got these big embeddings. Are they individual dimensions or are they pairs or tuples of dimensions that kind of carry the meaning here?

7:44That's a great question. So we started looking at individual dimensions and asking exactly as you say, like for pick a single element in the given embedding, what can we learn about what that embedding captures or what that dimension captures amongst all the dimensions? And we can talk about how we go about doing that, one of the things that we observed is that the kind of interpretability, you can almost call it like the purity of a given dimension is actually fairly low and only a small percentage of all the dimensions on average provide unto themselves an interpretable kind of set of things that you might look at.

8:20And actually what ends up happening is if you combine multiple dimensions and use that combination, you get these more interpretable substasis in a much larger in a much larger way. And I think it makes sense because of the way that neural architectures work. Oftentimes, that dimension will encode multiple different pieces of information. And it is only in combination with other dimensions that encode other pieces of information that the combination of two actually start to become interpretable. The model that you use for the embeddings here is SimCLR, it's contrastive, it's a contrastively learned embedding model.

9:02You also use the concept of contrastive concept extraction. Are those two, is the use of contrastive techniques kind of coupled between those two or are they kind of independent decisions? They're independent decisions. Contrastive learning obviously being a fairly powerful paradigm, both in how you generate representations. But also, when we think about interpretability and explanations, contrasts are a common way that humans understand the world. It's not just this and this, it's this and this and not that is an easy way to go about differentiating between things. I guess I'll walk through exactly what is going on so we can describe what these contrastive images are.

9:49I recognize that I've jumped into the weeds here, so definitely do that. That's fine. Yeah. So at a high level, the way this technique works is that we're focusing for now on vision models, but this will potentially be informative for a wide range of other modalities besides vision, for instance, graph structured data and others. But you start with a pre-trained representation learning model for vision for images and you take the images that are in that data set and for each dimension in the representations that are generated for those images you identify the most highly activating images for that dimension so you essentially use a traditional in this case we use grad cam but you can use other methodologies for local explanations to understand like images where that dimension has a much higher capture of the activations of that image.

10:50And then even within those images, we want to then get down to a feature crop. So we don't want to look at the whole image because different things in the image might activate different dimensions in the embedding space. And so we take the images that are most activated for that dimension. And then we look within that image and find actually which pixels are most activated. And then we crop around those. And so now you have for that dimension, a set of cropped images. So in the example we have, it's a bunch of pictures of like a shaggy dog, but it's not the whole picture. It's like, you know, little subspaces like the nose and the eyes instead of the whole picture.

11:30And then you take another, you take another model, And this is where the benefit of pre-training becomes pretty powerful. So we have these language image models, these captioning models like clip, and you can take the feature crop and you can pass it into a language image model and it can generate some text, some descriptive language about what it's seeing in that feature crop. So that essentially becomes the natural language descriptor of the subset of the image that corresponded to the representation dimension that you started with and so that's how you map from representation dimension down to a natural language interface you can understand what's going on in that dimension now one of the things we found is that just doing this effort to go from highly activating to natural language alone did not provide super interpretable kind of explanations of that dimension.

12:29And that's where contrastive explanations come in. So instead of just relying on the things for that dimension that are highly activating, we also use images that are lowly activated for that dimension. So these are the things that that dimension don't encode or that aren't important to that dimension. And we go through that same process. And essentially, what you can then do is subtract from the descriptions of the highly activated images, the descriptions of the lowly activated images, and get to a more refined understanding of what is going on in that dimension, essentially saying, let's remove anything from our description of this image, or of this dimension, that would generally just appear in the data set, or that would randomly appear in the dataset.

13:19And so that's where contrastive explanations come from. And it really boosts the quality and the interpretability of the descriptions that we're looking at. How is that subtraction happening? Is that, are you giving both to a LLM and asking it to differentiate the concepts or is it something simpler, removing keywords or something? It's something simpler. It's removing keywords. So the, you know, when you think about hundreds or thousands of image cropped images that you've attributed to a given dimension a lot of inferences on an llm it's a lot of inferences on an llm and and it's also like it essentially becomes a book right ever so you know you add a thousand sentences and just to understand that one dimension you have to read through all of them and so we took a bit of a bag of words approach so you know removing stop words removing essentially getting down to the highest frequency words that the LLM was generating.

14:12And so then the same thing as you remove from that list of highest frequency words that the LLM is generating for the activated images, you subtract from that the highest frequency words for the lowly activating ones. And so you essentially get this removal of descriptors or descriptive tags that are more common in the overall data set or more common in the things that shouldn't be in that representation space. This process of identifying the highly activating images, is that starting from an individual dimension and working backwards? Or is it a batch task that happens on the training data set that identifies what dimensions an individual image activates?

14:55Right. You can still use the training data set or you can use any other data set after the model's been trained. It doesn't have to be the same data set that was used in the pre-training, but the idea is you compute the gradients with respect to that one dimension. And as you compute the gradients with respect to that one dimension, you'll see the highest gradients correspond to a subset of images. And those are the images that you select. You're essentially working your way back through the model that created the representations to get there. Talk a little bit about the prior work in this space and, you know, against which you benchmarked your results.

15:33Yeah, I think a couple of things that were relatively new, actually I should say are very new, the contrastive concept extraction is novel in this space. and some of the other methodologies often will do a much simpler approach where they just take images that correspond to given dimensions and try to assign some kind of language tag to them. The process of doing the contrastive subtraction is, I think, the thing that was most beneficial here. And I think more generally, there's a broad set of research on explainability that we've been looking at that goes beyond just individual dimensions. And I think this speaks to what this research really is beneficial for us when I step back and think about it.

16:22You know, a lot of explainability research, there's two tiers of explainability types. There's either, how do I explain a given prediction that a model's made? How do I explain why we classified this transaction is fraudulent or not fraudulent. And that'll be essentially a vector over the inputs to that model and weights that correspond to which of the inputs matter the most. And then there's global feature importance, which is this other paradigm, which is on average, what are the things that this model uses to make decisions? And that leaves this big gap in how models behave, where we look at different slices, either slices horizontally at what is the model doing for certain subpopulations, or in this case slices vertically or different dimensions in the input and i think we've been using representation learning in a few cases but this idea that you can then map backwards from a given dimension actually plays very well with how we think about their usage in these downstream models because we're often asking these local questions around for a given prediction what mattered the most and so therefore having the ability to say not just these dimensions are important, but then to go backwards and say, okay, from the inputs, these are the subset of inputs that correspond to that dimension.

17:41And so then we can inspect those inputs for a given model. One thing I will add is the usage of these different subspaces for error correction in model diagnosis. So by having a more intermediate picture of what's going on in model behavior, in this case, looking at individualized dimensions or subspaces across different dimensions, you can then draw correspondence between model errors in the downstream model and those given subspaces that are important or you know if there's a i guess let me explain it this way so say you have a group of images that your model classifies incorrectly you can then see okay which subspaces do those images belong to almost treating it as the inverse problem so take those feature crops and say, okay, which of these existing groups that I've already identified, does this belong to?

18:35Which dimensions does this correspond to? And start to identify model behavior that you could then correct. You could say, well, it's actually attending to this cropped feature within this one. And it's actually learning from the wrong information about this image. And there's this whole notion of spurious features, spurious correlations in a lot of deep learning where models will, especially computer vision models, will learn from anything you give them, whether you intend for them to learn from it or not. And if in the training set, the things that are most predictive of the thing you're trying to predict are not actually indicative of that outcome, the model will still learn it.

19:17So the classic example is the husky on the white background is classified as a wolf because the model is picking up on snow, not on anything thing about the dog. In our case, we show some examples where the models classified a flock of geese as an airplane wing because of the blue sky. Using this methodology, you can not only see, okay, that the model got the prediction wrong, but you can actually see what the model thought was in the picture, right? So, you can see the model thought that this was an airplane wing or something like that. And then you can visually inspect it and say, okay, that's actually a goose.

19:55And so then you can go back and tweak your training data set, or you can tweak your learning procedure to make sure that your model's not picking up on those kinds of spurious correlations. So, you know, from a model diagnostics perspective... It becomes a bit of a debug tool. Exactly. Which is very hard in some of these deep learning methodologies. Awesome. And then was there a particular benchmark that you compared performance against? There are a few benchmarks. Pull up the exact reference to the benchmark. One of the comparisons is against a Milan. Yes. Is that another approach? Hernandez and all?

20:33That is another one. And then one of the other interesting things that our partners at Maryland did in this particular case is we also did a human study, which I think is actually quite important in explainability research. And so what they did is they actually had individuals look and rate the different explanations and the different descriptions of the dimensions and say how accurate they were just by what they're seeing in the images, as well as kind of how useful they are in understanding what's going on with the model. Awesome. So the next paper that we wanted to talk about is GOAT, a global transformer on large scale graphs.

21:22It is a planting a stake in the ground, making a statement with the name of that paper. Yeah, it's a great name. I guess the alternative would have been GLOAT, which may be even more of a statement. The many years that I've kind of followed some of the work at Capital One, graphs and graph neural networks has always been kind of a common theme and something that the teams there are very interested in, presumably related to their utility and fraud use cases and others. And just based on the title here, you're applying the hot new thing in Transformers to the thing that you've been concerned about for a very long time.

22:05So with that, give us the high level overview of this one. Yeah, that's exactly right. financial services, by and large, operates on top of very large natural graphs or networks. In our case, you can imagine every time that you swipe a credit card, you are establishing an edge between you and the merchant. That scales up to functionally the many tens of millions of a node network with billions of edges. And it contains a lot of very useful structural information that we can use in a wide range of applications. So we've been looking at, as you say, we've been looking at how do we apply graph neural networks and graph representation learning for a number of years.

22:47This is not a new area for us, but with the dominance of transformers in recent years, it makes sense to look at them in the space of graphs. And one of the big challenges is that transformer architectures particularly self-attention is on long sequences in the language space is something that has a polynomial time complexity it's n squared so you're essentially taking this n by n matrix of all the elements in a given sequence and you have to compute that that's a very computationally expensive task and then on a graph the corresponding notion of a sequence on a graph is a neighborhood right so we think about a neighborhood what we have is all the if you're a single node in the graph, your neighborhood is all the nodes that you're connected to.

23:35That is the size of the N when you are computing attention. If you go two hops, which means that you go to your neighbor's neighbors on the graph, then the N is now not just your neighbors, but all of their neighbors. So now that's the square of all of your neighbor's neighbors and your neighbors. And as you go outwards to three hops, four hops, however, however large your graph is so now you not only have a square in the size of the number of elements that you have to compute attention over that is itself growing exponentially because graphs obey these power law in the relationship between neighborhoods over the scale of the graph and so it's it becomes computationally intractable to actually do attention on anything larger than very small networks.

24:24Now, it has been shown that attention is quite powerful in the graph domain. And so there have been a number of papers over the years that have looked at how do we apply this notion of a graph attention. And what has primarily happened is that the focus of those efforts has been on very small graphs. I think a lot of things like molecules or other, you know, very small synthetic graphs where you're doing maybe not large scale node classification, you're doing like a graph classification. You're asking, what is this graph that I'm looking at? And in those domains on those small graphs is actually quite powerful, but to scale it up to the size of a graph that we have is computational intractable.

25:06And so this research set out to propose some methodologies to get around that computational barrier so that we can do attention on a larger scale. The other benefit of doing attention on a larger scale, besides just being able to do attention on a larger scale is that most graph neural networks are built on the notion of homophily. Homophily being the principle that you are more like the things that you are connected to, or a node is more like its neighbors than it is nodes that it is further away from. And that is true of many networks, particularly in the social sciences. However, there are also a lot of networks where that is not true.

25:49And there is a notion of heterophily, meaning that you are dissimilar from your immediate neighborhood. And then I think the reality is that a lot of tasks are some mix of the two, where in certain situations, a node, you might want to have information about the node's neighborhood because it's highly predictive. And in other situations, you're actually looking for relationships between that node and nodes that are far away from it on the natural graph. And so having methodologies that don't require you to a priori pick, I'm going to use this because I think it's a homophily principle, or I'm going to use this because I think it's a heterophily principle, but can then just work as a whole becomes quite beneficial from a practitioner's perspective, because you don't always know, and you don't always want to tune that.

26:36It's not a matter of just like picking a parameter and a model, what you're picking is an entirely different model class to have to switch back and forth. And that can be quite challenging. And so this paper sets out to solve both of those things. Can we scale to very large graphs and still apply attention across the entire graph instead of just a local neighborhood? But then can we also retain certain localized information? And so the way that the research takes this problem and makes it tractable is through this notion of a codebook, which is essentially taking the entire graph and doing a k-means aggregation into smaller clusters of the graph and assigning every single node in the graph into one of these clusters.

27:24and then doing rather than having to do your global attention over every node in the entire graph when you're making a prediction now you only have to do your attention over the space of clusters that you have defined which you can fix in dimensionality it becomes much more tractable if you want to make it really really efficient you could say there's only three rows in this code book and i just want you to to compute the global attention on those three or if i want to expand and look at a much finer grained perspective you can go all the way out and so essentially that becomes a hyper parameter that you can tune to understand how much global information you want to collapse into this lower dimensional space or how much you want to leave in there and that allows you to attend to every node in the graph of course it's an approximation of every node in the graph but you can attend to every approximate node in the graph through this codebook and then there's an additional self-attention mechanism which allows you to attend similarly to at least a couple hop neighborhood around a given node and so then that gives the node classification task the ability both to go global in this approximate way as well as to go local in a more traditional way and see significant performances and now what's really interesting about the results is that they show that this model is capable of performing very, very well, not just on homophilic graphs and not just on heterophilic graphs, but on both.

29:02And so it balances performance. It can learn from the structural properties of the graph and from the task at hand, which pieces of information it needs to attend to in order to perform optimally at the node classification. And so a lot of that structure is not baked into the model. The model gets to learn it directly from the data, which is a really neat approach. When I think about that n-squared problem and the nature of these graph connectivity matrices, one of the things that jumps out at me that you might want to exploit is sparsity. What are the properties of these matrices that the codebook, that codebook approach is exploiting?

29:44It is. In a way, it's exploiting sparsity. In a way, it's actually making the space denser. You think about projecting this graph into a lower dimensional space and then running your attention over that projection is essentially what is happening. That's going, you know, your original graph is very sparse and you're going into something denser, but it's fixed in dimensionality and the dimensions that you want it to be fixed in. and that I think is what makes it more beneficial from an attention perspective. I tend to find sparsity to be a big problem for machine learning. I don't usually think of it as a benefit.

30:23In a lot of cases, models can't handle sparsity particularly well and we often try to use approximate techniques to remove the sparsity and to make things more dense. Even if you think about traditional embeddings like you know going way back to word to vec the goal there was to move away from one hot encoded vectors where you could think of your document word matrix as this giant matrix of zeros and ones much like a graph is a you know an adjacency matrix is a giant matrix of zeros and ones and what researchers found very early on was that when you started to work with large vocabularies, it was impossible to model them effectively.