💻 computer science

Learning To Focus: Anatomy-Guided Attention Regularization for Medical Image Classification

The paper proposes Locus, a framework that leverages pretrained segmentation foundation models to introduce an adaptive attention regularization term, guiding medical image classifiers toward diagnostically relevant anatomical regions without requiring manual segmentation annotations.

Original authors: Tonmoy Hossain, Atiqur Rahman, Farhana Hossain Swarnali, Miaomiao Zhang

Published 2026-07-14

📖 6 min read🧠 Deep dive

CC BY 4.0

Original authors: Tonmoy Hossain, Atiqur Rahman, Farhana Hossain Swarnali, Miaomiao Zhang

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you're trying to teach a robot to spot a specific type of tumor in a medical scan. You show it thousands of pictures and say, "This one has a tumor, this one doesn't." But here's the catch: you never tell the robot where to look.

Without a guide, the robot gets a bit lazy. Instead of hunting for the actual tumor, it might start guessing based on weird clues, like the color of the scanner's border, a shadow in the corner, or even the type of paper the image was printed on. It's like a student taking a test who ignores the math problem and just guesses the answer because the teacher's handwriting looks familiar. The robot might get the right answer by luck, but it's not actually learning what a tumor looks like.

This is the problem the authors of this paper, Locus, are trying to fix.

The "Magic Eye" Problem

In the past, to teach the robot where to look, scientists had to hire expert doctors to draw detailed outlines around every single tumor in every single picture. This is like hiring an artist to trace every leaf on a tree before you can teach a squirrel to find the nuts. It takes forever, costs a fortune, and is hard to do for huge piles of data.

Some other researchers tried to teach the robot to draw its own outlines while it learned to find the tumor. But this is like asking a student to solve a math problem and draw a perfect picture of the answer at the same time. The two tasks often get in each other's way, and the robot ends up confused.

The New Trick: Borrowing a Superpower

The authors had a clever idea. Instead of hiring artists or making the robot draw, they decided to borrow a "magic eye" that already exists. There are powerful AI models (called Segmentation Foundation Models) that are already incredibly good at finding body parts like lungs, bones, or skin lesions. These models are like a super-organized librarian who can instantly point to the "Anatomy Section" in a messy library without needing to be trained on every single book.

The authors' framework, Locus, uses this pre-trained librarian to give the robot a gentle nudge. It doesn't force the robot to stare exactly at the outline the librarian drew. Instead, it just says, "Hey, make sure you're looking mostly at the body part, not the empty space around it."

How It Works: The "Don't Look at the Background" Rule

Here is the secret sauce: The robot usually pays attention to the whole image. The authors added a special rule to the robot's training. They measure how much attention the robot is paying to the "foreground" (the body part) versus the "background" (the empty space).

If the robot starts paying more attention to the background than the body part, the rule slaps its hand (mathematically speaking) and says, "Nope, focus on the anatomy!" But if the robot is already looking at the body part, the rule stays quiet. This is a "hinge" rule—it only kicks in when things go wrong. It's like a coach who only yells when the player starts running in the wrong direction, letting them play freely when they're on track.

To make sure the librarian doesn't miss any parts (like a tumor that's split into two pieces), the authors chop the image into little squares and ask the librarian to look at each square individually. This ensures no tiny tumor gets left behind.

Does It Actually Work?

The authors tested this idea on eight different medical imaging datasets. These included:

Skin photos (dermoscopy) with over 10,000 images.
Chest X-rays with over 112,000 images.
Bone X-rays, heart MRI videos, and microscope slides of tissue.

They compared their method against three other approaches:

Standard Training: Just guessing without any help.
Masked Input: Only showing the robot the body part and hiding the rest (which can accidentally hide important clues).
Strict Alignment: Forcing the robot to match the librarian's drawing perfectly (which is too rigid).

The Results:
The Locus method consistently won. On the skin cancer dataset, it reached 86.53% accuracy (compared to 83.99% for the standard method). On the bone tumor dataset, it hit 87.54% accuracy. On the heart MRI videos, it reached 83.15% accuracy.

But the real win wasn't just the score; it was the "eyes" of the robot. When the authors looked at where the robot was focusing, the standard robots were often looking at shadows or the edges of the image. The Locus robots, however, were staring right at the lungs, the tumors, or the heart muscle, just like a human doctor would.

Why This Matters

The authors suggest that this method makes the robots more reliable. When they tested the robots with blurry or noisy images (like a photo taken in bad light), the standard robots got confused and their scores dropped sharply. The Locus robots, however, stayed steady. Because they were trained to ignore the background noise and focus on the anatomy, they didn't get tricked by the blur.

What They Didn't Do

It's important to note what this paper doesn't claim. They didn't say this is a cure for cancer or that it replaces doctors. They didn't say it works perfectly on every single type of disease (they specifically excluded cases where the diagnosis isn't about a specific shape, like "no finding" in an X-ray). They also didn't say the robot is perfect; they just showed it's better at focusing than the robots we had before.

The Bottom Line

The authors propose a way to teach medical AI to "look where it should" by borrowing the eyes of a pre-trained model, without the huge cost of hiring humans to draw outlines. It's a "plug-and-play" tool that makes the robot smarter, more focused, and less likely to get distracted by the background. They suggest this could be a big step forward for making medical AI more trustworthy, but they admit there's still work to do, like figuring out how to handle cases where the "magic librarian" might get a little confused.

In short: They taught the robot to stop looking at the shadows and start looking at the tumor, and the robot got much better at its job.

Technical Summary: Locus – Anatomy-Guided Attention Regularization for Medical Image Classification

1. Problem Statement

Medical image classification models are typically trained using only image-level labels, lacking explicit spatial supervision. Consequently, these models often learn to rely on statistical correlations with anatomically irrelevant regions, such as imaging artifacts, surrounding tissue, or dataset-specific contextual cues, rather than the diagnostically relevant anatomical structures. This behavior limits model interpretability and generalizability.

While previous approaches have attempted to address this by incorporating anatomical shape information (e.g., segmentation masks), they face significant limitations:

Manual Annotation: Relying on expert-delineated masks is costly and difficult to scale.
Joint Training: Training a segmentation branch alongside a classifier introduces optimization challenges, such as competing gradients and noisy predictions that can corrupt guidance.
CAM-based Pseudo-masks: While avoiding annotation costs, these methods often activate only discriminative regions correlated with the label rather than the full relevant anatomical structure.
Dedicated Segmentation Models: Training a new segmentation model from scratch reintroduces the annotation burden and computational overhead.

2. Methodology: The Locus Framework

The authors propose Locus, a framework that leverages pretrained Segmentation Foundation Models (SFMs) (e.g., Medical SAM, BiomedParse) to guide classifier attention without requiring dedicated segmentation training or manual annotations.

2.1 Anatomy-Centric Attention Regularization

Instead of enforcing rigid pixel-wise alignment between the classifier's attention map and the anatomical mask (which can be overly restrictive), Locus introduces a hinge-based regularization term ( $R$ ).

Mechanism: The framework calculates the weighted average attention energy within the anatomical foreground ( $A_{fg}$ ) and the background ( $A_{bg}$ ) derived from the SFM mask.
Loss Function: The regularizer is defined as $R(M, A) = \max(0, A_{bg} - A_{fg})$ .
Logic: This term penalizes the classifier only when background attention dominates foreground attention. If the foreground attention is already sufficient ( $A_{fg} \geq A_{bg}$ ), the regularizer remains inactive. This approach avoids over-suppressing potentially diagnostically relevant contextual information while ensuring the model does not rely primarily on background noise.

2.2 Patch-wise Anatomical Prompting

Standard prompting of SFMs on full images often fails to capture all relevant anatomical structures, particularly when they are small or disjoint (e.g., scattered lesions or bilateral organs), often returning only the most dominant region.

Strategy: Locus partitions the input image into a grid of non-overlapping patches. Each patch is independently prompted with the same class-derived text descriptor (e.g., "lung," "tumor").
Aggregation: The resulting patch-level masks are aggregated via union to form a full-resolution anatomical mask ( $M$ ).
Goal: This strategy maximizes anatomical coverage to ensure all relevant regions contribute to attention supervision, even if the resulting mask is coarse. The authors emphasize that precise pixel-level segmentation is not required; the goal is sufficient anatomical grounding to guide attention.

2.3 Network Loss

The total network loss combines the standard task loss (cross-entropy) with the anatomy-centric regularizer:
$L_{net} = L_{task} + \lambda \cdot R(M, A)$

Training Strategy: A warm-up strategy is employed where the network is trained with $L_{task}$ alone for a few epochs before the regularization strength ( $\lambda$ ) is gradually ramped up.
Flexibility: The framework is agnostic to the specific attention extraction method (e.g., Grad-CAM, Grad-CAM++, Vanilla CAM) and the choice of SFM.

3. Key Contributions

Novel Regularization Framework: Proposes an anatomy-centric attention regularizer that adaptively balances attention between foreground and background, penalizing the classifier only when background attention dominates, thereby removing the need for dedicated segmentation training.
Plug-and-Play Adaptability: The framework is easily adaptable to various classifier backbones (CNNs and Transformers) and SFMs without modifying the classifier architecture or fine-tuning the foundation model.
Comprehensive Validation: Validated on eight diverse medical imaging datasets (dermoscopy, X-ray, histopathology, ultrasound, and cardiac MRI), demonstrating consistent performance gains and improved interpretability.

4. Experimental Results

The authors evaluated Locus against three baselines: standard Cross-Entropy (CE), Masked input (using SFM masks to discard background), and MGA (Dice-based attention alignment requiring expert annotations).

Classification Performance: Locus achieved the highest accuracy and F1-scores in nearly all settings across six benchmark datasets (HAM10000, PneumoniaMNIST, BTXRD, PanNuke, ChestX-ray8, BUSI) and a cardiac cine MRI dataset. It consistently outperformed unconstrained training and masked-input baselines.
Attention Quality: Visualizations (Grad-CAM) showed that Locus produced more focused and anatomically coherent attention maps. Notably, for architectures prone to background dominance (e.g., ConvNeXt, ViT), Locus significantly sharpened attention onto relevant structures (e.g., lung fields, lesions) compared to standard CE training.
Robustness: Under test-time perturbations (brightness, contrast, Gaussian blur, noise), models trained with Locus maintained higher accuracy than CE-only baselines, suggesting the regularization discourages reliance on spurious features sensitive to distortion.
Generalizability: The benefits of the regularizer were consistent across different attention extraction methods (Grad-CAM, Grad-CAM++, Vanilla CAM, Gradient Norm) and different network backbones (ResNet, ConvNeXt, ViT).

5. Significance and Claims

The paper claims that Locus offers a practical solution to the "black box" nature of medical image classifiers by steering them toward diagnostically relevant anatomical structures without the prohibitive cost of manual annotation or the complexity of joint segmentation-classification training.

By leveraging the "out-of-the-box" localization capabilities of foundation models and employing a flexible, hinge-based regularizer, the framework achieves:

Improved Interpretability: Models focus on the correct anatomical regions (e.g., myocardium in cardiac MRI) rather than artifacts or background.
Enhanced Generalizability: Better performance across diverse modalities and robustness against image perturbations.
Efficiency: It avoids the computational overhead of training a dedicated segmentation model while providing reliable anatomical guidance.

The authors conclude that this approach represents a significant step toward making medical AI more reliable and clinically interpretable, particularly in scenarios where pixel-level annotations are unavailable. Future work is suggested to explore uncertainty-aware weighting and end-to-end learning of patch configurations.