ERISLab/TokenReductionPT

TokenReductionPT: token reduction and cross-layer cache aggregation for fine-grained recognition

These are the checkpoints behind Cross-Layer Cache Aggregation for Token Reduction in Ultra-Fine-Grained Image Recognition (arXiv:2501.00243), ICASSP 2025, and behind its extension, in preparation. Vision transformers are fine-tuned with a token-reduction method (EViT, TopK, DynamicViT, ATS, ToMe, PatchMerger, SiT, DPC-KNN and the FGIR-aware MAWS, DMAWS, EDAR, NFEDAR and GLSF criteria), with and without cross-layer cache aggregation (CLCA). CLCA adds a cross-layer aggregation head (CLA) and a cross-layer cache (CLC) that keeps information from discarded tokens. Code: arkel23/TokenReductionPT.

2656 checkpoints, one per configuration, each the last epoch of one training run. Each file is a torch.save dict with config (the full training configuration), model (the state dict), accuracy and epoch, with no optimizer state. File names are the runs' experiment-log names: dataset, model (reduction method as prefix), keep rate, cla for the aggregation head, clc_1 for the cache, and the serial. Load them with fgir-zoo. The collection groups this repo with the ICASSP 2025 paper.

Layout

One folder per serial. setting is bl (no CLCA), cla (aggregation head only) or clca (head and cache). "Extension" is the extension of the ICASSP 2025 paper, in preparation.

manifest.csv lists every file with its dataset, model, method, backbone, keep rate, setting, serial, seed, image size, class count, accuracy, the papers that use it, SHA-256 and size. The 224 px ablations of the ICASSP 2025 paper have no released checkpoints.

Load a checkpoint and classify an image

import torch
from PIL import Image
from torchvision import transforms
from fgir_zoo import tokenred

model = tokenred.create_model('serial_62/cub_evit_deit_tiny_patch16_224.fb_in1k_0.7_cla_clc_1_62')
cfg = model.config
# evaluation at 448 px resizes to a square without cropping
tf = transforms.Compose([
    transforms.Resize((cfg.input_size, cfg.input_size),
                      interpolation=transforms.InterpolationMode.BICUBIC),
    transforms.ToTensor(),
    transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
])
# a CUB-200-2011 test image, class index 50 (051.Horned_Grebe)
x = tf(Image.open('Horned_Grebe_0050_34561.jpg').convert('RGB')).unsqueeze(0)
with torch.no_grad():
    logits = model(x)
print(logits.argmax(-1).item(), logits.softmax(-1).max().item())  # 50 0.9315

Accuracy of the released checkpoints

Top-1 accuracy (%) is the run's own test accuracy after the last epoch, as logged in training. The papers average three seeds; where a configuration ran with several seeds, the released file is the seed with the highest accuracy, so the papers' tables can differ slightly. Per-file accuracy is in manifest.csv; 1 file has no recorded accuracy (serial_16/cub_edar_deit_small_patch16_224.fb_in1k_0.7_16). The table below covers serial_30, the ICASSP 2025 main results; method none is the plain ViT without token reduction.

Requirements

• —fgir-zoo (pip install git+https://github.com/arkel23/fgir-zoo.git), which pins timm==0.9.12
• —torch (checked with 2.5.1)

Citation

@inproceedings{rios_cross-layer_2025,
  title         = {Cross-Layer Cache Aggregation for Token Reduction in Ultra-Fine-Grained Image Recognition},
  author        = {Rios, Edwin Arkel and Yuanda, Jansen Christopher and Ghanz, Vincent Leon and
                   Yu, Cheng-Wei and Lai, Bo-Cheng and Hu, Min-Chun},
  booktitle     = {IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP)},
  year          = {2025},
  eprint        = {2501.00243},
  archivePrefix = {arXiv},
  primaryClass  = {cs.CV},
  doi           = {10.48550/arXiv.2501.00243},
  url           = {https://arxiv.org/abs/2501.00243}
}