LAYA: Layer-wise Attention Aggregation for Interpretable Depth-Aware Neural Networks
This addresses the issue of interpretability and information loss in neural networks for researchers and practitioners, though it is incremental as it builds on existing attention mechanisms.
The paper tackles the problem of deep neural networks discarding intermediate layer information by introducing LAYA, a layer-wise attention aggregation output head that dynamically combines internal representations, resulting in performance gains of up to about one percentage point in accuracy on vision and language benchmarks while providing interpretable layer-attribution scores.
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.