Unified representation of sequence models

• Quantifying Memory Utilization with Effective State-Size

Introduction

• The majority of sequence models of practical interest can formally be expressed as either linear systems ($y=Tu$) or systems with input-varying linear operators $(y=f_T(u)u$), the latter of which they abbreviate to LIV.

• The input-varying linear operator framework decouples the input-varying featurization $u\to T:=f_T(u)$ and the linear mapping $y=Tu$ required to construct and apply the operator respectively.

• This decomposition enables a wide array of deep learning primitives to be uniformly formulated as linear systems, including models like convolutions , linear state-transition recurrences and attention variants.

Background (recurrences)

Consider a general linear recurrence formulated as follows:

$$s_{i+1}=A_i{s}_i+B_i{u}_i,y_i=C_i{s}_i+D_i{u}_i,$$

where $(A_i\in {\mathbb{R}}^{n_{i+1}\times n_i},B_i\in {\mathbb{R}}^{n_{i+1}\times d},C_i\in {\mathbb{R}}^{d\times n_i},D_i\in {\mathbb{R}}^{d\times d}{)}_{i\in [\ell ]}$;
$s_i$ and $n_i$ are the state and state-size at sequence index $i$ respectively.

Creating the operator T

• In this section, we begin by showing that most modern sequence models can effectively materialize a linear operator $T$.

• Using the flattened notation, we let $T\in {\mathbb{R}}^{d\ell \times d\ell }$, $u,y\in {\mathbb{R}}^{d\ell }$ denote the operator, inputs, and outputs respectively, $\ell$ denote the sequence length and $d$ denote the channel dimension. Here, we index sequence indices with subscripts, i.e. $T_{ij}\in {\mathbb{R}}^{d\times d}$, $u_i\in {\mathbb{R}}^d$ and channels (and other non-temporal dimensions) with superscripts, i.e. $T^{\alpha \beta }\in {\mathbb{R}}^{\ell \times \ell }$, $u^{\alpha }\in {\mathbb{R}}^{\ell }$. For additional details on notation, refer to Section D.1.

A unified representation of sequence models

While typically nonlinear, most sequence models of interest can effectively materialize a linear operator $T$, where the equation $y=Tu$ faithfully expresses the computation performed by the model (see Section E.2 for further elaboration):

$$\begin{array}{rlrl}{T}_{ij}& =C_i{B}_j& & \text{linear attention,}\\ T_{ij}& =C_i{A}_{i-1}\cdots A_{j+1}{B}_j& & \text{recurrence,}\\ T_{ij}& =K_{i-j}& & \text{convolution,}\\ T_{ij}& =C_i{K}_{i-j}{B}_j& & \text{gated convolution,}\\ T_{ij}& =\sigma (C_i{B}_j)& & \text{attention.}\end{array}$$

• They make a distinction between linear systems (such as convolutions) and LIVs (such as attention and gated convolutions).
• In the former, the operator $T$ is input-invariant whereas the latter are constructed via causal featurizers that map past inputs into features, i.e. $f_B:u_{:i}\mapsto B_i$, which are then used to construct the elements of $T$ as outlined above.