End-to-end Feature Alignment: A Simple CNN with Intrinsic Class Attribution

TL;DR

This paper introduces FA-CNN, a CNN architecture with intrinsic class attribution via end-to-end feature alignment, improving interpretability and providing theoretical insights into feature evolution.

Contribution

The paper proposes new order-preserving layers and a theoretical framework linking penultimate features to Grad-CAM saliency maps, enhancing interpretability.

Findings

FA-CNN's penultimate features are identical to Grad-CAM saliency maps.

The model maintains end-to-end feature alignment from input to logits.

FA-CNN performs well on benchmark image classification datasets.

Abstract

We present Feature-Align CNN (FA-CNN), a prototype CNN architecture with intrinsic class attribution through end-to-end feature alignment. Our intuition is that the use of unordered operations such as Linear and Conv2D layers cause unnecessary shuffling and mixing of semantic concepts, thereby making raw feature maps difficult to understand. We introduce two new order preserving layers, the dampened skip connection, and the global average pooling classifier head. These layers force the model to maintain an end-to-end feature alignment from the raw input pixels all the way to final class logits. This end-to-end alignment enhances the interpretability of the model by allowing the raw feature maps to intrinsically exhibit class attribution. We prove theoretically that FA-CNN penultimate feature maps are identical to Grad-CAM saliency maps. Moreover, we prove that these feature maps slowly morph layer-by-layer over network depth, showing the evolution of features through network depth toward penultimate class activations. FA-CNN performs well on benchmark image classification datasets. Moreover, we compare the averaged FA-CNN raw feature maps against Grad-CAM and permutation methods in a percent pixels removed interpretability task. We conclude this work with a discussion and future, including limitations and extensions toward hybrid models.