SAGE: Streaming Agreement-Driven Gradient Sketches for Representative Subset Selection

Official implementation of the paper:

SAGE: Streaming Agreement-Driven Gradient Sketches for Representative Subset Selection Ashish Jha, Salman Ahmadi-Asl NeurIPS 2025 Workshop on Reliable ML from Unreliable Data arXiv:2510.02470

SAGE is a streaming data-subset selection method. It maintains a compact Frequent Directions (FD) sketch of gradient geometry in O(lD) memory and prioritizes examples whose sketched gradients align with a consensus direction. This avoids N x N pairwise similarity computations and explicit N x D gradient stores, yielding a simple two-pass, GPU-friendly pipeline whose memory footprint is independent of the dataset size.

Method Overview

1. Phase I (sketch): Stream per-sample gradients through an FD sketch S of size l x D. Each time the sketch fills, an SVD-based shrinkage step compresses it while preserving dominant gradient directions with deterministic error guarantees.
2. Phase II (scoring): Project each sample's gradient to l dimensions (z_i = S g_i), normalize, and score it by cosine agreement with the consensus direction (the mean of the normalized projections).
3. Selection: Keep the top-k samples by agreement score, either globally or per class (CB-SAGE, the class-balanced variant used for imbalanced datasets).

Repository Structure

src/sage/
  core.py            FD streaming sketch, agreement scoring, subset selection
  train_main.py      End-to-end training script with full CLI
  data.py            Dataset loading, transforms, label-corruption utilities
  models.py          Model factory (CIFAR ResNeXt-29, torchvision, timm)
  embedding/train.py Experimental: gradient-free selection from embeddings
  llm/prompt_trimming.py Experimental: agreement-based LLM prompt trimming
scripts/
  train.py           Main training entrypoint (forwards all CLI flags)
  train_imagenet.py  ImageNet entrypoint with sensible defaults
ablations.py         Ablation: projection methods (FD, random projection, PCA)
ablations_sketch_size.py Ablation: sketch size (l)
plot_sketch_size_curves.py Plotting for sketch-size ablation results
notebooks/
  sage_active_learning_eurosat.ipynb Active learning prototype on EuroSAT
  active_learning/   Script versions of the active learning experiments

Installation

Requires Python 3.8+ and PyTorch with TorchVision (CUDA recommended).

pip install -r requirements.txt
# or
conda env create -f environment.yml && conda activate sage

Quick Start

Train on CIFAR-100 with a 5% subset selected by SAGE:

python scripts/train.py --dataset cifar100 --subset_fraction 0.05 --epochs 200

Class-balanced selection (CB-SAGE) with an FD sketch of size 256:

python scripts/train.py --dataset cifar100 --subset_fraction 0.05 \
  --per_class --fd_method fd --sketch_size 256

Train on ImageNet (defaults: 10% subset, ResNet-50, per-class selection):

python scripts/train_imagenet.py --data_path ./data
# expected layout: ./data/imagenet/{train,val}/ or ./data/{train,val}/

Full-data baseline without selection:

python scripts/train.py --dataset cifar100 --no_selection

See python scripts/train.py --help for all options, including:

• --fd_method {fd, random_projection, pca, top_k_grads, no_projection} - gradient dimensionality reduction backend
• --sketch_size - FD sketch size l
• --per_class - class-balanced selection (CB-SAGE)
• --selection_method {sage, random} - SAGE or the random baseline
• --refresh_epochs, --single_selection, --warmup_epochs - selection schedule
• --corrupt_labels, --minority_downsample - robustness experiments

Supported datasets: cifar10, cifar100, tiny_imagenet, imagenet.

Ablations

# Projection methods (FD vs. random projection vs. PCA)
python ablations.py --quick

# Sketch size sweep
python ablations_sketch_size.py --quick
python plot_sketch_size_curves.py

Experimental Extensions

These are exploratory extensions beyond the paper's experiments:

• Active learning (notebooks/): SAGE agreement scores as an acquisition function on EuroSAT.
• Embedding-based selection (src/sage/embedding/train.py): builds the FD sketch from pretrained embeddings (ResNet, CLIP, DINOv2) instead of gradients. Run with python -m src.sage.embedding.train --help.
• LLM prompt trimming (src/sage/llm/prompt_trimming.py): applies agreement-based selection to sentence embeddings for context compression.

Notes

• Datasets, results, checkpoints, and TensorBoard runs are git-ignored by default (data/, results/, checkpoints/, runs/).
• CIFAR datasets are downloaded automatically. TinyImageNet and ImageNet must be provided under --data_path.

License

This project is released under the MIT License. The ResNeXt-29 implementation in src/sage/models.py is adapted from kuangliu/pytorch-cifar (MIT).

Citation

@inproceedings{jha2025sage,
  title     = {{SAGE}: Streaming Agreement-Driven Gradient Sketches for Representative Subset Selection},
  author    = {Jha, Ashish and Ahmadi-Asl, Salman},
  booktitle = {NeurIPS 2025 Workshop on Reliable ML from Unreliable Data},
  year      = {2025},
  eprint    = {2510.02470},
  archivePrefix = {arXiv},
  primaryClass  = {cs.LG}
}