An interactive simulation tool to analyze and compare first-order optimization algorithms (SGD, Momentum, Nesterov, AdaGrad, RMSprop, Adam) across non-convex and ill-conditioned loss surfaces.
Built with Python, NumPy, Plotly, and Streamlit.
Most machine learning libraries treat optimizers like Adam or RMSprop as plug-and-play black boxes. In practice, the local geometry of the loss surface (such as high condition number ravines, flat plateaus, or saddle points) dictates whether an algorithm converges, oscillates, or diverges.
I built this project to study and visualize:
- Ill-conditioned ravines: Why vanilla gradient descent oscillates across steep valley walls without making forward progress, and how momentum and adaptive learning rates fix this.
- Saddle points: How momentum-based methods build velocity to escape zero-gradient saddle points where standard SGD stalls.
- Adaptive learning rates: The practical differences between coordinate-wise accumulators (AdaGrad) and exponential moving averages (RMSprop, Adam).
- Benchmark Loss Landscapes: Analytical gradients for 7 non-convex test surfaces (Rosenbrock, Beale, Rastrigin, Saddle Point, Himmelblau, Ill-conditioned Quadratic Bowl, and Sphere).
- First-Order Optimizers: Implementations of SGD, Polyak Momentum, Nesterov Accelerated Gradient (NAG), AdaGrad, RMSprop, and Adam.
- Interactive Visualizations:
- 2D Contour Plot: High-resolution contour maps, gradient direction vectors, and colored trajectory paths.
- 3D Loss Landscape: Rotatable 3D surface mesh with trajectories plotted directly onto the surface.
- Convergence Curves: Log-scaled loss vs. step count and distance to known global minima.
- Multi-Optimizer Comparison: Run multiple algorithms simultaneously from identical starting coordinates to compare path efficiency and speed.
- CLI and Headless Runs: Run simulations from the terminal and export interactive HTML reports.
- Unit Tests: 27 unit tests validating analytical gradients against central finite difference approximations and checking optimizer convergence.
| Optimizer | Core Update Rule | Key Hyperparameters | Intuition |
|---|---|---|---|
| SGD |
|
Basic steepest descent. Oscillates in steep ravines. | |
| Momentum |
|
Heavy ball mechanics. Cancels orthogonal oscillations, builds speed down ravines. | |
| NAG |
|
Lookahead momentum. Anticipatory braking prevents overshooting narrow basins. | |
| AdaGrad |
|
Adapts learning rate per parameter. Can stall prematurely due to monotonic growth of |
|
| RMSprop |
|
Uses Exponential Moving Average (EMA) of squared gradients to keep step sizes responsive. | |
| Adam | $\hat{m}_t = \frac{m_t}{1 - \beta_1^t}, ; \hat{v}t = \frac{v_t}{1 - \beta_2^t}$ $\theta{t+1} = \theta_t - \frac{\alpha}{\sqrt{\hat{v}_t} + \epsilon} \hat{m}_t$ |
Combines momentum with adaptive learning rates and initialization bias correction. |
-
Rosenbrock ("Banana") Function:
$f(x, y) = (1 - x)^2 + 100(y - x^2)^2$ - Valley minimum at
$(1, 1)$ . Tests whether optimizers can follow a sharp curved ridge without diverging.
- Valley minimum at
-
Beale's Function:
$f(x, y) = (1.5 - x + xy)^2 + (2.25 - x + xy^2)^2 + (2.625 - x + xy^3)^2$ - Global minimum at
$(3, 0.5)$ . Tests navigation around sharp corners and flat plateaus with vanishing gradients.
- Global minimum at
-
Escapable Saddle Point:
$f(x, y) = x^2 - y^2 + \frac{1}{4}y^4$ - Saddle at
$(0, 0)$ , minima at$(0, \pm\sqrt{2})$ . Tests escape velocity along negative curvature directions.
- Saddle at
-
Ill-Conditioned Quadratic:
$f(x, y) = \frac{1}{2}(x^2 + 20y^2)$ - Condition number
$\kappa = 20$ . Demonstrates how anisotropic Hessian eigenvalues cause vanilla gradient descent to bounce between walls.
- Condition number
-
Himmelblau's Function:
$f(x, y) = (x^2 + y - 11)^2 + (x + y^2 - 7)^2$ - 4 identical global minima separated by ridges. Shows basin selection based on initialization.
-
Rastrigin Function:
$f(x, y) = 20 + (x^2 - 10\cos(2\pi x)) + (y^2 - 10\cos(2\pi y))$ - Multimodal landscape with many local minima. Demonstrates vulnerability of first-order methods to local traps.
git clone https://github.com/adem-temam/gradient-descent-visualizer.git
cd gradient-descent-visualizer
python3 -m venv .venv
source .venv/bin/activate
pip install -r requirements.txtstreamlit run app.pyOpen your browser at http://localhost:8501.
Compare algorithms directly from the terminal:
# Benchmark on Rosenbrock
python cli.py --surface rosenbrock --optimizers sgd momentum rmsprop adam --steps 120
# Export interactive HTML contour plot
python cli.py --surface quadratic_bowl --optimizers sgd momentum adam --export-html report.htmlRun the automated test suite with pytest:
pytest -vThe test suite covers:
- Gradient Accuracy: Central finite difference numerical verification across 5 points per surface.
- Global Minimum Checks: Verifies that analytical minima produce near-zero gradient and expected loss.
- Convergence Verification: Checks that each optimizer converges to within 5% of starting loss on standard convex objectives.
- Divergence Handling: Tests handling of extreme learning rates and invalid hyperparameters.
gradient-descent-visualizer/
├── src/
│ └── gd_visualizer/
│ ├── __init__.py
│ ├── functions.py # Loss landscapes, analytical gradients, domains, minima
│ ├── optimizers.py # SGD, Momentum, NAG, AdaGrad, RMSprop, Adam
│ ├── runner.py # Trajectory simulation engine & metrics collection
│ └── visualizer.py # Plotly 2D contour & 3D surface plot generators
├── tests/
│ ├── __init__.py
│ ├── test_functions.py # Finite difference gradient checks
│ ├── test_optimizers.py # Convergence tests on convex quadratics
│ ├── test_runner.py # Divergence and simulation checks
│ └── test_visualizer.py # Plotly figure construction tests
├── app.py # Streamlit interactive application
├── cli.py # CLI tool for headless benchmarks & HTML export
├── pyproject.toml # Standard modern packaging metadata
├── requirements.txt # Pinned dependencies
├── .gitignore # Standard Python gitignore
└── README.md # Project documentation