LAYA: Layer-wise Attention Aggregation for Interpretable Depth-Aware Neural Networks

Technical Summary: LAYA (Layer-wise Attention Aggregation)

Problem Statement
Deep neural networks (DNNs) typically rely exclusively on the representation produced by the final hidden layer ($h_L$) to generate predictions. This design implicitly assumes that the deepest embedding fully encapsulates all relevant semantic information extracted through preceding transformations. However, empirical evidence suggests that intermediate layers contain rich, complementary information ranging from low-level patterns to high-level abstractions. By discarding these intermediate representations at the output stage, standard architectures not only potentially lose discriminative information but also obscure the contribution of different abstraction levels to the final decision. Furthermore, existing methods for interpreting these contributions are largely post-hoc, requiring external analysis after inference, rather than being intrinsic to the model's prediction mechanism.

Methodology
The paper introduces LAYA (Layer-wise Attention Aggregator), a lightweight, architecture-agnostic output module designed to replace the standard classifier head. Instead of projecting only the final embedding, LAYA dynamically aggregates internal representations from all hidden layers $\{h_i\}_{i=1}^L$.

The mechanism operates as follows:

1. Projection: Each hidden representation $h_i$ is mapped into a shared latent space of dimension $d^*$ via a learnable adapter $g_i(\cdot)$, resulting in $z_i = g_i(h_i)$.
2. Scoring: These projected representations are optionally transformed by a shared mapping $\psi(\cdot)$ (either identity or a lightweight MLP) and concatenated. A scoring MLP processes this concatenation to produce unnormalized relevance scores $s(x) \in \mathbb{R}^L$.
3. Attention: The scores are converted into input-conditioned attention coefficients $\alpha_i(x)$ using a temperature-scaled softmax:
   $$\alpha_i(x) = \frac{\exp(s_i(x)/\tau)}{\sum_{j=1}^L \exp(s_j(x)/\tau)}$$
   where $\tau$ controls the sharpness of the distribution.
4. Aggregation & Prediction: The final aggregated representation is computed as a weighted sum $h_{agg} = \sum_{i=1}^L \alpha_i(x) z_i$, which is then passed to the task-specific output head for classification.

Crucially, the attention weights $\alpha_i(x)$ are generated during the forward pass as part of the prediction process, serving as intrinsic layer-attribution scores without requiring additional backward computations or external explanation models.

Key Contributions

1. Revisiting the Output Layer: The work highlights the limitation of conventional output heads that ignore intermediate representations and proposes a shift toward depth-aware aggregation.
2. Intrinsic Interpretability: LAYA provides a mechanism where the model's decision process is inherently interpretable. The learned attention coefficients explicitly quantify the contribution of each layer to the final prediction, eliminating the need for post-hoc explanation tools like Integrated Gradients or Grad-CAM for layer-level attribution.
3. Architecture Agnosticism: The module functions as a drop-in replacement for standard output heads, leaving the backbone network (e.g., CNNs, Transformers, MLPs) completely unchanged.
4. Comprehensive Evaluation: The paper presents a rigorous evaluation across three distinct representational hierarchies: a Multi-Layer Perceptron (MLP) on Fashion-MNIST, a Convolutional Neural Network (CNN) on CIFAR-10, and a frozen Vision Transformer (ViT) on an artistic style classification task.

Experimental Results

• Predictive Performance: LAYA achieves competitive, and in some cases superior, predictive performance compared to standard baselines (LastLayer), static concatenation (Concat), and global mixing (ScalarMix).
  • On Fashion-MNIST and CIFAR-10, LAYA achieved the highest mean accuracy among all tested heads.
  • On the Best Artworks dataset, while static concatenation performed best, LAYA significantly outperformed the standard LastLayer baseline, demonstrating that intermediate representations contain valuable information even in frozen pre-trained backbones.
• Faithfulness of Explanations: Quantitative analysis using Faithfulness (Spearman correlation between attribution scores and leave-one-layer-out perturbation effects) and Deletion@1 (confidence drop upon removing the top-ranked layer) shows that LAYA's intrinsic scores are highly faithful.
  • On Fashion-MNIST, LAYA's faithfulness score (0.8575) approached that of gradient-based post-hoc methods (approx. 0.89) while significantly outperforming static global weights.
  • On the Best Artworks dataset, LAYA achieved the highest Deletion@1 score (0.5787), indicating that the layer it identified as most relevant was indeed the most influential for the prediction.
• Interpretability Patterns: Qualitative analysis revealed task-dependent depth utilization. For instance, on CIFAR-10, attention concentrated almost entirely on the final layer, whereas on Best Artworks, attention was sparsely distributed across specific intermediate and late ViT blocks (e.g., blocks 7, 9, 11, 12), varying by artistic style.

Significance and Claims
The paper posits that LAYA offers a principled alternative to the standard "collapse" of hierarchical representations at the output stage. Its primary significance lies not in universally dominating classification accuracy, but in its ability to combine competitive predictive performance with directly inspectable, input-dependent layer attribution.

The author claims that LAYA:

• Provides intrinsic interpretability, making the model's reliance on specific abstraction levels transparent without external tools.
• Reveals structured, task-dependent patterns of depth utilization, showing that different tasks and even different classes within a task rely on different combinations of layer depths.
• Offers a foundation for future applications such as attention-guided model compression (pruning layers with negligible attention), adaptive early-exit mechanisms, and diagnostic analysis of model failures based on attention shifts.

The work concludes that while the focus is on computer vision, the framework is applicable to other sequence architectures, potentially offering new insights into how information is refined across depth in models like Transformers.