SIT-SAM

SIT-SAM is a small research package for adding trainable or random key/value token selection to BERT-like Hugging Face Transformers models. The main implementation target is DistilBERT sequence classification, matching the experiments in the paper draft.

The method replaces selected self-attention layers with a token-selection variant. In the trainable mode, a lightweight scoring network selects a fixed fraction of key/value tokens using a Gumbel-Softmax top-k relaxation during training and real key/value gathering during evaluation.

Repository Layout

• src/sit_sam/ - Python package and CLI.
• experiments/results/ - archived CSV/log outputs from the paper experiments.
• notebooks/ - exploratory plotting and token-score notebooks.
• models/ - local model checkpoints and Hugging Face model cache.
• datasets/ - local Hugging Face dataset cache.

Installation

conda activate ml_labs
pip install -e .

All CLI commands default to models/ and datasets/ for model and dataset caches, so downloads do not go to the global Hugging Face cache.

CLI Examples

Fine-tune a base DistilBERT classifier:

sit-sam train-base \
  --dataset imdb \
  --model-name-or-path distilbert-base-uncased \
  --output-dir models/imdb-distilbert

Train a learned token-selection module on a fine-tuned base model:

sit-sam train-selector \
  --dataset imdb \
  --model-name-or-path models/imdb-distilbert \
  --strategy learned \
  --mask-size 0.6 \
  --num-layers 6 \
  --temperature 10 \
  --output-dir models/imdb-sit-sam-06

This command writes sit_sam_config.json, model_state.pt, tokenizer files, and eval_metrics.csv. The CLI can load that directory again for evaluation.

Run an article-style sweep over mask sizes and replaced layer counts:

sit-sam sweep \
  --dataset imdb \
  --model-name-or-path models/imdb-distilbert \
  --strategy learned \
  --mask-sizes 0.4 0.5 0.6 0.7 0.8 0.9 \
  --layer-counts 1 2 3 4 5 6 \
  --output-dir experiments/runs/imdb-learned

For random-mask baselines, use --strategy random --epochs 0.

Article Tables

To reproduce the full experiment grid used for the article tables, run:

bash scripts/reproduce_article_tables.sh

The script trains a base DistilBERT classifier for each dataset, then runs learned-mask and random-mask sweeps for mask sizes 0.4 ... 0.9 and replaced layer counts 1 ... 6. Outputs are written to experiments/runs/article/; the aggregated table is written to experiments/runs/article/tables/mask_results.csv.

This is a server-oriented run. On a laptop it is useful for checking the pipeline, but not for timing claims.

Library Use

from transformers import AutoModelForSequenceClassification, AutoTokenizer
from sit_sam import apply_token_selection, freeze_backbone_except_selectors, load_dataset

tokenizer = AutoTokenizer.from_pretrained("models/imdb-distilbert")
model = AutoModelForSequenceClassification.from_pretrained("models/imdb-distilbert")
apply_token_selection(
    model,
    selection_size=0.6,
    num_layers=6,
    position="first",
    strategy="learned",
    temperature=10.0,
)
freeze_backbone_except_selectors(model)

train_dataset, eval_dataset, data_collator, spec = load_dataset(
    "imdb",
    tokenizer,
    datasets_dir="datasets",
)

apply_token_selection supports DistilBERT and BERT/RoBERTa-style encoder classification models. DistilBERT is the configuration used for the reported experiments.

Scope

The reported experiments use DistilBERT for binary sentiment classification. The BERT/RoBERTa-style replacement API is provided for reuse in related encoder models, but it should be treated as an extension point rather than a separately validated benchmark result. The notebooks are exploratory/legacy analysis files; the supported reproducibility path is the package CLI plus the article-table script.