Zero-Shot Predictive Models for Relational Databases

Ankit Gupta 0:00 [upbeat music]

Ankit Gupta 0:07 Uh, Mark, tell us what you're working on.

Mark Žnidar 0:08 Cool. So, um, we built the first, um, zero-shot predictive models for relational databases. So let's say you're a company, you have some data, and, um, how do you predict the model on top of it? You can essentially put it in a table, but doing this, you lose some information, which is present, like foreign key, primary key relationships in your relational database, and this is flattening to do some feature engineering on top of it. And we say, perfect. Um, like this is-

Ankit Gupta 0:35 We want the model to have access to those.

Mark Žnidar 0:37 That's right, yeah. Essentially, like what, what I described is what data scientists would do. But we can say, "Hey, what about if we can make a model that doesn't require a data scientist, which kind of speeds up the exploration of your data, um, and maybe even gives you better predictions?" Like, so essentially, yeah, we develop a model which consumes the relational database as a primary modality. And the first thing you can notice, notice is that relational database can essentially be seen as a graph, where here, let's say, uh, a purple color means a

Mark Žnidar 1:07 transaction, transaction table. So this kind of abstraction, which we used to work on. And then from this context, um, context graph around the entity we want to predict for, so we want to predict for a particular, let's say, article in the table. We kind of collect the neighboring context, um, by sampling a neighbor- neighborhood graph. And then how, given this graph of like other tables and information in those tables,

Mark Žnidar 1:37 how can we produce, uh, uh, let's say, uh, appropriate format to do predictive, uh, modeling on top of it? It's like you can just kind of walk along this graph and collect the features along the way. So-

Ankit Gupta 1:47 Nice

Mark Žnidar 1:48 ... essentially DFS, BFS. So we, we, we, um, use this and-

Ankit Gupta 1:53 So that basically is your context into the transformer.

Mark Žnidar 1:56 That's right. So yeah, let's say this is our graph context, and we serialize this. Serialize means like here we produce a sequence of cells. Sequence of cells. Um, so for example, this would be like a gender of a user, let's say salary or something. And, um, then we embed those using, um, different encoders, like number encoder, text encoder, Boolean encoder, so on to like, uh, get an embedding presentation. And this is where the, let's say, the most novel part of

Mark Žnidar 2:26 our, of our method comes into play, is that we deploy these special attention mechanisms. It's kind of a mask attention, which allows our model to, for example, column attention allows it to attend over a particular column, which gives a distribution, meaning, let's say, hey, there is 90% of women and 10% of men in, in our data, which is useful for models to do predictions, as we know.

Ankit Gupta 2:49 Sorry, it's attending over the column's values or over the column, the column options?

Mark Žnidar 2:55 It is attending essentially like, let's say, um, you sample the user you want to predict for, but you also sample in a neighborhood, let's say three of his friends. And all those friends maybe also, uh, also like have a gender column or, I don't know, salary column, and having those values in context, you can maybe, for example, attend over only the salary column. So the model, like it learns the distribution of that column, which is used to do like better predictions down the, down the line.

Ankit Gupta 3:22 Cool. Makes sense.

Mark Žnidar 3:23 And then feature attention kind of just collects all of the information about user you want to predict for. So to kind of like have the, yeah, user information present here. And then neighborhood attention is like essentially you collect information from your, like, neighbors, which you kind of know that is very rele- relevant to, to do, maybe extrapolate or like extract information from the neighbors. And from there on, we also apply, uh, pool attention, which is like non-mask attention, but it was shown that it doesn't really add any significant boost.

Ankit Gupta 3:56 Significant boost.

Mark Žnidar 3:56 Yes, yes. So it was dropped. Um, so these are kind of like maybe the main contribution of like this special modeling approach with attention. And then from there on, you can just like do the codes that, that encoding and, um, do classification or, or regression on top of it.

Ankit Gupta 4:13 Makes sense. And how well does this actually do?

Mark Žnidar 4:15 It does very, very well. Like in comparison, for example, um, and this is also saying that LLMs are very bad-

Ankit Gupta 4:23 Sure

Mark Žnidar 4:23 ... working with relational data. Um, so for example, let's say, let's pick here, um, item-churn. So what is the kind of probability that, um, an item will churn is like the classification performance like 62 AUROC for Gemma, a four billion model, while it's at like 73 for our model.

Ankit Gupta 4:41 Ah.

Mark Žnidar 4:42 And-

Ankit Gupta 4:42 And it's like way smaller, your model.

Mark Žnidar 4:43 Yes. Like, like we can see here, it's like 22 million parameter-

Ankit Gupta 4:46 Yeah

Mark Žnidar 4:46 ... which is much, much smaller. Um-

Ankit Gupta 4:48 So you would imagine someone would use this specifically for the database access queries, for example. You know, you could have a big model query this model to do those really high performance-ly and then-

Mark Žnidar 4:58 Yes

Ankit Gupta 4:58 ... continue to use the big model for everything else.

Mark Žnidar 5:00 Essentially, this is something that was developed by my professor in his startup, Kumu.ai, is essentially like they developed an agentic MCP protocol that uses very similar model to this, but let's say industry, industry scale, um, to kind of use agent to ask the right questions and, um, write a query like SQL query to query the right, the right predictive quantity.

Ankit Gupta 5:22 Nice. Awesome. Thanks so much, Mark.

Mark Žnidar 5:24 Yeah. Thank you.

Ankit Gupta 5:24 This is, this is really great work. Thanks for showing it to us. I'm excited, uh, to see all the work you do next.

Mark Žnidar 5:28 Yes. Perfect. Thank you.

Ankit Gupta 5:29 Cool. Thanks so much.

Mark Žnidar 5:30 Thanks.

Ankit Gupta 5:30 Thanks.