提出START模型,用线性复杂度提升跨域泛化能力
START: A Generalized State Space Model with Saliency-Driven Token-Aware Transformation
- 设计注意力感知的令牌变换机制,抑制领域特异性特征
- 在5个基准上超越现有最先进方法,保持线性计算复杂度
- 适合需要高效跨域泛化且资源受限的场景
领域泛化(DG)旨在通过多源域学习使模型适应未见目标域。现有方法多依赖卷积神经网络(CNN),因感受野有限易学习纹理偏差,导致过拟合源域。虽有基于视觉变换器(ViTs)的方法利用全局感受野,但自注意力的二次复杂度带来高计算开销。最近,状态空间模型(SSM)如Mamba在监督任务中展现潜力,训练时序列长度呈线性复杂度,推理类似RNN快速。受此启发,我们研究Mamba在域偏移下的泛化能力,发现输入相关矩阵会累积并放大领域特异性特征,阻碍泛化。为此,提出新型基于SSM的架构START,实现当前最优性能,并为CNN和ViTs提供高效替代方案。START可选择性地扰动与抑制输入中显著令牌内的领域特异性特征,有效降低不同域间的差异。大量实验在五个基准上验证,START优于现有最先进方法,且具备线性复杂度。代码已公开于https://github.com/lingeringlight/START。
原文摘要 · Abstract (English)
Domain Generalization (DG) aims to enable models to generalize to unseen target domains by learning from multiple source domains. Existing DG methods primarily rely on convolutional neural networks (CNNs), which inherently learn texture biases due to their limited receptive fields, making them prone to overfitting source domains. While some works have introduced transformer-based methods (ViTs) for DG to leverage the global receptive field, these methods incur high computational costs due to the quadratic complexity of self-attention. Recently, advanced state space models (SSMs), represented by Mamba, have shown promising results in supervised learning tasks by achieving linear complexity in sequence length during training and fast RNN-like computation during inference. Inspired by this, we investigate the generalization ability of the Mamba model under domain shifts and find that input-dependent matrices within SSMs could accumulate and amplify domain-specific features, thus hindering model generalization. To address this issue, we propose a novel SSM-based architecture with saliency-based token-aware transformation (namely START), which achieves state-of-the-art (SOTA) performances and offers a competitive alternative to CNNs and ViTs. Our START can selectively perturb and suppress domain-specific features in salient tokens within the input-dependent matrices of SSMs, thus effectively reducing the discrepancy between different domains. Extensive experiments on five benchmarks demonstrate that START outperforms existing SOTA DG methods with efficient linear complexity. Our code is available at https://github.com/lingeringlight/START.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。