Model optimization (distillation, quantization) in NLP

Introduction

Model optimization helps make AI models smaller and faster without losing much accuracy. This is useful to run models on devices with less power or memory.

When you want to run a language model on a smartphone with limited memory.

When you need faster responses from a chatbot by making the model smaller.

When deploying AI models on edge devices like smart cameras or IoT gadgets.

When reducing cloud computing costs by using smaller models.

When you want to keep the model's accuracy but make it easier to share or download.

Syntax

NLP

from transformers import DistilBertForSequenceClassification, DistilBertTokenizer
import torch

# Load original model and tokenizer
model_name = 'bert-base-uncased'
tokenizer = DistilBertTokenizer.from_pretrained(model_name)

# Distillation example: load a smaller distilled model
distilled_model = DistilBertForSequenceClassification.from_pretrained('distilbert-base-uncased')

# Quantization example (PyTorch dynamic quantization)
quantized_model = torch.quantization.quantize_dynamic(
    distilled_model, {torch.nn.Linear}, dtype=torch.qint8
)

Distillation means training a smaller model to mimic a bigger model's behavior.

Quantization means reducing the precision of numbers in the model to save space and speed up.

Examples

This loads a smaller version of BERT that is faster and lighter.

NLP

# Distillation: Load a smaller pretrained distilled model
from transformers import DistilBertForSequenceClassification
model = DistilBertForSequenceClassification.from_pretrained('distilbert-base-uncased')

This reduces model size and speeds up inference by using 8-bit integers instead of 32-bit floats.

NLP

# Quantization: Apply dynamic quantization to a PyTorch model
import torch
quantized_model = torch.quantization.quantize_dynamic(
    model, {torch.nn.Linear}, dtype=torch.qint8
)

Sample Model

This code loads a small distilled BERT model, runs a sample sentence through it, then applies quantization to make it smaller and faster. It compares the outputs before and after quantization.

NLP

from transformers import DistilBertForSequenceClassification, DistilBertTokenizer
import torch

# Load tokenizer and distilled model
model_name = 'distilbert-base-uncased'
tokenizer = DistilBertTokenizer.from_pretrained(model_name)
model = DistilBertForSequenceClassification.from_pretrained(model_name)

# Sample text
text = "I love learning about AI!"
inputs = tokenizer(text, return_tensors='pt')

# Run original model
with torch.no_grad():
    original_outputs = model(**inputs)
    original_logits = original_outputs.logits

# Apply dynamic quantization
quantized_model = torch.quantization.quantize_dynamic(
    model, {torch.nn.Linear}, dtype=torch.qint8
)

# Run quantized model
with torch.no_grad():
    quantized_outputs = quantized_model(**inputs)
    quantized_logits = quantized_outputs.logits

# Print logits from both models
print(f"Original logits: {original_logits}")
print(f"Quantized logits: {quantized_logits}")

Important Notes

Quantization may slightly reduce accuracy but improves speed and size.

Distillation requires training or using a pretrained smaller model.

Always test optimized models to ensure they still work well for your task.

Summary

Model optimization makes AI models smaller and faster.

Distillation trains a smaller model to copy a bigger one.

Quantization reduces number precision to save space and speed up.