Multi-Task Learning for RecSys

The Big Picture

Multi-Task Learning (MTL) is a paradigm where a single neural network learns to optimize multiple objectives simultaneously, sharing a common representation while maintaining task-specific output layers. Instead of training separate models for each objective, MTL leverages the shared structure between related tasks to improve generalization and efficiency.

Why Multi-Task Learning?

Consider an ad platform that needs to predict three outcomes:

• Click-Through Rate (CTR): Will the user click the ad?
• Conversion Rate (CVR): Will the user convert after clicking?
• Dwell Time (Satisfaction): How long will the user engage?

Training three separate models is expensive: redundant feature engineering, training compute, and inference latency. More importantly, the tasks share underlying patterns. A user who clicks ads is likely to be in a high-intent state. Both CTR and CVR depend on user-item similarity and contextual factors. Dwell time correlates with content quality, which CTR and CVR also depend on. MTL captures these shared representations in a single network.

MTL Paradigms

Beyond these, we have Mixture of Experts (MoE) architectures where expert sub-networks specialize, and a gate selects which experts to use. This can scale to many tasks without parameter explosion.

The Business Problem

Recommendation and ranking systems optimize multiple objectives, but naively optimizing CTR alone creates clickbait. The real business metric is Gross Merchandise Value (GMV), which factors in multiple stages of the conversion funnel.

The Objectives

In an e-commerce or ad platform, consider these objectives:

• Click-Through Rate (CTR): P(click | impression). Measures visibility and relevance.
• Conversion Rate (CVR): P(conversion | click). Measures purchase intent and product-market fit.
• Average Order Value (AOV): Average price or margin per conversion. Maximizes revenue.
• Satisfaction / Engagement: Dwell time, repeat rate, long-term retention. Prevents user fatigue and burnout.
• Diversity: Avoid showing the same content repeatedly. Improves user experience and coverage.

The Tension: Why Not Just Optimize CTR?

Optimizing CTR alone leads to clickbait: sensational, low-quality items that users click but don't convert. The true business metric is:

GMV = CTR × CVR × Price × Quantity

If you maximize only CTR, you get more clicks but lower conversion quality, hurting overall GMV. The challenge is that we don't see all outcomes equally:

• We observe CTR for all impressions (you see whether a user clicked or not).
• We observe CVR only for clicked items (sample selection bias).
• We observe satisfaction only after a long-term engagement.

Sample Selection Bias & Data Sparsity in CVR

This is the heart of the problem that ESMM solves. If we train a CVR model using only clicked items, we're training on biased data:

Hard Parameter Sharing

Hard parameter sharing is the simplest MTL architecture: a shared trunk (set of layers) followed by task-specific head networks. All tasks see the same learned features, then specialize in their output layers.

Architecture

Input x
  ↓
Shared Embedding Layer (d_emb dimensions)
  ↓
Shared Tower (3-5 hidden layers, ReLU)
  ↓
Task 1 Head (2-3 layers) → pCTR
Task 2 Head (2-3 layers) → pCVR
Task 3 Head (2-3 layers) → pSatisfaction

Benefits

• Parameter efficiency: Shared trunk reduces total parameters compared to separate models. One embedding, one backbone, multiple heads.
• Regularization: Sharing features acts as an implicit regularizer. The model is forced to learn representations useful for all tasks, reducing task-specific overfitting.
• Single forward pass: All predictions are made in one forward pass, reducing inference latency.
• Inductive bias: If tasks truly share underlying features, the shared trunk captures the common structure.

Challenges

Hard sharing has a major limitation: task conflict. If one task's gradient points in a direction opposite to another's, the model gets confused. This is called negative transfer.

When It Works Best

• Tasks have high affinity: Strong correlation in outcomes. CTR and CVR are correlated; both favor relevant, high-quality items.
• Tasks share the same feature space: Same user, item, and context features matter for all tasks.
• Training data is imbalanced: Hard sharing provides regularization to smaller datasets.

When It Fails

• Tasks have conflicting objectives: Optimizing for short-term engagement (dwell time) may conflict with long-term retention.
• Tasks have different feature importance: One task relies on categorical features; another on dense features.
• Gradient magnitudes differ greatly: One task's loss dominates, overshadowing others.

Task Weighting

In practice, you need to balance the losses from different tasks. The total loss is:

L_total = λ_ctr * L_CTR + λ_cvr * L_CVR + λ_sat * L_Sat

Weight choices matter significantly. If λ_cvr is too large, the model optimizes CVR at the expense of CTR. Techniques like uncertainty weighting and GradNorm (covered in Section 9) adaptively balance these losses during training.

Soft Parameter Sharing

Soft parameter sharing allows each task to have its own parameters while regularizing them to stay close to each other. This provides flexibility to specialize while maintaining the benefits of shared inductive bias.

Approach 1: L2 Distance Regularization

Each task has its own tower (embedding + hidden layers). The loss includes an L2 regularization term that penalizes differences between task-specific weights:

L_total = Σ_k L_k(θ_k) + λ * Σ_{k,j} ||θ_k - θ_j||²

Here, θ_k are the parameters for task k, and the regularization encourages all task parameters to be similar. This is softer than hard sharing: tasks can diverge if it helps their individual objectives, but at a regularization cost.

Approach 2: Cross-Stitch Networks

Cross-stitch networks learn a combination matrix at each layer that blends representations from multiple task-specific towers:

Layer output for task k:
h_k^(l) = α_kk * g_k^(l) + α_kj * g_j^(l) + α_km * g_m^(l)

where [α] is a learned mixing matrix.
Each row sums to 1 (convex combination).

At each layer, the network learns how much of each task's representation to use. This is more expressive than hard sharing (each layer can specialize) while maintaining coupling through the mixing.

Approach 3: Sluice Networks

Sluice networks use task-specific gating on shared representations. A shared tower feeds all tasks, but each task applies learned masks (gates) to the shared features:

h_shared = Shared_Tower(x)
h_k = Task_Gate_k(h_shared) ⊙ h_shared  # Element-wise gating

where ⊙ is element-wise multiplication.

This is more parameter-efficient than full cross-stitch networks but more flexible than hard sharing.

Benefits & Trade-offs

Soft Sharing Reduces Negative Transfer

Because each task can have specialized parameters, soft sharing handles task conflict better than hard sharing. If tasks have conflicting objectives, they can learn divergent representations while still benefiting from the regularization signal that keeps them somewhat aligned.

Mixture of Experts (MoE)

Mixture of Experts (MoE) is a scaling technique where instead of one shared tower, we use N specialized expert networks, and a gate function learns which experts to use for each input. This allows the model to scale to many tasks without parameter explosion.

How MoE Works

Experts: E₁, E₂, ..., E_N (small neural networks)
Gate: G(x) = softmax(W_g · x) → [w₁, w₂, ..., w_N]

Output: o = Σ wᵢ · Eᵢ(x)   (weighted sum of expert outputs)

Each expert is a small neural network specializing in a sub-problem. The gate receives the input and produces a probability distribution over experts. The final output is a weighted combination of all expert outputs.

Why It Works

MoE enables specialization without parameter explosion:

• Experts specialize: Each expert learns to handle a subset of the input space (e.g., "long-form videos" vs "short ads").
• Sparse activation: Though all experts are evaluated, you can make inference sparse (only top-k experts). Or train with all, allowing learned routing.
• Scalable: Adding tasks doesn't require retraining the shared expert pool; just add a new gate and head.

Load Balancing: The Key Challenge

A major issue with MoE is expert collapse: the gate converges to always using one expert, ignoring others. This wastes capacity. The solution is an auxiliary loss that encourages balanced expert usage:

Load Balancing Loss:
L_aux = CV(load) * CV(importance)

where:
  load_i = average gate value for expert i
  importance_i = sum of gate values selecting expert i
  CV = coefficient of variation (std / mean)

This loss penalizes experts that are used too much or too little, encouraging all experts to contribute.

Computational Considerations

With N experts, the forward pass requires computing all N experts (though you can prune in inference). The dense gate scales O(d_input × N), which is small. If you have 8 experts of dimension 1024 each, and 64 tasks, the capacity scales much better than 64 separate models.

Multi-Gate Mixture of Experts

Multi-gate Mixture of Experts (MMoE) extends the basic MoE architecture by giving each task its own gate over the shared expert pool. This enables task-specific routing: different tasks can weight experts differently based on their objectives.

Architecture

For task k:
  Gate_k(x) = softmax(W_k · x) → [w_k1, w_k2, ..., w_kN]
  o_k = Σ w_ki · E_i(x)

Total output: [o_1, o_2, ..., o_K] for K tasks

The key difference from single-gate MoE: instead of one gate shared by all tasks, each task has its own gate. This allows:

• Task-specific routing: CTR task might route heavily to "engagement experts," while CVR routes to "intent experts."
• Reduced negative transfer: Tasks don't fight over expert capacity; each learns its preferred expert weights.
• Specialization: Experts can learn representations useful to multiple tasks simultaneously.

Why Task-Specific Gates Help

When tasks have conflicting objectives, a shared gate forces compromise. With task-specific gates, each task independently selects which experts matter most:

Load Balancing in MMoE

Load balancing becomes more important with MMoE. Since multiple gates compete for expert capacity, you must prevent one gate from monopolizing an expert. The auxiliary loss is modified to average across tasks:

L_aux = Σ_k (CV(load_k) * CV(importance_k)) / K

where load_k_i = average gate_k value for expert i

This ensures balanced usage across all task gates, not just globally.

Empirical Performance

In recommendation systems, MMoE typically outperforms single-gate MoE when tasks have moderate to high conflict. Studies show:

• MMoE improves CVR AUC by 1-2% over hard sharing when CTR and CVR objectives conflict.
• Reduces negative transfer: CVR no longer hurt by CTR optimization.
• Scales to 10+ tasks without severe parameter overhead (only K additional gate networks).

Key Advantage: Each task learns its own preferred routing through the expert pool. CTR might heavily use certain experts, while CVR uses different experts, reducing task interference and improving both metrics.

ESMM: Entire Space Multitask Model

ESMM (Entire Space Multi-Task Model) is a specific MTL architecture designed to solve the CVR sample selection bias problem. Proposed for e-commerce and ad ranking, it models the entire conversion funnel in a unified framework.

The Problem ESMM Solves

Naively training a CVR model on only clicked items introduces sample selection bias. We want P(conversion | impression), but we only have labels for P(conversion | click).

Traditional approach:

CTR Model: predict P(click | impression)
CVR Model: predict P(conversion | click) using only clicked items
Issue: CVR training data excludes items user never clicked (but would convert)

ESMM's insight:

P(conversion | impression) = P(click | impression) × P(conversion | click, impression)

Therefore:
P(CTCVR) = P(CTR) × P(CVR)

Train:
- CTR model on all impressions (unbiased)
- pCVR indirectly via pCTCVR = pCTR × pCVR (no direct CVR label needed)

Mathematical Formulation

Let's denote:

• pCTR = P(click = 1 | impression)
• pCVR = P(conversion = 1 | click = 1, impression)
• pCTCVR = P(click = 1 AND conversion = 1 | impression) = pCTR × pCVR

The key relationship:

P(conversion | impression) = P(click | impression) × P(conversion | click, impression)
                             = pCTR × pCVR

This is valid under the assumption that conversion can only happen if user clicks (which is reasonable in e-commerce).

Why This Avoids Bias

Instead of directly learning from {clicked, converted} labels, ESMM uses:

• For CTR: All impressions with click labels (unbiased).
• For CVR: All impressions with pCTCVR computed from user's actual click-and-conversion behavior. Never uses a direct CVR label that would be biased toward high-quality clicked items.

Data Requirements

ESMM requires:

• All impressions: For CTR supervision and pCTCVR supervision.
• Click events: Users' actual clicks, used to compute clicks.
• Conversion events: Users' actual conversions (which can only happen after a click).

Note: We don't construct a direct "CVR label" from conversion data. The model learns pCVR implicitly from the pCTCVR loss.

ESMM Architecture & Training

Network Architecture

ESMM has two towers: a CTR tower and a CVR tower. Both share an embedding layer but have separate MLPs:

Shared Layer:
  x_emb = Embedding(x)

CTR Tower:
  h_ctr = MLP_CTR(x_emb) → pCTR ∈ (0, 1)

CVR Tower:
  h_cvr = MLP_CVR(x_emb) → pCVR ∈ (0, 1)

Combined:
  pCTCVR = pCTR × pCVR

Loss Function

ESMM uses two loss terms:

L_total = L_CTR + λ * L_CTCVR

where:
L_CTR = BCE(pCTR, y_click)       # trained on all impressions
L_CTCVR = BCE(pCTCVR, y_click∧y_conversion)  # trained on all impressions

Note: y_click∧y_conversion is 1 only if user both clicked AND converted,
      0 otherwise (including non-clicks and clicks without conversion)

The key insight: pCVR is never directly supervised. It's learned implicitly through the constraint that pCTCVR = pCTR × pCVR.

Training Procedure

1. Forward pass: Compute pCTR from CTR tower, pCVR from CVR tower, multiply to get pCTCVR.
2. Loss computation: Compute L_CTR using pCTR and all clicks. Compute L_CTCVR using pCTCVR and all click-and-conversion pairs.
3. Backward pass: Gradients flow through both towers. CVR tower gradients come entirely from the L_CTCVR term, which depends on pCTCVR = pCTR × pCVR.
4. Update: Update both towers' weights using the total loss gradient.

Why It Works

• Unbiased CTR: CTR is trained on all impressions, so it's unbiased.
• Entire-space CVR: pCVR is adjusted to make pCTCVR match the true click-and-conversion rate. This pCVR estimate is valid for the entire impression space, not just clicked items.
• Shared representations: The shared embedding captures common user-item interactions, benefiting both towers.
• No direct CVR label: We never feed a biased {clicked, converted} / {clicked, not converted} label to a CVR model.

Inference

At inference time, you have two choices:

• Use pCTCVR: The combined probability, directly comparable to GMV metrics.
• Use pCVR only for ranking: For users who clicked, use pCVR to predict conversion. But ESMM's pCVR is most accurate for P(conversion | impression), not P(conversion | click).

Most systems use pCTCVR for ranking, as it captures the full funnel in one number.

Advanced MTL Techniques

Beyond basic hard and soft sharing, several advanced techniques address task conflict and improve MTL training.

1. Uncertainty Weighting (Kendall et al.)

Instead of fixed task weights (λ_1, λ_2, ...), learn task-dependent weights during training:

For each task k, learn log variance σ_k (in log space for stability)
L_weighted = Σ_k (1 / (2σ_k²)) * L_k + log(σ_k)

Higher σ implies lower confidence in that task → lower weight on its loss.
The log(σ_k) term prevents σ from growing to infinity.

Benefit: Automatically balances losses based on task uncertainty. Tasks with noisier labels get lower weights. Simple and effective in practice.

2. GradNorm (Chen et al.)

Dynamically balance task loss scales so all tasks learn at similar rates:

For each task k:
  G_k = ||∇_θ L_k||      # gradient norm of task loss

Normalize:
  Ḡ_k = G_k / mean(G)

Compute loss balance scalars α_k so all Ḡ_k converge to same target rate.
Meta-loss: L_meta = Σ_k ||Ḡ_k - mean(Ḡ)||²
Gradient from L_meta scales the loss α_k dynamically.

Benefit: Ensures tasks train at balanced rates. Prevents one task's gradients from dominating. More sophisticated than uncertainty weighting but can be unstable.

3. Gradient Surgery (PCGrad, GradVac)

Explicitly prevent negative transfer by modifying gradients. When gradient directions conflict, remove the conflicting component:

PCGrad approach:
for each task k:
  for each previously processed task j:
    if grad_k · grad_j < 0:    # conflicting gradients
      grad_k ← grad_k - (dot product) * grad_j / ||grad_j||²

Effect: Projects grad_k onto the tangent space orthogonal to conflicting grad_j.

Benefit: Eliminates negative transfer by construction. Expensive (O(K²) gradient comparisons) but very effective for high-conflict tasks.

4. Progressive Training

Train tasks in sequence, starting with "easy" tasks that act as regularization for harder tasks:

1. Pre-train on main task (e.g., CTR) for N epochs
2. Unfreeze auxiliary task head, train MTL model for M epochs
3. Fine-tune on harder tasks (e.g., CVR) with higher loss weight

Benefit: Simpler tasks help learn good shared representations, then harder tasks specialize. Reduces training instability.

5. Task-Adaptive Feature Selection

Learn per-task feature masks or attention weights over the shared representation:

h_shared = Shared_Tower(x)
a_k = softmax(W_k · h_shared)     # task-specific attention over features
h_k = a_k ⊙ h_shared              # apply mask
o_k = Task_Head_k(h_k)

Benefit: Each task selects relevant features from shared representation. Soft version of specialization without separate towers.

6. Loss Scale Tuning

A critical but often overlooked detail: the scale of each loss matters more than the λ weights:

L_k is in range [0, 1] for sigmoid loss, [0, ∞) for MSE loss.
If L_k and L_j have different scales, even λ_k = λ_j gives unequal weighting.

Solution:
1. Normalize losses: L̃_k = L_k / mean(L_k)
2. Then apply λ weights: L_total = Σ λ_k * L̃_k

When to Use Which Approach

Choosing the right MTL architecture depends on your specific scenario. This section provides decision frameworks.

Decision Table: Architecture Selection

Diagnosis: Detecting Negative Transfer

How do you know if your tasks are conflicting? Several metrics help:

1. Compare to Separate Models

Train separate model for each task, compare metrics:
- If MTL_AUC_k < Single_AUC_k for some task k → negative transfer
- Magnitude of difference indicates severity

2. Task Affinity (Gradient Correlation)

Compute dot product of task gradients at shared layer:
affinity(k, j) = (grad_k · grad_j) / (||grad_k|| * ||grad_j||)

affinity ≈ 1    → aligned, tasks help each other
affinity ≈ 0    → independent
affinity < 0    → conflicting, negative transfer likely

3. Loss Evolution During Training

Plot individual task losses. If one loss increases while another decreases (and stays increased), you have conflict. With proper weighting, all losses should decrease (or plateau) together.

Guidelines for Task Addition

• Add a task if: It shares features with main task, has abundant data, or helps regularize (e.g., predicting auxiliary target that improves main task).
• Don't add a task if: It has conflicting objectives and small dataset, or requires completely different feature engineering.
• Test empirically: Train MTL model and compare main task AUC with single-task baseline. If worse, the auxiliary task hurts.

Ad Platform Case Study

For an ad platform with CTR, CVR, and satisfaction objectives:

1. Step 1: Train separate CTR, CVR, satisfaction models. Measure baseline AUCs.
2. Step 2: Try hard sharing. If CVR AUC drops > 0.1%, consider soft sharing.
3. Step 3: Implement ESMM for CVR to handle sample selection bias.
4. Step 4: Use task-specific gates or soft sharing for satisfaction vs CTR (satisfaction prioritizes dwell time; CTR prioritizes clicks).
5. Step 5: Apply uncertainty weighting to balance losses at different scales.
6. Step 6: Evaluate on offline metrics and run A/B test for online impact.

Python Implementation

This section provides implementations of three MTL models: hard sharing, MMoE, and ESMM. We'll use a synthetic ad dataset with CTR and CVR labels (with sample selection bias introduced intentionally).

Dataset Preparation

import numpy as np
import pandas as pd
from sklearn.metrics import roc_auc_score
from sklearn.preprocessing import MinMaxScaler

# Generate synthetic ad platform dataset
np.random.seed(42)
n_samples = 50000

# Features: user quality (0-1), item quality (0-1), context (0-1)
user_quality = np.random.uniform(0, 1, n_samples)
item_quality = np.random.uniform(0, 1, n_samples)
context = np.random.uniform(0, 1, n_samples)

# CTR depends on user quality + item quality + context
ctr_score = 0.3 * user_quality + 0.4 * item_quality + 0.2 * context + np.random.normal(0, 0.1, n_samples)
ctr_score = np.clip(ctr_score, 0, 1)
y_click = (np.random.uniform(0, 1, n_samples) < ctr_score).astype(int)

# CVR depends on item quality + context (NOT user quality strongly)
# Sample selection bias: CVR is only observed for clicked items
cvr_score = 0.5 * item_quality + 0.3 * context + np.random.normal(0, 0.1, n_samples)
cvr_score = np.clip(cvr_score, 0, 1)
y_conversion = np.zeros(n_samples)
y_conversion[y_click == 1] = (np.random.uniform(0, 1, np.sum(y_click)) < cvr_score[y_click == 1]).astype(int)

# Combine features
X = np.column_stack([user_quality, item_quality, context])

# CTR/CVR rates for reference
print(f"Overall CTR: {y_click.mean():.4f}")
print(f"Overall CVR (entire space): {y_conversion.mean():.4f}")
print(f"CVR given click: {y_conversion[y_click==1].mean():.4f}")

Model 1: Shared Bottom MTL (Hard Sharing)

import torch
import torch.nn as nn
import torch.optim as optim

class SharedBottomMTL(nn.Module):
    """Hard parameter sharing: shared trunk + task-specific heads"""
    def __init__(self, input_dim=3, hidden_dim=64):
        super().__init__()
        # Shared trunk
        self.shared = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim // 2),
            nn.ReLU()
        )
        # Task 1: CTR head
        self.ctr_head = nn.Sequential(
            nn.Linear(hidden_dim // 2, 32),
            nn.ReLU(),
            nn.Linear(32, 1),
            nn.Sigmoid()
        )
        # Task 2: CVR head
        self.cvr_head = nn.Sequential(
            nn.Linear(hidden_dim // 2, 32),
            nn.ReLU(),
            nn.Linear(32, 1),
            nn.Sigmoid()
        )

    def forward(self, x):
        shared_out = self.shared(x)
        pCTR = self.ctr_head(shared_out)
        pCVR = self.cvr_head(shared_out)
        return pCTR, pCVR

# Training
model = SharedBottomMTL(input_dim=3, hidden_dim=64)
optimizer = optim.Adam(model.parameters(), lr=0.001)
criterion = nn.BCELoss()

X_tensor = torch.FloatTensor(X)
y_click_tensor = torch.FloatTensor(y_click).reshape(-1, 1)
y_conversion_tensor = torch.FloatTensor(y_conversion).reshape(-1, 1)

losses_ctr, losses_cvr = [], []
for epoch in range(100):
    optimizer.zero_grad()
    pCTR, pCVR = model(X_tensor)

    loss_ctr = criterion(pCTR, y_click_tensor)
    loss_cvr = criterion(pCVR, y_conversion_tensor)
    loss_total = loss_ctr + 0.5 * loss_cvr

    loss_total.backward()
    optimizer.step()

    losses_ctr.append(loss_ctr.item())
    losses_cvr.append(loss_cvr.item())

# Evaluation
with torch.no_grad():
    pCTR, pCVR = model(X_tensor)
    auc_ctr = roc_auc_score(y_click, pCTR.numpy())
    auc_cvr = roc_auc_score(y_conversion, pCVR.numpy())
    print(f"SharedBottom - CTR AUC: {auc_ctr:.4f}, CVR AUC: {auc_cvr:.4f}")

Model 2: Mixture of Experts

class MMoEModel(nn.Module):
    """Multi-gate Mixture of Experts with task-specific gates"""
    def __init__(self, input_dim=3, num_experts=4, expert_dim=32, num_tasks=2):
        super().__init__()
        self.num_experts = num_experts
        self.num_tasks = num_tasks

        # Shared embedding
        self.embedding = nn.Linear(input_dim, 32)

        # Expert networks
        self.experts = nn.ModuleList([
            nn.Sequential(
                nn.Linear(32, expert_dim),
                nn.ReLU(),
                nn.Linear(expert_dim, expert_dim)
            ) for _ in range(num_experts)
        ])

        # Task-specific gates (one gate per task)
        self.gates = nn.ModuleList([
            nn.Sequential(
                nn.Linear(32, num_experts),
                nn.Softmax(dim=1)
            ) for _ in range(num_tasks)
        ])

        # Task-specific heads
        self.task_heads = nn.ModuleList([
            nn.Sequential(
                nn.Linear(expert_dim, 16),
                nn.ReLU(),
                nn.Linear(16, 1),
                nn.Sigmoid()
            ) for _ in range(num_tasks)
        ])

    def forward(self, x):
        # Shared embedding
        emb = torch.relu(self.embedding(x))

        # Expert outputs
        expert_outputs = [expert(emb) for expert in self.experts]
        expert_outputs = torch.stack(expert_outputs, dim=1)  # (batch, num_experts, expert_dim)

        # Task outputs with task-specific gates
        outputs = []
        for task_id in range(self.num_tasks):
            gate_out = self.gates[task_id](emb)  # (batch, num_experts)
            gate_out = gate_out.unsqueeze(2)  # (batch, num_experts, 1)

            # Weighted sum of expert outputs
            task_input = (expert_outputs * gate_out).sum(dim=1)  # (batch, expert_dim)
            task_output = self.task_heads[task_id](task_input)
            outputs.append(task_output)

        return tuple(outputs)

# Training MMoE
mmoe_model = MMoEModel(input_dim=3, num_experts=4, expert_dim=32, num_tasks=2)
optimizer_mmoe = optim.Adam(mmoe_model.parameters(), lr=0.001)

losses_mmoe_ctr, losses_mmoe_cvr = [], []
for epoch in range(100):
    optimizer_mmoe.zero_grad()
    pCTR_mmoe, pCVR_mmoe = mmoe_model(X_tensor)

    loss_ctr_mmoe = criterion(pCTR_mmoe, y_click_tensor)
    loss_cvr_mmoe = criterion(pCVR_mmoe, y_conversion_tensor)
    loss_total_mmoe = loss_ctr_mmoe + 0.5 * loss_cvr_mmoe

    loss_total_mmoe.backward()
    optimizer_mmoe.step()

    losses_mmoe_ctr.append(loss_ctr_mmoe.item())
    losses_mmoe_cvr.append(loss_cvr_mmoe.item())

# Evaluation
with torch.no_grad():
    pCTR_mmoe, pCVR_mmoe = mmoe_model(X_tensor)
    auc_ctr_mmoe = roc_auc_score(y_click, pCTR_mmoe.numpy())
    auc_cvr_mmoe = roc_auc_score(y_conversion, pCVR_mmoe.numpy())
    print(f"MMoE - CTR AUC: {auc_ctr_mmoe:.4f}, CVR AUC: {auc_cvr_mmoe:.4f}")

Model 3: ESMM (Entire Space Multitask Model)

class ESMMModel(nn.Module):
    """ESMM: CTR tower + CVR tower, pCTCVR = pCTR × pCVR"""
    def __init__(self, input_dim=3, hidden_dim=32):
        super().__init__()
        # Shared embedding
        self.embedding = nn.Sequential(
            nn.Linear(input_dim, 32),
            nn.ReLU()
        )

        # CTR tower
        self.ctr_tower = nn.Sequential(
            nn.Linear(32, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, 1),
            nn.Sigmoid()
        )

        # CVR tower
        self.cvr_tower = nn.Sequential(
            nn.Linear(32, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, 1),
            nn.Sigmoid()
        )

    def forward(self, x):
        emb = self.embedding(x)
        pCTR = self.ctr_tower(emb)
        pCVR = self.cvr_tower(emb)
        pCTCVR = pCTR * pCVR
        return pCTR, pCVR, pCTCVR

# Training ESMM
esmm_model = ESMMModel(input_dim=3, hidden_dim=32)
optimizer_esmm = optim.Adam(esmm_model.parameters(), lr=0.001)

y_ctcvr = (y_click * y_conversion).astype(float)  # 1 if clicked AND converted
y_ctcvr_tensor = torch.FloatTensor(y_ctcvr).reshape(-1, 1)

losses_esmm_ctr, losses_esmm_ctcvr = [], []
for epoch in range(100):
    optimizer_esmm.zero_grad()
    pCTR_esmm, pCVR_esmm, pCTCVR_esmm = esmm_model(X_tensor)

    loss_ctr_esmm = criterion(pCTR_esmm, y_click_tensor)
    loss_ctcvr_esmm = criterion(pCTCVR_esmm, y_ctcvr_tensor)
    loss_total_esmm = loss_ctr_esmm + 0.5 * loss_ctcvr_esmm

    loss_total_esmm.backward()
    optimizer_esmm.step()

    losses_esmm_ctr.append(loss_ctr_esmm.item())
    losses_esmm_ctcvr.append(loss_ctcvr_esmm.item())

# Evaluation
with torch.no_grad():
    pCTR_esmm, pCVR_esmm, pCTCVR_esmm = esmm_model(X_tensor)
    auc_ctr_esmm = roc_auc_score(y_click, pCTR_esmm.numpy())
    auc_cvr_esmm = roc_auc_score(y_conversion, pCVR_esmm.numpy())
    auc_ctcvr = roc_auc_score(y_ctcvr, pCTCVR_esmm.numpy())
    print(f"ESMM - CTR AUC: {auc_ctr_esmm:.4f}, CVR AUC: {auc_cvr_esmm:.4f}, CTCVR AUC: {auc_ctcvr:.4f}")

Comparison & Results

After training all three models, compare their performance:

import matplotlib.pyplot as plt

results = {
    'SharedBottom': (auc_ctr, auc_cvr),
    'MMoE': (auc_ctr_mmoe, auc_cvr_mmoe),
    'ESMM': (auc_ctr_esmm, auc_cvr_esmm)
}

print("\nModel Comparison:")
for model_name, (ctr_auc, cvr_auc) in results.items():
    print(f"{model_name:15} CTR AUC: {ctr_auc:.4f}  CVR AUC: {cvr_auc:.4f}")

# Plot loss curves
fig, axes = plt.subplots(1, 3, figsize=(15, 4))

axes[0].plot(losses_ctr, label='CTR Loss', alpha=0.7)
axes[0].plot(losses_cvr, label='CVR Loss', alpha=0.7)
axes[0].set_title('SharedBottom MTL')
axes[0].set_ylabel('Loss')
axes[0].set_xlabel('Epoch')
axes[0].legend()

axes[1].plot(losses_mmoe_ctr, label='CTR Loss', alpha=0.7)
axes[1].plot(losses_mmoe_cvr, label='CVR Loss', alpha=0.7)
axes[1].set_title('MMoE')
axes[1].set_ylabel('Loss')
axes[1].set_xlabel('Epoch')
axes[1].legend()

axes[2].plot(losses_esmm_ctr, label='CTR Loss', alpha=0.7)
axes[2].plot(losses_esmm_ctcvr, label='CTCVR Loss', alpha=0.7)
axes[2].set_title('ESMM')
axes[2].set_ylabel('Loss')
axes[2].set_xlabel('Epoch')
axes[2].legend()

plt.tight_layout()
plt.show()

Expected Observations

• SharedBottom: Good CTR AUC but may suffer CVR degradation due to task conflict (CTR prefers sensational items, CVR prefers quality items).
• MMoE: Better balance between CTR and CVR due to task-specific gates. Each task routes to preferred experts.
• ESMM: Best CVR AUC because it models the entire impression space, avoiding sample selection bias. pCVR is more accurate for ranking.