Untangling Neural Network Mechanisms: Goodfire's Lee Sharkey on Parameter-based Interpretability

Hello, and welcome back to The Cognitive Revolution.

Today, I'm speaking with Lee Sharkey, principal investigator at mechanistic interpretability startup GoodFire, about fascinating recent work that he and co-authors have done to start moving beyond analysis of the concepts that neural networks represent between layers and actually begin decoding how they compute within and across layers. We begin by discussing why an understanding of concepts isn't enough. On one level, this needs no justification. An approach that explains only the meaning of the intermediate results between layers— while it does take a serious bite out of the black box problem— leaves the layers themselves as smaller but still unexplained black boxes. But more concretely, research has also highlighted important weaknesses of the feature-centric approach. Conceptually, a sparse representation of features inherently loses a lot of potentially important information, which is encoded in the structure of how features are clustered together and otherwise meaningfully arranged in space. For example, in some networks the days of the week are represented not by 7 random directions in activation space, but by a set of directions that lie together in a plane, such that a simple rotation operation can act as a sort of next-day function that converts a given day of the week into the next. If history is any guide, there is presumably a lot more such critical complexity to be discovered. And so with this motivation in mind, and taking inspiration from sparse autoencoder and similar techniques that effectively separate clean concepts out of their usual state of superposition, Li and team have similarly set out to decompose a neural network's parameters— which of course are used to process layer inputs into outputs— into simpler subcomponents that they hope will correspond to interpretable mechanisms that the network has learned. Their first approach, published earlier this year, was called attribution-based parameter decomposition. Li describes this as a sort of model unmerging. If you imagine counterfactually having started with a giant mixture-of-experts model and then merging all the experts into a single instance of the architecture, parameter decomposition would be the process of unmerging to recover those hypothetical original experts. Personally, I visualize it as expanding a neural network, which of course has width and depth defined by its architecture, into a 3rd vertical dimension, where each vertical slice consists of a simpler subnetwork that's presumably needed only a small fraction of the time, such that when all the vertical dimensions are again collapsed down to a single model, they all add up to the original network weights. Whatever mental model you prefer, amazingly, by constructing a loss function that incentivizes 3 things: first, faithfulness, or the accuracy of the sum of the sparse networks back to the original full network, 2, minimality, or the idea that as few network slices as possible should be active for any given input. And 3, simplicity of the discrete subnetworks derived from the process. You can actually get this to work! I really do think it's worth taking a moment to appreciate how incredible this is. As Lee says, channeling Ilya, sometimes the models really do seem to want to learn. That said, given the complexity of the target, you shouldn't be surprised to learn that this first method was still far from ideal. In addition to being computationally expensive and requiring massive memory for all those parameter copies, training success was extremely sensitive to hyperparameter choices, and there were conceptual issues with the gradient-based way that performance was attributed to parameters. Specifically, in some cases, critical components that had approached a local performance maximum would have a near-zero gradient and would be misclassified as unimportant. These weaknesses naturally inspired the second method that we discussed: stochastic parameter decomposition. For efficiency, instead of describing network subcomponents as sparse instantiations of the original network architecture, the new approach breaks each weight matrix into rank-1 components, which are matrix operations that read from one direction in activation space and write to another, and which can later be regrouped into larger units. Also, to better identify which subcomponents really matter, they replaced the gradient-based approach with a novel approach that uses stochastic masking to help the network learn to predict each component's causal importance, effectively identifying which elements the overall network can't do without. Overall, it's a more scalable, stable, and accurate approach, and on toy problems with known ground truth answers, it is able to successfully recover the expected mechanisms. Of course, as with any branch of interpretability, plenty of work remains. One big question is how can we algorithmically group the rank-1 subcomponents into semantically meaningful mechanisms? Much as sparse autoencoders require a feature labeling process, we still need to figure out how the rank-1 components ladder up to conceptually intuitive transformations. And of course, there will be plenty more challenges with scale. But as Lee explains, the promise of this work is tremendous. This kind of understanding could enable everything from surgical unlearning of capabilities to the identification and interpretation of novel scientific insights that models may have learned, including and perhaps especially from non-language data. All things considered, this conversation is a great chance to visualize what's going on inside neural networks, to ponder how much we've learned and how little that still seems to explain, to marvel at the fact that gradient descent and backpropagation can effectively optimize such complex architectures, even under such complicated multi-part constraints, and to appreciate the combination of brilliance and plain hard work that goes into figuring all this out. Right up there with understanding our own biology, and arguably no less important or challenging, understanding how and why neural networks do what they do is one of the grand scientific challenges of our times. So I really hope you enjoy this conversation about decomposing neural networks In Parameter Space with Lee Sharkey of Goodfire.

Lee Sharkey, principal investigator at mechanistic interpretability startup Goodfire. Welcome to The Cognitive Revolution.

It's great to be here, Nathan. Thanks so much. A huge fan of the show.

Thank you. That's an honour.

I'm excited for this conversation. I always love to learn, and I expect to learn a lot over the next hour and a half or so. Real briefly, before we get into the technical work, that's going to be the primary focus, You've moved recently. We've had 2 guests from your previous organisation, Apollo Research, on the podcast over time, focusing on risks of deception and identification of deceptive behaviours in frontier models. And we've also had 2 different episodes with folks from Goodfire. Dan and Tom have done a great job orienting us to everything that's going on in mechanistic interpretability. You were doing Mech-Interp at Apollo and recently moved to Goodfire. So maybe give us just a little bit of an update on how that came about, and then we'll get into the work itself.

Yeah, absolutely. So yeah, Apollo Research and I, I started Apollo Research with a bunch of folks, including Marius, Marius Hoban, who you mentioned you had on. And yeah, we started and we were focused on Detecting and ideally mitigating deceptive behaviors in frontier AI systems. One of the goals there for my team was to focus on the mechanistic interpretability side of that, where if we can read the thoughts, so to speak, of these frontier AI systems, and maybe that will give us a leg up on detecting deceptive behaviors without necessarily needing to rely on their outputs. And yeah, Yeah, I think we did. I was really happy with the work we got done there and really love the organization we built and really continue to support the work that Apollo does. I continue to be somewhat involved there. But yeah, we figured that it was probably better at this time for Apollo to really double down on the evals side of things, which left my team which primarily focused on mechanistic interpretability rather than evals, a bit in the lurch. And so we all agreed that this was the right step. And meanwhile, say maybe 6 months before we actually decided to make this decision, I had— it seems like ancient history now, but I ended up connecting Tom McGrath and Eric Ho, who went on to co-found Goodfire together. And so I had a B character role to play in the founding story of Goodfire. And so I was obviously very aware of Goodfire and massively supported the work they're doing. I think the team is absolutely amazing and it just felt like a really natural fit to move over from Apollo to Goodfire. And it made sense for some of the team to come along as well. So we can basically continue the stuff that we were working on there.

Was it 3 of you that made the move?

Yeah. So me, Dan Braun, and Lucius Bucsnak.

Gotcha. Cool.

Well, that's great. And I assume the $50 million Goodfire raise to support both team and compute certainly didn't hurt the The value proposition either, right? And we'll get into a bit more, you know, where the compute costs are for some of this work. But everything seems to be, you know, dependent on a healthy dose of compute. So the ability to draw those resources in from the private sector, you know, definitely, as opposed to, you know, Apollo being philanthropically funded.

Yeah.

Also, you know, it seems like it could be a big differentiator long term for the ability to scale what you're doing.

Yep, 100%.

Cool. Well, we're going to, I think, primarily focus on 2 papers today, one of which came out still under the Apollo banner and the more recent one under the Goodfire header. First one is called Interpretability in Parameter Space: Minimising Mechanistic Description Length with Attribution-Based Parameter Decomposition. I think there's a— try not to retread Well-covered ground too much. So I definitely would refer listeners who want to go deeper on that particular paper and have a really thorough exploration of it to the AXRP podcast that you did not too long ago. I thought Daniel did a great job there, and I listened to the full thing to help inform myself coming into this. I guess it seems like maybe I'll just ask you the first question. What is attribution-based parameter decomposition? Or maybe even more narrowly, What is parameter decomposition and how does it differ from the schools of interpretability that people are probably familiar with? And you can probably assume people, if they've listened to our feed much, at least are aware of the basics of SAEs.

Yeah. So what is parameter decomposition then? Well, let's contrast it then with activation-based decomposition. So in interpretability, typically in the past, what we've done is we've wanted to understand what is going on inside a neural network. And so what we've done is we've collected lots of activations of how the network has processed its input data, and you've got these intermediate activations and then it spits out an output. And what we've typically done is we started with the idea that, well, these activations should in some sense be involved in different things that the network is doing in order compute its intelligent output. There is some sort of learned algorithm that the network has learned basically to exhibit its intelligent behavior. And activation decomposition basically is the idea that we can look at these intermediate activations, piece them apart, and say this part does this and this part does that. By contrast, parameter decomposition, it is philosophically very similarly motivated, but instead of picking apart the activations, the idea here is to pick apart the parameters. Now, why might you want to look at the parameters? Well, in some sense, the thing that we're really interested in is the neural network, and the neural network is not its its input data, although it interacts in very special ways with it. But the neural network is in some sense the parameters, the architecture that pieces them together, the nonlinearities that connect them. And it uses all these components to transform the input data into activations and transform those activations until it gets to the output. And in some sense then, The parameters and the architecture and the nonlinearities are the thing that we really want to understand. They are the algorithm that— they are implementing the algorithm that the network has learned. And the activations in some sense are along for the ride, but they do interact in very particular ways where it's not quite as separable as that. But The basic idea here then is that the network is using its parameters in different ways for different inputs. And for an input where there is a cat, it will take these inputs and spit out a cat label, say. And for another input, maybe a picture of the Eiffel Tower, it may not use very many of the same parts of the network as it did when looking at a cat. And so in some sense, there's this idea of modularity. There's like parts of the network are doing one thing, other parts are doing some other thing. There's like a specialization in what the parameters are doing. And we don't really— the aim for a parameter decomposition is to find these modules that are doing specific specific jobs, specific computations. And the term that we're— that we, you know, say we're looking for here is we're looking for the mechanisms, like the mechanisms that the network is using in order to compute its behavior.

So to just flesh out that cat example a little bit more, and to say some really basic stuff, all these neural networks are composed of layers. Often the layers are the same exact structure repeated over and over again, although not always. We've observed by doing interpretability on the activations that we can see things like fur has been detected and eyes have been detected and a tail has been detected. And we sort of see this gradual move from low-level features to higher order features as we go through the layers. And so there's this story of kind of understanding that kind of mirrors, especially in these visual cases, kind of mirrors what I understand to also be happening in the human visual system. And so that's a remarkable coincidence. But when we are purely looking at the results of the intermediate calculations, we can sort of say what concepts are active at any particular time, with lots of caveats. Refer to earlier episode with Dan and Tom about the sort of philosophical gaps between the labels of these concepts and the underlying sort of what exactly is being activated there and how does it relate to the labels is a bit fraught. But nevertheless, bracketing that for the moment, we have these sort of labelled concepts that we can say, okay, this seems to This feature always lights up for all these inputs. All these inputs seem to be fur. And so that seems to be the fur concept that's activated. When that and the tail and the whiskers and the eyes are all— and the pointy ears or whatever are all activated, then in future layers we get the cat activated. But we're not saying anything there about how the transformations are happening from layer to layer, right? That's the big gap that has been sort of left by the sparse autoencoder work for future work. And this parameter decomposition is basically, is that future work that is now saying, okay, how do we actually move from these concepts through the layers one to another?

I agree with that. I'll say also, it's not the only approach that you might consider using in order to piece together how these representations at one part of the network may become other representations at other parts. There are other approaches that remain in activation space that you might consider also using. But yeah, they both try and achieve this idea of characterizing the computation between representations rather than just identifying the representations themselves. Maybe one way of thinking about this is like, identifying the variables used in computations rather than identifying the computations. And parameter decomposition and other approaches aim to find the computations rather than the variables.

So why do we need that? Why isn't it enough to say, well, we've got concepts A, B, and C in layer 1, and then in layer 2 we see concept D, and isn't that like Doesn't that tell us what the neural network is thinking, so to speak? What is not answered by that level of analysis that we still need to get clarity on?

It gets somewhat philosophical. However, I'll try to convey kind of how I think about this. I would say that we don't— what does it mean for the network to have to use particular variables. We can go in and we can say, look, there's this set of activations and this set of activations seems to correspond to data points on which there are cats, but maybe we can also find groups of inputs where there's a cat, but it's a cat standing in a particular position. And there's, you know, one group of your inputs is a cat standing in one position and another in another position. Now, does the network use the fact that there are, you know, 2 different positions of these cats, or does it just have a cat variable? Like, what would it mean to be able to, like, distinguish between a network that just has a cat variable and a cat— and a network that has, you know, multiple different cat variables. My proposition here is that what it means for a network to use particular variables is that these variables are— they are like the variables that the network does its computations over. Now, it may be the case that this network that had multiple cat in different positions variables actually might not use those things. They all just, to the network, they're all just cats. But you may nevertheless be able to look really in detail at the dataset and find lots of little distinctions within the activations, but the network might not itself actually use these. In some sense, we want some way to identify the set of things that the network is using. And to do this then, like, in some sense, You might say, well, this is putting computations first and representations second, rather than putting representations first and then computations second. This is— it's like— yeah, it's like saying the fundamental, so to speak, variables that the network is using are those over which it does computation. And we can maybe find things that that, I don't know, the network seems to represent. But there's not a particular reason to say that this is a feature that the network uses just because we can find it in there. Does that make some sense?

Yes, I think so. Though, as you said, it does get a little philosophical. I was also thinking of the finding of multidimensional features as sort of another window into why the assumptions that underlie the SAE paradigm may need to be elaborated and may need a little richer treatment. You want to tackle that as well? Hey, we'll continue our interview in a moment after a word from our sponsors.

Sure.

Yeah. So I guess since some of the earlier SAE work in the latest bout of sparse coding work in the deep learning space, which kind of maybe started late 2022 and kind of picked up in 2023, some of the observations that people had since that time were things like, well, it looks like the representations that we're finding seem to be spread across multiple different layers. And it's kind of confusing to think about what that might mean. There were also observations that these things that we were finding, well, they were directions in activation space, they're particular patterns of neural activation, but that didn't seem to be the whole story. There seemed to be some sort of multidimensionality to them. it wasn't just one direction, it was perhaps a plane of directions or even higher-dimensional spaces potentially. And it was kind of unclear how to think about these kind of higher-dimensional features, so to speak, as— like, it just was philosophically somewhat confusing to think about this. And one of the reasons that it was indeed— these were several of the reasons that motivated some of the thinking that led to parameter— or primary— sorry, attribution-based parameter decomposition, which was that, well, what if, say, a multidimensional feature just is an input variable to a multidimensional computation? And that's indeed why the network has bothered to go the length to structure its representations in this way. Without that, without there being some sort of multidimensional computation applied to it, you might ask, well, wow, why did the network bother to structure its activations in this really ordered way? There must be something for which it is using that order. And this was why we were thinking, well, maybe it's the computations that are giving this, that are kind of defining the structure of these features. And maybe then the same kind of thing might apply for the multi-layer case. You might think, well, it doesn't— for certain networks like transformers, which use a residual stream, it's very easy for computations to be spread over multiple layers because it just works out straightforwardly for them to do this. And it was confusing to think about this in terms of SAEs as well. And it just made some natural sense to think about, well, maybe it's the same feature if it interacts with a a part of the parameter vector that happens to span multiple layers. And for this kind of edge case as well, this multi-layer edge case, it just made a bit more sense to think about the computations as well, rather than representations as primal.

And so an example of that is like the days of the week, or the, I think, the original Grokking result with modular addition would probably fall under this heading where there's a— it's not like there is a— I guess it could vary obviously across different networks, but at least in some settings it has been found that it's not like there are 7 days of the week features that light up independently and are all like, this is a really strong Monday signal or a really strong Tuesday signal. But rather, there seems to be a plane in space through which the day of the week kind of rotates and remarkably, like, seems to sort of crystallize into this, like, you know, the graphics on this are, like, pretty amazing sometimes where it's like you pretty much divided the circle into 7 equal shares and each day of the week has a different direction in this plane as opposed to each of them having their own independent linear directions. And basically, I think the same mechanism seems to have, or a similar structure, if not the same mechanism, has been found in the Grokking result, where to do this modular addition, things are first kind of translated to an angle, then the angles are added trigonometrically. And then the result of that trigonometric calculation gets mapped back onto the final number. And that is seemingly hard to— maybe you can give a little more technical intuition for why that would be really hard to figure out just with features. But it seems like on some level it may just be that— I guess I don't know exactly what you would find if you applied the standard SAE methodology to one of those things. Because if I'm not oversimplifying the standard SAE methodology, you would be expecting to see these 7 distinct features. Or in the modular addition case, you'd expect to see— I'm not even sure exactly what you'd expect to see, but I don't know that you would expect to see a sort of cyclic trigonometric structure if you were looking for all independent features.

Yeah, I think it's probably easiest to think about the days of the week case. So to be clear, it can be both, that there is both a multidimensional aspect to these representations as well as individual days aspects to these representations. And the reason is that Sometimes you might want to— in some tasks, you might want to say, well, the day after Tuesday is, and you might just apply the computation that rotates the day of the week feature, because now you've just got one mechanism that you can apply to every single day of the week. But you might just In some settings, you might really just want one day. You might just say, I was on my way to the shop on— and there's really just one direction that you want to say here. Say it's Monday. And so in some settings, you may just want computations that apply to these individual days, and some settings where you might want these, like, higher-dimensional computation, this kind of 2-dimensional one that might rotate it, right? And this is, I think, you know, an important aspect of parameter decomposition. There's not like one basis that you might use to interpret, you know, the whole model in. It's more that there are steps in an algorithm, and, you know, sometimes these steps will Yeah, I guess just sometimes these— some steps will be useful and others will not. And yeah, I guess, let's see. So the SAE kind of approach might find— might not find this shared computation that just rotates from one day to the next. It might find— conditioned on prompts that say the day after or the day following or the day, so on, following Tuesday is, you might find such— sorry, you might find one for Tuesday, one for Wednesday, one for Thursday, just to rotate around every single time. This is just because it's trying to find sparsely activating parts of the dataset that have a particular shared property. And in particular, parts of the dataset that have a particular shared direction in activation space, which is not necessarily all the ways in which inputs might share properties. One of the ways in which they might share properties is they use similar mechanisms. And these mechanisms may do the same operation but to different inputs depending on what the prompt is. I'm not sure if that was clear, but yeah, this is one of the differences you might expect from parameter decomposition and an SAE-like approach. Now, this is not— there are other approaches that might get closer to what you might expect or want to find with a approach that identifies the computations. For instance, Anthropic recently released an update on MOLT. M-O-L-T. I will need to remind myself what that stands for, but there's a similar principle here where it's, instead of looking at individual directions, it's looking at more multidimensional computations.

Yeah. So I guess Maybe another way, a very simple way to say it, is we want to also understand the functions. Like, if the whole project here is to understand why the AIs do what they do and possibly be able to intervene in certain ways, we're just going to have to dig in and figure out the nature of the transformations as well as the intermediate results. And the intermediate results can shed some light on that. But at times there may be assumptions that in that approach that leave us blind to— I thought a key point you said was functions that are reused on different inputs and do the same semantically meaningful transformation on a given input. That is something that the SAE approach could be entirely blind to, but which obviously is a pretty important aspect and potentially happens a whole lot. I guess we don't really know yet how much that may happen, but certainly seems like it probably should happen a lot, right? I mean, the higher, especially as we get to more and more powerful networks, it would seem like we are moving away from or graduating perhaps in a sense from sort of simple rule-based combinations of distinct features and probably toward higher and higher level abstractions, which would presumably then correspond to more functions that can operate over ranges of inputs and do the same useful transformation to some space of inputs. In a way, that's chunking reality, right? And that's like an important aspect of chunking reality that it seems like the SAEs are sort of hinting at, but not really directly characterising.

This feels reasonable. Yeah. I think one of the cases that I like to think about is in an SAE, or say, a transcoder, which is an SAE that takes as input some activation and spits out a prediction of what the activation will be at the next layer. Both of these will have a problem where suppose the network is just doing a simple transformation, suppose it's a rotation, or suppose indeed that it's actually just an identity transformation, you might get— you're going to have to spend a lot of A lot of representational capacity to represent the computations that the network is doing at this particular transformation.

Why?

Because, well, you've got, say, you know, a million different features, and there's a million different features in the output, especially, you know, in the case of the identity transformation, there's going to be a million there. And you would need one, you know, input feature and one output feature for each each such input-output pair, whereas in essence what the network is doing is actually very simple. It's just a transformation of this type, whether it's an identity or a rotation, and you've just got one object doing that transformation. That's the kind of thing that we want parameter decomposition to be able to find, if it is indeed the case that networks are doing that kind of thing.

Yeah.

Okay.

So, we've got some examples of that coming up later. Hey, we'll continue our interview in a moment after a word from our sponsors. So, let's bracket that for a second and maybe just talk about sort of how it works and the sort of intuition for how the whole setup works there. I think there are a lot of similarities to SAEs. I mean, I'll do the quick recap of the SAE. So, we To do an SAE, you insert this very wide layer between 2 layers of a network. The thinking being that we know that the normal width of a network is way smaller than the number of concepts that it can handle. So clearly there's some superposition of concepts. It's not like one neuron corresponds to one concept. It's sort of Any direction in this 4 or 8 or 16 or however many thousand dimensional space, any direction in that space can represent a concept. And so the goal is to untangle those into a sparse thing. And for that, you need this really wide layer. So you train through this reconstruction loss where the goal is to project out into this really wide space, but have a sparsity term in the loss function so that some small, reasonable number of truly relevant features are activating, and then project back into the dense space so that you are recovering the behaviour of the model, but then can look at which things are activated. And then you go through this labelling process at the end of that to say, okay, here are all the inputs that led to this particular thing lighting up. What does that in fact seem to be? And with that, we now have this thing that we can use hopefully for monitoring detection of potentially bad concepts being activated, potentially steering. That's where Golden Gate Claude comes from, et cetera, et cetera. It seems like there's a pretty similar motivation here saying, like, okay, if concepts are densely packed in this superposition way at the activation as the results of the intermediate calculations, which we call activations, then perhaps the same is true about the computations themselves. The network itself represents or combines a huge number of probably much simpler, more semantically sort of intuitive calculations, computations, transformations, functions, whatever the exactly right word is there. And because they're all so densely packed and overlapping, and the sparsity of the dataset also is a key thing here, right? Because the reason that this can work is because certain concepts that may be pointing in very similar directions almost never occur together. And so even though they may look very similar, Since they don't occur together almost ever, that's fine. The network can sort of get away with that. And there's probably that sort of happening in computation space as well. And so the same question kind of applies. Can we break this up into a really wide version of it that separates the flattened layer or the flattened network that we originally trained into all these more atomic units of computation? And if we can do that, then does it appear that they actually have intuitive meaning so we can start to go through them painstakingly and understand what each one is doing? How am I doing there in terms of just creating the motivation? Anything I'm missing or anything you would Correct or complicate?

There's plenty that would complicate, but I think it's a really good intro. Yeah, I think that's just the right way of thinking about it. We think that there are many different computations going on inside this network, more indeed than it has neurons. And the way we think networks can do this is that they spread their computations out over there, like computational units, where these computational units are the neurons. And so this, by spreading it out, you kind of get some nice properties such as avoid, even though you might overlap with other computations, and because you've spread this out, you can kind of suppress this noise that comes from this overlap. And you don't necessarily need to, you can basically fit more in because you're kind of like silencing the stuff that you don't really want to be there, the kind of computations you don't really want to be active right now. It's very similar principles to the idea of representation and superposition, but yeah, just brought into computational space. And yeah, I think the other, thing that you mentioned. Yeah, it's kind of like there's this extra layer to what the network is doing. We've kind of split it out into this. It's not just, you know, we don't just have layers and width dimension anymore. We kind of have this extra dimension where, you know, you're asking like, what computation is the, or what computations is the network using at this particular layer, at this particular, you know, in this particular width dimension as well. And like, this does have precedent in other areas of deep learning. Mixture of experts is a kind of example of this where, you know, you have layers, you have individual networks, but you're— there's kind of— you can do different computations at the same layer, at the same, you know, neuron dimension, so to speak. depending on which expert you use. And so there are senses in which this is kind of like, well, suppose every single network is just a big mixture of experts that's been smushed together. Can we separate out the experts? One way I've heard it described is it's like model unmerging, where model merging is just like you've combined these 2 networks that do different things into one network that can do both. And we kind of want to do the opposite. We want to split this network out, this network that can do many different things, out into networks that can only do one thing, or as small a number of things as possible. Yeah.

So the setup is to actually do the training of this, to start to actually do this splitting, I think your work does a great job of really emphasising the loss functions. And so the constraints that you put on this are, first of all, that the— and I think there's a couple of different versions of it. We'll maybe start with the first one and we can explain how the second one is new and improved and better in a couple of different key ways. But if I understand the first one right, it's basically saying, okay, Let's make a bunch of copies of the network with essentially the same footprint to start. And we'll have the constraint that all these different copies must sum to the original. So if you had, whatever, a billion parameters in your original model and you decided you were going to split this into a million different subcomponents, you would then have a quadrillion parameters, which we can talk maybe about the computational challenges and the memory requirements for things like this. But okay, now we've got the constraint though, that those million copies, they must add up to the original. We then want to say, okay, akin to the SAE sparsity requirement, To make this, to make what comes out of it hopefully intuitive and tractable and semantically natural, we want to have as few of those be active as possible for any given input, which you call minimality. And then the final one is for each of those things, we also want to make them as simple as possible. So that they individually are interpretable. So again, we can split this into— and I sort of think of this as giving the network height. Obviously, we have width of the network is sort of how many neurons it has at each layer, and depth is like how many layers it has. And think of this as sort of splitting it into a vertical height dimension where instead of having all this stuff happening, in the same track of computations in a way that we can't untangle. We're now untangling it into all these different, hopefully atomic computational subcomponent units, only a few of which will be in use at any given time. And ideally, those can be kind of distilled to their simplest form while still, of course, reproducing the original network behaviour. A visual I have of this comes from the Tegmark group. I think it was Ziming Liu was the lead author on this paper, Seeing Is Believing, where they just trained toy, very small models on simple problems, but with a strong sparsity feature. And you could watch, they had these great graphics where you could watch and initially randomised network learn a particular function. But also you could see that most of the weights would drop out to zero as the thing sort of crystallised into its simplest form that could still do the task that it needed to do. Those are relatively simple toy problems. So again, there, what would you add to my setup? Of the— it's those. And so now we're going to do this giant training, optimising all of that together, which is kind of an amazing thing to me. I felt this way about SAEs, and I think I feel it even more strongly here that you can put all those constraints into a single optimisation problem and it works. If you told me that in advance, if you said, like, here's my idea, I'm going to put all these different terms into one single joint loss function, try to optimise it all together at the same time, I would say good luck. But that seems like it seems hard to find something that's actually going to work in that space. So, yeah, tell me if I'm missing anything, or again, anything you think is important to add for understanding there. And then I'm really interested in how do you account for the fact that this actually Yeah. Yeah.

So yeah, I think you did a great job on introducing the various losses. So just to recap, we have the faithfulness one that makes them sum up to the parameters of the— first of all, what are these things that we're summing up? These are— we're calling them parameter components. And these are the things we want to approximate one individual job or one individual mechanism that the network has learned. So we have the faithfulness loss, which makes all these parameter components sum up to the parameters of the target model. We have the minimality one that basically makes it Do its job, do the same job as the original network, but using as few of these parameter components as possible. And then the simplicity one, which is like, we want these things to use as little computational machinery as possible. And that one is pretty important because, well, they're all important, but this one's important because, well, a great way to satisfy the first 2 losses is just to use the original network itself. You're only using one thing. It sums up to the parameters of the original model, but you haven't done any work to decompose this. So you need them all to be simple. And the way we operationalize simple is that they should be low rank and involve as few layers as possible. And we can get into the technicalities of that in a bit. But yeah, I think the idea, I think I hadn't actually made the connection between, say, the work that the— Ximing Lu, did you say the name?

Yeah, I think so.

My apologies for messing up their name. But yeah, so I hadn't really made the connection between this kind of work and parameter decomposition, but the idea here is that suppose you aren't optimizing a network with the kind of sparsity constraints, and I believe some sort of locality constraint that they might have used so that there was kind of like neighboring neurons did similar things, if I recall correctly. But suppose you weren't optimizing it with these constraints. Well, the network is in some sense still able to learn. I would guess that the network is able to learn a very similar algorithm, no matter whether or not it's trained to be sparse in a particular basis and constrained so that these neurons are close together. I would guess that it's like we just need to figure out the way that the network is doing the same algorithm, but in some sort of basis that we don't have direct access to. The computations, if the computations are the same, the computations both have to be sparse, and we just need to find the basis in which they are sparse. And so yeah, it's just, we are assuming that the network is doing sparse computation, even if we haven't optimized for it, and then we're just trying to find the ways in which— the basis in which it is indeed sparse. With regard to whether or not— sorry, how do networks manage to do these things? Well, the networks, I think it was Ilyas Siskever who said, the networks, they just want to learn. And I think there's something really deep to this in that there's just many different ways for The larger your network, the more ways it has to get things right. And this is a pretty powerful principle, I think, just because it means it's actually just much, much easier to find the right way to do something if you're simultaneously looking in many different directions, many different directions in parameter space. And this So it's kind of, in some sense, somewhat unsurprising that you can add— we could add many more different constraints and it'll find some reasonably satisfying way to satisfy all of them. But it is kind of counterintuitive to our low-dimensional brains to think in these terms, and it's still kind of amazing. But yeah, I should also qualify this and say, well, The algorithm that we're talking about, attribution-based parameter decomposition, where your parameter components are very large and it's very computationally expensive and where the algorithm basically has got a bunch of problems, it is still somewhat hard to get this particular algorithm working, which is why we ended up doing future work. And so even though it is still kind of amazing that they can find some solution that satisfies all these things, it's caveated with the idea that, well, at least for this algorithm, it was at least somewhat hard for that algorithm. But Less so for others.

Yeah.

One practical question on the faithfulness concept, which is the idea that all the different vertical dimensions, let's say, of the identified subcomponents, that they all have to sum up to the original. Can you help me just, I guess, develop my intuition for exactly what's going on there? And I should say, most of the problems that you've worked on so far, I'd say we're kind of at roughly— you could again complicate this, but it seems like we're roughly at the toy models of superposition phase of this work, where the models that we're studying are rather small. And then a big question, of course, is going to be like, what does it look like to scale this up? But when— so it's not like a million layers yet is the key point there. But when one parameter in one particular layer gets turned up, that means that that same parameter position in all the other layers has to be turned down by the corresponding amount, right?

In these different parameter components, yes, as long as they all sum to the The parameter.

Yeah. And I said layer there, but I meant vertical layer, which is a concept maybe we should retire. But I am thinking about it in this vertical visualization, I guess, for myself.

Yeah.

One thing I wasn't clear on there, and I was thinking about this concept of attention sinks that I've seen in the past where I'm recalling this on the fly. But it's been found in certain cases at least that having a few sort of junk token positions at the beginning of a transformer can be really useful because in the absence of that, aka in the normal approach, you got to look back to something in attention. And so whatever tokens happen to be in those first 5 to 10 token positions, become really important because everything kind of has to look back to that. And so the introduction of this buffer, which they called an attention sink, was a way to say, let's not overweight the initial tokens of this sequence. In some cases, you might just want to look back to sort of nothing and kind of recognise that these initial tokens aren't actually super critical to predicting what comes next. So let's not force all attention to kind of land on those initial tokens. And that attention sink concept, I was kind of wondering, could there be sort of a junk sink concept in this setup where, like, how do I not end up in a situation where the new subcomponents that I've created are learning whatever they're learning, which is recreating the initial model behaviour, But maybe some of the other layers, which would never need to get used, are in fact just sort of taking on the opposite of whatever the ones that are actually doing the computations are learning. And could that create a sort of conceptual disconnect where my new vertical layers are learning computations?

Yes.

But how do I know that those are faithful to— I know that they still sum in that sense of faithfulness. They still sum. But maybe the computations that have been learned in these layers could perhaps be quite distinct from what is happening in the original network. And maybe that's being hidden by the fact that some of the other new vertical layers are just absorbing all these gradient changes in ways that don't actually mean anything because those never get used at all. And you've maybe just kind of trained some new stuff that didn't actually exist in perhaps, or doesn't necessarily correspond to a mechanism in the original.

Let me see if I have understood what you're saying. So we have these parameter components and they satisfy the constraint that they all sum to the parameters of the original model. And you're worried about the possibility that even though they all sum, there's a lot of different ways in which you can sum. And if one of these, we're worried about the case where there are some of the computations that, say, one of these parameter components is doing is kind of cancelled out by another one. Is it the case that both of these would be active?

Well, I'm assuming that if there is sort of a junk sink, component that it would presumably very rarely be active.

So we do in many networks expect there to be a junk component and a junk component that does actually do nothing. And this is just because networks, neural networks are degenerate. There are many different ways to implement not just something that works well, but implement the same algorithm. One degeneracy you might think of is if you scale up the weight before a ReLU and scale down the weight after a ReLU, well, this is very much the same algorithm, the same kind of amount of activation gets put back into the residual stream, but there are, you know, there's a 2-dimensional space here in which we can, well, in fact, it's a 1-dimensional space here, but in which we can move parameters such that exactly the same thing is done. And there are many more such degeneracies in neural networks. And so there are, you know, we basically expect one of the degeneracies that networks might use is that, well, they say, well, none of the— some of the— sorry, all of the activations that the network ever sees are, say, orthogonal to this direction in parameter space. So this direction in parameter space is in some sense never used, and you could indeed ablate this direction entirely in parameter space, and it would just not affect the outputs of the original model. Now, this would need to be an overparameterized model. You might not expect this kind of thing in the underparameterized case, such as language models, but in, say, an MNIST model, you might expect a junk component an MNIST model of a sufficient size, you might expect a junk component that is just not really used for the algorithm. Now, would you be able to— yeah, I guess, but I'm not sure this is exactly the same kind of junk component that you're talking about. Maybe it's something else in mind.

Well, I guess maybe another way to frame the question is, how do you know where the real learning or algorithm? Where has the real learning happened? Where is the real algorithm implemented? Another kind of inspiration for this question is going back to some of the Neil Nanda work on the Othello board state type thing. I recall that there is at least the possibility of confusion where you want to know what is being represented inside the model, But if you train a big detector model on the internal states and try to predict board state from the internal state of the model, at some point you're like, well, maybe my new network has learned to do that, but that doesn't really mean that my original network had a sort of semantic understanding of this concept. maybe that kind of only came online in the interpreter model that I trained.

Yeah.

And so I guess I'm worried about a disconnect between the original model. We've split it into these many copies. Those individual layers, those vertical layers, those subcomponents are evolving through this training process as they learn to be simple and reproduce and be only a few of them activated at a time.

Yeah.

But I guess I'm just wondering, like, what is tethering those resulting subcomponents to the original? Because I could imagine that I can change, change, change, change, change in a functional way that gets me the desired behavior.

Yep.

And, you know, I can offset that in some other one of these subcomponents, which maybe never gets used. And then I could potentially, if I do that long enough, hard enough, I could arrive at a spot where potentially the algorithm that the new thing implements is quite distinct from the original, and I'm not sure how I would know that, or if there's something that I'm not catching that is preventing that from happening conceptually.

So for attribution-based parameter decomposition, that is, I think, a fairly reasonable concern. I don't know if it happens in practice. I don't know if the other losses in some sense implicitly penalise it. It will Whenever we move on to the stochastic parameter decomposition, I think it will become a lot more obvious why this can't be the case. We can talk about that now, or I'm happy to park it until we get there.

Yeah, well, let's do it. I think we're pretty much there. I mean, so we can add this to the list of possible problems and whether or not it's an actual problem with the original method. But let's just run through maybe, what made this not the last word, or why you needed an upgrade. You mentioned that it's not super easy to get working. There's a lot of hyperparameter dependency. It's not super easy or stable to train. So that's one. There's computational expense, which maybe you can unpack a little bit more. There's this conceptual drift possibility that I'm flagging.

Yeah.

And then I also wanted to hear a little bit about reliance on attribution methods and why that's a problem as well. So maybe you can just run through those and then we'll move on to the new hotness.

Yeah, this algorithm was extremely janky. It was just very difficult to get it to do what we felt was a sensible thing to do. Now, a later algorithm has much less of this issue. But for instance, you mentioned there was a bunch of different hyperparameters that were necessary to tune in order to get a decomposition that made sense to us. And there's just— one of the reasons here was that, well, we're using a top-k parameter where you're only allowed to activate the top-k most attributed parameter components for any for any given input, and you want to train these attributed ones to be able to do better at the task on which they were attributed. And so you're basically training them to reconstruct their output. And so yeah, this top-k parameter is one of the very fiddly things to get right here. It's very It's very discontinuous, basically. You can— because you can very quickly go from one component being active to, you slightly update the parameters ever so slightly such that another one comes online instead of that one. Well, now you've, for a very small nudge in your parameter space, you've got a very different output. Now you're implementing a very different function for a very small change in parameter space. So that's just— inherently hard to optimize with gradient descent. Another issue is just the sheer number of different hyperparameters. But yeah, I think the other main one is these, you're really relying on your attributions to be right. And if they're in some sense wrong or in some sense systematically biased, you might end up— there may not even be a kind of stable optimum in your training landscape. Because if your biases, suppose you were at a good point in your training landscape, but your biases were slightly off, they might nudge you somewhere else. And then you might move over there and your biases are systematically biased in some other— sorry, your attributions are systematically biased in some other way. And so it might nudge you somewhere else. Sorry, the attributions that we were using were gradient-based, which there's some intuitive sense in which this is kind of capturing what we wanted. What we wanted was some number that told us this was how important this parameter component was on this input. And gradients are a kind of proxy for this because Well, if your output changes a lot for a small change in this parameter component, that is to say, if the gradient of the output with respect to this parameter component is large, well, you can say that it's an important parameter component. But there's some cases where this is just straightforwardly not the case. One example might be attention, where if you've got a really— suppose you have a transformer and it's attending very strongly to one particular sequence, index. And attention there is almost 1 and everywhere else 0. Well, if you nudge the parameter component that is responsible for this very strong attention to this sequence index, you nudge it by some small amount, well, the attention's basically saturated here. And so the gradient is actually very small. It's actually not going to move your attention by very much, despite this being a very mechanistically important parameter component because it's implementing such strong attention. And so this is just like one of the ways in which gradients may not be a good attribution method. We actually just want some other method that tells us more accurately how important was this parameter component. And This was one of the developments that we introduced in the follow-up paper on stochastic parameter decomposition.

I think that's a perfect transition. So that brings us to stochastic parameter decomposition. My general read of the paper is it's like very much a natural follow-on, right? In fact, I actually read— I went through the axrp podcast and the first paper and then wrote the first half of this outline. And then I read the introduction to the new paper and I was like, basically the first half of this outline sort of mirrors the introduction to the new paper. So I felt good about that. So it's basically same goal, same concepts, same— not exactly the same loss functions, but sort of same conceptual constraints, just operationalised differently this time and addressing the weaknesses that you just ran through in the first case with better results. So maybe take us through it. I do have some questions on kind of the intuitions for some of the decisions, but take us through the set, what's new and improved first, and then we can kind of unpack exactly why you did it.

Yep. Yep. So one of the new things is in the headline, stochastic parameter decomposition. And this stochasticness relates to the replacement for the attributions in the previous paper. Now, it's still reasonable to think about, you know, we still need in some sense these attributions, but we will, you know, just to avoid confusion with the previous method, we'll call them something else. We'll call them causal importances. If something is causally important for the network's algorithm, it should be attributed. And we just want something that will let us approximate how causally important this subcomponent is. And before we get into the details on the causal importance calculations, well, it's also important to note that one of the other major differences with the previous method was that we are now not using this frankly ridiculous size for each of the parameter components. They are not just a randomly initialized copy of the original model. They are in fact just a randomly initialized rank 1 matrix for each of these layers. And we're basically like, the way we think about this is that if in the previous method we had these parameter components, that spanned all layers and spanned all potential ranks in these layers? Well, we are basically just pre-splitting it up. We are using— we're saying, well, pretend everything was just localized in rank 1 matrices in each— in one layer. And later on, the idea will be that we group these things together back into full parameter components such that if we did indeed have like rank 2 parameter components or parameter components that span multiple layers, we could indeed find them, but post hoc after we'd grouped these things together. So those were 2 of the major differences between the APD, attribution-based parameter decomposition, and SPD, stochastic parameter decomposition. And so it's probably now reasonable to get into the the causal importance calculations, which is where the stochasticness comes in.

Can we take just one more second on how you're breaking the original network down into these rank 1 components? So again, in the original, you've got a network, it has a certain architecture, it has a certain parameter structure, and you make n of those with— is it really like random initialization on all those copies?

Yeah, yeah, yeah.

Interesting.

And it's similar in the subcomponent that we're now using in stochastic parameter decomposition. These are called subcomponents rather than components because we will group them together later on. But yeah, these subcomponents and parameter components are, they're all randomly initialized, and it's this faithfulness loss that makes them all sum to the parameters of the original model. And that faithfulness loss is pretty strict, so they don't really get to deviate very much from the parameters of the original model as you sum them.

But they do at first. They just converge quickly to sum back to the—

yeah, interesting.

I guess I would have naively thought in the first case, I would have thought more along the lines of, like, and I don't know why this would wouldn't work, you could probably tell me. Why not, if you're going to divide the thing into 1,000 vertical layers, just take all the parameter values and divide them by 1,000 and then have sort of 1,000 weak copies of the thing, some of which can then be turned up, down, whatever? Why randomly initialise instead of doing something more principled or sort of a simple transformation of the— train network itself?

Yeah, I think we explored a bunch of different initializations, and I actually don't recall the exact one that we used in the APD paper. This was the example that you gave of dividing it down and making just exact copies is, I think, something that we explored, and we may even have used a randomized version of that. But you can't really do this in the stochastic parameter decomposition case. Why? Because, well, you have, say, a— I don't know, a— I guess one of the ways you might consider doing this is having a, say, a 1,000 by 1,000 matrix and dividing it up into, say, 10,000 subcomponents. Now, I don't know if there's a guarantee that if you sum up all but one of these parameter components, that the remainder is also rank 1. Does that make any sense? So there's like, you assume that your 10,000 subcomponents sum to the parameters of the Original model. And the idea would be, well, we might as well just constrain this sum to be exactly the parameters of the original model. And so what we'll do is we'll just take the sum of the first 9,999, and then for the remainder, we will just let that be whatever exact makes exact this sum to be the parameters of the original model. But I think then this final one may be of arbitrary rank. That feels intuitive to me, but I'm not sure if it's 100% true. But yeah, it feels like that one may be of arbitrary rank. So if we want them all to be rank 1, then we can kind of just randomly initialize. There is a slightly more You know, involved randomization— sorry, slightly more involved initialization such that they do point in similar directions to the original model, such that their cosine similarity with the original model is not, you know, it's at least not negative at initialization, despite being rank 1. But yeah, it would— the details of that are probably not, you know, super important. It should just work with random initialization. Random initialization. Yeah.

So how should we think about these? The original setup is pretty intuitive to me to sort of envision, okay, I've got this thing, I split it into all these vertical layers. Each one gradually over time becomes simple. Not that many are active. And so I sort of can visualise, instead of this sort of dense and uninterpretable computation, I now have like n of 100n or n of 1,000n that are actually active, and those are each relatively simple and they create the same output. I can sort of visualise that and I'm like, yeah, okay, so there's sort of these, and they feel like circuits. Each of those things feels like a circuit. And it's kind of, okay, if I activate these 8 circuits for this input, then it works. And these other, this one takes 12 for a different input and it works. But each one sort of intuitively feels like it's doing some sort of information processing through the original architecture. And so that's kind of intuitive here. I'm a little less able to tell that story because I'm like, at each layer, we are breaking up whatever matrix exists at that layer into a bunch of rank 1 parts.

Right.

How should I think about those rank 1 parts? I mean, I had to go refamiliarize myself a little bit with what exactly is rank and what does it mean. And in the simplest terms, it's like just a much simpler, smaller part of the overall transformation. Specifically referring to dimensionality. You have some n-dimensional space. Something that is rank 1 basically exists as a line in that space, right? And something that is rank 2 exists as a plane in that space and so on.

Yeah.

So we've got now a bunch of, like, instead of a transformation that can operate in n dimensions all at once, we've got a bunch of transformations that all operate on one linear direction in that broader space. And now we're composing those sequentially, I guess, in a way that recreates the original. But how should I be thinking about those rank 1 things? Like, why do that, I guess?

Yeah, I guess, like, I don't really view them as very fundamentally different. So suppose, you know, consider the case in the old algorithm where we have this parameter component, and this parameter component can span all layers and all potential ranks. But one of the penalties that we were emphasizing was this, like, simplicity penalty, where we don't want it to involve too much computational machinery. We want to be able to study it we want to be able to study as simple objects as possible. And the way we operationalized this was that we penalized indeed the rank of the matrices in this parameter component, and also penalized it for existing in multiple layers. And so what this, you know, the idea here being that if this parameter component really was just localized in a few layers and in a few ranks, the rest should become The rest of its parameters should become 0 in all these different layers, and there should just be a low-rank set of matrices in the layers in which it is involved. And so we are hoping that— and just to maybe go into a little bit more detail about why that is simpler, if it's rank 1, as you mentioned, this is just a rank 1 matrix, basically, looks for and it reads and writes in a particular, in one single direction in activation space. And the only direction in which it reads is the direction defined on its, the right-hand singular vector, and the direction to which it writes is the left singular vector. You know, this matrix will not be used if the activations don't, in some sense, overlap with that read direction. And if the write direction just has no relevance downstream, it has no causal importance downstream, it doesn't do anything, then this rank 1 matrix won't do anything. So given that we're already looking for these rank 1, these low rank like matrices in these full parameter components. The idea of stochastic parameter decomposition is that we can— well, one of the ideas we introduced in stochastic parameter decomposition, it's not a fundamental part of the algorithm, but the idea that we introduced here was that we'll just start with these low-rank versions of the parameter— sorry, of the subcomponents, And then later on, we can aggregate them such that if we want to combine 2 in order to create a rank 2 parameter component, then we can do that. And if we wanted to have it exist in multiple layers, we can aggregate them together as well. And so they're really just chunked versions of the same fundamental object, this parameter component, this vector in parameter space. One is just something that we'll group together later on.

Okay, well, let's come back to how the grouping is done, because I think that's another quite vexing, potentially quite deep question. But for the moment, okay, so we've broken the matrix into— if it's a rank n or an n-dimensional space, we've now got not just n, but actually could be more than n rank 1? So how about a little help on that intuition as well?

Yeah.

So yeah, it's one of the core ideas in the representations in superposition line of work is that there are these features in activation space that, and there's more, that each individually represent a single feature, a single thing in the world, but there are more of them than there are neurons that represent them. And the same idea holds true here. In parameter space, we can represent more directions in activation space by using just more rank 1 matrices. matrices. And this gets back to the read and write direction. So if we want to read from more directions than we have neurons, well, we just have more ranks than we have dimensions in our activation space, more ranks than we have columns in the matrix. And so we're like, it's possible for this to read directions, read these representations in superposition, if we have more of these rank-1 matrices in our individual matrix. The idea then is that these rank-1 matrices are what will implement individual or parts of individual computations that the network is doing to compute its behavior. But only whenever— the challenge is basically to identify which ones are important for a given input data point and when it's not important. And that's where the stochastic calculation of the causal importance comes in, which I can get into if—

Yeah. Maybe one more question before we go to the causal importance. We've broken these things into all these And there's an interesting question too, of what is the ratio of how many of these rank 1 subcomponents do you need for different size and complexity of original network? In the relatively simple problems studied so far, that ratio isn't super high. I'm kind of wondering, is that ratio going to become extremely high if we get into the language modelling space? It seems like Probably yes.

Potentially. We will probably have more computations than we have neurons. That's just the name of the game in taking things out of superposition. But one of the questions that I ask myself often is, well, maybe, is it going to be more or less efficient than, say, the dominant approaches At the moment. And these dominant approaches may be, say, sparse autoencoders, transcoders, cross-layer transcoders, and so on. And the basic question here is, well, is it going to be more or less efficient? And it's unclear, but I think there's reason to believe that if indeed this is a reasonable way to decompose networks, there's probably going to be like fewer of these subcomponents than there are, say, latents in a sparse autoencoder or a transcoder or so on. The reason being that there's no upper bound to how many latents you can have in your sparse autoencoder or your transcoder, and whereas there is a kind of an upper bound on the number of parameter components that you might have Why? Because, well, they have to all sum to the parameters of the original model. And this prevents you from just adding in an extra component or an extra latent and just cranking up your reconstruction loss. Basically, there is kind of like a— it should level off a bit faster, I would say. But this is an empirical question that— remains to be answered for the kind of models that we care about, like large language models or other similarly interesting models.

Is that sort of like saying that while reality is complicated, the rules of physics are simple? Like the sort of the number of fundamental transformations, the number of sort of fundamental functions that are used to carve up reality should, in theory, just be a much reduced space that can then operate on many more inputs. Maybe thinking like Platonic representation hypothesis type notions.

It feels like it's getting at something true, although, yeah, It's not obvious to me right now how to connect it to the way in which I think of things. But yeah, it feels like it's getting at something true.

Yeah. Well, more work to be done for sure. Is the— and how expressive can these rank 1 things be? Like, they can represent, for example, a rotation. So is that right?

I think they can represent a rotation of one direction to another. But I'm hesitating because I don't— I'm trying to recall the technical definition of rotation, but I think it is true, yes, that it basically can do a rotation and also a scaling. The rotation is a very simple kind of rotation. It is just like, to the extent that an activation or an input activation projects into one direction, it will now project into this other direction. you can get arbitrarily high-dimensional rotations up to the dimensionality of the input. But yeah, so I think it is true to say that it can do a rotation, but it's a very simple, low-dimensional kind of rotation.

But that would be enough, for example, to handle the transformation of like Tuesday into Wednesday with a sort of next-day operator.

I think we basically need 2 For this particular transformation, because these variables live on a plane. And so you basically would need to project 2 directions in a particular way, whereas if it were— yeah, you're basically leveraging 2 such 1-dimensional rotations in order to do this 2-dimensional. Okay.

So this gets us back also to the grouping concept being quite important. So let's come back to that again in a second. I think I've belaboured the rank 1 thing enough. We've broken this thing up into some ratio greater rank 1 subcomponents from the original n-dimensional matrix. Now tell us about this swap out of the attribution method to the causal importance learned, the learned estimation of the causal importance. Yeah.

So one of the ways you might think about what the attribution method was doing was that we wanted something that said, well, if we turn this parameter component off, it shouldn't really affect the output. And so the way we did that was just we literally did turn them off and only kept on this top k parameter components. And what this is really saying is that it shouldn't really matter if this parameter component is on or off. It could be on to its full extent, or it could be off to its full extent, or anywhere in between. And it really should not affect the output. The top-k approach was just not really great at optimizing for this such that if you did turn it on a little amount, it probably still would do stuff. Whereas what we're basically doing in stochastic parameter decomposition is we are in fact learning a function that predicts for a given input, If, like, how ablatable is this? How turn-offable is this subcomponent? And if it is turn-offable, then again, it shouldn't matter how much we turn it off by. It is just causally irrelevant. And so we have this causal importance function that predicts this one number, how causally, you know, we call it the causal importance value, for this subcomponent on this data point. And it is just a number between— supposed to be between 0 and 1. And to the extent that it is 0, we can basically randomly mask this subcomponent on this input anywhere between 0 and 1. That is to say, we can turn it on randomly anywhere between 0 and 1. And if this causal importance value is 1, well, then we don't really get to randomly modify it, randomly mask it at all. It just has to be on because it's really, really causally important. And so the way— yeah, we literally just have a neural network that spits out this one number for each subcomponent. And if the subcomponent is very causally important and it gets the right causal importance value of 1, well then great. But if it is causally important and accidentally the network spits out, well, 0.5. Well, now this parameter subcomponent can be masked anywhere between 1 and 0.5, and if it's masked close to anywhere that's not 1, so anywhere along this random distribution, well, it's going to damage the loss. And so what the— this means that gradients will be able to flow into this Causal importance function that produced this number using, yeah, basically the reparameterization trick to let us basically pass gradients through this random sample and basically output this random distribution such that it actually does learn to approximate how ablatable this parameter component is for a given input. And I think this was not my idea, but I think it's a really great idea that I actually haven't seen elsewhere in the attribution literature. And I think it has applications far beyond just the approach where the interpretability approach we're having here. It doesn't have to be applied to parameter components. It could indeed be applied to parts of an input image and in the kind of way that attribution methods have been done in the past. We haven't done that ourselves. We'd be keen to see someone do it. But yeah, it's just an attribution method that actually gets to the— really captures just how causally important this component was for this input.

And so is that First of all, that's a pretty big matrix, right, that's been added for that, or network that's been added for that learning, because it's got to go from the input space to the output space of all the subcomponents, which are all the rank 1 breakouts at all the layers.

This causal importance function can basically be of arbitrary architecture. It can indeed, as you suggested, be a map from the input to this, the number of components per layer times the number of layers. That would be a very large network. And this is not in fact how we implemented it in the paper. In the paper, we just had a very, very small 16-parameter, I think it was 16 or it varied somewhere between say 16 and 128 parameter thresholding network that took as input the— I mean, it gets a little bit mathematical here, but it's not super complicated. It's just like you multiply the activations by the right singular vectors of the parameter components, which is basically you're asking how much are these parameter components reading from the activations? Is this one reading a lot? Is this one reading not very much? Just measuring— you're measuring the overlap between the activations and the right singular vectors of these components. And to the extent that the parameter subcomponent has a lot to read there, to the extent that it's active, this is like a number that you can basically do a fancy thresholding with a small neural network. And that's just 16 times however many parameter components you have, which in the scale of the overall network is not very large. And it doesn't really even need to be a very fancy neural network at all. It can just be some sort of simple learned threshold. We found in practice it helped it learn nice kind of thresholding. functions that were not completely straightforward, and they had some kind of structure to them. But overall, it was just like a slightly fancy threshold.

That's interesting. And then when the whole thing is run, when it's trained and when it's run, that's now part of the reconstruction, right? It is, because again, all this is jointly optimized, this whole prediction of which component is, or which subcomponents are going to be active and the structure of the subcomponents themselves, that's all being jointly optimised in one big compound loss function. And then when you're actually running the replacement network for study on different inputs, that is still part of it, right?

It can and usually is part of it, but you can throw away this causal importance network and just sum up the parameters of the— sum up all these subcomponents together and just run it like the original model. And if training has gone well, well, this should look very much like just running the original model. Now, it's like running everything with a mask of 1 instead of 0. It should do exactly the same thing. And the causal importance function is not necessary. What the causal importance function is letting you do is basically telling you which parts are actually being used on this input. And you can imagine a world where you can throw away— suppose there was a subtask of particular interest, you might be able to throw away most of the network and only have— and only keep the parameter subcomponents that are used on this distribution of interest. And that may indeed, depending on how large the model is, may indeed be a much smaller set of parameters than the whole model. But yeah, that's future work.

Yeah, interesting. Okay. So, Why does this work better? I guess, how does it work better? Tell us about the results of maybe go problem by problem that we had in the last version and talk about the updates. You can get into the maybe toy problem, compare and contrast. And then again, I'm really interested in sort of, do you have any intuition for why this works? It still seems very magical that all this just shakes out.

Yeah. So I think with regard to attributions, this is just much closer to what we actually wanted from gradient attributions. It doesn't have the kind of biases that gradients would have. It should, for instance, be able to tell us, it should be able to cope with the saturated attention case that I mentioned earlier when we were talking about gradient attributions. I think it might also be time to revisit the the garbage component that you mentioned, where, because basically, I think we were talking about a case where maybe the network has learned these subcomponents. And in fact, maybe the subcomponents are actually doing slightly different algorithm from the original network. But because in APD, you could switch them off and never activate them, well, you can kind of compensate for one component that is active doing a different thing while still summing up to the parameters of the original network. In stochastic parameter decomposition, you can't really do that anymore because you're just randomly activating every part of the network basically all the time. And if it is doing something that it shouldn't be doing, That was going to modify the activations in a way that the original network just didn't modify the activations. And so it's going to be bad for the loss. And it's important to realize that there's just— we are not actually able to turn off subcomponents in SPD. There's no dead components or anything. What it means for it to be a dead component is that you, not that you never activate, but that you activate with a Always randomly. And I think this is just fundamentally different from just being completely silenced all the time and gets rid of a lot of pathologies like the one that you mentioned. And it also gets around this idea of the discontinuity that I mentioned in the top-k optimization, where it's no longer the case that you're only activating the top k and the rest are silent. Everything again is active all the time, which means that if you are in APD where you're using top k, well, if you're silenced, well, gradients aren't really flowing through you. You can't really learn to be— your parameters of your parameter components can't really learn to be better. Whereas if you're always on randomly all the time, well, gradients are always flowing through every subcomponent. And so it feels just like there's a lot more continuity in the learning process. And this pans out in terms of, it's not so much, the algorithm is still somewhat, you still need to get hyperparameters somewhat right. However, it's just not like in attribution-based parameter decomposition where Not only was it sensitive, but it was just not well behaved. You changed it one way and it changed in a kind of unpredictable fashion, where at least in stochastic parameter decomposition, fine, there's a range in which you need your hyperparameters to exist. But because things are well behaved, you can actually just find them without having to— it's just easier to understand which thing you need to change in order to get to the right. right solution. So it's just much nicer to work with than the previous algorithm. Now, that said, the method that we introduced in the paper is just not perfect. There's plenty of things that even since publishing that paper, we've focused on changing and found better versions for sampling the this randomness and the gates we used in the causal importance functions and so on, just to make it a better behaved algorithm overall. And there's still optimization on algorithms to do, but it is just significantly better to work with than the previous method, which makes us excited for scaling this up to larger models, which we've begun to do for, e.g., like small transformers, including language models.

Okay, let's talk about then feature splitting. You had said that some of the results here show that the new method avoids feature splitting, I guess. And then coming from an SAE perspective, I sort of think feature splitting is good. I think of it as being just like, as I make my SAE bigger and bigger, then I get finer-grained resolution on the features, and I expect to see these sort of vanilla get resolved into French vanilla and normal vanilla or whatever. And that seems good. But tell me, in what sense is it bad? And in what sense, why are we glad? And how has this helped us avoid that problematic version of feature splitting?

Yeah. So we've talked already about one of the examples with the cat and the cat in a number of different positions. And it's not clear whether or not the network really is actually using, in some sense, these more fine-grained features that you can identify. And so this is just, from a practical point of view, we don't want our explanations to explain more than they need to. Now, we can explain potentially how the network takes one of these cats in one position and applies circuits to it, and then we can do it for the next cat in another position and so on. But we're wasting our time if the network is just not really using these variables in its algorithms. And so we basically want to avoid doing unnecessary explanation. We want our descriptions of what the network is doing to be minimal in some sense. So another particularly pathological case of feature splitting is, say, feature absorption. Where you might get— say you have a— in one SAE, a smaller SAE, you have a feature for words that begin with the letter E. But then, you know, in a larger SAE, one that has more fine-grained features, we have a feature for elephant and then a feature for everything else that begins with the letter E. Or it splits into other words that begin with the letter E. And in some sense, this other decomposition, the one with more features, is just less interpretable because, well, there was no need for us to split out elephant maybe. It was just, maybe we want— we just, again, just want the variables on which the network is working. And this is just— there's no obvious way to avoid this using, say, SAEs. Now, it may be the case that there's a perfect stopping point for SAEs, such that beyond this point, you're using variables that the network is using, and before that, you're not. Sorry. Yeah, yeah. There may be a point at which you go into a regime where you've got SAE features that correspond to variables that the network is not using, and before that you've got only variables that the network is using. But that's not obvious either, because, well, fine, you could find these sparsely activating directions in activation space, but there's no real reason to expect— well, no guaranteed reason to expect that these are indeed the way in which the network is chunking up its computations. Maybe maybe it is in fact a group of these SAE features over here that the network is actually using, and maybe you actually did need to split further for this group over here. There's no guaranteed reason that SAE features will find variables for computations. Now this is not necessarily a criticism that applies to other various other variants of the sparse dictionary learning paradigm. And I do want to caveat everything we've been talking about here in that there are people whose views I respect a lot that see issues with our approach. We see issues with theirs. But overall, it— I don't want to give the impression that the algorithm that we're putting forward here is the final word on anything in interpretability or that it's definitely finding the right thing. I think there's a lot more research to be done on this. And yeah, I just wanted to make clear that there's still more to be done.

So then how about the Actual grouping, labelling, and sort of semantic understanding of all these rank 1 units. There was one bit, and this is one of the toy problems. We don't have time to do the full setup, but there was this sentence that SPD splits its corresponding W out into 50 subcomponents. In other words, this one matrix that's been broken down into into rank 1. It has 50 rank 1 things that all remain active. And then the paper says that these 50 subcomponents appear to be effectively part of a single rank 50 component comprising the entire Wout matrix. And there I'm like, how do I know? What does that mean? And how do I know it's right? Like, I'm And doesn't it leave us kind of with a similar problem to where we started, which is like, what's going on there? If we've got 50 things and we group them back into one thing, how did you know to do that versus determining that there are 50 different things happening? That part seemed to go by very quickly, and I wasn't quite clear on it, let alone how to extrapolate it into larger and larger models in the future?

Yep. Yep. No, it's a really great question. So I think what's important to appreciate is that we are kind of, whenever we're optimizing these things and whenever we're saying that this is the right decomposition, there's an asterisk here. We're saying this is, for a given amount of decomposition, this is the the decomposition that breaks it up into as few components as is necessary to explain a given level of functionality of the network. And so we can basically scale up and down the amount of decomposition that we want to do and break up the network into larger or smaller parts. There is some reason to prefer some of these larger— there are some reasons to prefer some regions of this spectrum over others. One of these reasons might be, for instance, that the reconstruction loss is actually, as you scale things and decompose them more and more and more, There comes a point where you actually just, your reconstruction loss kind of gets a lot worse. And correspondingly, there may be a point where your causal importances, these values that are supposed to tell us how important each subcomponent are, they kind of behave a bit pathologically. For instance, they might all typically take fractional values rather than values that should be like 1 or 0, because Ultimately, if something's causally important, fine, it can kind of be partly causally important, which will let us— which will shed light on how much loss we might lose by ablating it. But by and large, we kind of prefer regimes where it's a bit more binary. And there's points along this spectrum as we modify these hyperparameters that control how much decomposition or activation sparsity is going on, that give us clues as to what decompositions make sense. And in some sense, none of these decompositions are right. This holds true for sparse dictionary learning, for parameter decomposition. None of these decompositions are right in some sense, because Well, we're just dealing with one fundamental truth that is the original network. And we're just really trying to get a lens onto what it is doing under the hood. But we're always going to— none of our lenses are going to be super perfect. They're always going to introduce approximations of some sorts. This is just a necessity. We're all trying to discretize a fundamentally continuous object in various different ways. And whenever you do that, You're just simply going to throw out some non-zero amount of approximation quality. But we are willing to pay these costs to get tractable explanations, even though there is not one actual ground truth true decomposition. Yeah.

Okay. Maybe last question. maybe bundle a couple last questions into one, and you can take as long as you have. Where do we go from here? How expensive is this going to be to scale up to the, let's say, small but meaningful language model regime, your sort of 3Bs or your 7Bs? Is that the next leap that you're going to try to make? And what sort of practical applications do you find to be most exciting? Obviously, again, SAEs can do the sort of monitoring and steering. Are there fundamentally new use cases that you see opening up here? And maybe if you have time for it, is there an intersection between this and the science models? I know Goodfire has been doing some interp work on science models, and that's one of the things that excites me most, in addition to just keeping the general purpose models on the rails, the idea that we could make scientific discoveries by understanding what the models have learned is pretty exciting. So, yeah, what's the next chapter? What's it cost scale-up-wise? What are the new capabilities that that unlocks? And any comments on science? And then we'll let you go. Cool.

So yeah, the next steps are definitely to scale these up to more respectably sized models, eventually getting— I think we've tried out in the region of low-digit billions. We've scaled this down to even smaller language models in the millions just to get some early traction before scaling back up to those sizes. There's a lot of interesting empirical questions that we want to answer there. Like, how much do we actually get to decompose these language models into low or high? Maybe the things that they decompose into are too high rank for us to understand, which would be— we would basically need to do extra work after this first step of decomposition, but we'd at least get a decomposition into a bunch of components that are themselves easier to isolate and understand. Then, yeah, I think some of the applications that I think this opens up, it could open up, are I think it should make, say, unlearning more straightforward than, say, in SAEs, because you're already working in parameter space, you already know this is the direction that does x, we can just modify this vector in parameter space and we have a more straightforward lens onto functions as they relate to parameter space. Now, whether or not this beats some gradient descent guided approach is a separate question, But I think monitoring is also something that I think it will be useful for to the same or a greater extent than, say, sparse autoencoders or other approaches would be. Why? Because, well, sparse autoencoders might be detecting a particular region of the training dataset. It's a too narrow region perhaps. And instead, what we really wanted was like a Whenever the network was using function X, that's when we want to pay a bit more attention and go into detail on what it's doing, rather than was there an activation in this particular direction? There's going to be some correlation between these 2, but there will be just some things we want to monitor for that might be better thought of in terms of functions. To get less abstract, maybe there are maybe there's a function for deception and there's many different situations in which a model might be doing deception. The input directions might be different. And so the intermediate computations might point in different directions in a way that SAEs might struggle to identify as this is the one direction that does deception. It may instead have many different directions that do this, but in sparse parameter— sorry, stochastic parameter decomposition and other, we hope, later versions of this approach might be able to identify the individual components that are doing this functionality. And then lastly, yeah, we really want to extract knowledge from these models, some of which are superhuman at some tasks that we've trained them to do. And we want to be able to actually think in the terms that the network itself is thinking in, because that is the knowledge in some sense that the network is using. And so we hope that this may be a way to get a lens onto that in a way that may or may not generalize better than other approaches. But yeah, that remains to be seen. Cool.

Well, we'll definitely be following your work here. It's fascinating stuff, and I'm glad to have you on, you know, while we're still in the relatively early phases of it. I think the, you know, this mega project of figuring out why these things do what they do, and, you know, if there's any way to really make ourselves confident that we're going to be able to predict what they're going to do in novel circumstances, is about as high a priority as any in my mind. So I appreciate that you're—

Thank you.

working hard and have taken at least one good bite out of that problem. And I'm really looking forward to hopefully seeing this line of work mature into something that can tackle the biggest models and the hardest problems. It'd be great to ablate that deception circuit if we can find it. So any other thoughts you want to leave people with before we break?

I think that's honestly it, Nathan. This has been a really great conversation. Yeah, we had a lot of fun. And yeah, just super pleased to be on.

Lee Sharkey, principal investigator at mechanistic interpretability startup Goodfire. Thank you for being part of the Cognitive Revolution. If you're finding value in the show, we'd appreciate it if you'd take a moment to share it with friends, post online, write a review on Apple Podcasts or Spotify, or just leave us a comment on YouTube.

Of course, we always welcome your feedback, guest and topic suggestions, and sponsorship inquiries. Either via our website, cognitiverevolution.ai, or by DMing me on your favorite social network. The Cognitive Revolution is part of the Turpentine Network, a network of podcasts where experts talk technology, business, economics, geopolitics, culture, and more, which is now a part of a16z. We're produced by AI Podcasting. If you're looking for podcast production help for everything from the moment you stop recording to the moment your audience starts listening, Check them out and see my endorsement at aipodcast.ing. And finally, I encourage you to take a moment to check out our new and improved show notes, which were created automatically by Notion's AI Meeting Notes. AI Meeting Notes captures every detail and breaks down complex concepts so no idea gets lost. And because AI Meeting Notes lives right in Notion, everything you capture, whether that's meetings, podcasts, interviews, or conversations, Lives exactly where you plan, build, and get things done.

No switching, no slowdown.

Check out Notion's AI meeting notes if you want perfect notes that write themselves, and head to the link in our show notes to try Notion's AI meeting notes free for 30 days.