From Layers to States: A State Space Model Perspective to Deep Neural Network Layer Dynamics

TL;DR

This paper introduces a novel continuous state space model approach to layer aggregation in deep neural networks, leveraging SSMs and S6 to improve feature extraction and performance in vision tasks.

Contribution

It proposes a new layer aggregation method using state space models, specifically the Selective State Space Model Layer Aggregation (S6LA), for very deep neural networks.

Findings

S6LA improves image classification accuracy.

S6LA enhances object detection performance.

The approach effectively models long-range dependencies.

Abstract

The depth of neural networks is a critical factor for their capability, with deeper models often demonstrating superior performance. Motivated by this, significant efforts have been made to enhance layer aggregation - reusing information from previous layers to better extract features at the current layer, to improve the representational power of deep neural networks. However, previous works have primarily addressed this problem from a discrete-state perspective which is not suitable as the number of network layers grows. This paper novelly treats the outputs from layers as states of a continuous process and considers leveraging the state space model (SSM) to design the aggregation of layers in very deep neural networks. Moreover, inspired by its advancements in modeling long sequences, the Selective State Space Models (S6) is employed to design a new module called Selective State Space Model Layer Aggregation (S6LA). This module aims to combine traditional CNN or transformer architectures within a sequential framework, enhancing the representational capabilities of state-of-the-art vision networks. Extensive experiments show that S6LA delivers substantial improvements in both image classification and detection tasks, highlighting the potential of integrating SSMs with contemporary deep learning techniques.