Table of Contents
Fetching ...

Hard-Attention Gates with Gradient Routing for Endoscopic Image Computing

Giorgio Roffo, Carlo Biffi, Pietro Salvagnini, Andrea Cherubini

TL;DR

The paper addresses overfitting and limited generalization in automated gastroenterological polyp size assessment by introducing Hard-Attention Gates (HAG) for online feature selection and Gradient Routing (GR) for independent optimization of gating. HAG reweights features through sigmoid gates so that the gated input X_att = F ⊙ X emphasizes informative cues, while GR updates gating and main model parameters in a dual-pass scheme with distinct optimizers and gradient clipping, formalized by updates for θ_att and θ_main. Empirically, HAG improves performance on CIFAR-100 (e.g., ViT-T+HAG reaching 83.8% accuracy) and yields superior polyp sizing metrics on REAL-Colon, Misawa, and SUN datasets, including binary F1 scores around 87.8% and triclass F1 around 76.5% for key configurations. The authors also provide a public codebase and standardized dataset splits to enhance reproducibility in polyp size estimation and potentially extend to broader medical imaging tasks.

Abstract

To address overfitting and enhance model generalization in gastroenterological polyp size assessment, our study introduces Feature-Selection Gates (FSG) or Hard-Attention Gates (HAG) alongside Gradient Routing (GR) for dynamic feature selection. This technique aims to boost Convolutional Neural Networks (CNNs) and Vision Transformers (ViTs) by promoting sparse connectivity, thereby reducing overfitting and enhancing generalization. HAG achieves this through sparsification with learnable weights, serving as a regularization strategy. GR further refines this process by optimizing HAG parameters via dual forward passes, independently from the main model, to improve feature re-weighting. Our evaluation spanned multiple datasets, including CIFAR-100 for a broad impact assessment and specialized endoscopic datasets (REAL-Colon, Misawa, and SUN) focusing on polyp size estimation, covering over 200 polyps in more than 370,000 frames. The findings indicate that our HAG-enhanced networks substantially enhance performance in both binary and triclass classification tasks related to polyp sizing. Specifically, CNNs experienced an F1 Score improvement to 87.8% in binary classification, while in triclass classification, the ViT-T model reached an F1 Score of 76.5%, outperforming traditional CNNs and ViT-T models. To facilitate further research, we are releasing our codebase, which includes implementations for CNNs, multistream CNNs, ViT, and HAG-augmented variants. This resource aims to standardize the use of endoscopic datasets, providing public training-validation-testing splits for reliable and comparable research in gastroenterological polyp size estimation. The codebase is available at github.com/cosmoimd/feature-selection-gates.

Hard-Attention Gates with Gradient Routing for Endoscopic Image Computing

TL;DR

The paper addresses overfitting and limited generalization in automated gastroenterological polyp size assessment by introducing Hard-Attention Gates (HAG) for online feature selection and Gradient Routing (GR) for independent optimization of gating. HAG reweights features through sigmoid gates so that the gated input X_att = F ⊙ X emphasizes informative cues, while GR updates gating and main model parameters in a dual-pass scheme with distinct optimizers and gradient clipping, formalized by updates for θ_att and θ_main. Empirically, HAG improves performance on CIFAR-100 (e.g., ViT-T+HAG reaching 83.8% accuracy) and yields superior polyp sizing metrics on REAL-Colon, Misawa, and SUN datasets, including binary F1 scores around 87.8% and triclass F1 around 76.5% for key configurations. The authors also provide a public codebase and standardized dataset splits to enhance reproducibility in polyp size estimation and potentially extend to broader medical imaging tasks.

Abstract

To address overfitting and enhance model generalization in gastroenterological polyp size assessment, our study introduces Feature-Selection Gates (FSG) or Hard-Attention Gates (HAG) alongside Gradient Routing (GR) for dynamic feature selection. This technique aims to boost Convolutional Neural Networks (CNNs) and Vision Transformers (ViTs) by promoting sparse connectivity, thereby reducing overfitting and enhancing generalization. HAG achieves this through sparsification with learnable weights, serving as a regularization strategy. GR further refines this process by optimizing HAG parameters via dual forward passes, independently from the main model, to improve feature re-weighting. Our evaluation spanned multiple datasets, including CIFAR-100 for a broad impact assessment and specialized endoscopic datasets (REAL-Colon, Misawa, and SUN) focusing on polyp size estimation, covering over 200 polyps in more than 370,000 frames. The findings indicate that our HAG-enhanced networks substantially enhance performance in both binary and triclass classification tasks related to polyp sizing. Specifically, CNNs experienced an F1 Score improvement to 87.8% in binary classification, while in triclass classification, the ViT-T model reached an F1 Score of 76.5%, outperforming traditional CNNs and ViT-T models. To facilitate further research, we are releasing our codebase, which includes implementations for CNNs, multistream CNNs, ViT, and HAG-augmented variants. This resource aims to standardize the use of endoscopic datasets, providing public training-validation-testing splits for reliable and comparable research in gastroenterological polyp size estimation. The codebase is available at github.com/cosmoimd/feature-selection-gates.
Paper Structure (10 sections, 4 equations, 7 figures, 6 tables)

This paper contains 10 sections, 4 equations, 7 figures, 6 tables.

Figures (7)

  • Figure 1: Feature Selection-Attention Gates (HAG) Integration in Deep Learning Models. (a) Conceptual design of HAG. (b) Application of HAG in a multi-stream CNN architecture, with each stream being optional. (c) Embedding of HAG within a ViT block, positioned after the multihead attention and the MLP for enhanced feature re-weighting.
  • Figure 2: ViT/R18 Performance on CIFAR-100 compared with SotA TinyViT2022Deng2020R18Cifar.
  • Figure 3: CIFAR-100: FSG-GR weight distributions in the ViT.
  • Figure 4: Polyp Sizing: FSG-GR weight distributions in the ViT.
  • Figure 5: Distribution of GT Classes in each Fold.
  • ...and 2 more figures