E39: Seeing is Believing with MIT’s Ziming Liu

Ziming Liu is a Physics PhD student at MIT and IAIFI, advised by Prof. Max Tegmark. Ziming’s research is at the intersection of AI and physics. Today’s discussion goes in-depth on Liu’s paper “Seeing is Believing” where he presents Brain-Inspired Modular Training (BIMT), a method for making neural networks more modular and interpretable. The ability to see modules visually can complement current mechanistic interpretability strategies. We're hiring across the board at Turpentine and for Erik's personal team on other projects he's incubating. He's hiring a Chief of Staff, EA, Head of Special Projects, Investment Associate, and more. For a list of JDs, check out: eriktorenberg.com. RECOMMENDED PODCAST: The HR industry is at a crossroads. What will it take to construct the next generation of incredible businesses – and where can people leaders have the most business impact? Hosts Nolan Church and Kelli Dragovich have been through it all, the highs and the lows – IPOs, layoffs, executive turnover, board meetings, culture changes, and more. With a lineup of industry vets and experts, Nolan and Kelli break down the nitty-gritty details, trade offs, and dynamics of constructing high performing companies. Through unfiltered conversations that can only happen between seasoned practitioners, Kelli and Nolan dive deep into the kind of leadership-level strategy that often happens behind closed doors. Check out the first episode with the architect of Netflix’s culture deck Patty McCord. https://link.chtbl.com/hrheretics The Cognitive Revolution is a part of the Turpentine podcast network. To learn more: Turpentine.co Watch this episode on YouTube to see research figures and graphics: https://www.youtube.com/watch?v=VKxwlxb2XWw Links: Seeing is Believing, authors: Ziming Liu, Eric Gan, and Max Tegmark https://arxiv.org/pdf/2305.08746.pdf Ziming Liu: https://kindxiaoming.github.io/ TIMESTAMPS: (00:00) Episode preview (01:16) Nathan introduces Ziming Liu and his research (04:50) What motivated Ziming's research into neural networks and modularity (08:06) Key Visuals of "Seeing is Believing" (08:44) Biology analogy & paper overview (15:07) Sponsor: Omneky (26:08) Research precedents (30:00) Locality vs modularity (37:40) Shoutout Tiny Stories : Episode 33 (39:00) Quantization Model of Neuro Scaling. (1:00:00) How can this work impact mechanistic interpretability strategies (1:27:00) Training (1:46:30) Downscaling large-scale models TWITTER: @labenz (Nathan) @ZimingLiu11 (Ziming) @cogrev_podcast SPONSORS: Shopify is the global commerce platform that helps you sell at every stage of your business. Shopify powers 10% of ALL eCommerce in the US. And Shopify's the global force behind Allbirds, Rothy's, and Brooklinen, and 1,000,000s of other entrepreneurs across 175 countries. From their all-in-one e-commerce platform, to their in-person POS system – wherever and whatever you're selling, Shopify's got you covered. With free Shopify Magic, sell more with less effort by whipping up captivating content that converts – from blog posts to product descriptions using AI. Sign up for $1/month trial period: https://shopify.com/cognitive Thank you Omneky (www.omneky.com) for sponsoring The Cognitive Revolution. Omneky is an omnichannel creative generation platform that lets you launch hundreds of thousands of ad iterations that actually work, customized across all platforms, with a click of a button. Omneky combines generative AI and real-time advertising data. Mention "Cog Rev" for 10% off. This show is produced by Turpentine: a network of podcasts, newsletters, and more, covering technology, business, and culture — all from the perspective of industry insiders and experts. We’re launching new shows every week, and we’re looking for industry-leading sponsors — if you think that might be you and your company, email us at [email protected]. MUSIC CREDIT: MusicLM

Speaker names marked ? are inferred from speaking time and the show’s known hosts — they are not voice-verified.

• 0:00Speaker 00

  If we want neural networks to be interpretable, we should in training explicitly encourage them to be interpretable. That's when we posed this research question, like what kind of training techniques we need to induce modularity in otherwise non-modular neural networks? Actually, all the results I got, I'm just using my Mac M1. There are no GPU. For the module addition, it takes around 20 minutes. For symbolic formulas, just less than one minute. The extinction risk is a common threat to all human beings, and no one can survive if we don't collaborate.

• 0:42Speaker 00

  So with this, actually, this can bring together US and China collaboration on this AI safety thing.

• 0:50Nathan Labenz?

  Hello and welcome to The Cognitive Revolution, where we interview visionary researchers, entrepreneurs, and builders working on the frontier of artificial intelligence. Each week, we'll explore their revolutionary ideas, and together, we'll build a picture of how AI technology will transform work, life, and society in the coming years. I'm Nathan Labenz, joined by my co-host, Eric Torenberg. Hello and welcome back to The Cognitive Revolution. Today, my guest is Ziming Lu, grad student in physicist-turned-AI-safety researcher at Max Tegmark's Group at MIT, and first author of the recent paper, "Seeing is Believing: Brain-Inspired Modular Training for Mechanistic Interpretability."

• 1:33Nathan Labenz?

  This paper immediately stood out to me for three big reasons. First, any work that makes deep learning neural networks easier to reverse engineer and interpret is, in my view, worth celebrating. The infamous black box problem, which in its simplest form simply means that we don't really know why AIs do what they do, is one of the biggest reasons to worry about what more powerful AI systems might do in the future. And personally, I see mechanistic interpretability work as one of the most promising paths to AI safety. Second, the technique in this paper is conceptually simple, but nevertheless profound. Taking inspiration from biology and the practical physical constraints that favor local connections

• 2:17Nathan Labenz?

  and modular structures within animal brains, Ziming devised a straightforward modification to the loss function that encourages the development of similarly modular structures in digital neural networks. I suspect that this technique, perhaps more than any other I've seen so far in 2023, left many in the AI research community asking themselves, "Why didn't I think of that?"

• 2:40Nathan Labenz?

  And finally, as the title "Seeing is Believing" suggests, this paper includes phenomenal visual aids, which very quickly and intuitively communicate both the core ideas and the results, while also inviting deeper study. Given how much I understood at a glance, I was really pleased with how much more I learned by digging in and exploring some of the important design decisions and trade-offs that Xie Ming made along the way. Because this work is so visual, we've taken the extra step of showing the figures in the YouTube version. So if it's convenient, I would encourage you to watch this one on YouTube. If not, you can definitely still get the bulk of the benefit by following the links in the show notes to the GitHub page, which has all of the important animations, to Max Tegmark's announcement tweet, and of course, to the paper itself.

• 3:25Nathan Labenz?

  Just a couple minutes spent looking at the animations and the images will do wonders for your understanding of this discussion.

• 3:33Nathan Labenz?

  As always, if you're finding value in the show, we'd appreciate it if you'd take a moment to share it with friends or leave us a review. We'd also welcome your guest suggestions, either by email at [email protected] or via Twitter DM, where I am @labenz. Now, I hope you enjoy this illuminating conversation about neural networks made sparse, modular, and interpretable by design with Xie Minglu from MIT. Xie Minglu, welcome to the cognitive revolution.

• 4:01Speaker 02

  Yeah, thank you for the invitation. It's my pleasure.

• 4:04Nathan Labenz?

  Very excited to have you. The pleasure's all mine. You guys have published this really interesting work, "Seeing is Believing: Brain-Inspired Modular Training for Mechanistic Interpretability." And you had me at mechanistic interpretability, but...

• 4:20Speaker 02

  Especially when I saw the visual that you guys posted on Twitter. It's so rare to see in AI, you know, mechanistic work that you can kind of absorb and really start to get an intuition for in just a few seconds. Immediately, I was like, I need to read this paper in its entirety. And obviously now here we are to dive deep into it. So great job, very exciting work. You started off with this question. What training techniques can induce modularity in otherwise non-modular networks? Maybe we could just start there and you can tell us about kind of the motivation and how you came into this line of research.

• 4:56Speaker 00

  We all know that deep neural networks are successful, but hard to understand. So to interpret neural networks, we kind of hope that neural networks can be decomposed into different modules. And like we divide and conquer. Once we decompose that we divide networks into smaller parts like modules, then they are more amenable to interpretability. But we kind of rely on the analogy to human brains where in brains, different parts of our brains are responsible for different functions. But what's different for biological versus artificial neural networks is that artificial neural networks do not have incentives to become modular.

• 5:41Speaker 00

  By contrast, biological neural networks have evolutionary reasons to become modular. Because modular brains are more energy efficient than non-modular brains. That's why modular brains have a selective advantage in evolution. And we human beings today have modular brains. So the philosophy is if we want to make neural networks, artificial neural networks modular, we should in training explicitly make them modular. Or to even generalize beyond modularity, if we want neural networks to be interpretable, we should in training explicitly encourage them to be interpretable.

• 6:27Speaker 00

  That's when we ask this research question, we pose this research question. Like what kind of training techniques we need to induce modularity in otherwise non-modular neural networks? And also borrowing the lesson from biology, we ask why our human brains are remarkably modular? It's because this is closely tied to locality. Because, because our brains, of course, we all live in a three dimensional space. Where you can define distances. And to create neurons in your brains, you need energy to create neurons to maintain them and also transmit signals.

• 7:10Speaker 00

  So, so local neuron connections tend to be more energy efficient than non-local ones. So the lesson we learned from our human brains is that locality gives rise to modularity. And that is the key trick we used. That's a key strategy we used in the BIM paper where we explicitly embed the network into a geometric space where you can define distances. And we have a penalty proportional to the length of the neural connections. We penalize non-local or long neural connections more than, more than short or local connections.

• 7:52Speaker 00

  Just to mimic these evolutionary effects to some extent. That's the basic idea.

• 7:58Speaker 02

  You know, you definitely want to look at some of the artifacts from this paper. As the title suggests, seeing is believing. So the audio only form here, you know, almost can't possibly do this work full justice. But the good news is you can probably spend five minutes looking at the key figures. And then, you know, you'll, they're simple enough and kind of striking enough that you can absorb that information and kind of keep that, you know, that image in mind. And then go, you know, listen to the rest in, in audio form, if that's, if that's what you want to do. So definitely do that. But okay, so just to restate kind of key ideas. So obviously we, you know, as people and as biological, you know, beings exist in space.

• 8:41Speaker 02

  We have this three dimensional problem to solve. How do we get kind of the most? Bang for our buck out of a certain sized skull, right? And that naturally means like it's gonna be, there's gonna be a cost imposed on having super long range connections across the brain. If nothing else, it's gonna just like clog up space to do that. But also like the cell has to get bigger. It seems like it might, you know, be more susceptible to like damage or kind of supply chain issues, you know, so to speak. And so, you know, it makes sense intuitively that most of the brain activity would be local. And then we sort of also see that the brain activity is modular

• 9:23Speaker 02

  in that there's these like particular parts which seem to be like fairly clearly, you know, the, the boundaries are fairly clearly defined between them in many cases, where like this part does something and, you know, this other part does something else.

• 9:38Speaker 02

  In contrast, though, neural networks, as they've generally existed, just kind of have this concept of layers, but there's not any sort of spatial meaning that is given to those layers. So whether, you know, two neurons or activations or positions in the, in the, you know, in the network are, you have the same index or are, you know, are very far from each other in a layer. That has never really been something that has been taken into account in the process of training. So whatever kind of connections happen to form at the beginning, those are kind of the ones that continue to be built around and refined. And that's why we get kind of spaghetti looking networks.

• 10:24Speaker 00

  Yes, we would, normally the artificial neural networks we used normally before is, we treat them as topological objects. So sorry for being a little bit mathematical, but for, for topological objects, there's no numeric measure to evaluate how close these two points are. We only care about whether two things are connected, whether two nodes are connected by an edge or not.

• 10:49Speaker 00

  However, if we go to geometric space, for example, the 3D Euclidean space that we live in, we, we can, we have the privilege to define distances. So actually the switching from normal neural networks to BIMT is just switching from a topological space to a geometric space, so to speak.

• 11:11Speaker 02

  Yeah, I thought this was really interesting. Concept in that it's ultimately a pretty simple and kind of elegant one where, you know, you you describe it so simply there, right? It's moving from this kind of topological conception to a more, you know, literal spatial conception. And then if I understand correctly, like the modifications that you're making to the loss function are also pretty simple, right? Like in terms of the code, you're we're probably talking about just like, not even that many characters of additional code to to kind of ultimately define that, you know, there is a link and now it's going to be counted against you.

• 11:54Speaker 00

  Yes, so basically we want to sparsify the network. So there is this L one regularization. So naturally, even if we do just do L one regularization without this locality trick, you would require one hyperparameter, which is the penalty strength. So on top of that, we just have one additional more hyperparameter, which is like how much you want to encourage locality. So it's simply just L one regularization, but it's distance dependent. And the only extra hyperparameter is the locality strength. When you put that locality strength to zero, you recover the Vellena L one regularization.

• 12:37Speaker 00

  When you put A to be large, encourage more locality. That's really simple implementation code.

• 12:44Speaker 02

  L one regularization is, again, motivated by the same thing. We want to make the network sparse. We don't want to have like wasteful, kind of needless extra connections. So let's just penalize, in this case, the weight itself, right? Like, we want to optimize a loss function that includes both like prediction accuracy, but now we're adding on to it. And we also want to minimize the total weights in the network overall. So all of the weights themselves get summed up into that loss function. And therefore, when we take the Derivative, or you know, we find the gradient, they're all naturally incentivized to go down, all the weights naturally tend down at every step.

• 13:26Speaker 00

  Yes, so the loss landscape would look like, uh, for each weight, it looks like an absolute value for the, for the L1 penalty. So it, it has this incentive to go down to zero magnitude, except that the prediction loss can also drag you somewhere else, but the penalty terms drags you down to zero.

• 13:44Speaker 02

  Right, so only over rounds of gradient descent, only those weights that actually contribute to prediction accuracy enough to offset the, you know, the weight that they, no pun intended, the weight that they are carrying against the loss function actually survive and everything else just gets gradually pruned back to zero. So that's kind of traditional L1. And now you're adding on this additional parameter on top of that, that is, and you're doing both, right? So you're, you're penalizing the weight itself. Is it multiplied by the length? I was, I'm, I kind of, I'm not the best notation guy. So as I was reading the paper, I wasn't a hundred percent sure of all that.

• 14:25Speaker 00

  It's multiplied by the distance. Say for layer i and layer i plus one, the closest two neurons, the closest, uh, uh, connection are still non-zero. So it's like every weight has a non-zero distance. So got pruned even if, you know, it's the most local connection.

• 14:47Speaker 02

  Gotcha, so like the minimum would be one if it's the same, like index in the layer from one to the next. But the more distance there is across the layer, um, and again, go look at the figures in the paper. It will, I guarantee you, it will really help make this much, much clearer.

• 15:05Speaker 03

  Hey, we'll continue our interview in a moment after a word from our sponsors.

• 15:09Speaker 02

  So basically they can get, the penalty can get stronger as the proportional to the length going up.

• 15:15Speaker 00

  Yeah, so, so if you go straight, for example, if you go straight from this neuron to the nearest neuron, uh, in the next layer, the penalty is like one. But if you go, uh, if you have some certain angle, then probably the penalty is two. And then if you further tilted away, maybe it becomes three. So the network have the, uh, incentive to, uh,

• 15:38Speaker 00

  To make only the most local connections non-zero.

• 15:43Speaker 02

  And then just for completeness in in this regularization, there's also L two regularization. And that is just a little bit less aggressive version because instead of adding the weight itself, you're adding the square of the weight. And that would tend to make a smoother curve around zero. So you're not kind of pulled as hard to zero. And so you get maybe a little less sparsity with L two than with L one.

• 16:16Speaker 00

  Yeah, my experience with L two is that it will shrink the overall weight norm of all the parameters, but not necessarily push down one single parameter to, you know, really close to zero. So that you can really, you can really ignore it, prune it away. So like they are, well, I think they're engineering tricks that if you really care about pruning, you should try both. They have their own advantages and limitations. Like L one sometimes have optimization issues because when you get really get down to zero, you bounce back and forth. But with L two, as long as your step size is smaller than some threshold, you converge steadily towards zero.

• 16:57Speaker 00

  Uh, but but also for the reason I I mentioned before, like L two does not, you know, encourage one single parameter, really drive it to zero. Uh, but encourage all the parameters overall have a smaller weight norm. So they are like this kind of trade offs you uh want to try when you do engineering.

• 17:18Speaker 02

  So you run the process with everything that we've described so far. So there's this balance between all the weights are tending to zero, the weights that are. That have a long distance on the connection are particularly trending to zero. But of course, the prediction accuracy itself is maintaining some weights up because you need to make good predictions. And so you run this and you do indeed find and you show these images of sparse networks. But then you add this additional concept of neuron swapping. And that kind of takes the the outcome from like sparse to sort of sparse, but also much

• 18:02Speaker 02

  more symmetrical, much more aesthetic, you know, much more elegant. How did you come up with this idea of swapping? And then we can get more into like, understanding it as well.

• 18:16Speaker 00

  Sure, so there's an anecdote behind it. So when I first tried modular vision without swapping, the networks looks a little bit like 3D DNA, the double helix, but projected to 2D. So so it looks something like this. It just swaps and then swaps back. So I was about to try. So I was thinking, if I make these neurons movable, make their embeddings trainable, I should be able to avoid this topological, so-called topological problem. But then I asked my advisor, Max Tegmark, and he asked me to try something simpler, as simple as possible.

• 19:00Speaker 00

  And so he suggested that I try the swapping thing. That sounds like the simplest thing we could try that can solve this double helix problem. So it turned up working pretty well. So that's why we end up using swapping. But potentially, one could also imagine try make those embeddings trainable. And also, in hindsight, we think that sometimes the your inputs and outputs do not have a natural, you don't have a preferred ordering of your inputs and outputs.

• 19:35Speaker 00

  You want an automatic way, you want your neural network to determine which kind of order you want. You don't, perhaps you don't want to pre-specify the ordering. So this is a really contrived example, but if your task is just swapping two numbers. So your input is X and Y, X on the left, Y on the right. And then your output are Y and X, Y on the left, X on the right. Then, uh, your neural network would have two really long connections and which also cross each other. Well, this is okay, right? Because you can ima- because you can understand what's happening. Um, but it's not aesthetic. Uh, we, we want everything to become local and modular.

• 20:18Speaker 00

  So, uh, naturally, we want to swap these two things, uh, either the input or, or the output. To make them, to make the connections as short as possible. Just go straight to where it should go. Yeah, that's the philosophy, I think, uh, behind swap. But, um, for some tasks, we find that swapping does not help that, that much. It's basically just reorganize things a little bit so that the pictures look better. But for other cases, it indeed help, even help neural network training. It help us get lower training loss. Because you can imagine if without swapping, you can, your neural network get stuck at those weird double helix, uh, structure.

• 21:04Speaker 00

  And so you have a large penalty. Uh, a large penalty can, uh, you know, there's a trade-off between prediction loss and your penalty. A large penalty means that you need to sacrifice the prediction loss. So if you can swap these neurons, uh, you know, cleverly to get a more local, con- local connected configuration, your penalty dropped down.

• 21:29Speaker 00

  Which means uh given the same uh penalty strength, you can get a better uh prediction loss. So um there's this hyperparameter like uh how frequently you do the swapping. Uh I think I just tried it to be, I I just set it to two hundred or set it to infinity. But you could imagine that maybe there's a face face diagram, face transitions like as in physics. When we talk about ice, water and gas, right? When we tune these hyperparameters, maybe there are different phases that can emerge. Yeah, but but right now I just uh set them to be something simple, uh but uh more things to be done in the future.

• 22:13Speaker 02

  And again, just to be super explicit, the the process of the swapping is basically, okay, we've been, we, every so often, every however many steps in our training process, we

• 22:30Speaker 02

  add an additional step, which just goes down the layers and compares, not even compares, but takes two different neuron positions. And essentially says, would loss go down if I just flipped the position of these? Where loss is already reflecting the the length penalty. So you can just do like a bunch of pairwise comparisons and find those enhancements. And essentially, it sounds like that is, it's more than aesthetic, as you said, it's also kind of helping you not get stuck in local minimum. Because, and I guess how those local minimum would potentially form would be like

• 23:15Speaker 02

  you have, you reach some kind of stable uh place where the the prediction is like, you know, pretty good. And the length, you know, of whatever the structure has kind of settled into, the length like could be shortened, but there's not really a mechanism other than the swapping to shorten the length if the other weight isn't helping with prediction just yet, right? Like a a weight that's not active. Like I guess a weight that's not active can basically not really be turned back on in this scenario, right? Is there a?

• 23:49Speaker 00

  We only care about important weights in the sense that have a large value. So actually, we we have, we attach each neuron a significant score, an important score. And we only care about swapping the top 10, say, important neurons with other neurons to see if such swapping can reduce the locality loss. Yeah, and also we constrain, like, we only swap two neurons in the same layer such that the prediction loss remains the same, but the locality loss may change. And we want the swapping that we reduce the locality loss.

• 24:32Speaker 02

  And the the prediction loss would stay the same by definition, or no? Am I missing something there?

• 24:38Speaker 00

  Oh, the prediction loss remains the same by definition. When I say I swap two neurons, I actually also swapping the incoming and outgoing weights. So basically, the computational graph, topologically speaking, remains the same. So the function, so the whole neural network as a function remains the same. But it's just geometrically speaking, it changes.

• 25:02Speaker 02

  So that's true at that step, but within the broader optimization process, it can still help achieve better end performance because you are applying less of it. You're helping it kind of apply as little length penalty as it can along the way, essentially.

• 25:23Speaker 00

  Yes, so that at single point, when we do the swapping, the prediction loss remains the same. But in the long term, having, you know, having connections in the right place, we'll encourage, you know, the network to to decrease prediction loss. Otherwise, if the network got stuck at local, some bad local minimal configurations, this would, you know, have a negative effect on the prediction loss.

• 25:50Speaker 02

  This feels like something when I saw it, I was like, why did nobody do that, you know, years ago? Do you see this as like something you're like? Also feel that way, like, why did nobody else think of this sooner? Or was there some unlock that, you know, you think made this kind of an idea whose time has come now?

• 26:09Speaker 00

  Actually, there are, there are literature, although when we wrote, when we wrote the paper, we are not aware of those related papers. So like, like 10 years ago, Jeff Kuhn is, is an expert in this. But like their papers, he and collaborators, their papers basically focus on the biological side. And they evolved their tiny networks, like with evolutionary organism, just to mimic what's really happening in biological evolution. Well, they have some follow up, which, you know, extended to gradient based method, but still for, like, like in different setups.

• 26:51Speaker 00

  So first of all, like, the previous works did not properly leverage, or at that time, the machine learning or gradient based methods are not as popular as it is now. So like, we're integ- we're integrating the current machine learning tools. Secondly, like, we have different focus. Our focus is like for doing mechanistic interpretability to understand how neural net work internally. And again, as I said, they care more about how modular it emerged from biological systems. So I would not shamelessly say that we totally, this is a totally novel idea. Well, we always stand on the shoulders of, of, of, of, of people before.

• 27:35Speaker 00

  And they're like clever minds, we always learn lessons from. But still, they are, we're in this different setup. And we also integrate like this new tools available only recently in the machine learning community.

• 27:50Speaker 00

  Uh, as long as it's a useful tool, I mean, there's no harm. We rediscover things, right? Because we find it's valuable things, uh, only valuable, only really valuable things can be rediscovered and rediscovered over and over again. Yeah, that's it. I wouldn't say that it's totally novel, but it still has some contribution, I guess.

• 28:12Speaker 02

  Everything's a remix, uh, but this had a feeling of kind of, I bet a lot of people saw it and thought, why didn't I think of that? Uh, because it just feels so, once you, especially when you see that animation, it's like, man, I definitely should have thought of that. So maybe, maybe it's more that I'm reacting to how compellingly the final result has kind of been packaged up and communicated in terms of the, the density of that realization than anything else. But it certainly creates that, creates that feeling.

• 28:40Speaker 00

  I see, I see. Yeah, I was really, I myself was really shocked when I got the results. When I actually see those connectivity graphs, I was shocked. How possibly a neural network can do this complicated thing with such, you know, spars and modular networks? And everything, actually, we, because we are actually, uh, using this locality trick. And I'm a little bit worried that locality is not, you know, exactly modularity, but it works. So there's still a lot of things unexpected here, even if I, uh, we first came up with the BIMT idea. We're still not sure whether it can work on these examples, what we are end up.

• 29:21Speaker 00

  And when I first see that, you know, the really sparse and modular graphs of symbolic, for those symbolic regression tasks, uh, I'm like, oh, we are really on to something and we should really dig deep into it and try a lot of, uh, examples that people already have tested in Mac and Terp and make it a real thing. So, yeah, it's a lot of surprise along the way.

• 29:46Speaker 02

  So a couple little follow-ups here, I, because the, in particular, I definitely want to dig in a little bit on the distinction between locality and modularity. And then I definitely want to get into also like the things you actually did and, and kind of understanding, you know, now what can we interpret out of all this? Just one more little follow-up on the swapping. Does it seem likely to you that? The swapping is like more valuable at the beginning of the training process. Would I, would you expect that like front loading those swaps would be more useful than like late in training swaps?

• 30:24Speaker 00

  That's exactly the opposite of what I'm hoping for. Yeah, so right now I'm thinking of applying BIM to pre-trained large language models. So what I hope is that I can just take pre-trained large language models like at the end of the train, at the end of the pre-training and apply BIM to just, you know, do a little bit fine tuning. Because maybe we don't have the computer to do that. So hopefully, uh, the swapping would be most valuable for, you know, at the end of the training. But um, since you, since you brought up the point, I got a little bit concerned. Maybe I tend to agree with you that maybe swapping is most valuable for, uh, you know,

• 31:09Speaker 00

  at the beginning of the training when, you know, at the beginning of the training, there are like many lottery ticket directions you can go. And swapping, it's the time that the neural network needs to decide which way you need to branch, which lottery ticket you want to branch into. And after you branch into that basin of attraction or lottery ticket, so to speak, swapping becomes no longer uh important because you're already in that basin. You only need to, you know, uh, roll down to that uh, to that bottom. You don't need to select uh which basin. So, yeah, but, but that's just my conjecture.

• 31:53Speaker 00

  Uh, we'll see uh how it works on uh language models.

• 31:58Speaker 02

  I'm sure you've looked at like the get rebasin paper had sort of a similar concept of, I think they called it kind of confusingly aligning the model and there was some sort of.

• 32:16Speaker 02

  Swapping mechanism that kind of, you know, was meant to like take two, you know, same shape networks trained on different things and somehow kind of align them to each other so they could be merged?

• 32:27Speaker 00

  Yeah, I'm a bit like for the GIBT rebasing paper, I think um there's a follow up paper which contradict the idea of GIBT rebasing called mechanistic mode connectivity. So if a network has multiple mechanisms to do the same thing, for example, modular addition, you have uh many group representations that can achieve the same goal. Then actually you cannot make two, you know, by just swapping two neurons uh to uh by swapping neurons to uh make them equivalent. And that's also observed in in my code uh well, this may not be something relevant. But it's interesting to note that when I uh run the modular addition example,

• 33:12Speaker 00

  I should to make my code reproducible, I should really set the precision to be double, to be double precision numbers. Even for float precision numbers, like the truncation arrows can domi- can dominate. Okay, like like this kind of noise is large enough to make you go to another basin. Uh this is, yeah, so like there are exponentially many equivalent basins there. So this is like this permutation symmetry uh that makes the um you know, the neural loss landscape very complicated.

• 33:52Speaker 00

  Yeah, just a note, uh if you want to make the codes uh reproducible, you need to use the double flows instead of the, you know, it is a double precision numbers.

• 34:01Speaker 02

  And is that sixty four? How many how many digits is is in that?

• 34:04Speaker 00

  Sixty four, yeah.

• 34:05Speaker 02

  Yeah, so if you were to round down to, you're saying even at thirty two digits, it still has this problem or?

• 34:13Speaker 00

  Yeah, that's problem specific, like only for modular addition and permutation symmetry, like for the algorithmic data sets, you need to set it to sixty four uh digits because additional to.

• 34:28Speaker 00

  Neuron permutation, you also have this, you know, representation equivalence. But for other tasks like image classification or symbolic regression, there is no such problem because you don't have this, you know, representation problem. But unfortunately, for language models, you only need to learn the token embeddings. And there they have such problem. I'm not sure if there is any bug in my code, but this is what I found. I need to, for trainable embedding tasks like modular addition or in the future or language models, you need to use sixty-four bits to make sure reproducibility.

• 35:12Speaker 00

  Yeah, that's a little bit unfortunate, but.

• 35:15Speaker 02

  Modularity and locality then. I was thinking, I was, I was thinking about that. You know, locality obviously doesn't necessarily imply like clear boundaries. I guess that's, that's the simple way I was framing the question. So if you wanted to be more encouraging of boundaries or sort of, you know, separation of concerns, is there a loss that you could conjure up to do that to? Could you sort of count like the number of non-zero weights coming in to a position as part of, as part of the loss? Like I, I'm sure you've thought about this.

• 35:58Speaker 00

  Yeah, so there are a bunch of off-the-shelf modularity measure in the literature, especially in graph theory. So the idea basically you want to detect a community where you want to separate a whole graph into different sub-grabs where intra-subgraphs, the connection are quite sparse, but intra inside the subgraph, the connections can be dense, but across different parts of the subgraph, the connection can, should be sparse, something like that. We already have such measures for modularity, but they, again, they're just defined on graphs, which are topological objects.

• 36:41Speaker 00

  They're good, but it's hard, it's again, hard to visualize them. Yeah, this is just some preliminary thought. I'm what I'm thinking is maybe we can combine the BIMT penalty loss with uh some modularity loss. People use in uh graph theory, something like that. But I I but I I I'm a little skeptical of this because uh it's it's like uh it's it's more like an engineering trick. There it it's not um, maybe we can have something more elegant than this. But uh it's the it's the simplest thing that writes on top of my mind.

• 37:19Speaker 02

  As cool as this is, you know, we're not at the end of this line of uh research, that's for sure. Maybe worth checking out the Tiny Stories project.

• 37:28Speaker 00

  I I noticed that I noticed that project that's uh that's super cool.

• 37:32Speaker 02

  You know, you mentioned like the compute budget obviously is is a meaningful constraint when you are trying to run this research. Probably in any setting these days, but certainly, you know, in an academic setting can be a a challenge. Their explicit goal was to create a data set that would capture, you know, the sort of. I don't think this is a well-defined concept, but you know, the full complexity, let's say, or almost the full complexity of natural language. While also being like reduced vocabulary, relatively simple concepts, you know, three-year-old can understand. And then they were able to get a lot more of the kind of reasoning sort of behavior that we see in larger language models.

• 38:16Speaker 02

  They were able to get it at, as far as I know, like the smallest scale models that I've seen. Because I think essentially they sort of narrowed the universe of what had to be learned while still kind of preserving. You know, the value of learning reasoning in that context. So there could be something perhaps interesting.

• 38:36Speaker 00

  Yeah, this is actually related to uh one paper uh by our group led by uh Eric Michaud. Called the quantization model of neuroscaling. So we have this quantization hypothesis where we conjecture that uh a prediction task. Say, language modeling can actually be decomposed into many subtasks.

• 38:59Speaker 00

  Uh called quanta as in physics, we we call we call them energy qua- quantum. So uh those quantas, they have different importance, like they appear with different frequency in, you know, in lo- in natural language. So presumably, we can order this quanta from the most significant to least significant. So we have a quanta sequence. And we conjecture that neural network would first, based on its capacity, would would first learn the most significant quanta. And then learn, learn it in sequence with this decreasing order of uh of importance.

• 39:42Speaker 00

  So I would imagine that to say, to speak English coherently as a three-year-old child would only require, say, the first-hand quanta. But to know about, you know, to role-play a physicist, or say, something more fancy than just, you know, a three-year-old can, uh you need less frequent uh quanta that appear in the tails. So that's why you need a very large language model, simply because you need to, you need to memorize, so to speak, that kind of facts. But the grammar and the, yeah, but but to speak English coherently, you don't need that uh language model.

• 40:26Speaker 00

  So, yeah, I I just want to mention that the tiny stories also has a con- connected back to uh our quantum model.

• 40:34Speaker 02

  That's cool. I could definitely imagine a, you know, reconvergence of these uh ideas, and it would be really interesting to see if you could actually start to get a visual on like the different, you know, I've been using this term um reasoning micro skills to kind of emphasize, you know, that they're, at least as I'm imagining them, they're, you know, very discrete, you know, probably very small little modules that, you know, do very specific tasks and, you know, kind of ladder up to more general-purpose reasoning. But it seems like what they see, and I guess this, you know, makes sense given what we understand about the mechanisms of training,

• 41:14Speaker 02

  is that like, super-specific things like negation, you know, kind of seems to come online as like a discrete.

• 41:23Speaker 02

  Skill, sort of a discrete skill in superposition, obviously with like a bunch of just memorized stuff and, you know, kind of baseline priors, but it seems to kind of rock these very, very specific little skills in kind of a, you know, a discrete way. So yeah, fascinating. Let's get to then some of the results. So just to recap, uh, you know, see if I can do this briefly. You say, hey, you know, we've built neural networks, uh, digital neural networks for the last like 10 years on a sort of very loose biological inspiration. But one of the biological realities

• 42:03Speaker 02

  that we've not really taken inspiration from is the fact that the larger the distance between two neurons in like a human brain, you know, the harder it's going to be for those cells to connect, uh, for very just practical, physical reasons. What if we kind of take that same concept and bring it back to the machine learning context? That turns into an additional penalty that corresponds to the length of a connection between neurons that gets thrown into the loss function. You reoptimize around that. You find that you're getting great, you know, much better sparsity, but then you also realize, hey, I can sometimes get stuck in these kind of local minima. So I'll add this swapping concept to come along and kind of untangle

• 42:48Speaker 02

  and help me kind of skip out of some of those local minima. And therefore we can both get great performance and kind of aesthetic organized, you know, almost crystal like structures that we can visualize at the end. You start off, there's kind of a series of experiments. The first one, uh, and again, make sure you look at these, uh, things. It will, it will help tremendously to see them. The first things that you're doing are these symbolic formulas where basically you're taking a few variables coming in. And then it's the job of the network to predict function output values based on those input values. And it's presented as symbolic functions, but I was a little confused by that because I was.

• 43:32Speaker 02

  To the best of my understanding, at the end of the process, the network itself is just numbers in, numbers out. Is that right?

• 43:39Speaker 00

  Right. So we we synthesize the data set. You know, we have symbolic functions, but we actually input X and Y numeric values and output F also in numeric values. But F is pre-computed as a symbolic formula of X and Y. So yeah, number seen, numbers out, everything numeric.

• 44:01Speaker 02

  Again, look at the graphs. But there's like, in each case, you're kind of demonstrating this conceptual sort of congruence, I guess, between the nature of the functions that the network is learning. And then the structure of the network that you can look at. So just to be super, super concrete, there is a function or there's a network that's trained on X. One, two, three, and four, so four numeric input values. And then it's trained to predict two functional output values. And I was surprised that they're not like super simple. One is X two squared plus sine of pi X four.

• 44:46Speaker 02

  That's one prediction. And then the other prediction is X one plus X three squared. So in other words, take two numbers, add them and square them. And then take the other two numbers, square one, take the sine of the other and add that. Okay, that's kind of random. Is there any? Should I understand that in any deeper way other than like you pick two random functions and that's kind of that?

• 45:09Speaker 00

  Yeah, yeah, independence is the main thing. We we we want to see that the network can split into two parts, two independent parts. And for the functions, we just pick. Yeah, we just random take it like on the top of my head or like squared functions, cubic functions. A sine, a square root, something like that, as simple as possible.

• 45:31Speaker 02

  Okay, so just, I feel like I didn't maybe make that as clear as I could have. The independence concept is that because the two output functions that you're training the network to learn and predict numerically are such that.

• 45:46Speaker 02

  One of those only depends on two of the inputs and the other one depends depends only on the other two inputs. Then if everything is working according to our initial theory, then we should see that there's just two parallel paths through the network that don't interact at all. And because they don't need to, because information doesn't need to cross in that way. And indeed, that is what you see. So you have presumably due to the swapping, as it's shown in the paper, X four and X two end up together on one side and they feed directly up to their output function. And then X one and X three feed directly up to their output function. I guess you probably chose it that way so that they were forced to be crossing initially.

• 46:30Speaker 02

  And then you you see that they uncross as you go.

• 46:33Speaker 00

  Yes, we we deliberately put like X one to X four, but like the outputs, the first output depend on X one and X three, the second output depend on X two and four. We deliberately want to test if the swapping is effective enough to, you know, to swap the inputs such that they group into the correct group. We are happy to see that it indeed work.

• 46:56Speaker 02

  And they're they're extremely sparse. So these are these are literal graphs, right? There are just two layers in the network that is trained to do this task.

• 47:08Speaker 00

  Yes, and I I did do the task, so so so I indeed do some pruning task. One might be, one might suspect that those connections you cannot see can actually contribute a lot. But that's not the case. What you see, like I actually, I literally plot every weight. But the the thickness of that connection is proportional to the magnitude of the is proportional to the magnitude of the weights. So if you don't see it, this means that the weight is too small for you for for the naked eyes to spot. And I, if I deliberately set those small weights to zero, the network, the output of the network is on, is not affected at all.

• 47:54Speaker 00

  This means that this is not a visualization effect, uh, this is not a visualization, you know, stuff that makes you feel like they're small, but they're still uh contributing. It's not that case. They are not contributing and they're, they're indeed zero, close to zero.

• 48:10Speaker 02

  That's fascinating and I'm, I'm left feeling like this is a eureka breakthrough in the sense that, oh my God, you know, look how simple the structure is. And, and in this particular case of X two squared plus sine of pi X four, it is really just a couple of neurons that are doing the job, right? There's the two inputs, there's three active neurons in the, in the first layer. There's two active neurons in the second layer and there's only two layers. So essentially there's only five total active neurons

• 48:52Speaker 02

  that are needed to translate these two inputs to that functional form output. So I'm like, again, that seems like eureka, but then, you know, if I could be vulnerable with you for a second, I still don't really get how it's doing it. Like I'm looking at the graph and I'm like, it's still, what's not clicking in to my brain is like, oh yeah, now I certainly I don't feel like I could diagram a five node thing and like know how to predict that function. So how do you interpret that now that we have that like super lean, super sparse thing? Does that, is that like very meaningful to you?

• 49:32Speaker 00

  So, so, so like I can make some explanation what's the neural network did. My, my like, I, I can actually write down the symbolic formulas and trying to figure out what's, you know, the neural network are trying to figure out, are trying to figure out. So my take, so, so I, I was really shocked when I see the results too. Uh, but in hindsight, I think, uh, it's understandable in the sense that, well, just like physicists always are fascinated by the unreasonable effectiveness of mathematics. Here I would frame our surprise as the unreasonable effectiveness of smooth activation functions.

• 50:12Speaker 00

  So I'm using Sylu or Seilu as activations. They're both smooth functions. And in applied math, we know that if we want to approximate a smooth function with Fourier basis, the approximation error drops exponentially as we add more and more higher, you know, higher frequency modes. Um, and the statements can also generalize to other smooth basis functions, not just Fourier modes. So Sylu or Seilu in our case. That's, well, I understand that's not a very satisfying explanation.

• 50:49Speaker 00

  But the, the, but my take is that, wow, this uh smooth activation functions are uh remarkably are unreasonably effective.

• 50:59Speaker 02

  Are there basically two nonlinearities, like you have a nonlinearity at each layer?

• 51:05Speaker 00

  Yes, it has two hidden layers. At each hidden layer, it has a nonlinearity.

• 51:10Speaker 02

  If people know any activation function, they might know the ReLU function, which is zero if the value is negative and then just, you know, y equals x if the value is positive. So it looks like a straight line and then a sharp corner at zero and then a straight line going up. And the Sylu function is essentially a curvy approximation of that, which remains differentiable at zero for one thing.

• 51:37Speaker 00

  Right, right. So basically you have a ReLU function, but you drag, you drag around zero uh to make a, you know, to, to make a well there. Uh, so Sylu now becomes non-monotonic. But asymptotically, when x, when your input is very small, negative, or very large, positive, the asymptotics are the same as ReLU, but it's only around, it's just that around zero, it's, it's differentiable.

• 52:07Speaker 02

  So if you want to use ReLU, did you, I assume you tried ReLU and it doesn't work? Is that a safe assumption?

• 52:14Speaker 00

  That's a good question. Actually, somehow I never tried ReLU. Uh, yeah, it is, it's my unreasonable craze for Sylu or Seilu.

• 52:25Speaker 02

  When you said. It's the unreasonable effectiveness of the smooth activations. You're confident enough that this is the better activation function, you don't even need to try the old one.

• 52:37Speaker 00

  Right, right. And there are papers like proving that with Sylou or Selou, these kind of smooth functions, you can construct

• 52:46Speaker 00

  construct like quadratic functions or multiplying two numbers with just two or three neurons. So for example, for the squared function, if you, you know, the Sylou or Selou function is non-monotonic. So there's a bottom. So if you tailor expand around the bottom, you got a parabola, the quadratic function. So this may sound like you only need one neuron to approximate the quadratic function, which is actually true in construction. But in practice, what really happens, because the chance of being initialized around the bottom is very low. So what really happens in practice is that neural network have two neurons.

• 53:31Speaker 00

  Both neurons have their own first order terms. But somehow in the later layers, they they weighted them such that their first order terms cancel, but the second order term survives in tail expansion. So that's where, you know, the the cubic, uh, sorry, that's where the squared function comes from. By leveraging this sneaky tail expansion trick that,

• 53:59Speaker 00

  you know, I myself did not think about it. Neural network can be this sneaky. But they, but the network just discovered this themselves. So it's a little bit shocking to me, like in some sense, uh, neural networks are more clever than myself.

• 54:18Speaker 02

  Okay, I'm looking at this.

• 54:22Speaker 02

  Train network. It takes these four inputs. It predicts these two functions.

• 54:28Speaker 02

  You know, there's this remarkably sparse structure that is able to achieve this like pretty, you know, non trivial functional form with just a few neurons. And then on that, you report the loss. And the loss is on this particular thing. I guess it's actually this. It's the joint task of the two predictions. But nevertheless, the loss is 7.4 E minus three or 7.4 times 10 to the minus three. I guess we can jump down to the bottom, right? There you guys have some, some graphs in the appendix that kind of showed this. I was just trying to figure out like, what does that

• 55:12Speaker 02

  loss mean in terms of, is it like really tightly fitting the functional form? Or is it kind of loose around it? Like, how, how close does it actually come to learning the functional form?

• 55:24Speaker 00

  Yeah, so I have the same concern when I just look at the loss function. Yeah, that's why in appendix, we also plot the scatter plots to see how well the predicted results aligned with the ground truth results. And you see, basically, they lie on the line and the R squared is like 0.999 or something. So, so that's pretty good. But set, set, set that still, including this penalty can also still degrade the performance. Because as I said, to approximate the quadratic function, well, you, you can approximate a quadratic function acceptably well with just two neurons. But if you include more neurons, the approximation

• 56:09Speaker 00

  could be better and better. So there is this trade off between accuracy and sparsity. Presumably, there's a Pareto frontier, something like that. There's no best solution. There's a Pareto frontier of the solution trading off between accuracy and sparsity and choosing lambda. I mean, choosing the penalty strength actually make, make us move, make us move along that Pareto frontier. Maybe I set lambda to be small. I didn't try it. Maybe I should. If I set lambda to be small, maybe you'll see that there are three or four neurons, but the prediction loss could be better.

• 56:44Speaker 02

  If you turn the prediction loss up to infinity, then presumably everything just goes to zero and your predictions are all terrible.

• 56:51Speaker 00

  Right, yeah.

• 56:52Speaker 02

  And then if you turned it down to no penalty, then you just are back to the beginning where you have a no incentive to sparse five in the first place.

• 57:02Speaker 00

  Yes, yes.

• 57:03Speaker 02

  Yeah, it is pretty. So this is figure 10 in the paper and the you know, the the scatter plot is tight to the line. I can say, you know, it's it's basically you know, at least for that one, it's almost indistinguishable from the line. And a couple of these others, you can see a little bit of wobble around the line, but it's still, you know, very close. Do you have any sense for how this works? When it comes to like, all of these inputs are on the on the interval minus one to one, right? Did you look at what would happen if you just started to go a little outside of domain on the input? Could you like, just

• 57:45Speaker 02

  would that like ruin everything?

• 57:47Speaker 00

  Yeah, good question. I honestly, I didn't try that. I would imagine it will fail. Like, I don't believe in like I don't think in this problem, there are like systematic generalization. It's not, it's like, we have no information about the outside. So the network at most can, you know, just do interpolation in in the range we trained on. So that's my guess.

• 58:12Speaker 02

  That's maybe a good bridge to start to move into.

• 58:17Speaker 02

  You know, because we're still kind of climbing the ladder of complexity in these experiments, right? The first one is purely, it's, you know, synthetic data set derived by these symbolic functional forms, but the network itself is just taking numbers from some narrow interval predicting essentially the curve. We see that amazingly, it can learn to do that with just very few active neurons, and we attribute that to the unreasonable effectiveness of the smooth activation curve. But then we could still ask like

• 58:52Speaker 02

  Okay, we're seeing these kind of conceptual notions that we designed into the problem reflected back to us where like, you know, when we set one up so that in theory, we could have two completely distinct sub graphs, indeed, that's what we get. And when we set one up with feature sharing, where like, you know, one of the outputs only depends on one of the inputs and the next one depends on two and the next one depends on three. Again, we see like the right information flow that, you know, aligns to our expectations. And then you've got the third one, which is compositionality, which, you know, starts to look a little bit, you know, even more tangled, but like, you know, clearly has some, some parts that I see as kind of interpretable where like,

• 59:34Speaker 02

  it appears to me that like a square root function, you know, it kind of consists of like, much like a square function consists of like one neuron going to two and then back to one. So we see these like little motifs that kind of pop up. And that's all super interesting, but then you could still ask, has this gropped anything, you know, that's kind of can, you know, more conceptual than learning, you know, the explicit shape of a curve. And your expectation is no, but maybe in some of the later experiments, that starts to become more relevant, right? So.

• 1:00:10Speaker 00

  I mean, even if you cannot generalize outside the training data, it's still interesting to understand, like, so I guess what physicists believe that's all the theories we have are effective theories. It's only valuable, it's only valid within, you know, in certain energy scales, so to speak. Yeah, even we, even if we cannot generalize, there's still a lot of interesting things we can say about the, at least for the task.

• 1:00:37Speaker 02

  Cool. So the next one, I think, again, is a really nice little visual, the two moon classification problem. You've got a bunch of points in a 2D space. They're color coded for the viewers' help in interpreting them. The model itself doesn't get that in, you only knows like the, the position, right? And then has to use the position ultimately to make a prediction of which class the data point belongs to. And it has to learn this kind of curve that separates these two, you know, almost overlapping, but not quite overlapping shapes. And so indeed, it does that. And again, you see like a pretty sparse result. I was a little bit just struggling to interpret the last graph there where

• 1:01:22Speaker 02

  there's two numbers in, right, the x, y coordinate, and then there's the out of like, but it seems like all the, all the nodes are going to class one and like class two has been sort of left.

• 1:01:31Speaker 00

  Yeah, that's a very interesting observation. I was surprised by how clever neural networks can be. When I saw the connectivity graph. So for binary classification, what really matters is like the difference of the largest of the two classes. That is the relative magnitude, but not, but not the absolute magnitude. So an efficient strategy for neuroclassifiers is to simply set class A to always have zero logic. While only learning the logic function for another class B. And a positive logic for class B means that the classification is B.

• 1:02:15Speaker 00

  And a negative logic for class B means that the classification result is class A. So if we also look at the evolution of the whole thing, the evolution dynamics is actually very interesting. There is this intermediate phase where you see that both classes have outputs largest, they have network connections to the outputs. And they're almost symmetric. But as training goes on and the pruning, more and more weights are got pruned, you see that the network learns to be, to transit from this symmetric phase to the asymmetric phase. Because the asymmetric phase is more energy efficient, requiring fewer neuron connections.

• 1:03:00Speaker 00

  Yeah, so, so, so, so it's very interesting to see that there's actually this phase transition. At first it's messy, everything fully connected, and then in the middle, in the middle state, they're like this sparse network and also symmetric for the two classes. And finally, the network realizes that it only needs to predict the logic for one class while pruning away totally for the other class. So it's a very clever strategy that I learned from my neural networks.

• 1:03:31Speaker 02

  You know, I'm studying that final visual and that definitely jumps out.

• 1:03:37Speaker 02

  It's basically, you know, reduce the dimensionality of the problem on its own to just having now to choose, make one prediction instead of two effectively. I don't have any other immediate intuitions for the shape of what I'm looking at. Is there anything else you could say about that?

• 1:03:53Speaker 00

  So it, they're sparse. What's good about sparsity is that now you can, you know, that there are just seven, if I remember correctly, weights in the graph. So you can just intervene any important weights or neurons to see what it did to the prediction results. I'm guessing because the, the problem is too simple, it does not have any meaningful structure to emerge. So maybe the whole thing itself is a module. So that's why we don't have very good explanation for, uh, for it. Because itself, the whole thing is a module and we don't have a good, we don't expect to have a good explanation for internal activations inside modules.

• 1:04:34Speaker 02

  Yeah, especially for something as kind of arbitrary as learn to separate two.

• 1:04:39Speaker 00

  Right, right. Exactly. But in PANDIX, we did the intervention experiments. You can see which, uh, what each weights and neurons are doing to the.

• 1:04:50Speaker 02

  Yeah, it's amazing. There's only six active neurons across the two layers, four in the first, two in the second. Oh, you're actually eliminating connections in this case, is that right?

• 1:05:01Speaker 00

  Yeah, I'm neutralizing individual way, but yeah, yeah, but one in principle can neutralizing the neurons. Yeah, and, and, and also, uh, one thing to note is that, well, if you change random Cs, the, the graph would look completely different. That's another sign that there is no consistent modularity. Uh, but if we move on to the modular addition case, uh, in the, the, the, the, the next example, you would see that no matter random seed you have, you always, you almost always have three parallel modules emerging. So there, uh, so, so this is a more consistent thing you can say about the task.

• 1:05:46Speaker 00

  Yeah, random seed also play a huge role here, but my take is that if there's a consistent structure in the task, no matter what your random seed is, you should, for the most time, you should be able to find it. And in the least, uh. In the most difficult case where like only some of the random seed can find the structure, that's also fine. You can just select the most interpretable, you can just run a hundred models with different random seeds in parallel and then select the one that you think you you feel most interpretable to you and you go from there to do map inter.

• 1:06:25Speaker 02

  It seems like we probably could graph, maybe even in just closed form, again, I'm not amazing with all the linear algebra notation, but either either in closed form or at worst, like with just a mesh of points, you could kind of graph the value of class one as kind of a z-axis over the two, you know, inputs and kind of get a visualization for like, oh, essentially you're learning a sort of elevation landscape.

• 1:06:56Speaker 00

  Yeah, that's correct. So in appendix, we write down the symbolic formulas explicitly, well, basically just to extract those weights and biases and and write them down into a symbolic formula. So yeah, as you said, in principle, you can plot the, you know, 3D surface plot like that. That may be more intuitive to see, but this can all, you know, in principle, one can do that.

• 1:07:26Speaker 02

  So yeah, there's no, it makes sense then that there's not any super interpretable, you know, there's not like a two-sentence summary of this because it's an arbitrary shape. It's almost like you just kind of sprinkled some stuff out there and it had to kind of learn this particular shape, but that's not like a super simple problem in the first place.

• 1:07:51Speaker 00

  Yeah, or using the language of the quantization model I mentioned before, the task itself has only one subtask. So there's no need to modularize itself. If we're dealing with some compositional task, like language modeling, it actually involves many subtasks. Then there you have the incentive to like grow those modules for different abilities. But this two moon classification or later, the transformer example on linear regression, they're just, you know, one. Single task, uh at least I myself cannot imagine their subtasks, subtasks that underneath this uh whole task.

• 1:08:32Speaker 00

  So that's why uh the graph look a little bit, you know, look less interpretable, look less modular, but that's because the whole network is presumably a modules, it's probably a module and there are like this degree of freedom kind of thing that make the graph looks messy.

• 1:08:50Speaker 02

  If you were to add a third moon to this, or you know, a third region, would you expect that that would then create some sort of subtask that would, that you could see reflected in modularity?

• 1:09:08Speaker 00

  Yeah, maybe we can look at the Amnist figures, uh there, I'm not sure if I can say confidently there I see something modular, but but there's something like this pattern mismatching thing emerging that also interesting, but I, I can't say for sure that there are meaningful modularity there. Uh but that's a reasonable conjecture, maybe like classifying between Y and two is a subtask and then two and three is another subtask and yeah, and then you can do the, you know, classify three things by combining this binary classification task. I think that's a reasonable conjecture, but still need to be tested.

• 1:09:52Speaker 02

  So next, modular addition, one of my favorite problems in the world due to all of the mechanistic interpretability focus on it. The big questions I had on this one were, is this doing the same thing that we've seen kind of prior interpretability work demonstrate that?

• 1:10:13Speaker 02

  Um, you know, a GROCK network is doing, is it solving the problem in the same way?

• 1:10:19Speaker 00

  Uh, that's a very good question. Uh, actually, we have a follow-up work that's still not, we're still working on it. So the short answer is yes and no. Uh, but yes, I mean, the circles, the, the embeddings of those numbers are still four-way basis. That's the, uh, that's the same thing as in previous MacIntyre works. But by no, I mean that the model into the internal computations are different. We actually find a new, a different algorithm than what Neil and collaborators' paper described. So, so we call Neil's algorithm the clock algorithm. And while we discovered a new algorithm called the pizza algorithm.

• 1:11:02Speaker 00

  Uh, which algorithm the neural network ends up learning depends on architectures and uh, hyperparameters, and this can be very subtle. Um, and there could be phase transitions uh, from clock to pizza and also pizza to clock. So this is really funny to say out of context, but this is doing an ongoing work with my MIT colleagues, uh, Zi Qianzhong, Jacob Andreas and Max Tegmark. Uh, we'll post the pre-print to archive soon, but like this is for, so we discovered clock and pizza algorithm and uh, we passed the algorithm, the network we got with BIMT, and we find that BIMT actually ends up getting pizza algorithms.

• 1:11:45Speaker 00

  The new algorithm we discovered, not what Nandan collaborators described as the clock algorithm. But, but no one is wrong. It's simply that, you know, MacIntyre can be much more subtle than we, we, we could imagine. It's a very subtle depending on the architectures and also hyperparameters, all kinds of stuff.

• 1:12:06Speaker 02

  So the lesson here is there's a lot of possible mechanisms to interpret. And which one you end up needing to interpret is kind of cast in the upstream decisions of exactly how you lay out your network and training process, and in this case. Localization incentive as well.

• 1:12:27Speaker 00

  Right, yes, yes. Uh, yeah, yeah. In short, the clock algorithm is more accurate. But it's requires more resources. By contrast, the pizza uh, the pizza algorithm is less accurate. But it's takes a fewer resources. So that's why we see in in our BIM paper, there are three shapes. Uh, maybe I can explain more later because each shape, each parallel module is an imperfect pizza algorithm. So the network needs to ha- come up with some kind of error correction to make each imperfect algorithm to become a perfect one.

• 1:13:05Speaker 00

  Aggregate uh, this, the results from uh, each imperfect algorithm to to aggregate some of them cleverly. Uh, such that the final outputs are perfect. So that's the uh, takeaway. We hope to, we we we hope to post a paper into archive soon. And uh.

• 1:13:24Speaker 02

  Is there a simple, so I mean, the clock algorithm is you're doing modular addition. You take advantage of the cyclicality of, you know, the modular math. And so you sort of say, okay, I'm going to rotate some and then rotate some more. And then if I get past the origin, you know, I'm naturally kind of the position that I'm at. Doesn't matter how many times I went around, it's just kind of the final position, right, that I need to look at.

• 1:13:48Speaker 00

  Right, right, right. But by contrast, the pizza algorithm adds, you know, two numbers. So that uh, so so the final prediction is more like slicing pizza, slicing pizza. And determining which pizza slice the outcome uh, lies on. So the frequency of the pizza algorithm is like, is doubling the frequency of the clock algorithm. Or the other way around, I I can't, I can't remember exactly. But like there's this this uh, tiny detail there that actually distinguish between clock and pizza. If you don't look carefully enough. For example, if we use the metric, uh, Nanda and collaborators uh used in their paper,

• 1:14:32Speaker 00

  we cannot distinguish uh clock and pizza algorithm. We came up other metrics that can distinguish between these two algorithms. So there's still like this kind of metrics. Uh, each metric is a coarse grain or drop some information of the whole system. So even if two things are are different, but you coarse grain too much, you end up getting the same thing, you cannot distinguish two things. So so my takeaway from this, there's still a lot to be done in Macintosh. At least uh we need to come up with more, you know, metrics to really distinguish those algorithms.

• 1:15:11Speaker 02

  Yeah, I'd say we are just scratching the surface. So that's really interesting. When you say it is less accurate, is that

• 1:15:26Speaker 02

  a matter of like getting problems wrong? Or is it a matter of being less confident with the right answer? So your loss is higher, but you're still getting all the answers right.

• 1:15:38Speaker 00

  So when I say clock algorithm and or pizza algorithm, both algorithms have their perfect versions and imperfect versions. So what we end up getting in experiments are their imperfect versions. And the imperfect versions can uh make wrong answers on some samples. But it can be supplemented by and complemented by another uh imperfect algorithm which make wrong answers on another subset subset of samples. Uh, if you have enough number of sufficient number of uh imperfect copies, at least one copy would give the right answer on every sample.

• 1:16:22Speaker 00

  So by aggregating the results, finally, the final aggregated result would make the correct prediction on every sample. But uh for one uh head or one, you know, uh, parallel module, it can only make prediction correct on some subset of all the samples.

• 1:16:44Speaker 02

  So you essentially train.

• 1:16:47Speaker 02

  N copies of this tiny little network. And then you find that, I guess, due to random seeds, they sort of cohere in different ways. Each one represents an imperfect approximation of the ideal algorithm. And then in aggregating them, you can kind of get good performance, even if they have their individual weaknesses.

• 1:17:09Speaker 00

  I mean, we don't explicitly have this N copies. We just have a whole network. But we can somehow, this in Kango, and find that automatically their N copies emerge from training, which is all really fascinating. And this is again, an example of where I think, oh, maybe neural network is more clever than myself. It leverages kind of

• 1:17:35Speaker 02

  Is that what is meant by voting?

• 1:17:37Speaker 00

  Yes, exactly.

• 1:17:38Speaker 02

  Okay, so as we're looking at this graph, and again, it's not that big, right? We're talking about two hidden layers. So that is to say, inputs, two hidden layers. Again, I assume each with a non-linearity as part of that layer and then outputs. So there's not really that much room to do the work here. I've always kind of, through every graph in this paper, it leaves me feeling the same way about that. And then each of these, as you look at it, it's just clear that there are like three distinct modules that go from the inputs up to layer two. And then at layer two, it's basically all kind of fully connected, it looks like, again, to the output.

• 1:18:18Speaker 02

  But those three distinct sections are, you're saying, now I'm understanding these shapes more deeply. So each of those is an approximation. And I guess you can see that because you can do the ablation, you're still like.

• 1:18:35Speaker 00

  Yes, so each copy is an algorithm that tries to perfectly do the modular addition, but fails. So somehow the neural network, at some points, the neural network figures out that, oh, it's more expensive to form one perfect copy of a single algorithm. So instead, I resort, I form like three copies of imperfect algorithms and somehow aggregate or some sort of majority voting to get the end results.

• 1:19:05Speaker 02

  And how literally do you understand that concept of? Voting. I'm always, I always push really hard on analogies. Cause I find I confuse myself as often as I clarify things for myself.

• 1:19:17Speaker 00

  Yeah, so the terminology voting uh actually comes from, so we borrowed the term from uh error correction. Uh in in um as in information theory. So let's say you want to communicate a classical bit, zero or one, over some channel. The channel is imperfect. It can randomly flip the bit with some probability. So to achieve better accuracy, there's something called the repetition code. Which basically uh you repeat the bit for three times or five times or even more. So then the receiver then can infer the bit by doing a mo- majority voting of the three bits.

• 1:20:01Speaker 00

  The probability of making an error reduces exponentially as the repetition times grows larger. So this analogy of repetition code also applies to module addition here. Uh each shape or each module is an imperfect algorithm of module addition. Uh one module alone cannot, can make mistakes on some examples. But by aggregating these uh three modules together, or when I say vote, I actually mean aggregating. Just the last linear layer aggregate their results, but just by aggregating or voting their uh results. They can, you know, end up with a perfect classification uh results.

• 1:20:45Speaker 02

  In the kind of bit level example from the information theory. You basically repeat the thing, you repeat the signal three times. And you trust that, you know, it's unlikely that two out of three get degraded. So I can trust two out of three. That's the basic intuition.

• 1:21:05Speaker 00

  Yes, uh and the I, yeah, and also there is this independent voter. Uh like the flipping thing are like independent.

• 1:21:14Speaker 02

  So I'm trying to map that concept onto the data that I see here. And it's like, I would expect, and I'm being very literal. So I'm, I guess my approach in this is trying to be very literal and see if like I get confused. And I am getting a little confused because. I guess if if I was thinking of it in like a strict analogy and I said, let's imagine that I'm using this voting approach, but then I just eliminate one of my inputs. Now I've got two inputs. Now, like, if they agree, great. If they disagree, I mean, like, 50/50 chance. But what it seems here is like you knock out one and your performance is like much more degraded than 50/50.

• 1:22:00Speaker 00

  Yeah, so I think the majority voting thing is more like, it's more like a metaphor. What really happens is the linear aggregation of of the outputs. And it's it's more like cooperation. It's it's it's more like so. So in majority voting, we assume the voters to be independent to each other. But here, we sort of like, since we're adding this kind of penalty to the network, the three modules really need to cooperate with each other.

• 1:22:33Speaker 00

  So they rely on each other. Somehow they talk to each other in training. So, for example, the first module say, I take sample A. And then the the second module say, I take, then I will take sample B, something like that. But but even more complicated because it's a linear combination of things. We cannot partition things that clearly. So in principle, we can understand, we can try to understand this structure more clearly, but but I haven't done that. But I think it's something to investigate in details.

• 1:23:06Speaker 02

  So the major basis, I guess, then for the kind of notion of like

• 1:23:15Speaker 02

  a voting metaphor or the the reason that you're saying that you interpret these three modules as kind of each an independent approximation.

• 1:23:29Speaker 02

  Is more anchored to those shapes, I guess, like you you show kind of the the spatial representation and that definitely looks like something.

• 1:23:37Speaker 00

  Yes, yes, I think it's more anchored on the shapes. And it's surprising to me, at least to see that three modules emerge. And finally, I understand that, oh, there is this error correction mechanism that we didn't think about. Yeah, it's amazing. We learned something from neural networks.

• 1:23:58Speaker 02

  So when I looked at the the next problem around permutations, anything you want to say on that that kind of jumps out, you know, as the most important takeaway for you?

• 1:24:08Speaker 00

  So the best part I love about the permutation example is that Neo and the collaborators found that the network can automatically learn the group representations. And for for S four permutation group, they it has a nine dimensional representation. And but they still need to, you know, probe and find where to look at. But this nine dimensional representation just naturally emerged on privileged basis, on like aligned with the neurons after training. And there are exactly nine active neurons, and there are exactly one sign neuron, the

• 1:24:52Speaker 00

  twenty, the magic twenty second neuron corresponding to the sign of the representation. And all the other neurons, we can also explain them based on calligraphs. And it it tells us that the neural network indeed leverages the structure of the of the group data sets. So I think the main takeaway is that, well, with BIM, you can more, you can clearly know what you're, where you should be looking at by just looking at the graph. But with other methods like probing or intervention, at first you have no idea where to look at,

• 1:25:33Speaker 00

  or you need to really understand your problem such so that you know where to look at. But BIM, you don't need to know anything about the task. You just look at the graph and you can say something useful about it.

• 1:25:45Speaker 02

  Yeah, I have this. Visualization sometimes in my mind of like shrink wrapping conceptual reality.

• 1:25:54Speaker 02

  And I don't know, there's probably a lot of, that's kind of an analogy. So, you know, my analogy detector goes off. But it does feel like you've, you've kind of created another way to create additional kind of negative pressure to like suck all of the extraneous activity out of the network. And then, you know, it kind of snaps or, you know, coheres right into the appropriate dimension.

• 1:26:21Speaker 00

  So the paper, we only have those static images, but if you look at the videos I posted on Twitter and also on GitHub, this is exactly what you described. You seem to have some kind of pressure that push inwards and the shrink prune, the whole thing that, uh, yeah, really amazing. I encourage everyone to get the video, to watch the videos.

• 1:26:43Speaker 02

  Okay, so we're getting into the home stretch. So the next thing you do is then extend this to the transformers. Everything we've talked about so far. I wanted to ask too, like, how long does it take to train these really simple networks? Like if I, if you were to set up, you know, one run to train the modular addition network and it goes through how many steps?

• 1:27:09Speaker 00

  To 20,000, I guess.

• 1:27:11Speaker 02

  So how long does that actually take just on, you know, a fairly standard issue computer?

• 1:27:15Speaker 00

  Yeah, so actually all the results I got, I'm just using my Mac M1. There are no GPU.

• 1:27:26Speaker 00

  So for the modular addition, it takes around 20 minutes. For symbolic formulas, just one, less than one minute. The transformer example is the most, uh, is the most time consuming one, almost, uh, well, ranging from one hour to two hour, but. Hours at most. It's a so, so, so if you want to play with BIM, I encourage you just to start playing with the symbolic regression example. Uh, you can even make modification to the code is very fast, you know, iteration because to get the results, you only it only take less than one minute to get the results. So it's a it's a good point to start.

• 1:28:07Speaker 02

  So it basically can feel with, I guess, with these later projects, not quite, but with the the simplest cases, it's almost like kind of real time. Uh, I presume you could even just kind of, although, yeah, you, I guess, if you don't go all the steps, you don't see all the sparsity, so you do kind of have to run it to, you know, some sort of completion. You mentioned the the transformer portion. Uh, it wasn't entirely clear to me. Are you modifying the attention mechanism as well as the fully connected portion of the transformer?

• 1:28:42Speaker 00

  Yeah, yeah, yeah, that's a very good question. So think about attention layers. If we put aside the softmax part of the tension layer. So the tension layer is just three linear, you know, just three matrix multiplication, the uh key query and value. So this linear matrix, uh, we can understand them as a linear layer with zero bias. So we apply the same trick we used in MLP, just treat it as a linear layer. And uh, yeah, so so attention layer, the softmax has no parameter, so we don't have to worry about that. We just need to deal with this linear matrix as in MLPs.

• 1:29:28Speaker 02

  So the representation of this then for kind of purposes of the loss function is in the like fully expanded form? Like you're applying, you have this sort of three multiplications, but you're treating those as essentially one giant.

• 1:29:46Speaker 00

  That's a good question. So there are three matrices. I overlapped them. I stacked them along a third dimension. I actually did not plot. And the shift along the third dimension is very small, just to, you know, just to visualize it. But like because query key and value, they're actually in some sense, they're equivalent. So there's no sense to separate them in space. So they're overlapping each other. But visually, you can think of them separated by a small amount along the third dimension.

• 1:30:19Speaker 02

  Would this have a, for kind of scaling up purposes, I'm not great with this sort of optimization of this, but I definitely followed Neil Nada's work quite a bit. And what I understand from, I think, taking away from some of his YouTube tutorials, is that there's a divergence between the representation of the attention mechanism that is most compute efficient and therefore actually like gets used in code from a more, you know, purely analytical representation that's like easier to interpret. And if I understand correctly,

• 1:31:02Speaker 02

  you're using, you're applying the penalty in a way that works on that more analytically sound, like not separated, but maybe like less compute efficient representation? Do I have that right? Or am I going off track somewhere?

• 1:31:19Speaker 00

  Sorry, what do you mean by low representation here?

• 1:31:23Speaker 02

  So you had said a second ago that the sort of multiple matrix operations that constitute the transformer mechanism in analytical terms are not really separable, but they are, in practice, they are sort of coded separately, as I understand it, for efficiency reasons?

• 1:31:45Speaker 00

  I see, I see. So I treat them independently?

• 1:31:48Speaker 02

  So the penalty is added in and calculated in such a way that it would, that the process of calculating that penalty would scale? It would, the whole process would still scale just like a normal transformer scales?

• 1:32:03Speaker 00

  Yeah, yeah, exactly. There's no additional overhead in this.

• 1:32:07Speaker 02

  Can you maybe just help develop a little bit more of the intuition of?

• 1:32:12Speaker 02

  How exactly should I understand the locality penalty in the context of attention? In the, in the above like simple graphs, it makes total sense. It's just okay, 2D space, boom, I got it. Is there anything more that I should kind of be intuiting for the attention version?

• 1:32:31Speaker 00

  Yeah, so that's a very good question. Actually, it takes a lot of, it takes at least some uh sort of modifications to make it work, to make BIM work on uh transformers. Uh I don't have an intuitive picture for that, but like the modifications are uh we need to consider the, the head, the multi-head attention. And uh that affects swapping. We cannot just swap in two neurons in the same layer. We need, we only need, we can only swap uh the two neurons in the same layer and also in the same head. But we can also swap two heads as a whole. We can only swap two heads.

• 1:33:13Speaker 00

  And the residual stream is another pain in the ass. Because with the, so, so, so with the residual stream, this uh the residual stream, the, the neurons got aligned. Uh so the permutation in, along the residual stream are, are, are in some sense locked. So their permutations are not independent. The residual stream in different layers, they share the same permutation. Uh they, they share the same permutation. And that may also be uh one thing that makes the network uh less interpretable or less uh local.

• 1:33:54Speaker 02

  And that's kind of presumably related to like? Superposition or not, not really, I guess. Basically, the idea is that there you would have to untangle. If you're gonna swap, you'd have to swap at every layer of the residual stream.

• 1:34:10Speaker 00

  Yes, if you swap, you need to swap every layer of the residual stream. So that's uh sounds to me like um I like a constraint. It it's it's I I think it's undesirable. Because if you have this constraints, you you have to consider all across all the layers. So you cannot untangle some non-local connections in some certain layer. Because you need to consider everything uh globally. Uh maybe there are better ways to to do this. So so again, I think BIMT is still, there's still a lot of things to be done to improve BIMT. And also, another pain to ask is is the layer norm.

• 1:34:50Speaker 00

  So I think Anthropic published their uh results also arguing that layer norm, layer norm did something sneaky. So in in in our case, we also dropped. So they didn't drop layer norm. They just say, oh, layer norm is um is a is a pain. Yes, so but I dropped layer norm entirely. Because it's normalized things, sometimes even it has a small input, it's still normalized it to zero to one uh range. And that's something we don't want with sparse uh neural networks. Uh yeah, so again, so so there are a lot, there's still uh engineering tricks that we need to incorporate to make to make BIMT work on on transformers.

• 1:35:35Speaker 00

  This is just the first step that I I want to show that, well, in principle, BIMT is a high-level idea that you can apply to any architectures as long as you evaluate your architecture in a geometric space. Uh and this is just a prototype example, but a lot of things can be done in the future.

• 1:35:54Speaker 02

  That brings us, I think, to the last experiment, the last visualization, which is applying this to 3D space instead of 2D space. And here you're using the classic MNIST handwritten data sets, handwritten uh handwritten digit data set.

• 1:36:10Speaker 02

  And you know, this of all the visualizations, I don't know, there are, there's some good ones, but this one might be the best because you see at the beginning just a massive, you know, 3D tangle of intersecting weights. And then over the course of all the steps, you know, it gets, you know, just much more, first of all, kind of cropped in to the actual, you know, space where the images actually are. And then just obviously a lot more sparse. It seems like in this one, again, in the, you know, what it, what is it really doing is sort of limited, so we still need the kind of, you know, we still need the Neil Nanda breakdown.

• 1:36:53Speaker 02

  But you've kind of run the first leg of the relay in that you have, you know, by sparsifying things, made it pretty clear where to look, as you say. And then, you know, next step is for somebody to kind of figure out to what degree can we kind of assign any meaning to this? And maybe none, I guess, right? It could just be that this is like a super tight module that does this really weird task. And like the shape of the numerals themselves, you know, is kind of a,

• 1:37:24Speaker 02

  especially as like, you know, played out across a ton of handwritten versions is just weird. And there's not necessarily more to say about it.

• 1:37:32Speaker 00

  Yeah, that's a reasonable conjecture. It could be that the task could be just too simple. The whole thing viewed by the network is just a whole task. There's no subtask. Like even linear classifiers can get about, can get above 90% accuracy on MNIST. But it could also be that vision tasks, at least for image classification, demonstrates less modularity than a language task. As I imagine, the language task presumably have many quanta, have many subtasks. And also the chain of thought reasoning thing in language models might be related to modularity.

• 1:38:16Speaker 00

  You need, you need each step, you have a, have a, a concrete reasoning step. So I, I personally uh think that maybe BIM's ability to make language models more modular uh more than uh vision models, uh that's where I'm imagining. So, so, so, so, so about next step, I think I would try BIMs on language models. But also, I think it's important to, like lately, AI for science, try, try all kinds of scientific problems. Where interpretability, people also care a lot about interpretability in scientific problems. So I think that's where uh BIM can also play a huge role.

• 1:38:58Speaker 00

  So my next step would be a language model and, and scientific problems.

• 1:39:03Speaker 02

  Cool. Do you wanna comment just for a second on the pattern anti-matching? Followed by pattern matching, and so it's kind of like from, there's the input layer. And then in kind of the, there's again just two hidden layers, right? Between the first and the second hidden layers, looks like basically all the weights are negative. And then from the second layer to the output, all the weights are positive. And you're interpreting that as pattern anti-matching. I guess my answer is probably gonna be the non-linearity, cuz I was initially thinking, if we just change the sign of all the weights, would that be, would it still just work? But the answer, I guess, is no because of the non-linearity.

• 1:39:49Speaker 00

  Yeah, I, I, I think linearity, linearity also plays a role here. Um, so, so in appendix, we also tried uh one hidden layer and three hidden layer. We find this anti uh pattern mismatching thing are consistent. No matter what, yeah, no matter how deep your neural networks are, at least for one or two or three hidden layers, they're consistent. So I'm guessing there's something deep or fundamental about it, like uh this kind of strategies more biologically or fit, or energy efficient than what we would imagine, like how we classify, h- h- how the image recognition, we would, we would, uh.

• 1:40:32Speaker 00

  So previously, we would imagine that we, we want to do pattern matching. We have a template of things. Say, uh, we have a, we have a, we have a dog uh feature. And we scan over the whole image and try to max polling to see if there is a dog uh on the image. But it seems like by having this BIMT strategy, you end up doing something opposite. Like you want to tell, oh, this image does not have a dog. Uh, so it's more likely to a cat, something like that.

• 1:41:11Speaker 00

  So I haven't, I, I can't say that I have fully understand this, there, there, there's still things to be understood.

• 1:41:19Speaker 02

  Anything that you would add, you know, just to kind of wrap things up on the paper front before just asking a few more general questions?

• 1:41:26Speaker 00

  Yeah, so one thing I still want to know is that about, about the next direction, you mentioned tiny stories. So that's one direction I'm, uh, I, I'm looking to lately. Uh, one thing I'm hoping that uh BIMT can achieve on tiny stories, which is already achieved for module addition, is that the learnable token embeddings have privileged spaces or better line with coordinate basis. For example, uh, there is presumably a direction in the embedding space corresponding to color. Say, uh, with unprivileged spaces, you need to try really hard to search for that color direction and then project token embeddings to that direction.

• 1:42:10Speaker 00

  But there are actually infinitely many possible directions. But with privileged spaces or with BIMT, very likely there's just one dimension of the token embeddings corresponds to the color dimension. So you ju- you can just obtain that direction by enumerating the finite number of uh the embedding dimensions. So for, but yeah, one thing I'm thinking about whether it can scale, what, whether BIMT can scale to large uh networks. So my only concern so far is that it might be harder to visualize the whole network.

• 1:42:46Speaker 00

  This is apparent difficulty. But other than that, the way I'm seeing it, I don't see any, you know, deal breaker here. I don't, I don't necessarily see any factor that blocking it from scaling.

• 1:43:01Speaker 02

  I'm actually doing another interview with one of the authors of the Megabyte paper that has kind of the hierarchical approach. This seems like it could sort of play nicely with that as well, perhaps, although I don't quite know how. But, you know, it, it has these sort of a global model and then a bunch of these like local models. And I don't know, have you thought about any sort of trying to apply this to any sort of like more hierarchical structure such that like modules could sort of naturally cohere in different places?

• 1:43:37Speaker 00

  Yeah, that's a good question. Actually, uh, like, like I mentioned, um, outstanding researcher Jeff Kuhn on a connection cost thing like ten, even ten years ago, but in biology, right? They're just, uh, they discovered that by imposing this connection cost thing, you not only induce modularity, you also induce hierarchy. And also allows you to do lifelong learning or continual learning. So a lot of promising directions, uh, but the regularity thing. Uh, as you said, like, i- is it possible to have one module that can be copied all over the network?

• 1:44:22Speaker 00

  For example, there's a basic unit that you need recorded multiple times. Unfortunately, BIMT or this kind of connection cost techniques itself cannot induce this kind of repetition or regularity. And you need some. So so this is another thing I'm thinking about lately. I I want to draw inspiration from biology. Uh and some kind of stuff. And there there's some literature like trying hyper networks or some uh compositional producing networks, something like that. Uh to make sure the regularity or the repetition of the mod, to reuse the modules.

• 1:45:05Speaker 00

  But that's something orthogonal to modularity and hierarchy. I think there are many traits. I'm all very interested in like modularity, hierarchy, locality, uh module reusability, maybe they're connected somehow. But at least beams can cannot solve the last one, cannot solve module reusability. Uh that's one thing I'm thinking, maybe I can uh improve. Integrate some kind of hyper network technique into BIM to make it also leverage uh this kind of module reusability. Structure and also this is for better efficiency.

• 1:45:44Speaker 00

  If we can decompose a network into different motifs, we only need to restore those motifs. And how this motifs connected to each other, we don't need to, you know, store the motifs. Store the same motifs for thousands of times. So that's also for, you know, uh you know, more efficient storage and more efficient, you know, inference. For uh especially for large scale models, if we can downsize large scale models in this way. Decomposing them into motifs, that would be awesome.

• 1:46:21Speaker 02

  And the inference time efficiency gains, at least on these like toy models are pretty huge, right? I mean, if if you you run the fully connected version versus the. Like say you take the the sparse version and actually prune it and you just only have the fewer weights at the end. Presumably, it's like under ten percent the compute cost at inference time.

• 1:46:43Speaker 00

  So if we explicitly sparsify the network after training, the inference time would be very much reduced. Like in the modular addition example, like there are three parallel modules emerge from training. We don't know prior that there are three parallel modules, but they just emerge. With that, we can make them parallelizable. For example, we can run these three modules in parallel and then use one machine in principle and then use one machine to aggregate the whole results. So if this kind of parallelization can also generalize to language models, that would be fascinating because because then the inference time can be saved.

• 1:47:29Speaker 00

  You can leverage the like the the all sorts of multi-fredding or parallelization techniques we have at hand to make sure that the you know, the inference, the inference time and inference memory can be saved by a large margin.

• 1:47:50Speaker 02

  That seems extremely compelling, actually. You so Max Tegmark is your advisor. He is not, you know, for most of his career has not been in AI. Did you, did he already like make the switch into AI and then you joined the group? You know, after that switch or did you? Tell, tell me the story of like working with somebody who's obviously got some special ability, but is new to AI.

• 1:48:13Speaker 00

  I, I, I cannot recall the exact year Max switched to AI, but I joined the group after Max switched to AI. And when I joined the group, we are still focusing on AI for physics. Because, because Max used to be a cosmologist. I had strong background in physics, has a the position in Department of Physics. So and, and also I got admitted to the Department of Physics at MIT. But uh like I did research on AI for physics with Max uh for two or three years, uh my, yeah, I'm in my third year uh right now. And then Max

• 1:48:53Speaker 00

  sort of decided to to switch entirely the focus to AI safety, Mac and Turf specifically. Because like trained as physicists, we think that we need to first fully understand something so that we can, you know, fully control it. We can fully keep AI systems under control. It it did no harm to our hu- to our human beings. Uh but the first step is to basically to understand it.

• 1:49:22Speaker 02

  Yeah, I've been a admirer of his work for a number of years and it's been cool to see that he's been so flexible and you know, adaptable and obviously has, you know, been able to uh attract some bright minds to the group to uh do some exciting work. I'm struck by the fact that from what I understand, just on your website and, you know, briefly looking into your background, you grew up in China, went to university in China and then came here after university. Is that right?

• 1:49:49Speaker 00

  Yeah, so I went to uh T K University uh in the physics department, school of physics. And then Max hired me as his PhD student. So uh yeah, I've been here in the US for like two years.

• 1:50:06Speaker 02

  So does the general, like broader context of sort of US-China tension and the increasingly like center stage that AI has in that broader debate, is that like relevant to you personally? Do you just try to stay out of it or how how do you engage with that subject, if at all?

• 1:50:25Speaker 00

  I think I agree with Max on this point, like uh the existential risk is a common threat to all human beings. So if we see this as kind of extinction risk and no one can survive if we don't collaborate. So with this, uh actually this can bring together US and China uh collaboration on this AI safety thing. I think collaboration is the right way, uh not, you know, you know, opposing each other.

• 1:50:57Speaker 02

  Yeah, totally. I try to highlight any.

• 1:51:01Speaker 02

  Instance of positive US-China relationship and collaboration that I can. So your time here and work on this project is a great example of that. And I'm really appreciative that you've spent so much time walking through it with us. Keep up the great work, whether it's here or whether it's one day back in China, or hopefully you can continue to span the two indefinitely. We certainly need a lot more of this ability to understand. Because as I think you guys put it very well, understanding is usually a precondition for control.

• 1:51:42Speaker 02

  And everybody has a shared interest in figuring out how to control these systems so we can get the best from them and avoid the worst.

• 1:51:51Speaker 00

  Yeah, sure. Thank you so much for the invitation. Yeah, it was great speaking to you and have fun with your family.

• 1:52:00Speaker 02

  Thank you very much. Ziming Liu, thank you for being part of the cognitive revolution.