Tiny ViT Distillation

A 5.7M-parameter ultrasound encoder for point-of-care deployment: simple masked latent distillation, and where compression fails.

The gist

Ultrasound foundation models keep converging on the same shape. A Vision Transformer with roughly 86 million parameters. That is fine inside a research lab. It will never run on the hardware point-of-care ultrasound actually happens on, a laptop, a phone, or a handheld probe's own onboard compute. A prior paper, TinyUSFM, already showed these models can be compressed. But it took an elaborate pipeline to get there: a curated 200,000-image subset, a domain-separated masking scheme, and dynamically weighted consistency losses. This project asks a more basic question first. How much of that machinery is actually necessary?

The answer it tries is almost stubbornly simple. Freeze a large teacher model. Show a small student model a masked ultrasound image. Have the student predict what the teacher would have seen at the masked positions. Not the pixels themselves, but the teacher's own internal representation of them. Add one small auxiliary head on top of that, and stop there. No curated subset. No dynamic weighting. Train on the full, uncurated corpus.

Research question

Point-of-care ultrasound is exactly the setting where automated interpretation would matter most. It is also exactly the setting current ultrasound foundation models cannot run in. TinyUSFM already proved that compression itself is possible. That is not the question this project is asking.

The real question is which parts of a heavier compression pipeline were genuinely necessary, and which downstream tasks quietly stop working the moment the model gets smaller. A recipe that is simpler to build and train is only worth adopting if it does not silently break something a clinician would actually rely on.

Solution

A frozen 86-million-parameter ViT-B teacher, USF-MAE, supervises a 5.7-million-parameter ViT-Tiny student. Input patches are masked at a ratio between 0.3 and 0.5. The student encodes the masked image, and a cross-attention predictor regresses the teacher's embeddings at exactly the masked positions, under a Smooth-L1 loss. One auxiliary head sits on top of the pooled student feature, adding a small supervised signal.

That is the entire recipe. Two components, not the prior pipeline's several. Trained on an uncurated corpus rather than a hand-picked subset. The bet behind this design is that most of what made the earlier, heavier pipeline work was never the curation or the dynamic weighting in the first place. It was the basic act of making a small model predict a large model's internal representation.

Results

Training used 532,752 B-mode ultrasound images spanning 18 organs and 54 acquisition sources, all from public datasets, deduplicated by perceptual hash. Evaluation covered seven classification datasets and nine segmentation datasets, each checked for overlap with the pretraining corpus first.

Classification is where the recipe's simplicity pays off most clearly. The numbers below are reported relative to baselines.

At the pooled-feature level, the compact student is not just close to the much larger models it learned from. Fine-tuned, it is statistically indistinguishable from four separate 86-million-parameter baselines.

That parity does not extend to segmentation, where the compact model genuinely trails the larger encoders it was distilled from.

So compression is nearly free for whole-image classification, but it genuinely costs something for dense, pixel-level prediction. A separate analysis looked at the shape of what each model actually learned, and it turned up a real surprise. A model's usable representational rank has no measurable correlation with its downstream accuracy. Stranger still, the relationship between rank and robustness to image corruption flips sign depending on whether you are comparing different models to each other or shrinking one model's own representation. Some representation-quality metrics that sound intuitive simply do not transfer the way you would expect them to.

What next

The practical case here is straightforward. A model that matches much larger encoders on classification, at a fraction of the size, is exactly what point-of-care hardware needs for the tasks it is most often asked to do, like triage and quick screening reads. That is a direct path to real deployment on a phone or a handheld probe.

Segmentation is the open problem. The gap there suggests dense, pixel-level prediction needs either a larger student or a distillation signal that targets spatial detail directly, rather than a single pooled representation. Closing that gap, or clearly scoping the compact model to tasks where it already matches its teacher, is the natural next step.

Status & links

Manuscript targeting Computers in Biology and Medicine.