The Annotated S4

Notice that this challenge makes use of JAX with the Flax NN library. Whereas we
personally primarily use Torch, the purposeful nature of JAX is an effective match
for a number of the complexities of S4. We make heavy use of vmap,
scan,
their NN
cousins, and most significantly jax.jit
to compile quick and environment friendly S4 layers.

from functools import partial
import jax
import jax.numpy as np
from flax import linen as nn
from jax.nn.initializers import lecun_normal, regular
from jax.numpy.linalg import eigh, inv, matrix_power
from jax.scipy.sign import convolve

if __name__ == "__main__":
    # For this tutorial, assemble a worldwide JAX rng key
    # However we do not need it when importing as a library
    rng = jax.random.PRNGKey(1)

Half 1: State Area Fashions

Let’s get began! Our objective is the environment friendly modeling of lengthy
sequences. To do that, we’re going to construct a brand new neural community layer
primarily based on State Area Fashions. By the top of this part we will probably be ready
to construct and run a mannequin with this layer. Nevertheless, we’re going to want
some technical background. Let’s work our approach by way of the background of
the paper.

The state
space model is outlined by this easy equation. It maps a 1-D enter
sign u(t) to an N-D latent state x(t) earlier than projecting to a 1-D output sign
y(t).
start{aligned}
x'(t) &= boldsymbol{A}x(t) + boldsymbol{B}u(t)
y(t) &= boldsymbol{C}x(t) + boldsymbol{D}u(t)
finish{aligned}
Our objective is to easily use the SSM as a black-box illustration
in a deep sequence mannequin, the place boldsymbol{A}, boldsymbol{B}, boldsymbol{C},
boldsymbol{D} are parameters realized by gradient descent. For
the rest, we’ll omit the parameter boldsymbol{D} for exposition (or
equivalently, assume boldsymbol{D} = 0
as a result of the time period boldsymbol{D}u will be
considered as a skip connection and is simple to compute).

An SSM maps a enter u(t) to a state
illustration vector x(t) and an
output y(t). For simplicity, we assume
the enter and output are one-dimensional, and the state illustration
is N-dimensional. The primary equation
defines the change in x(t) over
time.

Our SSMs will probably be outlined by three matrices – boldsymbol{A}, boldsymbol{B},
boldsymbol{C} – which we’ll study. For now we start with a
random SSM, to outline sizes,

def random_SSM(rng, N):
    a_r, b_r, c_r = jax.random.cut up(rng, 3)
    A = jax.random.uniform(a_r, (N, N))
    B = jax.random.uniform(b_r, (N, 1))
    C = jax.random.uniform(c_r, (1, N))
    return A, B, C

Discrete-time
SSM: The Recurrent Illustration

To be utilized on a discrete enter sequence (u_0, u_1, dots ) as a substitute of steady
operate u(t), the SSM have to be
discretized by a step dimension Delta that represents the decision of the
enter. Conceptually, the inputs u_k can
be considered as sampling an implicit underlying steady sign u(t), the place u_k =
u(okay Delta).

To discretize the continuous-time SSM, we use the bilinear
method, which converts the state matrix boldsymbol{A} into an approximation boldsymbol{overline{A}}. The discrete SSM
is:
start{aligned}
boldsymbol{overline{A}} &= (boldsymbol{I} – Delta/2 cdot
boldsymbol{A})^{-1}(boldsymbol{I} + Delta/2 cdot boldsymbol{A})
boldsymbol{overline{B}} &= (boldsymbol{I} – Delta/2 cdot
boldsymbol{A})^{-1} Delta boldsymbol{B}
boldsymbol{overline{C}} &= boldsymbol{C}
finish{aligned}

def discretize(A, B, C, step):
    I = np.eye(A.form[0])
    BL = inv(I - (step / 2.0) * A)
    Ab = BL @ (I + (step / 2.0) * A)
    Bb = (BL * step) @ B
    return Ab, Bb, C

This equation is now a sequence-to-sequence map u_k mapsto y_k as a substitute of
function-to-function. Furthermore the state equation is now a recurrence in
x_k, permitting the discrete SSM to be
computed like an RNN. Concretely, x_k in
mathbb{R}^N will be considered as a hidden state with
transition matrix boldsymbol{overline{A}}.
start{aligned}
x_{okay} &= boldsymbol{overline{A}} x_{k-1} +
boldsymbol{overline{B}} u_k
y_k &= boldsymbol{overline{C}} x_k
finish{aligned}

Because the paper says, this “step” operate does look superficially like
that of an RNN. We will implement this with a scan
in JAX,

def scan_SSM(Ab, Bb, Cb, u, x0):
    def step(x_k_1, u_k):
        x_k = Ab @ x_k_1 + Bb @ u_k
        y_k = Cb @ x_k
        return x_k, y_k

    return jax.lax.scan(step, x0, u)

Placing all the things collectively, we are able to run the SSM by first
discretizing, then iterating step-by-step,

def run_SSM(A, B, C, u):
    L = u.form[0]
    N = A.form[0]
    Ab, Bb, Cb = discretize(A, B, C, step=1.0 / L)

    # Run recurrence
    return scan_SSM(Ab, Bb, Cb, u[:, np.newaxis], np.zeros((N,)))[1]

Tangent: A Mechanics Instance

To realize some extra instinct and check our SSM implementation, we pause
from machine studying to implement a classic
example from mechanics.

On this instance, we take into account the ahead place y(t) of a mass connected to a wall with a
spring. Over time, various drive u(t)
is utilized to this mass. The system is parameterized by mass (m), spring fixed (okay), friction fixed (b). We will relate these with the next
differential equation:

start{aligned}
my”(t) = u(t) – by'(t) – ky(t)
finish{aligned}

Rewriting this in matrix kind yields an SSM within the following
kind:

start{aligned}
boldsymbol{A} &= start{bmatrix} 0 & 1 -k/m & -b/m
finish{bmatrix}
boldsymbol{B} & = start{bmatrix} 0 1/m finish{bmatrix} &
boldsymbol{C} = start{bmatrix} 1 & 0 finish{bmatrix}
finish{aligned}

def example_mass(okay, b, m):
    A = np.array([[0, 1], [-k / m, -b / m]])
    B = np.array([[0], [1.0 / m]])
    C = np.array([[1.0, 0]])
    return A, B, C

Wanting on the boldsymbol{C}, we
ought to be capable of persuade ourselves that the primary dimension of the
hidden state is the place (since that turns into y(t)). The second dimension is the rate,
as it’s impacted by u(t) by way of boldsymbol{B}. The transition boldsymbol{A} relates these phrases.

We’ll set u to be a steady
operate of t,

@partial(np.vectorize, signature="()->()")
def example_force(t):
    x = np.sin(10 * t)
    return x * (x > 0.5)

Let’s run this SSM by way of our code.

def example_ssm():
    # SSM
    ssm = example_mass(okay=40, b=5, m=1)

    # L samples of u(t).
    L = 100
    step = 1.0 / L
    ks = np.arange(L)
    u = example_force(ks * step)

    # Approximation of y(t).
    y = run_SSM(*ssm, u)

    # Plotting ---
    import matplotlib.pyplot as plt
    import seaborn
    from celluloid import Digicam

    seaborn.set_context("paper")
    fig, (ax1, ax2, ax3) = plt.subplots(3)
    digicam = Digicam(fig)
    ax1.set_title("Pressure $u_k$")
    ax2.set_title("Place $y_k$")
    ax3.set_title("Object")
    ax1.set_xticks([], [])
    ax2.set_xticks([], [])

    # Animate plot over time
    for okay in vary(0, L, 2):
        ax1.plot(ks[:k], u[:k], colour="purple")
        ax2.plot(ks[:k], y[:k], colour="blue")
        ax3.boxplot(
            [[y[k, 0] - 0.04, y[k, 0], y[k, 0] + 0.04]],
            showcaps=False,
            whis=False,
            vert=False,
            widths=10,
        )
        digicam.snap()
    anim = digicam.animate()
    anim.save("pictures/line.gif", dpi=150, author="imagemagick")

if False:
    example_ssm()

Neat! And that it was simply 1 SSM, with 2 hidden states over 100
steps. The ultimate mannequin could have had 100s of stacked
SSMs over 1000’s of steps. However first – we
must make these fashions sensible to coach.

Coaching SSMs:
The Convolutional Illustration

The punchline of this part is that we are able to flip the “RNN” above
right into a “CNN” by unrolling. Let’s undergo the derivation.

The recurrent SSM shouldn’t be sensible for coaching on trendy {hardware}
resulting from its sequential nature. As a substitute, there’s a well-known connection
between linear time-invariant (LTI) SSMs and steady convolutions.
Correspondingly, the recurrent SSM can truly be written as a discrete
convolution.

For simplicity let the preliminary state be x_{-1} = 0. Then unrolling explicitly
yields:

start{aligned}
x_0 &= boldsymbol{overline{B}} u_0 &
x_1 &= boldsymbol{overline{A}} boldsymbol{overline{B}} u_0 +
boldsymbol{overline{B}} u_1 &
x_2 &= boldsymbol{overline{A}}^2 boldsymbol{overline{B}} u_0 +
boldsymbol{overline{A}} boldsymbol{overline{B}} u_1 +
boldsymbol{overline{B}} u_2 & dots

y_0 &= boldsymbol{overline{C}} boldsymbol{overline{B}} u_0 &
y_1 &= boldsymbol{overline{C}} boldsymbol{overline{A}}
boldsymbol{overline{B}} u_0 + boldsymbol{overline{C}}
boldsymbol{overline{B}} u_1 &
y_2 &= boldsymbol{overline{C}} boldsymbol{overline{A}}^2
boldsymbol{overline{B}} u_0 + boldsymbol{overline{C}}
boldsymbol{overline{A}} boldsymbol{overline{B}} u_1 +
boldsymbol{overline{C}} boldsymbol{overline{B}} u_2
& dots
finish{aligned}

This may be vectorized right into a convolution with an express method
for the convolution kernel.

start{aligned}
y_k &= boldsymbol{overline{C}} boldsymbol{overline{A}}^okay
boldsymbol{overline{B}} u_0 + boldsymbol{overline{C}}
boldsymbol{overline{A}}^{k-1} boldsymbol{overline{B}} u_1 + dots +
boldsymbol{overline{C}} boldsymbol{overline{A}}
boldsymbol{overline{B}} u_{k-1} +
boldsymbol{overline{C}}boldsymbol{overline{B}} u_k

y &= boldsymbol{overline{Ok}} ast u
finish{aligned}

start{aligned}
boldsymbol{overline{Ok}} in mathbb{R}^L =
(boldsymbol{overline{C}}boldsymbol{overline{B}},
boldsymbol{overline{C}}boldsymbol{overline{A}}boldsymbol{overline{B}},
dots,
boldsymbol{overline{C}}boldsymbol{overline{A}}^{L-1}boldsymbol{overline{B}})
finish{aligned}
We name boldsymbol{overline{Ok}} the SSM
convolution kernel or filter.

Notice that this can be a large filter. It’s the dimension of the
whole sequence!

def K_conv(Ab, Bb, Cb, L):
    return np.array(
        [(Cb @ matrix_power(Ab, l) @ Bb).reshape() for l in range(L)]
    )

Warning: this implementation is naive and unstable. In apply it
will fail to work for greater than very small lengths. Nevertheless, we’re
going to exchange it with S4 in Half 2, so for now we simply preserve it round
as a placeholder.

We will compute the results of making use of this filter both with a
customary direct convolution or by utilizing convolution theorem with Fast Fourier
Transform (FFT). The discrete convolution theorem – for round
convolution of two sequences – permits us to effectively calculate the
output of convolution by first multiplying FFTs of the enter sequences
after which making use of an inverse FFT. To make the most of this theorem for
non-circular convolutions as in our case, we have to pad the enter
sequences with zeros, after which unpad the output sequence. Because the size
will get longer this FFT technique will probably be extra environment friendly than the direct
convolution,

def causal_convolution(u, Ok, nofft=False):
    if nofft:
        return convolve(u, Ok, mode="full")[: u.shape[0]]
    else:
        assert Ok.form[0] == u.form[0]
        ud = np.fft.rfft(np.pad(u, (0, Ok.form[0])))
        Kd = np.fft.rfft(np.pad(Ok, (0, u.form[0])))
        out = ud * Kd
        return np.fft.irfft(out)[: u.shape[0]]

The CNN technique and the RNN technique yield (roughly) the identical
outcome,

def test_cnn_is_rnn(N=4, L=16, step=1.0 / 16):
    ssm = random_SSM(rng, N)
    u = jax.random.uniform(rng, (L,))
    jax.random.cut up(rng, 3)
    # RNN
    rec = run_SSM(*ssm, u)

    # CNN
    ssmb = discretize(*ssm, step=step)
    conv = causal_convolution(u, K_conv(*ssmb, L))

    # Test
    assert np.allclose(rec.ravel(), conv.ravel())

An SSM Neural Community.

We now have all the equipment wanted to construct a fundamental SSM neural
community layer. As outlined above, the discrete SSM defines a map from
mathbb{R}^L to mathbb{R}^L, i.e. a
1-D sequence map. We assume that we’re going to be studying the
parameters B and C, in addition to a step dimension Delta and a scalar D parameter. The HiPPO matrix is used for the
transition A. We study the step dimension in
log house.

def log_step_initializer(dt_min=0.001, dt_max=0.1):
    def init(key, form):
        return jax.random.uniform(key, form) * (
            np.log(dt_max) - np.log(dt_min)
        ) + np.log(dt_min)

    return init

For the SSM layer a lot of the work is to construct the filter. The precise
name to the community is simply the (enormous) convolution we specified
above.

Notice for Torch customers: setup in Flax is known as every time
the parameters are up to date. That is just like the Torch
parameterizations.

As famous above this similar layer can be utilized both as an RNN or a CNN.
The argument decode determines which path is used. Within the
case of RNN we cache the earlier state at every name in a Flax variable
assortment known as cache.

class SSMLayer(nn.Module):
    N: int
    l_max: int
    decode: bool = False

    def setup(self):
        # SSM parameters
        self.A = self.param("A", lecun_normal(), (self.N, self.N))
        self.B = self.param("B", lecun_normal(), (self.N, 1))
        self.C = self.param("C", lecun_normal(), (1, self.N))
        self.D = self.param("D", nn.initializers.ones, (1,))

        # Step parameter
        self.log_step = self.param("log_step", log_step_initializer(), (1,))

        step = np.exp(self.log_step)
        self.ssm = discretize(self.A, self.B, self.C, step=step)
        self.Ok = K_conv(*self.ssm, self.l_max)

        # RNN cache for lengthy sequences
        self.x_k_1 = self.variable("cache", "cache_x_k", np.zeros, (self.N,))

    def __call__(self, u):
        if not self.decode:
            # CNN Mode
            return causal_convolution(u, self.Ok) + self.D * u
        else:
            # RNN Mode
            x_k, y_s = scan_SSM(*self.ssm, u[:, np.newaxis], self.x_k_1.worth)
            if self.is_mutable_collection("cache"):
                self.x_k_1.worth = x_k
            return y_s.reshape(-1).actual + self.D * u

Since our SSMs function on scalars, we make H totally different, stacked copies (H totally different SSMs!) with totally different parameters.
Right here we use the Flax
vmap technique to simply outline these copies,

def cloneLayer(layer):
    return nn.vmap(
        layer,
        in_axes=1,
        out_axes=1,
        variable_axes={"params": 1, "cache": 1, "prime": 1},
        split_rngs={"params": True},
    )

SSMLayer = cloneLayer(SSMLayer)

This SSM Layer can then be put into an ordinary NN. Right here we add a
block that pairs a name to an SSM with dropout and a linear
projection.

class SequenceBlock(nn.Module):
    layer_cls: nn.Module
    layer: dict  # Hyperparameters of inside layer
    dropout: float
    d_model: int
    prenorm: bool = True
    glu: bool = True
    coaching: bool = True
    decode: bool = False

    def setup(self):
        self.seq = self.layer_cls(**self.layer, decode=self.decode)
        self.norm = nn.LayerNorm()
        self.out = nn.Dense(self.d_model)
        if self.glu:
            self.out2 = nn.Dense(self.d_model)
        self.drop = nn.Dropout(
            self.dropout,
            broadcast_dims=[0],
            deterministic=not self.coaching,
        )

    def __call__(self, x):
        skip = x
        if self.prenorm:
            x = self.norm(x)
        x = self.seq(x)
        x = self.drop(nn.gelu(x))
        if self.glu:
            x = self.out(x) * jax.nn.sigmoid(self.out2(x))
        else:
            x = self.out(x)
        x = skip + self.drop(x)
        if not self.prenorm:
            x = self.norm(x)
        return x

We will then stack a bunch of those blocks on high of one another to
produce a stack of SSM layers. This can be utilized for classification or
technology in the usual approach as a Transformer.

class Embedding(nn.Embed):
    num_embeddings: int
    options: int

    @nn.compact
    def __call__(self, x):
        y = nn.Embed(self.num_embeddings, self.options)(x[..., 0])
        return np.the place(x > 0, y, 0.0)

class StackedModel(nn.Module):
    layer_cls: nn.Module
    layer: dict  # Additional arguments to go into layer constructor
    d_output: int
    d_model: int
    n_layers: int
    prenorm: bool = True
    dropout: float = 0.0
    embedding: bool = False  # Use nn.Embed as a substitute of nn.Dense encoder
    classification: bool = False
    coaching: bool = True
    decode: bool = False  # In all probability ought to be moved into layer_args

    def setup(self):
        if self.embedding:
            self.encoder = Embedding(self.d_output, self.d_model)
        else:
            self.encoder = nn.Dense(self.d_model)
        self.decoder = nn.Dense(self.d_output)
        self.layers = [
            SequenceBlock(
                layer_cls=self.layer_cls,
                layer=self.layer,
                prenorm=self.prenorm,
                d_model=self.d_model,
                dropout=self.dropout,
                training=self.training,
                decode=self.decode,
            )
            for _ in range(self.n_layers)
        ]

    def __call__(self, x):
        if not self.classification:
            if not self.embedding:
                x = x / 255.0  # Normalize
            if not self.decode:
                x = np.pad(x[:-1], [(1, 0), (0, 0)])
        x = self.encoder(x)
        for layer in self.layers:
            x = layer(x)
        if self.classification:
            x = np.imply(x, axis=0)
        x = self.decoder(x)
        return nn.log_softmax(x, axis=-1)

In Flax we add the batch dimension as a lifted transformation. We
must route by way of a number of variable collections which deal with RNN and
parameter caching (described beneath).

BatchStackedModel = nn.vmap(
    StackedModel,
    in_axes=0,
    out_axes=0,
    variable_axes={"params": None, "dropout": None, "cache": 0, "prime": None},
    split_rngs={"params": False, "dropout": True},
)

General, this defines a sequence-to-sequence map of form (batch
dimension, sequence size, hidden dimension), precisely the signature uncovered
by associated sequence fashions corresponding to Transformers, RNNs, and CNNs.

Full code for coaching is outlined in training.py.

Whereas we now have our essential mannequin, there are two core issues
with SSMs. First, the randomly initialized SSM truly doesn’t
carry out very properly. Moreover, computing it naively like we’ve accomplished so
far is absolutely sluggish and reminiscence inefficient. Subsequent, we’ll full our
dialogue of the modeling facet of S4 by defining a particular
initialization for long-range dependencies, after which determine easy methods to
compute this SSM Layer sooner – rather a lot sooner
(Part 2)!

Half 1b:
Addressing Lengthy-Vary Dependencies with HiPPO

Prior work discovered that
the essential SSM truly performs very poorly in apply. Intuitively,
one rationalization is that they endure from gradients scaling exponentially
within the sequence size (i.e., the vanishing/exploding gradients
drawback). To handle this drawback, earlier work developed the HiPPO
principle of continuous-time memorization.

HiPPO specifies a category of sure matrices boldsymbol{A} in mathbb{R}^{N instances N}
that when integrated, enable the state x(t) to memorize the historical past of the enter
u(t). An important matrix on this
class is outlined by the HiPPO matrix.

start{aligned}
(textual content{textbf{HiPPO Matrix}})
qquad
boldsymbol{A}_{nk}
=
start{instances}
(2n+1)^{1/2}(2k+1)^{1/2} & textual content{if } n > okay
n+1 & textual content{if } n = okay
0 & textual content{if } n < okay
finish{instances}
finish{aligned}

Earlier work discovered that merely modifying an SSM from a random matrix
boldsymbol{A} to HiPPO improved its
efficiency on the sequential MNIST classification benchmark from 60% to 98%.

This matrix goes to be actually vital, however it’s a little bit of
magic. For our functions we primarily must know that: 1) we solely must
calculate it as soon as, and a pair of) it has a pleasant, easy construction (which we’ll
exploit partially 2). With out going into the ODE math, the primary takeaway
is that this matrix goals to compress the previous historical past right into a state that
has sufficient data to roughly reconstruct the historical past.

def make_HiPPO(N):
    P = np.sqrt(1 + 2 * np.arange(N))
    A = P[:, np.newaxis] * P[np.newaxis, :]
    A = np.tril(A) - np.diag(np.arange(N))
    return -A

Diving a bit deeper, the intuitive rationalization of this matrix is that
it produces a hidden state that memorizes its historical past. It does this by
protecting observe of the coefficients of a Legendre
polynomial. These coefficients let it approximate all the
earlier historical past. Allow us to have a look at an instance,

def example_legendre(N=8):
    # Random hidden state as coefficients
    import numpy as np
    import numpy.polynomial.legendre

    x = (np.random.rand(N) - 0.5) * 2
    t = np.linspace(-1, 1, 100)
    f = numpy.polynomial.legendre.Legendre(x)(t)

    # Plot
    import matplotlib.pyplot as plt
    import seaborn

    seaborn.set_context("discuss")
    fig = plt.determine(figsize=(20, 10))
    ax = fig.gca(projection="3d")
    ax.plot(
        np.linspace(-25, (N - 1) * 100 + 25, 100),
        [0] * 100,
        zs=-1,
        zdir="x",
        colour="black",
    )
    ax.plot(t, f, zs=N * 100, zdir="y", c="r")
    for i in vary(N):
        coef = [0] * N
        coef[N - i - 1] = 1
        ax.set_zlim(-4, 4)
        ax.set_yticks([])
        ax.set_zticks([])
        # Plot foundation operate.
        f = numpy.polynomial.legendre.Legendre(coef)(t)
        ax.bar(
            [100 * i],
            [x[i]],
            zs=-1,
            zdir="x",
            label="xpercentd" % i,
            colour="brown",
            fill=False,
            width=50,
        )
        ax.plot(t, f, zs=100 * i, zdir="y", c="b", alpha=0.5)
    ax.view_init(elev=40.0, azim=-45)
    fig.savefig("pictures/leg.png")

if False:
    example_legendre()

The purple line represents that curve we’re approximating, whereas the
black bars signify the values of our hidden state. Every is a
coefficient for one factor of the Legendre collection proven as blue
capabilities. The instinct is that the HiPPO matrix updates these
coefficients every step.

Half 2: Implementing S4

Warning: this part has a whole lot of math. Roughly it boils right down to
discovering a approach to compute the filter from Half 1 for “HiPPO-like”
matrices actually quick. In case you are , the main points are
actually neat. If not, skip to Half 3 for some cool purposes like
MNIST completion.

Skip Button

To set the stage, recall that S4 has two essential variations from a
fundamental SSM. The primary addresses a modeling problem –
long-range dependencies – by utilizing a particular method for the boldsymbol{A} matrix outlined within the earlier
half. These particular SSMs had been thought of in predecessor works to S4.

The second essential function of S4 solves the computational
problem of SSMs by introducing a particular illustration and
algorithm to have the ability to work with this matrix!

The elemental bottleneck in computing the discrete-time SSM is that
it includes repeated matrix multiplication by boldsymbol{overline{A}}. For instance,
computing naively includes L successive
multiplications by boldsymbol{overline{A}}, requiring O(N^2 L) operations and O(NL) house.

Particularly, recall this operate right here:

def K_conv(Ab, Bb, Cb, L):
   return np.array(
       [(Cb @ matrix_power(Ab, l) @ Bb).reshape() for l in range(L)]
   )

The contribution of S4 is a secure technique for rushing up this
specific operation. To do that we’re going to give attention to the case the place
the SSM has particular construction: particularly, Diagonal Plus Low-Rank
(DPLR) in complicated house.

A DPLR SSM is (boldsymbol{Lambda} –
boldsymbol{P}boldsymbol{Q}^*, boldsymbol{B}, boldsymbol{C})
for some diagonal boldsymbol{Lambda}
and matrices boldsymbol{P}, boldsymbol{Q},
boldsymbol{B}, boldsymbol{C} in mathbb{C}^{N instances 1}. We
assume with out lack of generality that the rank is 1, i.e. these matrices are vectors.

Beneath this DPLR assumption, S4 overcomes the pace bottleneck in
three steps

1. As a substitute of computing boldsymbol{overline{Ok}} immediately, we
   compute its spectrum by evaluating its truncated
   generating function . This now includes a matrix
   inverse as a substitute of energy.
2. We present that the diagonal matrix case is equal to the
   computation of a Cauchy
   kernel frac{1}{omega_j –
   zeta_k}.
3. We present the low-rank time period can now be corrected by making use of the
   Woodbury
   identity which reduces (boldsymbol{Lambda} +
   boldsymbol{P}boldsymbol{Q}^*)^{-1} by way of boldsymbol{Lambda}^{-1}, actually decreasing to
   the diagonal case.

Step 1. SSM Producing
Capabilities

The principle step will probably be switching from computing the sequence to
computing its producing operate. From the paper’s appendix:

To handle the issue of computing powers of boldsymbol{overline{A}}, we introduce
one other method. As a substitute of computing the SSM convolution filter boldsymbol{overline{Ok}} immediately, we
introduce a producing operate on its coefficients and compute
evaluations of it.

The truncated SSM producing operate at node z with truncation L is
hat{mathcal{Ok}}_L(z; boldsymbol{overline{A}},
boldsymbol{overline{B}}, boldsymbol{overline{C}}) in mathbb{C} :=
sum_{i=0}^{L-1} boldsymbol{overline{C}} boldsymbol{overline{A}}^i
boldsymbol{overline{B}} z^i

def K_gen_simple(Ab, Bb, Cb, L):
    Ok = K_conv(Ab, Bb, Cb, L)

    def gen(z):
        return np.sum(Ok * (z ** np.arange(L)))

    return gen

The producing operate primarily converts the SSM convolution
filter from the time area to frequency area. This transformation is
additionally known as z-transform (as much as
a minus signal) in management engineering literature. Importantly, it
preserves the identical data, and the specified SSM convolution filter
will be recovered. As soon as the z-transform of a discrete sequence recognized, we
can get hold of the filter’s discrete fourier rework from evaluations of
its z-transform
at the roots of unity Omega = {
exp(2pi frac{okay}{L} : okay in [L] }. Then, we are able to apply inverse
fourier transformation, stably in O(L log
L) operations by making use of an FFT, to
get better the filter.

def conv_from_gen(gen, L):
    # Consider at roots of unity
    # Producing operate is (-)z-transform, so we consider at (-)root
    Omega_L = np.exp((-2j * np.pi) * (np.arange(L) / L))
    atRoots = jax.vmap(gen)(Omega_L)
    # Inverse FFT
    out = np.fft.ifft(atRoots, L).reshape(L)
    return out.actual

Extra importantly, within the producing operate we are able to exchange the
matrix energy with an inverse!
hat{mathcal{Ok}}_L(z) = sum_{i=0}^{L-1} boldsymbol{overline{C}}
boldsymbol{overline{A}}^i boldsymbol{overline{B}} z^i =
boldsymbol{overline{C}} (boldsymbol{I} – boldsymbol{overline{A}}^L
z^L) (boldsymbol{I} – boldsymbol{overline{A}} z)^{-1}
boldsymbol{overline{B}} = boldsymbol{widetilde{C}} (boldsymbol{I}
– boldsymbol{overline{A}} z)^{-1} boldsymbol{overline{B}}

And for all z in Omega_L, we have now
z^L = 1 in order that time period is eliminated. We
then pull this fixed time period into a brand new boldsymbol{widetilde{C}}. Critically, this
operate doesn’t name K_conv,

def K_gen_inverse(Ab, Bb, Cb, L):
    I = np.eye(Ab.form[0])
    Ab_L = matrix_power(Ab, L)
    Ct = Cb @ (I - Ab_L)
    return lambda z: (Ct.conj() @ inv(I - Ab * z) @ Bb).reshape()

However it does output the identical values,

def test_gen_inverse(L=16, N=4):
    ssm = random_SSM(rng, N)
    ssm = discretize(*ssm, 1.0 / L)
    b = K_conv(*ssm, L=L)

    a = conv_from_gen(K_gen_inverse(*ssm, L=L), L)
    assert np.allclose(a, b)

In abstract, Step 1 permits us to exchange the matrix energy with an
inverse by using a truncated producing operate. Nevertheless this
inverse nonetheless must be calculated L
instances (for every of the roots of unity).

Step 2: Diagonal Case

The subsequent step to imagine particular construction on the matrix
boldsymbol{A} to compute the inverse
sooner than the naive inversion. To start, allow us to first convert the
equation above to make use of the unique SSM matrices. With some algebra you
can increase the discretization and present:

start{aligned}
boldsymbol{widetilde{C}}left(boldsymbol{I} –
boldsymbol{overline{A}} proper)^{-1} boldsymbol{overline{B}}
=
frac{2Delta}{1+z} boldsymbol{widetilde{C}} left[ {2
frac{1-z}{1+z}} – Delta boldsymbol{A} right]^{-1} boldsymbol{B}
finish{aligned}

Now think about boldsymbol{A}=boldsymbol{Lambda} for a
diagonal boldsymbol{Lambda}.
Substituting within the discretization method the authors present that the
producing operate will be written within the following method:

start{aligned}
boldsymbol{hat{Ok}}_{boldsymbol{Lambda}}(z) & = c(z) sum_i cdot
frac{boldsymbol{widetilde{C}}_i boldsymbol{B}_i} {(g(z) –
boldsymbol{Lambda}_i)} = c(z) cdot k_{z,
boldsymbol{Lambda}}(boldsymbol{widetilde{C}}, boldsymbol{B})
finish{aligned} the place c is a
fixed, and g is a operate of z.

We’ve successfully changed an inverse with a weighted dot product.
Let’s make a small helper operate to compute this weight dot product
to be used.

def cauchy_dot(v, omega, lambd):
    return (v / (omega - lambd)).sum()

Whereas not vital for our implementation, it’s price noting that
this can be a Cauchy
kernel and is the topic of many different fast
implementations.

Step 3: Diagonal Plus
Low-Rank

The ultimate step is to calm down the diagonal assumption. Along with
the diagonal time period we enable a low-rank element with boldsymbol{P}, boldsymbol{Q} in
mathbb{C}^{Ntimes 1} such that:

boldsymbol{A} = boldsymbol{Lambda} – boldsymbol{P} boldsymbol{Q}^*

The Woodbury
identity tells us that the inverse of a diagonal plus rank-1 time period is
equal to the inverse of the diagonal plus a rank-1 time period. We write it out
right here including the low-rank time period.

start{aligned}
(boldsymbol{Lambda} + boldsymbol{P} boldsymbol{Q}^*)^{-1} &=
boldsymbol{Lambda}^{-1} – boldsymbol{Lambda}^{-1} boldsymbol{P} (1
+ boldsymbol{Q}^* boldsymbol{Lambda}^{-1} boldsymbol{P})^{-1}
boldsymbol{Q}^* boldsymbol{Lambda}^{-1}
finish{aligned}

There’s a bunch of algebra within the appendix. It largely consists of
substituting this element in for A, making use of the Woodbury identification and
distributing phrases. We find yourself with 4 phrases that every one seem like Step 2
above:

start{aligned}
boldsymbol{hat{Ok}}_{DPLR}(z) & = c(z) [k_{z,
Lambda}(boldsymbol{widetilde{C}}, boldsymbol{boldsymbol{B}}) –
k_{z, Lambda}(boldsymbol{widetilde{C}}, boldsymbol{boldsymbol{P}})
(1 + k_{z, Lambda}(boldsymbol{q^*}, boldsymbol{boldsymbol{P}})
)^{-1} k_{z, Lambda}(boldsymbol{q^*}, boldsymbol{boldsymbol{B}}) ] finish{aligned}

The code consists of amassing up the phrases and making use of 4 weighted
dot merchandise,

def K_gen_DPLR(Lambda, P, Q, B, C, step, unmat=False):
    aterm = (C.conj(), Q.conj())
    bterm = (B, P)

    def gen(o):
        g = (2.0 / step) * ((1.0 - o) / (1.0 + o))
        c = 2.0 / (1.0 + o)

        def okay(a):
            # Checkpoint this calculation for reminiscence effectivity.
            if unmat:
                return jax.remat(cauchy_dot)(a, g, Lambda)
            else:
                return cauchy_dot(a, g, Lambda)

        k00 = okay(aterm[0] * bterm[0])
        k01 = okay(aterm[0] * bterm[1])
        k10 = okay(aterm[1] * bterm[0])
        k11 = okay(aterm[1] * bterm[1])
        return c * (k00 - k01 * (1.0 / (1.0 + k11)) * k10)

    return gen

That is our ultimate model of the Ok
operate. As a result of conv_from_gen is all the time known as collectively
with a producing operate (e.g. K_gen_DPLR), we’ll fuse
them into outline a devoted operate to compute the DPLR SSM kernel
from all of its parameters. (With fewer layers of indirection, this
might additionally make it simpler for XLA compiler to optimize.)

@jax.jit
def cauchy(v, omega, lambd):
    """Cauchy matrix multiplication: (n), (l), (n) -> (l)"""
    cauchy_dot = lambda _omega: (v / (_omega - lambd)).sum()
    return jax.vmap(cauchy_dot)(omega)

def kernel_DPLR(Lambda, P, Q, B, C, step, L):
    # Consider at roots of unity
    # Producing operate is (-)z-transform, so we consider at (-)root
    Omega_L = np.exp((-2j * np.pi) * (np.arange(L) / L))

    aterm = (C.conj(), Q.conj())
    bterm = (B, P)

    g = (2.0 / step) * ((1.0 - Omega_L) / (1.0 + Omega_L))
    c = 2.0 / (1.0 + Omega_L)

    # Discount to core Cauchy kernel
    k00 = cauchy(aterm[0] * bterm[0], g, Lambda)
    k01 = cauchy(aterm[0] * bterm[1], g, Lambda)
    k10 = cauchy(aterm[1] * bterm[0], g, Lambda)
    k11 = cauchy(aterm[1] * bterm[1], g, Lambda)
    atRoots = c * (k00 - k01 * (1.0 / (1.0 + k11)) * k10)
    out = np.fft.ifft(atRoots, L).reshape(L)
    return out.actual

Now we are able to examine whether or not it labored. First, let’s generate a random
Diagonal Plus Low Rank (DPLR) matrix,

def random_DPLR(rng, N):
    l_r, p_r, q_r, b_r, c_r = jax.random.cut up(rng, 5)
    Lambda = jax.random.uniform(l_r, (N,))
    P = jax.random.uniform(p_r, (N,))
    Q = jax.random.uniform(q_r, (N,))
    B = jax.random.uniform(b_r, (N, 1))
    C = jax.random.uniform(c_r, (1, N))
    return Lambda, P, Q, B, C

We will examine that the DPLR technique yields the identical filter as computing
boldsymbol{A} immediately,

def test_gen_dplr(L=16, N=4):
    I = np.eye(4)

    # Create a DPLR A matrix and discretize
    Lambda, P, B, _ = make_DPLR_HiPPO(N)
    A = np.diag(Lambda) - P[:, np.newaxis] @ P[:, np.newaxis].conj().T
    _, _, C = random_SSM(rng, N)

    Ab, Bb, Cb = discretize(A, B, C, 1.0 / L)
    a = K_conv(Ab, Bb, Cb.conj(), L=L)

    # Examine to the DPLR producing operate method.
    C = (I - matrix_power(Ab, L)).conj().T @ Cb.ravel()
    b = kernel_DPLR(Lambda, P, P, B, C, step=1.0 / L, L=L)
    assert np.allclose(a.actual, b.actual)

Diagonal Plus Low-Rank RNN.

A secondary good thing about the DPLR factorization is that it permits us to
compute the discretized type of the SSM with out having to invert the
A matrix immediately. Right here we return to
the paper for the derivation.

Recall that discretization computes,
start{align*}
bm{overline{A}} &= (bm{I} – Delta/2 cdot bm{A})^{-1}(bm{I} +
Delta/2 cdot bm{A})
bm{overline{B}} &= (bm{I} – Delta/2 cdot bm{A})^{-1} Delta
bm{B}
.
finish{align*}

We simplify each phrases within the definition of bm{overline{A}} independently. The primary
time period is:
start{align*}
bm{I} + frac{Delta}{2} bm{A}
&= bm{I} + frac{Delta}{2} (bm{Lambda} – bm{P} bm{Q}^*)
&= frac{Delta}{2} left[ frac{2}{Delta}bm{I} + (bm{Lambda}
– bm{P} bm{Q}^*) right] &= frac{Delta}{2} bm{A_0}
finish{align*}
the place bm{A_0} is outlined as
the time period within the ultimate brackets.

The second time period is called the Backward Euler’s technique. Though
this inverse time period is generally troublesome to take care of, within the DPLR case
we are able to simplify it utilizing Woodbury’s Id as described above.
start{align*}
left( bm{I} – frac{Delta}{2} bm{A} proper)^{-1}
&=
left( bm{I} – frac{Delta}{2} (bm{Lambda} – bm{P} bm{Q}^*)
proper)^{-1}
&=
frac{2}{Delta} left[ frac{2}{Delta} – bm{Lambda} + bm{P}
bm{Q}^* right]^{-1}
&=
frac{2}{Delta} left[ bm{D} – bm{D} bm{P} left( 1 + bm{Q}^*
bm{D} bm{P} right)^{-1} bm{Q}^* bm{D} right] &= frac{2}{Delta} bm{A_1}
finish{align*}
the place bm{D} = left(
frac{2}{Delta}-bm{Lambda} proper)^{-1} and bm{A_1} is outlined because the time period within the ultimate
brackets.

The discrete-time SSM turns into
start{align*}
x_{okay} &= bm{overline{A}} x_{k-1} + bm{overline{B}} u_k
&= bm{A_1} bm{A_0} x_{k-1} + 2 bm{A_1} bm{B} u_k
y_k &= bm{C} x_k
.
finish{align*}

def discrete_DPLR(Lambda, P, Q, B, C, step, L):
    # Convert parameters to matrices
    B = B[:, np.newaxis]
    Ct = C[np.newaxis, :]

    N = Lambda.form[0]
    A = np.diag(Lambda) - P[:, np.newaxis] @ Q[:, np.newaxis].conj().T
    I = np.eye(N)

    # Ahead Euler
    A0 = (2.0 / step) * I + A

    # Backward Euler
    D = np.diag(1.0 / ((2.0 / step) - Lambda))
    Qc = Q.conj().T.reshape(1, -1)
    P2 = P.reshape(-1, 1)
    A1 = D - (D @ P2 * (1.0 / (1 + (Qc @ D @ P2))) * Qc @ D)

    # A bar and B bar
    Ab = A1 @ A0
    Bb = 2 * A1 @ B

    # Recuperate Cbar from Ct
    Cb = Ct @ inv(I - matrix_power(Ab, L)).conj()
    return Ab, Bb, Cb.conj()

Turning HiPPO to DPLR

This method applies to DPLR matrices, however keep in mind we want it
to additionally apply to the HiPPO matrix. Whereas not DPLR in its present kind,
the HiPPO matrix does have particular construction. It’s Regular Plus
Low-Rank (NPLR). As a result of normal matrices
are precisely the category of matrices which can be unitarily diagonalizable,
NPLR matrices are primarily equal to DPLR matrices from the
perspective of SSM fashions. that is simply pretty much as good as DPLR for the needs
of studying an SSM community.

The S4 methods can apply to any matrix boldsymbol{A} that may be decomposed as
Regular Plus Low-Rank (NPLR).
boldsymbol{A} = boldsymbol{V} boldsymbol{Lambda} boldsymbol{V}^*
– boldsymbol{P} boldsymbol{Q}^high = boldsymbol{V} left(
boldsymbol{Lambda} – boldsymbol{V}^* boldsymbol{P}
(boldsymbol{V}^*boldsymbol{Q})^* proper) boldsymbol{V}^*
for unitary boldsymbol{V} in mathbb{C}^{N instances N},
diagonal boldsymbol{Lambda}, and
low-rank factorization boldsymbol{P},
boldsymbol{Q} in mathbb{R}^{N instances r}. An NPLR SSM is
subsequently unitarily equal to some DPLR matrix.

For S4, we have to work with a HiPPO matrix for boldsymbol{A}. This requires first writing
it as a traditional plus low-rank time period, after which diagonalizing to extract
boldsymbol{Lambda} from this
decomposition. The appendix of the paper reveals how by writing the traditional
half as a skew-symmetric
(plus a continuing instances the identification matrix), that are a particular class
of regular matrices.

An extra simplification is that there’s truly a
illustration that ties the low-rank elements phrases boldsymbol{P} = boldsymbol{Q}, which was
proven in follow-up work
to be vital for stability.

def make_NPLR_HiPPO(N):
    # Make -HiPPO
    nhippo = make_HiPPO(N)

    # Add in a rank 1 time period. Makes it Regular.
    P = np.sqrt(np.arange(N) + 0.5)

    # HiPPO additionally specifies the B matrix
    B = np.sqrt(2 * np.arange(N) + 1.0)
    return nhippo, P, B

After extracting the traditional half, we are able to diagonalize to get out the
DPLR phrases. As a result of the traditional half is definitely skew-symmetric, we are able to
extract the actual and complicated elements of boldsymbol{Lambda} individually. This serves
two functions. First, this provides us finer-grained management over the actual
and imaginary elements, which can be utilized to enhance stability. Second,
this lets us use extra highly effective diagonalization algorithms for Hermitian
matrices – the truth is, the present model of JAX doesn’t help GPU
diagonalization for non-Hermitian matrices!

def make_DPLR_HiPPO(N):
    """Diagonalize NPLR illustration"""
    A, P, B = make_NPLR_HiPPO(N)

    S = A + P[:, np.newaxis] * P[np.newaxis, :]

    # Test skew symmetry
    S_diag = np.diagonal(S)
    Lambda_real = np.imply(S_diag) * np.ones_like(S_diag)
    # assert np.allclose(Lambda_real, S_diag, atol=1e-3)

    # Diagonalize S to V Lambda V^*
    Lambda_imag, V = eigh(S * -1j)

    P = V.conj().T @ P
    B = V.conj().T @ B
    return Lambda_real + 1j * Lambda_imag, P, B, V

Sanity examine simply to ensure these identities maintain,

def test_nplr(N=8):
    A2, P, B = make_NPLR_HiPPO(N)
    Lambda, Computer, Bc, V = make_DPLR_HiPPO(N)
    Vc = V.conj().T
    P = P[:, np.newaxis]
    Computer = Computer[:, np.newaxis]
    Lambda = np.diag(Lambda)

    A3 = V @ Lambda @ Vc - (P @ P.T)  # Take a look at NPLR
    A4 = V @ (Lambda - Computer @ Computer.conj().T) @ Vc  # Take a look at DPLR
    assert np.allclose(A2, A3, atol=1e-4, rtol=1e-4)
    assert np.allclose(A2, A4, atol=1e-4, rtol=1e-4)

Closing Test

This exams that all the things works as deliberate.

def test_conversion(N=8, L=16):
    step = 1.0 / L
    # Compute a HiPPO NPLR matrix.
    Lambda, P, B, _ = make_DPLR_HiPPO(N)
    # Random complicated Ct
    C = regular(dtype=np.complex64)(rng, (N,))

    # CNN kind.
    Ok = kernel_DPLR(Lambda, P, P, B, C, step, L)

    # RNN kind.
    Ab, Bb, Cb = discrete_DPLR(Lambda, P, P, B, C, step, L)
    K2 = K_conv(Ab, Bb, Cb, L=L)
    assert np.allclose(Ok.actual, K2.actual, atol=1e-5, rtol=1e-5)

    # Apply CNN
    u = np.arange(L) * 1.0
    y1 = causal_convolution(u, Ok.actual)

    # Apply RNN
    _, y2 = scan_SSM(
        Ab, Bb, Cb, u[:, np.newaxis], np.zeros((N,)).astype(np.complex64)
    )
    assert np.allclose(y1, y2.reshape(-1).actual, atol=1e-4, rtol=1e-4)

Half 3: S4 in Apply

That was a whole lot of work, however now the precise mannequin is concise. In actual fact
we’re solely utilizing 4 capabilities:

1. K_gen_DPLR → Truncated producing operate when boldsymbol{A} is DPLR (S4-part)
2. conv_from_gen → Convert producing operate to
   filter
3. causal_convolution → Run convolution
4. discretize_DPLR → Convert SSM to discrete kind for
   RNN.

S4 CNN / RNN Layer

A full S4 Layer is similar to the straightforward SSM layer above. The
solely distinction is within the the computation of boldsymbol{Ok}. Moreover as a substitute of
studying boldsymbol{C}, we study boldsymbol{widetilde{C}} so we keep away from
computing powers of boldsymbol{A}.
Notice as properly that within the authentic paper boldsymbol{Lambda}, boldsymbol{P},
boldsymbol{Q} are additionally realized. Nevertheless, on this submit, we go away
them fastened for simplicity.

class S4Layer(nn.Module):
    N: int
    l_max: int
    decode: bool = False

    # Particular parameters with multiplicative issue on lr and no weight decay (dealt with by essential prepare script)
    lr = {
        "Lambda_re": 0.1,
        "Lambda_im": 0.1,
        "P": 0.1,
        "B": 0.1,
        "log_step": 0.1,
    }

    def setup(self):
        # Realized Parameters (C is complicated!)
        init_A_re, init_A_im, init_P, init_B = hippo_initializer(self.N)
        self.Lambda_re = self.param("Lambda_re", init_A_re, (self.N,))
        self.Lambda_im = self.param("Lambda_im", init_A_im, (self.N,))
        # Guarantee the actual a part of Lambda is unfavorable
        # (described within the SaShiMi follow-up to S4)
        self.Lambda = np.clip(self.Lambda_re, None, -1e-4) + 1j * self.Lambda_im
        self.P = self.param("P", init_P, (self.N,))
        self.B = self.param("B", init_B, (self.N,))
        # C ought to be init as customary regular
        # This does not work resulting from how JAX handles complicated optimizers https://github.com/deepmind/optax/points/196
        # self.C = self.param("C", regular(stddev=1.0, dtype=np.complex64), (self.N,))
        self.C = self.param("C", regular(stddev=0.5**0.5), (self.N, 2))
        self.C = self.C[..., 0] + 1j * self.C[..., 1]
        self.D = self.param("D", nn.initializers.ones, (1,))
        self.step = np.exp(self.param("log_step", log_step_initializer(), (1,)))

        if not self.decode:
            # CNN mode, compute kernel.
            self.Ok = kernel_DPLR(
                self.Lambda,
                self.P,
                self.P,
                self.B,
                self.C,
                self.step,
                self.l_max,
            )

        else:
            # RNN mode, discretize

            # Flax trick to cache discrete kind throughout decoding.
            def init_discrete():
                return discrete_DPLR(
                    self.Lambda,
                    self.P,
                    self.P,
                    self.B,
                    self.C,
                    self.step,
                    self.l_max,
                )

            ssm_var = self.variable("prime", "ssm", init_discrete)
            if self.is_mutable_collection("prime"):
                ssm_var.worth = init_discrete()
            self.ssm = ssm_var.worth

            # RNN Cache
            self.x_k_1 = self.variable(
                "cache", "cache_x_k", np.zeros, (self.N,), np.complex64
            )

    def __call__(self, u):
        # That is similar to SSM Layer
        if not self.decode:
            # CNN Mode
            return causal_convolution(u, self.Ok) + self.D * u
        else:
            # RNN Mode
            x_k, y_s = scan_SSM(*self.ssm, u[:, np.newaxis], self.x_k_1.worth)
            if self.is_mutable_collection("cache"):
                self.x_k_1.worth = x_k
            return y_s.reshape(-1).actual + self.D * u

S4Layer = cloneLayer(S4Layer)

We initialize the mannequin by computing a HiPPO DPLR initializer

# Manufacturing facility for fixed initializer in Flax
def init(x):
    def _init(key, form):
        assert form == x.form
        return x

    return _init

def hippo_initializer(N):
    Lambda, P, B, _ = make_DPLR_HiPPO(N)
    return init(Lambda.actual), init(Lambda.imag), init(P), init(B)

Sampling and Caching

We will pattern from the mannequin utilizing the RNN implementation. Right here is
what the sampling code appears like. Notice that we preserve a operating cache to
keep in mind the RNN state.

def pattern(mannequin, params, prime, cache, x, begin, finish, rng):
    def loop(i, cur):
        x, rng, cache = cur
        r, rng = jax.random.cut up(rng)
        out, vars = mannequin.apply(
            {"params": params, "prime": prime, "cache": cache},
            x[:, np.arange(1, 2) * i],
            mutable=["cache"],
        )

        def replace(x, out):
            p = jax.random.categorical(r, out[0])
            x = x.at[i + 1, 0].set(p)
            return x

        x = jax.vmap(replace)(x, out)
        return x, rng, vars["cache"].unfreeze()

    return jax.lax.fori_loop(begin, finish, jax.jit(loop), (x, rng, cache))[0]

To get this in a great kind, we first precompute the discretized
model of the the RNN for every S4 layers. We do that by way of the
“prime” assortment of variables.

def init_recurrence(mannequin, params, init_x, rng):
    variables = mannequin.init(rng, init_x)
    vars = {
        "params": params,
        "cache": variables["cache"].unfreeze(),
        "prime": variables["prime"].unfreeze(),
    }
    print("[*] Priming")
    _, prime_vars = mannequin.apply(vars, init_x, mutable=["prime"])
    return vars["params"], prime_vars["prime"], vars["cache"]

Placing this altogether we are able to pattern from a mannequin immediately.

def sample_checkpoint(path, mannequin, size, rng):
    from flax.coaching import checkpoints

    begin = np.zeros((1, size, 1), dtype=int)

    print("[*] Initializing from checkpoint %s" % path)
    state = checkpoints.restore_checkpoint(path, None)
    assert "params" in state
    params, prime, cache = init_recurrence(mannequin, state["params"], begin, rng)
    print("[*] Sampling output")
    return pattern(mannequin, params, prime, cache, begin, 0, size - 1, rng)

Experiments: MNIST

Now that we have now the mannequin, we are able to attempt it out on some MNIST
experiments. For these experiments we linearize MNIST and simply deal with
every picture as a sequence of pixels.

The primary experiments we ran had been on MNIST classification. Whereas not
in principle a tough drawback, treating MNIST as a linear sequence
classification activity is a bit unusual. Nevertheless in apply, the mannequin
with H=256 and 4 layers appears to get
up close to 99% straight away.

A extra visually fascinating activity is producing MNIST digits, by
predicting whole sequences of pixels! Right here, we merely feed in a
sequence of pixels into the mannequin and have it predict the following one like
language modeling. With slightly tweaking, we’re capable of get the mannequin
to an NLL of 0.36 on this activity with dimension 512 and 6 layers (~4m
parameters).

The metric often used for this activity is bits
per dimension which is NLL in base 2 for MNIST. A lack of 0.36
is ~0.52 BPD which is SOTA based on PapersWithCode.

We will additionally do prefix-samples – given the primary 300 pixels, attempt to
full the picture. S4 is on the left, true on the suitable.

def sample_image_prefix(
    params,
    mannequin,
    # size,
    rng,
    dataloader,
    prefix=300,
    # bsz=32,
    imshape=(28, 28),
    n_batches=None,
    save=True,
):
    """Pattern a grayscale picture represented as intensities in [0, 255]"""
    import matplotlib.pyplot as plt
    import numpy as onp

    # from .information import Datasets
    # BATCH = bsz
    # begin = np.zeros((BATCH, size), dtype=int)
    # begin = np.zeros((BATCH, size, 1), dtype=int)
    begin = np.array(subsequent(iter(dataloader))[0].numpy())
    begin = np.zeros_like(begin)
    # params, prime, cache = init_recurrence(mannequin, params, begin[:, :-1], rng)
    params, prime, cache = init_recurrence(mannequin, params, begin, rng)

    BATCH = begin.form[0]
    START = prefix
    LENGTH = begin.form[1]
    assert LENGTH == onp.prod(imshape)

    # _, dataloader, _, _, _ = Datasets["mnist"](bsz=BATCH)
    it = iter(dataloader)
    for j, im in enumerate(it):
        if n_batches shouldn't be None and j >= n_batches:
            break

        picture = im[0].numpy()
        picture = np.pad(
            picture[:, :-1, :], [(0, 0), (1, 0), (0, 0)], constant_values=0
        )
        cur = onp.array(picture)
        # cur[:, START + 1 :, 0] = 0
        # cur = np.pad(cur[:, :-1, 0], [(0, 0), (1, 0)], constant_values=256)
        cur = np.array(cur[:, :])

        # Cache the primary `begin` inputs.
        out, vars = mannequin.apply(
            {"params": params, "prime": prime, "cache": cache},
            cur[:, np.arange(0, START)],
            mutable=["cache"],
        )
        cache = vars["cache"].unfreeze()
        out = pattern(mannequin, params, prime, cache, cur, START, LENGTH - 1, rng)

        # Visualization
        out = out.reshape(BATCH, *imshape)
        ultimate = onp.zeros((BATCH, *imshape, 3))
        final2 = onp.zeros((BATCH, *imshape, 3))
        ultimate[:, :, :, 0] = out
        f = ultimate.reshape(BATCH, LENGTH, 3)
        i = picture.reshape(BATCH, LENGTH)
        f[:, :START, 1] = i[:, :START]
        f[:, :START, 2] = i[:, :START]
        f = final2.reshape(BATCH, LENGTH, 3)
        f[:, :, 1] = i
        f[:, :START, 0] = i[:, :START]
        f[:, :START, 2] = i[:, :START]
        if save:
            for okay in vary(BATCH):
                fig, (ax1, ax2) = plt.subplots(ncols=2)
                ax1.set_title("Sampled")
                ax1.imshow(ultimate[k] / 256.0)
                ax2.set_title("True")
                ax1.axis("off")
                ax2.axis("off")
                ax2.imshow(final2[k] / 256.0)
                fig.savefig("impercentd.%d.png" % (j, okay))
                plt.shut()
                print(f"Sampled batch {j} picture {okay}")
    return ultimate, final2

Experiments: QuickDraw

Subsequent we tried coaching a mannequin to generate drawings. For this we used
the QuickDraw
dataset. The dataset features a model of the dataset downsampled
to MNIST dimension so we are able to use roughly the identical mannequin as above. The dataset
is far bigger although (5M pictures) and extra complicated. We solely skilled for
1 epoch with a H=256, 4 layer mannequin.
Nonetheless, the method was capable of generate comparatively coherent
completions. These are prefix samples with 500 pixels given.

Experiments: Spoken Digits

Lastly we performed with modeling sound waves immediately. For these, we
use the Free
Spoken Digits Datasets an MNIST like dataset of varied audio system
studying off digits. We first skilled a classification mannequin and located
that the method was capable of attain 97% accuracy simply from the uncooked soundwave.
Subsequent we skilled a technology mannequin to provide the sound wave immediately.
With H=512 the mannequin appears to select up
the info comparatively properly. This dataset solely has round 3000 examples,
however the mannequin can produce moderately good (cherry-picked) continuations.
Notice these sequences are 6400 steps lengthy at an 8kHz sampling fee,
discretized to 256 lessons with Mu Law
Encoding.

Our full code
base accommodates extra examples and infrastructure for coaching fashions
for generations and classification.

Conclusion

Placing collectively this submit impressed a lot of ideas about future
work on this space. One apparent conclusion is that long-range fashions have
all kinds of future purposes from acoustic modeling to genomic
sequences to trajectories (to not point out our shared space of NLP).
One other is a few shock that linear fashions will be so efficient right here,
whereas additionally opening up a spread of environment friendly methods. Lastly from a
sensible degree, the transformations in JAX make it very nice to
implement complicated fashions like this in a really concise approach (~200 LoC),
with related effectivity and efficiency!

We finish by thanking the authors Albert Gu and Karan Goel, who had been tremendous
useful in placing this collectively, and pointing you once more to their paper and codebase. Thanks
to Ankit Gupta, Ekin Akyürek, Qinsheng Zhang, Nathan Yan, and Junxiong
Wang for contributions. We’re additionally grateful for Conner Vercellino and
Laurel Orr for offering useful suggestions on this submit.

/ Cheers – Sasha & Sidd

Changelog

• v3
  • Main modifying go from Albert.
  • Repair bug in HiPPO calculation.
  • Added coaching of all S4 parameters.
  • Repair studying fee / initialization points.
• v2
  • Added RNN decoding
  • Added Speech examples
• v1 – Authentic model