Skip to content

autodiff_optimize.kt

src/jvmMain/kotlin/afield/math/autodiff/autodiff_optimize.kt · 623 lines

package and imports — 53 lines, from line 1 
@file:Suppress("MemberVisibilityCanBePrivate")
@file:afield.runtime.completion.CompletionCategory("Math")

package afield.math.autodiff

import afield.data.collections.get
import afield.data.collections.times
import afield.data.graph.rules
import afield.document.graph.minusAssign
import afield.document.graph.plusAssign
import afield.document.properties.compareTo
import afield.document.workspace.file
import afield.document.workspace.path
import afield.geometry.mesh.normalized
import afield.geometry.mesh.scaled
import afield.graphics.render.rangeTo
import afield.math.functions.FnEvolveRun
import afield.math.functions.FnFitRun
import afield.math.functions.evolve
import afield.math.functions.fitGradient
import afield.math.linalg.bar
import afield.math.linalg.f
import afield.math.linalg.remAssign
import afield.math.linalg.set
import afield.math.linalg.times
import afield.math.units.div
import afield.math.units.minus
import afield.math.units.plus
import afield.math.units.times
import afield.media.image.Gradient
import afield.media.image.normalized
import afield.runtime.execution.end
import afield.runtime.execution.go
import afield.runtime.execution.start
import afield.simulation.automata.history

import afield.audio.synth.*
import afield.data.collections.alphas
import afield.data.collections.invoke
import afield.data.collections.m
import afield.data.collections.s
import afield.geometry.mesh.gradient
import afield.math.functions.DenseDoubleOps
import afield.math.linalg.i
import afield.math.linalg.invoke
import afield.math.linalg.x
import afield.math.linalg.yy
import afield.simulation.agents.step
import afield.simulation.automata.step
import afield.simulation.automata.trajectory
import kotlin.math.absoluteValue
import kotlin.math.sqrt
import kotlin.random.Random

Gradient-descent optimizers that consume the exact gradients from autodiff.kt — sgd, adam and lbfgs, behind one minimize front door.

All three are deterministic given their seed and all three record a full OptimizeRun.losses trajectory which you can hand straight to showSignals("loss" to run.losses) and look at. FnEvolveRun and FnFitRun carry a losses list of the same shape, so a fit, an evolve and an optimisation all replay and plot the same way.

val run = minimize(doubleArrayOf(-1.2, 1.0), OptimizeOptions(method = OptimizeMethod.LBFGS)) { p ->
    (1.0 - p[0]).sq() + 100.0 * (p[1] - p[0].sq()).sq()   // Rosenbrock
}
run.x          // ~[1, 1]
run.losses     // the trajectory, ready to plot

Which one

  • lbfgs is the right default for a smooth deterministic objective with a handful to a few hundred parameters: it builds a curvature model from the last m steps and takes Newton-shaped strides, so it reaches machine-precision minima in tens of iterations where Adam is still crawling. It needs the objective to actually be smooth; a kinked or noisy loss will make its line search flail.
  • adam is the robust workhorse for a rough, kinked or stochastic landscape, and the one that tolerates a badly scaled parameter vector without preconditioning.
  • sgd with momentum is here because it is sometimes what you want artistically — a slower, visibly settling trajectory with a tunable inertia, which reads as motion rather than as a solver converging.

The line search inside lbfgs is backtracking Armijo, which guarantees monotone descent but not the curvature (Wolfe) condition; the L-BFGS memory therefore skips any update pair with non-positive curvature rather than letting an indefinite model through.

enum class OptimizeMethod

Which descent to run. See the file KDoc for when to use each.

enum class OptimizeMethod {

Plain gradient descent with (optionally Nesterov) momentum.

    SGD,

Adaptive moment estimation — per-parameter step sizes from the gradient's running moments.

    ADAM,

Limited-memory BFGS with a backtracking Armijo line search.

    LBFGS
}

data class OptimizeOptions

Everything that shapes a minimize run, kept as one immutable value so a result can carry its own provenance.

The defaults are a 500-iteration Adam at lr = 0.05 — the same learning rate fitGradient uses, so switching a fit from finite differences to exact gradients does not also silently change the step size.

data class OptimizeOptions(
    /** Which optimizer to run. */
    val method: OptimizeMethod = OptimizeMethod.ADAM,
    /** Maximum number of iterations. */
    val iterations: Int = 500,
    /** Step size. Ignored by [OptimizeMethod.LBFGS], whose line search chooses its own. */
    val learningRate: Double = 0.05,
    /** Momentum coefficient for [OptimizeMethod.SGD]. */
    val momentum: Double = 0.9,
    /** Use the Nesterov look-ahead form of momentum. */
    val nesterov: Boolean = false,
    /** Adam's first-moment decay. */
    val beta1: Double = 0.9,
    /** Adam's second-moment decay. */
    val beta2: Double = 0.999,
    /** Adam's denominator floor. */
    val epsilon: Double = 1e-8,
    /** How many `(s, y)` curvature pairs L-BFGS remembers. */
    val memory: Int = 8,
    /** Armijo sufficient-decrease constant for the L-BFGS line search. */
    val armijo: Double = 1e-4,
    /** How many halvings the line search will try before giving up on an iteration. */
    val maxLineSearchSteps: Int = 32,
    /** Stop when the gradient's infinity norm falls to this. */
    val gradientTolerance: Double = 1e-9,
    /** Stop when a step moves the parameters less than this. */
    val stepTolerance: Double = 1e-14,
    /** If > 0, jitter the start by `±initJitter·(1+|x|)` using [rngSeed] before descending. */
    val initJitter: Double = 0.0,
    /** The single source of randomness; with `initJitter = 0` the run is randomness-free anyway. */
    val rngSeed: Long = 0L,
    /** Whether to offer a progress bar. Turn it off for an inner loop (a restart sweep, a demo). */
    val showProgress: Boolean = true,
    /** Retain a copy of every parameter vector. Disable for large simulations; scalar history remains. */
    val recordTrajectory: Boolean = true
) {

init

    init {
        require(iterations >= 0) { "iterations must be non-negative" }
        require(learningRate > 0.0) { "learningRate must be positive" }
        require(memory >= 1) { "memory must be >= 1" }
        require(armijo > 0.0 && armijo < 1.0) { "armijo must be in (0,1)" }
    }
}

data class OptimizeStep

One recorded iteration: the objective and the gradient's infinity norm after the step.

data class OptimizeStep(
    /** 1-based iteration index; iteration 0 is the starting point and lives only in [OptimizeRun.losses]. */
    val iteration: Int,
    /** The objective value after this step. */
    val value: Double,
    /** The gradient's infinity norm after this step — the convergence signal that counts. */
    val gradientInfNorm: Double
)

data class OptimizeRun

The result and provenance of a minimize run, shaped like FnFitRun so the two read the same way at a call site.

losses is start-inclusive, so losses.first() is the objective before any step and losses.last() is value; that is what makes it plottable without an off-by-one.

declaration — 41 lines, from line 171 
data class OptimizeRun(
    /** The parameters at the end of the run. */
    val x: DoubleArray,
    /** The objective at [x]. */
    val value: Double,
    /** The exact gradient at [x]. */
    val gradient: DoubleArray,
    /** The starting parameters, after any [OptimizeOptions.initJitter]. */
    val initialX: DoubleArray,
    /** The objective at [initialX]. */
    val initialValue: Double,
    /** Per-iteration records. */
    val steps: List<OptimizeStep>,
    /** The objective trajectory including the starting point. */
    val losses: List<Double>,
    /**
     * The **parameter** trajectory, indexed exactly like [losses]: `trajectory[i]` is the point at
     * which `losses[i]` was measured, so `trajectory.first()` is [initialX] and `trajectory.last()`
     * is [x].
     *
     * Recorded because a finished run otherwise cannot tell you what the answer *looked like* on the
     * way — only how good it was. Reconstructing an intermediate point means re-running the descent
     * truncated at that iteration, which is both slower and a different thing: a re-run is only the
     * same point if the descent is deterministic, and nothing forces it to be.
     *
     * Each entry is a defensive copy, because the search mutates its point in place.
     * Empty when [OptimizeOptions.recordTrajectory] is false; [initialX], [x], [losses]
     * and [steps] remain available without storing every intermediate parameter vector.
     */
    val trajectory: List<DoubleArray>,
    /** How many iterations actually ran. */
    val iterations: Int,
    /** How many times the objective (and its gradient) was evaluated. */
    val evaluations: Int,
    /** Whether a tolerance was met, as opposed to running out of iterations. */
    val converged: Boolean,
    /** One of "observer", "gradient", "step", "line search", "max iterations", "no iterations". */
    val reason: String,
    /** The options this run was made with. */
    val options: OptimizeOptions
) {

val improved: Boolean get() = value < initialValue

True when the run ended strictly below where it started.

val gradientInfNorm: Double get()

The gradient's infinity norm at x.

    val gradientInfNorm: Double get() = if (gradient.isEmpty()) 0.0 else gradient.maxOf { it.absoluteValue }

override fun equals(other: Any?): Boolean

Two runs are equal when x, value, losses, iterations, evaluations, converged, reason and options all match. trajectory is left out: it is a List<DoubleArray> and DoubleArray.equals is identity, so including it would make any two separately-computed runs unequal whatever their contents. It is fully determined by initialX, options and the objective, and x and losses compare its endpoint and its whole value trace.

    override fun equals(other: Any?): Boolean {
        if (this === other) return true
        if (other !is OptimizeRun) return false
        return x.contentEquals(other.x) && value == other.value && losses == other.losses &&
            iterations == other.iterations && evaluations == other.evaluations &&
            converged == other.converged && reason == other.reason && options == other.options
    }

override fun hashCode(): Int

    override fun hashCode(): Int {
        var result = x.contentHashCode()
        result = 31 * result + value.hashCode()
        result = 31 * result + losses.hashCode()
        result = 31 * result + reason.hashCode()
        return result
    }
}

fun minimize

Minimises an objective written in AD arithmetic, using exact gradients.

Pass the loss; get the answer and its whole trajectory.

val run = minimize(doubleArrayOf(0.0, 0.0)) { p -> (p[0] - 3.0).sq() + 2.0 * (p[1] + 1.0).sq() }
fun minimize(
    initial: DoubleArray,
    options: OptimizeOptions = OptimizeOptions(),
    objective: (List<AD>) -> AD
): OptimizeRun = minimizeWithGradient(initial, options, valueAndGradOf(objective))

data class OptimizeProgress(val iteration: Int, val x: DoubleArray, val value: Double,

An accepted optimizer state (including iteration zero), with line-search diagnostics. x and gradient are borrowed read-only arrays, valid only during the callback. Copy them only if retaining them. Rejected trial points are never reported.

data class OptimizeProgress(val iteration: Int, val x: DoubleArray, val value: Double,
    val gradient: DoubleArray, val evaluations: Int, val stepNorm: Double,
    val lineSearchReductions: Int, val curvatureAccepted: Boolean?)

fun minimizeWithGradient

Minimises an objective you already have a value-and-gradient function for: the entry point for a gradient that does not come from this file's AD (a hand-derived one, or the exact Fn-graph gradient in autodiff_fn.kt).

fun minimizeWithGradient(
    initial: DoubleArray,
    options: OptimizeOptions = OptimizeOptions(),
    valueAndGradient: (DoubleArray) -> ADGradient
): OptimizeRun = minimizeWithGradient(initial, options, null, valueAndGradient)

fun minimizeWithGradient

minimizeWithGradient with an accepted-state observer. Return true when the caller's convergence criterion is satisfied; the run stops with reason observer. Exceptions propagate (including cancellation). The default optimizer stopping rules still apply.

body — 110 lines, from line 280 
fun minimizeWithGradient(
    initial: DoubleArray,
    options: OptimizeOptions,
    observer: ((OptimizeProgress) -> Boolean)?,
    valueAndGradient: (DoubleArray) -> ADGradient
): OptimizeRun {
    require(initial.isNotEmpty()) { "initial must not be empty" }

    val x = initial.copyOf()
    if (options.initJitter > 0.0) {
        val rng = Random(options.rngSeed)
        for (i in x.indices) {
            x[i] += (rng.nextDouble() * 2.0 - 1.0) * options.initJitter * (1.0 + x[i].absoluteValue)
        }
    }
    val startX = x.copyOf()

    var evaluations = 0
    fun eval(p: DoubleArray): ADGradient {
        evaluations++
        return valueAndGradient(p)
    }

    var current = eval(x)
    val losses = mutableListOf(current.value)
    val trajectory = mutableListOf<DoubleArray>()
    if (options.recordTrajectory) trajectory += x.copyOf()
    val steps = mutableListOf<OptimizeStep>()

    val initiallyAccepted = observer?.invoke(OptimizeProgress(0,x,current.value,current.gradient,
        evaluations,0.0,0,null)) == true
    if (options.iterations == 0 || initiallyAccepted) {
        return OptimizeRun(
            x, current.value, current.gradient, startX, losses.first(), steps, losses,
            trajectory, 0, evaluations, converged = initiallyAccepted,
            reason = if (initiallyAccepted) "observer" else "no iterations", options = options
        )
    }

    val state: DescentState = when (options.method) {
        OptimizeMethod.SGD -> SgdState(x.size)
        OptimizeMethod.ADAM -> AdamState(x.size)
        OptimizeMethod.LBFGS -> LbfgsState(options.memory)
    }

    // Bounded like evolve/fitGradient — a descent can never disappear into an animation loop.
    val reporter = if (options.showProgress) afield.ui.output.progressReporter("minimize") else null
    var converged = false
    var reason = "max iterations"
    var iterationsRun = 0

    try {
        for (iteration in 1..options.iterations) {
            if (current.gradientInfNorm <= options.gradientTolerance) {
                converged = true
                reason = "gradient"
                break
            }

            val outcome = state.step(iteration, x, current, options, ::eval)
            if (outcome == null) {
                converged = true
                reason = "line search"
                iterationsRun = iteration - 1
                break
            }
            current = outcome.next
            iterationsRun = iteration
            losses += current.value
            if (options.recordTrajectory) trajectory += x.copyOf()
            steps += OptimizeStep(iteration, current.value, current.gradientInfNorm)
            reporter?.update(iteration.toLong(), options.iterations.toLong())

            if (observer?.invoke(OptimizeProgress(iteration,x,current.value,current.gradient,
                    evaluations,outcome.stepNorm,outcome.lineSearchReductions,outcome.curvatureAccepted)) == true) {
                converged = true
                reason = "observer"
                break
            }
            if (outcome.stepNorm <= options.stepTolerance) {
                converged = true
                reason = "step"
                break
            }
        }
    } finally {
        reporter?.finish()
    }

    if (!converged && current.gradientInfNorm <= options.gradientTolerance) {
        converged = true
        reason = "gradient"
    }

    return OptimizeRun(
        x = x,
        value = current.value,
        gradient = current.gradient,
        initialX = startX,
        initialValue = losses.first(),
        steps = steps,
        losses = losses,
        trajectory = trajectory,
        iterations = iterationsRun,
        evaluations = evaluations,
        converged = converged,
        reason = reason,
        options = options
    )
}

fun sgd

minimize with OptimizeMethod.SGD — plain descent with momentum, and a visibly settling path.

fun sgd(
    initial: DoubleArray,
    iterations: Int = 500,
    lr: Double = 0.05,
    momentum: Double = 0.9,
    nesterov: Boolean = false,
    objective: (List<AD>) -> AD
): OptimizeRun = minimize(
    initial,
    OptimizeOptions(
        method = OptimizeMethod.SGD, iterations = iterations, learningRate = lr,
        momentum = momentum, nesterov = nesterov
    ),
    objective
)

fun adam

minimize with OptimizeMethod.ADAM — the robust default for a rough or badly scaled landscape.

fun adam(
    initial: DoubleArray,
    iterations: Int = 500,
    lr: Double = 0.05,
    beta1: Double = 0.9,
    beta2: Double = 0.999,
    epsilon: Double = 1e-8,
    objective: (List<AD>) -> AD
): OptimizeRun = minimize(
    initial,
    OptimizeOptions(
        method = OptimizeMethod.ADAM, iterations = iterations, learningRate = lr,
        beta1 = beta1, beta2 = beta2, epsilon = epsilon
    ),
    objective
)

fun lbfgs

minimize with OptimizeMethod.LBFGS — the line-searched quasi-Newton, for smooth objectives.

fun lbfgs(
    initial: DoubleArray,
    iterations: Int = 200,
    memory: Int = 8,
    gradientTolerance: Double = 1e-9,
    objective: (List<AD>) -> AD
): OptimizeRun = minimize(
    initial,
    OptimizeOptions(
        method = OptimizeMethod.LBFGS, iterations = iterations,
        memory = memory, gradientTolerance = gradientTolerance
    ),
    objective
)

The three descents

What one optimizer iteration produced: the new value/gradient and how far the step moved.

private class StepOutcome(val next: ADGradient, val stepNorm: Double,
    val lineSearchReductions: Int = 0, val curvatureAccepted: Boolean? = null)

The mutable per-run state of one descent; step mutates x in place and returns the outcome.

source — 172 lines, from line 451 
private interface DescentState {
    fun step(
        iteration: Int,
        x: DoubleArray,
        current: ADGradient,
        options: OptimizeOptions,
        eval: (DoubleArray) -> ADGradient
    ): StepOutcome?
}

private class SgdState(n: Int) : DescentState {
    private val velocity = DoubleArray(n)

    override fun step(
        iteration: Int,
        x: DoubleArray,
        current: ADGradient,
        options: OptimizeOptions,
        eval: (DoubleArray) -> ADGradient
    ): StepOutcome {
        var moved = 0.0
        for (i in x.indices) {
            val g = current.gradient[i]
            if (!g.isFinite()) continue
            val previous = velocity[i]
            velocity[i] = options.momentum * previous - options.learningRate * g
            // Nesterov's look-ahead: apply the *new* velocity plus the momentum-weighted change.
            val delta = if (options.nesterov) {
                options.momentum * velocity[i] - options.learningRate * g
            } else {
                velocity[i]
            }
            if (delta.isFinite()) {
                x[i] += delta
                moved += delta * delta
            }
        }
        return StepOutcome(eval(x), sqrt(moved))
    }
}

private class AdamState(n: Int) : DescentState {
    private val m = DoubleArray(n)
    private val v = DoubleArray(n)

    override fun step(
        iteration: Int,
        x: DoubleArray,
        current: ADGradient,
        options: OptimizeOptions,
        eval: (DoubleArray) -> ADGradient
    ): StepOutcome {
        var moved = 0.0
        val correction1 = 1.0 - Math.pow(options.beta1, iteration.toDouble())
        val correction2 = 1.0 - Math.pow(options.beta2, iteration.toDouble())
        for (i in x.indices) {
            val g = current.gradient[i]
            if (!g.isFinite()) continue
            m[i] = options.beta1 * m[i] + (1.0 - options.beta1) * g
            v[i] = options.beta2 * v[i] + (1.0 - options.beta2) * g * g
            val update = options.learningRate * (m[i] / correction1) / (sqrt(v[i] / correction2) + options.epsilon)
            if (update.isFinite()) {
                x[i] -= update
                moved += update * update
            }
        }
        return StepOutcome(eval(x), sqrt(moved))
    }
}

private class LbfgsState(private val memory: Int) : DescentState {
    private val sHistory = ArrayDeque<DoubleArray>()
    private val yHistory = ArrayDeque<DoubleArray>()
    private val rhoHistory = ArrayDeque<Double>()

    override fun step(
        iteration: Int,
        x: DoubleArray,
        current: ADGradient,
        options: OptimizeOptions,
        eval: (DoubleArray) -> ADGradient
    ): StepOutcome? {
        val n = x.size
        val direction = twoLoopRecursion(current.gradient, n)

        // The model can go uphill after a skipped curvature pair; fall back to steepest descent.
        var slope = DenseDoubleOps.dot(direction, current.gradient)
        if (slope >= 0.0) {
            for (i in 0 until n) direction[i] = -current.gradient[i]
            slope = DenseDoubleOps.dot(direction, current.gradient)
        }
        if (slope >= 0.0 || !slope.isFinite()) return null

        // Backtracking Armijo. The first iteration starts small because the initial inverse-Hessian
        // guess is the identity and the raw gradient can be enormous.
        var alpha = if (iteration == 1 && sHistory.isEmpty()) {
            val norm = sqrt(current.gradient.sumOf { it * it })
            if (norm > 1.0) 1.0 / norm else 1.0
        } else {
            1.0
        }

        val candidate = DoubleArray(n)
        var accepted: ADGradient? = null
        var reductions = 0
        for (attempt in 0 until options.maxLineSearchSteps) {
            DenseDoubleOps.combine(candidate, x, alpha, direction)
            val probe = eval(candidate)
            if (probe.value.isFinite() && probe.value <= current.value + options.armijo * alpha * slope) {
                accepted = probe
                break
            }
            alpha *= 0.5
            reductions++
        }
        val next = accepted ?: return null

        val s = DoubleArray(n) { candidate[it] - x[it] }
        val y = DoubleArray(n) { next.gradient[it] - current.gradient[it] }
        val sy = DenseDoubleOps.dot(s, y)
        // A relative curvature test is invariant to objective units; an absolute
        // cutoff discards useful pairs when energies/gradients are normalized.
        val sn = sqrt(DenseDoubleOps.dot(s, s))
        val yn = sqrt(DenseDoubleOps.dot(y, y))
        val usableCurvature = sy > 0.0 && sy.isFinite() && (1.0 / sy).isFinite() && sn > 0.0 && yn > 0.0 && sy / sn / yn > 1e-10
        if (usableCurvature) {
            sHistory.addLast(s)
            yHistory.addLast(y)
            rhoHistory.addLast(1.0 / sy)
            if (sHistory.size > memory) {
                sHistory.removeFirst()
                yHistory.removeFirst()
                rhoHistory.removeFirst()
            }
        }

        for (i in 0 until n) x[i] = candidate[i]
        return StepOutcome(next, sqrt(DenseDoubleOps.dot(s, s)), reductions, usableCurvature)
    }

    /** Nocedal's two-loop recursion: the L-BFGS descent direction `-H∇f` without forming `H`. */
    private fun twoLoopRecursion(gradient: DoubleArray, n: Int): DoubleArray {
        val q = DoubleArray(n) { -gradient[it] }
        val k = sHistory.size
        val alphas = DoubleArray(k)
        for (index in k - 1 downTo 0) {
            val s = sHistory[index]
            val y = yHistory[index]
            val sq = DenseDoubleOps.dot(s, q)
            val a = rhoHistory[index] * sq
            alphas[index] = a
            DenseDoubleOps.combine(q, q, -a, y)
        }
        if (k > 0) {
            val s = sHistory[k - 1]
            val y = yHistory[k - 1]
            val sy = DenseDoubleOps.dot(s, y)
            val yy = DenseDoubleOps.dot(y, y)
            val gamma = if (yy > 0.0) sy / yy else 1.0
            DenseDoubleOps.scale(q, gamma)
        }
        for (index in 0 until k) {
            val s = sHistory[index]
            val y = yHistory[index]
            val yq = DenseDoubleOps.dot(y, q)
            val beta = rhoHistory[index] * yq
            val diff = alphas[index] - beta
            DenseDoubleOps.combine(q, q, diff, s)
        }
        return q
    }
}