Files
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
chore: import upstream snapshot with attribution
2026-07-13 12:49:20 +08:00

584 lines
19 KiB
Plaintext

{
"cells": [
{
"cell_type": "markdown",
"id": "title-cell",
"metadata": {},
"source": [
"# Optimizer Comparison: Schedule-Free, Muon, Adafactor, and More\n",
"\n",
"[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/ludwig-ai/ludwig/blob/main/examples/optimizers/optimizer_comparison.ipynb)\n",
"\n",
"Optimizer choice is one of the most impactful but often overlooked hyperparameters in deep learning.\n",
"This notebook benchmarks five optimizers available in Ludwig on the UCI Wine Quality dataset:\n",
"\n",
"| Optimizer | Key idea |\n",
"|---|---|\n",
"| **AdamW** (baseline) | Adam with decoupled weight decay — the standard choice |\n",
"| **RAdam** | Rectified Adam — warm-up-free, stable at the start of training |\n",
"| **Adafactor** | No second-moment buffer — large models, small memory footprint |\n",
"| **Schedule-Free AdamW** | No LR scheduler needed — the optimizer is the schedule |\n",
"| **Muon** | Newton-Schulz orthogonalisation of weight-matrix gradients |\n",
"\n",
"The task is binary classification: predict whether a wine has quality >= 7.\n",
"Everything runs on CPU and completes in a few minutes."
]
},
{
"cell_type": "markdown",
"id": "install-header",
"metadata": {},
"source": [
"## Setup"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "install-cell",
"metadata": {},
"outputs": [],
"source": [
"!pip install ludwig scikit-learn --quiet"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "imports-cell",
"metadata": {},
"outputs": [],
"source": [
"import tempfile\n",
"import time\n",
"import warnings\n",
"\n",
"warnings.filterwarnings(\"ignore\")\n",
"\n",
"import matplotlib.pyplot as plt\n",
"import matplotlib.ticker as ticker\n",
"import pandas as pd\n",
"\n",
"from ludwig.api import LudwigModel"
]
},
{
"cell_type": "markdown",
"id": "dataset-header",
"metadata": {},
"source": [
"## Dataset\n",
"\n",
"We use the **UCI Wine Quality** dataset (red wines). \n",
"The target is binarised: `quality >= 7` → True, otherwise False."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "dataset-cell",
"metadata": {},
"outputs": [],
"source": [
"DATA_URL = \"https://archive.ics.uci.edu/ml/machine-learning-databases/wine-quality/winequality-red.csv\"\n",
"\n",
"df = pd.read_csv(DATA_URL, sep=\";\")\n",
"df.columns = [c.strip().replace(\" \", \"_\") for c in df.columns]\n",
"df[\"quality\"] = (df[\"quality\"] >= 7).astype(str) # binary: \"True\" / \"False\"\n",
"\n",
"print(f\"Shape: {df.shape}\")\n",
"print(f\"Class distribution:\\n{df['quality'].value_counts()}\")\n",
"df.head()"
]
},
{
"cell_type": "markdown",
"id": "shared-config-header",
"metadata": {},
"source": [
"## Shared configuration\n",
"\n",
"All models share the same architecture and 30 training epochs.\n",
"Only the `optimizer` block (and `learning_rate_scheduler`) differs."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "shared-config-cell",
"metadata": {},
"outputs": [],
"source": [
"FEATURE_NAMES = [\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",
"INPUT_FEATURES = [{\"name\": n, \"type\": \"number\"} for n in FEATURE_NAMES]\n",
"OUTPUT_FEATURES = [{\"name\": \"quality\", \"type\": \"binary\"}]\n",
"\n",
"BASE_TRAINER = {\n",
" \"epochs\": 30,\n",
"}\n",
"\n",
"# Storage for training curves and summary metrics\n",
"all_results = {}\n",
"summary = []"
]
},
{
"cell_type": "markdown",
"id": "adamw-header",
"metadata": {},
"source": [
"## Optimizer 1 — AdamW (baseline)\n",
"\n",
"AdamW adds **decoupled weight decay** to Adam. \n",
"It is the default optimizer in Ludwig and a solid baseline for most tasks.\n",
"\n",
"```yaml\n",
"trainer:\n",
" epochs: 30\n",
" optimizer:\n",
" type: adamw\n",
" lr: 0.001\n",
" learning_rate_scheduler:\n",
" type: cosine\n",
"```"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "adamw-cell",
"metadata": {},
"outputs": [],
"source": [
"config_adamw = {\n",
" \"model_type\": \"ecd\",\n",
" \"input_features\": INPUT_FEATURES,\n",
" \"output_features\": OUTPUT_FEATURES,\n",
" \"trainer\": {\n",
" **BASE_TRAINER,\n",
" \"optimizer\": {\"type\": \"adamw\", \"lr\": 0.001},\n",
" \"learning_rate_scheduler\": {\"type\": \"cosine\"},\n",
" },\n",
"}\n",
"\n",
"with tempfile.TemporaryDirectory() as tmpdir:\n",
" model = LudwigModel(config_adamw, logging_level=30)\n",
" t0 = time.time()\n",
" train_stats, _, _ = model.train(\n",
" dataset=df,\n",
" output_directory=tmpdir,\n",
" skip_save_model=True,\n",
" skip_save_progress=True,\n",
" skip_save_log=True,\n",
" )\n",
" elapsed_adamw = time.time() - t0\n",
"\n",
"all_results[\"adamw\"] = train_stats.validation\n",
"summary.append(\n",
" {\n",
" \"optimizer\": \"adamw\",\n",
" \"final_val_loss\": train_stats.validation[\"quality\"][\"loss\"][-1],\n",
" \"final_val_accuracy\": train_stats.validation[\"quality\"][\"accuracy\"][-1],\n",
" \"training_time_s\": round(elapsed_adamw, 1),\n",
" }\n",
")\n",
"print(f\"AdamW done in {elapsed_adamw:.1f}s\")"
]
},
{
"cell_type": "markdown",
"id": "radam-header",
"metadata": {},
"source": [
"## Optimizer 2 — RAdam\n",
"\n",
"**Rectified Adam** (Liu et al., 2019) automatically warms up the adaptive learning rate\n",
"by rectifying the variance of the second moment. \n",
"Unlike AdamW, it does not need a manual warm-up schedule.\n",
"\n",
"```yaml\n",
"trainer:\n",
" epochs: 30\n",
" optimizer:\n",
" type: radam\n",
" lr: 0.001\n",
" learning_rate_scheduler:\n",
" type: cosine\n",
"```"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "radam-cell",
"metadata": {},
"outputs": [],
"source": [
"config_radam = {\n",
" \"model_type\": \"ecd\",\n",
" \"input_features\": INPUT_FEATURES,\n",
" \"output_features\": OUTPUT_FEATURES,\n",
" \"trainer\": {\n",
" **BASE_TRAINER,\n",
" \"optimizer\": {\"type\": \"radam\", \"lr\": 0.001},\n",
" \"learning_rate_scheduler\": {\"type\": \"cosine\"},\n",
" },\n",
"}\n",
"\n",
"with tempfile.TemporaryDirectory() as tmpdir:\n",
" model = LudwigModel(config_radam, logging_level=30)\n",
" t0 = time.time()\n",
" train_stats, _, _ = model.train(\n",
" dataset=df,\n",
" output_directory=tmpdir,\n",
" skip_save_model=True,\n",
" skip_save_progress=True,\n",
" skip_save_log=True,\n",
" )\n",
" elapsed_radam = time.time() - t0\n",
"\n",
"all_results[\"radam\"] = train_stats.validation\n",
"summary.append(\n",
" {\n",
" \"optimizer\": \"radam\",\n",
" \"final_val_loss\": train_stats.validation[\"quality\"][\"loss\"][-1],\n",
" \"final_val_accuracy\": train_stats.validation[\"quality\"][\"accuracy\"][-1],\n",
" \"training_time_s\": round(elapsed_radam, 1),\n",
" }\n",
")\n",
"print(f\"RAdam done in {elapsed_radam:.1f}s\")"
]
},
{
"cell_type": "markdown",
"id": "adafactor-header",
"metadata": {},
"source": [
"## Optimizer 3 — Adafactor\n",
"\n",
"**Adafactor** (Shazeer & Stern, 2018) avoids storing a full second-moment matrix\n",
"by using a factored approximation. \n",
"This makes it memory-efficient for large models (e.g., T5 was trained with Adafactor).\n",
"It includes internal step-size scaling, so an external LR scheduler is optional.\n",
"\n",
"```yaml\n",
"trainer:\n",
" epochs: 30\n",
" optimizer:\n",
" type: adafactor\n",
" lr: 0.001\n",
" # no learning_rate_scheduler needed\n",
"```"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "adafactor-cell",
"metadata": {},
"outputs": [],
"source": [
"config_adafactor = {\n",
" \"model_type\": \"ecd\",\n",
" \"input_features\": INPUT_FEATURES,\n",
" \"output_features\": OUTPUT_FEATURES,\n",
" \"trainer\": {\n",
" **BASE_TRAINER,\n",
" \"optimizer\": {\"type\": \"adafactor\", \"lr\": 0.001},\n",
" # No learning_rate_scheduler — Adafactor handles step-size internally\n",
" },\n",
"}\n",
"\n",
"with tempfile.TemporaryDirectory() as tmpdir:\n",
" model = LudwigModel(config_adafactor, logging_level=30)\n",
" t0 = time.time()\n",
" train_stats, _, _ = model.train(\n",
" dataset=df,\n",
" output_directory=tmpdir,\n",
" skip_save_model=True,\n",
" skip_save_progress=True,\n",
" skip_save_log=True,\n",
" )\n",
" elapsed_adafactor = time.time() - t0\n",
"\n",
"all_results[\"adafactor\"] = train_stats.validation\n",
"summary.append(\n",
" {\n",
" \"optimizer\": \"adafactor\",\n",
" \"final_val_loss\": train_stats.validation[\"quality\"][\"loss\"][-1],\n",
" \"final_val_accuracy\": train_stats.validation[\"quality\"][\"accuracy\"][-1],\n",
" \"training_time_s\": round(elapsed_adafactor, 1),\n",
" }\n",
")\n",
"print(f\"Adafactor done in {elapsed_adafactor:.1f}s\")"
]
},
{
"cell_type": "markdown",
"id": "sfa-header",
"metadata": {},
"source": [
"## Optimizer 4 — Schedule-Free AdamW\n",
"\n",
"**Schedule-Free AdamW** (Defazio et al., 2024) eliminates the need for a separate\n",
"learning-rate schedule by building the schedule into the optimizer itself via\n",
"a momentum-based interpolation between two internal sequences.\n",
"\n",
"> **Important:** do _not_ add a `learning_rate_scheduler` block when using this optimizer.\n",
"> Adding one would counteract the built-in schedule and harm convergence.\n",
"\n",
"```yaml\n",
"trainer:\n",
" epochs: 30\n",
" optimizer:\n",
" type: schedule_free_adamw\n",
" lr: 0.001\n",
" # No learning_rate_scheduler — this is the whole point!\n",
"```"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "sfa-cell",
"metadata": {},
"outputs": [],
"source": [
"config_sfa = {\n",
" \"model_type\": \"ecd\",\n",
" \"input_features\": INPUT_FEATURES,\n",
" \"output_features\": OUTPUT_FEATURES,\n",
" \"trainer\": {\n",
" **BASE_TRAINER,\n",
" \"optimizer\": {\"type\": \"schedule_free_adamw\", \"lr\": 0.001},\n",
" # No learning_rate_scheduler — Schedule-Free AdamW handles it internally\n",
" },\n",
"}\n",
"\n",
"with tempfile.TemporaryDirectory() as tmpdir:\n",
" model = LudwigModel(config_sfa, logging_level=30)\n",
" t0 = time.time()\n",
" train_stats, _, _ = model.train(\n",
" dataset=df,\n",
" output_directory=tmpdir,\n",
" skip_save_model=True,\n",
" skip_save_progress=True,\n",
" skip_save_log=True,\n",
" )\n",
" elapsed_sfa = time.time() - t0\n",
"\n",
"all_results[\"schedule_free_adamw\"] = train_stats.validation\n",
"summary.append(\n",
" {\n",
" \"optimizer\": \"schedule_free_adamw\",\n",
" \"final_val_loss\": train_stats.validation[\"quality\"][\"loss\"][-1],\n",
" \"final_val_accuracy\": train_stats.validation[\"quality\"][\"accuracy\"][-1],\n",
" \"training_time_s\": round(elapsed_sfa, 1),\n",
" }\n",
")\n",
"print(f\"Schedule-Free AdamW done in {elapsed_sfa:.1f}s\")"
]
},
{
"cell_type": "markdown",
"id": "muon-header",
"metadata": {},
"source": [
"## Optimizer 5 — Muon\n",
"\n",
"**Muon** (Kosson et al.) applies Nesterov momentum updates orthogonalised via\n",
"Newton-Schulz iterations to weight matrices, while using AdamW for embeddings\n",
"and other non-matrix parameters. \n",
"The orthogonalisation acts as a preconditioner with minimal overhead.\n",
"\n",
"```yaml\n",
"trainer:\n",
" epochs: 30\n",
" optimizer:\n",
" type: muon\n",
" lr: 0.001\n",
" learning_rate_scheduler:\n",
" type: cosine\n",
"```"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "muon-cell",
"metadata": {},
"outputs": [],
"source": [
"config_muon = {\n",
" \"model_type\": \"ecd\",\n",
" \"input_features\": INPUT_FEATURES,\n",
" \"output_features\": OUTPUT_FEATURES,\n",
" \"trainer\": {\n",
" **BASE_TRAINER,\n",
" \"optimizer\": {\"type\": \"muon\", \"lr\": 0.001},\n",
" \"learning_rate_scheduler\": {\"type\": \"cosine\"},\n",
" },\n",
"}\n",
"\n",
"with tempfile.TemporaryDirectory() as tmpdir:\n",
" model = LudwigModel(config_muon, logging_level=30)\n",
" t0 = time.time()\n",
" train_stats, _, _ = model.train(\n",
" dataset=df,\n",
" output_directory=tmpdir,\n",
" skip_save_model=True,\n",
" skip_save_progress=True,\n",
" skip_save_log=True,\n",
" )\n",
" elapsed_muon = time.time() - t0\n",
"\n",
"all_results[\"muon\"] = train_stats.validation\n",
"summary.append(\n",
" {\n",
" \"optimizer\": \"muon\",\n",
" \"final_val_loss\": train_stats.validation[\"quality\"][\"loss\"][-1],\n",
" \"final_val_accuracy\": train_stats.validation[\"quality\"][\"accuracy\"][-1],\n",
" \"training_time_s\": round(elapsed_muon, 1),\n",
" }\n",
")\n",
"print(f\"Muon done in {elapsed_muon:.1f}s\")"
]
},
{
"cell_type": "markdown",
"id": "results-header",
"metadata": {},
"source": [
"## Results"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "results-table-cell",
"metadata": {},
"outputs": [],
"source": [
"results_df = pd.DataFrame(summary)\n",
"results_df = results_df.set_index(\"optimizer\")\n",
"results_df[\"final_val_loss\"] = results_df[\"final_val_loss\"].round(4)\n",
"results_df[\"final_val_accuracy\"] = results_df[\"final_val_accuracy\"].round(4)\n",
"print(results_df.to_string())"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "results-plot-cell",
"metadata": {},
"outputs": [],
"source": [
"COLORS = {\n",
" \"adamw\": \"#1f77b4\",\n",
" \"radam\": \"#ff7f0e\",\n",
" \"adafactor\": \"#2ca02c\",\n",
" \"schedule_free_adamw\": \"#d62728\",\n",
" \"muon\": \"#9467bd\",\n",
"}\n",
"\n",
"fig, axes = plt.subplots(1, 2, figsize=(13, 4))\n",
"fig.suptitle(\"Optimizer Comparison — Wine Quality (30 epochs, CPU)\", fontsize=13)\n",
"\n",
"for opt_name, val_stats in all_results.items():\n",
" losses = val_stats[\"quality\"][\"loss\"]\n",
" accs = val_stats[\"quality\"][\"accuracy\"]\n",
" epochs = range(1, len(losses) + 1)\n",
" color = COLORS.get(opt_name)\n",
" axes[0].plot(epochs, losses, label=opt_name, color=color)\n",
" axes[1].plot(epochs, accs, label=opt_name, color=color)\n",
"\n",
"axes[0].set_title(\"Validation Loss\")\n",
"axes[0].set_xlabel(\"Epoch\")\n",
"axes[0].set_ylabel(\"Loss\")\n",
"axes[0].legend(fontsize=8)\n",
"axes[0].xaxis.set_major_locator(ticker.MaxNLocator(integer=True))\n",
"\n",
"axes[1].set_title(\"Validation Accuracy\")\n",
"axes[1].set_xlabel(\"Epoch\")\n",
"axes[1].set_ylabel(\"Accuracy\")\n",
"axes[1].legend(fontsize=8)\n",
"axes[1].xaxis.set_major_locator(ticker.MaxNLocator(integer=True))\n",
"\n",
"plt.tight_layout()\n",
"plt.savefig(\"optimizer_comparison.png\", dpi=120, bbox_inches=\"tight\")\n",
"plt.show()\n",
"print(\"Figure saved to optimizer_comparison.png\")"
]
},
{
"cell_type": "markdown",
"id": "when-to-use-header",
"metadata": {},
"source": [
"## When to use each optimizer\n",
"\n",
"| Optimizer | Best for | Watch out for |\n",
"|---|---|---|\n",
"| **AdamW** | General-purpose default; well-understood | Needs a good LR schedule and warm-up |\n",
"| **RAdam** | Short runs or when warm-up tuning is painful | Slightly more compute per step |\n",
"| **Adafactor** | Large models (LLMs, T5-scale); memory-constrained training | Can be slower to converge on small models |\n",
"| **Schedule-Free AdamW** | When you want to skip LR schedule tuning entirely | Must not add a `learning_rate_scheduler`; needs correct `warmup_steps` |\n",
"| **Muon** | Training deep networks with many weight matrices | Falls back to AdamW for embeddings/biases; newer, less battle-tested |\n",
"| **SOAP** | When you want Shampoo-quality second-order convergence with Adam memory | Higher per-step cost than Adam |"
]
},
{
"cell_type": "markdown",
"id": "sfa-tip-header",
"metadata": {},
"source": [
"## Schedule-Free AdamW — important tip\n",
"\n",
"> **Do not add `learning_rate_scheduler` when using `schedule_free_adamw`.**\n",
">\n",
"> Schedule-Free AdamW internalises the learning-rate schedule via a\n",
"> momentum-based averaging trick. Adding an external cosine or step-decay\n",
"> scheduler on top fights the internal schedule and typically hurts final\n",
"> accuracy.\n",
">\n",
"> Correct config:\n",
"> ```yaml\n",
"> trainer:\n",
"> optimizer:\n",
"> type: schedule_free_adamw\n",
"> lr: 0.001\n",
"> ```\n",
">\n",
"> You may still set `warmup_steps` inside the optimizer block to give the\n",
"> internal schedule a short ramp-up."
]
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3",
"language": "python",
"name": "python3"
},
"language_info": {
"name": "python",
"version": "3.10.0"
}
},
"nbformat": 4,
"nbformat_minor": 5
}