Decoder-Free Segmentation

A decoder-free soft-token segmenter for point-of-care breast ultrasound.

The gist

Segmentation papers almost always assume the same shape. An encoder that understands the image, and a decoder that turns that understanding back into a pixel-by-pixel mask. The decoder is usually taken for granted. Of course you need one, the thinking goes. How else would you get back to pixel resolution?

This project pokes at that assumption directly. Instead of a decoder, it uses a single linear layer on top of the Vision Transformer's own patch tokens. One number per patch, trained to predict what fraction of that patch is tumour, not just whether it is tumour or is not. At inference, that reduces to a sigmoid, a reshape into a grid, and a bilinear upsample. No decoder at all.

An oracle experiment backs this up. If you compute the best-possible continuous per-patch label directly from the ground-truth mask, and just upsample that, you already recover almost all of the achievable segmentation quality. That result suggests the usual objection, that token classification is too coarse, is really about label softness and patch size. It is not really about needing a decoder at all.

Research question

Encoder-decoder segmentation adds real parameter count and real training overhead on top of the encoder alone. Point-of-care deployment feels that cost acutely, since it is exactly the setting with the least room for extra parameters and extra latency. So the question is a direct one. Is that decoder actually earning its place, or can a much smaller mechanism do nearly as well, once patch size and label softness are handled properly?

Solution

The segmenter is a ViT encoder plus a linear head, 194 parameters at ViT-Tiny scale, producing one logit per patch token. It is trained with binary cross-entropy against the continuous tumour-area fraction inside that patch, not a binarised label. That continuous target is the part that lets a single linear layer stand in for a whole decoder.

Patch size is set to 8 instead of the usual 16, for finer spatial granularity, which works out to 784 tokens instead of 196. Patch and position embeddings are resized from a patch-16 checkpoint rather than reinitialised from scratch, so the finer grid does not have to relearn everything from nothing.

At inference, the whole pipeline is four steps. Take the sigmoid of each token's logit. Reshape the result into a grid. Bilinearly upsample it to full resolution. Threshold.

Results

Training pooled 8 public breast-ultrasound sources, 3,373 images from 3,263 patients. A ninth source, BUS-BRA, was deliberately excluded from training entirely and used only as the test population. That gave two separate tests: an in-domain test, training and testing on BUS-BRA, and a cross-population test, training on the other 8 sources and testing on BUS-BRA with no BUS-BRA labels ever seen.

The oracle analysis is the clearest evidence for the core idea. At the same patch size, continuous labels beat binarised ones by a wide margin, and finer patches push the ceiling further still.

In-domain, the actual 5.53-million-parameter decoder-free model lands within one to two Dice points of ViT-B+FPN baselines that are roughly 15 times its size.

That gap all but disappears once you look at compute instead of accuracy.

On the cross-population test, training on the other 8 sources and testing on BUS-BRA with no BUS-BRA labels ever seen, the decoder-free model holds a Dice of 0.7915. That is a meaningful drop from its in-domain score, as expected for any model moved to an unseen population, but the model still produces a usable segmentation without ever training on that population's data.

What next

The core benefit here is a much smaller model with almost no accuracy cost and a real speed advantage. On raw throughput, the decoder-free model already beats a ResNet50 with an FPN decoder while using roughly a fifth of the parameters. For point-of-care hardware, where every extra layer is a real cost in latency and battery life, that trade is close to free.

The pattern behind this result is not specific to breast ultrasound either. Anywhere a segmentation task can tolerate a continuous, soft label instead of a hard one, and can use a fine enough patch grid, the same decoder-free approach should apply. That is a fairly common situation in medical imaging, where ground-truth boundaries are already somewhat uncertain by nature.

The natural next step is testing this same decoder-free recipe on other point-of-care imaging tasks beyond breast ultrasound, and running the cross-population comparison against a conventional encoder-decoder baseline tested the same way, to see whether the two architectures lose accuracy at the same rate when they leave their training population.

Status & links

Manuscript under review at the Journal of Imaging Informatics in Medicine. Read the manuscript.