NSA-Diff

Null-Space Absorbing Compression for Diffusion U-Nets. Compresses DDPM denoisers via CP decomposition + source-wise conditional null-space loss at skip connections + knowledge distillation.

Teacher: google/ddpm-cifar10-32 (35.7M params). Paper runs use rank_ratio=0.7 unless stated otherwise; inspect other ranks with scripts/inspect_model.py.

Method

Standard NSA (Ozdemir et al.) drives teacher–student activation mismatch into the null space of low-rank student weights: ||W_eff · e||² → 0. This fails at U-Net skip connections where the error has two independent sources (decoder + encoder skip).

Conditional NSA prevents decoder and skip errors from cancelling each other in the row space:

W_eff = [W_dec | W_skip]
L_cond = ||W_dec · e_dec||² + ||W_skip · δ_skip||²

Each source must place its own error in the null space of its corresponding channel block.

Full Loss

L = L_ε + α·L_null + α_s·L_cond + β·L_KD + λ·L_orth

Component	Description
L_ε	Noise prediction MSE (diffusion objective)
L_null	Standard NSA on non-skip conv layers
L_cond	Conditional NSA on skip-receiving conv1 layers
L_KD	Output-level distillation (noise pred matching)
L_orth	Orthogonality regularization on CP output factors

Baselines

Method	Losses Used
lowrank_kd	L_ε + β·L_KD + λ·L_orth
standard_nsa	L_ε + α·L_null (all layers) + β·L_KD + λ·L_orth
nsa_diff	L_ε + α·L_null + α_s·L_cond + β·L_KD + λ·L_orth
nsa_diff --disable_orth true	NSA-Diff without L_orth
fitnets	L_ε + β·L_FitNets
gramian	L_ε + β·L_Gram + β·L_KD

Setup

conda env create -f environment.yml
conda activate nsa_diff
pip install -e .

Or manually:

conda create -n nsa_diff python=3.10 -y
conda activate nsa_diff
pip install torch torchvision diffusers accelerate tensorly pytorch-fid wandb pytest

Usage

Inspect model architecture

python scripts/inspect_model.py --rank_ratio 0.25

Train

# NSA-Diff (proposed method)
python scripts/train.py --method nsa_diff --num_steps 100000 --batch_size 64

# Smoke test on MPS
python scripts/train.py --method nsa_diff --num_steps 100 --batch_size 8

# Train on a fixed number of samples (same data for any batch size)
python scripts/train.py --method nsa_diff --total_samples 6400000 --batch_size 64
python scripts/train.py --method nsa_diff --total_samples 6400000 --batch_size 128  # 50k steps instead of 100k

# All 5 baselines
bash scripts/run_baselines.sh

CoLorAI paper preset

The paper preset matches the current 4-page CoLorAI submission plan. Run this on the server; it launches the main baselines, the no-orth ablation, and the NSA-Diff rank sweep.

conda activate nsa_diff
RANK=0.7 NUM_STEPS=50000 BATCH_SIZE=256 bash scripts/run_paper_experiments.sh

Equivalent single-run command:

python scripts/train.py \
    --method nsa_diff \
    --rank_ratio 0.7 \
    --num_steps 50000 \
    --batch_size 256 \
    --lr 1e-4 \
    --alpha 1.0 \
    --alpha_s 1.0 \
    --beta 0.5 \
    --lam 0.01 \
    --warmup_steps 1000 \
    --ema_decay 0.9999 \
    --run_name nsa_diff_r0.7

Key arguments:

Arg	Default	Description
--method	nsa_diff	One of: lowrank_kd, standard_nsa, nsa_diff, fitnets, gramian
--rank_ratio	0.25	CP rank / min(C_in, C_out). Lower = more compression
--num_steps	100000	Training iterations (ignored if --total_samples is set)
--total_samples	—	If set, steps = total_samples // batch_size so the same number of samples is seen for any batch size
--batch_size	64	Batch size
--lr	1e-4	Learning rate
--alpha	0.1	NSA loss weight
--alpha_s	0.1	Conditional NSA weight
--beta	0.1	KD/FitNets/Gramian weight
--lam	0.1	Orthogonality weight
--disable_orth	false	Disable L_orth for no-orth ablations
--use_wandb	false	Enable W&B logging
--device	auto	Force device (cuda/mps/cpu)

Evaluate

python scripts/evaluate.py \
    --checkpoint outputs/nsa_diff/checkpoint_100000.pt \
    --num_samples 50000 \
    --scheduler ddpm \
    --benchmark

Qualitative paper grid

Generate the fixed-seed main-paper sample grid after checkpoints are available:

python scripts/make_sample_grid.py \
    --teacher \
    --cp_init "CP init=0.7" \
    --checkpoint "Low-rank KD=outputs/paper/lowrank_kd_r0.7/checkpoint_50000.pt" \
    --checkpoint "Standard NSA=outputs/paper/standard_nsa_r0.7/checkpoint_50000.pt" \
    --checkpoint "NSA-Diff=outputs/paper/nsa_diff_r0.7/checkpoint_50000.pt" \
    --num_samples 16 \
    --cols 8 \
    --num_steps 1000 \
    --seed 42 \
    --output latex_report/figures/qualitative_grid.png

Result collection

After FID is computed externally with pytorch-fid, collect rows for the paper tables:

python scripts/collect_results.py \
    --teacher \
    --checkpoint lowrank_kd=outputs/paper/lowrank_kd_r0.7/checkpoint_50000.pt \
    --checkpoint standard_nsa=outputs/paper/standard_nsa_r0.7/checkpoint_50000.pt \
    --checkpoint nsa_diff=outputs/paper/nsa_diff_r0.7/checkpoint_50000.pt \
    --checkpoint nsa_diff_no_orth=outputs/paper/nsa_diff_no_orth_r0.7/checkpoint_50000.pt \
    --fid teacher=13.71 \
    --fid nsa_diff=16.12 \
    --output results/colorai_main.csv

Tests

# Unit tests (fast, no model download)
pytest tests/test_cp_decompose.py tests/test_losses.py tests/test_hooks.py -v

# Integration tests (requires model download)
pytest tests/test_student_builder.py tests/test_training_step.py -v

# All tests
pytest -v

Project Structure

configs/
  __init__.py          ExperimentConfig dataclass + CLI parser
  defaults.py          Per-method default loss weights
src/
  decomposition/
    cp_decompose.py    CP decomposition → 4-conv sequence (pw_in, dw_h, dw_v, pw_out)
    student_builder.py Deep-copy teacher, replace convs with CP sequences
  losses/
    nsa_loss.py        ||W_eff · e||²  on non-skip layers
    conditional_nsa.py Source-wise loss for skip receivers
    distillation.py    KD (MSE), FitNets (hint MSE), Gramian (F·Fᵀ matching)
    orthogonality.py   ||UᵀU - I||_F²  on CP output factors
    composite.py       Per-method loss aggregator
  hooks/
    activation_capture.py  Forward pre-hooks on conv layers
  training/
    trainer.py         Training loop (noise sampling, forward, loss, backward, EMA)
    ema.py             Exponential moving average
    data.py            CIFAR-10 dataloader ([-1,1], flip augmentation)
  evaluation/
    sample.py          DDPM/DDIM generation
    fid.py             FID via pytorch-fid
    benchmark.py       Latency, memory, param count
  utils/
    device.py          CUDA/MPS/CPU auto-detection, AMP config
    unet_inspect.py    Discover compressible layers + skip connection info
    logging_utils.py   W&B + console logger
scripts/
  train.py             CLI entry point
  evaluate.py          Generate + FID evaluation
  make_sample_grid.py  Fixed-seed qualitative paper figure
  collect_results.py   CSV/JSON result table collector
  inspect_model.py     Print architecture + compression stats
  run_paper_experiments.sh  Server launcher for paper runs
  run_baselines.sh     Launch all 5 methods
tests/
  test_cp_decompose.py, test_student_builder.py, test_losses.py,
  test_hooks.py, test_training_step.py

Compression Configurations

Rank Ratio	Student Params	Compression	Target
0.70	Mild	Main paper operating point
0.50	Moderate	Rank/compression tradeoff
0.25	Aggressive	Rank/compression tradeoff

Device Support

• CUDA (RTX 3090): Full training with AMP. Recommended for full runs.
• MPS (Apple Silicon): Works for dev/smoke tests. AMP auto-disabled.
• CPU: Functional but slow. Integration tests run here.