{ "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 }