# Identifying Sub-networks in Neural Networks via Functionally Similar Representations

arXiv: https://arxiv.org/abs/2410.16484 · HTML: https://mlsanity.com/paper/2410.16484/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 40 of its 57 references that are in the map.
- First posted on arXiv 21 Oct 2024.

## Builds on (most cited first, 30 of 40)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [GLUE: A Multi-Task Benchmark and Analysis Platform for Natural Language Understanding](https://mlsanity.com/paper/1804.07461/) (2018)
- [Locating and Editing Factual Associations in GPT](https://mlsanity.com/paper/2202.05262/) (2022)
- [Sinkhorn Distances: Lightspeed Computation of Optimal Transportation Distances](https://mlsanity.com/paper/1306.0895/) (2013)
- [Similarity of Neural Network Representations Revisited](https://mlsanity.com/paper/1905.00414/) (2019)
- [Interpretability in the Wild: a Circuit for Indirect Object Identification in GPT-2 small](https://mlsanity.com/paper/2211.00593/) (2022)
- [Progress measures for grokking via mechanistic interpretability](https://mlsanity.com/paper/2301.05217/) (2023)
- [Towards Automated Circuit Discovery for Mechanistic Interpretability](https://mlsanity.com/paper/2304.14997/) (2023)
- [Dissecting Recall of Factual Associations in Auto-Regressive Language Models](https://mlsanity.com/paper/2304.14767/) (2023)
- [Insights on representational similarity in neural networks with canonical correlation](https://mlsanity.com/paper/1806.05759/) (2018)
- [TRAK: Attributing Model Behavior at Scale](https://mlsanity.com/paper/2303.14186/) (2023)
- [Convergent Learning: Do different neural networks learn the same representations?](https://mlsanity.com/paper/1511.07543/) (2015)
- [Model Fusion via Optimal Transport](https://mlsanity.com/paper/1910.05653/) (2019)
- [Linearity of Relation Decoding in Transformer Language Models](https://mlsanity.com/paper/2308.09124/) (2023)
- [ZipIt! Merging Models from Different Tasks without Training](https://mlsanity.com/paper/2305.03053/) (2023)
- [Optimal Transport for structured data with application on graphs](https://mlsanity.com/paper/1805.09114/) (2018)
- [The Geometry of Categorical and Hierarchical Concepts in Large Language Models](https://mlsanity.com/paper/2406.01506/) (2024)
- [The Clock and the Pizza: Two Stories in Mechanistic Explanation of Neural Networks](https://mlsanity.com/paper/2306.17844/) (2023)
- [Similarity of Neural Network Models: A Survey of Functional and Representational Measures](https://mlsanity.com/paper/2305.06329/) (2023)
- [AtP*: An efficient and scalable method for localizing LLM behaviour to components](https://mlsanity.com/paper/2403.00745/) (2024)
- [Representation Similarity Analysis for Efficient Task taxonomy & Transfer Learning](https://mlsanity.com/paper/1904.11740/) (2019)
- [RAVEL: Evaluating Interpretability Methods on Disentangling Language Model Representations](https://mlsanity.com/paper/2402.17700/) (2024)
- [Towards Vision-Language Mechanistic Interpretability: A Causal Tracing Tool for BLIP](https://mlsanity.com/paper/2308.14179/) (2023)
- [Towards a Mechanistic Interpretation of Multi-Step Reasoning Capabilities of Language Models](https://mlsanity.com/paper/2310.14491/) (2023)
- [Differentiable Subset Pruning of Transformer Heads](https://mlsanity.com/paper/2108.04657/) (2021)
- [Competition of Mechanisms: Tracing How Language Models Handle Facts and Counterfactuals](https://mlsanity.com/paper/2402.11655/) (2024)
- [Learning to grok: Emergence of in-context learning and skill composition in modular arithmetic tasks](https://mlsanity.com/paper/2406.02550/) (2024)
- [The Shape of Data: Intrinsic Distance for Data Distributions](https://mlsanity.com/paper/1905.11141/) (2019)

## Abstract (arXiv)

Providing human-understandable insights into the inner workings of neural networks is an important step toward achieving more explainable and trustworthy AI. Existing approaches to such mechanistic interpretability typically require substantial prior knowledge and manual effort, with strategies tailored to specific tasks. In this work, we take a step toward automating the understanding of the network by investigating the existence of distinct sub-networks. Specifically, we explore a novel automated and task-agnostic approach based on the notion of functionally similar representations within neural networks to identify similar and dissimilar layers, revealing potential sub-networks. We achieve this by proposing, for the first time to our knowledge, the use of Gromov-Wasserstein distance, which overcomes challenges posed by varying distributions and dimensionalities across intermediate representations, issues that complicate direct layer to layer comparisons. On algebraic, language, and vision tasks, we observe the emergence of sub-groups within neural network layers corresponding to functional abstractions. Through downstream applications of model compression and fine-tuning, we show the proposed approach offers meaningful insights into the behavior of neural networks with minimal human and computational cost.