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
584 lines
19 KiB
Plaintext
584 lines
19 KiB
Plaintext
{
|
|
"cells": [
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "title-cell",
|
|
"metadata": {},
|
|
"source": [
|
|
"# Optimizer Comparison: Schedule-Free, Muon, Adafactor, and More\n",
|
|
"\n",
|
|
"[](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
|
|
}
|