ZACH-ViT: A Zero-Token Vision Transformer with ShuffleStrides Data Augmentation for Robust Lung Ultrasound Classification

AI summary

Overview

Research area: Medical image analysis / computer vision for critical care — deep learning applied to lung ultrasound video classification, with an architectural contribution to Vision Transformer design.

Technical level: Advanced. The paper assumes familiarity with Vision Transformers, positional embeddings, residual connections, attention, and multiple-instance learning, and it includes formal notation for permutation invariance and per-layer transformer operations.

Scope: The paper proposes ZACH-ViT, a Vision Transformer variant with no positional embeddings and no class token, plus a structured augmentation scheme (ShuffleStrides Data Augmentation) for classifying cardiogenic pulmonary oedema versus a heterogeneous non-cardiogenic class in 380 lung ultrasound videos from 95 critically ill patients.

What This Paper Is About

In intensive care, doctors need to tell cardiogenic pulmonary oedema (fluid build-up driven by heart failure) apart from non-cardiogenic lung problems, because the treatments differ. Lung ultrasound is the standard bedside tool, but its interpretation is qualitative and varies between operators and centres, which makes automated reading desirable but hard.

The authors' classification problem is deliberately made harder than a simple two-way split: their non-cardiogenic group (Class 0) mixes non-cardiogenic inflammatory pathology (NCIP/ARDS-like patterns), interstitial lung disease (ILD), and structurally normal or healthy lungs, all of which share overlapping B-line and pleural artefacts. The goal is a model that stays stable and generalises on a small, heterogeneous, domain-shifted clinical dataset where conventional deep architectures fail.

Key Contributions

1. ZACH-ViT architecture (Zero-token Adaptive Compact Hierarchical Vision Transformer): a Vision Transformer variant that removes positional embeddings and the [CLS] token entirely, making it permutation-invariant, and replaces fixed-dimension residuals with adaptive residual projections that project features to matching dimensions when they differ.

2. ShuffleStrides Data Augmentation (SSDA): a structured augmentation framework that permutes probe-view sequences and frame orderings while keeping the augmented samples anatomically valid, with a baseline regime of 24 probe-view permutations and an extended regime using ten prime-number random seeds (2, 3, 5, 7, 11, 13, 17, 19, 23, 29) for intra-view frame shuffling and up to twelve-fold data expansion.

3. A harder, more realistic benchmark formulation: distinguishing cardiogenic pulmonary oedema from a composite non-cardiogenic class containing NCIP/ARDS-like pathology, ILD, and healthy lungs, rather than the simpler binary CPE–NCIP separation.

4. Open-source release: the full implementation, preprocessing and training pipelines are published at https://github.com/Bluesman79/ZACH-ViT and installable via pip install zachvit.

Main Findings

• ZACH-ViT was the only model that learned non-trivial representations. It reached validation and test ROC-AUC of 0.80 and 0.79 with sensitivity 0.60 and specificity 0.91 (test), while all nine baselines collapsed to trivial classification (sensitivity = 0.00, specificity = 1.00) under the structured SSDA regimes.

• Detailed validation/test metrics for ZACH-ViT (threshold 0.50): validation — specificity 0.85, sensitivity 0.60, accuracy 0.72, F1 0.60, ROC-AUC 0.80; test — specificity 0.91, sensitivity 0.60, accuracy 0.75, F1 0.67, ROC-AUC 0.79.

• Results varied by augmentation regime. Under VIS, ZACH-ViT's own ROC-AUC was 0.69/0.70 with sensitivity 0.00 and specificity 1.00 (trivial). Under 0-SSDA it reached 0.87/0.79 with sensitivity 0.60/0.40 and specificity 0.69/0.64. Under 0_2-SSDA it reached 0.80/0.79 with the best balance (sensitivity 0.60/0.60, specificity 0.85/0.91, F1 0.60/0.67).

• Baselines failed across the board. All other models ranged from 0.37–0.70 (validation) and 0.33–0.68 (test) ROC-AUC under VIS, 0.46–0.71 / 0.45–0.70 under 0-SSDA, and 0.48–0.70 / 0.47–0.69 under 0_2-SSDA, always with sensitivity 0.00 and specificity 1.00. The single exception noted in the text is ABMIL, which reached moderate generalisation (AUC 0.70 validation, 0.68 test).

• Peak validation ROC-AUC (Table 5): ZACH-ViT 0.70 (VIS), 0.93 (0-SSDA), 0.94 (0_2-SSDA); ABMIL 0.71 / 0.70 / 0.70; ResNet50 0.65 / 0.54 / 0.54; Minimal ViT 0.62 / 0.71 / 0.71; Swin-Tiny 0.68 / 0.65 / not reported; TransMIL 0.62 / 0.65 / 0.65; others below 0.65 under VIS and 0.50 under both SSDA regimes. (The Section 3.3 text instead reports a peak validation ROC-AUC of 0.87 for 0-SSDA; the two figures differ between the text and Table 5 as published.)

• Parameter efficiency. ZACH-ViT has 0.25 M parameters versus 0.62 M for the Minimal ViT configuration, a 60% reduction, while achieving a higher validation ROC-AUC (0.80 vs. 0.58). Table 7 lists the Minimal ViT at 0.61 M, and Table 1 lists it at 0.62 M. For context, ResNet50 has 23.85 M, ConvNeXt-Tiny 27.92 M and Swin-Tiny 28.29 M parameters.

• Speed. ZACH-ViT trains 1.35× faster than Minimal ViT with 2.5× fewer parameters. Table 6 reports ZACH-ViT at 2.03 minutes training time and 60 ms/batch inference, versus Minimal ViT at 2.75 minutes and 201 ms/batch.

• Heavier augmentation eventually hurt. ZACH-ViT held up to 0_2_3-SSDA but showed mild overfitting from 0_2_3_5-SSDA onward, where validation sensitivity peaked at 1.00 while test sensitivity fell to 0.40. At 0_2_3_5-SSDA the validation figures were sensitivity 1.00, specificity 0.62, ROC-AUC 0.86, F1 0.67, and test figures were sensitivity 0.40, specificity 0.73, AUC 0.76, F1 0.40. At 0_2_3_5_7-SSDA validation dropped to sensitivity 0.40, specificity 0.69, AUC 0.66, while test stayed at sensitivity 0.40, specificity 0.91, AUC 0.78. At the most extreme SSDA_10 configuration, validation overfit (specificity 0.62) but test discrimination remained balanced (sensitivity 0.60, specificity 0.91, AUC 0.79, F1 0.67).

• MIL and efficient-transformer baselines still failed. ABMIL and TransMIL are designed for unordered sets, and efficient transformers such as MobileViT and TinyViT reduce parameter counts, but the authors report that these retain positional embeddings and [CLS] tokens or remain sensitive to spurious frame-level variation, and they collapsed under SSDA like the rest.

Methodology in Plain English

The study is retrospective, using lung ultrasound data from 95 critically ill patients admitted to Amsterdam University Medical Centers between 2016 and 2020, approved under ethics reference 2021.0102. Each patient contributed up to four videos of anterolateral thoracic regions, giving 380 videos total, recorded on a FUJIFILM SonoSite Edge II at 30 frames per second with a convex probe. The data were split by patient into 61 for training, 18 for validation, and 16 for testing, with at least five cardiogenic cases in each evaluation set. The reported class distribution is 43 non-cardiogenic and 18 cardiogenic in training, 13 and 5 in validation, and 11 and 5 in test.

Each frame was converted to grayscale, normalised to the [0,1] range, rescaled to 224 × 224 pixels, had all pixels below intensity 93 (on the 0–255 scale) zeroed to suppress background noise, and was cropped to a region of interest centred on the pleural line, retaining high-frequency features such as B-lines, pleural thickening and subpleural consolidations.

Because videos differ in frame count and probe position, the authors defined five ways to turn a video into a 2D image: VI (horizontal concatenation of all frames), SVI (same but with shuffled frame order), VIS (vertical stacking of four probe-view videos into a stride image), 0-SSDA (VIS plus all 24 permutations of the four probe views), and 0_2-SSDA (0-SSDA plus intra-view frame shuffling across the ten prime-number seeds). The key idea is that every permutation stays anatomically plausible while creating extra training variety.

The model itself is a compact transformer. Input patches are extracted and reshaped into a 196 × 768 representation, reduced by a Dense layer to 128 dimensions, with no positional embeddings added. Each transformer block applies layer normalisation, multi-head attention, and a dense layer; whenever the residual path's feature dimension does not match the branch output, a dynamic projection layer adjusts it before the addition. Instead of a learnable [CLS] token, the final representation is a global average pooling over patches. The authors argue this matches ultrasound, where the diagnostic signal is local texture plus a clinical reading of which view was captured, not a fixed spatial layout.

Training used the Adam optimiser with a learning rate of 1×10⁻⁴, binary cross-entropy loss, class weights where needed for imbalance, and early stopping on validation loss with a maximum of 23 epochs. Results were averaged over three random seeds at a classification threshold of 0.50. The experiments ran on Ubuntu Linux (kernel 5.15.0-67-generic) with Python 3.10.16, TensorFlow 2.19.0 and PyTorch 2.3.1+cu121 on a single NVIDIA GeForce RTX 3060 GPU, although Table 6 reports efficiency measurements on a single NVIDIA RTX 3090 GPU. (Section 2.9 mentions TensorFlow 2.11 and mixed-precision training, which differs from the version given in Section 3.1.)

Nine baselines were trained under identical preprocessing, hyperparameters and early stopping, with hyperparameters chosen by grid search on the validation set: ResNet50, DenseNet121 and EfficientNetB0 (CNNs); Minimal ViT, Swin-Tiny and ConvNeXt-Tiny (transformers); and ABMIL, CNN-ABMIL and TransMIL (multiple-instance learning). Metrics were sensitivity, specificity, accuracy, F1-score and ROC-AUC.

Why This Matters

The result challenges a common assumption in medical AI: that more capacity or more generic architecture is the route to better performance. Here, a 0.25 M-parameter model with positional embeddings and the class token deliberately removed outperformed models with two orders of magnitude more parameters, because its inductive bias matches the structure of the data. If that pattern holds, it has implications for how small clinical datasets should be modelled and how augmentation should be designed (here, constrained by anatomical validity rather than applied arbitrarily).

Real-world applications:

• Bedside decision support in intensive care, helping distinguish cardiogenic from non-cardiogenic pulmonary oedema when treatment pathways diverge.
• Deployment on resource-limited or real-time clinical hardware, given the reported 1.35× faster training than Minimal ViT, 2.5× fewer parameters, and 60 ms/batch inference.
• Extension to other ultrasound modalities the authors name, such as cardiac and abdominal imaging.
• Transfer to non-ultrasound domains with unordered image sets, including multi-view satellite mosaics, bag-of-patches histopathology, robotic tactile sensing arrays, and computed tomography-based representation learning.

Industry relevance: medical device and ultrasound vendors looking for lightweight on-device inference, hospital IT teams with constrained compute, and health-AI developers who need reproducible preprocessing and training pipelines (the code and package are openly released). The authors position architectural parsimony as a practical advantage for real-time clinical deployment, not just an academic curiosity.

Future Directions

• Multi-class or hierarchical modelling. The paper frames the task as binary even though Class 0 contains distinct subtypes (NCIP/ARDS-like, ILD, healthy lungs). The authors call for explicitly modelling that subtype structure.
• Multicentre validation. The dataset comes from a single centre

Authors’ abstract

Differentiating cardiogenic pulmonary oedema (CPE) from non-cardiogenic and structurally normal lungs in lung ultrasound (LUS) videos remains challenging due to the high visual variability of non-cardiogenic inflammatory patterns (NCIP/ARDS-like), interstitial lung disease, and healthy lungs. This heterogeneity complicates automated classification as overlapping B-lines and pleural artefacts are common. We introduce ZACH-ViT (Zero-token Adaptive Compact Hierarchical Vision Transformer), a 0.25 M-parameter Vision Transformer variant that removes both positional embeddings and the [CLS] token, making it fully permutation-invariant and suitable for unordered medical image data. To enhance generalization, we propose ShuffleStrides Data Augmentation (SSDA), which permutes probe-view sequences and frame orders while preserving anatomical validity. ZACH-ViT was evaluated on 380 LUS videos from 95 critically ill patients against nine state-of-the-art baselines. Despite the heterogeneity of the non-cardiogenic group, ZACH-ViT achieved the highest validation and test ROC-AUC (0.80 and 0.79) with balanced sensitivity (0.60) and specificity (0.91), while all competing models collapsed to trivial classification. It trains 1.35x faster than Minimal ViT (0.62M parameters) with 2.5x fewer parameters, supporting real-time clinical deployment. These results show that aligning architectural design with data structure can outperform scale in small-data medical imaging.