Charge Neural Nets by Connection Length, and the Black Box Splits Into Modules You Can See

Nathan Labenz introduces Ziming Liu, a physics PhD student at MIT and IAIFI in Max Tegmark's group, and first author of *Seeing is Believing: Brain-Inspired Modular Training for Mechanistic Interpretability*. Labenz treats mechanistic interpretability as one of the most promising paths to AI safety, because the black-box problem, in its simplest form, is that we do not know why AIs do what they do. The paper stood out to him for three reasons: it makes networks easier to reverse-engineer, the method is conceptually simple, and the figures communicate the idea in seconds. He tells listeners to spend a few minutes on the GitHub animations before they keep listening, because this work is built to be seen.

The core trick is a small change to the loss, inspired by the physical cost of long-range connections in animal brains. Labenz suspects that, more than any other idea he has seen so far in 2023, it left researchers asking why they had not thought of it. Co-authors are Eric Gan and Tegmark, the cosmologist who wrote *Life 3.0* and has lately turned the group toward AI safety. Watch it on YouTube if you can, Labenz says; the figures are simple enough to hold in mind for the rest of the conversation.

Liu is blunt about the hardware. Every result in the paper was trained on his MacBook M1, with no GPU. Modular addition takes about 20 minutes; the symbolic-regression examples finish in less than a minute; the transformer example, the slowest, runs between one and two hours. That is the entire experimental budget for a method aimed at one of the hardest problems in AI safety. The gap between the laptop and the ambition is the first thing he wants people to notice.

Liu starts from a hope shared across interpretability research: if a deep net can be decomposed into modules, you can divide and conquer. Human brains already look like that, with different regions carrying different functions. Biological networks had an evolutionary reason to become modular, he argues, because modular brains are more energy-efficient and therefore had a selective advantage. Artificial networks have no such incentive. The training loss only asks whether predictions are accurate, not which wires are long.

His philosophy collapses to one line: "If we want neural networks to be interpretable, we should in training explicitly encourage them to be interpretable." The research question is then practical: what training techniques induce modularity in otherwise non-modular nets? Biology's lesson, he says, is that locality gives rise to modularity. Brains live in three-dimensional space, where distance is real; creating neurons, maintaining them, and transmitting signals all cost energy, so local connections win.

Labenz restates the contrast. A skull has finite volume, long-range cables clog space and are fragile to damage, and most activity ends up local, with fairly sharp regional boundaries. Standard nets have layers but no spatial meaning inside a layer: two activations with nearby indices are no more related, during training, than two at opposite ends. Whatever connections form at initialization tend to be the ones that get refined, which is why the graphs look like spaghetti.

Liu puts the same point in mathematical language. Ordinary nets are treated as topological objects: you care whether two nodes share an edge, not how close they are. BIMT's move is to embed the network in a geometric space—the 3D Euclidean space we live in, or a 2D plane for the toy figures—where distances can be defined. Then the loss adds a penalty proportional to connection length, charging long wires more than short ones, as evolution would. That, he says, is the whole idea.

The implementation is almost embarrassingly small. To sparsify, they use L1 regularization, which already needs one hyperparameter for penalty strength. BIMT multiplies that penalty by distance, adding a second hyperparameter: how strongly to encourage locality. Set locality strength to zero and you recover vanilla L1; set it high and the net is pushed toward the shortest surviving wires. Every connection has a nonzero minimum distance, so even the most local weight still pays a base fee of one, with more diagonal connections paying two or three.

L2, the usual cousin, shrinks overall weight norm without reliably driving any single parameter to zero, which makes pruning harder. L1 is more aggressive but can bounce around zero; L2 converges steadily if the step size is small enough. Liu's engineering advice is to try both if you actually care about pruning. Prediction accuracy keeps some weights alive; everything that does not earn its keep is gradually pressed to zero. Labenz restates the picture as a three-way balance: all weights trend toward zero, long ones especially, while the data holds a few of them up.

The aesthetic jump comes from neuron swapping, and it started as an anecdote. Without swaps, Liu's first modular-addition nets looked like a 3D double helix projected onto 2D, crossing and crossing back. He wanted to make neuron coordinates trainable; Tegmark told him to try something simpler. The simple thing is to pick two neurons in the same layer, swap them along with their incoming and outgoing weights, and keep any swap that shortens the locality loss. The function computed by the net is identical at that step—the computational graph does not change—but the geometry does. They score neurons by importance and only try swapping the top ten or so, every 200 steps or continuously.

Swapping is not only cosmetic. Stuck in a helix, the net carries a large length penalty and has to sacrifice prediction loss to afford it; a clever swap drops the penalty, and the same penalty strength then buys a better fit. Liu has not mapped a phase diagram: he set the swap frequency to 200 or to infinity and left the rest for later. He hopes swaps still help at the end of training, because he wants to fine-tune pretrained language models he cannot afford to retrain from scratch. Labenz's lottery-ticket intuition points the other way, and Liu admits that early on, when the net is choosing which basin to enter, swapping may matter most. That, he says, is a conjecture to test on language models.

The first experiment is symbolic regression on four numeric inputs and two outputs: one target is x2 squared plus sin(π x4), the other is (x1 + x3) squared. The functions were chosen so each output depends on a disjoint pair of inputs, and Liu deliberately ordered the inputs 1–2–3–4 so the first output needed 1 and 3, forcing a cross that swapping would have to uncross. The trained net splits into two independent pathways and regroups the inputs automatically. Final loss is about 7.4×10⁻³, with predicted-versus-true scatter sitting on a line at R² around 0.99. Invisible edges are not a plotting trick: setting the small weights to zero leaves the output unchanged.

Two-moons classification is a 2D point cloud that a linear separator cannot split. The sparse solution surprised Liu: for binary classification only the difference of logits matters, so the net learned to pin one class's logit at zero and train a single signed score for the other. Training shows a phase transition from a messy dense graph to a symmetric two-class skeleton, then to this cheaper asymmetric scheme. The final graph has about seven weights and six active neurons, sparse enough that the appendix can intervene on each one. Change the random seed and the picture looks entirely different—a sign, he says, that the task has no stable internal modules, just one blob that is itself the module.

Modular addition is the community's favorite test, studied by Neel Nanda and colleagues. Embeddings still use a Fourier basis, matching prior work, but the internal algorithm is not Nanda's "clock," which rides the circularity of modular arithmetic. Liu's group found a "pizza" algorithm that slices the disk into sectors and asks which slice the sum lands on; which algorithm you get depends on architecture and hyperparameters, with phase transitions either way. BIMT grows three parallel, imperfect pizzas: each fails on some examples, and a last linear layer aggregates them so every input is covered. A follow-up with Zichen Zhong, Jacob Andreas, and Tegmark notes that Nanda's original metric cannot tell clock from pizza—each metric is a coarse-graining, and if you drop too much information, different algorithms look the same.

On the S4 permutation group, Nanda's probing had found that nets learn the group's nine-dimensional representation. After BIMT, that representation sits on a privileged basis: exactly nine active neurons, plus a "magic" 22nd neuron for the sign, and the rest can be read from Cayley graphs. You do not need to know the task to know where to look, Liu says—you look at the graph. Labenz likens the penalty to shrink-wrapping conceptual reality until the extra air pops out; Liu says the training videos on GitHub are exactly that inward pressure. MNIST embedded in 3D adds a pattern nobody fully understands: first-to-second-layer weights are almost all negative, second-to-output almost all positive, anti-matching then matching, stable across one, two, or three hidden layers.

Labenz, staring at the symbolic-regression graph, still cannot see how five active neurons compute a square-plus-sine. Liu was shocked too, then reached for a physicist's analogy: the unreasonable effectiveness of smooth activation functions. He uses SiLU (and mentions SELU), everywhere-differentiable curves, and in approximation theory a smooth target approximated by smooth bases has error that falls exponentially as you add modes. He has not tried ReLU. Asked whether he assumed it would fail, he admits an "unreasonable craze" for SiLU; papers exist showing that two or three such neurons can construct a quadratic.

Construction is not what training finds. The chance of initializing a SiLU neuron at its trough, where a Taylor expansion is a parabola, is tiny. In practice two neurons keep first-order terms, and later layers weight them so the linear parts cancel and the second-order parts survive. Nobody taught the net that cancellation. "In some sense," Liu says, "neural networks are more clever than myself." He does not expect the fit to generalize outside the training interval of roughly [−1, 1]; like an effective theory in physics, it is valid on its own scale, and interpolation can still be worth understanding.

Reproducibility has a sharp edge. For modular addition he has to use 64-bit doubles; at 32-bit, truncation noise is large enough to knock training into another basin. Permutation symmetry creates exponentially many equivalent basins, and algorithmic tasks with learned embeddings inherit the same problem—he expects language models will too. Image classification and symbolic regression do not. If only some seeds find structure, his advice is to train a hundred models in parallel and pick the most interpretable one. Git Re-Basin-style alignment also fails when two nets implement different mechanisms for the same task, a point made by work on mechanistic mode connectivity and visible in his own modular-addition runs.

Why did nobody do this years ago? Related work existed. About a decade earlier, Jeff Clune and collaborators studied connection costs, evolving tiny networks to recover biological modularity, with later extensions to gradient methods. Those papers were aimed at biology, and gradient tools were less mature. BIMT uses current training machinery and aims at mechanistic interpretability. Liu will not call the idea wholly novel: we stand on earlier shoulders, and "only really valuable things can be rediscovered over and over again." Labenz adds that everything is a remix; what feels new is how densely the result is packaged in an animation.

Training cost is modest enough to play with. Modular addition, at 20,000 steps, is about 20 minutes on the M1; symbolic regression is under a minute; the transformer example is one to two hours. Liu tells people to start with symbolic regression, change the code, and watch the shape of the net move. Attention is treated as three bias-free linear maps—query, key, and value—stacked along a third dimension and penalized like MLP weights. Softmax has no parameters. He says there is no extra scaling overhead versus a normal transformer.

Making that picture actually work took several hacks. Swaps may only happen inside a head, or as whole-head swaps. Residual streams lock permutations across layers: swap once and you must swap every residual copy, a global constraint that can leave non-local wires stuck. LayerNorm is dropped entirely, because it rescales even tiny activations into a full range and fights sparsity; Anthropic has also called LayerNorm a nuisance, though they did not remove it. BIMT, he says, is a high-level idea—embed any architecture in a geometric space. The transformer results are a prototype with a long to-do list.

Other limits are conceptual. Two-moons and MNIST may be too simple to split—linear classifiers already exceed 90 percent accuracy on MNIST—so the whole net is one module. He has thought about adding graph-theoretic modularity losses on top of the length penalty and dislikes that as an engineering patch; he wants something more elegant. Connection cost also does not induce module reuse, the "one block copied a thousand times" pattern; that may need hypernetworks or compositional networks, which he treats as orthogonal. Clune's older biology work suggested hierarchy and continual learning can come with connection cost too, another direction still open.

The experiments he actually wants next are language models, including TinyStories, and AI for science, where people already care about explanations. TinyStories connects to a quantization model of neural scaling from his group, led by Eric Michaud: a task such as language decomposes into "quanta," subtasks ordered by frequency, and a net of a given capacity learns them in that order. Fluent three-year-old English might be the first ten quanta; role-playing a physicist lives in the tail. He also hopes BIMT on small language models will give token embeddings a privileged basis, so a feature such as color sits on one dimension you can enumerate rather than a direction you have to search among infinitely many.

On scaling, Liu's only immediate worry is visualization: a large net may be hard to draw. He does not see a principled blocker beyond that. If modules can be stored as motifs plus a wiring diagram, you would not keep a thousand copies of the same block, which would help both memory and inference. In modular addition the three parallel modules can already be run separately and aggregated; if that pattern reaches language models, wall-clock time and memory could drop by a large margin. Labenz notes that a pruned sparse net is already a large inference win on the toys, likely an order of magnitude cheaper than the dense graph.

Labenz asks how a cosmologist's student ended up here. Liu joined after Tegmark had already moved into AI, first spending two or three years on AI for physics; he is in his third year. Tegmark then shifted the group entirely toward AI safety and mechanistic interpretability. The physicist's creed, as Liu restates it, is that you need to fully understand something before you can fully control it, and control is how you keep systems from harming people. Understanding, in this frame, is not a luxury. It is the first step.

Liu did his undergraduate physics at Tsinghua, then came to MIT; he has been in the United States about two years. On US–China tension and AI's place in it, he sides with Tegmark: existential risk is a common threat to all human beings, and "no one can survive if you don't collaborate." That, he thinks, is a reason to put the two countries at the same table on AI safety rather than in opposition. Labenz treats the collaboration as a small instance of the kind of relationship he tries to highlight. Both of them end on the same precondition: understanding before control.

For anyone who trains models, the practical path is short. The code is on GitHub; start with symbolic regression on CPU and watch the graph change when you tweak locality strength. If you train your own nets, a distance-weighted L1 term is one extra hyperparameter that can be set to zero. If you want to interpret a trained model, sparsify and modularize it first, then look: BIMT's claim is that you no longer have to probe blindly. And if the task has heavy permutation symmetry, turn on 64-bit before you chase a random seed.