Attention-based MIL Classification with ViT Embeddings

This project performs bag-level binary classification using Multiple Instance Learning (MIL) with Vision Transformer (ViT) embeddings. Each bag contains multiple image instances but only the bag label is provided (no instance-level supervision). The model learns to identify key instances via attention and classify bags accordingly.

---

Features

• Supports bag-level weak supervision
• Uses ViT (Base-16/224) as feature extractor
• Attention MIL for instance importance modeling
• Mixed-precision training (torch.cuda.amp) for efficiency
• Early stopping + checkpoint saving
• Bag-level inference for test data

---

Project Structure

.
├── preprocess_data.py       # Convert raw bags into ViT embeddings
├── train_mil.py             # Train MIL + classifier
├── test_mil.py              # Bag-level inference / evaluation
└── requirements.txt         # Dependencies

---

Installation

pip install -r requirements.txt

---

Workflow

This project consists of three main stages:

---

1. Preprocessing (Feature Extraction)

Run preprocessing to extract ViT CLS embeddings for each image instance within a bag:

python preprocess_data.py

This converts each raw bag into a matrix of shape (K, 768) where:

• K = number of images in the bag
• 768 = ViT CLS embedding dimension

Output saved as pickle:

• bag_data: list of numpy arrays (K × 768)
• labels: bag-level labels (0 or 1)

---

2. Training (Attention MIL)

Train the MIL classifier using instance attention aggregation:

python train_mil.py

Training includes:

• Attention MIL
• BCEWithLogitsLoss
• Adam optimizer
• AMP mixed precision
• Early stopping & checkpointing

Best checkpoint saved to:

weights/mil_model_best.pth

---

3. Testing (Bag-Level Inference)

Run inference on test bags:

python test_mil.py

Inference path:

bag → ViT embeddings → attention MIL → bag embedding → classifier → predicted label

---

Model Design

Instance Embedding (ViT)
Each image → CLS token embedding (768-d)

Attention MIL Aggregation

$$ v = \sum_{i=1}^{K} \alpha_i x_i $$

where:

• $x_i$ = instance embedding
• $\alpha_i = \text{softmax}(f(x_i))$

Bag Classification

MLP → sigmoid → bag label

---

Results

Training Behavior

• Training converged stably

• Validation loss slightly increased around epoch 13, indicating early signs of overfitting

• Early stopping prevented performance degradation

Performance

• Achieved 100% accuracy on the test set

• Model successfully identified key instances within bags to determine bag labels

---