PS-MAE

A clinically-informed peritumoral masked autoencoder for axillary lymph node status assessment from breast ultrasound.

The gist

Here is the clinical fact this whole project starts from. Whether breast cancer has spread to the axillary lymph nodes is not written only in the tumour itself. Radiologists also read the tissue right around it. Margin irregularity matters. Subtle changes in the surrounding tissue matter too.

Most imaging models never get the chance to learn any of that. They get pretrained by masking random patches anywhere in the image and reconstructing raw pixels. That process spends most of its capacity on ultrasound speckle. It spends almost none of it on anything actually diagnostic.

PS-MAE changes two things at once. It masks a ring of tissue drawn only around the tumour. Not the lesion itself. Not random patches. And instead of asking the model to redraw pixels, it asks the model to reconstruct the edges inside that ring. Edges survive speckle noise. Raw pixels do not. On their own, neither change moves the needle much. Put together, they do.

Research question

Confirming whether cancer has already reached the axillary lymph nodes currently means a biopsy. That is invasive, so predicting the answer from an ultrasound scan instead would spare a lot of patients an unnecessary procedure. Two things get in the way of doing that well.

Labelled data for this exact question is scarce, because collecting it means running the invasive test in the first place. Ultrasound is also a noisy signal on its own. Its low signal-to-noise ratio means most self-supervised pretraining ends up learning the noise, not the peritumoral tissue that actually carries the prognostic signal. So the real question here is not just whether ultrasound can predict nodal status. It is whether a pretraining scheme can be pointed at the right tissue from the start, instead of learning whatever pattern happens to be easiest to reconstruct.

Solution

PS-MAE starts by teaching a U-Net to find the tumour-like region on a large pool of unlabeled breast-ultrasound images. Once that boundary exists, it gets dilated outward by a few patches. That dilated band is the peritumoral ring, and it is the only part of the image PS-MAE masks during pretraining. Random patches never get touched.

From tumour mask to masking regions, across three samples

Column (b) is the tumour segmentation the U-Net produces. Column (c) expands that boundary into the perimeter and background ring PS-MAE actually masks. Column (d) is what the encoder sees: everything inside that ring blanked out, the rest left as context.

Instead of reconstructing raw pixel values inside that ring, the model reconstructs the Sobel edge-gradient response of the masked region. That shift matters, because edges emphasise boundary and texture structure, which is closer to what a radiologist actually looks at. Raw pixel values mostly just chase speckle.

What the model is asked to reconstruct: pixels versus edges

The pixel target (b) keeps every grey-level value inside the ring, speckle included. The Sobel target (c) keeps only edge structure, which is what actually distinguishes tumour margin from surrounding tissue.

The encoder does not start from ImageNet weights either. It starts from USFM, a model already trained on ultrasound. That choice means training time goes toward learning the peritumoral signal, not toward first unlearning natural-image statistics.

Once pretraining finishes, the decoder is discarded. The frozen encoder is evaluated with a linear probe on labelled nodal-status data. Attention rollout is then used as a check afterward. It shows how much of the model's attention actually lands inside the peritumoral ring it was trained to look at.

Results

Pretraining pooled 83,271 unlabeled frames from 11 public breast-ultrasound sources. Downstream evaluation used an internal cohort of 698 images from 175 patients, plus a fully held-out external cohort from AIIMS New Delhi, 88 patients and 176 images, that never appeared anywhere in pretraining.

On the internal cohort, PS-MAE was tested against seven baselines. Two were generic self-supervised models, DINOv2 and an ImageNet-pretrained network. Two more were already trained specifically on ultrasound, USFM and USF-MAE. PS-MAE came out ahead of all of them.

The more interesting test is what happens on data PS-MAE has never seen. AIIMS New Delhi is a hospital that never appeared anywhere in pretraining. This is exactly where generic self-supervision tends to fall apart.

Two ablations explain where the gain actually comes from. First, the masking region and the reconstruction target only help when they work together. Sobel edges alone barely move the needle. The peritumoral ring alone barely moves it either. Combine the two, and the result clearly beats either piece on its own.

Second, attention rollout confirms the model is actually looking where it was trained to look. More of its attention lands inside the peritumoral ring than any baseline's does.

Where each model's attention actually lands

PS-MAE's attention concentrates inside the outlined peritumoral ring more consistently than the generic baselines, whose brightest regions often sit outside it.

What next

The clearest benefit here is that PS-MAE holds up somewhere it was never trained. Most models lose most of their edge the moment they leave the hospital that produced their training data. PS-MAE keeps a meaningful share of its advantage on a hospital it never saw during pretraining, and the attention-rollout result explains why. It is genuinely looking at the tissue radiologists look at, not just tissue that happens to resemble its training distribution.

That property matters most anywhere a diagnostic model has to work across hospitals with different scanners and different patient populations, which describes most real deployment settings. A model that quietly degrades the moment it changes hospitals is a much bigger practical risk than one whose accuracy ceiling is a point or two lower.

Only one external institution has validated this so far, so the next real test is broader multi-centre validation across more than one unseen hospital. The ring itself is also still a fixed width right now. Letting that width adapt to tumour size and shape, instead of using one dilation amount for every case, is the next step before this moves closer to something usable in the clinic.

Status & links

First-author manuscript under review at IEEE Journal of Biomedical and Health Informatics. Read the manuscript. Code and weights will be released on publication.