Learn synaptic time-constants via backpropagation

[arvoelke]

(Feature request) Since the synaptic operations are differentiable with respect to the coefficients in their difference equations, they can also be optimized via backpropagation.

After learning, the synaptic coefficients can be mapped back onto the time-constants of the synapse via the poles of the new discrete transfer function, as follows:

from nengolib.synapses import DoubleExp
from nengolib.signal import cont2discrete

dt = 0.001
tau1 = 0.05
tau2 = 0.03

# LTI difference equation is obtained from the ZOH discretization tf coefficients
# https://en.wikipedia.org/wiki/Digital_filter#Difference_equation
disc = nengolib.signal.cont2discrete(DoubleExp(tau1, tau2), dt=dt, method='zoh')

# These coefficients can be related back to the time-constant by the following formula
# derived by zero-pole matching
assert np.allclose(
    np.sort(- dt / np.log(disc.poles)),
    np.sort([tau1, tau2]))

This should be useful any time the data set has some temporal dynamics. These dynamics can be learned through not only the recurrent connection, but the dynamics of the synapses (which are like miniature recurrent connections).

As a contrived yet simple example, suppose our input data is all ones, and the output data is a step-response with some exponential time-constant. If our network is feed-forward with a specific time-constant, then backpropagation could in theory could minimize the MSE by optimizing the time-constant on the synapse. However, nengo_dl is currently only able to reduce the error by scaling the static gain:

import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns

import nengo
import nengo_dl
import tensorflow as tf

tau_actual = 0.005
tau_ideal = 0.1

length = 1000
dt = 0.001

t = np.arange(length)*dt
y_ideal = nengo.Lowpass(tau_ideal).filt(np.ones_like(t), y0=0, dt=dt)

with nengo.Network() as model:
    u = nengo.Node(output=1)
    x = nengo.Node(size_in=1)
    nengo.Connection(u, x, synapse=tau_actual)
    p = nengo.Probe(x, synapse=None)

inputs = {u: np.ones((1, length, 1))}
outputs = {p: y_ideal[None, :, None]}

with nengo_dl.Simulator(model, dt=dt, minibatch_size=1) as sim:
    optimizer = tf.train.AdamOptimizer()
    sim.train(inputs, outputs, optimizer, n_epochs=1000, objective='mse')
    sim.run_steps(length)

plt.figure()
plt.plot(t, y_ideal, label="Ideal")
plt.plot(t, sim.data[p].squeeze(), label="Actual")
plt.xlabel("Time (s)")
plt.legend()
plt.show()

Note: if you set tau_actual == tau_ideal, then the MSE becomes zero. And so the optimal solution with this architecture is to modify the time-constant on the synapse.

[drasmuss]

Yes, this would definitely be a cool thing to support, and very doable from a technical standpoint. It's basically just a matter of changing one flag from True to False. The main issue is a user interface one: how to let the user easily control which parts of the model they want to be trainable. This is part of a larger issue, as there are lots of other parameters that we could in theory allow to be trained, but currently don't (e.g., tau_rc of LIF neurons).

Right now this is all controlled through the trainable system. This works fairly well for the set of parameters that are trainable right now, but doesn't scale up that well. We could actually extend it fairly easily to support Synapse objects, which would be a good intermediate step. But, for example, how could we specify that we want to make the biases of an Ensemble trainable but not the RC constant? I suspect that long-term we'll need a custom solution, rather than piggy-backing on the Nengo config system, but I/we need to do some thinking about what that system might look like.

[arvoelke]

That's an interesting challenge from a design stand-point. One possibility could be to introduce an operator to mark specific values as either learnable or static. For sake of having something to look at, let's say L(·) for trainable and S(·) for static. Under the hood, this could recast values to a special object that the builder would consider when setting the flag. To use the default parameter values the operator could be provided a special value (e.g., Default). Example syntax:

from nengo_dl import L, S, Default

# Marking a transform as fixed (e.g., passthrough nodes)
nengo.Connection(..., transform=S(1))

# or, since 1 is the default:
nengo.Connection(..., transform=S(Default))

# Learning tau_rc while keeping tau_ref fixed (may not be possible right now?)
x = nengo.Ensemble(..., neuron_type=nengo.LIF(tau_rc=L(0.02), tau_ref=S(0.002)))

# Learning the neuron model, and the gains, but not the biases
x = nengo.Ensemble(..., neuron_type=L(Default), gain=L(Default), bias=S(Default))

# Marking a synapse as learnable
nengo.Connection(..., synapse=L(0.01))

This could further scale to learning only a subset of coefficients in a discretized transfer function, although the syntax starts to become unwieldy and the abstraction becomes somewhat leaky. But if this can be done in a way that keeps the builder extensible, then people could potentially roll their own solutions for these special use-cases.

[arvoelke]

Quick update: This feature request is made somewhat obsolete by the new keras_spiking.Lowpass layer which learns the time-constant(s) (and initial state) of the lowpass filter for each dimension. There is also a trainable alpha filter in the same repo.

The caveat is that this layer currently cannot be converted via the NengoDL converter if there is more than one time-constant in the layer. Related issue: nengo/nengo#1636.