ICD & AICD: saliency-guided masking augmentation for vision models (method preprint + PyTorch impl)

vision
1

Hi all,

I wanted to share two data augmentation methods we’ve been working on, ICD (Intelligent Coarse Dropout) and AICD (Anti-ICD), together with a reference PyTorch implementation.

The idea. Standard masking augmentations like Cutout or Random Erasing remove rectangular regions at random, so they erase background as often as the object. ICD and AICD instead condition the masking on the model’s own Grad-CAM saliency:

  • ICD masks the highest-saliency tiles (the regions the model already relies on), forcing it to learn from other cues.
  • AICD masks the lowest-saliency tiles (background, irrelevant texture), sharpening focus on the object.

Both split the image into a tile grid, score each tile by mean saliency, threshold by a percentile, and fill masked tiles with a soft strategy (blur, local mean, or noise) instead of a hard black box.

Implementation. It builds on pytorch-grad-cam and drops into a standard training loop:

from bnnr.icd import ICD
from bnnr.xai_cache import XAICache

# precompute saliency once (after a short warm-up), on the same loader
# whose dataset yields (image, label, sample_index) tuples
cache = XAICache("./xai_cache")
cache.precompute_cache(
    model=model,
    train_loader=train_loader,
    target_layers=[model.layer4[-1]],
    method="gradcam",  
)

icd = ICD(
    model=model,
    target_layers=[model.layer4[-1]],
    cache=cache,
    explainer="gradcam",          # keep in sync with cache method (used on cache miss)
    threshold_percentile=70.0,
    tile_size=8,
    fill_strategy="gaussian_blur",
)

# in the training step, on a uint8 NHWC batch (numpy);
# `index` are the per-sample dataset indices for cache lookup
aug = icd.apply_batch_with_labels(images, labels, sample_indices=index)

Status. This is a method preprint: formal mask construction, fill strategies, and hyperparameters. No benchmark numbers yet. A controlled evaluation is the next paper. I’d really appreciate feedback on the formulation or the implementation.

Thanks!