OMENN: One Matrix to Explain Neural Networks

TL;DR

OMENN tackles the opacity of deep networks by reformulating a trained model for a given input as a single input-dependent affine transformation, yielding an explanation matrix $\mathbf{C} = \mathbf{C_w} \odot \mathbf{X} + \mathbf{C_b}$ whose sum equals the model output. Grounded in the dynamic linearity property, the method derives $\mathbf{C_w}$ and $\mathbf{C_b}$ by fusing layer weights and biases and augmenting the input with a ones channel, enabling exact, locally faithful attributions for modern architectures including ViTs and CNNs. The authors provide theoretical justification and validate OMENN on the FunnyBirds benchmark and the Quantus faithfulness metric, showing state-of-the-art or competitive performance across backbones and tasks, with qualitative analyses illustrating pixel-level contributions aligned with class-specific regions. The work has practical implications for explanations, potential benefits for knowledge distillation and continual learning, and highlights remaining limitations related to non-affine operations and user-centric visualization aspects.

Abstract

Deep Learning (DL) models are often black boxes, making their decision-making processes difficult to interpret. This lack of transparency has driven advancements in eXplainable Artificial Intelligence (XAI), a field dedicated to clarifying the reasoning behind DL model predictions. Among these, attribution-based methods such as LRP and GradCAM are widely used, though they rely on approximations that can be imprecise. To address these limitations, we introduce One Matrix to Explain Neural Networks (OMENN), a novel post-hoc method that represents a neural network as a single, interpretable matrix for each specific input. This matrix is constructed through a series of linear transformations that represent the processing of the input by each successive layer in the neural network. As a result, OMENN provides locally precise, attribution-based explanations of the input across various modern models, including ViTs and CNNs. We present a theoretical analysis of OMENN based on dynamic linearity property and validate its effectiveness with extensive tests on two XAI benchmarks, demonstrating that OMENN is competitive with state-of-the-art methods.

Figures (9)

• Figure 1: This figure illustrates OMENN, a novel XAI method that provides locally exact explanations by transforming the neural network's processing of an input into a linear operation comprising two matrices: one representing the model's weights and the other representing its biases. When combined with the input, these matrices yield a single, interpretable matrix that forms the explanation. Note that the sum of the values within this explanation matrix directly corresponds to the logit value of the class being explained.
• Figure 2: OMENN method provides an explanation in the form of a single matrix, $\bold{C}$. For this purpose, it represents a neural network as a single, input-dependent affine transformation using dynamic linearity. This transformation includes two components: a bias representation denoted as $\bold{C_b}$ and a weight representation $\bold{C_w}$. These components are applied to the input as an affine transform to form the overall explanation matrix $\bold{C} = \bold{C_w}\odot\bold{X} + \bold{C_b}$. Notably, the sum of all elements in $\bold{C}$ matches the score produced by the original neural network, ensuring the explanation is faithful to the network’s output.
• Figure 3: An example highlighting the difference between our OMENN ($\Omega_w(x)$, red) and gradient ($\nabla_x$, green) methods for GELU. One can observe that gradient incorrectly returns the same contribution ($0$) for both considered points ($x_0=-4$ and $x_1=-0.75$). While our OMENN accurately provides higher contribution for $x_1$, which when multiplied by $x_1$ is equal to the output of GELU ($-0.17$), in contrast to output obtained from the gradient ($0$).
• Figure 4: Results on the FunnyBirds benchmark for the ViT-B/16 backbone show that OMENN achieves the highest score among various explainability methods. Notably, Chefer LRP ranks as the second-best approach, with only a marginal difference from OMENN, indicating that these methods are closely competitive.
• Figure 5: Examples of explanations generated by various XAI methods, including our OMENN approach, for ViT-B/16 on ImageNet. While OMENN’s contributions may appear noisy, they precisely reveal how each pixel influences the final logit value.