大语言模型轻量化：知识蒸馏的范式迁移与工程实践

class HierarchicalDistiller(nn.Module):
    def __init__(self, teacher, student):
        super().__init__()
        self.teacher = teacher
        self.student = student
        
    def forward(self, inputs):
        # 分层知识提取
        t_hidden_states = self.teacher(**inputs, output_hidden_states=True).hidden_states
        s_hidden_states = self.student(**inputs, output_hidden_states=True).hidden_states
        
        # 多尺度损失计算
        loss = 0
        for t_hid, s_hid in zip(t_hidden_states[::2], s_hidden_states):  # 分层采样
            loss += F.kl_div(
                F.log_softmax(s_hid / self.temp, dim=-1),
                F.softmax(t_hid.detach() / self.temp, dim=-1),
                reduction='batchmean'
            ) * (self.temp ** 2)
        return loss

from transformers import BertTokenizer, BertForSequenceClassification
from datasets import load_dataset

# 加载预训练模型
teacher = BertForSequenceClassification.from_pretrained('bert-base-uncased')
student = TinyBertForSequenceClassification(
    config=TinyBertConfig(
        num_labels=2,
        num_hidden_layers=4,
        intermediate_size=512
    )
)

# 准备GLUE数据集
dataset = load_dataset('glue', 'sst2')
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')

def preprocess(examples):
    return tokenizer(examples['sentence'], truncation=True, padding='max_length')
dataset = dataset.map(preprocess, batched=True)

import torch
from torch.optim import AdamW
from accelerate import Accelerator

accelerator = Accelerator()
device = accelerator.device

optimizer = AdamW(student.parameters(), lr=5e-5)
teacher, student, optimizer = accelerator.prepare(teacher, student, optimizer)

for epoch in range(10):
    for batch in train_dataloader:
        with torch.no_grad():
            teacher_outputs = teacher(**batch)
            
        student_outputs = student(**batch)
        
        # 分层蒸馏损失
        loss = hierarchical_distill_loss(
            student_outputs.hidden_states,
            teacher_outputs.hidden_states,
            temperature=current_temp(epoch)
        )
        
        accelerator.backward(loss)
        optimizer.step()
        optimizer.zero_grad()

class FederatedDistiller:
    def aggregate(self, client_models):
        # 模型参数安全聚合
        secure_params = homomorphic_encryption(
            [model.state_dict() for model in client_models]
        )
        self.global_model.load_state_dict(secure_params)
        
    def client_update(self, local_data):
        # 本地差分隐私训练
        noise = laplace_noise(scale=1.0/self.epsilon)
        return local_model.state_dict() + noise

import torch
import torch.nn as nn
import torch.optim as optim
import torchvision.transforms as transforms
from torch.utils.data import DataLoader
from torchvision import datasets

# 定义教师模型（Teacher Model）
class TeacherModel(nn.Module):
    def __init__(self):
        super(TeacherModel, self).__init__()
        self.fc = nn.Sequential(
            nn.Linear(784, 512),
            nn.ReLU(),
            nn.Linear(512, 256),
            nn.ReLU(),
            nn.Linear(256, 10)
        )
    
    def forward(self, x):
        return self.fc(x)

# 定义学生模型（Student Model）
class StudentModel(nn.Module):
    def __init__(self):
        super(StudentModel, self).__init__()
        self.fc = nn.Sequential(
            nn.Linear(784, 128),
            nn.ReLU(),
            nn.Linear(128, 10)
        )
    
    def forward(self, x):
        return self.fc(x)

# 蒸馏损失函数
def distillation_loss(student_logits, teacher_logits, labels, T=2.0, alpha=0.5):
    soft_loss = nn.KLDivLoss()(nn.functional.log_softmax(student_logits / T, dim=1),
                               nn.functional.softmax(teacher_logits / T, dim=1)) * (T * T)
    hard_loss = nn.CrossEntropyLoss()(student_logits, labels)
    return alpha * soft_loss + (1 - alpha) * hard_loss

# 训练过程
def train_model():
    teacher = TeacherModel()
    teacher.load_state_dict(torch.load('teacher_model.pth'))  # 预训练的教师模型
    teacher.eval()  # 设置为评估模式

    student = StudentModel()
    optimizer = optim.Adam(student.parameters(), lr=0.001)
    
    for epoch in range(10):
        for images, labels in train_loader:
            optimizer.zero_grad()
            teacher_logits = teacher(images.view(-1, 784)).detach()  # 不更新教师模型参数
            student_logits = student(images.view(-1, 784))
            loss = distillation_loss(student_logits, teacher_logits, labels)
            loss.backward()
            optimizer.step()

# 数据加载与训练
train_loader = DataLoader(datasets.MNIST('.', train=True, download=True, transform=transforms.ToTensor()), batch_size=32, shuffle=True)
train_model()