22. Multi-Task Fine-Tuning

Chapter 22 of 24 · 20 min

KEY INSIGHT

Multi-task fine-tuning trains a single model on multiple tasks simultaneously, potentially enabling generalization and reducing deployment complexity. The challenge lies in balancing task contributions and preventing negative transfer where one task degrades another. The dataset format must accommodate task identifiers or the model must learn to route based on input patterns. Each approach has tradeoffs in flexibility and complexity. ```python class MultiTaskDataset(Dataset): def __init__(self, task_configs): """ task_configs: list of dicts with 'task_name', 'dataset', 'prompt_template' """ self.datasets = [] for config in task_configs: dataset = load_dataset(**config["dataset"]) self.datasets.append({ "name": config["task_name"], "dataset": dataset, "template": config["prompt_template"] }) def __len__(self): return sum(len(ds["dataset"]) for ds in self.datasets) def __getitem__(self, idx): # Find which dataset contains this index for ds in self.datasets: if idx < len(ds["dataset"]): item = ds["dataset"][idx] return { "task": ds["name"], "input": ds["template"].format(**item), "label": item.get("label", item.get("output")) } idx -= len(ds["dataset"]) ``` **Task weighting strategies** address imbalance between large and small task datasets: ```python def create_balanced_batch(dataset, batch_size, temperature=2.0): """ Sample tasks with temperature-based weighting. Smaller datasets get higher probability. """ task_sizes = {ds["name"]: len(ds["dataset"]) for ds in dataset.datasets} weights = {k: v ** (1/temperature) for k, v in task_sizes.items()} total = sum(weights.values()) weights = {k: v/total for k, v in weights.items()} # Sample task, then sample from that task's dataset tasks = list(weights.keys()) probs = list(weights.values()) chosen_task = random.choices(tasks, weights=probs, k=batch_size) # Build batch from chosen task batch = [] for task_name in chosen_task: sample = dataset.sample_from_task(task_name) batch.append(sample) return collate_fn(batch) ``` **Gradient balancing** prevents dominant tasks from overwhelming small ones: ```python class GradientBalancer: def __init__(self, task_weights, alpha=0.5): self.task_weights = task_weights # Dynamic weights self.alpha = alpha # Balancing strength def compute_balanced_loss(self, losses, task_names): """ Apply gradient regularization to balance task contributions. """ base_loss = sum(losses) / len(losses) # Encourage task-specific heads to learn independently regularization = self.alpha * sum( torch.std(l) for l in losses if l.requires_grad ) return base_loss + regularization ```

EXERCISE

: Implement Task Routing

def evaluate_multitask(model, task_datasets):
    """Report per-task performance."""
    results = {}
    for task_name, dataset in task_datasets.items():
        metrics = evaluate_model(model, dataset)
        results[task_name] = metrics
        print(f"{task_name}: {metrics}")
    return results

Run this before and after multi-task training to measure generalization.