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

TL;DR

This paper presents Brain-Inspired Modular Training (BIMT), a novel method that enhances neural network interpretability by embedding neurons in a geometric space and encouraging modular structures, revealing insights into the network's decision-making process.

Contribution

BIMT introduces a brain-inspired approach to train more modular and interpretable neural networks, enabling direct visualization of modules and understanding of their functions.

Findings

BIMT discovers useful modular structures in neural networks for various tasks.

Modules reveal compositional and mathematical structures in data.

The method enhances interpretability beyond traditional techniques.

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.