LLMs and generative AI / LLM architecture and theory
The Exponential Growth of Large Language Models
Contents
- Model Pruning Fundamentals
- Advanced Gradient-Based Pruning
- Quantization Implementation
- Quantization-Aware Training
- Low-Rank Matrix Decomposition Implementation
- Knowledge Distillation Framework
- Feature-Based Knowledge Distillation
- Lightweight Model Design - MobileNet Implementation
- Model Compression Performance Analysis
- Real-World Application - Computer Vision Model Compression
- Real-World Application - NLP Model Compression
- Results Analysis and Visualization
- Additional Resources
Model Pruning Fundamentals
Neural network pruning aims to reduce model size by systematically removing weights based on their importance. The gradient-based approach evaluates each parameter's contribution to the loss function through first-order derivatives, enabling identification of non-critical weights for removal.
import torch
import torch.nn as nn
class PruningExample:
def __init__(self, model, pruning_percentage=0.3):
self.model = model
self.pruning_percentage = pruning_percentage
def compute_weight_importance(self):
importance_scores = {}
for name, param in self.model.named_parameters():
if 'weight' in name:
# Calculate importance using gradient magnitude
if param.grad is not None:
importance = torch.abs(param.grad * param)
importance_scores[name] = importance
return importance_scores
def prune_weights(self):
importance_scores = self.compute_weight_importance()
for name, importance in importance_scores.items():
# Calculate threshold for pruning
threshold = torch.quantile(importance.abs().flatten(),
self.pruning_percentage)
# Create binary mask
mask = (importance.abs() > threshold).float()
# Apply mask to weights
self.model.get_parameter(name).data *= mask
# Example usage
model = nn.Sequential(
nn.Linear(784, 256),
nn.ReLU(),
nn.Linear(256, 10)
)
pruner = PruningExample(model)
Advanced Gradient-Based Pruning
This implementation extends basic pruning by incorporating second-order derivatives through the Hessian matrix approximation, providing more accurate importance scoring for weights while maintaining computational efficiency through diagonal approximation.
import torch
import torch.nn as nn
import torch.nn.functional as F
class HessianPruning:
def __init__(self, model, dataloader, criterion):
self.model = model
self.dataloader = dataloader
self.criterion = criterion
def compute_hessian_diag(self):
hessian_diag = {}
for name, param in self.model.named_parameters():
if 'weight' in name:
hessian_diag[name] = torch.zeros_like(param)
for batch_x, batch_y in self.dataloader:
self.model.zero_grad()
output = self.model(batch_x)
loss = self.criterion(output, batch_y)
grad_dict = torch.autograd.grad(loss, self.model.parameters(),
create_graph=True)
for (name, param), grad in zip(self.model.named_parameters(),
grad_dict):
if 'weight' in name:
hessian_diag[name] += grad.pow(2)
return hessian_diag
def prune_weights(self, pruning_percentage=0.3):
hessian = self.compute_hessian_diag()
for name, param in self.model.named_parameters():
if 'weight' in name:
importance = torch.abs(param) / torch.sqrt(hessian[name] + 1e-8)
threshold = torch.quantile(importance.flatten(),
pruning_percentage)
mask = (importance > threshold).float()
param.data *= mask
Quantization Implementation
Quantization is implemented through a custom quantizer class that converts floating-point weights to fixed-point integers, significantly reducing model size while maintaining numerical stability through proper scaling factors.
import numpy as np
import torch
import torch.nn as nn
class WeightQuantizer:
def __init__(self, bits=8):
self.bits = bits
self.qmin = -(1 << (bits - 1))
self.qmax = (1 << (bits - 1)) - 1
def quantize(self, tensor):
# Compute scale factor
scale = (tensor.max() - tensor.min()) / (self.qmax - self.qmin)
zero_point = self.qmin - torch.round(tensor.min() / scale)
# Quantize
qtensor = torch.round(tensor / scale + zero_point)
qtensor = torch.clamp(qtensor, self.qmin, self.qmax)
# Dequantize for computing gradient
dqtensor = (qtensor - zero_point) * scale
return dqtensor, scale, zero_point
class QuantizedLinear(nn.Module):
def __init__(self, in_features, out_features, bits=8):
super().__init__()
self.in_features = in_features
self.out_features = out_features
self.quantizer = WeightQuantizer(bits)
self.weight = nn.Parameter(torch.randn(out_features, in_features))
self.bias = nn.Parameter(torch.zeros(out_features))
def forward(self, x):
qweight, scale, zp = self.quantizer.quantize(self.weight)
return F.linear(x, qweight, self.bias)
Quantization-Aware Training
Quantization-aware training integrates the quantization process during model training, allowing the network to adapt its weights to compensate for quantization errors. This implementation shows how to create a training loop that simulates quantized inference.
class QuantizationAwareTrainer:
def __init__(self, model, optimizer, bits=8):
self.model = model
self.optimizer = optimizer
self.quantizer = WeightQuantizer(bits)
def train_step(self, inputs, targets, criterion):
self.optimizer.zero_grad()
# Forward pass with quantized weights
with torch.no_grad():
for name, param in self.model.named_parameters():
if 'weight' in name:
# Store original weights
param.data_backup = param.data.clone()
# Quantize and replace weights
q_data, _, _ = self.quantizer.quantize(param.data)
param.data = q_data
# Compute loss with quantized weights
outputs = self.model(inputs)
loss = criterion(outputs, targets)
# Backward pass
loss.backward()
# Restore original weights
with torch.no_grad():
for name, param in self.model.named_parameters():
if 'weight' in name:
param.data = param.data_backup
delattr(param, 'data_backup')
self.optimizer.step()
return loss.item()
# Example usage
model = nn.Sequential(
QuantizedLinear(784, 256),
nn.ReLU(),
QuantizedLinear(256, 10)
)
optimizer = torch.optim.Adam(model.parameters())
trainer = QuantizationAwareTrainer(model, optimizer)
Low-Rank Matrix Decomposition Implementation
Low-rank decomposition reduces model size by factorizing weight matrices into products of smaller matrices. This implementation demonstrates SVD-based decomposition with automatic rank selection based on singular value importance.
import torch
import torch.nn as nn
from typing import Tuple
class LowRankLinear(nn.Module):
def __init__(self, in_features: int, out_features: int, rank_ratio: float = 0.25):
super().__init__()
self.in_features = in_features
self.out_features = out_features
self.rank = int(min(in_features, out_features) * rank_ratio)
# Initialize decomposed matrices
self.U = nn.Parameter(torch.randn(out_features, self.rank))
self.V = nn.Parameter(torch.randn(self.rank, in_features))
self.bias = nn.Parameter(torch.zeros(out_features))
def forward(self, x: torch.Tensor) -> torch.Tensor:
# Compute W = UV
reconstructed_weight = torch.mm(self.U, self.V)
return F.linear(x, reconstructed_weight, self.bias)
@staticmethod
def from_dense(linear_layer: nn.Linear, rank_ratio: float = 0.25) -> 'LowRankLinear':
U, S, V = torch.svd(linear_layer.weight.data)
rank = int(min(U.shape[1], V.shape[0]) * rank_ratio)
low_rank_layer = LowRankLinear(
linear_layer.in_features,
linear_layer.out_features,
rank_ratio
)
# Initialize with SVD results
low_rank_layer.U.data = U[:, :rank] * torch.sqrt(S[:rank])
low_rank_layer.V.data = (V[:rank, :] * torch.sqrt(S[:rank].unsqueeze(1)))
low_rank_layer.bias.data = linear_layer.bias.data
return low_rank_layer
Knowledge Distillation Framework
Knowledge distillation transfers knowledge from a large teacher model to a smaller student model. This implementation includes temperature scaling and various distillation objectives including response-based and feature-based approaches.
class DistillationLoss(nn.Module):
def __init__(self, temperature=2.0, alpha=0.5):
super().__init__()
self.temperature = temperature
self.alpha = alpha
self.kl_div = nn.KLDivLoss(reduction='batchmean')
self.ce_loss = nn.CrossEntropyLoss()
def forward(self,
student_logits: torch.Tensor,
teacher_logits: torch.Tensor,
labels: torch.Tensor) -> torch.Tensor:
# Compute soft targets
soft_targets = F.softmax(teacher_logits / self.temperature, dim=1)
soft_prob = F.log_softmax(student_logits / self.temperature, dim=1)
# Distillation loss
distillation_loss = self.kl_div(soft_prob, soft_targets) * (self.temperature ** 2)
# Standard cross-entropy loss
student_loss = self.ce_loss(student_logits, labels)
# Combined loss
total_loss = (self.alpha * distillation_loss +
(1 - self.alpha) * student_loss)
return total_loss
class KnowledgeDistillation:
def __init__(self, teacher_model, student_model, temperature=2.0, alpha=0.5):
self.teacher = teacher_model
self.student = student_model
self.criterion = DistillationLoss(temperature, alpha)
def train_step(self, inputs, labels, optimizer):
optimizer.zero_grad()
# Teacher predictions (no grad needed)
with torch.no_grad():
teacher_logits = self.teacher(inputs)
# Student predictions
student_logits = self.student(inputs)
# Compute loss
loss = self.criterion(student_logits, teacher_logits, labels)
# Backprop and optimize
loss.backward()
optimizer.step()
return loss.item()
Feature-Based Knowledge Distillation
Feature-based distillation extends basic knowledge transfer by matching intermediate representations between teacher and student networks. This implementation includes attention transfer and feature regression mechanisms.
class FeatureDistillation(nn.Module):
def __init__(self, student_channels, teacher_channels):
super().__init__()
self.adapter = nn.Conv2d(student_channels, teacher_channels, 1)
self.attention_loss = nn.MSELoss()
def attention_map(self, features):
# Convert feature maps to attention maps
b, c, h, w = features.size()
attention = features.pow(2).mean(1, keepdim=True)
attention = F.normalize(attention.view(b, -1), p=2, dim=1)
return attention.view(b, 1, h, w)
def forward(self, student_features, teacher_features):
# Adapt student features to match teacher dimensions
adapted_student = self.adapter(student_features)
# Compute attention maps
student_attention = self.attention_map(adapted_student)
teacher_attention = self.attention_map(teacher_features)
# Feature regression loss
regression_loss = F.mse_loss(
adapted_student,
teacher_features.detach()
)
# Attention transfer loss
attention_loss = self.attention_loss(
student_attention,
teacher_attention.detach()
)
return regression_loss + attention_loss
class FeatureDistillationTrainer:
def __init__(self, teacher_model, student_model, feature_pairs):
self.teacher = teacher_model
self.student = student_model
self.distillers = nn.ModuleList([
FeatureDistillation(s_ch, t_ch)
for s_ch, t_ch in feature_pairs
])
def compute_feature_loss(self, student_features, teacher_features):
total_loss = 0
for distiller, s_feat, t_feat in zip(
self.distillers, student_features, teacher_features):
total_loss += distiller(s_feat, t_feat)
return total_loss
Lightweight Model Design - MobileNet Implementation
MobileNet architecture demonstrates efficient model design through depthwise separable convolutions, significantly reducing parameter count while maintaining performance. This implementation shows the core building blocks.
class DepthwiseSeparableConv(nn.Module):
def __init__(self, in_channels, out_channels, stride=1):
super().__init__()
self.depthwise = nn.Conv2d(
in_channels, in_channels, kernel_size=3, stride=stride,
padding=1, groups=in_channels, bias=False
)
self.bn1 = nn.BatchNorm2d(in_channels)
self.pointwise = nn.Conv2d(
in_channels, out_channels, kernel_size=1,
stride=1, padding=0, bias=False
)
self.bn2 = nn.BatchNorm2d(out_channels)
def forward(self, x):
x = self.depthwise(x)
x = self.bn1(x)
x = F.relu(x)
x = self.pointwise(x)
x = self.bn2(x)
return F.relu(x)
class LightweightMobileNet(nn.Module):
def __init__(self, num_classes=1000, width_multiplier=1.0):
super().__init__()
def conv_bn(inp, oup, stride):
return nn.Sequential(
nn.Conv2d(inp, oup, 3, stride, 1, bias=False),
nn.BatchNorm2d(oup),
nn.ReLU(inplace=True)
)
input_channel = int(32 * width_multiplier)
self.features = nn.Sequential(
conv_bn(3, input_channel, 2),
DepthwiseSeparableConv(input_channel, 64),
DepthwiseSeparableConv(64, 128, stride=2),
DepthwiseSeparableConv(128, 128),
DepthwiseSeparableConv(128, 256, stride=2),
DepthwiseSeparableConv(256, 256),
DepthwiseSeparableConv(256, 512, stride=2),
DepthwiseSeparableConv(512, 512),
nn.AdaptiveAvgPool2d(1)
)
self.classifier = nn.Linear(512, num_classes)
def forward(self, x):
x = self.features(x)
x = x.view(x.size(0), -1)
x = self.classifier(x)
return x
Model Compression Performance Analysis
This implementation provides a comprehensive framework for measuring and comparing different model compression techniques, including memory usage, inference speed, and accuracy metrics across various compression methods.
class CompressionAnalyzer:
def __init__(self, original_model, compressed_model):
self.original_model = original_model
self.compressed_model = compressed_model
def count_parameters(self, model):
return sum(p.numel() for p in model.parameters())
def measure_inference_time(self, model, input_tensor, num_runs=100):
times = []
model.eval()
with torch.no_grad():
# Warmup
for _ in range(10):
_ = model(input_tensor)
# Actual measurement
for _ in range(num_runs):
start_time = time.time()
_ = model(input_tensor)
times.append(time.time() - start_time)
return {
'mean': np.mean(times),
'std': np.std(times),
'median': np.median(times)
}
def measure_memory_usage(self, model):
memory_params = sum([param.nelement() * param.element_size()
for param in model.parameters()])
memory_buffers = sum([buf.nelement() * buf.element_size()
for buf in model.buffers()])
return {
'parameters_bytes': memory_params,
'buffers_bytes': memory_buffers,
'total_bytes': memory_params + memory_buffers
}
def evaluate_accuracy(self, model, dataloader, criterion):
model.eval()
total_loss = 0
correct = 0
total = 0
with torch.no_grad():
for inputs, targets in dataloader:
outputs = model(inputs)
loss = criterion(outputs, targets)
total_loss += loss.item()
_, predicted = outputs.max(1)
total += targets.size(0)
correct += predicted.eq(targets).sum().item()
return {
'accuracy': 100. * correct / total,
'avg_loss': total_loss / len(dataloader)
}
def generate_report(self, input_tensor, dataloader, criterion):
orig_params = self.count_parameters(self.original_model)
comp_params = self.count_parameters(self.compressed_model)
report = {
'parameter_reduction': {
'original': orig_params,
'compressed': comp_params,
'compression_ratio': orig_params / comp_params
},
'memory_usage': {
'original': self.measure_memory_usage(self.original_model),
'compressed': self.measure_memory_usage(self.compressed_model)
},
'inference_speed': {
'original': self.measure_inference_time(
self.original_model, input_tensor),
'compressed': self.measure_inference_time(
self.compressed_model, input_tensor)
},
'accuracy': {
'original': self.evaluate_accuracy(
self.original_model, dataloader, criterion),
'compressed': self.evaluate_accuracy(
self.compressed_model, dataloader, criterion)
}
}
return report
Real-World Application - Computer Vision Model Compression
This implementation demonstrates a complete pipeline for compressing a pre-trained vision model while maintaining acceptable performance on a specific task.
class VisionModelCompressor:
def __init__(self, model_name='resnet18', dataset='cifar10'):
self.original_model = models.__dict__[model_name](pretrained=True)
self.dataset = dataset
self.setup_data()
def setup_data(self):
transform = transforms.Compose([
transforms.Resize(224),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
])
if self.dataset == 'cifar10':
self.train_data = datasets.CIFAR10(
root='./data', train=True,
download=True, transform=transform)
self.test_data = datasets.CIFAR10(
root='./data', train=False,
download=True, transform=transform)
def apply_compression(self, compression_config):
compressed_model = copy.deepcopy(self.original_model)
if compression_config.get('pruning', False):
pruner = PruningExample(
compressed_model,
pruning_percentage=compression_config['pruning_ratio']
)
pruner.prune_weights()
if compression_config.get('quantization', False):
for name, module in compressed_model.named_children():
if isinstance(module, nn.Linear):
setattr(compressed_model, name,
QuantizedLinear(
module.in_features,
module.out_features,
bits=compression_config['quantization_bits']
))
if compression_config.get('distillation', False):
distiller = KnowledgeDistillation(
self.original_model,
compressed_model,
temperature=compression_config['temperature']
)
self.train_distillation(distiller, epochs=5)
return compressed_model
def train_distillation(self, distiller, epochs):
optimizer = torch.optim.Adam(distiller.student.parameters())
train_loader = DataLoader(
self.train_data, batch_size=32, shuffle=True)
for epoch in range(epochs):
for inputs, targets in train_loader:
loss = distiller.train_step(inputs, targets, optimizer)
def evaluate_compression(self, compressed_model):
analyzer = CompressionAnalyzer(
self.original_model, compressed_model)
test_loader = DataLoader(
self.test_data, batch_size=32, shuffle=False)
dummy_input = torch.randn(1, 3, 224, 224)
report = analyzer.generate_report(
dummy_input,
test_loader,
nn.CrossEntropyLoss()
)
return report
Real-World Application - NLP Model Compression
This implementation showcases compression techniques specifically tailored for large language models, demonstrating practical approaches to reduce transformer model sizes while preserving performance.
class TransformerCompressor:
def __init__(self, model_name='bert-base-uncased'):
self.tokenizer = AutoTokenizer.from_pretrained(model_name)
self.original_model = AutoModel.from_pretrained(model_name)
def apply_attention_pruning(self, threshold=0.1):
def prune_attention_heads(layer, important_heads):
# Reshape attention weights to identify head importance
attention_weights = layer.attention.self.query.weight.view(
config.num_attention_heads, -1)
head_importance = torch.norm(attention_weights, dim=1)
mask = head_importance > threshold
return mask
pruned_heads = {}
for name, layer in self.original_model.encoder.layer.named_children():
if hasattr(layer, 'attention'):
head_mask = prune_attention_heads(layer)
pruned_heads[name] = head_mask
return pruned_heads
def quantize_embeddings(self, bits=8):
embeddings = self.original_model.embeddings.word_embeddings
quantizer = WeightQuantizer(bits)
# Quantize embedding weights
quantized_weights, scale, zero_point = quantizer.quantize(
embeddings.weight)
embeddings.weight.data = quantized_weights
return {'scale': scale, 'zero_point': zero_point}
def compress_model(self, config):
compressed_model = copy.deepcopy(self.original_model)
if config.get('prune_attention', False):
pruned_heads = self.apply_attention_pruning(
config['attention_threshold'])
if config.get('quantize_embeddings', False):
embedding_params = self.quantize_embeddings(
config['embedding_bits'])
if config.get('low_rank', False):
for name, module in compressed_model.named_modules():
if isinstance(module, nn.Linear):
low_rank_module = LowRankLinear.from_dense(
module, rank_ratio=config['rank_ratio'])
setattr(compressed_model, name, low_rank_module)
return compressed_model
def evaluate_on_task(self, model, task_dataset, batch_size=32):
model.eval()
total_loss = 0
accuracies = []
dataloader = DataLoader(task_dataset, batch_size=batch_size)
with torch.no_grad():
for batch in dataloader:
inputs = self.tokenizer(
batch['text'],
padding=True,
truncation=True,
return_tensors='pt'
)
outputs = model(**inputs)
# Task-specific evaluation
if 'classification' in task_dataset.task_type:
logits = outputs.logits
predictions = torch.argmax(logits, dim=1)
accuracies.append(
(predictions == batch['labels']).float().mean())
return {
'accuracy': torch.tensor(accuracies).mean().item()
if accuracies else None,
'avg_loss': total_loss / len(dataloader)
}
Results Analysis and Visualization
A comprehensive framework for analyzing and visualizing the impact of different compression techniques, helping in making informed decisions about compression tradeoffs.
class CompressionAnalyzer:
def __init__(self):
self.compression_results = {}
def add_result(self, model_name, compression_type, metrics):
if model_name not in self.compression_results:
self.compression_results[model_name] = {}
self.compression_results[model_name][compression_type] = metrics
def plot_compression_comparison(self):
import matplotlib.pyplot as plt
import seaborn as sns
# Prepare data for plotting
models = []
compression_types = []
compression_ratios = []
accuracies = []
for model, results in self.compression_results.items():
for comp_type, metrics in results.items():
models.append(model)
compression_types.append(comp_type)
compression_ratios.append(metrics['compression_ratio'])
accuracies.append(metrics['accuracy'])
# Create comparison plot
plt.figure(figsize=(12, 6))
# Plot compression ratio vs accuracy
plt.subplot(1, 2, 1)
sns.scatterplot(
x=compression_ratios,
y=accuracies,
hue=compression_types,
style=models
)
plt.xlabel('Compression Ratio')
plt.ylabel('Accuracy')
plt.title('Compression vs Accuracy Trade-off')
# Plot relative performance
plt.subplot(1, 2, 2)
performance_data = pd.DataFrame({
'Model': models,
'Compression': compression_types,
'Accuracy': accuracies
})
sns.barplot(
data=performance_data,
x='Model',
y='Accuracy',
hue='Compression'
)
plt.xticks(rotation=45)
plt.title('Relative Performance by Compression Method')
plt.tight_layout()
return plt.gcf()
def generate_summary_report(self):
report = []
for model, results in self.compression_results.items():
model_summary = f"\nModel: {model}\n"
model_summary += "-" * 40 + "\n"
for comp_type, metrics in results.items():
model_summary += f"\nCompression Method: {comp_type}\n"
model_summary += f"Compression Ratio: {metrics['compression_ratio']:.2f}x\n"
model_summary += f"Accuracy: {metrics['accuracy']:.2%}\n"
model_summary += f"Inference Time: {metrics['inference_time']:.2f}ms\n"
report.append(model_summary)
return "\n".join(report)
Additional Resources
- "Pruning Neural Networks at Initialization: Why are we missing the mark?"
- "The State of Sparsity in Deep Neural Networks"
- "Knowledge Distillation: A Survey"
- "Movement Pruning: Adaptive Sparsity by Fine-Tuning"
- "What is the State of Neural Network Pruning?"