Learning Features, Not Layouts: Feature-Axis Attention as a Transfer Prior for Medical Imaging

Aon Safdar; Mohamed Saadeldin

cs.CVOct 9, 2026

Learning Features, Not Layouts: Feature-Axis Attention as a Transfer Prior for Medical Imaging

Authors: Aon Safdar, Mohamed Saadeldin

Organizations: University College Dublin, Ireland

Abstract

Natural-image-pretrained foundation models are increasingly adapted to medical imaging, yet transfer remains unreliable under domain shift, modality heterogeneity, and label scarcity. We trace part of this fragility to an overlooked design choice. Standard vision transformers place all global modeling on self-attention over spatial tokens, an object-layout prior well suited to natural scenes but not to medical images, where the discriminative signal lies in feature correlations and texture that recur across scanners while layout does not. We ask whether the axis of global attention, spatial versus feature, is itself a lever for transferable medical representations. Existing channel-aware hybrids confound this question, adding channel attention to a spatially-mixed backbone and leaving the axes coupled and the feature axis unmeasurable in isolation. We design FAD-ViT, a backbone that strictly decouples the two axes, making feature-axis attention the primary global operator and the attention axis a controlled variable. Evaluating comprehensively across 23 classification, segmentation, and low-data tasks under matched tokenization, capacity, and protocol, we find that shifting global attention from the spatial to the feature axis leaves from-scratch performance essentially unchanged, yet markedly strengthens transfer. Identical ImageNet pretraining yields up to roughly three times the transfer gain, with consistent improvements on every held-out segmentation task at up to a quarter the parameter count and far lower compute than standard and channel-attention baselines. The advantage is largest where texture rather than geometry is discriminative, tracking greater reuse of pretrained features. The axis of global attention is thus a design lever for medical transfer, and making it channel-primary is a simple, cheap, and principled choice.

Figures & tables

Figure 1: Overview of the studied block. Spatial Feature Extraction (SFE) preserves the 2D grid with depthwise convolution; Feature Self-Attention (FSA) transposes tokens ( $B{\times}N{\times}C\rightarrow B{\times}C{\times}N$ ) and attends over feature tokens. Global interaction thus scales with width $C$ rather than token count $N$ .

Figure 2: Study design. Under a fixed protocol (convolutional tokenization, matched $\sim$ 23M-parameter capacity, shared U-Net / linear-head decoders, identical schedule, and ImageNet-1K pretraining with a from-scratch control) we vary only the global-mixing operator across four backbone families: a spatial-axis baseline (ViT/DINOv2), two channel-hybrids (XCiT, DaViT), and the strictly decoupled feature axis. Transfer is evaluated on classification, segmentation, and low-data probing and attributed with transfer-gap, CKA, modality, and multi-seed analyses.

Dataset	CvT-13		ViT-S/16		FAD-ViT (feat-axis)
Dataset	Acc	AUC	Acc	AUC	Acc	AUC
PathMNIST	98.50	99.90	99.00	99.90	99.40	99.90
DermaMNIST	73.77	91.32	76.50	93.00	77.20	90.70
OCTMNIST	91.36	97.27	94.39	98.42	93.44	98.30
RetinaMNIST	54.17	74.46	56.67	78.89	57.50	79.54
BreastMNIST	67.52	77.95	87.18	93.60	89.74	92.54

Table 1: MedMNIST-2D classification transfer (ImageNet-1K pretraining). Test Accuracy/AUC (%). Bold : best; underline : second.

Encoder	Params	BUSI		B
	(M)	Acc	AUC	Acc	AUC
ResNet50	23.5	85.64	93.23	78.33	90.05
VGG16	134.3	85.47	93.69	80.00	88.24
ViT-S/16	21.7	86.50	93.60	82.50	90.68
ViT-S/32	22.5	84.10	93.59	79.17	87.55
CvT-13	20.1	75.50	88.90	75.22	89.17

Table 2: Breast-ultrasound transfer on BUSI and B (shared fine-tuning). Params in millions; Accuracy/AUC (%). Bold : best; underline : second.

Figure 3: Accuracy–efficiency trade-off under the matched transfer protocol. (a) Segmentation transfer: faint points are the 11 per-task Dice scores per encoder, large markers the mean; (b) BUSI classification accuracy versus parameters (log scale) across eight encoders spanning $20\text{--}134$ M. Feature-axis FAD-ViT ( $\star$ ) is Pareto-dominant in both, matching or exceeding encoders up to $6\times$ larger (DINOv2-B, VGG16; all evaluated at the same $224$ px input) at a fraction of the parameters.

Dataset	Model	Dice	mIoU	HD95	Dataset	Model	Dice	mIoU	HD95
BUSI	ViT-UNet	53.1	43.9	47.7	Lung	ViT-UNet	95.9	95.9 †	0.1 †
	DINOv2-UNet	53.7	43.2	50.2		DINOv2-UNet	96.0	96.0 †	0.0 †
	FAD-ViT	77.7	70.5	21.4		FAD-ViT	96.3	96.2	0.0
Brain	ViT-UNet	80.2	76.7	3.3	Pancreas	ViT-UNet	82.9	81.2	2.3
	DINOv2-UNet	77.4	74.0	3.9		DINOv2-UNet	81.5	80.0	2.4
	FAD-ViT	83.0	79.6	2.5		FAD-ViT	86.1	83.9	1.8

Table 3: Segmentation transfer (

256{\times}256

, matched U-Net decoder) on the held-out test split. Dice/mIoU (%)

\uparrow

; HD95 (px)

\downarrow

. Bold : best per task. † Near-saturated binary CT tasks (Lung, Colon) yield degenerate HD95 near zero for several models; see text. ‡ Prostate has a small test set (

n{=}77

slices) with high seed variance (std

15\text{--}18

Dice); we draw no task-level conclusion from it.

Modality (#tasks)	ViT-UNet	DINOv2-UNet	FAD-ViT	$\Delta$ vs ViT
Ultrasound (1)	53.1	53.7	77.7	+24.6
MRI (4)	75.8	71.4	78.6	+2.8
CT (6)	87.0	85.4	89.1	+2.1

Table 4: Segmentation transfer by modality (mean Dice %, matched U-Net decoder). The feature-axis gain over spatial-axis ViT is largest on texture-dominant, low-contrast ultrasound and smallest on high-contrast structural CT.

Figure 4: Qualitative segmentation transfer. Feature-axis FAD-ViT-UNet produces cleaner boundaries and fewer spurious regions than spatial-attention baselines; error maps (TP green, FP red, FN blue) concentrate residual error at ambiguous boundaries.

Backbone	Global operator	Params	Natural	Medical
Backbone	Global operator	(M)	mIoU	BUSI Dice
ViT-UNet	spatial-token only	90.6	41.9	53.1
XCiT-UNet	channel + spatial (columnar)	27.1	43.9	75.1
DaViT-UNet	channel/spatial interleaved	34.4	42.8	77.8
FAD-ViT-UNet	feature-axis only (decoupled)	23.1	42.6	77.7

Table 5: Channel-primary decoupling versus coupled channel attention, matched U-Net decoder. Natural-image competence (ADE20K mIoU) versus medical transfer (held-out BUSI test Dice %). Coupled backbones (XCiT, DaViT) match or beat FAD-ViT on natural images but only tie it on medical transfer, which the decoupled design reaches at the lowest parameter count. Bold : best.

Dataset	FAD-ViT (feat-axis)			ViT-S/16			DINOv2-B
Dataset	1%	10%	100%	1%	10%	100%	1%	10%	100%
PathMNIST	79.8	85.3	86.2	74.6	82.3	84.5	80.0	70.5	80.3
DermaMNIST	66.7	69.2	68.8	66.7	69.2	70.1	66.9	68.8	70.1
OCTMNIST	73.0	76.3	76.5	71.2	75.3	75.8	68.3	74.7	75.7
PneumoniaMNIST	87.1	88.7	89.1	71.2	81.0	80.5	74.4	89.4	89.6
RetinaMNIST	44.2	46.7	44.6	48.8	50.8	52.5	42.1	43.3	46.7

Table 6: Low-data linear probing (frozen backbone). Test Accuracy (%) at 1/10/100% labels (single-seed; 3-seed means in Table S9). Bold : best per row and fraction.

Variant	BUSI Dice	$\Delta$	Heart Dice	$\Delta$
Full FAD-ViT (feat-axis, dw- $Q/K/V$ , +tail)	64.2	–	82.6	–
$\rightarrow$ swap to spatial-axis attn. (23.1M)	63.1	$-1.1$	80.8	$-1.8$
$\rightarrow$ remove spatial tail (16.4M)	63.1	$-1.1$	81.3	$-1.3$

Table 7: Medical ablation on BUSI and Heart, from-scratch Dice (%), matched decoder and schedule.

\Delta

is relative to the full model on the same task. From scratch the axis swap and the spatial tail each cost only

1\text{--}2

Dice at identical or lower capacity; the decisive axis effect appears under transfer ( Tab. 8 ), not here. Bold : full model.

Backbone	Params (M)	Scratch	Pretrained	Gap
Feature-axis FAD-ViT	23.1	61.3	76.4	+15.1
Spatial ViT-S (matched)	22.9	48.9	53.2	+4.3
Spatial ViT-B	90.6	49.9	55.7	+5.7
Convolutional ResNet50	71.7	61.7	67.7	+6.0

Table 8: Transfer-specificity on breast ultrasound (held-out test Dice %): from scratch vs. ImageNet-1K pretraining, matched U-Net decoder and schedule. Gap

=

pretrained

-

scratch. Feature-axis, ViT-S, and ViT-B values are averaged over seeds (ViT-S over three); ResNet50 is single-seed.

Modality (task)	Feature axis (ours)			Spatial axis
Modality (task)	Scratch	Pretr.	Gap	Scratch	Pretr.	Gap
Ultrasound (BUSI)	61.3	76.4	+15.1	63.1	65.3	+2.2
CT (Spleen)	90.6	95.3	+4.7	91.1	94.3	+3.2
MRI (Heart)	80.9	84.5	+3.6	80.5	82.2	+1.7

Table 9: Cross-modality transfer gap under the strict in-place axis swap (identical

23.1

M skeleton, only the global-mixing axis flipped). Held-out test Dice (%); Gap

=

pretrained

-

scratch. CT/MRI are multi-seed means (std

\leq\!2.8

); ultrasound reproduces the primary run ( Tab. 8 ). The feature axis wins the gap on every modality, most on texture-dominated ultrasound.

Figure 5: Feature reuse tracks the transfer gap. (a) Linear CKA between the ImageNet-pretrained and from-scratch encoders on BUSI (higher = more reuse): the feature axis reuses more early/mid structure and specializes hardest late, whereas the spatial axis decays monotonically. (b) Their transfer gap (pretrained $-$ scratch Dice): higher mid-level reuse coincides with a larger gain from identical pretraining.

Explore similar work

Mar 1, 2026

LUMOS: Latent Universal Medical Priors for Segmentation

General vision foundation models (VFMs) have been primarily developed on natural images, and their utility for medical image segmentation is therefore often considered to depend on costly adaptation or domain-specific fine-tuning. In this paper, we revisit this assumption from a different perspective: rather than requiring VFM segmentors to relearn visual regularities, we investigate whether the low-level visual priors necessary for anatomical delineation already lie dormant within general VFMs. We observe that frozen VFMs, despite lacking medical supervision, encode transferable visual regularities. These properties are not exclusive to natural images but are also fundamental to medical image understanding. Motivated by this observation, we propose Latent Universal Medical PriOrs for Segmentation (LUMOS), a novel framework that amplifies general VFM priors to conventional medical segmentors. LUMOS consists of two key components: (1) Pathfinder that distills visual cues from a frozen vision foundation model, and (2) Inspiror that sparks the conventional medical networks with spatial guidance from distilled visual regularities. In this way, the segmentor is relieved from learning complex visual regularities entirely from limited medical annotations and can instead focus on task-specific anatomical delineation. Across diverse medical datasets and token-based VFMs, LUMOS shows that general VFMs can serve as spatial prior generators when their frozen token spaces preserve patch-level pattern relevance. DINO provides stable matched-backbone gains, while SigLIP exposes VFM-specific sensitivity caused by its different token granularity and representation objective.

Zhuonan Liang, Wei Guo, Jie Gan +4

The University of Sydney, Sydney, NSW 2006, Australia. · Biological System & Engineering Division, Lawrence Berkeley National Laboratory, Berkeley, CA 94720, USA. · Berkeley Biomedical Data Science Center, Lawrence Berkeley National Laboratory, Berkeley, CA 94720, USA.

Medical Image Segmentation Transfer Learning Vision Foundation Model Adaptation

Ask PDF

Aug 2, 2026

Location-Aware Fine-Grained Representation Learning for Medical Vision Foundation Models

Fine-grained visual representations are essential for medical image analysis, particularly when diagnostically relevant evidence is subtle and spatially localized. Modern transformer-based medical vision encoders must therefore learn patch-level representations that are both clinically meaningful and spatially consistent. Without these properties, large vision-language models (LVLMs) operate on an ambiguous visual foundation, limiting their ability to generate clinically reliable and spatially grounded responses. However, existing training strategies for medical vision encoders rarely achieve both objectives. Image-text alignment provides clinically meaningful supervision primarily at the image level, leaving the spatial localization of diagnostic evidence weakly constrained. In contrast, self-supervised learning promotes spatial consistency but lacks the semantic supervision needed to distinguish visually similar yet clinically distinct regions. To address this gap, we present LoFi, a medical vision foundation model built on location-aware fine-grained representation learning. LoFi trains a vision encoder with a lightweight large language model under grounding and grounded captioning objectives. Because these objectives require predicting location from clinical text and vice versa, spatial consistency emerges without any explicit patch-level regularization. To enable training at scale, we construct MedG, a large-scale medical grounding dataset of 4.48M image-text-box triplets curated from 84 datasets spanning 7 modalities. Across phrase grounding, visual question answering, and region-based organ classification under perturbations, LoFi consistently outperforms general-purpose and medical vision foundation models as well as state-of-the-art LVLMs. Code is available at https://github.com/myeongkyunkang/lofi-medg.

Myeongkyun Kang, Yanting Yang, Xiaoxiao Li

The University of British Columbia · Vector Institute

Medical Imaging Foundation Models Medical Image Grounding Visual Grounding Medical Image Representation Learning Fine-Grained Visual Perception

Ask PDF

Feb 20, 2026

ZACH-ViT: Regime-Dependent Inductive Bias in Compact Vision Transformers for Medical Imaging

Vision Transformers rely on positional embeddings and class tokens encoding fixed spatial priors. While effective for natural images, these priors may be suboptimal when spatial layout is weakly informative, a frequent condition in medical imaging. We introduce ZACH-ViT (Zero-token Adaptive Compact Hierarchical Vision Transformer), a compact Vision Transformer that removes positional embeddings and the [CLS] token, achieving permutation-invariant patch processing via global average pooling. Zero-token denotes removal of the dedicated aggregation token and positional encodings. Patch tokens remain unchanged. Adaptive residual projections preserve training stability under strict parameter constraints. We evaluate ZACH-ViT across seven MedMNIST datasets under a strict few-shot protocol (50 samples/class, fixed hyperparameters, five seeds). Results reveal regime-dependent behavior: ZACH-ViT (0.25M parameters, trained from scratch) achieves strongest advantage on BloodMNIST and remains competitive on PathMNIST, while relative advantage decreases on datasets with stronger anatomical priors (OCTMNIST, OrganAMNIST), consistent with our hypothesis. Component and pooling ablations show positional support becomes mildly beneficial as spatial structure increases, whereas reintroducing a [CLS] token is consistently unfavorable. These findings support that architectural alignment with data structure can outweigh universal benchmark dominance. Despite minimal size and no pretraining, ZACH-ViT achieves competitive performance under data-scarce conditions, relevant for compact medical imaging and low-resource settings. Code: https://github.com/Bluesman79/ZACH-ViT

Athanasios Angelakis

BioML Lab, RI CODE, UniBw, Munich, Germany · EDS, Amsterdam UMC, Amsterdam, Netherlands

Vision Transformer Efficient ViTs Medical Imaging

Ask PDF