feat(optim): mutable learning rate and warmup+cosine LR schedule

skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/exec/optim/SgdOptimizerTest.kt

@@ -86,4 +86,28 @@ class SgdOptimizerTest {
         opt.step()
         assertEquals(3.9025f, w.value.data[0, 0], 1e-3f)
     }
+
+    @Test
+    fun lr_can_be_rescheduled_between_steps() {
+        val ctx = DirectCpuExecutionContext(phase = Phase.TRAIN)
+        val w = param1x1(10f)
+
+        val opt = sk.ainet.lang.nn.optim.SgdOptimizer(lr = 0.1)
+        opt.addParameter(w)
+
+        // Step 1 at lr = 0.1: w = 10 - 0.1 * 2 = 9.8
+        val g1 = ctx.full<FP32, Float>(Shape(1, 1), FP32::class, 2.0)
+        w.value.accumulateGrad(g1 as sk.ainet.lang.tensor.Tensor<FP32, Float>)
+        opt.step()
+        opt.zeroGrad()
+        assertEquals(9.8f, w.value.data[0, 0], 1e-6f)
+
+        // Reschedule to lr = 0.5 (as a warmup/decay schedule would):
+        // w = 9.8 - 0.5 * 2 = 8.8
+        opt.lr = 0.5
+        val g2 = ctx.full<FP32, Float>(Shape(1, 1), FP32::class, 2.0)
+        w.value.accumulateGrad(g2 as sk.ainet.lang.tensor.Tensor<FP32, Float>)
+        opt.step()
+        assertEquals(8.8f, w.value.data[0, 0], 1e-6f)
+    }
 }

skainet-lang/skainet-lang-core/api/android/skainet-lang-core.api

@@ -1897,11 +1897,22 @@ public final class sk/ainet/lang/nn/optim/AdamOptimizer : sk/ainet/lang/nn/optim
 	public synthetic fun <init> (DDDDDZZILkotlin/jvm/internal/DefaultConstructorMarker;)V
 	public fun addParameter (Lsk/ainet/lang/nn/topology/ModuleParameter;Z)V
 	public fun addParameter (Lsk/ainet/lang/nn/topology/Parameter;Z)V
+	public final fun getLr ()D
 	public final fun reset ()V
+	public final fun setLr (D)V
 	public fun step ()V
 	public fun zeroGrad ()V
 }
 
+public abstract interface class sk/ainet/lang/nn/optim/LrSchedule {
+	public abstract fun lrAt (I)D
+}
+
+public final class sk/ainet/lang/nn/optim/LrScheduleKt {
+	public static final fun linearWarmupCosineDecay (IIDDD)Lsk/ainet/lang/nn/optim/LrSchedule;
+	public static synthetic fun linearWarmupCosineDecay$default (IIDDDILjava/lang/Object;)Lsk/ainet/lang/nn/optim/LrSchedule;
+}
+
 public abstract interface class sk/ainet/lang/nn/optim/Optimizer {
 	public abstract fun addParameter (Lsk/ainet/lang/nn/topology/ModuleParameter;Z)V
 	public abstract fun addParameter (Lsk/ainet/lang/nn/topology/Parameter;Z)V
@@ -1928,6 +1939,8 @@ public final class sk/ainet/lang/nn/optim/SgdOptimizer : sk/ainet/lang/nn/optim/
 	public synthetic fun <init> (DDDILkotlin/jvm/internal/DefaultConstructorMarker;)V
 	public fun addParameter (Lsk/ainet/lang/nn/topology/ModuleParameter;Z)V
 	public fun addParameter (Lsk/ainet/lang/nn/topology/Parameter;Z)V
+	public final fun getLr ()D
+	public final fun setLr (D)V
 	public fun step ()V
 	public fun zeroGrad ()V
 }

skainet-lang/skainet-lang-core/api/jvm/skainet-lang-core.api

@@ -2125,11 +2125,22 @@ public final class sk/ainet/lang/nn/optim/AdamOptimizer : sk/ainet/lang/nn/optim
 	public synthetic fun <init> (DDDDDZZILkotlin/jvm/internal/DefaultConstructorMarker;)V
 	public fun addParameter (Lsk/ainet/lang/nn/topology/ModuleParameter;Z)V
 	public fun addParameter (Lsk/ainet/lang/nn/topology/Parameter;Z)V
+	public final fun getLr ()D
 	public final fun reset ()V
+	public final fun setLr (D)V
 	public fun step ()V
 	public fun zeroGrad ()V
 }
 
+public abstract interface class sk/ainet/lang/nn/optim/LrSchedule {
+	public abstract fun lrAt (I)D
+}
+
+public final class sk/ainet/lang/nn/optim/LrScheduleKt {
+	public static final fun linearWarmupCosineDecay (IIDDD)Lsk/ainet/lang/nn/optim/LrSchedule;
+	public static synthetic fun linearWarmupCosineDecay$default (IIDDDILjava/lang/Object;)Lsk/ainet/lang/nn/optim/LrSchedule;
+}
+
 public abstract interface class sk/ainet/lang/nn/optim/Optimizer {
 	public abstract fun addParameter (Lsk/ainet/lang/nn/topology/ModuleParameter;Z)V
 	public abstract fun addParameter (Lsk/ainet/lang/nn/topology/Parameter;Z)V
@@ -2156,6 +2167,8 @@ public final class sk/ainet/lang/nn/optim/SgdOptimizer : sk/ainet/lang/nn/optim/
 	public synthetic fun <init> (DDDILkotlin/jvm/internal/DefaultConstructorMarker;)V
 	public fun addParameter (Lsk/ainet/lang/nn/topology/ModuleParameter;Z)V
 	public fun addParameter (Lsk/ainet/lang/nn/topology/Parameter;Z)V
+	public final fun getLr ()D
+	public final fun setLr (D)V
 	public fun step ()V
 	public fun zeroGrad ()V
 }

skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/nn/optim/Adam.kt

@@ -33,7 +33,12 @@ import kotlin.math.pow
  * @param amsgrad If true, uses the AMSGrad variant that maintains the maximum of all v_t (default: false)
  */
 public class AdamOptimizer @kotlin.jvm.JvmOverloads constructor(
-    private val lr: Double = 0.001,
+    /**
+     * Learning rate. Mutable so learning-rate schedules (warmup, cosine
+     * decay, …) can adjust it between steps without recreating the optimizer
+     * and losing the moment estimates. See [LrSchedule].
+     */
+    public var lr: Double = 0.001,
     private val beta1: Double = 0.9,
     private val beta2: Double = 0.999,
     private val epsilon: Double = 1e-8,

skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/nn/optim/LrSchedule.kt

@@ -0,0 +1,56 @@
+package sk.ainet.lang.nn.optim
+
+import kotlin.math.PI
+import kotlin.math.cos
+
+/**
+ * A learning-rate schedule: maps a global training step to a learning rate.
+ *
+ * Schedules are stateless and decoupled from optimizers — a training loop
+ * assigns the value to the optimizer's mutable `lr` before each step:
+ *
+ * ```kotlin
+ * val schedule = linearWarmupCosineDecay(totalSteps, warmupSteps, peakLr = 5e-4)
+ * // inside the loop:
+ * optimizer.lr = schedule.lrAt(globalStep)
+ * ```
+ */
+public fun interface LrSchedule {
+    /** The learning rate to use for [step] (0-based, `0 <= step < totalSteps`). */
+    public fun lrAt(step: Int): Double
+}
+
+/**
+ * Linear warmup followed by cosine decay — the schedule used by GPT-style
+ * pretraining recipes.
+ *
+ * - steps `0 ..< warmupSteps`: linear ramp from [initialLr] to [peakLr]
+ * - steps `warmupSteps ..< totalSteps`: cosine decay from [peakLr] to [minLr]
+ *
+ * @param totalSteps total number of optimization steps
+ * @param warmupSteps steps spent ramping up (must be in `1 ..< totalSteps`)
+ * @param peakLr learning rate at the end of the warmup
+ * @param initialLr learning rate at step 0
+ * @param minLr learning rate approached at the end of training
+ */
+public fun linearWarmupCosineDecay(
+    totalSteps: Int,
+    warmupSteps: Int,
+    peakLr: Double,
+    initialLr: Double = 3e-5,
+    minLr: Double = 1e-6,
+): LrSchedule {
+    require(totalSteps > 0) { "totalSteps must be positive, was $totalSteps" }
+    require(warmupSteps in 1 until totalSteps) {
+        "warmupSteps must be in 1 until totalSteps=$totalSteps, was $warmupSteps"
+    }
+    return LrSchedule { step ->
+        require(step in 0 until totalSteps) { "step $step out of range [0, $totalSteps)" }
+        if (step < warmupSteps) {
+            initialLr + (peakLr - initialLr) * step / warmupSteps
+        } else {
+            val progress = (step - warmupSteps).toDouble() / (totalSteps - warmupSteps)
+            minLr + (peakLr - minLr) * 0.5 * (1 + cos(PI * progress))
+        }
+    }
+}

skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/nn/optim/SGD.kt

@@ -9,7 +9,11 @@ import sk.ainet.lang.types.DType
  * Stochastic Gradient Descent optimizer with optional momentum and weight decay.
  */
 public class SgdOptimizer @kotlin.jvm.JvmOverloads constructor(
-    private val lr: Double,
+    /**
+     * Learning rate. Mutable so learning-rate schedules can adjust it between
+     * steps. See [LrSchedule].
+     */
+    public var lr: Double,
     private val momentum: Double = 0.0,
     private val weightDecay: Double = 0.0,
 ) : Optimizer {

skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/nn/optim/LrScheduleTest.kt

@@ -0,0 +1,60 @@
+package sk.ainet.lang.nn.optim
+
+import kotlin.math.abs
+import kotlin.test.Test
+import kotlin.test.assertEquals
+import kotlin.test.assertFailsWith
+import kotlin.test.assertTrue
+
+class LrScheduleTest {
+
+    private val total = 100
+    private val warmup = 20
+    private val peak = 1e-3
+    private val initial = 3e-5
+    private val min = 1e-6
+
+    private val schedule = linearWarmupCosineDecay(total, warmup, peak, initial, min)
+
+    @Test
+    fun warmup_starts_at_initialLr_and_ramps_linearly() {
+        assertEquals(initial, schedule.lrAt(0), 1e-12)
+        val lrs = (0 until warmup).map { schedule.lrAt(it) }
+        lrs.zipWithNext().forEach { (a, b) -> assertTrue(b > a, "warmup must strictly increase") }
+        // Linear: constant increments
+        val increments = lrs.zipWithNext().map { (a, b) -> b - a }
+        increments.zipWithNext().forEach { (a, b) -> assertEquals(a, b, 1e-12) }
+    }
+
+    @Test
+    fun peak_is_reached_at_the_end_of_warmup() {
+        assertEquals(peak, schedule.lrAt(warmup), 1e-12)
+    }
+
+    @Test
+    fun decay_is_monotonic_and_approaches_minLr() {
+        val lrs = (warmup until total).map { schedule.lrAt(it) }
+        lrs.zipWithNext().forEach { (a, b) -> assertTrue(b <= a, "decay must not increase") }
+        assertTrue(abs(lrs.last() - min) < peak * 0.01, "final lr ${lrs.last()} should approach $min")
+    }
+
+    @Test
+    fun invalid_arguments_are_rejected() {
+        assertFailsWith<IllegalArgumentException> { linearWarmupCosineDecay(0, 1, peak) }
+        assertFailsWith<IllegalArgumentException> { linearWarmupCosineDecay(10, 0, peak) }
+        assertFailsWith<IllegalArgumentException> { linearWarmupCosineDecay(10, 10, peak) }
+        assertFailsWith<IllegalArgumentException> { schedule.lrAt(-1) }
+        assertFailsWith<IllegalArgumentException> { schedule.lrAt(total) }
+    }
+
+    @Test
+    fun optimizer_lr_is_assignable_between_steps() {
+        val adam = AdamOptimizer(lr = 0.001)
+        adam.lr = schedule.lrAt(warmup)
+        assertEquals(peak, adam.lr, 1e-12)
+
+        val sgd = SgdOptimizer(lr = 0.1)
+        sgd.lr = 0.05
+        assertEquals(0.05, sgd.lr, 1e-12)
+    }
+}