MASS: Multi-scale Attention-Augmented SSM

• MASS is a model architecture that combines multi-scale adaptive attention with efficient state space evolution to capture both local details and global context.
• In time series forecasting, it integrates a global SSM branch with a local Transformer branch to achieve 3–8% lower MSE and 2–5× faster processing.
• For visual recognition, MASS uses cross-attention and adaptive gating to boost top-1 accuracy and segmentation metrics while maintaining linear-time efficiency.

The Multi-scale Attention-augmented State Space Model (MASS) is an advanced architectural module that fuses multi-scale attention mechanisms with state space models (SSMs) to address both spatial and temporal dependencies across modalities, including time series forecasting and visual recognition. By coupling adaptive self-attention with efficient, learned SSM dynamics—and enabling cross-modal, cross-scale communication—MASS bridges the gap between long-range global modeling and fine-grained local context, achieving state-of-the-art empirical results with high efficiency in both memory and computation (2502.11340, Lou et al., 22 Jul 2025).

1. Core Principles and High-level Structure

MASS integrates two foundational ideas: (1) multi-scale adaptive attention for capturing local and global affinities, and (2) efficient, linear-time state evolution through selective state space models. In both time series and vision, it extends prior hybrids of SSMs and Transformers by deep cross-attention or re-mixing mechanisms, tightly coupling these components within a single token-mixer block.

In time series, as in S2TX (2502.11340), MASS comprises a global SSM (e.g., Mamba) path for long-range context aggregation and a local windowed attention branch for variable-level, high-resolution modeling. In vision, as realized in A2Mamba (Lou et al., 22 Jul 2025), MASS consists of a two-branch multi-scale self-attention (Adaptive Multi-scale Attention, AMA) whose attention maps spatially aggregate SSM hidden states, yielding 2D spatially coherent and dynamically informed outputs.

2. Mathematical Formulation

State Space Evolution

Let $X$ denote the input sequence (for time series: scalar- or vector-valued; for images: feature maps $X \in \mathbb{R}^{C \times H \times W}$). The SSM dynamics are given—per token or patch—by a linear recurrence (vectorized here):

$$x_0 = 0,\qquad x_k = A x_{k-1} + B y_k,\qquad z_k = C x_k + D y_k,\qquad k = 1, \dots, N,$$

where $y_k$ is the $k$-th (possibly projected) input, and the SSM parameters $A, B, C, D$ are either structured (e.g., HiPPO-derived, as in S2TX (2502.11340)) or learned end-to-end (as in A2-SSM (Lou et al., 22 Jul 2025)).

Multi-scale Attention and Cross-Attention

In vision (A2Mamba), channels are split and attended locally (Sliding Local Attention, SLA) and at a dilated spatial scale (Dilated Local Attention, DLA). Each branch computes:

$$A_1[i,:] = \mathrm{softmax}\left(\frac{Q_1[i] \cdot K_1[\mathrm{neigh}(i)]}{\sqrt{d}}\right), \quad A_2[i,:] = \mathrm{softmax}\left(\frac{Q_2[i] \cdot K_2[\mathrm{dil{-}neigh}(i)]}{\sqrt{d}}\right).$$

These attention maps are reused to spatially reproject the SSM hidden states:

$$S_1'[i,:] = \sum_{j \in \mathrm{neigh}(i)} A_1[i, j] S_1[j], \qquad S_2'[i,:] = \sum_{j \in \mathrm{dil{-}neigh}(i)} A_2[i, j] S_2[j].$$

The outputs are then gated and further processed with a per-channel residual:

$$Z = S' \odot \mathrm{SiLU}(C'),\qquad Z_{\mathrm{out}} = D_r \odot Z + Y_{\mathrm{flat}},$$

with $D_r$ a learnable residual gating vector.

For time series (S2TX), cross-attention fuses global SSM outputs into local Transformer streams: given global tokens $X \in \mathbb{R}^{C \times H \times W}$0 and local tokens $X \in \mathbb{R}^{C \times H \times W}$1, cross-attention is computed as:

$X \in \mathbb{R}^{C \times H \times W}$2

where queries $X \in \mathbb{R}^{C \times H \times W}$3 derive from local Transformer outputs and keys/values from global SSM tokens for each variable.

3. Architectural Components and Data Flow

Time Series Forecasting (S2TX) (2502.11340)

• Global SSM Branch: Processes downsampled, multi-variate input patches through bidirectional Mamba modules (forward/backward) and aggregates outputs for global context.
• Local Transformer Branch: Applies windowed self-attention and feed-forward layers to high-resolution local patches per variable.
• Cross-Attention Fusion: Injects SSM-derived global context into the Transformer branch for each variable via dedicated cross-attention blocks.
• Projection and Output: Variable-level representations are flattened and projected by a single linear head to the forecast horizon.

Visual Recognition (A2Mamba) (Lou et al., 22 Jul 2025)

• 3×3 Depthwise Convolution: Adds positional biases.
• AMA (Adaptive Multi-scale Attention): Splits input channels for learning both fine-grained and global context, via both sliding-window and dilated attention.
• Attention-augmented SSM (A2-SSM): Evolves the flattened, attended features via SSM, with spatial cross-mixing by the computed attention maps.
• Gating and Residual: Applies SiLU-based channel gating and learnable per-channel residual weights, followed by reconstruction of spatial dimensions.

The following table summarizes the primary architectural flows in time series and vision instantiations:

Domain	Global Path	Local Path	Fusion Mechanism
Time Series (S2TX)	Mamba SSM (selective, global)	Local Transformer (windowed attention)	Cross-attention (global→local)
Vision (A2Mamba)	SSM (A2-SSM, linear-time)	AMA (local & dilated self-attention)	Attention-reprojection of SSM states

4. Computational Complexity and Optimization

The design of MASS ensures that its computational and memory footprints scale linearly with the sequence/spatial length, avoiding the quadratic cost of global attention, except in small local windows:

• Global SSM branch: $X \in \mathbb{R}^{C \times H \times W}$4 (linear in sequence length or pixel count and hidden dimension).
• Local attention: $X \in \mathbb{R}^{C \times H \times W}$5 per variable or spatial window, with $X \in \mathbb{R}^{C \times H \times W}$6 (or window size $X \in \mathbb{R}^{C \times H \times W}$7) chosen small (e.g., 16–32 for time series, $X \in \mathbb{R}^{C \times H \times W}$8–$X \in \mathbb{R}^{C \times H \times W}$9 for vision), resulting in negligible overhead.
• End-to-end training: Uses AdamW optimizer with cosine learning-rate decay. No additional regularization beyond weight decay is typically required (2502.11340).

Empirically, as input length increases, the stride for global SSM patches can be dynamically chosen to keep the cost constant in practice.

5. Empirical Performance and Ablations

Time Series (S2TX (2502.11340))

On seven multivariate benchmarks (ETTh1, ETTh2, ETTm1, ETTm2, Weather, Electricity, Exchange), MASS achieves or matches the lowest mean squared error across all forecast horizons (96, 192, 336, 720), averaging 3–8% lower MSE versus prior best (SST), using only $$x_0 = 0,\qquad x_k = A x_{k-1} + B y_k,\qquad z_k = C x_k + D y_k,\qquad k = 1, \dots, N,$$0 memory and 2–5× faster than full-sequence Transformers. MASS is robust to missing data: with up to 40% random masking and imputation, MSE increases by only $$x_0 = 0,\qquad x_k = A x_{k-1} + B y_k,\qquad z_k = C x_k + D y_k,\qquad k = 1, \dots, N,$$15%, compared to $$x_0 = 0,\qquad x_k = A x_{k-1} + B y_k,\qquad z_k = C x_k + D y_k,\qquad k = 1, \dots, N,$$215% for SST.

Visual Recognition (A2Mamba (Lou et al., 22 Jul 2025))

MASS demonstrably improves top-1 accuracy and segmentation/detection metrics relative to both prior SSM-based and Transformer-based backbones, e.g.:

• ImageNet-1K classification (224²): A2Mamba-L (MASS) achieves 86.1% top-1 accuracy.
• Semantic segmentation (ADE20K/mIoU): A2Mamba-S yields 51.6% mIoU, $$x_0 = 0,\qquad x_k = A x_{k-1} + B y_k,\qquad z_k = C x_k + D y_k,\qquad k = 1, \dots, N,$$3 over the best non-MASS prior.
• Object detection/segmentation (COCO, AP^b/AP^m): A2Mamba-S attains 54.0% AP^b / 46.6% AP^m, outperforming MambaVision-B with 40% fewer parameters.

Ablation studies indicate the largest single improvement ($$x_0 = 0,\qquad x_k = A x_{k-1} + B y_k,\qquad z_k = C x_k + D y_k,\qquad k = 1, \dots, N,$$4 top-1, $$x_0 = 0,\qquad x_k = A x_{k-1} + B y_k,\qquad z_k = C x_k + D y_k,\qquad k = 1, \dots, N,$$5 mIoU) derives from attention-augmented crossmixing of SSM states, with further gains from per-channel residual gating.

6. Significance and Limitations

MASS represents the first architectural token mixer to deeply intertwine multi-scale attention with SSM evolution. Its unique adaptive attention → SSM → attention feedback loop introduces 2D spatial dependencies (addressing limitations of strictly 1D scans), while preserving the linear-time scalability and global context aggregation of SSMs. For multimodal and multiscale sequence modeling, this design combines the expressivity of global attention with the efficiency and memorization strength of SSMs.

A plausible implication is that future architectures that integrate more advanced cross-modal and multi-scale feedback between attention mechanisms and recurrent or SSM dynamics may inherit similar benefits in both accuracy and efficiency.

7. Hyperparameters and Implementation Considerations

Typical hyperparameters for MASS include:

• Window sizes per stage (e.g., $$x_0 = 0,\qquad x_k = A x_{k-1} + B y_k,\qquad z_k = C x_k + D y_k,\qquad k = 1, \dots, N,$$6 for vision tasks).
• Adaptive dilation rates to cover complete spatial maps.
• Equal channel splitting between parallel attention branches.
• Gated activations (e.g., SiLU) and learnable per-channel residuals.
• Linear SSM dimension set to the token or channel dimension.

Implementation utilizes depthwise convolutions, LayerNorm, local and/or dilated attention, learnable state space matrices (structured or fully learned), efficient flattening/restoration of spatial dimension, and compositional gating mechanisms.

Public code for the A2Mamba variant is available at https://github.com/LMMMEng/A2Mamba (Lou et al., 22 Jul 2025).

Both S2TX and A2Mamba architectures implement MASS as a modular block, enabling integration into larger hybrid or hierarchical architectures for sequence and vision modeling, respectively.