Identifying Sub-networks in Neural Networks via Functionally Similar Representations

TL;DR

This paper introduces an automated, task-agnostic method using Gromov-Wasserstein distance to identify sub-networks in neural networks based on functionally similar representations, enhancing interpretability.

Contribution

It pioneers the use of Gromov-Wasserstein distance for comparing neural network layers, enabling automated detection of sub-networks across diverse tasks.

Findings

Sub-networks correspond to functional abstractions in different tasks.

The approach provides meaningful insights with minimal human effort.

Applications include model compression and fine-tuning improvements.

Abstract

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.