nengo/keras-lmu

Keras implementation of Legendre Memory Units

KerasLMU: Recurrent neural networks using Legendre Memory Units

Paper

This is a Keras-based implementation of the Legendre Memory Unit (LMU). The LMU is a novel memory cell for recurrent neural networks that dynamically maintains information across long windows of time using relatively few resources. It has been shown to perform as well as standard LSTM or other RNN-based models in a variety of tasks, generally with fewer internal parameters (see this paper for more details). For the Permuted Sequential MNIST (psMNIST) task in particular, it has been demonstrated to outperform the current state-of-the-art results. See the note below for instructions on how to get access to this model.

The LMU is mathematically derived to orthogonalize its continuous-time history – doing so by solving d coupled ordinary differential equations (ODEs), whose phase space linearly maps onto sliding windows of time via the Legendre polynomials up to degree d − 1 (the example for d = 12 is shown below).

A single LMU cell expresses the following computational graph, which takes in an input signal, x, and couples a optimal linear memory, m, with a nonlinear hidden state, h. By default, this coupling is trained via backpropagation, while the dynamics of the memory remain fixed.

The discretized A and B matrices are initialized according to the LMU's mathematical derivation with respect to some chosen window length, θ. Backpropagation can be used to learn this time-scale, or fine-tune A and B, if necessary.

Both the kernels, W, and the encoders, e, are learned. Intuitively, the kernels learn to compute nonlinear functions across the memory, while the encoders learn to project the relevant information into the memory (see paper for details).

Note

The paper branch in the lmu GitHub repository includes a pre-trained Keras/TensorFlow model, located at models/psMNIST-standard.hdf5, which obtains a psMNIST result of 97.15%. Note that the network is using fewer internal state-variables and neurons than there are pixels in the input sequence. To reproduce the results from this paper, run the notebooks in the experiments directory within the paper branch.

Nengo Examples

• LMUs in Nengo (with online learning)
• Spiking LMUs in Nengo Loihi (with online learning)
• LMUs in NengoDL (reproducing SotA on psMNIST)

Citation

@inproceedings{voelker2019lmu,
  title={Legendre Memory Units: Continuous-Time Representation in Recurrent Neural Networks},
  author={Aaron R. Voelker and Ivana Kaji\'c and Chris Eliasmith},
  booktitle={Advances in Neural Information Processing Systems},
  pages={15544--15553},
  year={2019}
}

keras-lmu's Issues

Please expand time series prediction examples

LMU hasn't received the attention due it since it's introduction. Nor has Nengo's spiking approach received due attention. IMHO, ABR's priorities have gotten the spiking cart before the LMU horse. LMU's superiority to LSTM is a strength of that horse that may well draw greater attention to spiking deployments.

Since the LMU's primary strength is dynamical systems modeling, (and having casually worked with the LMU code for a couple of years off and on) one of the things that becomes apparent to me (as a casual user) is the need for more examples involving dynamics and, more specifically, time series prediction. The most obvious omission is the original repository's Mackey-Glass example. But even that example, although it does demonstrate superiority to LSTM (other than in the hybrid architecture), it doesn't really get to the heart of dynamical systems identification for which LMU is likely to really shine:

Online identification of nonstationary dynamical systems.

Something that would accomplish this is an algorithm generating multiple, dynamically interdependent, wave forms (generated in the CPU fed to the GPU(s)), with dependency parameters changing continuously in time for the LMU to learn, online, in the GPU/TPU, and predict.

Particular attention to illustrating the function of LMU-unique parameters (theta, etc.) -- especially in contrast to the LSTM in this environment -- would help the outreach a great deal.

PS: Something to avoid in this kind of outreach is reliance on interpolative test sets -- that is to say, avoid the normal Keras training/testing mode involving chopping up time series data into training and test sets where what the model actually learns to do is interpolate rather than extrapolate.

mackey-glass problem with LMU giving out the error with the recent updates of keras-lmu 0.2.0 and 0.3.0 versions

<ipython-input-16-f9340b4f7a97> in <module>()
     50 layers = 4
     51 lstm_model = make_lstm(units=25, layers=layers)
---> 52 lmu_model = make_lmu(units=49, layers=layers)
     53 hybrid_model = make_hybrid(units_lstm=25, units_lmu=40, layers=layers)

1 frames

<ipython-input-16-f9340b4f7a97> in delay_layer(units, **kwargs)
     18                        #memory_d=memory_d,
     19                        #hidden_cell=tf.keras.layers.SimpleRNNCell(212),
---> 20                        theta=4,
     21                       ),
     22                return_sequences=True,

TypeError: __init__() missing 2 required positional arguments: 'memory_d' and 'hidden_cell'
]

Here I passed the missing parameters..

def make_lstm(units, layers):
    model = Sequential()
    model.add(LSTM(units,
                   input_shape=(train_X.shape[1], 1),  # (timesteps, input_dims)
                   return_sequences=True))  # continuously outputs per timestep
    for _ in range(layers-1):
        model.add(LSTM(units, return_sequences=True))
    model.add(Dense(train_X.shape[-1], activation='tanh'))
    model.compile(loss="mse", optimizer="adam")
    model.summary()
    return model

def delay_layer(units, **kwargs):
    return RNN(LMUCell(units=units,
                       order=4,
                       memory_d=4,
                       hidden_cell=tf.keras.layers.Layer,
                       #memory_d=memory_d,
                       #hidden_cell=tf.keras.layers.SimpleRNNCell(212),
                       theta=4,
                      ),
               return_sequences=True,
               **kwargs)
def make_lmu(units, layers):
    model = Sequential()
    model.add(delay_layer(units,
                          input_shape=(train_X.shape[1], 1)))  # (timesteps, input_dims)
    for _ in range(layers-1):
        model.add(delay_layer(units))
    model.add(Dense(train_X.shape[-1], activation='linear'))
    model.compile(loss="mse", optimizer="adam")
    model.summary()
    
    return model


def make_hybrid(units_lstm, units_lmu, layers):
    assert layers == 4, "unsupported"
    model = Sequential()
    model.add(delay_layer(units=units_lmu,input_shape=(train_X.shape[1], 1)))
    model.add(LSTM(units=units_lstm, return_sequences=True))
    model.add(delay_layer(units=units_lmu))
    model.add(LSTM(units=units_lstm, return_sequences=True))
    model.add(Dense(train_X.shape[-1], activation='tanh'))
    model.compile(loss="mse", optimizer="adam")
    model.summary()
    return model


layers = 4
lstm_model = make_lstm(units=25, layers=layers)
lmu_model = make_lmu(units=49, layers=layers) 
hybrid_model = make_hybrid(units_lstm=25, units_lmu=40, layers=layers)

Even after passing out the missing parameters giving another error with units, even I installed the latest keras-lmu but same problem I think because version 0.2.0 removed the units and hidden_activation parameters of LMUCell (these are now specified directly in the hidden_cell. (#22) So please provide the updated codes for mackey-glass prediction problem also. Thanks in advance

TypeError                                 Traceback (most recent call last)

<ipython-input-17-1c880cdb37fc> in <module>()
     50 layers = 4
     51 lstm_model = make_lstm(units=25, layers=layers)
---> 52 lmu_model = make_lmu(units=49, layers=layers)
     53 hybrid_model = make_hybrid(units_lstm=25, units_lmu=40, layers=layers)

6 frames

/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/utils/generic_utils.py in validate_kwargs(kwargs, allowed_kwargs, error_message)
    776   for kwarg in kwargs:
    777     if kwarg not in allowed_kwargs:
--> 778       raise TypeError(error_message, kwarg)
    779 
    780 

TypeError: ('Keyword argument not understood:', 'units')

Opting in to supports_masking

https://www.tensorflow.org/guide/keras/masking_and_padding#opting-in_to_mask_propagation_on_compatible_layers:

Most layers don't modify the time dimension, so don't need to modify the current mask. However, they may still want to be able to propagate the current mask, unchanged, to the next layer. This is an opt-in behavior. By default, a custom layer will destroy the current mask (since the framework has no way to tell whether propagating the mask is safe to do).

If you have a custom layer that does not modify the time dimension, and if you want it to be able to propagate the current input mask, you should set self.supports_masking = True in the layer constructor. In this case, the default behavior of compute_mask() is to just pass the current mask through.

Here's an example of a layer that is whitelisted for mask propagation:

class MyActivation(keras.layers.Layer):
    def __init__(self, **kwargs):
        super(MyActivation, self).__init__(**kwargs)
        # Signal that the layer is safe for mask propagation
        self.supports_masking = True

    def call(self, inputs):
        return tf.nn.relu(inputs) 

I believe mask propagation is safe for the LMUCell, as it does not modify the time dimension, so we should be able to opt-in by setting self.supports_masking = True.

`LMUFFT` cannot be created if `hidden_cell` is `Dense`

Versions:

• keras_lmu==0.3.2.dev0
• tensorflow==2.4.1

Taking this example from the unit tests:

keras-lmu/keras_lmu/tests/test_layers.py

Lines 158 to 164 in ab07757

	out = layers.LMUFFT(
	1,
	2,
	3,
	tf.keras.layers.SimpleRNNCell(4),
	return_sequences=True,
	)(inp)

and modifying it as follows:

        out = layers.LMUFFT(
            1,
            2,
            3,
            tf.keras.layers.Dense(4),  # tf.keras.layers.SimpleRNNCell(4),
            return_sequences=True,
        )(inp)

results in the error:

ValueError: in user code:

    /home/arvoelke/git/keras-lmu/keras_lmu/layers.py:660 call  *
        h = tf.keras.layers.TimeDistributed(self.hidden_cell)(
    /home/arvoelke/anaconda3/envs/*/lib/python3.8/site-packages/tensorflow/python/keras/engine/base_layer.py:1012 __call__  **
        outputs = call_fn(inputs, *args, **kwargs)
    /home/arvoelke/anaconda3/envs/*/lib/python3.8/site-packages/tensorflow/python/keras/layers/wrappers.py:244 call
        output_shape = self.compute_output_shape(input_shape).as_list()
    /home/arvoelke/anaconda3/envs/*/lib/python3.8/site-packages/tensorflow/python/keras/layers/wrappers.py:188 compute_output_shape
        child_output_shape = self.layer.compute_output_shape(child_input_shape)
    /home/arvoelke/anaconda3/envs/*/lib/python3.8/site-packages/tensorflow/python/keras/layers/core.py:1218 compute_output_shape
        raise ValueError(

    ValueError: The innermost dimension of input_shape must be defined, but saw: (None, None)

../../anaconda3/envs/*/lib/python3.8/site-packages/tensorflow/python/autograph/impl/api.py:670: ValueError

Show more spiking examples

Such as the one here: https://arvoelke.github.io/nengolib-docs/nengolib.networks.RollingWindow.html
-- or preferably something scaled up and structured the same way as the LMUCell.

Support sequences with irregular time-steps

Common feature request. Trivial solution is to solve the ODE for each new sample (i.e., fast-forward to the next point in time), which is a matrix exponential for ZOH. Should try the naive approach first, and then speed it up with some approximation if it's too slow.

The Mackey Glass experiment seems to be predicting the past

Could it be that in the Mackey Glass experiment, the network is asked to approximate a 15-step delay, instead of simulating a complex system?

The definition of the X and Y data is as follows:

Y = X[:, :-predict_length, :]
X = X[:, predict_length:, :]

This implies the X data starts at tilmestep predict_length (15) and ends at the end of the series. Y starts at t=0 and runs until 15 steps before the end of the series. From what I understand, this means the network has seen all the values it is takes to predict beforehand. Is this intentional?

LMU allows the input to have ndim>3 (and then fails later)

Minimal reproducer:

inp = tf.keras.layers.Input(shape=(1000, 8, 128))

out = keras_lmu.LMU(
    memory_d=128,
    order=6,
    theta=10,
    hidden_cell=tf.keras.layers.Dense(100),
)(inp)

model = tf.keras.Model(inp, out)
model.summary()

Output:

_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
input_0 (InputLayer)         [(None, 1000, 8, 128)]    0         
_________________________________________________________________
lmu_0 (LMU)                  (None, 100)               93326     
=================================================================
Total params: 93,326
Trainable params: 93,284
Non-trainable params: 42

It looks like the extra dimensions are being collapsed somewhere internally through a reshape, but attempting to use the model can result in shaping errors.

If you try this same thing with other RNNs, such as out = tf.keras.layers.SimpleRNN(128)(inp) then you will get the error message:

ValueError: Input 0 of layer simple_rnn_0 is incompatible with the layer: expected ndim=3, found ndim=4. Full shape received: [200, 1000, 8, 128]

I think we should perform the same kind of validation to disallow ndim>3 explicitly (unless of course there's a way to make this work correctly)?

Signatures in docs are *args, **kwargs which hides all default values

https://www.nengo.ai/keras-lmu/api-reference.html

It's currently not possible to see what the defaults are in the documentation because they are hidden by *args, **kwargs. At first glance I'm not sure why this is happening because the signatures in the source code (e.g. https://github.com/nengo/keras-lmu/blob/master/keras_lmu/layers.py#L11-L285) are made explicit.

allow_growth to permit parallel notebook kernels?

In order to run multiple notebooks at once, I've found it necessary to preface them with:

import tensorflow as tf
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
tf.Session(config=config)

phased_lstm_keras problem? add_weight () got multiple values for argument 'name'

See fferroni/PhasedLSTM-Keras#19

Error when loading model with tf.keras.models.load_model

Minimal example to reproduce the issue:
This was tested with tf version 2.4.1 and keras_lmu version 0.3.1

import tensorflow as tf
import keras_lmu
import numpy as np
from tensorflow.keras.layers import Input
from tensorflow.keras.models import Model

inp = Input(shape=(None, 10))

out = keras_lmu.LMU(
    memory_d=1,
    order=4,
    theta=5,
    hidden_cell=tf.keras.layers.LSTMCell(units=50)
)(inp)

model = Model(inputs=inp, outputs=out)

optimizer = tf.keras.optimizers.Adam(learning_rate=1e-3)
loss = tf.keras.losses.BinaryCrossentropy()

model.compile(
    optimizer=optimizer,
    loss=loss,
    metrics=loss,
)


history = model.fit(
    np.zeros((32, 16, 10)),
    np.zeros((32, 50)),
    steps_per_epoch=1,
    epochs=2,
)

model.save("my_model")

del model

model = tf.keras.models.load_model(
    "my_model", custom_objects={"LMU": keras_lmu.LMU},
    compile=False
)

Note that if model.fit is not called, the model will be loaded fine.

The error that occurs is:

Traceback (most recent call last):
  File "loading_minimal_example.py", line 39, in <module>
    model = tf.keras.models.load_model(
  File "/home/bjkomer/anaconda3/envs/myenv/lib/python3.8/site-packages/tensorflow/python/keras/saving/save.py", line 212, in load_model
    return saved_model_load.load(filepath, compile, options)
  File "/home/bjkomer/anaconda3/envs/myenv/lib/python3.8/site-packages/tensorflow/python/keras/saving/saved_model/load.py", line 138, in load
    keras_loader.load_layers(compile=compile)
  File "/home/bjkomer/anaconda3/envs/myenv/lib/python3.8/site-packages/tensorflow/python/keras/saving/saved_model/load.py", line 380, in load_layers
    self.loaded_nodes[node_metadata.node_id] = self._load_layer(
  File "/home/bjkomer/anaconda3/envs/myenv/lib/python3.8/site-packages/tensorflow/python/keras/saving/saved_model/load.py", line 417, in _load_layer
    obj, setter = self._revive_from_config(identifier, metadata, node_id)
  File "/home/bjkomer/anaconda3/envs/myenv/lib/python3.8/site-packages/tensorflow/python/keras/saving/saved_model/load.py", line 431, in _revive_from_config
    obj = self._revive_metric_from_config(metadata)
  File "/home/bjkomer/anaconda3/envs/myenv/lib/python3.8/site-packages/tensorflow/python/keras/saving/saved_model/load.py", line 540, in _revive_metric_from_config
    obj = metrics.deserialize(
  File "/home/bjkomer/anaconda3/envs/myenv/lib/python3.8/site-packages/tensorflow/python/keras/metrics.py", line 3446, in deserialize
    return deserialize_keras_object(
  File "/home/bjkomer/anaconda3/envs/myenv/lib/python3.8/site-packages/tensorflow/python/keras/utils/generic_utils.py", line 346, in deserialize_keras_object
    (cls, cls_config) = class_and_config_for_serialized_keras_object(
  File "/home/bjkomer/anaconda3/envs/myenv/lib/python3.8/site-packages/tensorflow/python/keras/utils/generic_utils.py", line 311, in class_and_config_for_serialized_keras_object
    deserialized_objects[key] = deserialize_keras_object(
  File "/home/bjkomer/anaconda3/envs/myenv/lib/python3.8/site-packages/tensorflow/python/keras/utils/generic_utils.py", line 360, in deserialize_keras_object
    return cls.from_config(cls_config)
  File "/home/bjkomer/anaconda3/envs/myenv/lib/python3.8/site-packages/tensorflow/python/keras/metrics.py", line 642, in from_config
    return super(MeanMetricWrapper, cls).from_config(config)
  File "/home/bjkomer/anaconda3/envs/myenv/lib/python3.8/site-packages/tensorflow/python/keras/engine/base_layer.py", line 720, in from_config
    return cls(**config)
TypeError: __init__() got an unexpected keyword argument 'reduction'

Add software license

from utils.utils import create_config "No Module named 'utils'"

http://localhost:8888/notebooks/experiments/psMNIST-phased-lstm.ipynb

Does:

from utils.utils import create_config

resulting in an error:

ModuleNotFoundError: No module named 'utils'

It is unclear what python module this refers to.

acc -> accuracy

In https://github.com/abr/neurips2019/blob/master/experiments/psMNIST-standard.ipynb I had to change:

plt.plot(result.history['val_acc'], label="Validation")
plt.plot(result.history['acc'], label="Training")

to

plt.plot(result.history['val_accuracy'], label="Validation")
plt.plot(result.history['accuracy'], label="Training")

This may be another versioning issue of some sort.

In order to get neurips2019 running at all on Windows 10, I had to perform some installation gymnastics.

Legendre memory unit (LMU) for normal MNIST

Hi, all

I wanted to know is it possible to train the proposed model on normal MNIST dataset without permutation? and how is it possible?

Yours Sincerely,
Ehsan

Register LMUCell with Keras

If the LMUCell is wrapped in another layer (e.g. RNN) then it cannot be serialized since LMUCell is a custom object unknown to Keras. For example:

# Build an LMU layer
dt = 1e-3
activation = "tanh"
dropout=0.2

lmu_layer = RNN(
    keras_lmu.LMUCell(
        memory_d=10,
        order=8,
        theta=10 / dt,
        hidden_cell=Dense(1024, activation),
        hidden_to_memory=False,
        memory_to_memory=False,
        input_to_hidden=False,
        dropout=dropout,
    ),
    return_sequences=True,
)

# Test serialization
lmu_layer.from_config(
    lmu_layer.get_config(),
)

This fails with ValueError: Unknown layer: LMUCell.

The quick fix is to tell Keras about the LMUCell via custom_objects:

# Test serialization
lmu_layer.from_config(
    lmu_layer.get_config(),
    custom_objects={"LMUCell":keras_lmu.LMUCell},  # <-- This is key
)

Although this allows the LMUCell to be properly (de)serialized, this requires direct access and may be challenging if using additional scripts on top of the RNN.

It seems like there is a way to register custom objects with Keras and that may be the proper general solution, just don't have time to test that out right now!

---

aside
For completeness/reference, using theLMU layer (instead of the LMUCell wrapped in an RNN, for example) serializes fine:

lmu_layer_builtin = keras_lmu.LMU(
    memory_d=10,
    order=8,
    theta=10 / dt,
    hidden_cell=Dense(1024, activation),
    hidden_to_memory=False,
    memory_to_memory=False,
    input_to_hidden=False,
    dropout=dropout,
    return_sequences=True,
)
lmu_layer_builtin.from_config(
    lmu_layer_builtin.get_config(),
)

Support stateful=True on the LMU (and other RNN flags)

The LMU layer currently accepts return_sequences and passes that through to tf.keras.layers.RNN with the created LMUCell. However, there are a number of other RNN flags that are ignored (https://www.tensorflow.org/api_docs/python/tf/keras/layers/RNN):

• return_state
• go_backwards
• stateful
• unroll
• time_major

In order to use any of these additional flags, the LMUCell must be invoked directly then and passed to tf.keras.layers.RNN alongside the flags. Supporting the flags at the layer level would mirror the pattern for other recurrent layers such as tf.keras.layers.LSTM.

Create simple LMUCell example with in-line explanations

Upgrade to TensorFlow 2

and ensure all examples still perform well

Make a pytorch implementation

Common feature request.

AttributeError: 'LMUFFT' object has no attribute 'kernel'

Hello，I'm a beginner in tensorflow2，and recently I'm doing some stuff on "Time Series Forcasting"
I've read your paper, and simply want to replace your implimentation with tf.keras.layers.LSTM

But when I'm testing function keras_lmu.LMU , I got such errors:

AttributeError: in user code:

    C:\ProgramData\Miniconda3\lib\site-packages\tensorflow\python\keras\engine\training.py:806 train_function  *
        return step_function(self, iterator)
    C:\ProgramData\Miniconda3\lib\site-packages\keras_lmu\layers.py:439 call  *
        return self.fft_layer.call(inputs, training=training)
    C:\ProgramData\Miniconda3\lib\site-packages\keras_lmu\layers.py:619 call  *
        u = tf.matmul(inputs, self.kernel, name="input_encoder_mult")
    
    AttributeError: 'LMUFFT' object has no attribute 'kernel'

model.summary() still goes well, so I think the model build successfully but unsuccessfully initial weight?
May be this line did not run? https://github.com/nengo/keras-lmu/blob/master/keras_lmu/layers.py#L153

My Testing code

import tensorflow as tf
import tensorflow.keras as keras
############################################## from your doc
import keras_lmu
from tensorflow.keras import Input, Model
from tensorflow.keras.layers import Dense
lmu_layer = keras_lmu.LMU(
    memory_d=1,
    order=256,
    theta=784,
    hidden_cell=tf.keras.layers.SimpleRNNCell(units=10),
)

inputs = Input((None, 10))
lmus = lmu_layer(inputs)
outputs = Dense(1)(lmus)

model = Model(inputs=inputs, outputs=outputs)
#################################################### from your doc

model.summary()

x_train = tf.ones((5,5,10))
y_train = tf.ones((5,5,10))
x_test = tf.ones((1,))
y_test = tf.ones((1,))
model.compile(
    loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
    optimizer=keras.optimizers.RMSprop(),
    metrics=["accuracy"],
)

history = model.fit(x_train, y_train, epochs=2, validation_split=0.2)

test_scores = model.evaluate(x_test, y_test, verbose=2)
print("Test loss:", test_scores[0])
print("Test accuracy:", test_scores[1])

My python version is 3.8.3

tensorflow version is 2.3.1

keras-lmu version is 0.3.0

Am I writting these code correct? , or is there a bug in keras_lmu\layers.py?

THANKS!!!

Link to published PDF in readme

Add page range to bibtex instruction

Blocked on publication in proceedings.

Feature - parallel training

Are there any plans to implement training in parallel manner as shown in

https://arxiv.org/pdf/2102.11417.pdf

Paths to DropoutRNNCellMixin and BaseRandomLayer not found (keras-lmu 0.7.0)

This is likely due to Keras version changes. Currently using Keras 3.2.1 with Tensorflow 2.16.1.

So far, I've changed the following in layers.py:

from keras.src.layers.rnn.dropout_rnn_cell_mixin import DropoutRNNCellMixin
to:
from tensorflow.python.keras.layers.recurrent import DropoutRNNCellMixin

and:
from keras.src.engine.base_layer import BaseRandomLayer
to:
from tensorflow.keras.layers import Layer as BaseRandomLayer

and:
training = tf.keras.backend.learning_phase()
to:
training = True

Remove/replace LMUCellODE and LMUCellGating

We have better versions of these, and they are not currently documented or used by any of the examples.

Reformat with flake8+black