From cf653f7664d5d4ef636ee2f2abd842bc2d3f4aff Mon Sep 17 00:00:00 2001
From: Lars Bogner <ich@larsbogner.de>
Date: Fri, 19 Jun 2026 10:43:31 +0200
Subject: [PATCH] Add linear warmup before cosine LR decay

Replaces CosineAnnealingLR with a LambdaLR that linearly ramps the LR
from lr/warmup_epochs to lr over the first warmup_epochs steps, then
applies cosine decay for the remainder. Default warmup_epochs=5;
overridable via --warmup-epochs CLI flag.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
---
 giant/config.py   |  1 +
 giant/pipeline.py |  1 +
 giant/train.py    | 12 +++++++++++-
 scripts/train.py  |  2 ++
 4 files changed, 15 insertions(+), 1 deletion(-)

diff --git a/giant/config.py b/giant/config.py
index 9724f88..5fbb127 100644
--- a/giant/config.py
+++ b/giant/config.py
@@ -19,6 +19,7 @@ DEFAULT_CONFIG: dict = {
         "seed": 0,
         "validate_every": 10,
         "validate_steps": 10,
+        "warmup_epochs": 5,
     },
     "model": {
         "hidden_dim": 256,
diff --git a/giant/pipeline.py b/giant/pipeline.py
index 895b244..5976ae4 100644
--- a/giant/pipeline.py
+++ b/giant/pipeline.py
@@ -144,6 +144,7 @@ def run_train_job(
         mode=t["mode"],
         epochs=t["epochs"],
         lr=t["lr"],
+        warmup_epochs=t["warmup_epochs"],
         device=device,
         out_dir=out_dir,
         normalizer_dict={"cond": cond_norm.to_dict(), "target": tgt_norm.to_dict()},
diff --git a/giant/train.py b/giant/train.py
index f904fc5..539691c 100644
--- a/giant/train.py
+++ b/giant/train.py
@@ -1,4 +1,5 @@
 import csv
+import math
 import os
 import signal
 import time
@@ -62,6 +63,7 @@ def train(
     mode: str,
     epochs: int,
     lr: float,
+    warmup_epochs: int,
     device: torch.device,
     out_dir: str | Path,
     normalizer_dict: dict | None = None,
@@ -77,7 +79,15 @@ def train(
 
     model = model.to(device)
     optimizer = optim.AdamW(model.parameters(), lr=lr)
-    lr_sched = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs)
+
+    def _lr_lambda(epoch: int) -> float:
+        if warmup_epochs > 0 and epoch < warmup_epochs:
+            return (epoch + 1) / warmup_epochs
+        t = epoch - warmup_epochs
+        T = max(epochs - warmup_epochs, 1)
+        return 0.5 * (1.0 + math.cos(math.pi * t / T))
+
+    lr_sched = optim.lr_scheduler.LambdaLR(optimizer, _lr_lambda)
 
     ddpm_schedule = CosineSchedule().to(device) if mode == "ddpm" else None
 
diff --git a/scripts/train.py b/scripts/train.py
index 067ceec..899f380 100644
--- a/scripts/train.py
+++ b/scripts/train.py
@@ -17,6 +17,7 @@ def main() -> None:
     parser.add_argument("--epochs", type=int)
     parser.add_argument("--batch-size", type=int)
     parser.add_argument("--lr", type=float)
+    parser.add_argument("--warmup-epochs", type=int, dest="warmup_epochs")
     parser.add_argument("--hidden-dim", type=int)
     parser.add_argument("--n-blocks", type=int)
     parser.add_argument("--emb-dim", type=int)
@@ -57,6 +58,7 @@ def main() -> None:
             "epochs": args.epochs,
             "batch_size": args.batch_size,
             "lr": args.lr,
+            "warmup_epochs": args.warmup_epochs,
             "val_fraction": args.val_fraction,
             "num_workers": args.num_workers,
             "seed": args.seed,