Seeing is Believing: Brain-Inspired Modular Training for Mechanistic Interpretability

TL;DR

This work introduces Brain-Inspired Modular Training (BIMT), a method that injects modularity into neural networks by embedding neurons in a geometric space, enforcing locality via distance-based regularization, and allowing neuron swapping to improve locality. BIMT is demonstrated across symbolic regression, simple classification, and algorithmic tasks, revealing clear modular and tree-like structures, and is extended to transformers and image-like tensor data. The results show that BIMT often yields interpretable, modular circuits with only modest performance penalties, and reveal group-theoretical and compositional structures in learned representations. These findings suggest BIMT as a practical tool for mechanistic interpretability and for visualizing how modules cooperate to solve complex tasks, with potential extensions to larger models.

Abstract

We introduce Brain-Inspired Modular Training (BIMT), a method for making neural networks more modular and interpretable. Inspired by brains, BIMT embeds neurons in a geometric space and augments the loss function with a cost proportional to the length of each neuron connection. We demonstrate that BIMT discovers useful modular neural networks for many simple tasks, revealing compositional structures in symbolic formulas, interpretable decision boundaries and features for classification, and mathematical structure in algorithmic datasets. The ability to directly see modules with the naked eye can complement current mechanistic interpretability strategies such as probes, interventions or staring at all weights.

Figures (19)

• Figure 1: Top: Brain-inspired modular training (BIMT) contains three ingredients: (1) embedding neurons into a geometric space (e.g., 2D Euclidean space); (2) training with regularization which penalizes non-local weights more; (3) swapping neurons during training to further enhance locality. Bottom: Zoo of modular networks obtained via BIMT (see experiments for details).
• Figure 2: The connectivity graphs of neural networks when trained with different techniques for a regression problem (blue/red denote positive/negative weights). Our proposed BIMT = $L_1$ regularization (not novel) + local regularization (novel) + swap (novel). BIMT finds the simplest circuit (e) which clearly contains two parallel modules, with a small sacrifice in test loss compared to vanilla (a), but with lower loss than for mere $L_1$ regularization (b). Note that swapping aims to reduce the local connection cost, so all of (c)(d)(e) encourage locality.
• Figure 3: The connectivity graphs of neural networks trained with BIMT to regress symbolic formulas (blue/red lines stand for positive/negative weights). For symbolic formulas with modular properties, e.g., independence, shared features or compositionality, the connectivity graphs display modular structures revealing these properties.
• Figure 4: Top: Evolution of network structures trained with BIMT on the two moon dataset. Bottom: Evolution of decision boundaries.
• Figure 5: MLP trained with BIMT for modular addition. Left: the final connectivity graph is tree-like, demonstrating three parallel modules (voters); middle: the representations of each module in the input layer; right: ablation results, which imply a voting mechanism. The input layer contains embeddings of two tokens, which overlap each other but are drawn to be vertically separated.