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

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. Beyond improving feature aggregation, the learned attention coefficients provide intrinsic layer-attribution scores that explicitly quantify the contribution of each representation to the final decision, without requiring external post-hoc explanation methods. Experiments on image classification datasets show that LAYA achieves competitive predictive performance while producing meaningful depth-aware explanations. Furthermore, quantitative and qualitative interpretability analyses demonstrate that LAYA's attention scores closely reflect the actual contribution of individual layers, revealing structured, task-dependent patterns of depth utilization while providing intuitive explanations of how different abstraction levels contribute to each prediction.

Machine Learning (cs.LG)

arXiv:2511.12723 [cs.LG]

(or arXiv:2511.12723v2 [cs.LG] for this version)

https://doi.org/10.48550/arXiv.2511.12723 Focus to learn more arXiv-issued DOI via DataCite

arXivLabs is a framework that allows collaborators to develop and share new arXiv features directly on our website.

Both individuals and organizations that work with arXivLabs have embraced and accepted our values of openness, community, excellence, and user data privacy. arXiv is committed to these values and only works with partners that adhere to them.

Have an idea for a project that will add value for arXiv's community? Learn more about arXivLabs.