Skip to content

Welcome to SwarmPy

What is SwarmPy?

SwarmPy is a Python library that implements swarm-based gradient descent for non-convex optimization introduced in arXiv:2211.17157. It is a lightweight library that is easy to use and extend. It can be used with NumPy or PyTorch.

Animation of SBGD on the Rastrigin function

Why SwarmPy?

The primary purpose of this library is to reproduce the results of arXiv:2211.17157. However, it is also designed to be a useful tool for researchers and practitioners who want to test swarm-based gradient descent or use it in their problems.

Swarm-Based Gradient Descent

Swarm-based gradient descent (SBGD) is a multi-agent, non-convex optimization method. The swarm consists of agents, each with position, \(x\), and mass, \(m\). They search for the global optimum using gradient descent with step size, \(h(x,m)\), that is adjusted to their relative mass. Heavier agents proceed with small steps in the direction of the local gradient, while lighter agents take larger steps based on the backtracking protocol implemented in line_search. The agents transfer mass from higher agents to lower ones using mass_transition. Accordingly, the crowd of agents is dynamically divided between heavier leaders gradually approaching local minima with small steps and lighter explorers surveying the loss landscape with large steps.

As opposed to other multi-agent methods, the agents in SBGD do not try to find a common minimum. Instead, light explorers survey the landscape independently, thereby increasing their chances of encountering improved positions for the swarm, upon which they assume the role of heavy leaders. At every step, the agent with the worst position is removed from the swarm. This dynamic process facilitated by communication via masses and step sizes allows the swarm to escape local minima.

Installation

It is recommended to install SwarmPy in a virtual environment. You can use virtualenv or conda. Create a virtual environment and activate it. With virtualenv:

virtualenv -p python3 venv
source venv/bin/activate

or with conda:

conda create -n swarmpy python=3.10
conda activate swarmpy

Then install the requirements using pip

pip install -r requirements.txt

Clone the repository

git clone git@github.com:anilzen/swarmpy.git

You can use the library by adding the path to the repository to your PYTHONPATH environment variable

export PYTHONPATH=$PYTHONPATH:/path/to/swarmpy

and importing the library in your code.

import swarmpy

Requirements

SwarmPy is currently tested on 3.10.6. It can be used with NumPy or PyTorch.