Scientist Stories: Jan Stühmer, Interpretable Representations in Deep Learning

Good morning, everybody. Welcome online. We are very happy to welcome Jan to Heidelberg. Jan studied computer science, went on to earn his PhD in Munich in Daniel Kremers' lab, and then worked as a scientist first with Microsoft Research and later with Samsung AI in Cambridge. And then decided to join HITS, the Heidelberg Institute for Theoretical Studies, up on the hill. And we are super happy to have such a high-caliber PI in Heidelberg. We're sure he brings new perspectives, and we are very happy to have you here. And also happy that Jan will give us an overview today of some of the things he did in the past, but also of where he's going with his lab. Jan, very welcome.

Thanks, Felix. Thanks for the introduction. Yeah, and thank you for coming here in person and also for joining online, of course. Today, I want to talk about two topics that I think are very related, and then also will give some outlook into future work that is planned. Planned for the lab. And those two topics are interpretability, first of all, in deep learning, and then also our most recent work that will be presented at ICLR in two weeks at the ICNR conference on neurosymbolic methods. And yeah, first of all, why should you be interested in interpretability and in this field of representation learning? So there are some applications also in biomedical research. So I thought as I give the talk here, at the Cancer Research Institute, I should maybe connect a bit to biomedical research. And so one of these applications would be, for example, when you learn a representation, a feature representation of a data set, for example, here in histopathology, where you would now just pre-train a model with some self-supervised learning task. For example, you crop a little region of the image away, and you ask the model what was here. And then it should just fill in the missing part. And in order to do so, the model needs to learn some kind of understanding about the data. And now the question is, what kind of understanding does the model learn? Did it capture the relevant features? And so for this, there might be some interest in understanding these feature representations better. Also, in biomedical research, very abundant is also that you might have data that comes in the form of a graph, and then you would like to learn a representation of this graph-based data. And I will show some early work that we did on this as well when I was at Microsoft Research. And this data that comes in the graph in biomedical research could be, for example, protein interaction networks or some other healthcare knowledge networks like patient data, for example. And now you would like to integrate all that or maybe not integrate everything in one model, but at least you would like to digest all this graph data that is organized in graphs with your machine learning models. So it might be also interesting to learn a representation for graphs in the end. Now, how do we gain interpretability? So in general, there are two approaches to this, which is that you could do it either intrinsically by using a model which allows you an interpretation by the way that the model was built, like from its intrinsic features, or you could do it post-hoc. You would look at the feature embedding apply statistical methods to gain some insights into what is captured actually in the features. And today, I'll give first some examples for the intrinsic interpretability. So this is from our work at ICLR this year, where we worked on solving Sudoku as an example task. And we used a transformer model for that. And what you now see here is the summed up aggregated attention layers of the transformer when solving the Sudoku, like a certain field of the Sudoku. I hope you all know the Sudoku rules, right? May only be one number per column and per row, and also only one number per 3×3 block of the Sudoku board. And what you now see is that when the model should find the answer for the top left square here on the board, then it understood that it has to pay attention to the column, the row, and this 3×3 block. And also you see that this kind of consistently carries through to the other fields. Interestingly, when we were working on this project, we were using some datasets from the Kaggle benchmark. These Sudoku boards there were generated by some algorithm that was taking a full Sudoku board and then would remove one number one by one. I think this algorithm must have been going through the Sudoku board from the top left to the lower right because our model learned that it only has to care about the 3×3 blocks, when it solves the first 3 rows. And then after that, the model used only row and column constraints for the last 2 to 3 block rows. And I was like, ooh, that's kind of funny, right? I don't think that this works as a general strategy. And then we could also make a counterexample that the model couldn't solve. So that was kind of funny. And in the end, yeah, by just reshuffling these datasets by applying permutations onto the Sudoku boards, now you see that this bias is kind of removed, right? So interpretability now also offers you some ways to analyze what is my model actually doing. Did it take some shortcut somewhere, like here, for example, in this example? So it remains a very important part when working with these deep learning models to gain understandability for the quality of the prediction maybe as well and to get better understanding when you can use these models and what the limitations are. Then the post-hoc interpretability example would be, for example, that you map your dataset into feature space and you perform some statistical statistics models, right? Like maybe most simple one would be PCA, principal component analysis. And if you know, for example, trained a classifier, then you would assume that this data somehow is then segregated in this PCA space, right? There are of course better methods to do that. I will speak a bit about this. And now what we were looking at in our work was some method that would more fall into a combination of both categories. So it combines this idea of building a model that is intrinsically interpretable in the way that we built the model in a certain way so that it has characteristics, certain characteristics, which allow us to perform a better post-hoc analysis, actually. And here, I will now show you some examples where PCA, for example, would not work. And this is if you had 3 linearly dependent factors. And now this data here was then generated by students' t-distributions around these 3 latent factors. And what you then can see is that our model could reconstruct these 3 factors, although they were linearly dependent. This is something that would not work with PCA. In PCA, you always assume that the factors are orthogonal to each other. And so, yeah, now I want to give you some more details about this work. And this was— yeah, the title of the book is Independent Subspace Analysis for Ansatz-Over-Resolvent Motifs and Tangled Representations. This was a collaboration while I was at Microsoft Research with Richard Turner of Cambridge University and Sebastian Nowotzin from MSR who is now at Google DeepMind. And first of all, what is representation learning? Well, I gave you already a little introduction to this and maybe more formally is that you want to actually in representation learning want to reconstruct underlying generative factors of of observed data. And a commonly used toy dataset for this task is the DSpides dataset from— I think this also comes from Google DeepMind originally, where you have like sprites of 3 different shapes, a heart, an ellipse, and a square. And those now rotate and move around in the image. And this is to make some controlled experiments. You can think of this dataset more for academic research, of course. But because these sprites undergo these different transformations, shape, scale, rotation, and the location in the image, now we can try to build models that now offer us actually to interpret the latent space and look for where these features here are represented in the space. And if we can find a representation which then factorizes in those, right? And one model that we're going to use for that is the Variational autoencoder is a model that combines some latent variable representation to which the encoder maps from the data and the decoder that maps back from the latent representation onto the data. And training is performed by optimizing the so-called evidence lower bound, the ELBO, which is then averaged over the empirical distribution. So we are computing an empirical loss here over the training data and the loss involves the Kullback-Leibler divergence between an approximate posterior of the latent variable given the input data and a prior for the latent variable. And now often in variational autoencoders when people use them, for some reason almost always the standard normal prior is being applied. And this has historical reason because the first paper used the standard normal prior, although the author said that you can actually use any prior. But for some reason, then these things carry on. Maybe there was some example code as well to download, and then you could just use that implementation. And this work then was extended by Higgins et al. that proposed to add a tunable weight into the ELBO so that we can now increase the weight of the cardiac lipolytic agents. And I will show you later what this then results in. But you can also think about this that when you increase the beta to larger values, then this means that we are less confident about our data and want to pay more attention onto the prior. So this means that then it becomes even more important what you choose as a prior then. And we see this in this example here. So for beta 1, we are able to reconstruct the generating factors, the generative factors. But if you increase the beta weight, then you get certain biases of the variation inference method that you use. That means that now you trust more and more to the prior, and because your prior is a diagonal Gaussian with a diagonal covariance matrix where all the different dimensions are orthogonal to each other, right? We have no statistically independent of the different components. And this now means that we orthogonalize completely and we we can now explain away these linear dependent factors by just fitting one Gaussian to it. And of course, this is not the true underlying statistics of the data, right? But we became some kind of blind to that by increasing this weight. So it seems that this Beta-VE framework is maybe not well suited for cases where your data actually has dependence, interdependence. Yeah, not to collect what we've covered so far. So your variational autoencoder typically uses Gaussian prior for the variable. But when our goal is interpretability of the latent space, then in general, this may be not a good prior because it's also rotationally symmetric, which means we could now align our prior in any direction, like all the axes could be rotated freely and you would still get the same probability, the same likelihood of our latent variable. And this problem is well known for latent variable approaches. For example, it's also called probabilistic PCA because when we use the normal distribution, then we assume independency of the different components like it's done in PCA. And there is a nice quote from Chris Bishop, of course, because he He was leading the lab in Cambridge when I was there. So Chris Bishop already in his book wrote that historically factor analysis has been the subject of controversy. Many attempts have been made to place interpretation on the individual factors, which has proven problematic due to the non-identifiability of factor analysis associated with rotations in the space. And if we wish to remove this degeneracy associated with latent space rotations, we must consider non-Gaussian latent variable distributions, which gives rise to the so-called independent component analysis methods where you use different distributions, not the Gaussian distribution. And the distributions that you use for— now I see the fonts are not working really well here. OK, so all these equations that you see that are popped down by one row, they are supposed to be there. The font of the slides is not properly installed on this laptop, it seems. So a generalized version of ICA uses now a prior from the family of exponential power distributions, which are of this form here. So the Gaussian distribution, as you might notice, fit— also is part of this family by choosing p and q for p and q— sorry, no, it's not p and q, it's q both times— for choosing the value of 2 for q. So then we get the square norm of the L2 norm. And by using now— by parameterizing this Q in this way, 2 by 1 plus beta, this is now a different beta than in the beta VAE. This measure of beta is now a measure of the kurtosis of the distribution. And for beta 0, we get 2 divided by 1, which is 2, so that's the normal distribution. For beta 1, we get 2 divided by 2, so that's a 1, that's the Laplacian distribution. And now in general, we get for beta larger than 0, leptokurtic, and for beta beta smaller, you would have a platykurtic distribution. This is interesting. Leptokurtic and platykurtic sounds like a flat distribution, right? But instead, the tail is flat. So it's always used the other way around. So here, this is now a leptokurtic distribution for— or am I not the wrong way around? And this is beta plus. Maybe I did it wrong around here. Now I'm confused myself. So this would be the platykurtic distribution. The left-diskurtic distribution. You see that when you reduce the beta even further here, then this here approaches the uniform distribution. And when we reduce this further here, then you end up with a Dirac impulse, right? And for beta 2, we get the normal distribution. And for 1, we get the Laplace distribution. And when you use these distributions and now don't use the special case of beta 0 but one of the others, then you get actually rotationally invariant priors. And we can now extend this framework. So there's work by Fabian Sins, which now extended this idea of using these Q-norms, and he extended it to PQ-norms. And now we can group these PQ-norms and organize our latent distribution in different vector spaces. And that works the following way. When you look at this equation here, we took the second component, third component, took it to the exponent and then took the p_1 root out of it, right? So this is now p_1 norm. But those now together goes with the first component into p_0 norm, which we do here. We take it to the p_0 exponent and take the p_0 root out of it. And this now means we— when you now want to write this down in some kind of graph representation, right, you see that here we have one vector space where z_2 and z_3 are taken into the p_1 norm, and here we have one vector space where the result of this norm and said, does that one go into one vector space? And I was wondering, why should I do something like that? So it seems to be less intuitive than just using the family of exponential distribution themselves. And the nice thing is we can build on these rotationally invariant vector spaces. And you also see that you get some corner cases. Those are maybe the most interesting. So when we take when we take a distribution that approaches the uniform distribution of them, we get this box. But when we make now one of them the L2 distribution of the Gaussian, then we get this, this cylindrical shape here. You can now organize the dimensions of your vector space depending on what you think, what kind of statistics are there for these components of your vector space. Another thing that you can do, and this now becomes more more tractable and the purpose is maybe more obvious as well, is that you can organize these different vector spaces in this way here, that you have all these sub-vector spaces that are now independent from each other. And those then all go into one norm. So we can now show that each of the vector spaces, V1, V2, V3, V4, are independent. So those are orthogonal to each other. But within the components of each vector space, we can have covariance. And we will now use this form of this prior, which then gives rise to a method called independent subspace analysis. So instead of assuming that our components of the distribution are independent, we now assume that some vector subspaces are independent. And yeah, this is a proper distribution. We can compute all of the quantities we need, like surface area and and normalize the distribution. But yeah, I'll not go into detail here. But we can compute that, right? Some nice connections with a gamma function to compute volumes in higher dimensions. If you're interested in these things, I think it's always nice to see the connection to geometry there. And the data set that we use are two. This one is the Display Rights data set that I already explained a bit. And this is a data set of faces, synthetic faces rendered from different different directions, different angles. And for this dataset, now the goal would be to define what was the viewing angle that we looked at. And also, I think the angle of the light is changing as well. So we want to now factorize the images regarding those features. And so now we use Beta VAE on the DSpice dataset. And how you read this plot here means that Here we vary the position of the sprite in the data. And when we now change this position, we plot how the different components in the latent space change. And we see that now this is in the Z6 component is now encoding the y-coordinate, and Z2 is encoding the x-coordinate. And we do not observe much change for Z1 and Z5. But we see that Z5 changes when the scale is being changed and that one changes when the rotation changed. So, Beta-VAE is able now to represent this dataset in the latent space in an interpretable way so that all these components are now independent factors that are encoding a certain feature, a certain property of our input data. And this is a competing method that was published after Beta-VAE, which even got got even better results. And this is now our method. And what you see, like when you go— maybe I go back to BetaVE— what you see here is that a lot of noise is present in the reconstructions. So this orthogonalization of the latent space actually resulted in a very noisy reconstruction. So the likelihood of the data given the prior is not very high anymore because we we were paying more attention by increasing the beta, putting more attention onto the prior. So our data, the fit to the data was then not so accurate anymore. Whereas in our method, because we don't have to change the beta parameter at all, we just leave beta at the value of 1. So we just use the standard error, but we use our prior that factorizes into independent vector spaces. We also get an interpretable representation and we get a much better reconstruction of the data in the end. And an intuition for this is that in the original work, the model had to compress the data into the latent space so that all the information about rotation or scale, for example, had to fit into one single component onto one-dimensional variable or this had to be mapped, right? Whereas in our case, all this information can be So in our case, all this information can be mapped into a whole vector space, and this gives the model much more capacity and better reconstruction of the data. And here are some quantitative results. So we are in the top right, where here we now learned the exponents of these different vector spaces. So we now fit the statistics of the subspaces to the latent distribution given the data. And you see that for beta 1, we now get the best reconstruction. So you want to be in the top right. This here is the reconstruction error that you make that gets better when you go to the right, the negative log likelihood. And this is the mutual information gap. This is some kind of measure how well the different features are represented in the latent space, or whether they are orthogonalized and how well they align with the actual properties of the dataset that we want to represent. So the higher the score, the more interpretable your latent space becomes. And you see that we end up in the top right here, whereas Beta VAE also reaches quite a high interpretability, but you pay for this with a worse reconstruction. For beta 2, you see this is the result. And for beta 1, we outperform all of the other methods. Then there are some more datasets that you can try this out on. There's this Cars3D dataset, which are now views of cars from different dimensions and different colors. And you see that our model now can— yeah, here you see a plot of some of the latent dimensions. And you see this— This is now changing the size of the car. It goes from a personal car to more like something like a van. And here we see that the rotation of the car is changing. Some of this should be some color change. Yeah, here's color encoded. So even on this more complex data, you can now use these models, which is nice because anything that you can digest with a deep learning encoder, any data like images or possibly even text— I didn't try this yet— can be now approached with these methods. Then I promised you that I wanted to talk about graph representation learning as well. This is what a research intern did. Oh, no, it was a master's thesis. It's a master's student of mine when I was at Microsoft Research. And we worked on so-called Erdős-Rényi random graphs. Which have two parameters. The one is the number of nodes and the other one is the density of the edges. And you see that our model now— so this is now the graphs grouped by the latent dimensions that the model learned from the dataset. You see that in the top right, we end up with a graph with a lot of nodes which is fully connected, and down here we have just one node which is not connected. Up here we get a ring, so we have as many nodes as in top right, but just one edge in between. Here we get all the complete graphs for different numbers of vertices. And this is then how the adjacency matrices look like of these graphs in latent space. And you see that our model learned the underlying generative factors of this distribution fairly well. OK, so much for this work. Now I want to speak about our work on neurosymbolic methods that will be published in— yeah, will be presented in 2 weeks. We got a top 5% rating, it's now called. It's not our presentation anymore, but now it's called top 5% of the ICLR. That's International Conference on Learning Representations, I think. So that kind of fits as well, right? Fits to our whole talk. And this is work when I was at Samsung AI together with my colleague Cristina Corneli. She was leading the project. I'm second author. I forgot myself here. It's maybe not too— maybe modest here. Cheryl and Terese Hospedales. Oh, it's because of the width. Yeah. So I was the second author on this work. Here you see the reference. So yeah, and this method is a neurosymbolic approach, which means that we have some neural component that is now working like neural network works, that it takes some input data, maps onto some output data. And we build it in a way so that we can learn each component individually in supervised learning. We do supervised pre-training of the models. And then we have a symbolic component, which is now a logic solver. So a solver that, for example, knows the Sudoku rules. And given the board, the solver then can solve Sudoku. But of course, if this problem gets larger and larger, then the solver here will run also longer and longer, right? Solving these problems with the logic solver, maybe not for Sudoku, but for other problems, can become very hard when you have a lot of variables. And so the overall idea of this approach is that first we compute an approximate solution with a neural network. And then we have a neural network that looks at the output of the previous stage and in the end compute something like a confidence measure. So this network here now decides which of the results of the first stage do I trust. And the remaining values are then being solved by the logic solver. So by this, you reduce the complexity of the logic solver a lot. At the same time, you get predictions which are much better because this model here doesn't have to solve the original task anymore. It just has to look at some some intermediate result of the first stage and decide which ones of these do I trust and which ones don't I trust. And for multiclass classification, this model here would now see the overall distribution of all the different classes. So even by looking at that, you can easily decide where the uncertainty of this model was maybe too high and where you don't want to trust it. And this is now all in an end-to-end trainable fashion. And we can train this overall model with reinforcement reinforcement learning. And this is now an example of how this pipeline works. So first, you take a Sudoku board, which now has the blanks to fill in. Now you have your neural network, which tries to solve the Sudoku board. And we see it made a mistake here, right? There are 2 times the number 7 in the top left square. And so now these mistakes are now marked by the second stage. And then the resulting gaps in the solution are then filled in with the solver, with the symbolic solver. And here is finally what I showed you in the beginning. So this is now the attention layer of this first stage, of this solver stage. And we also applied this to a different problem than only Sudoku, but to scene graph estimation, where you were given an image. You should now mark all of the objects in the image and then reason about the relationship between those. So famous example is a rider is riding a horse. And I see that the neural model that is used for this task— so this is like a standard current state of the art is to use this kind of model, a neural network that gives you the relations between these objects. And you see that then you get sometimes mistakes. And now because we have this controlling stage in the second stage, these are now being rejected. And instead, then with a symbolic solver that knows about an ontology, right? So in the ontology now, it might be encoded that a person is not riding a jacket but instead wearing the jacket, right? And so by this, then the solver can solve this scene graph and results in a better result than if you would just use a neural network alone. Here are some results of our method for Sudoku, the quantitative results. These are SATNET, NeurASP, those are state-of-the-art neurosymbolic approaches. And you see that we are outperforming those in most cases. And interesting is also when you look at the results per modulator, we do supervised pre-training of the models. And depending on the difficulty of the dataset, this works without reinforcement but sometimes also fails. And now with reinforcement learning, the overall pipeline is tuned, right? Like the second stage of the second neural network stage learns which results to reject and which results not to reject in a better way. Then I want to give some outlook. So yeah. I promised that I also will speak about what I will work on. So far, we were looking into the past and now about my plans for the research group. So I would like to carry my research on interpretability further to causality. And I now want to just mention two existing works from other authors. So I don't have publications on this yet. And afterwards, I will also speak about some more plans in the biomedical research domain from my research group. One connection between— between interpretability and causality is this independent component analysis that I was already speaking about. And there's now an extension by Huwer Wienen from UCL and University of Helsinki to extend the guarantees that you get for independent component analysis also to the nonlinear case. And there are certain conditions, for example, if you have non-stationarity of the different components, that are generating your data, then actually with nonlinear ICA you can reconstruct these. And why non-stationarity? Well, you can come up now with a self-supervised learning task. That is that you take your data and you cut it into different pieces, and then the task for the model is to decide, was this data now permuted or was the data in the correct order? And this is also called time contrastive of learning by Huberin and Murayoka. And this is nice because it's self-supervised learning. Many models nowadays, for example, also GPT was trained by self-supervised learning where the task was to fill in a blank word in a sentence. And here, now, the task would be I just take some data set. I reshuffle it. I don't reshuffle it. And the model has to decide which of the two cases is now present here just from looking at the data. And why does this task now offer you interpretability? Well, in order to decide if the data was reshuffled or not, the model has to gain some insights into the generative process underneath. And if the condition is non-stationarity of the generative factors but smoothness over time, then by solving this task, if the data was shuffled or not, you understand something about how the vectors are actually moving through time that generated the data. And why not causality? So this is not causal yet, right? This is just factorization of more and more complex data sets. There's an extension to that by Monte et al. where— and now this is the limitation still— they can decide which of the two cases holds. Is x2 depending on x1 and some additional noise for each of them, or is x1 depending on x2? So that's the causal inference that you can gain so far with ICA. It's limited to two variables. And of course, now it would be interesting, can this be extended more broadly to inferring larger graphs? And another plan for my research group is in molecular modeling. So last year at NeurIPS, there was a workshop on graph represent— on machine learning for graph— with graphs and graphs. And a lot of the applications currently are in the biomedical domain. It's about molecular modeling, either— in this example, it's a diffusion model for molecules. So you would start from a Gaussian cloud of atoms, and now the model would because it learned how to— so during training, you inverse the process. During training, you start with a model, and then you apply more and more noise to it until you arrive at a normal distribution. And because the model was trained during the training to invert this process, you can now start with Gaussian noise and invert the process and end up with a molecule. So this would be diffusion models for molecular modeling. And the other project this collaboration with Frauke Greta would be to use more geometric methods even. So I was also teaching a geometric deep learning class at KIT last semester and also a compact course last week here in Heidelberg. And the idea of these approaches is that we use symmetry and group theory to build models that then become invariant to certain group transformations. And for molecules, one of the group transformations most people are looking at at the moment is rotation invariance. So independent how a molecule rotates, the resulting properties should stay the same. And here is even a nicer application, more geometric application. This is for creating backbones for proteins. So this now is a diffusion process that takes in a collection of coordinate frames that then later on become— when this denoising process goes further, then the result of the diffusion process is now a chain of these codon eggfrings that then form the backbone of a protein. And now, of course, it would be interesting, can we provide context to this? Can we build proteins with certain functions? Can we build proteins with certain binding sites? And yeah, this concludes my talk. And now I'm happy to take your questions.