Relational Foundation Models for Enterprise Data with Jure Leskovec - #768

There's always exciting research uh going on. Our uh uh focus recently has been I would say on two areas. First is uh AI for science, uh and in particular in uh we have a project that we call AI Virtual Cell.

where we are basically building next generation foundation models that allow us to represent human cells, patients, as well as individual molecules in cells and allow us to re to reason um across this complex biomedical data for, you know, discovering new new cancer therapies, molecule design.

reasoning about all different biomedical data modalities, objects, and how they interact with each other to help speed up science. So it's everything from... foundation models at the lower level of understanding proteins to then models that aggregate the um let's say the molecules in the cell to represent a single cell, and then the next level models that now say, oh, you know, a tissue or a patient is a collection of cells, cells are collection of molecules.

Let's build models that just aggregate all this knowledge in a very faithful representation, let's say, of a of a patient. Um and that helps uh a lot because now the representations we have are much more uh robust.

I can use the protein language ball to now represent every protein in the cell. And now the cell needs to aggregate information from all these po proteins to say, I am a cell, this is my Right, now that you have a representation of the cell, you can build let's say a patient level model that says a patient is a collection of cells in given state. that are composed from the proteins that are in there and are we able to kind of collect this information over this?

orders of magnitude different scales to get a strong data-driven representation of the underlying patient in this example. And the interesting thing is that this is purely doable. and it's uh trained purely in an unsupervised, uh self-supervised way, right? So you don't need to insert any human bias, any human knowledge of biology. The biology emerges From the data itself, right? Like cell types, cell states, relationships between them.

That kind of human biology, how we describe it, actually emerges directly from the data, right? So the model learns how to best describe the underlying processes and phenomena without us. Pushing it on it from the top. That's kind of the exciting, uh interesting, kind of emergent capability there.

uh and then for every cell in that sample you measure the the number of different protein molecules inside that cell. So every cell is now represented by a twenty thousand dimensional vector that tells me the abundance of that specific protein in that specific cell. Right, and every cell has different

So that's the that's the raw input data. And then of course because we know what the protein is, we can actually bring the protein information through through um ESM or through alpha fold, uh and now it gets and then it gets very

uh the immune state of the of the entire body. So we are able to uh detect diseases, uh understand patient trajectories and things like that just from this digital twin of a single drop of blood.

Um and you know, where I started doing this was actually in a in a in a third domain, which is um uh computational social science. Right. I was I was very excited about how do people interact with each other. And when I started my research career, kind of social media just just started as a phenomenon. And and and my view at that point was like I can use social media as a telescope.

Υπότιτλοι AUTHORWAVE through using cell phones through using social media and so on and it's all about networks graphs of people interacting with each other How the virus will spread as we reopen the economy. If we increase the occupancy levels at different, you know, restaurants, gyms, uh churches, wh whatever the locations to say this is how the virus would spread, this is what you can do. And and that underlying was an So now what is biology?

And then, you know, what is what is a what's a tissue, right? It's again cells coming together, talking to each other, organizing in a given way, so that, you know, my skin has a given structure, it has given set of layers and so on. So that essentially a network a graph of interactions as well. And then, you know, you mentioned relational data, right? So data that sits in tables in a database that every enterprise in the world has and it's kind of the most valuable data.

That's also capturing a graph of interactions of different entities inside that organization.

Then we said, oh, we have gradient boosted trees, they are better, and things like that, right? But fundamentally, it has always been, you have your data, You feature engineer this single table of your features, you add a label, and now you train some supervised model that from the features predicts that label.

Right. And it we've been doing that over and over again. And maybe this predictive model, you know, it's a it's a deep model, we would call uh but um it's a neural network. But what I would argue is it that AI has not transformed this structured data space in the same way as uh computer vision or natural language understanding have been fundamentally transformed by AI.

Okay, and let let me let me quantify what do I mean by that, right? Like what was the big breakthrough, both in computer vision as well as in natural language? It was about let's build neural networks that learn directly on the road. I think in the old days, you would do, in computer vision, you would do all kinds of feature engineering, shift features, Gabor filters, and it'd be like, I'll describe this image as well as I can, so I can then predict, you know, is there a car on the image?

The same thing is actually not happening on structured tabular data. Right there, we don't learn on raw data. We run all these SQL queries, all this ETL, all this feature engineering to then come up with a set of signals from which we, let's say, try to predict. And when we came up with this idea of relational deep learning, our goal was to fundamentally disrupt this and say, hey, why can't I just learn directly over all relational data?

Um and why why do we always have to learn over a sing data in a single table?

And the point is that as I take this multitabular data, and just to be very precise, what's a good example of this? It could be like, I have a set of customers. I have a set of products. So these are two tables. Each customer has an ID. Each product has an ID. And maybe I have a third table that's a set of transactions that says, Customer I did this, bought product I did that at this at this time for this product.

And that's a three-table, super simple schema. And of course, organizations have schemas of 50, 60 tables and more, depending on their complexity. So our question was, how could I just learn directly with a neural network over this multi-tabular thing? Um and the answer is, you know, and kind of surprisingly simple, is to say just think of the database, think of these tables as a graph of relationships between the entities in the database.

so this would mean inside in my you know i i'm a graph person so i like to think in terms of graphs right so graphs are composed of vertices the nodes this would be my users would be my Would be my products, would be my transactions, and so on. So this would be now the nodes. And then the connections are just saying.

This user ID was part of this transaction that was part of that product. And now we have a path from a user to the transaction to the product. And then, you know, another user or another transaction is another path in this very simplistic graph. Um and That we have a graph, we can basically apply graph deep learning, like graph neural networks, which is a way to generalize deep learning to graph structured data, and just train over that to get an accurate prediction.

And what happens is two things happen. The first thing that happens is you don't have to do manual feature engineering. Right, so it's much faster, it requires much less effort to train these models. And the second thing that happens is your models are more accurate. And then you say, why can my models be more accurate? And the answer is

very similar to what happens in computer vision, right? If you are saying, I am a human, I know what a car is, so I will build perfect features that detect whether there is a car on the image or not. I know cars, I drive them, I'm such a car expert, I can build the best features for detecting cars. Nobody in the right mind claims that.

right but you know in in machine learning data science prediction people are still saying no i'm the domain expert i'll engineer Your features are just some arbitrary human-biased summary statistic of your data. right and a neural network that trains with gradient descent is able to do so much more nuance almost like feature discovery by basically attending over this graph to extract much more signal.

So we see this double-digit increases in model accuracy because the neural network is able to extract more signal out of the road. and i will just you know full transparency right if you are working on a super simple problem that falls on a line then no neural network is ever going to be better than a linear model.

Right. So what I'm basically trying to say, I cannot guarantee that always you will get better performance because sometimes the data is linear. And if you happen to train the linear model to it, you already have good performance. There's nothing more you can do. Right. So that's kind of the key idea behind relational deep learning is that now we can have neural networks just learn directly on the raw database data.

Don't need to build these manual feature pipelines and feature stores that are super painful and lead to so many different kind of bugs and inconsistencies and information leakage and time travel makes putting models in production super hard. Um rather just bring the raw data, have a neural network, and and and and have you know get get better results that way. That's kind of the the let's say the philosophy and the reasons why we are doing

Uh depends on the data. If you have a transaction graph, for example, where we see uh great results in is on fraud, all kinds of fraud detection, anti money laundering. Account-level fraud, transaction-level fraud works beautifully, right? You just bring this heterogeneous multi tabular data together and just learn over it what fraud is. And fraud is interesting because it's so non-stationary.

You know, fraudsters are trying to game the system all the time. So as a as a machine learning engineer, you are always behind. Your model is always deteriorating. And you're like, okay, how do I design the next feature? How do I design the next? you have a neural network and you just pick the signal directly out of the raw data. So fraud is an example. fraud you can think of let's say as a classification uh task.

Uh then you can think a lot around uh regression type tasks, for example for customer behavior in terms of customer churn, next best action, um uh things like that. because it's predicting a link between the customer, the user, and the product. So we've seen great uses of this in recommender systems for ads, product recommendations, and things like that.

Right. So the hard part and where the information gets lost is when you go from this rich relational structure into the single table. And once you are in a single table, you know, then we are you know, then we are kind of talking

It's almost like second order effects. Did you use this architecture, did you use that architecture, did you use this tabular model or this tabular foundation model or not? Right? Like All the information is there in that single table, and all the methods are about equally good at extracting. Right. I think where the difference happens is if you actually make a step back and say, Hey, single table model

uh or i i is not the hard is not the hard part. It's not where w is not in a sense uh general or realistic enough. Where you need to go. You need to go to the multi table setting. Because that's truly now the raw data you have. It's not some summarized, featurized data. It's the raw data. And there is much more signal there that got dropped when the data got flattened into summarized into a single... So to me, single table problems.

are are you know, are solved. I think the differences are kind of second order effects. What is unsolved is the multi table problem. Th that's where the wins

Right. So uh what we did actually at Stanford, we were like, okay, so where is a multi-table benchmark? And there is no multi-table benchmark. And even if you look at Kegel, out of thousands of competitions of on Keggle, you know, there are four that are multi-table. All the others are features have already been engineered for you, there is a single table and you start you know

Bagging and boosting and and and creating tricks until until you win, right? Um so we created uh a benchmark at Stanford. Um by collecting and curating open, multitabular data sets that we were able to find on the web. We call it RelBand.

We have now two versions of Railbench. It's about 40 different predictive tasks over, I think, about 10, 15 different databases. And then what's also interesting is that SAP... the big German IT company, they released a benchmark, a multi-tabular benchmark of enterprise data called Sol. So those are I would say the two the two big tabular, tabular or multi-tabular, so relational uh benchmarks, uh SAP from Salt, uh Salt from SAP uh and uh Rail Band.

Line of work that they've been doing and promoting here through Stack.

And then, you know, some other data scientist says, no, no, let's take the cheapest price of everything you want, right? And then somebody says, no, no, you should take the most expensive. And then somebody says, No, it's the average. Another person says, Oh, but distributions are skewed. We should take the medium. Then another person wakes up and says, hey, it's about shopping in the morning. That's what's predictive of church. Let's add another feedback.

Um and that's a pre trained foundation model uh that can reason over structured relational data. Um and it's crazy what this model can do. So what this model can do, it can make accurate predictions on any database and any predictive task without any model training.

I give it the information, I give it a task, and then it gives me the answer. So what we do here is the system has several, several components. So there is the database. And then there needs to be a way for me to instruct the pre-trained foundation model what kind of prediction I want.

Right. I wanna say predict me the sum of purchase prices over the next one month for this particular customer. And that maybe is like how much pr I'm predicting how much the customer is going Or I'm saying predict me, you know, uh transaction dot is fraud equals true for transaction ID. Okay, so this would be like predict me whether the transaction is fraudulent for this particular transaction, I think. Right? So I have a way to specify my prediction.

And now what the system does, the system now goes into the database. It extracts a set of labeled in context examples that then get passed through a pre-trained neural network to make a prediction. Okay, so now when I say a set of labeled pre uh um in-context examples, this means that you can take the task.

Okay, so now what the relational foundation model gets on the input is a set of historical subgraphs of previous transactions and their fraud labels. The new transaction that is unlabelled, we don't know it's fraud. And then this is passed through the relational foundation model architecture forward to kind of label the unlabeled graph, right? Like the unlabeled uh transaction and the graph around it.

And then you have a pre trained model that is able to take these in context examples, these these subgraphs that have some certain columns and tables and so on is able to encode them in a domain-agnostic way, and then... the neural network is able to essentially build a predictive model in its brain in a forward path to give you accurate prediction.

And now that it's been encoded, then the second step is to perform in context learn. So it means that the model in its brain needs to be able to build a model, right? There's no training. There's no backpropagation. There's no gradient. It's just a single forward pass in which the neural network kind of

in itself builds builds I don't know the model in a sense that that gives us the accurate prediction. Right? So no training is necessary, no hyperparameter optimization is necessary, no feature engineering is necessary. All you need is a raw database and a way to specify the top.

A single forward pass of a pre-trained frozen neural network. There is no language model here, right? This is kind of technology that's parallel or or complementary to language models, right? You cannot Textify a database and then go to ChatGPT and say, hey, what do you think? How how likely is this transaction to be fraudulent? you get horrible results, right? So

Um and then if you fine tune the let's say the foundation model uh on on on on more data for the specific task, then you get to this superhuman accuracy performance that, you know, present manual or semi-manual or agentic solutions are just not able to attain.

You also recently published that iClear Relational Graph Transformer is the one based on the other or are they independent uh lines of research?

The architecture we are using is a bit different. It's composed of two different parts. The first part is basically the encoding or the attention mechanism over this. Set of uh set of tables. Um and then the second part is this in context learning uh type machine. Um there are two papers uh that are relevant here. One is the relational graph transformer that we only mon mainly use for supervised uh fine-tuning type tasks.

But then another paper we also published at iClear, it's called the Relational Transformer. And that one actually allows for in-context learning.

So that one does attention all the way at the individual cell levels of a database. And essentially, you have three types of uh attention. You have attention over a given column. So if you are interested in a cell by attending over other cells in that same column, it kind of gives you a sense of a distribution. Then we attend over the cells in a row. Uh and that kind of then gives you a sense of what's the information al in that room.

And then we also have a graph-based attention mechanism that allows you to say, oh, this is a user and these are all their transactions. And then each transaction is a row and each row has columns. So this means that we can be attending over millions, tens of millions of cells. And the beautiful thing is that our attention mechanism, because of the graph, has much more structure, so the attention mechanism is never quadratic.

And this means we can we can compute much more effectively. And to do good reasoning, you really need humongous context sizes, right? Even the largest LLMs today, I know, go to a million token.

this, you know, works well for or conversely doesn't work well for. It sounds like part of that is size. Like you need a lot of data in order for this to work. Is that fair?

You just tune the model for that single task, and then the model can be smaller, much more efficient to run, and also more accurate because it doesn't have to, you know, almost like re relearn the the task every single time because you give it the What we see works best is some mixture mixture of pre-training and in-context examples, because the the way you choose in-context examples can actually depend on what the what the target entity is.

So in a sense, you'd say, oh, if I'm predicting fraud for me, then you could say, oh, let me put some other Stanford professors in my in-context examples. Let me put some other Bay Area folks in here because, you know. That's kind of the peer group or the most useful examples from which you can learn to make accurate predictions about me being a...

Just build the best model you can and see how high you can get. Um uh uh the foundation modeling improves that I think for about five percent uh relative uh the accuracy. Um and then if you further tune The model, meaning if you would fine-tune it and do some gradient-based updates, then the performance goes to 12% over the state of the art. And those are quite sizable gains, especially if you think about putting this in production.

recommender systems or fraud detection where, you know, every single digit performance in increasing accuracy can mean millions, tens of millions in business. Maybe the second thing I would say is where we see these methods also shine is with noisy and incomplete data, cold start problems. because of the relationships, b uh b because of the relational structure, the model is able to much better kind of hone in and be much more robust.

to the data missingness, data corruption and things like that. So we've also done quite a lot of analysis around understanding and sh and uh like How this performs in on real world data, uh sparse data, small amounts of data, noise, incompleteness. Irrelevant columns and the same.

Uh you don't y you are data poor. That's what technically is called cold start problems. So I still need historical labels, but to make reliable predictions, I don't need much data.

Uh there's a bunch there's I would say a bunch of different deployments depending on what organizations uh like and and prefer. And then yeah, in terms of let's say deployments or use cases, right, that uh at DoorDash it's uh um restaurant recommendations and the notification system, which user gets what no notification at what time of day and things like that. And we've seen, you know, uh revenue impacting hundreds of millions of dollars. Another great client we work with is Reddit.

Advertising models on Reddit are built on top of, or are built with Kuma. nearly a double digit uh increase in at uh click through rates. So basically the the revenue the yeah, it's like unbelievable. Like usually an entire team You know, like increases maybe 1% that accuracy year over year, right? Because click through the radius.

appended to their to their own features. Right? And even with that, there was there was a huge increase in the click-through rate because this signal that the neural network learned was kind of complementary to what the human feature engineering uh already helped. So actually the the model that is in production is combined from the neural network embeddings by by graph embeddings by Kumo as well as the manual feature feature engineering.

Uh so that's been that's been a great uh great collaboration. So it's add add the recommendation, click through rate prediction if you wanna think of it that way.

pure optimization place. There is so many kind of other business rules that need to trigger for the ad to be actually shown to the user. So you kind of need sometimes those signals to be able to trigger business rules.

Um right. So also we can scale to, you know, the si to the size of the entire edit, to the size of the entire Bitcoin or Coinbase. Uh these methods really scale. But then you know, with some with some clients like Let's say Databricks and Snowflake, they are using us to run their sales mods, predicting what the customer is going to buy next, which customer is going to convert into a paying customer, and this allows.

them to to optimize their sales team, right? And if you think about sales team data, that data is smaller because the sales teams are, you know, with hundreds, maybe uh, you know, thousand, thousand thousand people. So you can do well on small data as well.

Um and and that's uh that's I would say what the uh what the co what the cost uh what the cost is uh in the end. Of course these models and uh are not in trillions of parameters. They are billion-parameter type ones. They can be quite small, so they're actually quite efficient and cheap to run. because the reason we are making predictions is to make decisions.

Right? Our commander system is making decisions what to show to every user. So we are like the speed of those decisions is at, you know, tens of thousands, hundreds of thousands, millions of times per second. So performance cost.

That's been, I think, the biggest bottleneck, right, in a sense that models are built, developed, they work great, but then engineering teams need to hook them up to actually, you know, surface those predictions or make decisions.

uh based uh based on on those uh on those predictions. Another another right, like maybe that's one use case. Another use case is sometimes where we shine or where the technology shines in these predictive, well-defined predictive type products that can be mathematically well formulated and optimized.

If it's more about, hey, we wanna understand the patterns, we wanna understand what is happening in the past, that is much more, you know, this kind of traditional data analytics type things or some pattern detection type thing that that our platform and what we discussed is maybe not a good use case. So that's, I would say.

Right. And and where we also see, I would say, a lot of traction recently is in agentic work. Because agents need to make decisions. To take actions. And now you can make decisions based on this, you know, LLM-based common sense. But the best way to make decisions is to estimate or predict their downstream effect.

Right? something I would say common-sensy, but that's far away from optimal. So these predictions, this reasoning over this structured relational data that captures the patterns, behavior of, let's say, customers inside the organization. is crucial to make accurate decisions. And as we are deploying these agents, we cannot be building now separate models for each of these and pre-anticipate the questions. The beauty of the foundation model is that you can ask any...

And one thing I wanna say here it's like just to show how big the problem is, right? Like if you think about a organiz like maybe like uh you know, think of let's say SAP, right? Uh SAP has I think I know, seventy, hundred thousand customs.

Each of their customers has structured data because it's an organization. Every one of them... changes the schema a bit so everyone has their own data, have their own schema, and every one of those wants to do a churn prediction, uh wants to do churn But every one of them has a bit different definition of what churn is.

So now can you hire seventy seventy thousand data scientists that are going to build per client churn model with the client's data and the client specification what the churn churn means? You can't. A foundation model An agent can just ask. Predict me probability of churn under this definition, under under this data, and you get the answer half a second later.

um, you know, separate from the individual customers' data or does the foundation, you know, the breadth of the foundation model already capture all of the information? you know, at that like use case level of abstraction and really you're only improving it if you're looking at a specific customer's data.

So that prediction then like the data gets better encoded and prediction later will be more So that's one reason you would want to, let's say, train even in a task agnostic way over the underlying data to better capture distributions, to better learn. Another reason why you would want to fine-tune is for cost reasons. Because if I fine-tune for a specific task, I don't need to do ICF.

Right, because now I don't now my context is is much smaller. I don't need to bring in the label data, I just bring in the entity I wanna predict on. Now the attention is smaller, it's faster, it's cheaper to run if I have large amounts of data.

You know, the mall the mall can learn uh can learn a lot from. So I would say there is there is, you know, there is a spectrum, there is a continuum of what you can do and why and the benefits of it is kind of different depending on where on this continuum are you seeing.

Right, and and for example if you s you could say, Hey, why don't I just, you know, give this modeling task to Clot Code and Clot Code will build the model for me, so you know, what's the big deal? Και όταν κάνουμε αυτό, όταν δείχνουμε, όταν δείχνουμε, όταν δείχνουμε, είναι ότι... these models write thousands of lines of code but there are these like super subtle data science mistakes

So for example we've we've done this uh together uh together uh uh with Expedia. Um and uh you know when when it was uh account level fraud. And mistakes, for example, the agent make was that when it created features for that given account, it created, it aggregated the transactions till midnight. Not till the current time.

Right. So it said, oh, today is I don't know uh uh April thirtieth, so we'll use the data up to midnight of April thirtieth, not actually saying, Hey, it's actually ten A ten ten AM on April thirtieth. We can only use data up to here. Right? So that's information leak. was a little mistake in that another mistake it made was that you know we did it at a transaction level instead of the account lab and these are like these subtle mistakes that really you need the human

But if you give it this more higher-level Kumo-like API, then it's able to do the same work in about 50 lines of code. No mistake.

And and you know. Uh or I can say uh using the Kuma API, go build me the model. That only takes fifty lines. And the the prob the the now if you think of this analogy of steps, you know, fifty lines of code is maybe like two steps. Five hundred uh uh lines uh lines of code is uh is twenty steps. Right and a lot. Uh navigating, I don't know, twenty steps. I think the observation is more general. be effective. They need APIs that are agent or agentic-friendly.