13. Multi-Task Learning — Single Model for Intent, Sentiment, and Entities

graph TB
    subgraph "Shared Encoder (DistilBERT, 66M params)"
        INPUT["Input: 'I want to return this damaged manga volume'"]
        TOK["Tokenizer → [CLS] I want to return this damaged manga volume [SEP]"]
        EMB["Embedding Layer (shared)"]
        L1["Transformer Layer 1 (shared)"]
        L2["Transformer Layer 2 (shared)"]
        L3["Transformer Layer 3 (shared)"]
        L4["Transformer Layer 4 (shared)"]
        L5["Transformer Layer 5 (shared)"]
        L6["Transformer Layer 6 (shared)"]

        INPUT --> TOK --> EMB --> L1 --> L2 --> L3 --> L4 --> L5 --> L6
    end

    subgraph "Task-Specific Heads (2M params total)"
        CLS["[CLS] token → 768-dim"]
        SEQ["All tokens → 768-dim each"]

        L6 --> CLS & SEQ

        H_INT["Intent Head<br>Linear(768→256→10)<br>Softmax<br>Prediction: return_refund"]
        H_SENT["Sentiment Head<br>Linear(768→256→6)<br>Sigmoid per label<br>Pred: frustration=0.82"]
        H_NER["NER Head<br>Linear(768→256→9)<br>Per-token softmax<br>Pred: manga=PRODUCT"]

        CLS --> H_INT & H_SENT
        SEQ --> H_NER
    end

    style L1 fill:#e3f2fd
    style L2 fill:#e3f2fd
    style L3 fill:#e3f2fd
    style L4 fill:#e3f2fd
    style L5 fill:#e3f2fd
    style L6 fill:#e3f2fd
    style H_INT fill:#c8e6c9
    style H_SENT fill:#fff9c4
    style H_NER fill:#ffccbc

graph TB
    subgraph "Gradient Conflict in Shared Layers"
        LOSS_I["Intent Loss = 0.42<br>∇W₃(intent)"]
        LOSS_S["Sentiment Loss = 0.68<br>∇W₃(sentiment)"]
        LOSS_N["NER Loss = 0.31<br>∇W₃(NER)"]

        CHECK["Conflict Detection:<br>g_intent · g_sentiment = -0.15 ❌ CONFLICT<br>g_intent · g_ner = +0.34 ✅ aligned<br>g_sentiment · g_ner = +0.08 ✅ aligned"]

        LOSS_I --> CHECK
        LOSS_S --> CHECK
        LOSS_N --> CHECK

        PCGRAD["PCGrad Resolution:<br>g'_intent = g_intent - proj(g_intent, g_sentiment)<br>Remove conflicting component<br><br>g'_sentiment = g_sentiment - proj(g_sentiment, g_intent)<br>Remove conflicting component<br><br>g_ner unchanged (no conflicts)"]

        CHECK --> PCGRAD

        FINAL["Final gradient for Layer 3:<br>g_shared = g'_intent + g'_sentiment + g_ner<br>No task dominance, no destructive interference"]

        PCGRAD --> FINAL
    end

    style CHECK fill:#ffcdd2
    style PCGRAD fill:#c8e6c9
    style FINAL fill:#c8e6c9

graph LR
    subgraph "Epoch 1: Equal weights"
        E1_I["Intent<br>w=0.33<br>Loss=2.30<br>σ²=2.30"]
        E1_S["Sentiment<br>w=0.33<br>Loss=0.85<br>σ²=0.85"]
        E1_N["NER<br>w=0.33<br>Loss=1.45<br>σ²=1.45"]
    end

    subgraph "Epoch 5: Weights adapted"
        E5_I["Intent<br>w=0.21 ⬇<br>Loss=0.95<br>Easy → downweight"]
        E5_S["Sentiment<br>w=0.38 ⬆<br>Loss=0.62<br>Hard → upweight"]
        E5_N["NER<br>w=0.41 ⬆<br>Loss=0.78<br>Still learning"]
    end

    subgraph "Epoch 15: Converged"
        E15_I["Intent<br>w=0.25<br>Loss=0.42<br>Converged"]
        E15_S["Sentiment<br>w=0.35<br>Loss=0.38<br>Near-converged"]
        E15_N["NER<br>w=0.40<br>Loss=0.31<br>Converged"]
    end

    E1_I -->|"σ adapts"| E5_I -->|"stabilizes"| E15_I
    E1_S -->|"σ adapts"| E5_S -->|"stabilizes"| E15_S
    E1_N -->|"σ adapts"| E5_N -->|"stabilizes"| E15_N

graph TB
    subgraph "GradNorm: Normalize gradient magnitudes"
        G1["Without GradNorm:<br>‖∇intent‖ = 2.5<br>‖∇sentiment‖ = 0.3<br>‖∇NER‖ = 1.1<br><br>Intent dominates 8.3×<br>over sentiment!"]

        GN["GradNorm target (α=1.5):<br>r_intent = L(t)/L(0) = 0.18 (fast)<br>r_sentiment = 0.73 (slow)<br>r_NER = 0.54<br><br>Target ‖∇intent‖ = 0.49<br>Target ‖∇sentiment‖ = 1.80<br>Target ‖∇NER‖ = 1.21"]

        G2["With GradNorm:<br>w_intent adjusted down → ‖g‖ ≈ 0.5<br>w_sentiment adjusted up → ‖g‖ ≈ 1.8<br>w_NER adjusted slightly → ‖g‖ ≈ 1.2<br><br>Balanced! Sentiment gets<br>more gradient signal"]

        G1 --> GN --> G2
    end

    style G1 fill:#ffcdd2
    style G2 fill:#c8e6c9

graph TB
    subgraph "Intent Head (Classification)"
        CLS1["[CLS] embedding: 768-dim"]
        INT_D["Dropout(0.1)"]
        INT_L1["Linear(768→256) + GELU"]
        INT_LN["LayerNorm(256)"]
        INT_L2["Linear(256→10)"]
        INT_SF["Softmax → 10 intent probabilities"]

        CLS1 --> INT_D --> INT_L1 --> INT_LN --> INT_L2 --> INT_SF
    end

    subgraph "Sentiment Head (Multi-Label)"
        CLS2["[CLS] embedding: 768-dim"]
        SENT_D["Dropout(0.1)"]
        SENT_L1["Linear(768→256) + GELU"]
        SENT_LN["LayerNorm(256)"]
        SENT_L2["Linear(256→6)"]
        SENT_SIG["Sigmoid per label → 6 independent probabilities"]

        CLS2 --> SENT_D --> SENT_L1 --> SENT_LN --> SENT_L2 --> SENT_SIG
    end

    subgraph "NER Head (Token Classification)"
        TOK["All token embeddings: [n×768]"]
        NER_D["Dropout(0.1)"]
        NER_L1["Linear(768→256) + GELU"]
        NER_LN["LayerNorm(256)"]
        NER_L2["Linear(256→9)"]
        NER_SF["Per-token Softmax → BIO tags"]

        TOK --> NER_D --> NER_L1 --> NER_LN --> NER_L2 --> NER_SF
    end

    style INT_SF fill:#c8e6c9
    style SENT_SIG fill:#fff9c4
    style NER_SF fill:#ffccbc

import torch
import torch.nn as nn
from transformers import DistilBertModel, DistilBertTokenizer


class MultiTaskHead(nn.Module):
    """A single task head with configurable output."""

    def __init__(self, input_dim: int, hidden_dim: int, output_dim: int, dropout: float = 0.1):
        super().__init__()
        self.head = nn.Sequential(
            nn.Dropout(dropout),
            nn.Linear(input_dim, hidden_dim),
            nn.GELU(),
            nn.LayerNorm(hidden_dim),
            nn.Linear(hidden_dim, output_dim),
        )

    def forward(self, x):
        return self.head(x)


class MultiTaskDistilBERT(nn.Module):
    """
    Single DistilBERT encoder with three task-specific heads.
    Replaces 3 separate models → 65% memory reduction.
    """

    def __init__(
        self,
        num_intents: int = 10,
        num_sentiments: int = 6,
        num_ner_tags: int = 9,
        model_name: str = "distilbert-base-uncased",
    ):
        super().__init__()
        self.encoder = DistilBertModel.from_pretrained(model_name)
        hidden = self.encoder.config.hidden_size  # 768

        self.intent_head = MultiTaskHead(hidden, 256, num_intents)
        self.sentiment_head = MultiTaskHead(hidden, 256, num_sentiments)
        self.ner_head = MultiTaskHead(hidden, 256, num_ner_tags)

    def forward(self, input_ids, attention_mask, task: str = "all"):
        encoder_output = self.encoder(input_ids=input_ids, attention_mask=attention_mask)
        hidden_states = encoder_output.last_hidden_state  # [batch, seq_len, 768]
        cls_output = hidden_states[:, 0, :]  # [batch, 768]

        results = {}
        if task in ("all", "intent"):
            results["intent_logits"] = self.intent_head(cls_output)
        if task in ("all", "sentiment"):
            results["sentiment_logits"] = self.sentiment_head(cls_output)
        if task in ("all", "ner"):
            results["ner_logits"] = self.ner_head(hidden_states)

        return results

class UncertaintyWeightedLoss(nn.Module):
    """
    Kendall et al. 2018: automatically learn task weights
    from homoscedastic uncertainty.
    """

    def __init__(self, num_tasks: int = 3):
        super().__init__()
        # Log-variance parameters (learnable)
        self.log_vars = nn.Parameter(torch.zeros(num_tasks))

    def forward(self, losses: list[torch.Tensor]) -> tuple[torch.Tensor, dict]:
        total_loss = 0
        weights = {}

        for i, loss in enumerate(losses):
            # L_total = (1/2σ²) * L_t + log(σ)
            # Using log_var = log(σ²):
            precision = torch.exp(-self.log_vars[i])  # 1/σ²
            total_loss += 0.5 * precision * loss + 0.5 * self.log_vars[i]
            weights[f"task_{i}_weight"] = precision.item()
            weights[f"task_{i}_sigma"] = torch.exp(0.5 * self.log_vars[i]).item()

        return total_loss, weights

import copy


class PCGrad:
    """
    Yu et al. 2020: Project Conflicting Gradients.
    Removes conflicting gradient components between tasks.
    """

    def __init__(self, optimizer):
        self.optimizer = optimizer

    def step(self, losses: list[torch.Tensor], shared_params: list[nn.Parameter]):
        """
        Compute per-task gradients, resolve conflicts, then update.
        """
        task_grads = []

        for loss in losses:
            self.optimizer.zero_grad()
            loss.backward(retain_graph=True)
            grads = [p.grad.clone() if p.grad is not None else torch.zeros_like(p)
                     for p in shared_params]
            task_grads.append(grads)

        # Resolve conflicts via projection
        num_tasks = len(losses)
        projected_grads = [list(g) for g in task_grads]  # Deep copy

        for i in range(num_tasks):
            for j in range(num_tasks):
                if i == j:
                    continue
                for k in range(len(shared_params)):
                    g_i = projected_grads[i][k]
                    g_j = task_grads[j][k]

                    dot = (g_i * g_j).sum()
                    if dot < 0:
                        # Conflict detected: project g_i onto plane perpendicular to g_j
                        projected_grads[i][k] = g_i - (dot / (g_j.norm() ** 2 + 1e-8)) * g_j

        # Sum projected gradients and apply
        self.optimizer.zero_grad()
        for k, param in enumerate(shared_params):
            param.grad = sum(projected_grads[i][k] for i in range(num_tasks))

        self.optimizer.step()


class GradNormLoss(nn.Module):
    """
    Chen et al. 2018: Dynamically balance gradient norms.
    """

    def __init__(self, num_tasks: int = 3, alpha: float = 1.5):
        super().__init__()
        self.num_tasks = num_tasks
        self.alpha = alpha
        self.task_weights = nn.Parameter(torch.ones(num_tasks))
        self.initial_losses = None

    def forward(
        self,
        losses: list[torch.Tensor],
        shared_layer: nn.Module,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """
        Returns:
            weighted_loss: The task-weighted MTL loss
            gradnorm_loss: The GradNorm regularization loss
        """
        # Record initial losses for training rate computation
        if self.initial_losses is None:
            self.initial_losses = [l.item() for l in losses]

        # Weighted task loss
        weighted_loss = sum(self.task_weights[i] * losses[i] for i in range(self.num_tasks))

        # Compute gradient norms for shared layer
        grad_norms = []
        for i in range(self.num_tasks):
            gw = torch.autograd.grad(
                self.task_weights[i] * losses[i],
                shared_layer.parameters(),
                retain_graph=True,
                create_graph=True,
            )
            grad_norms.append(torch.norm(torch.cat([g.flatten() for g in gw])))

        # Average gradient norm
        avg_norm = sum(grad_norms) / self.num_tasks

        # Training rate ratios
        loss_ratios = [losses[i].item() / (self.initial_losses[i] + 1e-8)
                       for i in range(self.num_tasks)]
        avg_ratio = sum(loss_ratios) / self.num_tasks
        relative_rates = [r / (avg_ratio + 1e-8) for r in loss_ratios]

        # Target gradient norms
        target_norms = [avg_norm * (r ** self.alpha) for r in relative_rates]

        # GradNorm loss
        gradnorm_loss = sum(
            torch.abs(grad_norms[i] - target_norms[i].detach())
            for i in range(self.num_tasks)
        )

        return weighted_loss, gradnorm_loss

class MultiTaskTrainer:
    """End-to-end MTL trainer with configurable loss balancing."""

    def __init__(
        self,
        model: MultiTaskDistilBERT,
        strategy: str = "uncertainty",  # "uncertainty", "pcgrad", "gradnorm"
    ):
        self.model = model
        self.strategy = strategy

        if strategy == "uncertainty":
            self.loss_module = UncertaintyWeightedLoss(num_tasks=3)
        elif strategy == "gradnorm":
            self.loss_module = GradNormLoss(num_tasks=3, alpha=1.5)

        all_params = list(model.parameters())
        if strategy == "uncertainty":
            all_params += list(self.loss_module.parameters())
        elif strategy == "gradnorm":
            all_params += list(self.loss_module.parameters())

        self.optimizer = torch.optim.AdamW(all_params, lr=2e-5, weight_decay=0.01)

        if strategy == "pcgrad":
            self.pcgrad = PCGrad(self.optimizer)

        self.intent_loss_fn = nn.CrossEntropyLoss()
        self.sentiment_loss_fn = nn.BCEWithLogitsLoss()
        self.ner_loss_fn = nn.CrossEntropyLoss()

    def train_step(self, batch: dict) -> dict:
        self.model.train()
        outputs = self.model(
            batch["input_ids"], batch["attention_mask"], task="all",
        )

        # Compute per-task losses
        intent_loss = self.intent_loss_fn(outputs["intent_logits"], batch["intent_labels"])
        sentiment_loss = self.sentiment_loss_fn(
            outputs["sentiment_logits"], batch["sentiment_labels"].float(),
        )
        ner_loss = self.ner_loss_fn(
            outputs["ner_logits"].view(-1, 9), batch["ner_labels"].view(-1),
        )

        losses = [intent_loss, sentiment_loss, ner_loss]

        if self.strategy == "uncertainty":
            total_loss, weight_info = self.loss_module(losses)
            self.optimizer.zero_grad()
            total_loss.backward()
            torch.nn.utils.clip_grad_norm_(self.model.parameters(), 1.0)
            self.optimizer.step()
            return {"loss": total_loss.item(), **weight_info}

        elif self.strategy == "pcgrad":
            shared_params = list(self.model.encoder.parameters())
            self.pcgrad.step(losses, shared_params)
            # Update task heads normally
            head_loss = sum(losses)
            self.optimizer.zero_grad()
            head_loss.backward()
            self.optimizer.step()
            return {"loss": sum(l.item() for l in losses)}

        elif self.strategy == "gradnorm":
            shared_layer = self.model.encoder.transformer.layer[-1]
            weighted_loss, gn_loss = self.loss_module(losses, shared_layer)
            self.optimizer.zero_grad()
            (weighted_loss + gn_loss).backward()
            self.optimizer.step()
            # Renormalize task weights
            with torch.no_grad():
                self.loss_module.task_weights.data = (
                    self.loss_module.task_weights / self.loss_module.task_weights.sum() * 3
                )
            return {"loss": weighted_loss.item(), "gn_loss": gn_loss.item()}

    def evaluate(self, val_loader) -> dict:
        self.model.eval()
        intent_correct = intent_total = 0
        sentiment_f1_samples = []
        ner_f1_samples = []

        with torch.no_grad():
            for batch in val_loader:
                outputs = self.model(
                    batch["input_ids"], batch["attention_mask"], task="all",
                )

                # Intent accuracy
                preds = outputs["intent_logits"].argmax(dim=-1)
                intent_correct += (preds == batch["intent_labels"]).sum().item()
                intent_total += batch["intent_labels"].size(0)

                # Sentiment F1 (per-label threshold)
                sent_preds = (torch.sigmoid(outputs["sentiment_logits"]) > 0.5).float()
                # Simplified per-sample F1
                for i in range(sent_preds.size(0)):
                    tp = (sent_preds[i] * batch["sentiment_labels"][i]).sum()
                    fp = (sent_preds[i] * (1 - batch["sentiment_labels"][i])).sum()
                    fn = ((1 - sent_preds[i]) * batch["sentiment_labels"][i]).sum()
                    prec = tp / (tp + fp + 1e-8)
                    rec = tp / (tp + fn + 1e-8)
                    f1 = 2 * prec * rec / (prec + rec + 1e-8)
                    sentiment_f1_samples.append(f1.item())

        return {
            "intent_accuracy": intent_correct / intent_total,
            "sentiment_f1": sum(sentiment_f1_samples) / len(sentiment_f1_samples),
        }