START: A Generalized State Space Model with Saliency-Driven Token-Aware Transformation

要約

ドメイン一般化 (DG) は、複数のソース ドメインから学習することで、モデルを目に見えないターゲット ドメインに一般化できるようにすることを目的としています。
既存の DG 手法は主に畳み込みニューラル ネットワーク (CNN) に依存していますが、CNN は限られた受容野により本質的にテクスチャ バイアスを学習するため、ソース ドメインの過学習が起こりやすくなります。
一部の研究では、DG がグローバル受容野を活用するためにトランスベースの手法 (ViT) を導入していますが、これらの手法では、自己注意の 2 次の複雑さのため、高い計算コストが発生します。
最近、Mamba に代表される高度な状態空間モデル (SSM) は、トレーニング中の系列長の線形複雑性と推論中の高速な RNN のような計算を達成することで、教師あり学習タスクにおいて有望な結果を示しています。
これに触発されて、私たちはドメイン シフト下での Mamba モデルの一般化能力を調査し、SSM 内の入力依存行列がドメイン固有の特徴を蓄積および増幅し、その結果モデルの一般化を妨げる可能性があることを発見しました。
この問題に対処するために、私たちは顕著性ベースのトークン認識変換 (つまり START) を備えた新しい SSM ベースのアーキテクチャを提案します。これは、最先端 (SOTA) のパフォーマンスを実現し、CNN や ViT に代わる競争力のある代替手段を提供します。
私たちの START は、SSM の入力依存マトリックス内の顕著なトークンのドメイン固有の特徴を選択的に摂動および抑制することができるため、異なるドメイン間の不一致を効果的に削減できます。
5 つのベンチマークに関する広範な実験により、START が効率的な線形複雑さで既存の SOTA DG メソッドよりも優れていることが実証されました。
私たちのコードは https://github.com/lingeringlight/START で入手できます。

要約(オリジナル)

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.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google