Balance Multiple Tasks with Multi-Task Finetuning

You've spent weeks perfecting a single-task model, and it performs brilliantly — until someone asks it to also handle a second, third, or tenth task. Fine-tuning an LLM for one narrow purpose is straightforward; but when you need a single model to summarize, classify, and extract entities without forgetting how to do any of them, you're suddenly juggling competing gradients and a sliding scale of priorities. The result? Catastrophic forgetting, skewed performance, and a model that excels at one task while stumbling on the others. This lesson teaches you how to balance multiple tasks with multi-task finetuning — the strategies, the math, and the practical code to train one model that handles many responsibilities gracefully.

The problem this lesson solves

Single-task finetuning feels like a safe bet: you pick one objective, gather focused data, and monitor a single loss curve. But production LLMs rarely serve one purpose. A customer-support bot might need to detect intent, extract order numbers, summarize tickets, and generate replies — all in one turn of dialogue. If you train separate models for each capability, you multiply infrastructure costs and latency. If you cram all tasks into one model naively, you face two notorious problems:

• Catastrophic forgetting: The model overwrites earlier task knowledge with new task data, so performance on the first task drops sharply after training on the second.
• Task interference: Gradients from different tasks pull the shared weights in conflicting directions, making optimization unstable and final performance worse than training each task separately.

Multi-task finetuning addresses this by training a single model on supervised examples from multiple tasks simultaneously, mixing data in each batch. But mixing data isn't enough — you must balance the tasks so that each contributes appropriately to learning. Without balancing, high-resource tasks dominate the loss, and low-resource tasks get starved. This lesson is your practical guide to that balancing act.

Core concept / mental model

Think of multi-task finetuning as a team project where the LLM is the shared brain, and each task is a teammate with its own dataset and goals. If you let one teammate talk nonstop, the project outcome skews heavily toward their priorities. You need a facilitator — the loss function and data sampler — to ensure every voice is heard.

More formally, you have tasks $T_1, T_2, \ldots, T_k$, each with its own dataset $D_i$. The total training objective is a weighted sum of individual task losses:

$$\mathcal{L}{MTL} = \sum_i(\theta)$$}^{k} \lambda_i \cdot \mathcal{L

Here, $\theta$ is the model parameters, and $\lambda_i$ are task weights. Balancing means choosing $\lambda_i$ and sampling strategies so that no task is neglected and the model learns a robust shared representation.

This mental model is also known as joint training vs. sequential training (fine-tune on task A, then task B). Joint training is your primary tool for avoiding catastrophic forgetting because all tasks are seen together throughout training.

How it works step by step

The multi-task finetuning pipeline follows these stages:

1. Prepare task datasets — each task must have a consistent input/output format, often with a task prefix or instruction to tell the model what to do.
2. Choose a balancing strategy — decide how to sample and weight tasks during training.
3. Mix data in batches — combine examples from all tasks in each training step.
4. Train with weighted loss — apply task-specific loss coefficients and monitor per-task metrics.
5. Evaluate per task — measure performance on a held-out set for each task separately, not just an aggregate score.

Balancing strategies in detail

• Equal sampling: Sample each task with equal probability per batch. Simple, but if task A has 1M examples and task B has 100, task B is massively oversampled relatively — that can be good or bad depending on data quality.
• Proportional sampling: Sample in proportion to dataset size. This follows natural distribution but can cause large datasets to dominate.
• Weighted loss: Scale each task's loss by $\lambda_i$. You can tune these weights, but manual tuning is tedious.
• Dynamic weighting: Adjust $\lambda_i$ during training based on loss scales or gradients — e.g., using uncertainty (homoscedastic) or gradient variance (GradNorm).
• Task grouping: Group similar tasks to reduce negative transfer — e.g., all text classification tasks share a branch.

Why balance matters

Each task's loss has a different scale. A sequence-generation task (like summarization) produces cross-entropy loss in the range of 1–5, while a binary classification task might be 0.1–0.5. If you simply average losses, the generation task dominates and the classifier never learns. Balancing counteracts this by normalizing loss scales.

Hands-on walkthrough

Let's implement multi-task finetuning with the Hugging Face transformers and datasets libraries. For this exercise, we'll fine-tune a small GPT-2 model for two tasks: sentiment classification (binary) and email intent classification (tri-class). We'll use the imdb dataset for sentiment and a custom mini-dataset for intent.

Step 1: Load and prepare datasets

from datasets import load_dataset

# Sentiment task (binary)
sentiment_ds = load_dataset("imdb", split="train[:2000]").map(
    lambda x: {"text": x["text"], "sentiment_label": 0 if x["label"] == 0 else 1, "task": "sentiment"}
).select_columns(["text", "sentiment_label", "task"])

# Intent task (tri-class) — create a tiny demo dataset
intent_data = [
    {"text": "Please cancel my order", "intent_label": 0, "task": "intent"},
    {"text": "Where is my package?", "intent_label": 1, "task": "intent"},
    {"text": "I want a refund", "intent_label": 2, "task": "intent"},
    {"text": "Can you help me with a return?", "intent_label": 2, "task": "intent"},
    {"text": "Track my shipment urgently", "intent_label": 1, "task": "intent"},
    # Add more examples...
]
intent_ds = Dataset.from_list(intent_data)

Step 2: Tokenize and combine

from transformers import AutoTokenizer
from datasets import concatenate_datasets

tokenizer = AutoTokenizer.from_pretrained("gpt2")
tokenizer.pad_token = tokenizer.eos_token

def tokenize_fn(batch):
    # Add a task-specific prefix to help the model distinguish tasks
    texts = [f"{task}: {text}" for task, text in zip(batch["task"], batch["text"])]
    encodings = tokenizer(texts, truncation=True, padding="max_length", max_length=128)
    encodings["task"] = batch["task"]
    return encodings

sentiment_tokenized = sentiment_ds.map(tokenize_fn, batched=True)
intent_tokenized = intent_ds.map(tokenize_fn, batched=True)

# Combine into one dataset
combined_ds = concatenate_datasets([sentiment_tokenized, intent_tokenized])

Step 3: Set up dual-head model and train

For simplicity, we'll use a single model with separate classification layers for each task.

import torch
import torch.nn as nn
from transformers import AutoModel, Trainer, TrainingArguments

class MultiTaskGPT2(nn.Module):
    def __init__(self, base_model_name):
        super().__init__()
        self.base = AutoModel.from_pretrained(base_model_name)
        self.sentiment_head = nn.Linear(self.base.config.hidden_size, 2)
        self.intent_head = nn.Linear(self.base.config.hidden_size, 3)

    def forward(self, input_ids, attention_mask, task=None, labels=None):
        outputs = self.base(input_ids, attention_mask=attention_mask)
        pooled = outputs.last_hidden_state[:, 0, :]  # take CLS (or first token)
        logits = None
        loss = None
        if task is not None:
            # Compute logits for all tasks but only use the relevant for loss
            sent_logits = self.sentiment_head(pooled)
            intent_logits = self.intent_head(pooled)
            if labels is not None:
                loss = 0.0
                for i, t in enumerate(task):
                    if t == "sentiment":
                        loss += nn.functional.cross_entropy(sent_logits[i:i+1], labels[i:i+1])
                    else:
                        loss += nn.functional.cross_entropy(intent_logits[i:i+1], labels[i:i+1])
            logits = {"sentiment": sent_logits, "intent": intent_logits}
        return {"logits": logits, "loss": loss}

Pro tip: For real projects, avoid writing custom trainer logic — use Hugging Face Trainer with compute_metrics per task. The above is illustrative.

Step 4: Weighted loss for balancing

Use the uncertainty weighting approach:

class UncertaintyWeightedLoss(nn.Module):
    def __init__(self, num_tasks=2):
        super().__init__()
        self.log_vars = nn.Parameter(torch.zeros(num_tasks))

    def forward(self, losses):
        # losses is a list of per-task losses (scalar tensors)
        total = 0
        for i, loss in enumerate(losses):
            precision = torch.exp(-self.log_vars[i])
            total += precision * loss + self.log_vars[i]
        return total

Apply this in the training loop by accumulating per-task losses before combining.

Expected output / behavior

After training for a few epochs, you should observe:

• Both sentiment and intent validation accuracy improve over time.
• The loss curve for each task should be roughly stable — if one task's loss flatlines, your weights are off.
• A final evaluation shows balanced accuracy (e.g., sentiment 82%, intent 78%) rather than one at 90% and the other at 50%.

Compare options / when to choose what

Choosing rule of thumb: start with proportional sampling and manual or uncertainty loss weighting. If you see imbalance, switch to uncertainty weighting.

Troubleshooting & edge cases

• One task's loss keeps decreasing while another stagnates — your loss weights are causing negative transfer. Try increasing the weight for the stagnating task or using uncertainty weighting.
• Model performs well on one task, poorly on the other in evaluation — this is the classic imbalance symptoms. Check your data sampler: if you used proportional sampling with a huge dataset, undersample the large task or oversample the small one.
• Catastrophic forgetting still occurs — you might be training too long or with too high a learning rate. Use a small learning rate (1e-5) and early stopping based on the worst-performing task.
• Task confusion — the model mixes task behaviors (e.g., outputs sentiment labels when asked intent). Ensure task prefixes are distinct and add a separate adapter or head per task if needed.

Pro tip: Always evaluate on a held-out set per task during training, not just at the end. This lets you catch imbalance early and adjust weights on the fly.

What you learned & what's next

In this lesson, you learned how to balance multiple tasks with multi-task finetuning — covering the core concept of weighted loss and sampling strategies, the step-by-step pipeline, and a hands-on Python implementation. You now understand how to identify and fix task interference and catastrophic forgetting. The next step in your LLM Finetuning track is mastering parameter-efficient fine-tuning — or if you're already there, dive into evaluation and validation to ensure your multi-task model truly performs. You're building the skills to deploy robust, multi-capable LLMs in production.