Perturbed Saddle-escape Descent (PSD)

Project Summary

This repository implements the Perturbed Saddle-escape Descent (PSD)
algorithm for escaping saddle points in non-convex optimisation problems, as described in Alpay and Alakkad (2025).
It contains reference NumPy implementations, framework specific optimisers
for PyTorch and TensorFlow, and utilities for reproducing the synthetic
experiments reported in the accompanying manuscript.

Features

• Reference implementations of PSD, PSD-Probe and baseline gradient descent
  variants in pure NumPy.
• Suite of analytic test functions with gradients and Hessians.
• Synthetic data generator producing the tables and figures used in the
  paper (experiments.py).
• Framework specific optimisers: PSDTorch, PSDTensorFlow and a
  PSDOptimizer/PerturbedAdam package for PyTorch.
• Example training scripts for MNIST and CIFAR-10.

Technology Stack

The core project depends on the following libraries:

Python 3.8 or later is required.

Installation

Install the published optimiser package:

pip install psd-optimizer

Or install the repository in editable mode for development:

git clone https://github.com/farukalpay/PSD.git
cd PSD
pip install -e ".[dev]"

Quick Start

import numpy as np
from psd import algorithms, functions

x0 = np.array([1.0, -1.0])
x_star, _ = algorithms.gradient_descent(x0, functions.SEPARABLE_QUARTIC.grad)

Further examples are available in the examples/ directory and the
documentation.

Usage

Using the Reference Algorithms

The core PSD routines and test functions can be imported from the
psd package:

import numpy as np
from psd import algorithms, functions

x0 = np.array([1.0, -1.0])
x_star, _ = algorithms.gradient_descent(x0, functions.SEPARABLE_QUARTIC.grad)

This structure allows you to experiment with the reference NumPy
implementations directly in your projects.

The PyTorch optimisers PSDOptimizer and PerturbedAdam are also
available directly via from psd import ....

All-in-One "Monster" Interface

For rapid experimentation without navigating submodules, import the aggregated
psd.monster module. It re-exports the core algorithms, analytic test
functions and framework-specific optimisers in a single namespace:

import numpy as np
from psd import monster

x0 = np.array([1.0, -1.0])
x_star, _ = monster.gradient_descent(x0, monster.SEPARABLE_QUARTIC.grad)

This unified view aims to be approachable for both humans and language models
exploring the project.

Generating Synthetic Data

python experiments.py

The command writes CSV summaries to results/ and training curves to
data/.

Performance

Profiling identified rosenbrock_hess as a hot path when computing the
Rosenbrock Hessian. Vectorising the computation removed explicit
Python loops and yielded the following improvements (dimension 1000):

Benchmarking is automated via pytest-benchmark using a fixed NumPy seed.
Hard time and memory thresholds guard against major regressions.

Training with the PyTorch Optimiser

from psd_optimizer import PSDOptimizer

model = ...
opt = PSDOptimizer(model.parameters(), lr=1e-3)

def closure():
    opt.zero_grad()
    output = model(x)
    loss = criterion(output, y)
    loss.backward()
    return loss

opt.step(closure)

Example scripts using this API are available in the notebooks/
directory.

Training a Small Language Model

An illustrative example for fine-tuning a compact transformer with
PSDOptimizer is provided in scripts/train_small_language_model.py.
The script downloads a tiny GPT-style model from the Hugging Face Hub and
optimizes it on a short dummy corpus.

Run the example with default settings:

python scripts/train_small_language_model.py

Specify a different pretrained model and number of epochs:

python scripts/train_small_language_model.py --model distilgpt2 --epochs 5

Documentation

Full API documentation and guides are available in the
docs/ directory.
Additional materials include:

• notebooks/10_minute_start.ipynb – an interactive notebook showcasing the optimiser.
• docs/section_1_5_extension.md – theoretical notes on extending PSD to stochastic settings.
• notebooks/navigation.ipynb – links to all example notebooks including advanced_usage.ipynb.

Testing

After installing the repository in editable mode, run the test suite to
verify that everything works:

pytest

The current suite is small but helps prevent regressions.

Repository Structure

psd/              # Reference implementations and framework-specific optimisers
    algorithms.py # PSD and baseline algorithms
    functions.py  # Analytic test functions and registry
psd_optimizer/    # PyTorch optimiser package
experiments.py    # Synthetic data generation

Contributing

Contributions are welcome! Please open an issue or pull request on GitHub
and see CONTRIBUTING.md for guidelines. By participating you agree to
abide by the CODE_OF_CONDUCT.md.

Citation

If you use PSD in your research, please cite the following:

@misc{alpay2025escapingsaddlepointscurvaturecalibrated,
      title={Escaping Saddle Points via Curvature-Calibrated Perturbations: A Complete Analysis with Explicit Constants and Empirical Validation},
      author={Faruk Alpay and Hamdi Alakkad},
      year={2025},
      eprint={2508.16540},
      archivePrefix={arXiv},
      primaryClass={cs.LG},
      url={https://arxiv.org/abs/2508.16540},
}

License

This project is released under the MIT License. See LICENSE for details.