Capability-Guided Compression (CGC) — v1

Paper: Capability-Guided Compression: Toward Interpretability-Aware Budget Allocation for Large Language Models
Author: Rishaank Gupta, Independent Researcher, 2026
arXiv: https://arxiv.org/abs/2603.16440

---

Overview

This repository contains the code for CGC v1 experiments on GPT-2 Medium (355M parameters).

Honest result summary:

Metric	Value
Model	GPT-2 Medium (24 layers, 16 heads, head_dim=64)
SAEs	Self-trained TopK SAEs (dict_size=512, k=25, 5 epochs)
Density range	0.63–0.80 (all 384 heads)
Orthogonality ρ (density vs Wanda)	−0.054
Ablation r (density vs ΔPPL)	−0.066 (p=0.20, not significant)
Dense PPL	26.68
Uniform PPL	27.57
CGC-L PPL	27.87 (worse than Uniform)
Inverted PPL	27.86

The v1 experiments yield a negative PPL result — CGC-L underperforms Uniform compression on GPT-2 Medium. This is reported honestly in the paper, along with a diagnosis of three reasons why GPT-2 Medium is an insufficient test bed (uniform head structure, PPL insensitivity, shallow SAE training). See Section 5 and Section 6.1 of the paper.

The orthogonality finding (Spearman ρ = −0.054, confirming density is independent of Wanda importance) is the primary empirical contribution of v1.

For CGC v2 experiments on Llama 3 8B with positive results, see the companion cgc_v2 repository.

---

Installation

git clone https://github.com/rishaankgupta/CGC.git
cd cgc
pip install -e .

GPU required. All experiments originally run on T4 (Google Colab / Lightning.ai free tier).

---

Reproducing Results

# Step 1: Compute capability density map (~15-20 min on T4)
python scripts/compute_density_map.py --output results/density_map.npz

# Step 2: Ablation + orthogonality + compression (~30 min on T4)
python scripts/run_compression.py \
    --density_map results/density_map.npz \
    --output_dir  results/

# Step 3: Generate figures
python scripts/plot_results.py \
    --density_map results/density_map.npz \
    --results_dir results/ \
    --output_dir  figures/

Alternatively, run notebooks/cgc_v1_colab.ipynb end-to-end on Google Colab (T4 GPU).

---

Two Wanda Implementations

This codebase provides two Wanda implementations, selectable via --wanda_mode:

--wanda_mode paper (default)

Per-head approximation intended to reproduce the conditions of the original v1 experiment. Computes ||W_h||_F × mean(||act_h||_2) using head output activations already stored during density map computation. This is what was run when the paper was written. Results may differ from the paper's reported Spearman ρ = −0.054 due to stochastic SAE training (no fixed seed in the original run).

--wanda_mode real

Correct per-weight Wanda criterion from Sun et al. (2024). Hooks the input to each layer's output projection, computes per-channel RMS norms across calibration tokens, and multiplies by per-weight magnitudes. Results may differ from the paper's reported ρ value. Recommended for future work building on CGC.

# Intended to reproduce original paper experiment conditions
python scripts/run_compression.py --density_map results/density_map.npz \
    --output_dir results/ --wanda_mode paper

# Correct per-weight Wanda criterion
python scripts/run_compression.py --density_map results/density_map.npz \
    --output_dir results/ --wanda_mode real

Reproducibility Note

The original experiments were run without a fixed random seed. This codebase uses SEED = 42 (set in every script) for reproducibility. Re-running will produce approximately the same results — same statistical patterns and same sign on all correlations — but not bit-for-bit identical numbers to the paper.

The structural findings (density range 0.63–0.80, non-significant ablation correlation, PPL ordering) are stable properties of GPT-2 Medium that reproduce consistently regardless of seed.

---

Repository Structure

cgc_v1/
├── cgc/
│   ├── sae.py        — TopK SAE architecture + training function
│   ├── density.py    — Activation collection, SAE training loop, β/H/Ψ/δ
│   ├── compress.py   — Real Wanda importance, CGC-L allocation, magnitude pruning
│   └── evaluate.py   — WikiText-2 perplexity evaluation
├── scripts/
│   ├── compute_density_map.py  — Phase 1+2 main script
│   ├── run_compression.py      — Phase 3 main script
│   └── plot_results.py         — Figures 1–4
├── notebooks/
│   └── cgc_v1_colab.ipynb     — Self-contained Colab notebook (17 cells)
├── results/                   — Output directory (git-tracked placeholder only)
├── requirements.txt
└── setup.py

---

Citation

@article{gupta2026cgc,
  title   = {Capability-Guided Compression: Toward Interpretability-Aware
             Budget Allocation for Large Language Models},
  author  = {Gupta, Rishaank},
  journal = {arXiv preprint},
  year    = {2026}
}