TLDR: A new research paper introduces Randomized-MLP (RMLP) regularization, a contrastive learning-based method that improves the interpretability and robustness of Vision Transformers (ViTs) like DINOv2. By replacing standard MLP heads with a randomized, non-trainable operator, RMLP encourages semantically aligned representations, reduces attention artifacts, and maintains or improves performance across natural and medical imaging domains, particularly achieving state-of-the-art results in ophthalmology datasets, all without adding trainable parameters.
Vision Transformers (ViTs) have become a cornerstone in computer vision, excelling in various tasks. However, models like DINOv2, despite their strong performance, often face challenges in interpretability, especially when applied to specialized fields like medical imaging. A new research paper introduces a novel approach called Randomized-MLP (RMLP) regularization, designed to enhance both the performance and transparency of these powerful models.
The core issue identified by the researchers is that ViTs can sometimes repurpose less informative parts of an image, known as patch tokens, in ways that make it harder to understand what the model is focusing on. This problem is particularly pronounced in medical imaging, where subtle domain shifts can lead to a degradation in both accuracy and the ability to interpret the model’s decisions. The paper highlights two types of artifacts: ‘first-order artifacts’ where the model pays too much attention to background or irrelevant regions, and ‘second-order artifacts’ related to how the model captures variance in image features.
To tackle these challenges, Joel Valdivia Ortega and his colleagues propose Randomized-MLP (RMLP) regularization. This method is rooted in contrastive learning and aims to encourage the model to create more semantically meaningful representations of images. When fine-tuning DINOv2, the standard learnable MLP (Multi-Layer Perceptron) heads are replaced with a randomized, non-trainable operator. This clever design ensures that the fundamental structure of the representation space is maintained, while guiding the model to learn more robust and interpretable features without introducing any new trainable parameters.
The RMLP regularization offers several key advantages. It either improves or maintains the model’s performance on various tasks while simultaneously generating attention maps that are much easier to interpret. This is crucial for applications like medical diagnostics, where understanding why a model makes a certain prediction is as important as the prediction itself. The method is lightweight, theoretically sound, and can be easily integrated into existing ViT-based models without requiring extensive retraining.
The researchers conducted extensive experiments, fine-tuning DINOv2-S and another model called SwAV with RMLP across both natural image datasets (like ImageNet-1k) and medical imaging modalities (such as Colour Fundus Photography and Optical Coherence Tomography). The results were compelling: RMLP consistently outperformed traditional MLP baselines. In the medical domain, it achieved state-of-the-art performance in OCT segmentation and significantly improved classification accuracy, even with limited training data.
A significant finding was RMLP’s ability to reduce attention artifacts. In medical images, models fine-tuned with RMLP showed better alignment with anatomical features, focusing on biologically relevant tissues rather than void or irrelevant regions. This leads to more reliable and trustworthy pathology detection. The paper also includes a mathematical analysis, explaining how RMLPs effectively transform individual data points into ‘probability balls’ in the representation space, enhancing robustness and promoting sparsity without distorting the overall structure of the learned representations.
This work represents a promising step forward in making advanced vision models more transparent and reliable, particularly for sensitive applications like medical imaging. The code and pre-trained models are available for further exploration at https://github.com/peng-lab/rmlp.
Also Read:
- Task-Aware Explanations Enhance AI Prediction and Understanding
- Accelerating Image AI Training with CoMA: A New Approach to Masked Autoencoders and Dynamic Vision Transformers
While the RMLP regularization shows great potential, the authors acknowledge some limitations. The optimal strength of the regularization might vary depending on the data or the model’s architecture, and currently relies on heuristic tuning. Future research will explore methods for principled selection of this strength and evaluate performance on larger-scale datasets and a wider array of domains.
