Repository navigation
feat(optimization): add stateful SGD/momentum and Adam step optimisers #343
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Merged
Merged
Changes from all commits
Commits
Show all changes
2 commits
Select commit
Hold shift + click to select a range
File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
There are no files selected for viewing
|
gabrielfrasantos marked this conversation as resolved.
|
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,117 @@ | ||
| # Stateful Step Optimisers (SGD / Adam) | ||
|
|
||
| ## Overview & Motivation | ||
|
|
||
| On-device neural network training and online system identification require an optimiser that accepts one externally computed gradient per call and maintains its own state between calls. Unlike batch optimisers that own the full objective, a **step optimiser** separates the gradient source from the update rule, enabling training loops where back-propagation runs elsewhere. | ||
|
|
||
| Two industry-standard rules are provided: **Stochastic Gradient Descent** (SGD) with optional momentum or Nesterov lookahead, and the **Adam** (Adaptive Moment Estimation) optimiser with bias correction. | ||
|
|
||
| ## Mathematical Theory | ||
|
|
||
| ### SGD | ||
|
|
||
| Vanilla SGD applies the raw gradient: | ||
|
|
||
| $$\theta_{t+1} = \theta_t - \eta \, g_t$$ | ||
|
|
||
| where $g_t = \nabla_\theta \mathcal{L}(\theta_t)$ and $\eta$ is the learning rate. | ||
|
|
||
| **Momentum** accumulates a velocity $v$, smoothing oscillations and accelerating progress in low-curvature directions: | ||
|
|
||
| $$v_{t+1} = \beta \, v_t + g_t, \qquad \theta_{t+1} = \theta_t - \eta \, v_{t+1}$$ | ||
|
|
||
| with $\beta \in [0, 1)$ the momentum coefficient (typically $0.9$). | ||
|
|
||
| **Nesterov momentum** incorporates a lookahead correction, replacing the plain velocity step with: | ||
|
|
||
| $$\theta_{t+1} = \theta_t - \eta \, \bigl(g_t + \beta \, v_{t+1}\bigr)$$ | ||
|
|
||
| This form evaluates the effective update at the anticipated next position, yielding faster convergence on smooth convex objectives. | ||
|
|
||
| ### Adam | ||
|
|
||
| Adam maintains exponential moving averages of the gradient (first moment $m$) and the squared gradient (second moment $v$): | ||
|
|
||
| $$m_{t+1} = \beta_1 \, m_t + (1 - \beta_1) \, g_t$$ | ||
|
|
||
| $$v_{t+1} = \beta_2 \, v_t + (1 - \beta_2) \, g_t^2 \quad (\text{element-wise})$$ | ||
|
|
||
| Both estimates are biased toward zero at initialisation. **Bias correction** removes this bias: | ||
|
|
||
| $$\hat{m}_{t+1} = \frac{m_{t+1}}{1 - \beta_1^{t+1}}, \qquad \hat{v}_{t+1} = \frac{v_{t+1}}{1 - \beta_2^{t+1}}$$ | ||
|
|
||
| The parameter update normalises the corrected first moment by the square root of the corrected second moment, providing a per-parameter adaptive step size: | ||
|
|
||
| $$\theta_{t+1} = \theta_t - \eta \, \frac{\hat{m}_{t+1}}{\sqrt{\hat{v}_{t+1}} + \varepsilon}$$ | ||
|
|
||
| Typical hyper-parameters: $\eta = 10^{-3}$, $\beta_1 = 0.9$, $\beta_2 = 0.999$, $\varepsilon = 10^{-8}$. | ||
|
|
||
| ## Complexity Analysis | ||
|
|
||
| | Algorithm | Time per step | Extra state (floats) | Notes | | ||
| |-----------|---------------|----------------------|---------------------------------------------------------------------| | ||
| | SGD | $O(N)$ | $N$ | One velocity vector (equals the gradient when momentum is 0) | | ||
| | Adam | $O(N)$ | $2N + 2$ | Two moment vectors plus the running powers $\beta_1^t$, $\beta_2^t$ | | ||
|
|
||
| $N$ is the number of parameters. All operations are in-place; no heap allocation is required. | ||
|
|
||
| ## Step-by-Step Walkthrough | ||
|
|
||
| **SGD with momentum** on $\mathcal{L}(\theta) = \frac{1}{2}\|\theta\|^2$, $N=1$, $\theta_0 = 1$, $\eta = 0.1$, $\beta = 0.9$: | ||
|
|
||
| | $t$ | $g_t = \theta_t$ | $v_t = 0.9 v_{t-1} + g_t$ | $\theta_{t+1} = \theta_t - 0.1 v_t$ | | ||
| |-----|------------------|---------------------------|-------------------------------------| | ||
| | 1 | 1.000 | 1.000 | 0.900 | | ||
| | 2 | 0.900 | 1.800 | 0.720 | | ||
| | 3 | 0.720 | 2.340 | 0.486 | | ||
|
|
||
| **Adam** on the same objective, $\beta_1 = 0.9$, $\beta_2 = 0.999$, $\varepsilon = 10^{-8}$, $\theta_0 = 0$, $g_1 = 1$: | ||
|
|
||
| | Quantity | Value | | ||
| |-------------|------------------------| | ||
| | $m_1$ | $0.1$ | | ||
| | $v_1$ | $0.001$ | | ||
| | $\hat{m}_1$ | $1.0$ | | ||
| | $\hat{v}_1$ | $1.0$ | | ||
| | $\theta_1$ | $-\eta \approx -0.001$ | | ||
|
|
||
| Bias correction is the critical step: without it, $m_1/\sqrt{v_1} \approx 3.16$, giving a first step roughly $\sqrt{1000}$ times larger than the bias-corrected value. | ||
|
|
||
| ## Pitfalls & Edge Cases | ||
|
|
||
| - **Learning rate too large.** SGD without momentum diverges for $\eta \geq 2/L$ on $L$-smooth losses. Momentum reduces the effective stability bound further; reduce $\eta$ or $\beta$ if oscillation is observed. | ||
| - **Adam with small $\varepsilon$.** Setting $\varepsilon$ too small causes division by near-zero when a parameter has zero gradient history, producing numerical instability. The default $10^{-8}$ is sufficient for `float`. | ||
| - **Momentum at reset.** When `Reset()` is called, the velocity (SGD) or moment estimates (Adam) return to zero. The first step after a reset behaves identically to starting from scratch, so the bias-correction denominator for Adam also restarts from $t=1$. | ||
| - **Nesterov with large $\beta$.** The lookahead correction adds $\beta \, v_{t+1}$ to the update, which can overshoot on sparse or noisy gradients. Prefer plain momentum when the gradient signal is noisy. | ||
| - **Float precision.** For Adam, the accumulated second moment $v$ can underflow toward zero for very small gradients and single-precision arithmetic. Increasing $\varepsilon$ mitigates this at the cost of less adaptivity. | ||
|
|
||
| ## Variants & Generalizations | ||
|
|
||
| | Variant | Change from base | | ||
| |---------|------------------------------------------------------------------------------------------------| | ||
| | AdaGrad | Non-decaying sum of squared gradients (no $\beta_2$ decay); aggressive learning rate shrinkage | | ||
| | RMSProp | Adam without first-moment tracking; lacks bias correction | | ||
| | AdamW | Decoupled weight-decay applied directly to $\theta$ before the gradient step | | ||
| | AMSGrad | Replaces $\hat{v}$ with the running maximum to guarantee monotone effective step-size | | ||
|
|
||
| ## Applications | ||
|
|
||
| - **On-device neural network training** — Updating weights after each batch of sensor data. | ||
| - **Online system identification** — Fitting model parameters in real time as measurements arrive. | ||
| - **Adaptive control** — Adjusting gain schedules or feed-forward maps without offline re-training. | ||
| - **Sensor calibration** — Minimising residual error by incrementally fitting a polynomial or affine model. | ||
|
|
||
| ## Connections to Other Algorithms | ||
|
|
||
| | Component | Relationship | | ||
| |------------------------------------------------|------------------------------------------------------------------------------------------------| | ||
| | [Gradient Descent](Optimizer.md) | Batch counterpart; shares the learning-rate update rule but recomputes the objective each time | | ||
| | [LMS Adaptive Filter](../estimators/README.md) | Equivalent to online SGD for a linear regression model under MSE loss | | ||
| | [Regularization](../regularization/README.md) | Adds a penalty gradient to $g_t$; compatible with any step optimiser | | ||
|
|
||
| ## References & Further Reading | ||
|
|
||
| - Ruder, S., "An overview of gradient descent optimization algorithms", *arXiv:1609.04747*, 2016. | ||
| - Kingma, D.P. and Ba, J., "Adam: A Method for Stochastic Optimization", *ICLR*, 2015. | ||
| - Nesterov, Y., "A method for solving the convex programming problem with convergence rate $O(1/k^2)$", *Soviet Mathematics Doklady*, 1983. | ||
| - Goodfellow, I., Bengio, Y., and Courville, A., *Deep Learning*, Chapter 8, MIT Press, 2016. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,6 @@ | ||
| #include "numerical/optimization/Adam.hpp" | ||
|
|
||
| namespace optimization | ||
| { | ||
| template class Adam<float, 2>; | ||
| } |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,92 @@ | ||
| #pragma once | ||
|
|
||
| #if defined(__GNUC__) && !defined(__clang__) | ||
| #pragma GCC push_options | ||
| #pragma GCC optimize("O3", "fast-math") | ||
| #endif | ||
|
|
||
| #include "numerical/math/CompilerOptimizations.hpp" | ||
| #include "numerical/math/Math.hpp" | ||
| #include "numerical/optimization/StepOptimizer.hpp" | ||
|
|
||
| namespace optimization | ||
| { | ||
| template<typename T, std::size_t N> | ||
| class Adam | ||
| : public StepOptimizer<T, N> | ||
| { | ||
| static_assert(std::is_floating_point_v<T>, "Adam supports floating-point types only"); | ||
|
|
||
| public: | ||
| using Vector = typename StepOptimizer<T, N>::Vector; | ||
|
|
||
| struct Parameters | ||
| { | ||
| T learningRate; | ||
| T beta1{ T{ 0.9 } }; | ||
| T beta2{ T{ 0.999 } }; | ||
| T epsilon{ T{ 1e-8 } }; | ||
| }; | ||
|
|
||
| explicit Adam(const Parameters& params); | ||
|
|
||
| void Step(Vector& theta, const Vector& gradient) override; | ||
| void Reset() override; | ||
|
|
||
| private: | ||
| Parameters parameters; | ||
| Vector firstMoment{}; | ||
| Vector secondMoment{}; | ||
| T beta1Power{ T{ 1 } }; | ||
| T beta2Power{ T{ 1 } }; | ||
| }; | ||
|
|
||
| template<typename T, std::size_t N> | ||
| Adam<T, N>::Adam(const Parameters& params) | ||
| : parameters{ params } | ||
| { | ||
| really_assert(params.learningRate > T{ 0 }); | ||
| really_assert(params.beta1 >= T{ 0 } && params.beta1 < T{ 1 }); | ||
| really_assert(params.beta2 >= T{ 0 } && params.beta2 < T{ 1 }); | ||
| really_assert(params.epsilon > T{ 0 }); | ||
| } | ||
|
|
||
| template<typename T, std::size_t N> | ||
| OPTIMIZE_FOR_SPEED void Adam<T, N>::Step(Vector& theta, const Vector& gradient) | ||
| { | ||
| beta1Power *= parameters.beta1; | ||
| beta2Power *= parameters.beta2; | ||
|
|
||
| firstMoment = firstMoment * parameters.beta1 + gradient * (T{ 1 } - parameters.beta1); | ||
|
|
||
| for (std::size_t i = 0; i < N; ++i) | ||
| secondMoment[i] = parameters.beta2 * secondMoment[i] + (T{ 1 } - parameters.beta2) * gradient[i] * gradient[i]; | ||
|
|
||
| const T beta1Correction = T{ 1 } - beta1Power; | ||
| const T beta2Correction = T{ 1 } - beta2Power; | ||
|
|
||
| for (std::size_t i = 0; i < N; ++i) | ||
| { | ||
| const T mHat = firstMoment[i] / beta1Correction; | ||
| const T vHat = secondMoment[i] / beta2Correction; | ||
| theta[i] -= parameters.learningRate * mHat / (math::Sqrt(vHat) + parameters.epsilon); | ||
| } | ||
| } | ||
|
|
||
| template<typename T, std::size_t N> | ||
| void Adam<T, N>::Reset() | ||
| { | ||
| firstMoment = Vector{}; | ||
| secondMoment = Vector{}; | ||
| beta1Power = T{ 1 }; | ||
| beta2Power = T{ 1 }; | ||
| } | ||
|
|
||
| #ifdef NUMERICAL_TOOLBOX_COVERAGE_BUILD | ||
| extern template class Adam<float, 2>; | ||
| #endif | ||
| } | ||
|
|
||
| #if defined(__GNUC__) && !defined(__clang__) | ||
| #pragma GCC pop_options | ||
| #endif |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,6 @@ | ||
| #include "numerical/optimization/Sgd.hpp" | ||
|
|
||
| namespace optimization | ||
| { | ||
| template class Sgd<float, 2>; | ||
| } |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,70 @@ | ||
| #pragma once | ||
|
|
||
| #if defined(__GNUC__) && !defined(__clang__) | ||
| #pragma GCC push_options | ||
| #pragma GCC optimize("O3", "fast-math") | ||
| #endif | ||
|
|
||
| #include "numerical/math/CompilerOptimizations.hpp" | ||
| #include "numerical/optimization/StepOptimizer.hpp" | ||
|
|
||
| namespace optimization | ||
| { | ||
| template<typename T, std::size_t N> | ||
| class Sgd | ||
| : public StepOptimizer<T, N> | ||
| { | ||
| static_assert(std::is_floating_point_v<T>, "Sgd supports floating-point types only"); | ||
|
|
||
| public: | ||
| using Vector = typename StepOptimizer<T, N>::Vector; | ||
|
|
||
| struct Parameters | ||
| { | ||
| T learningRate; | ||
| T momentum{ T{ 0 } }; | ||
| bool nesterov{ false }; | ||
| }; | ||
|
|
||
| explicit Sgd(const Parameters& params); | ||
|
|
||
| void Step(Vector& theta, const Vector& gradient) override; | ||
| void Reset() override; | ||
|
|
||
| private: | ||
| Parameters parameters; | ||
| Vector velocity{}; | ||
| }; | ||
|
|
||
| template<typename T, std::size_t N> | ||
| Sgd<T, N>::Sgd(const Parameters& params) | ||
| : parameters{ params } | ||
| { | ||
| really_assert(params.learningRate > T{ 0 }); | ||
| really_assert(params.momentum >= T{ 0 } && params.momentum < T{ 1 }); | ||
| } | ||
|
|
||
| template<typename T, std::size_t N> | ||
| OPTIMIZE_FOR_SPEED void Sgd<T, N>::Step(Vector& theta, const Vector& gradient) | ||
| { | ||
| velocity = velocity * parameters.momentum + gradient; | ||
| if (parameters.nesterov) | ||
| theta = theta - (gradient + velocity * parameters.momentum) * parameters.learningRate; | ||
| else | ||
| theta = theta - velocity * parameters.learningRate; | ||
| } | ||
|
|
||
| template<typename T, std::size_t N> | ||
| void Sgd<T, N>::Reset() | ||
| { | ||
| velocity = Vector{}; | ||
| } | ||
|
|
||
| #ifdef NUMERICAL_TOOLBOX_COVERAGE_BUILD | ||
| extern template class Sgd<float, 2>; | ||
| #endif | ||
| } | ||
|
|
||
| #if defined(__GNUC__) && !defined(__clang__) | ||
| #pragma GCC pop_options | ||
| #endif |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,20 @@ | ||
| #pragma once | ||
|
|
||
| #include "numerical/math/Matrix.hpp" | ||
|
|
||
| namespace optimization | ||
| { | ||
| template<typename T, std::size_t N> | ||
| class StepOptimizer | ||
| { | ||
| static_assert(std::is_floating_point_v<T>, "StepOptimizer supports floating-point types only"); | ||
|
|
||
| public: | ||
| virtual ~StepOptimizer() = default; | ||
|
|
||
| using Vector = math::Vector<T, N>; | ||
|
|
||
| virtual void Step(Vector& theta, const Vector& gradient) = 0; | ||
| virtual void Reset() = 0; | ||
| }; | ||
| } |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Uh oh!
There was an error while loading. Please reload this page.