HAMUX
A new class of deep learning library built around ENERGY
bbhamuxis a minimal version ofhamuxto construct Hierarchical Associative Memories (i.e., super-powered Hopfield Networks). It represents research code that is designed to be easy to use, easy to hack.
Getting started
This documentation walks through how to use the package as if installing from pip, but the main logic for this repo lives in a single .py file (src/bbhamux/bbhamux.py). See the original README for instructions on how to just copy and use the important file.
pip install bbhamuxHAMUX is a universal abstraction for Hopfield Networks
HAMUX fully captures the the energy fundamentals of Hopfield Networks and enables anyone to:
🧠 Build DEEP Hopfield nets
🧱 With modular ENERGY components
🏆 That resemble modern DL operations
Every architecture built using HAMUX is a dynamical system guaranteed to have a tractable energy function that converges to a fixed point. Our deep Hierarchical Associative Memories (HAMs) have several additional advantages over traditional Hopfield Networks (HNs):
| Hopfield Networks (HNs) | Hierarchical Associative Memories (HAMs) |
|---|---|
| HNs are only two layers systems | HAMs connect any number of layers |
| HNs model only simple relationships between layers | HAMs model any complex but differentiable operation (e.g., convolutions, pooling, attention, \(\ldots\)) |
| HNs use only pairwise synapses | HAMs use many-body synapses (which we denote HyperSynapses) |
How does HAMUX work?
HAMUX is a hypergraph` of 🌀neurons connected via 🤝hypersynapses, an abstraction sufficiently general to model the complexity of connections used in modern AI architectures.
HAMUX defines two fundamental building blocks of energy: the 🌀neuron layer and the 🤝hypersynapse (an abstraction of a pairwise synapse to include many-body interactions) connected via a hypergraph. It is a fully dynamical system, where the “hidden state” \(x_i^l\) of each layer \(l\) (blue squares in the figure below) is an independent variable that evolves over time. The update rule of each layer is entirely local; only signals from a layer’s connected synapses (red circles in the figure below) can tell the hidden state how to change. This is shown in the following equation:
\[\tau \frac{d x_{i}^{l}}{dt} = -\frac{\partial E}{\partial g_i^l}\]
where \(g_i^l\) are the activations (i.e., non-linearities) on each neuron layer, described in the section on Neuron Layers. Concretely, we implement the above differential equation as the following discretized equation (where the bold \({\mathbf x}_l\) is the collection of all elements in layer \(l\)’s state):
\[ \mathbf{x}_l^{(t+1)} = \mathbf{x}_l^{(t)} - \frac{dt}{\tau} \nabla_{\mathbf{g}_l}E(t)\]
HAMUX handles all the complexity of scaling this fundamental update equation to many layers and hyper synapses. In addition, it provides a framework to:
- Implement your favorite Deep Learning operations as a HyperSynapse
- Port over your favorite activation functions as Lagrangians
- Connect your layers and hypersynapses into a HAM (using a hypergraph as the data structure)
- Inject your data into the associative memory
- Automatically calculate and descend the energy given the hidden states at any point in time
Use these features to train any hierarchical associative memory on your own data! All of this made possible by JAX.
The examples/ subdirectory contains a (growing) list of examples on how to apply HAMUX on real data.
Contributing to the docs
This README is automatically generated from docs_src/nbs/index.qmd do NOT edit it directly.
From the root of this project:
. scripts/activate.sh # Activate the virtual environment
cd docs_src/
nbdev_preview # Live preview docs site
nbdev_test # Test the code examples in the docs
nbdev_docs # Generate static buildMerge the doc changes into main or dev (see /.github/workflows/pages.yml) to deploy the docs site to the gh-pages branch.