adu2021 ROCKET: Rapid Optimization via Calibration-guided Knapsack for Model Compression
Problem Context
Uniform compression (same quantization/sparsity across all layers) is suboptimal—attention layers are more sensitive than MLPs. Manual layer-wise allocation requires expertise. ROCKET automates this by profiling each layer's reconstruction error under different compression settings, then solving a constrained knapsack problem to distribute the global compression budget optimally.
Core Concept
ROCKET operates in two phases: (1) profile layers with multiple compression configurations (rank, sparsity pairs), measuring error per configuration, (2) solve knapsack via dynamic programming to minimize total reconstruction error while staying within target model size. The approach handles per-layer error caps to prevent pathological solutions.
Implementation
Step 1: Layer-wise compression profiling
import torch
from typing import Dict, List, Tuple
class LayerCompressionProfiler:
"""Profile layers under different compression settings."""
def __init__(self, model):
self.model = model
def profile_layer_compression(
self,
layer_name: str,
layer: torch.nn.Module,
calibration_data: torch.Tensor,
compression_configs: List[Dict]
) -> List[Dict]:
"""
Test multiple compression configs on layer.
Args:
compression_configs: List of {rank, sparsity} pairs
Returns:
results: [{'config': {rank, sparsity}, 'error': float, 'params': int}, ...]
"""
results = []
original_output = layer(calibration_data)
for config in compression_configs:
rank = config.get('rank', layer.out_features)
sparsity = config.get('sparsity', 0.0)
# Apply compression
compressed_layer = self._compress_layer(layer, rank, sparsity)
# Measure reconstruction error
with torch.no_grad():
compressed_output = compressed_layer(calibration_data)
error = torch.norm(original_output - compressed_output).item()
# Count parameters after compression
param_count = sum(p.numel() for p in compressed_layer.parameters())
results.append({
'config': config,
'error': error,
'param_count': param_count,
'compression_ratio': layer.out_features / max(1, rank)
})
return results
def _compress_layer(
self,
layer: torch.nn.Module,
rank: int,
sparsity: float
) -> torch.nn.Module:
"""Apply compression (low-rank + sparsity) to layer."""
# Low-rank decomposition
if hasattr(layer, 'weight'):
W = layer.weight.data
U, S, V = torch.linalg.svd(W, full_matrices=False)
S[rank:] = 0 # Zero out small singular values
W_compressed = U @ torch.diag(S) @ V
# Sparsity: zero out small weights
threshold = torch.quantile(torch.abs(W_compressed), sparsity)
W_compressed[torch.abs(W_compressed) < threshold] = 0
layer.weight.data = W_compressed
return layer
def profile_all_layers(
self,
calibration_data: torch.Tensor,
compression_configs: List[Dict]
) -> Dict[str, List[Dict]]:
"""Profile all layers in model."""
layer_profiles = {}
for name, module in self.model.named_modules():
if isinstance(module, (torch.nn.Linear, torch.nn.Conv2d)):
profile = self.profile_layer_compression(
name, module, calibration_data, compression_configs
)
layer_profiles[name] = profile
return layer_profiles
Step 2: Knapsack formulation and DP solution
class CompressionKnapsack:
"""Solve compression allocation as constrained knapsack."""
def __init__(
self,
layer_profiles: Dict[str, List[Dict]],
target_model_size: int,
error_cap_per_layer: float = float('inf')
):
"""
Args:
layer_profiles: Compression profiles per layer
target_model_size: Target parameter count
error_cap_per_layer: Max error allowed per layer (prevent pathological)
"""
self.layer_profiles = layer_profiles
self.target_size = target_model_size
self.error_cap = error_cap_per_layer
def solve_allocation(self) -> Dict[str, Dict]:
"""
Solve knapsack via dynamic programming.
Returns:
allocation: {layer_name: {config, error, param_count}, ...}
"""
layer_names = list(self.layer_profiles.keys())
num_layers = len(layer_names)
# DP table: dp[layer_idx][budget] = min_error
dp = {}
parent = {} # Track which config was chosen
# Base case
dp[(0, 0)] = 0
parent[(0, 0)] = None
# Fill DP table
for layer_idx in range(num_layers):
layer_name = layer_names[layer_idx]
profiles = self.layer_profiles[layer_name]
new_dp = {}
new_parent = {}
for (prev_idx, budget), min_error in dp.items():
if prev_idx != layer_idx:
continue
# Try each compression config for this layer
for config_idx, config_result in enumerate(profiles):
error = config_result['error']
param_count = config_result['param_count']
# Skip if error cap exceeded
if error > self.error_cap:
continue
new_budget = budget + param_count
if new_budget > self.target_size:
continue # Exceeds capacity
new_error = min_error + error
key = (layer_idx + 1, new_budget)
if key not in new_dp or new_error < new_dp[key]:
new_dp[key] = new_error
new_parent[key] = (layer_idx, budget, config_idx)
dp.update(new_dp)
parent.update(new_parent)
# Find best solution within target budget
best_error = float('inf')
best_state = None
for (layer_idx, budget), error in dp.items():
if layer_idx == num_layers and budget <= self.target_size:
if error < best_error:
best_error = error
best_state = (layer_idx, budget)
# Reconstruct allocation
allocation = {}
if best_state:
state = best_state
for layer_idx in range(num_layers - 1, -1, -1):
if state in parent and parent[state] is not None:
prev_layer_idx, prev_budget, config_idx = parent[state]
layer_name = layer_names[layer_idx]
allocation[layer_name] = self.layer_profiles[layer_name][config_idx]
state = (prev_layer_idx, prev_budget)
return allocation
Step 3: Compression execution
class CompressionExecutor:
"""Apply optimal compression to model."""
def __init__(self, model):
self.model = model
def apply_compression_allocation(
self,
allocation: Dict[str, Dict]
) -> torch.nn.Module:
"""Apply compression to each layer according to allocation."""
for layer_name, config_result in allocation.items():
layer = dict(self.model.named_modules())[layer_name]
config = config_result['config']
# Apply compression
self._compress_layer_inplace(layer, config)
return self.model
def _compress_layer_inplace(
self,
layer: torch.nn.Module,
config: Dict
):
"""Compress layer in-place."""
rank = config.get('rank', layer.weight.shape[0])
sparsity = config.get('sparsity', 0.0)
W = layer.weight.data
# Low-rank
U, S, V = torch.linalg.svd(W, full_matrices=False)
S[rank:] = 0
W_compressed = U @ torch.diag(S) @ V
# Sparsity
threshold = torch.quantile(torch.abs(W_compressed), sparsity)
W_compressed[torch.abs(W_compressed) < threshold] = 0
layer.weight.data = W_compressed
def evaluate_compression(
self,
test_data: torch.Tensor,
metric_fn=None
) -> Dict:
"""Evaluate compressed model."""
with torch.no_grad():
metric = metric_fn(self.model, test_data) if metric_fn else 1.0
total_params = sum(p.numel() for p in self.model.parameters())
return {
'metric': metric,
'total_params': total_params,
'compression_ratio': total_params / (sum(p.numel() for p in self.model.parameters()))
}
Step 4: Full compression pipeline
def compress_model_rocket(
model: torch.nn.Module,
calibration_data: torch.Tensor,
test_data: torch.Tensor,
target_size_fraction: float = 0.5,
num_compression_configs: int = 8
):
"""
Full ROCKET compression pipeline.
Args:
target_size_fraction: Fraction of original size (0.5 = 50% compression)
"""
# Generate compression configs to test
compression_configs = []
for rank_ratio in [0.25, 0.5, 0.75, 1.0]:
for sparsity in [0.0, 0.25, 0.5]:
compression_configs.append({
'rank': int(rank_ratio * 768), # Assume dim=768
'sparsity': sparsity
})
# Profile layers
print("Profiling layer compression...")
profiler = LayerCompressionProfiler(model)
layer_profiles = profiler.profile_all_layers(calibration_data, compression_configs)
# Compute target size
original_params = sum(p.numel() for p in model.parameters())
target_params = int(original_params * target_size_fraction)
print(f"Original params: {original_params}")
print(f"Target params: {target_params}")
# Solve knapsack
print("Solving knapsack optimization...")
knapsack = CompressionKnapsack(
layer_profiles,
target_params,
error_cap_per_layer=0.1
)
allocation = knapsack.solve_allocation()
# Execute compression
print("Applying compression...")
executor = CompressionExecutor(model)
compressed_model = executor.apply_compression_allocation(allocation)
# Evaluate
metrics = executor.evaluate_compression(test_data)
print(f"Compressed: {metrics['total_params']} params, "
f"Ratio: {metrics['compression_ratio']:.2f}x")
return compressed_model
Practical Guidance
When to use: LLM deployment with strict size budgets; heterogeneous layer importance known
Hyperparameters:
- target_size_fraction: 0.3-0.7 (compression ratio)
- error_cap_per_layer: 0.05-0.2 (pathology prevention)
- num_compression_configs: 8-20 (more = better solutions but slower)
- rank_fractions: [0.1, 0.25, 0.5, 0.75, 1.0]
- sparsities: [0.0, 0.2, 0.4, 0.6]
Key advantages:
- Optimal allocation avoids pathological solutions
- Attention layers naturally preserved
- Single unified optimization framework
- Calibration-guided (preserves model behavior)
Common pitfalls:
- error_cap too strict → infeasible solutions
- Calibration data not representative → poor allocation
- Not validating downstream task performance
- DP state space explosion → need pruning
Scaling: DP complexity is O(L·M·B̄) where L=layers, M=configs, B̄=avg budget.
Reference
Paper: https://arxiv.org/abs/2602.11008 Related work: Model compression, quantization, pruning, knapsack optimization Benchmarks: Language model compression, downstream task accuracy