|
| 1 | +""" |
| 2 | +RMSprop (Root Mean Square Propagation) optimizer implementation. |
| 3 | +
|
| 4 | +RMSprop is an adaptive learning rate optimizer that maintains a moving |
| 5 | +average of squared gradients to normalize the gradient. It was proposed |
| 6 | +by Geoffrey Hinton in his Coursera course on Neural Networks. |
| 7 | +
|
| 8 | +Key idea: Instead of using a fixed learning rate, RMSprop adapts the |
| 9 | +learning rate for each parameter by dividing by a running average of |
| 10 | +recent gradient magnitudes. |
| 11 | +
|
| 12 | +Update rules: |
| 13 | + v(t) = rho * v(t-1) + (1 - rho) * gradient^2 |
| 14 | + param = param - (learning_rate / sqrt(v(t) + epsilon)) * gradient |
| 15 | +
|
| 16 | +Where: |
| 17 | + v(t) = moving average of squared gradients |
| 18 | + rho = decay factor (typically 0.9) |
| 19 | + learning_rate = step size |
| 20 | + epsilon = small value to avoid division by zero |
| 21 | +
|
| 22 | +Reference: https://en.wikipedia.org/wiki/Stochastic_gradient_descent#RMSProp |
| 23 | +
|
| 24 | +>>> rmsprop([0.0], [0.1], 0.01) # doctest: +ELLIPSIS |
| 25 | +[-0.031...] |
| 26 | +>>> rmsprop([1.0, -1.0], [0.5, -0.5], 0.01) # doctest: +ELLIPSIS |
| 27 | +[0.968..., -0.968...] |
| 28 | +""" |
| 29 | + |
| 30 | +import math |
| 31 | + |
| 32 | + |
| 33 | +def rmsprop( |
| 34 | + params: list[float], |
| 35 | + gradients: list[float], |
| 36 | + learning_rate: float, |
| 37 | + rho: float = 0.9, |
| 38 | + epsilon: float = 1e-8, |
| 39 | + moving_avg: list[float] | None = None, |
| 40 | +) -> list[float]: |
| 41 | + """ |
| 42 | + Perform one step of the RMSprop optimization algorithm. |
| 43 | +
|
| 44 | + :param params: Current parameter values to be updated. |
| 45 | + :param gradients: Gradients of the loss with respect to each parameter. |
| 46 | + :param learning_rate: Step size for the update (must be positive). |
| 47 | + :param rho: Decay factor for the moving average (default 0.9). |
| 48 | + :param epsilon: Small constant to avoid division by zero (default 1e-8). |
| 49 | + :param moving_avg: Running average of squared gradients. Updated in |
| 50 | + place each call. Initialized to zeros if not provided. |
| 51 | + :return: Updated parameter values after one RMSprop step. |
| 52 | +
|
| 53 | + :raises ValueError: If params and gradients have different lengths. |
| 54 | + :raises ValueError: If learning_rate, rho, or epsilon are out of range. |
| 55 | +
|
| 56 | + >>> rmsprop([0.0], [0.0], 0.01) |
| 57 | + [0.0] |
| 58 | + >>> rmsprop([1.0], [0.0], 0.01) |
| 59 | + [1.0] |
| 60 | + >>> len(rmsprop([1.0, 2.0, 3.0], [0.1, 0.2, 0.3], 0.01)) == 3 |
| 61 | + True |
| 62 | + >>> rmsprop([1.0], [0.5], learning_rate=0.01) # doctest: +ELLIPSIS |
| 63 | + [0.968...] |
| 64 | + """ |
| 65 | + if len(params) != len(gradients): |
| 66 | + msg = ( |
| 67 | + f"params and gradients must have the same length, " |
| 68 | + f"got {len(params)} and {len(gradients)}" |
| 69 | + ) |
| 70 | + raise ValueError(msg) |
| 71 | + if learning_rate <= 0: |
| 72 | + msg = f"learning_rate must be positive, got {learning_rate}" |
| 73 | + raise ValueError(msg) |
| 74 | + if not 0.0 <= rho < 1.0: |
| 75 | + msg = f"rho must be in [0, 1), got {rho}" |
| 76 | + raise ValueError(msg) |
| 77 | + if epsilon <= 0: |
| 78 | + msg = f"epsilon must be positive, got {epsilon}" |
| 79 | + raise ValueError(msg) |
| 80 | + |
| 81 | + # Initialize moving average of squared gradients to zeros |
| 82 | + if moving_avg is None: |
| 83 | + moving_avg = [0.0] * len(params) |
| 84 | + |
| 85 | + updated_params = [] |
| 86 | + for i, (param, grad) in enumerate(zip(params, gradients)): |
| 87 | + # Update moving average IN PLACE so state persists across calls |
| 88 | + moving_avg[i] = rho * moving_avg[i] + (1 - rho) * grad**2 |
| 89 | + # Compute adaptive update |
| 90 | + param = param - (learning_rate / math.sqrt(moving_avg[i] + epsilon)) * grad |
| 91 | + updated_params.append(param) |
| 92 | + |
| 93 | + return updated_params |
| 94 | + |
| 95 | + |
| 96 | +if __name__ == "__main__": |
| 97 | + import doctest |
| 98 | + |
| 99 | + doctest.testmod() |
| 100 | + |
| 101 | + print("RMSprop Optimizer Demo") |
| 102 | + print("=" * 40) |
| 103 | + print("Minimizing f(x) = x^2 (minimum at x = 0)") |
| 104 | + print("Gradient: f'(x) = 2x | Learning rate: 0.1\n") |
| 105 | + |
| 106 | + param = [5.0] |
| 107 | + moving_avg = [0.0] |
| 108 | + |
| 109 | + print(f"{'Step':>6} | {'param':>12} | {'f(x)':>12}") |
| 110 | + print("-" * 38) |
| 111 | + print(f"{'0':>6} | {param[0]:>12.6f} | {param[0] ** 2:>12.6f}") |
| 112 | + |
| 113 | + for step in range(1, 101): |
| 114 | + gradient = [2 * param[0]] |
| 115 | + param = rmsprop(param, gradient, learning_rate=0.1, moving_avg=moving_avg) |
| 116 | + if step % 20 == 0: |
| 117 | + print(f"{step:>6} | {param[0]:>12.6f} | {param[0] ** 2:>12.8f}") |
| 118 | + |
| 119 | + print(f"\nConverged to x = {param[0]:.8f} (true minimum = 0.0)") |
0 commit comments