Reproducing CASS on Brain Tumor MRI Scans

This was a team project for CS598 (Deep Learning for Healthcare) at UIUC, done with Hanfei Deng and Stefan Wang. We tried to reproduce CASS, a self-supervised method that trains a CNN and a Vision Transformer side by side, and see how it held up on a brain tumor MRI dataset the original authors never touched.

Why this paper

Medical imaging is one of those areas where deep learning should help a lot, but there's rarely enough labeled data to make it work. Labeling needs a radiologist, patient privacy makes datasets hard to share, and for rare conditions there just aren't enough cases to learn from. Self-supervised methods like DINO get around the labeling problem by learning from raw images, but they're expensive to train — big batch sizes, long runs — which puts them out of reach if you don't have a lot of compute sitting around.

CASS (Singh & Cirrone, 2023) tries to fix that by training a CNN and a ViT at the same time and having them learn from each other, instead of using two copies of the same architecture the way DINO does. CNNs are good at local detail, ViTs are good at global context, and if you push their embeddings toward each other during training, each one seems to pick up something from the other. The authors reported CASS beating DINO on F1 and Recall, and doing it in roughly a third of the training time. We wanted to see if that held up on a dataset they hadn't tried.

CASS vs. DINO

CASS contrasts a CNN (R) and ViT (T) directly; DINO uses a student/teacher pair of the same architecture.

Self-Supervised Loss

Negative cosine similarity between the CNN embedding (R) and ViT embedding (T) for the same image.

What we were checking

We focused on two of the paper's claims:

• Hypothesis 1: CASS holds up better than DINO when you change the batch size or the number of pretraining epochs.
• Hypothesis 2: CASS trains faster than DINO.

To check these, we ran a batch-size ablation, a pretraining-epoch ablation, and a training-time comparison, all on the Brain Tumor MRI Dataset — glioma, meningioma, pituitary, no tumor — which the original CASS paper never evaluated on.

~7,000 MRI images, roughly balanced across 4 classes, split 81/19 train/test.

How it works

• Pretrain a CNN and ViT in parallel with the CASS self-supervised loss — no labels at this stage.
• Fine-tune each network separately for classification, using class-distribution-normalized Focal Loss since the dataset has a bit of class imbalance.
• Pretraining: Adam, learning rate 1e-3, cosine schedule, batch size 16 (we varied this for the ablations).
• Fine-tuning: Adam, learning rate 3e-4, cosine schedule, 50 epochs, batch size 16.
• Hardware: one NVIDIA V100 (16GB) on Google Cloud. About 3 minutes an epoch for pretraining, roughly 2 hours per fine-tuning run.

What we found

Batch size: 8, 12, 16 (100 pretrain epochs, 50 fine-tune epochs)

Left: our F1 scores. Right: the authors' F1 scores, for reference.

We wanted to test batch size 32 too, but our V100 only had 16GB of memory, which wasn't enough to fit it. Our CNN got better as batch size went up, which is the opposite of what the authors saw — their CNN actually dropped a bit at 32. The ViT matched their pattern though: better at small batch sizes, and a lot more stable than the CNN across the board (an 8% swing for the ViT vs. 40% for the CNN). So this one doesn't really back up Hypothesis 1 — our CNN was anything but robust to batch size — though the ViT being steadier than the CNN matches what the authors found.

Pretraining epochs: 50, 100, 200 (batch size 16, 50 fine-tune epochs)

Left: our F1 scores. Right: the authors' F1 scores, for reference.

This one held up better. Both our CNN and the authors' actually dipped at 200 epochs rather than continuing to improve, and both ViTs got steadily better with more pretraining. Biggest swings were 25% for the CNN and 15% for the ViT — smaller than what we saw across batch sizes, so CASS does seem a bit more forgiving about epoch count than about batch size.

Training time vs. DINO

We never got around to running DINO ourselves — not enough time — so we estimated its training time on our hardware based on what the authors reported. Even with that estimate, CASS pretraining came out to roughly a quarter of DINO's time. Out of the two hypotheses, this is the one that clearly held up.

One more check: 10% of the labels

Just to see what would happen, we also fine-tuned CASS (100 pretraining epochs, 50 fine-tuning epochs) using only 10% of the training labels. F1 dropped a lot — 0.4784 for the CNN, 0.2415 for the ViT — but that's still well above random guessing, so the model is learning something useful even when labels are scarce, just a lot less of it.

What didn't line up

Our F1 scores were lower than the authors' across the board, sometimes by a wide margin, even though we matched their hyperparameters and stuck close to their code. We emailed Pranav Singh, the paper's first author, and he looked over what we were doing but couldn't find anything wrong with it. His one suggestion was to use F1 instead of Recall as the training metric — we'd already tried that, with basically the same result. Our best guess is some mix of dataset differences (the authors evaluated on a different dataset than we did) and just not training as long as they did with better hardware, but we never actually pinned down the reason.

Still, the qualitative stuff mostly held: CNN beat ViT almost every time, ViT was steadier than CNN, and CASS trained way faster than DINO. Most of our real headaches weren't about the model at all — they were about Google Cloud. We started on K80s (too slow), tried pairing two of them (still too slow), then switched to V100s and ended up spinning up VMs in whatever region actually had one free, including a few in Europe, because availability was that inconsistent. We never landed an A100 despite trying multiple times a day for several days straight. Once we had something running reliably, we got through all the ablations, up to 200 epochs.

We don't have much to add for the original author — he was helpful and the setup instructions were clear. We're a little disappointed we couldn't reproduce his numbers, but the directional findings held up well enough that CASS still seems like a promising idea, if you can actually get it dialed in.

Code & repo

Code, notebook, and full report are on GitHub. The original CASS paper is on arXiv, and the authors' reference implementation is at pranavsinghps1/CASS. We also put together a short video walking through our approach and results.