wehub-resource-sync
593b94c120
pytest / Unit Tests (push) Has been cancelled
pytest / Integration (integration_tests_a) (push) Has been cancelled
pytest / Integration (integration_tests_b) (push) Has been cancelled
pytest / Integration (integration_tests_c) (push) Has been cancelled
pytest / Integration (integration_tests_d) (push) Has been cancelled
pytest / Integration (integration_tests_e) (push) Has been cancelled
pytest / Integration (integration_tests_f) (push) Has been cancelled
pytest / Integration (integration_tests_g) (push) Has been cancelled
pytest / Integration (integration_tests_h) (push) Has been cancelled
pytest / Integration (integration_tests_i) (push) Has been cancelled
pytest / Integration (integration_tests_j) (push) Has been cancelled
pytest / Distributed (distributed_a) (push) Has been cancelled
pytest / Distributed (distributed_b) (push) Has been cancelled
pytest / Distributed (distributed_c) (push) Has been cancelled
pytest / Distributed (distributed_d) (push) Has been cancelled
pytest / Distributed (distributed_e) (push) Has been cancelled
pytest / Distributed (distributed_f) (push) Has been cancelled
pytest / Minimal Install (push) Has been cancelled
pytest / Event File (push) Has been cancelled
pytest (slow) / py-slow (push) Has been cancelled
Publish JSON Schema / publish-schema (push) Has been cancelled
460 lines
16 KiB
Plaintext
460 lines
16 KiB
Plaintext
{
|
||
"cells": [
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "7fb27b941602401d91542211134fc71a",
|
||
"metadata": {},
|
||
"source": [
|
||
"# Multi-Task Learning with Nash-MTL Loss Balancing\n",
|
||
"\n",
|
||
"[](https://colab.research.google.com/github/ludwig-ai/ludwig/blob/main/examples/multi_task/multi_task.ipynb)\n",
|
||
"\n",
|
||
"This notebook demonstrates **multi-task learning** with Ludwig: training a single model to predict multiple outputs simultaneously, and using **loss balancing** to prevent one task from dominating training.\n",
|
||
"\n",
|
||
"**Use case:** The [UCI Wine Quality dataset](https://archive.ics.uci.edu/ml/datasets/wine+quality) has a 0–10 quality score. We simultaneously predict:\n",
|
||
"- `quality_score` — the raw numerical score (regression)\n",
|
||
"- `quality_binary` — whether the wine is good (quality ≥ 7, binary classification)\n",
|
||
"\n",
|
||
"These two tasks have different loss magnitudes. Without balancing, regression loss dominates and the classifier under-trains.\n",
|
||
"\n",
|
||
"**Methods compared:**\n",
|
||
"\n",
|
||
"| Method | Status | Description |\n",
|
||
"|---|---|---|\n",
|
||
"| `none` | Available | Static weighted sum — baseline |\n",
|
||
"| `famo` | Available | Fast Adaptive Multitask Optimization (Liu et al., NeurIPS 2023) |\n",
|
||
"| `uncertainty` | Available | Homoscedastic uncertainty weighting (Kendall et al., CVPR 2018) |\n",
|
||
"| `nash_mtl` | Requires PR #4092 | Nash bargaining solution (Navon et al., ICML 2022) |"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "acae54e37e7d407bbb7b55eff062a284",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"!pip install ludwig --quiet"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "9a63283cbaf04dbcab1f6479b197f3a8",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"import logging\n",
|
||
"import shutil\n",
|
||
"import warnings\n",
|
||
"\n",
|
||
"import pandas as pd\n",
|
||
"\n",
|
||
"logging.basicConfig(level=logging.WARNING)"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "8dd0d8092fe74a7c96281538738b07e2",
|
||
"metadata": {},
|
||
"source": [
|
||
"## Dataset\n",
|
||
"\n",
|
||
"We use the [UCI Wine Quality dataset](https://archive.ics.uci.edu/ml/datasets/wine+quality) (red wine, ~1 600 rows, 11 physicochemical input features).\n",
|
||
"\n",
|
||
"Two output columns are derived from the original `quality` score:\n",
|
||
"- `quality_score`: the raw numerical score (0–10) — a regression target\n",
|
||
"- `quality_binary`: 1 if quality ≥ 7 (\"good\"), else 0 — a classification target\n",
|
||
"\n",
|
||
"These tasks share the same inputs and are naturally correlated, making them a good multi-task setup."
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "72eea5119410473aa328ad9291626812",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"WINE_URL = \"https://archive.ics.uci.edu/ml/machine-learning-databases/wine-quality/winequality-red.csv\"\n",
|
||
"\n",
|
||
"WINE_FEATURES = [\n",
|
||
" \"fixed_acidity\",\n",
|
||
" \"volatile_acidity\",\n",
|
||
" \"citric_acid\",\n",
|
||
" \"residual_sugar\",\n",
|
||
" \"chlorides\",\n",
|
||
" \"free_sulfur_dioxide\",\n",
|
||
" \"total_sulfur_dioxide\",\n",
|
||
" \"density\",\n",
|
||
" \"pH\",\n",
|
||
" \"sulphates\",\n",
|
||
" \"alcohol\",\n",
|
||
"]\n",
|
||
"\n",
|
||
"print(\"Downloading wine quality dataset...\")\n",
|
||
"df = pd.read_csv(WINE_URL, sep=\";\")\n",
|
||
"df.columns = [c.replace(\" \", \"_\") for c in df.columns]\n",
|
||
"\n",
|
||
"df[\"quality_score\"] = df[\"quality\"].astype(float)\n",
|
||
"df[\"quality_binary\"] = (df[\"quality\"] >= 7).astype(int)\n",
|
||
"df = df.drop(columns=[\"quality\"])\n",
|
||
"\n",
|
||
"print(f\"Shape: {df.shape}\")\n",
|
||
"print(f\"Good wines (quality >= 7): {df['quality_binary'].mean():.1%}\")\n",
|
||
"print(f\"Quality score range: {df['quality_score'].min():.0f} – {df['quality_score'].max():.0f}\")\n",
|
||
"df.head()"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "8edb47106e1a46a883d545849b8ab81b",
|
||
"metadata": {},
|
||
"source": [
|
||
"### Why does loss balancing matter here?\n",
|
||
"\n",
|
||
"MSE for regression and binary cross-entropy for classification operate on different scales. Typically:\n",
|
||
"- Regression MSE: values in the range 0.5–2.0 (for a 0–10 quality score)\n",
|
||
"- Binary cross-entropy: values in the range 0.3–0.7\n",
|
||
"\n",
|
||
"Without balancing, the regression task provides larger raw gradients and the model updates are biased toward minimising the regression loss at the expense of classification performance."
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "10185d26023b46108eb7d9f57d49d2b3",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"from ludwig.api import LudwigModel\n",
|
||
"\n",
|
||
"\n",
|
||
"def input_features():\n",
|
||
" return [{\"name\": feat, \"type\": \"number\", \"preprocessing\": {\"normalization\": \"zscore\"}} for feat in WINE_FEATURES]\n",
|
||
"\n",
|
||
"\n",
|
||
"def make_config(loss_balancing: str) -> dict:\n",
|
||
" return {\n",
|
||
" \"model_type\": \"ecd\",\n",
|
||
" \"input_features\": input_features(),\n",
|
||
" \"output_features\": [\n",
|
||
" {\"name\": \"quality_score\", \"type\": \"number\"},\n",
|
||
" {\"name\": \"quality_binary\", \"type\": \"binary\"},\n",
|
||
" ],\n",
|
||
" \"combiner\": {\n",
|
||
" \"type\": \"concat\",\n",
|
||
" \"num_fc_layers\": 2,\n",
|
||
" \"output_size\": 128,\n",
|
||
" \"dropout\": 0.1,\n",
|
||
" },\n",
|
||
" \"trainer\": {\n",
|
||
" \"epochs\": 30,\n",
|
||
" \"learning_rate\": 0.001,\n",
|
||
" \"batch_size\": 128,\n",
|
||
" \"loss_balancing\": loss_balancing,\n",
|
||
" },\n",
|
||
" }\n",
|
||
"\n",
|
||
"\n",
|
||
"def train_and_evaluate(name: str, loss_balancing: str) -> dict | None:\n",
|
||
" result_dir = f\"./results/{name}\"\n",
|
||
" shutil.rmtree(result_dir, ignore_errors=True)\n",
|
||
" print(f\"\\n{'=' * 50}\")\n",
|
||
" print(f\"Training: {name} (loss_balancing={loss_balancing})\")\n",
|
||
" print(f\"{'=' * 50}\")\n",
|
||
" config = make_config(loss_balancing)\n",
|
||
" model = LudwigModel(config=config, logging_level=logging.WARNING)\n",
|
||
" train_stats, _, _ = model.train(\n",
|
||
" dataset=df,\n",
|
||
" experiment_name=\"multi_task\",\n",
|
||
" model_name=name,\n",
|
||
" output_directory=result_dir,\n",
|
||
" )\n",
|
||
" # Extract final validation metrics\n",
|
||
" vset = train_stats.get(\"validation\", {})\n",
|
||
" score_mae = _last(vset.get(\"quality_score\", {}).get(\"mean_absolute_error\", []))\n",
|
||
" binary_auc = _last(vset.get(\"quality_binary\", {}).get(\"roc_auc\", []))\n",
|
||
" print(f\" quality_score MAE : {score_mae:.4f}\")\n",
|
||
" print(f\" quality_binary ROC-AUC : {binary_auc:.4f}\")\n",
|
||
" return {\"method\": name, \"score_mae\": score_mae, \"binary_roc_auc\": binary_auc}\n",
|
||
"\n",
|
||
"\n",
|
||
"def _last(series):\n",
|
||
" if not series:\n",
|
||
" return float(\"nan\")\n",
|
||
" v = series[-1]\n",
|
||
" if isinstance(v, (list, tuple)):\n",
|
||
" v = v[-1]\n",
|
||
" return float(v)\n",
|
||
"\n",
|
||
"\n",
|
||
"results = []"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "8763a12b2bbd4a93a75aff182afb95dc",
|
||
"metadata": {},
|
||
"source": [
|
||
"## Baseline: No Loss Balancing\n",
|
||
"\n",
|
||
"The default `loss_balancing: none` computes the total loss as a static weighted sum:\n",
|
||
"\n",
|
||
"```\n",
|
||
"L_total = w_score * L_score + w_binary * L_binary\n",
|
||
"```\n",
|
||
"\n",
|
||
"where weights are set from the config (defaulting to 1.0 each). The regression task typically has a larger loss value, so it tends to dominate gradient updates.\n",
|
||
"\n",
|
||
"Config snippet:\n",
|
||
"```yaml\n",
|
||
"trainer:\n",
|
||
" loss_balancing: none\n",
|
||
"```"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "7623eae2785240b9bd12b16a66d81610",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"result = train_and_evaluate(\"none\", \"none\")\n",
|
||
"results.append(result)"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "7cdc8c89c7104fffa095e18ddfef8986",
|
||
"metadata": {},
|
||
"source": [
|
||
"## FAMO Balancing\n",
|
||
"\n",
|
||
"> **Available now** in the current Ludwig release.\n",
|
||
"\n",
|
||
"**FAMO** (Fast Adaptive Multitask Optimization, Liu et al., NeurIPS 2023) maintains an exponential moving average of each task's loss and updates task weights at every step to equalise the rate of loss decrease across tasks.\n",
|
||
"\n",
|
||
"Key properties:\n",
|
||
"- No gradient computation overhead (unlike gradient-based methods)\n",
|
||
"- Converges quickly; good default choice for most multi-task problems\n",
|
||
"- One hyperparameter: `loss_balancing_lr` (EMA learning rate, default 0.025)\n",
|
||
"\n",
|
||
"Config snippet:\n",
|
||
"```yaml\n",
|
||
"trainer:\n",
|
||
" loss_balancing: famo\n",
|
||
"```"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "b118ea5561624da68c537baed56e602f",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"result = train_and_evaluate(\"famo\", \"famo\")\n",
|
||
"results.append(result)"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "938c804e27f84196a10c8828c723f798",
|
||
"metadata": {},
|
||
"source": [
|
||
"## Uncertainty Weighting\n",
|
||
"\n",
|
||
"> **Available now** in the current Ludwig release.\n",
|
||
"\n",
|
||
"**Uncertainty weighting** (Kendall et al., CVPR 2018) treats each task's weight as a learned parameter representing homoscedastic (task-level) uncertainty. Tasks with higher uncertainty receive lower weight automatically.\n",
|
||
"\n",
|
||
"Key properties:\n",
|
||
"- Learns a single scalar per task — minimal parameter overhead\n",
|
||
"- Principled Bayesian interpretation\n",
|
||
"- Works best when tasks have stable, intrinsically different noise levels\n",
|
||
"\n",
|
||
"Config snippet:\n",
|
||
"```yaml\n",
|
||
"trainer:\n",
|
||
" loss_balancing: uncertainty\n",
|
||
"```"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "504fb2a444614c0babb325280ed9130a",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"result = train_and_evaluate(\"uncertainty\", \"uncertainty\")\n",
|
||
"results.append(result)"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "59bbdb311c014d738909a11f9e486628",
|
||
"metadata": {},
|
||
"source": [
|
||
"## Nash-MTL Balancing\n",
|
||
"\n",
|
||
"> **Requires PR #4092** (`future-capabilities` branch). Not yet available in the main Ludwig release.\n",
|
||
"> \n",
|
||
"> To try it now:\n",
|
||
"> ```bash\n",
|
||
"> pip install git+https://github.com/ludwig-ai/ludwig@future-capabilities\n",
|
||
"> ```\n",
|
||
"\n",
|
||
"**Nash-MTL** (Navon et al., ICML 2022) finds the [Nash bargaining solution](https://en.wikipedia.org/wiki/Nash_bargaining_solution) for task weight allocation. In game-theoretic terms, it selects a Pareto-optimal point where no task can improve its loss without worsening another task's loss.\n",
|
||
"\n",
|
||
"Concretely, Nash-MTL solves a quadratic programme at each step to find weights `alpha` such that:\n",
|
||
"\n",
|
||
"```\n",
|
||
"alpha = argmax sum_i log(alpha_i * loss_i) subject to sum alpha_i = 1, alpha_i >= 0\n",
|
||
"```\n",
|
||
"\n",
|
||
"This is equivalent to weights being inversely proportional to task losses — tasks with larger losses receive proportionally higher weight, equalising marginal contributions.\n",
|
||
"\n",
|
||
"Key properties:\n",
|
||
"- Theoretically principled (game theory, not heuristic)\n",
|
||
"- Most effective when tasks genuinely conflict (gradients point in different directions)\n",
|
||
"- Slightly higher per-step overhead than FAMO (solves a small QP per step)\n",
|
||
"\n",
|
||
"Config snippet:\n",
|
||
"```yaml\n",
|
||
"trainer:\n",
|
||
" loss_balancing: nash_mtl # Requires PR #4092 / Ludwig >= 0.14\n",
|
||
"```\n",
|
||
"\n",
|
||
"Reference: Navon et al., [Multi-Task Learning as a Bargaining Game](https://arxiv.org/abs/2202.01017), ICML 2022."
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "b43b363d81ae4b689946ece5c682cd59",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"# Nash-MTL is attempted; a warning is printed if PR #4092 is not yet available.\n",
|
||
"try:\n",
|
||
" result = train_and_evaluate(\"nash_mtl\", \"nash_mtl\")\n",
|
||
" results.append(result)\n",
|
||
"except Exception as exc:\n",
|
||
" warnings.warn(\n",
|
||
" f\"nash_mtl is not available in this Ludwig version ({exc}). \"\n",
|
||
" \"Install from 'future-capabilities' branch to enable it.\"\n",
|
||
" )\n",
|
||
" results.append({\"method\": \"nash_mtl\", \"score_mae\": float(\"nan\"), \"binary_roc_auc\": float(\"nan\")})"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "8a65eabff63a45729fe45fb5ade58bdc",
|
||
"metadata": {},
|
||
"source": [
|
||
"## Comparison Table\n",
|
||
"\n",
|
||
"Side-by-side comparison of all four methods on the validation split."
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "c3933fab20d04ec698c2621248eb3be0",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"import pandas as pd\n",
|
||
"\n",
|
||
"comparison = pd.DataFrame(results).set_index(\"method\")\n",
|
||
"comparison.columns = [\"Score MAE (lower better)\", \"Binary ROC-AUC (higher better)\"]\n",
|
||
"print(comparison.to_string())\n",
|
||
"comparison"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "code",
|
||
"execution_count": null,
|
||
"id": "4dd4641cc4064e0191573fe9c69df29b",
|
||
"metadata": {},
|
||
"outputs": [],
|
||
"source": [
|
||
"import matplotlib.pyplot as plt\n",
|
||
"\n",
|
||
"fig, axes = plt.subplots(1, 2, figsize=(12, 4))\n",
|
||
"\n",
|
||
"methods = comparison.index.tolist()\n",
|
||
"mae_vals = comparison[\"Score MAE (lower better)\"].tolist()\n",
|
||
"auc_vals = comparison[\"Binary ROC-AUC (higher better)\"].tolist()\n",
|
||
"\n",
|
||
"colors = [\"tab:gray\", \"tab:blue\", \"tab:orange\", \"tab:green\"]\n",
|
||
"\n",
|
||
"axes[0].bar(methods, mae_vals, color=colors[: len(methods)])\n",
|
||
"axes[0].set_title(\"Quality Score — MAE (lower is better)\")\n",
|
||
"axes[0].set_ylabel(\"MAE\")\n",
|
||
"axes[0].set_ylim(0, max(mae_vals) * 1.2)\n",
|
||
"\n",
|
||
"axes[1].bar(methods, auc_vals, color=colors[: len(methods)])\n",
|
||
"axes[1].set_title(\"Quality Binary — ROC-AUC (higher is better)\")\n",
|
||
"axes[1].set_ylabel(\"ROC-AUC\")\n",
|
||
"axes[1].set_ylim(0.5, 1.0)\n",
|
||
"\n",
|
||
"plt.tight_layout()\n",
|
||
"plt.savefig(\"results/comparison.png\", dpi=150)\n",
|
||
"plt.show()"
|
||
]
|
||
},
|
||
{
|
||
"cell_type": "markdown",
|
||
"id": "8309879909854d7188b41380fd92a7c3",
|
||
"metadata": {},
|
||
"source": [
|
||
"## When to Use Nash-MTL\n",
|
||
"\n",
|
||
"Use **Nash-MTL** when:\n",
|
||
"\n",
|
||
"- You have two or more output features with significantly different loss scales\n",
|
||
"- Your tasks conflict (improving one task hurts another), as measured by negative gradient cosine similarity\n",
|
||
"- You want the most principled, theoretically grounded balancing approach\n",
|
||
"- You can tolerate a small additional per-step compute cost (solving a small QP)\n",
|
||
"\n",
|
||
"Use **FAMO** when:\n",
|
||
"- You want a strong default with no tuning required\n",
|
||
"- Training speed is a priority (FAMO has lower overhead than Nash-MTL)\n",
|
||
"\n",
|
||
"Use **uncertainty weighting** when:\n",
|
||
"- Tasks have stable, intrinsically different noise levels\n",
|
||
"- You want a simple, interpretable approach with a Bayesian justification\n",
|
||
"\n",
|
||
"Use **none** (baseline) when:\n",
|
||
"- You have a single output feature\n",
|
||
"- Your tasks have similar loss scales and do not conflict\n",
|
||
"- You have manually tuned `loss_weight` values on each output feature\n",
|
||
"\n",
|
||
"### Further Reading\n",
|
||
"\n",
|
||
"- Navon et al., [Multi-Task Learning as a Bargaining Game](https://arxiv.org/abs/2202.01017), ICML 2022\n",
|
||
"- Liu et al., [FAMO: Fast Adaptive Multitask Optimization](https://arxiv.org/abs/2306.03792), NeurIPS 2023\n",
|
||
"- Kendall et al., [Multi-Task Learning Using Uncertainty to Weigh Losses](https://arxiv.org/abs/1705.07115), CVPR 2018\n",
|
||
"- Chen et al., [GradNorm: Gradient Normalization for Adaptive Loss Balancing](https://arxiv.org/abs/1711.02257), ICML 2018"
|
||
]
|
||
}
|
||
],
|
||
"metadata": {
|
||
"kernelspec": {
|
||
"display_name": "Python 3",
|
||
"language": "python",
|
||
"name": "python3"
|
||
},
|
||
"language_info": {
|
||
"name": "python",
|
||
"version": "3.12.0"
|
||
}
|
||
},
|
||
"nbformat": 4,
|
||
"nbformat_minor": 5
|
||
}
|