dkm-compression

DKM: Differentiable K-Means Clustering Layer for Neural Network Compression

PyTorch implementation of the ICLR 2022 paper by Cho et al.

📄 Paper (arXiv:2108.12659) | 🏛️ ICLR 2022

Overview

DKM casts k-means weight clustering as a differentiable attention problem, enabling joint optimization of DNN parameters and clustering centroids through standard backpropagation. Unlike prior weight-clustering methods that rely on hard assignments and approximated gradients, DKM uses soft attention-based assignment that is fully differentiable.

Key Innovation

Traditional: weights → hard k-means assignment → fixed centroids (not differentiable)
DKM:         weights → attention-based soft assignment → differentiable centroids

The DKM layer:

1. Computes a distance matrix D between weights W and centroids C
2. Applies softmax with temperature τ to get attention matrix A = softmax(D/τ)
3. Updates centroids: c_j = Σ_i(a_ij × w_i) / Σ_i(a_ij)
4. Iterates until convergence
5. Returns compressed weights: W̃ = A × C

Paper Results

Installation

git clone https://huggingface.co/syedmohaiminulhoque/dkm-compression
cd dkm-compression
pip install torch torchvision

Quick Start

import torch
import torch.nn as nn
from dkm import compress_model
from dkm.utils import print_compression_summary

# Load any pre-trained model
model = torchvision.models.resnet18(weights="DEFAULT")

# Compress with DKM (2-bit clustering)
compressor = compress_model(
    model, 
    bits=2,           # k = 2^bits = 4 clusters
    dim=1,            # scalar clustering (dim=1) or multi-dim
    tau=2e-5,         # temperature (controls softness of assignment)
    skip_first_last=True,  # skip first/last layers (per paper protocol)
)

# Print compression statistics
info = compressor.get_compression_info()
print_compression_summary(info)

# Train with standard PyTorch loop (paper: SGD, lr=0.008, momentum=0.9)
optimizer = torch.optim.SGD(compressor.parameters(), lr=0.008, momentum=0.9)
criterion = nn.CrossEntropyLoss()

compressor.train()
for images, labels in dataloader:
    optimizer.zero_grad()
    outputs = compressor(images)
    loss = criterion(outputs, labels)
    loss.backward()  # Gradients flow through DKM attention layers
    optimizer.step()

# Snap to nearest centroids for inference
compressor.snap_weights()

# Export compressed model (codebook + assignments)
export = compressor.export_compressed()
torch.save(export, "compressed_model.pt")

Multi-Dimensional Clustering (Section 3.3)

DKM supports multi-dimensional weight clustering for higher compression:

# Paper notation: "bits/dim" e.g., "4/4" means 4 bits, 4 dimensions
# Effective bits-per-weight = bits / dim

# Configuration cv:6/8, fc:6/4 (as in Table 3 of the paper)
compressor = compress_model(
    model,
    bits=6,
    conv_config={"bits": 6, "dim": 8},   # 6 bits, 8 dims → 0.75 bpw
    fc_config={"bits": 6, "dim": 4},      # 6 bits, 4 dims → 1.5 bpw
    tau=2e-5,
)

Temperature τ Guidelines (Appendix B)

The temperature controls the softness of cluster assignment:

• Smaller τ → harder assignment (near one-hot), closer to standard k-means
• Larger τ → softer assignment, more gradient flow, better for hard compression tasks

Architecture

dkm/
├── __init__.py          # Package exports
├── dkm_layer.py         # Core DKM layer (Section 3.2-3.3)
├── compressor.py        # Model wrapper with DKM layers (Section 4)
└── utils.py             # Compression analysis utilities

tests/
└── test_dkm.py          # 16 comprehensive test groups (all passing)

train.py                 # Full training pipeline (CIFAR-10 demo)

Core Components

• DKMLayer: The differentiable k-means clustering layer. Implements the iterative attention-based clustering from Fig. 2 of the paper, with k-means++ initialization, warm start across batches, and convergence checking.

• DKMCompressor: Wraps any PyTorch model by inserting DKM layers via forward pre-hooks. Handles per-layer configuration (different bits/dim for conv vs fc), the paper's protocol for small layers (<10K params → 8-bit), and first/last layer skipping.

• compress_model: High-level API matching the paper's notation (cv:bits/dim, fc:bits/dim).

Training Protocol (Section 4)

Following the paper exactly:

• Optimizer: SGD with momentum 0.9
• Learning rate: 0.008 (fixed, no per-layer tuning)
• Loss: Original task loss (no regularizers or modifications)
• Epochs: 200 for ImageNet, varies for GLUE
• Batch size: 128 per GPU (paper used 8× V100)
• Convergence: ε = 1e-4, max 5 DKM iterations per layer
• Small layers: Layers with <10,000 parameters get 8-bit clustering

Compressed Model Format

After training, export_compressed() returns:

• state_dict: Standard PyTorch state dict (with snapped weights)
• codebooks: Per-layer centroid tensors (k × d float32)
• assignments: Per-layer cluster index tensors (N/d integers, b bits each)
• layer_configs: Per-layer DKM configuration

The actual compressed size = Σ(codebook_bits + assignment_bits) per layer + uncompressed params.

Tests

All 16 test groups pass, covering:

1. Shape preservation (train & eval)
2. Distance matrix correctness
3. Attention matrix properties (row-sum=1, temperature effect)
4. Centroid convergence to cluster means
5. Gradient flow (differentiability — key paper contribution)
6. Multi-dimensional clustering
7. Iterative convergence
8. Full compressor pipeline
9. Weight snapping for inference
10. Model export
11. Multi-step training stability
12. Paper configurations (Table 1)
13. K-means++ initialization
14. Warm start across batches
15. Numerical stability (large/small/uniform weights)
16. ResNet-like model compression

python tests/test_dkm.py

Citation

@inproceedings{cho2022dkm,
  title={DKM: Differentiable k-Means Clustering Layer for Neural Network Compression},
  author={Cho, Minsik and Alizadeh-Vahid, Keivan and Adya, Saurabh and Rastegari, Mohammad},
  booktitle={International Conference on Learning Representations (ICLR)},
  year={2022},
  url={https://openreview.net/forum?id=J_F_qqCE3Z5}
}

License

This is a research implementation. The original paper is by Apple Research (Cho et al., ICLR 2022).