TRUST: Test-Time Refinement using Uncertainty-Guided SSM Traverses
Sahar Dastani, Ali Bahri, Gustavo Adolfo Vargas Hakim, Moslem Yazdanpanah, Mehrdad Noori, David Osowiechi, Samuel Barbeau, Ismail Ben Ayed, Herve Lombaert, Christian Desrosiers
Abstract
State Space Models (SSMs) have emerged as efficient alternatives to Vision Transformers (ViTs), with VMamba standing out as a pioneering architecture designed for vision tasks. However, their generalization performance degrades significantly under distribution shifts. To address this limitation, we propose TRUST (Test-Time Refinement using Uncertainty-Guided SSM Traverses), a novel test-time adaptation (TTA) method that leverages diverse traversal permutations to generate multiple causal perspectives of the input image. Model predictions serve as pseudo-labels to guide updates of the Mamba-specific parameters, and the adapted weights are averaged to integrate the learned information across traversal scans. Altogether, TRUST is the first approach that explicitly leverages the unique architectural properties of SSMs for adaptation. Experiments on seven benchmarks show that TRUST consistently improves robustness and outperforms existing TTA methods. The code is available at: https://github.com/Sahardastani/trust.
Ask about this paper
Your agent reads all of it.
Lune indexed this paper to the last equation, along with the top-tier papers that cite it. Ask a question and the answer quotes them.
Your agent calls
Luneget_paper_fulltext
Free to start. No credit card required.
Terminal
Install the CLIlune papers fulltext f0bac5d1-5a18-4f8e-b1da-bbb7d00750b6Builds on34
- Swin Transformer: Hierarchical Vision Transformer using Shifted WindowsZe Liu, Yutong Lin, Yue Cao, Han Hu et al.ICCV 2021 · 31,683 citations
- An Image is Worth 16x16 Words: Transformers for Image Recognition at ScaleAlexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn et al.ICLR 2021 · 21,477 citations
- Training data-efficient image transformers & distillation through attentionHugo Touvron, Matthieu Cord, Matthijs Douze, Francisco Massa et al.ICML 2021 · 8,974 citations
- A ConvNet for the 2020sZhuang Liu, Hanzi Mao, Chao-Yuan Wu, Christoph Feichtenhofer et al.CVPR 2022 · 6,782 citations
- Efficiently Modeling Long Sequences with Structured State SpacesAlbert Gu, Karan Goel, Christopher RéICLR 2022 · 3,482 citations
Related papers
- Spectral State Space Model for Rotation-Invariant Visual Representation LearningSahar Dastani, Ali Bahri, Moslem Yazdanpanah, Mehrdad Noori et al.CVPR 2025
- DGMamba: Domain Generalization via Generalized State Space ModelShaocong Long, Qianyu Zhou, Xiangtai Li, Xuequan Lu et al.ACM MM 2024 · 15 citations
- START: A Generalized State Space Model with Saliency-Driven Token-Aware TransformationJintao Guo, Lei Qi, Yinghuan Shi, Yang GaoNeurIPS 2024 · 6 citations
- MaskViM: Domain Generalized Semantic Segmentation with State Space ModelsJiahao Li, Yang Lu, Yuan Xie, Yanyun QuAAAI 2025 · 1 citation
- Partial Ring Scan: Revisiting Scan Order in Vision State Space ModelsYi-Kuan Hsieh, Kuan-Chuan Peng, Xin Li, Ming-Ching Chang et al.ICML 2026
