artubss Destilação de Conhecimento: Comprimindo LLMs
Quando Usar Esta Skill
Use Destilação de Conhecimento quando você precisa:
- Comprimir modelos de 70B → 7B retendo 90%+ de desempenho
- Transferir capacidades de modelos proprietários (GPT-4) para código aberto (LLaMA, Mistral)
- Reduzir custos de inferência implantando modelos alunos menores
- Criar modelos especializados destilando conhecimento específico de domínio
- Melhorar modelos pequenos usando dados sintéticos de professores grandes
Técnicas-chave: Escalamento de temperatura, alvos suaves, KLD reversa (MiniLLM), destilação de logits, destilação de resposta
Papers: Hinton et al. 2015 (arXiv 1503.02531), MiniLLM (arXiv 2306.08543), KD Survey (arXiv 2402.13116)
Instalação
# transformers padrão
pip install transformers datasets accelerate
# Para treinamento
pip install torch deepspeed wandb
# Opcional: implementação MiniLLM
git clone https://github.com/microsoft/LMOps
cd LMOps/minillm
pip install -e .
Início Rápido
Destilação de Conhecimento Básica
import torch
import torch.nn.functional as F
from transformers import AutoModelForCausalLM, AutoTokenizer, Trainer, TrainingArguments
# 1. Carregue modelos professor (grande) e aluno (pequeno)
teacher = AutoModelForCausalLM.from_pretrained(
"meta-llama/Llama-2-70b-hf", # Professor grande
torch_dtype=torch.float16,
device_map="auto"
)
student = AutoModelForCausalLM.from_pretrained(
"meta-llama/Llama-2-7b-hf", # Aluno pequeno
torch_dtype=torch.float16,
device_map="cuda:0"
)
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-70b-hf")
# 2. Defina perda de destilação
def distillation_loss(student_logits, teacher_logits, labels, temperature=2.0, alpha=0.5):
"""
Combine perda dura (cross-entropy) com perda suave (divergência KL).
Args:
temperature: Suaviza distribuições de probabilidade (maior = mais suave)
alpha: Peso para perda de destilação (1-alpha para perda dura)
"""
# Perda dura: Cross-entropy padrão com rótulos verdadeiros
hard_loss = F.cross_entropy(student_logits.view(-1, student_logits.size(-1)), labels.view(-1))
# Perda suave: Divergência KL entre aluno e professor
soft_targets = F.softmax(teacher_logits / temperature, dim=-1)
soft_student = F.log_softmax(student_logits / temperature, dim=-1)
soft_loss = F.kl_div(soft_student, soft_targets, reduction='batchmean') * (temperature ** 2)
# Perda combinada
return alpha * soft_loss + (1 - alpha) * hard_loss
# 3. Loop de treinamento
for batch in dataloader:
# Forward do professor (sem grad)
with torch.no_grad():
teacher_outputs = teacher(**batch)
teacher_logits = teacher_outputs.logits
# Forward do aluno
student_outputs = student(**batch)
student_logits = student_outputs.logits
# Compute perda de destilação
loss = distillation_loss(
student_logits,
teacher_logits,
batch['labels'],
temperature=2.0,
alpha=0.7 # 70% suave, 30% dura
)
# Backward e otimize
loss.backward()
optimizer.step()
optimizer.zero_grad()
MiniLLM (KLD Reversa)
Fonte: arXiv 2306.08543 (2024)
Inovação: Use KLD reversa em vez de KLD direta para melhor destilação de modelos generativos.
def reverse_kl_loss(student_logits, teacher_logits, temperature=1.0):
"""
Divergência KL reversa: KL(Professor || Aluno)
Melhor para modelos generativos que KLD direta.
"""
# Distribuição do professor (alvo)
p_teacher = F.softmax(teacher_logits / temperature, dim=-1)
# Distribuição do aluno (modelo)
log_p_student = F.log_softmax(student_logits / temperature, dim=-1)
# KL reversa: Soma sobre professor, aluno aprende a cobrir modos do professor
reverse_kl = -(p_teacher * log_p_student).sum(dim=-1).mean()
return reverse_kl * (temperature ** 2)
# Treinamento com MiniLLM
for batch in dataloader:
with torch.no_grad():
teacher_logits = teacher(**batch).logits
student_logits = student(**batch).logits
# KLD reversa (melhor para geração)
loss = reverse_kl_loss(student_logits, teacher_logits, temperature=1.0)
loss.backward()
optimizer.step()
Por que KL reversa?
- KL direta (padrão): Aluno aprende a corresponder à média do professor
- KL reversa (MiniLLM): Aluno aprende a cobrir todos os modos do professor
- Melhor para geração diversa de texto
Destilação de Resposta
# Gere dados sintéticos do professor, treine aluno para imitá-lo
# 1. Gere respostas sintéticas do professor
prompts = ["Explique IA:", "O que é ML?", "Defina PLN:"]
teacher_responses = []
for prompt in prompts:
inputs = tokenizer(prompt, return_tensors='pt').to(teacher.device)
outputs = teacher.generate(**inputs, max_new_tokens=256, do_sample=True, temperature=0.7)
response = tokenizer.decode(outputs[0], skip_special_tokens=True)
teacher_responses.append(response)
# 2. Treine aluno nas respostas do professor (fine-tuning padrão)
train_dataset = [
{"text": f"{prompt}\n{response}"}
for prompt, response in zip(prompts, teacher_responses)
]
# 3. Fine-tune do aluno
trainer = Trainer(
model=student,
args=TrainingArguments(output_dir="./student", num_train_epochs=3, learning_rate=2e-5),
train_dataset=train_dataset,
)
trainer.train()
Conceitos Fundamentais
1. Escalamento de Temperatura
Objetivo: Suavizar distribuições de probabilidade para expor incerteza do professor.
# Temperatura baixa (T=1): Distribuição pontuda
logits = [3.0, 2.0, 1.0]
probs_T1 = softmax(logits / 1.0) # [0.67, 0.24, 0.09]
# Temperatura alta (T=4): Distribuição suave
probs_T4 = softmax(logits / 4.0) # [0.42, 0.34, 0.24]
# T maior revela mais informações sobre rankings relativos
Regra: Use T=2-5 para destilação (2 é padrão comum).
2. Componentes da Função de Perda
# Perda total = alpha * perda_suave + (1 - alpha) * perda_dura
# Perda suave: Aprenda do conhecimento do professor
soft_loss = KL(aluno || professor)
# Perda dura: Aprenda dos rótulos verdadeiros
hard_loss = CrossEntropy(saída_aluno, rótulos_verdadeiros)
# Valores típicos:
alpha = 0.5 # Balanceado
alpha = 0.7 # Mais ênfase no professor
alpha = 0.3 # Mais ênfase nos rótulos
3. KLD Direta vs Reversa
# KL direta: KL(Aluno || Professor)
# - Aluno corresponde ao comportamento médio do professor
# - Mode-seeking: Aluno foca nos modos de maior probabilidade do professor
# - Bom para classificação
# KL reversa: KL(Professor || Aluno)
# - Aluno cobre todos os comportamentos do professor
# - Mode-covering: Aluno aprende comportamentos diversos
# - Bom para geração (MiniLLM)
Estratégias de Treinamento
Estratégia 1: Destilação de Logits
# Treine aluno para corresponder aos logits do professor diretamente
def logit_distillation_trainer(student, teacher, dataloader, temperature=2.0):
optimizer = torch.optim.AdamW(student.parameters(), lr=2e-5)
for epoch in range(3):
for batch in dataloader:
# Obtenha logits
with torch.no_grad():
teacher_logits = teacher(**batch).logits
student_logits = student(**batch).logits
# MSE em logits (alternativa a KLD)
loss = F.mse_loss(student_logits, teacher_logits)
# Ou use KLD
# loss = F.kl_div(
# F.log_softmax(student_logits/temperature, dim=-1),
# F.softmax(teacher_logits/temperature, dim=-1),
# reduction='batchmean'
# ) * (temperature ** 2)
loss.backward()
optimizer.step()
optimizer.zero_grad()
return student
Estratégia 2: Destilação em Duas Etapas
# Etapa 1: Destilar do professor
student = distill(teacher, student, epochs=5)
# Etapa 2: Fine-tune em dados específicos de tarefa
student = fine_tune(student, task_data, epochs=3)
# Resulta em melhor desempenho em tarefa que destilação de etapa única
Estratégia 3: Destilação Multi-Professor
# Aprenda de múltiplos professores especialistas
def multi_teacher_distillation(student, teachers, batch):
"""Destilar de conjunto de professores."""
teacher_logits_list = []
# Obtenha logits de todos os professores
with torch.no_grad():
for teacher in teachers:
logits = teacher(**batch).logits
teacher_logits_list.append(logits)
# Média das predições do professor
avg_teacher_logits = torch.stack(teacher_logits_list).mean(dim=0)
# Aluno aprende do conjunto
student_logits = student(**batch).logits
loss = F.kl_div(
F.log_softmax(student_logits, dim=-1),
F.softmax(avg_teacher_logits, dim=-1),
reduction='batchmean'
)
return loss
Implantação em Produção
Script de Treinamento Completo
from transformers import Trainer, TrainingArguments, DataCollatorForLanguageModeling
def train_distilled_model(
teacher_name="meta-llama/Llama-2-70b-hf",
student_name="meta-llama/Llama-2-7b-hf",
output_dir="./distilled-llama-7b",
temperature=2.0,
alpha=0.7,
):
# Carregue modelos
teacher = AutoModelForCausalLM.from_pretrained(teacher_name, torch_dtype=torch.float16, device_map="auto")
student = AutoModelForCausalLM.from_pretrained(student_name, torch_dtype=torch.float16)
tokenizer = AutoTokenizer.from_pretrained(teacher_name)
# Trainer customizado com destilação
class DistillationTrainer(Trainer):
def compute_loss(self, model, inputs, return_outputs=False):
# Forward do aluno
outputs_student = model(**inputs)
student_logits = outputs_student.logits
# Forward do professor (sem grad)
with torch.no_grad():
outputs_teacher = teacher(**inputs)
teacher_logits = outputs_teacher.logits
# Perda de destilação
soft_targets = F.softmax(teacher_logits / temperature, dim=-1)
soft_student = F.log_softmax(student_logits / temperature, dim=-1)
soft_loss = F.kl_div(soft_student, soft_targets, reduction='batchmean') * (temperature ** 2)
# Perda dura
hard_loss = outputs_student.loss
# Combinada
loss = alpha * soft_loss + (1 - alpha) * hard_loss
return (loss, outputs_student) if return_outputs else loss
# Argumentos de treinamento
training_args = TrainingArguments(
output_dir=output_dir,
num_train_epochs=3,
per_device_train_batch_size=4,
gradient_accumulation_steps=8,
learning_rate=2e-5,
warmup_steps=500,
logging_steps=100,
save_steps=1000,
bf16=True,
gradient_checkpointing=True,
)
# Treine
trainer = DistillationTrainer(
model=student,
args=training_args,
train_dataset=train_dataset,
data_collator=DataCollatorForLanguageModeling(tokenizer, mlm=False),
)
trainer.train()
student.save_pretrained(output_dir)
tokenizer.save_pretrained(output_dir)
# Uso
train_distilled_model(
teacher_name="meta-llama/Llama-2-70b-hf",
student_name="meta-llama/Llama-2-7b-hf",
temperature=2.0,
alpha=0.7
)
Melhores Práticas
1. Seleção de Hiperparâmetros
# Temperatura
T = 1.0 # Pontuda (menos transferência de conhecimento)
T = 2.0 # Padrão (bom equilíbrio)
T = 5.0 # Suave (mais transferência de conhecimento)
# Alpha (peso)
alpha = 0.5 # Balanceado
alpha = 0.7 # Enfatizar conhecimento do professor
alpha = 0.9 # Destilação forte
# Regra: T maior + alpha maior = destilação mais forte
2. Razão de Tamanho de Modelo
# Boas razões (professor/aluno)
70B / 7B = 10× # Excelente
13B / 1B = 13× # Bom
7B / 1B = 7× # Aceitável
# Evite diferença muito grande
70B / 1B = 70× # Muito grande, ineficaz
3. Qualidade dos Dados
# Melhor: Use dados gerados pelo professor + dados reais
train_data = {
"teacher_generated": 70%, # Diverso, alta qualidade
"real_data": 30% # Verdade fundamental
}
# Evite: Apenas dados reais (não aproveita professor totalmente)
Avaliação
from transformers import pipeline
# Compare aluno vs professor
teacher_pipe = pipeline("text-generation", model=teacher)
student_pipe = pipeline("text-generation", model=student)
prompts = ["Explique computação quântica:", "O que é IA?"]
for prompt in prompts:
teacher_out = teacher_pipe(prompt, max_new_tokens=100)
student_out = student_pipe(prompt, max_new_tokens=100)
print(f"Prompt: {prompt}")
print(f"Professor: {teacher_out[0]['generated_text']}")
print(f"Aluno: {student_out[0]['generated_text']}")
print(f"Qualidade de correspondência: {calculate_similarity(teacher_out, student_out):.2f}")
Recursos
- Hinton et al. 2015 (Fundacional): https://arxiv.org/abs/1503.02531
- MiniLLM (KLD Reversa): https://arxiv.org/abs/2306.08543
- KD Survey para LLMs (2024): https://arxiv.org/abs/2402.13116
- MiniLLM GitHub: https://github.com/microsoft/LMOps/tree/main/minillm