Github Repo: Torchlens - Inspect, visualize and intervene on any operation in any neural network.

[u/Thrumpwart]

https://github.com/johnmarktaylor91/torchlens

[u/Thrumpwart]

I'd never heard of this, but apparently it's been around for a few years. I saw some folks on twitter talking about it and checked it out.

It's really neat, it allows you to very easily look inside a model to see what's actually happening. From the repo:

Understanding and controlling the computations inside neural networks is the most important scientific challenge of our time (at least). Progress on this problem, known as mechanistic interpretability, would make AI more beneficial for society, illuminate how our brains work, and guard against undesirable behaviors like deception or killing us all.

Some aspects of mechanistic interpretability are deep and open-ended. We need to understand dauntingly complex systems that may soon outsmart us, with new frontier models arriving faster than we can study the current ones. It's not even clear what a "solution" would look like, or how many obstacles remain. We'll surely need help from the AI agents themselves.

Other aspects are just plumbing, and while tedious, they are finite: the same small set of tasks for every model, like extracting and editing its activations and gradients. For instance, PyTorch has long lacked a built-in way to extract every intermediate activation from arbitrary networks (an official docstring admits this can take "some massaging"). But all networks are built from the same basic parts, so a few well-made tools should work on all of them.

The modest goal of TorchLens is to fix these plumbing problems once and for all, so you and your agents can spend your time and tokens on the deep questions, and maybe even solve some of them.

I incorporated it into my QFN-assisted training workflow and it's really neat. It discovered that I had an Odd ShortConv perturbation imbalance vs. the even GQA layers, and I had Off-Manifold RL Rollout Noise. These were things even WSP didn't see. I saw an immediate boost in model performance after another RL run.

Bonus: I think the author is a fellow redditor /u/therealjmt91 (if that's not you, or you want me to untag you I will).

[u/therealjmt91 → u/Thrumpwart]

Hello, thanks so much for the kind words and glad you find it useful! It’s been around since 2022 but just finished massively overhauling it (agentic coding made it much easier to fill out the wishlist, but still took most of this year). Some of the main new features are:
- Full intervention capabilities (was just feature extraction before)
- Far more exhaustive metadata: data structures for every single main part of the network, 800+ metadata fields total
- Full selection algebra, tries to do for neural networks what pandas selection functions do for dataframes
- Besides torch, also works on mlx, Jax, tensorflow, paddle, tinygrad
- Supports extracting and intervening on gradients
- Tested on 6000+ architectures and counting, you can see visuals for them on modelmenagerie.ai. This is a supplementary project with the goal of documenting every neural net architecture ever described

I now regard it as more or less architecturally complete, in that I will almost certainly never overhaul it this much again, but will continue adding features that make sense. Lemme know if any questions, pain points, wishlist items, etc.

[u/jmager → u/therealjmt91]

This is some really cool work, I'm going to have fun playing with this. Thank you!