Calibration Attention (CalAttn)

This repository reproduces the experiments described in the paper: Calibration Attention: Representation-Conditioned Temperature Scaling for Vision Transformers. It supports:

• Datasets: CIFAR-10/100, MNIST, Tiny-ImageNet, ImageNet-1K
• Backbones: ViT-224, DeiT-S, Swin-S (via timm)
• Methods: CE, Brier, CE+BS, MMCE, Label Smoothing, Focal (FLSD-53-style), Dual Focal (DFL), Relaxed Softmax, SATS (post-hoc), and CalAttn (joint training).
• Post-hoc temperature scaling (TS): grid-search T in {0.1,0.2,...,10.0}, chosen by validation ECE.

• Full calibration metrics:

  • ECE, MCE, AdaECE, Classwise-ECE, Smooth-ECE, NLL, Brier, AUROC (OOD)

---

1. Installation

git clone https://github.com/yourname/calattn-repro.git
cd calattn-repro

python -m venv venv
source venv/bin/activate

pip install -r requirements.txt

requirements.txt contains:

torch>=2.2
torchvision>=0.17
timm>=0.9.16
numpy>=1.24
pyyaml>=6.0
tqdm>=4.66
scikit-learn>=1.3

---

2. Dataset Preparation

CIFAR-10 / CIFAR-100 / MNIST

Automatically downloaded.

Tiny-ImageNet

Download from:

http://cs231n.stanford.edu/tiny-imagenet-200.zip

Unzip into:

calattn-repro/data/tiny-imagenet-200/

ImageNet-1K

Set the dataset root in your config under data.data_dir and use dataset: imagenet1k.

Expected structure:

/path/to/imagenet/
 ├── train/
 └── val/

---

3. Training

Example: CIFAR-100, DeiT-Small + CalAttn

python src/train.py --config configs/cifar100_deits_calattn.yaml

Baseline (CE+Brier only, no CalAttn)

python src/train.py --config configs/cifar100_deits_ce_brier.yaml

SATS (post-hoc)

python src/calibrate_sats.py \
  --config configs/cifar100_deits_calattn.yaml \
  --ckpt your/path/to/outputs/cifar100_deits_calattn/seed0/best.pt

---

4. Post-hoc Temperature Scaling (TS)

python src/calibrate_ts.py \
  --config configs/cifar100_deits_calattn.yaml \
  --ckpt your/path/to/outputs/cifar100_deits_calattn/seed0/best.pt

This selects ( T^* \in {0.1, 0.2, ..., 10.0} ) to minimize validation ECE.

---

5. Evaluation

python src/eval.py \
  --config configs/cifar100_deits_calattn.yaml \
  --ckpt your/path/to/outputs/cifar100_deits_calattn/seed0/best.pt

Example Output

{
  "top1": 66.20,
  "nll": 2.51,
  "brier": 0.78,
  "ece": 1.42,
  "mce": 3.10,
  "adaece": 1.47,
  "classece": 0.27,
  "smece": 1.10
}

These values directly correspond to Tables in the paper.

---

6. OoD Robustness

python src/eval_ood.py \
  --config configs/cifar10_vit224_calattn.yaml \
  --ckpt your/path/to/outputs/cifar10_vit224_calattn/seed0/best.pt \
  --c10c_root /data/cifar10c \
  --c10c_severity 5 \
  --c10c_corruptions all \
  --save_json

Outputs AUROC for:

• CIFAR-10 → SVHN
• CIFAR-10 → CIFAR-10-C

---

7. Reproducing Main Tables

Paper Table	Command
ECE / smECE main results	train.py → calibrate_ts.py → eval.py
SATS comparison	calibrate_sats.py → eval.py
Tiny-ImageNet	configs/tinyimagenet_swinS_calattn.yaml
OoD robustness	src/eval_ood.py
λ-sensitivity	loop over configs with different λ
Dirichlet head ablation	not included

---

8. One-Command Reproduction

bash run.sh

This will:

1. Train CIFAR-100 DeiT-S CalAttn for seeds 0/1/2
2. Run TS and SATS on each checkpoint

---

9. Notes for Reviewers

• All experiments follow the protocol in Section 4 of the paper.
• No dataset-specific tuning is applied.
• Results may vary slightly across GPUs (±0.1 ECE).

---