TextBrewer Documentation

TextBrewer is a PyTorch-based toolkit designed for Knowledge Distillation (KD) in NLP tasks. It is model-agnostic (primarily targeting Transformer architectures), non-intrusive (no need to modify model structures), and supports various NLP tasks like text classification, reading comprehension, and sequence labeling.

The framework is organized into three main functional blocks:

1. Distillers: The core engine of the distillation process. Different distiller classes provide different distillation modes. Available distillers include GeneralDistiller, MultiTeacherDistiller, and MultiTaskDistiller.
2. Configurations and Presets: Handles training and distillation method configurations. It provides predefined distillation strategies and various knowledge distillation loss functions.
3. Utilities: Auxiliary tools, such as model parameter analysis.

To use TextBrewer, you need a trained teacher model, a student model (to be distilled), and your training dataset.

TextBrewer is a PyTorch-based toolkit for knowledge distillation of NLP models, specifically designed for transformer-based architectures. It supports various tasks like text classification, machine reading comprehension, and sequence labeling.

Core Components

1. Distillers: The core engines of distillation. Examples include GeneralDistiller, MultiTeacherDistiller, and BasicTrainer.
2. Configurations and presets: Classes for managing training and distillation settings, including predefined loss functions and strategies.
3. Utilities: Auxiliary tools, such as model parameter analysis.

Key Features

• Mixed soft-label and hard-label training.
• Dynamic loss weight and temperature adjustment.
• Various distillation loss functions (e.g., hidden states MSE, attention-matrix-based loss, neuron selectivity transfer).
• Support for intermediate feature matching losses.
• Multi-teacher distillation.
• Flexibility to use user-defined loss functions and modules without modifying model architectures.

To optimize training, you can provide pre-computed teacher outputs via a dataset. If your dataset returns a dictionary containing the key 'teacher_cache', TextBrewer will treat the contents of batch['teacher_cache'] as the output from the teacher and feed it directly to the teacher's adaptor.

When 'teacher_cache' is present, the teacher's forward method is not called, saving computation time.

Note: The items in teacher_cache should match the expected output format of the teacher model (e.g., a tuple of (logits, loss)).

import torch
from torch.utils.data import Dataset, TensorDataset, DataLoader

class TSDataset(Dataset):
    def __init__(self, teacher_dataset, student_dataset, teacher_cache):
        # teacher_dataset and student_dataset are normal datasets 
        # whose each element is a tuple or a dict.
        # teacher_cache is a list of items; each item is the output from the teacher.
        assert len(teacher_dataset) == len(student_dataset), \
            f"lengths of teacher_dataset {len(teacher_dataset)} and student_dataset {len(student_dataset)} are not the same!"
        assert len(teacher_dataset) == len(teacher_cache), \
            f"lengths of teacher_dataset {len(teacher_dataset)} and teacher_cache {len(teacher_cache)} are not the same!"
        self.teacher_dataset = teacher_dataset
        self.student_dataset = student_dataset
        self.teacher_cache = teacher_cache

    def __len__(self):
        return len(self.teacher_dataset)

    def __getitem__(self,i):
        return {'teacher' : self.teacher_dataset[i], 'student' : self.student_dataset[i], 'teacher_cache':self.teacher_cache[i]}

teacher_dataset = TensorDataset(torch.randn(32,3),torch.randn(32,3))
student_dataset = TensorDataset(torch.randn(32,2),torch.randn(32,2))

# We make some fake data and assume teacher model outputs are (logits, loss)
fake_logits = [torch.randn(3) for _ in range(32)]
fake_loss = [torch.randn(1)[0] for _ in range(32)]
teacher_cache = [(fake_logits[i],fake_loss[i]) for i in range(32)]

tssdataset = TSDataset(teacher_dataset=teacher_dataset,student_dataset=student_dataset, teacher_cache=teacher_cache)
dataloader = DataLoader(dataset=tsdataset, ... )

When distilling models with different vocabularies (e.g., RoBERTa to BERT), the student and teacher cannot share the same input. To handle this, you can create a custom Dataset that returns a dictionary containing separate keys for the student and teacher inputs.

TextBrewer will automatically unpack the dictionary and feed batch['student'] to the student (and its adaptor) and batch['teacher'] to the teacher (and its adaptor) following the standard forward conventions.

import torch
from torch.utils.data import Dataset, TensorDataset, DataLoader

class TSDataset(Dataset):
    def __init__(self, teacher_dataset, student_dataset):
        # teacher_dataset and student_dataset are normal datasets 
        # whose each element is a tuple or a dict.
        assert len(teacher_dataset) == len(student_dataset), \
            f"lengths of teacher_dataset {len(teacher_dataset)} and student_dataset {len(student_dataset)} are not the same!"

        self.teacher_dataset = teacher_dataset
        self.student_dataset = student_dataset

    def __len__(self):
        return len(self.teacher_dataset)

    def __getitem__(self,i):
        return {'teacher' : self.teacher_dataset[i], 'student' : self.student_dataset[i]}

teacher_dataset = TensorDataset(torch.randn(32,3),torch.randn(32,3))
student_dataset = TensorDataset(torch.randn(32,2),torch.randn(32,2))
tssdataset = TSDataset(teacher_dataset=teacher_dataset,student_dataset=student_dataset)
dataloader = DataLoader(dataset=tsdataset, ... )