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

TL;DR

The paper reframes very deep neural networks as continuous-state processes by treating layer outputs as latent states updated via a state-space model. It introduces Selective State Space Model Layer Aggregation (S6LA), a lightweight module that consolidates SSM-based state updates with layer outputs for CNNs and Vision Transformers, enabling effective long-range layer interactions. Through CNN and ViT experiments on ImageNet-1K and COCO, S6LA consistently improves classification, detection, and segmentation performance with modest increases in parameters and FLOPs, outperforming several prior layer-interaction methods. This work bridges statistical state-space modeling with deep learning to advance deep network representational power and lays groundwork for further integration of statistical models into computer vision architectures.

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.

Figures (9)

• Figure 1: Schematic diagram of a Network with Selective State Space Model Layer Aggregation.
• Figure 2: The correlation between accuracy and layer size with model PVT-v2.
• Figure 3: Schematic Diagram of a Network with Selective State Space Model Layer Aggregation. The green arrow represents the hidden state connection, while the grey arrow indicates communication between layers. The updated latent layer is derived from the previous latent layer $h^{t-1}$ and the last input layer $\boldsymbol{X}^{t}$. The output of the $t$-th layer is generated from the input $\boldsymbol{X}^{t-1}$ and the latent layer $h^{t}$.
• Figure 4: Detailed Operations in the S6LA Module with Convolutional Neural Network. The green arrow represents the hidden state connection, while the grey arrow indicates communication between layers. The temporary output $O^t$ is the concatenation of the latent layer $h^{t-1}$ and the input layer $\boldsymbol{X}^t$. After passing through the convolutional layers, the updates for the latent layer $h^t$ and the input layer $\boldsymbol{X}^{t+1}$ are derived from the last layers.
• Figure 5: The GFLOPs induced by our method with respect to input resolution with backbone PVT-v2-b1.