VariViT: A Vision Transformer for Variable Image Sizes

1. Problématique

Les Vision Transformers (ViT) sont devenus l'architecture de référence pour l'apprentissage de représentations d'images grâce à leurs mécanismes d'auto-attention. Cependant, leur application en imagerie médicale rencontre des obstacles majeurs :

• Contrainte de taille fixe : Les ViT standards divisent les images en patches de taille fixe, nécessitant un prétraitement (redimensionnement, remplissage ou recadrage) pour uniformiser les entrées.
• Perte d'information et artefacts : Le redimensionnement d'images médicales, en particulier celles contenant des structures irrégulières comme des tumeurs, peut dégrader les informations critiques et introduire des artefacts qui nuisent au diagnostic.
• Ratio signal/bruit variable : Un recadrage fixe crée des rapports variables entre le premier plan (la pathologie) et l'arrière-plan (tissu sain). Pour les petites lésions, cela inclut trop de contexte non pertinent, diluant l'attention du modèle.
• Coût computationnel : Traiter de grandes images avec une résolution fixe est coûteux, tandis que réduire la taille entraîne une perte d'information.

2. Méthodologie : VariViT

Les auteurs proposent VariViT, une architecture de Vision Transformer améliorée conçue pour gérer des images de tailles variables tout en conservant une taille de patch constante. L'approche repose sur deux piliers principaux :

A. Stratégie de Réduction de Position "Center and Select"

Contrairement aux méthodes traditionnelles qui interpolent les embeddings de position (ce qui peut entraîner une perte d'information ou un coût calcul élevé en 3D), VariViT utilise une approche géométrique :

• Alignement central : Dans les tâches de classification de tumeurs, les crops 3D sont centrés sur le centre de masse de la tumeur. Cela crée un point de référence cohérent quelle que soit la taille de la tumeur.
• Initialisation et Sélection : Le modèle initialise une grille d'embeddings de position sinusoïdaux fixes pour la taille d'image la plus grande possible. Pour les images plus petites, au lieu d'interpoler, le modèle sélectionne dynamiquement un sous-ensemble de cette grille centré sur le point de référence.
• Avantage : Cette méthode préserve l'information positionnelle originale sans interpolation complexe ni perte de précision, réduisant ainsi la charge computationnelle.

B. Stratégies de Batching (Regroupement)

Pour gérer l'hétérogénéité des tailles d'images au sein d'un même lot (batch), deux stratégies sont proposées :

1. Custom Batch Sampler (CBS) : Regroupe les images de la même taille dans un même lot. Cela permet une cohérence au sein du lot et accélère l'entraînement en évitant le gaspillage de ressources sur des images trop petites dans un lot standard.
2. Gradient Accumulation (GA) : Permet d'entraîner des lots contenant des images de tailles différentes en accumulant les gradients sur plusieurs mini-lots avant de mettre à jour les poids. Cela offre une flexibilité maximale tout en maintenant une taille de lot effective élevée.

3. Contributions Clés

1. Stratégie d'embedding flexible : Introduction d'une méthode de redimensionnement des embeddings de position basée sur la sélection centrale ("Center and Select"), adaptée aux tailles d'images variables sans interpolation.
2. Optimisation du batching : Proposition de nouvelles stratégies de regroupement (CBS et GA) pour améliorer l'efficacité computationnelle lors de l'entraînement sur des données de tailles hétérogènes.
3. Validation sur des tâches complexes : Démonstration de l'applicabilité de VariViT sur deux tâches difficiles de classification d'IRM cérébrale 3D : la prédiction du statut de mutation IDH (gliomes) et la distinction entre tumeurs cérébrales primaires et secondaires (métastases).

4. Résultats Expérimentaux

Les évaluations ont été menées sur deux jeux de données d'IRM cérébrale 3D (1856 et 1699 scans) en comparant VariViT avec un ResNet-18, un ViT standard (Vanilla ViT) et un modèle variable existant (Pix2Struct).

• Prédiction du statut IDH (Gliomes) :
  • VariViT-GA a obtenu les meilleurs scores avec un F1-score de 75,5 % et un AUC de 0,942, surpassant le ViT standard (F1: 74,4 %) et ResNet (F1: 71,6 %).
  • VariViT-CBS a également surpassé les modèles de base tout en réduisant encore davantage le temps d'entraînement.
• Classification des types de tumeurs (Primaire vs Secondaire) :
  • VariViT-GA a atteint un F1-score de 76,3 % et un MCC de 0,706, surpassant à nouveau les architectures de référence.
• Efficacité Computationnelle :
  • La stratégie de batching proposée permet de réduire le temps de calcul jusqu'à 30 % par rapport aux architectures conventionnelles.
  • Les visualisations t-SNE montrent une séparation de clusters plus nette pour VariViT, indiquant un apprentissage de caractéristiques plus discriminantes.
• Études d'ablation : La méthode "Center and Select" s'est révélée supérieure aux approches par interpolation (fixe ou apprise) pour la tâche de classification IDH.

5. Signification et Conclusion

Ce travail démontre que l'adaptation des Vision Transformers aux spécificités de l'imagerie médicale (tailles variables, importance du centre d'intérêt) est cruciale pour améliorer les performances.

• Impact clinique : En évitant le redimensionnement, VariViT préserve l'intégrité des données médicales, ce qui est essentiel pour le diagnostic précis de pathologies aux contours irréguliers.
• Efficacité : La combinaison d'une gestion intelligente des positions et de stratégies de batching innovantes permet d'entraîner des modèles plus performants et plus rapides.
• Généralisabilité : Bien que testé sur des IRM cérébrales, le cadre VariViT est conçu pour être adaptable à d'autres modalités d'imagerie et régions anatomiques, offrant une voie prometteuse pour l'apprentissage de représentations médicales à grande échelle.

En résumé, VariViT surmonte les limitations des ViT standards en imagerie médicale en alignant la flexibilité de l'architecture avec les contraintes physiques des données, offrant un compromis optimal entre précision et coût computationnel.