Package: hamux
HAMUXPart proof-of-concept, part functional prototype, HAMUX is designed to
bridge modern AI architectures and biologically plausible Hopfield
Networks.HAMUX: AHierarchicalAssociativeMemoryUser
eXperience🚧API is in rapid development. Remember to specify the version when building off of HAMUX.A Universal Abstraction for Hopfield NetworksHAMUX fully captures the the energy fundamentals of Hopfield Networks
and enables anyone to:🧠 BuildDEEPHopfield nets🧱 With modularENERGYcomponents🏆 That resemble modern DL operationsEveryarchitecture built using HAMUX is adynamical systemguaranteed to have atractable energyfunction thatconvergesto a
fixed point. Our deepHierarchical Associative
Memories(HAMs) have several
additional advantages over traditionalHopfield
Networks(HNs):Hopfield Networks (HNs)Hierarchical Associative Memories (HAMs)HNs are onlytwo layerssystemsHAMs connectany numberof layersHNs model onlylinear relationshipsbetween layersHAMs modelany differentiable operation(e.g., convolutions, pooling, attention, $\ldots$)HNs use onlypairwise synapsesHAMs usemany-body synapsesHow does HAMUX work?HAMUXis ahypergraphof
🌀neurons connected via 🤝synapses, an abstraction sufficiently
general to model the complexity of connections contained in the 🧠.HAMUX defines two fundamental building blocks of energy: the🌀neuron
layerand the🤝synapse, connected via ahypergraph🌀Neuron LayersNeuron layers are the recurrent unit of a HAM; that is, 🌀neurons keep a
state that changes over time according to the dynamics of the system.
These states always change to minimize the global energy function of the
system.For those of us familiar with traditional Deep Learning architectures,
we are familiar with nonlinear activation functions like theReLUandSoftMax. A neuron layer in HAMUX is exactly that: a nonlinear
activation function defined on some neuron. However, we need to express
the activation function as a convexLagrangian function$\mathcal{L}$ that is the integral of the desired non-linearity such
that thederivative of the Lagrangian function$\nabla \mathcal{L}$
is our desired non-linearity. E.g., consider the ReLU:$$
\begin{align*}
\mathcal{L}(x) &:= \frac{1}{2} (\max(x, 0))^2\
\nabla \mathcal{L} &= \max(x, 0) = \mathrm{relu}(x)\
\end{align*}
$$We need to define our activation layer in terms of theLagrangianof
the ReLU instead of the ReLU itself. Extending this constraint to other
nonlinearities makes it possible to define the scalar energy for any
neuron in a HAM. It turns out that many activation functions used in
today’s Deep Learning landscape are expressible as a Lagrangian. HAMUX
is “batteries-included” for many common activation functions includingrelus,softmaxes,sigmoids,LayerNorms, etc. See ourdocumentation on
Lagrangiansfor
examples on how to implement efficient activation functions from
Lagrangians in JAX. We show how to turn Lagrangians into usable energy
building blocks in ourdocumentation on neuron
layers.🤝SynapsesA 🤝synapse ONLY sees activations of connected 🌀neuron layers. Its one
job: report HIGH ⚡️energy if the connected activations are dissimilar
and LOW ⚡️energy when they are aligned. Synapses can resemble
convolutions, dense multiplications, even attention… Take a look at ourdocumentation on
synapses.🚨Point of confusion: modern AI frameworks haveAttentionLayers andConvolutionalLayers. In HAMUX, these would be more appropriately calledAttentionSynapses andConvolutionalSynapses.InstallFrom pip:pip install hamuxIf you are using accelerators beyond the CPU you will need to
additionally install the correspondingjaxandjaxlibversions
followingtheir
documentation. E.g.,pip install --upgrade "jax[cuda]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.htmlFrom source:After cloning:cd hamux
conda env create -f environment.yml
conda activate hamux
pip install --upgrade "jax[cuda]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html # If using GPU accelerator
pip install -e .
pip install -r requirements-dev.txt  # To run the examplesHow to Useimporthamuxashmximportjax.numpyasjnpimportjaximportjax.tree_utilasjtuWe can build a simple 4 layer HAM architecture using the following codelayers=[hmx.TanhLayer((32,32,3)),# e.g., CIFAR Imageshmx.SigmoidLayer((11,11,1000)),# CIFAR patcheshmx.SoftmaxLayer((10,)),# CIFAR Labelshmx.SoftmaxLayer((1000,)),# Hidden Memory Layer]synapses=[hmx.ConvSynapse((3,3),strides=3),hmx.DenseSynapse(),hmx.DenseSynapse(),]connections=[([0,1],0),([1,3],1),([2,3],2),]rng=jax.random.PRNGKey(0)param_key,state_key,rng=jax.random.split(rng,3)states,ham=hmx.HAM(layers,synapses,connections).init_states_and_params(param_key,state_key=state_key);Notice that we did not specify any output channel shapes in the
synapses. The desired output shape is computed from the layers connected
to each synapse duringhmx.HAM.init_states_and_params.We have two fundamental objects:statesandham. Thehamobject
contains the connectivity structure of the HAM (e.g.,
layer+synapse+hypergraph information) alongside theparametersof
the network. Thestatesobject is a list of lengthnlayerswhere
each item is a tensor representing the neuron states of the
corresponding layer.assertlen(states)==ham.n_layersassertall([state.shape==layer.shapeforstate,layerinzip(states,ham.layers)])We make it easy to run the dynamics of any HAM. Everyforwardfunction
is defined external to the memory and can be modified to extract
different memories from different layers, as desired. The general steps
for any forward function are:Initialize the dynamic statesInject an initial state into the systemRun dynamics, calculating energy gradient at every point in time.Return the layer state/activation of interestdeffwd(model,x,depth=15,dt=0.1):"""Assuming a trained HAM, run association with the HAM on batched inputs `x`"""# 1. Initialize model states at t=0. Account for batch sizexs=model.init_states(x.shape[0])# Inject initial statexs[0]=xenergies=[]foriinrange(depth):energies.append(model.venergy(xs))# If desired, observe the energydEdg=model.vdEdg(xs)# Calculate the gradientsxs=jtu.tree_map(lambdax,stepsize,grad:x-stepsize*grad,xs,model.alphas(dt),dEdg)# Return probabilities of our label layerprobs=model.layers[-2].activation(xs[-2])returnjnp.stack(energies),probsbatch_size=3x=jax.random.normal(jax.random.PRNGKey(2),(batch_size,32,32,3))energies,probs=fwd(ham,x,depth=20,dt=0.3)print(probs.shape)# batchsize, nclassesassertjnp.allclose(probs.sum(-1),1)(3, 10)The Energy Function vs the Loss FunctionWe use JAX’s autograd to descend the energy function of our system AND
the loss function of our task. The derivative of the energy is always
taken wrt to ourstates; the derivative of the loss function is always
taken wrt ourparameters. During training, we change our parameters to
optimize theLoss Function. During inference, we assume that
parameters are constant.Autograd for Descending EnergyEveryHAMdefines the
energy function for our system, which is everything we need to compute
memories of the system. Naively, we can calculate $\nabla_x E$: the
derivative of the energy function wrt thestatesof each layer:stepsize=0.01fscore_naive=jax.grad(ham.energy)next_states=jax.tree_util.tree_map(lambdastate,score:state-stepsize,states,fscore_naive(states))But it turns out we improve the efficiency of our network if we instead
take $\nabla_g E$: the derivative of the energy wrt. theactivationsinstead of thestates. They have the same local minima, even though
the trajectory to get there is different. Some nice terms cancel, and we
get:$$\nabla_g E_\text{HAM} = x + \nabla_g E_\text{synapse}$$stepsize=0.01deffscore_smart(xs):gs=ham.activations(xs)returnjax.tree_util.tree_map(lambdax,nabla_g_Esyn:x+nabla_g_Esyn,xs,jax.grad(ham.synapse_energy)(gs))next_states=jax.tree_util.tree_map(lambdastate,score:state-stepsize,states,fscore_smart(states))CitationWork is a collaboration between theMIT-IBM Watson AI
Laband thePoloClub@ GA Tech