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

2025-11-16

AI Summary

Deep neural networks typically rely solely on representations from the final hidden layer for prediction, thereby neglecting multi-granular semantic information—ranging from low-level patterns to high-level abstractions—encoded in intermediate layers. To address this, we propose LAYA, an input-conditioned inter-layer attention aggregation mechanism that dynamically weights and fuses representations across all hidden layers for adaptive multi-scale feature integration. Its key contributions are: (1) input-dependent layer-wise attention weights that enhance representation adaptivity; (2) direct generation of interpretable layer attribution scores during forward inference—eliminating the need for post-hoc explanation methods; and (3) architecture-agnostic design, compatible with mainstream vision and language models. Experiments across multiple benchmarks demonstrate that LAYA achieves an average accuracy gain of nearly 1 percentage point over baselines, while maintaining or exceeding baseline performance and providing intuitive, quantitative visualizations of layer-wise contribution.

Technology Category

• Computer Vision: Large Vision Models
• Machine Learning: Deep Neural Architectures and Foundation Models
• Natural Language Processing: Interpretability, Analysis, and Evaluation of NLP Models

Application Category

• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• Search and Retrieval-Augmented AI: Retrieval-Augmented Generation (RAG) and multi-modal RAG

Abstract

Deep neural networks typically rely on the representation produced by their final hidden layer to make predictions, implicitly assuming that this single vector fully captures the semantics encoded across all preceding transformations. However, intermediate layers contain rich and complementary information -- ranging from low-level patterns to high-level abstractions -- that is often discarded when the decision head depends solely on the last representation. This paper revisits the role of the output layer and introduces LAYA (Layer-wise Attention Aggregator), a novel output head that dynamically aggregates internal representations through attention. Instead of projecting only the deepest embedding, LAYA learns input-conditioned attention weights over layer-wise features, yielding an interpretable and architecture-agnostic mechanism for synthesizing predictions. Experiments on vision and language benchmarks show that LAYA consistently matches or improves the performance of standard output heads, with relative gains of up to about one percentage point in accuracy, while providing explicit layer-attribution scores that reveal how different abstraction levels contribute to each decision. Crucially, these interpretability signals emerge directly from the model's computation, without any external post hoc explanations. The code to reproduce LAYA is publicly available at: https://github.com/gvessio/LAYA.

Problem

• Aggregates intermediate layer features using attention mechanisms
• Enhances model interpretability through layer attribution scores
• Improves prediction accuracy while maintaining architecture independence

Innovation

• Layer-wise attention aggregation for feature synthesis
• Dynamic input-conditioned attention over layer features
• Interpretable architecture-agnostic prediction mechanism