Understanding Task Representations in Neural Networks via Bayesian Ablation

Andrew Joohun Nam, Declan Iain Campbell, Thomas L. Griffiths, Jonathan D. Cohen, Sarah-Jane Leslie

Keywords: representation learning, interpretability, neural networks, Bayesian inference

TL;DR: We introduce a Ablation Mask Distributions, a probabilistic framework for interpreting latent task representations in neural networks.

Abstract: Neural networks are powerful tools for cognitive modeling due to their flexibility and emergent properties. However, interpreting their learned representations remains challenging due to their sub-symbolic semantics. We introduce a novel probabilistic framework for interpreting latent task representations in neural networks. Inspired by Bayesian inference, our approach defines a distribution over representational units to infer their causal contributions to task performance. Using ideas from information theory, we propose a suite of tools and metrics to illuminate key model properties, including representational distributedness, manifold complexity, and polysemanticity.

Submission Number: 11