T RAINING AND G ENERATING N EURAL N ETWORKS IN C OMPRESSED W EIGHT S PACE
A BSTRACT
The inputs and/or outputs of some neural nets are weight matrices of other neural nets.
Indirect encodings or end-to-end compression of weight matrices could help to scale such approaches.
Our goal is to open a discussion on this topic, starting with recurrent neural networks for character-level language modelling whose weight matrices are encoded by the discrete cosine transform.
Our fast weight version thereof uses a recurrent neural network to parameterise the compressed weights.
We present experimental results on the enwik8 dataset
I NTRODUCTION
Many modern neural nets (NNs) are big.
Their number of trainable parameters may easily exceed a few hundred million.
This is a bottleneck for end-to-end differentiable systems that manipulate some NN’s weights as input or output variables of another NN, i.e., fast weight programmers (Schmidhuber, 1991; 1992b).
This concept has seen a revival under various names such as “hyper-networks” (Ha et al., 2017), and also in the context of Transformers with linear attention (Katharopoulos et al., 2020; Schlag et al., 2021).
The former were presented as NNs that produce a weight matrix of another NN, but practical implementations only generate a vector that scales the rows of a weight matrix, instead of generating a de novo weight matrix.
In the original work (Schmidhuber, 1991; 1992b) and related fast weight variants (Schmidhuber, 1993d; Ba et al., 2016; Schlag & Schmidhuber, 2017), the matrix size bottleneck is addressed through outer products of two NN-invented vectors of activations which generate rank-one weight matrices at each time step.
For small NNs, the whole weight matrix can be directly parametrised as the output of another NN (e.g. Schmidhuber (1991); Sitzmann et al. (2020)), which is a natural way of generating contextdependent weight matrices.
If we had a differentiable compression method for weight matrices, we could scale up such direct approaches.
Compact weight matrix representations could also be useful in scenarios which require a NN to read the weight matrix of another NN or itself (Schmidhuber, 1992a; 1993b;a;c; Unterthiner et al., 2020;
Harb et al., 2020; Faccio et al., 2021).
In this work we revisit indirect encodings (Koutn´ık et al., 2010a;b) based on the Discrete Cosine
Transform (DCT) (Ahmed et al., 1974).
DCT offers a simple and differentiable encoding of NN weight matrices.
The model can be trained from scratch and end-to-end (including the DCT transformations).
While NNs with DCT-encoded weight matrices have been successfully applied in evolutionary settings (Koutn´ık et al., 2010a;b), DCT’s effect on model performance has not been clearly quantiﬁed in prior work.
We evaluate recurrent neural networks (RNNs) with DCT-encoded weight matrices on a character level language modelling task with various compression rates.
We also report results of the fast weight version in which we augment the DCT-based model by another NN which generates the DCT-coefﬁcients of the main network.
N EURAL N ETWORKS WITH DCT-E NCODED W EIGHT M ATRICES
Here we review neural networks whose weight matrices are encoded as DCT coefﬁcients (Koutn´ık et al., 2010a;b; 2012), and deﬁne the type of models considered in this work. [Figure] [Figure] [Figure]
[Figure]
[Figure] [Figure] [Figure]
[Figure]
In contrast to regular neural networks whose trainable parameters are a set of weight matrices θ = { W 1 , ..., W N } , trainable parameters in a DCT-based neural network are a set of DCT frequency coefﬁcients G = { g 1 , ..., g N } which are used to generate the effective weight matrices { W ( g 1 ) , ..., W ( g N ) } .
In prior work (Koutn´ık et al., 2010a;b; 2012; Gomez et al., 2012; Srivastava et al., 2012; van Steenkiste et al., 2016), G is referred to as genome , and each g i as chromosome .
As a formal illustration let us consider a regular feed-forward layer (where we omit bias) which transforms an input vector x ∈ R m into an output y ∈ R n with an element-wise activation σ and a trainable weight matrix W ∈ R n × m as:
[Equation]
Its DCT-based counterpart is a layer with trainable parameters g ∈ R c with c (cid:28) n × m which transforms x ∈ R m to y ∈ R n as:
[Equation]
where operations in Eq. 2 and 3 correspond to a decompression process which transforms a DCTcoefﬁcient vector g ∈ R c into a weight matrix W ( g ) ∈ R n × m .
First, the TopLeft n,m operation generates a sparse matrix W f ∈ R n × m from g ∈ R c .
W f represents the weight matrix in the frequency domain.
Its coefﬁcients are all zero except the c highest frequency coefﬁcients (antidiagonals in the top left corner of the matrix) which are ﬁlled using the coefﬁcients of g .
This mapping is tied to the model deﬁnition, and the model is trained end-to-end under the corresponding weight pattern in the time domain.
The choice of high frequency domain (instead of low-frequency one) is empirically discussed in Appendix B.1. The inverse of DCT, DCT − 1 is then applied to the sparse frequency domain matrix W f to obtain a dense matrix W ( g ) ∈ R n × m .
Figure 1 illustrates these intermediate steps of the decompression process.
We also note that both DCT and its inverse DCT − 1 can be expressed as simple (differentiable) matrix multiplications using ﬁxed DCT matrices: [Equation]
[Equation]
where D n ∈ R n × n and D m ∈ R m × m are ﬁxed orthonormalised DCT matrices.
The size c of the DCT-coefﬁcient vector is a model hyper-parameter.
With c (cid:28) n × m , where n × m is the size of the original matrix, g effectively offers a compressed representation of W ( g ) .
In our experiments, instead of explicitly ﬁxing c , we ﬁx a compression rate and derive the value of c .
In practice, we allocate separate DCT-coefﬁcient vectors for different weight matrices (i.e. matrices which have different roles), e.g., in an RNN:
[Equation]
W ( g i ) and R ( g r ) are generated by parameters
g i and g r respectively via two separate transformations.
Similarly, for a Long Short-Term Memory (LSTM) (Hochreiter & Schmidhuber, 1997), separate DCT parameters are allocated for each weight matrix and decompressed separately.
This method has been previously investigated for RNNs in evolutionary settings (Koutn´ık et al., 2010a;b) and for small feedforward NNs in supervised settings of toy binary classiﬁcation tasks (Fabisch et al., 2013).
We revisit this method for RNNs in a language modelling setting.
In the experimental Sec. 3, we evaluate the performance of LSTM-RNNs with DCT-encoded weights under different compression rates.
[Figure]
Fast weight version.
Once we have a compact representation of weight matrices, an attractive extension is to generate the weights in the compressed space depending on the context by parameterising it by another neural network (Schmidhuber, 1991; 1992b; 1993d; Gomez & Schmidhuber, 2005;
Ha et al., 2017; Schlag et al., 2021).
Here we focus on systems whose main network (fast net) is a regular RNN (Eq. 11) and we parameterise its two weight matrices: input-to-hidden weights and hidden-to-hidden weights, using two separate LSTM-RNNs (Eqs.
This system thus transforms an input vector x t ∈ R m at time step t to a recurrent hidden state h t ∈ R n as:
[Equation]
where ◦ denotes the composition.
Figure 2 illustrates the corresponding model.
This approach can be considered as an alternative to outer-product based fast weights, with a potentially full rank weight matrix update at each time step.
Also, the slow network (e.g. LSTM θ i ) can see the current states of weight matrix ( g it − 1 ) to generate new weight states ( g it ) in the compressed space.
In Appendix B.2, we also discuss an alternative approach where a single LSTM is used to parameterise the two weight matrices.
E XPERIMENTAL RESULTS
We train character level LSTM-RNN language models on the standard enwik8 dataset (Hutter, 2012) which contains about 90 M characters for training with a character vocabulary size of 205.
We follow the base conﬁguration of Merity et al.
(2018) which has three LSTM-RNN layers with dimensions of 1840 for the two ﬁrst layers and 400 for the third one, with an embedding matrix of size 400 whose shared by the input and output layers.
In the DCT models, we encode each weight matrix in an LSTM layer by DCT as described in Sec. 2.
We apply the same compression rate for each weight matrix.
More on experimental details can be found in Appendix A.
Table 1 presents the effects of different compression rates on DCT-based LSTM model performance.
We obtain degradation of about 14% relative in bits-per-character with a compression rate of 90%.
This model with 4.8 M parameters has to be compared to the performance of baseline models with smaller layer sizes (the simplest way to reduce the model size) presented in Table 2.
The performance of the baseline model with a layer size of 465 and an embedding size of 400 from Table(also 4.8 M parameters) is similar to the performance of the DCT model with a compression rate of 90% in Table 1.
Apparently a large model with compressed weights does not improve much over a small model with a comparable number of parameters.
However, reducing the number of parameters by only reducing the recurrent layer size without reducing the embedding dimension has its limits.
[Table] [Table]
For example
[Table]
with an LSTM layer size of 16 (which is extremely small) and an embedding size of 400, the corresponding model size is 780 K. Further parameter reduction requires to reduce the embedding size.
In contrast, the DCT-based model can be conﬁgured to have an arbitrarily small number of parameters, without requiring to modify other model hyper-parameters, as shown in the lower part of Table 1.
Table 3 highlights performance of fast weight models.
The hyper-parameters are chosen such that the total number of model parameters is comparable to those we discuss in Table 1 and 2.
We found the performance of these models to vary a lot even with similar numbers of parameters. In general, we found that models with more fast variables (# Fast Params in tables) tend to perform better.
The best model achieves performance similar to the one of baseline models. More comparisons among the fast weight model variants can be found in Table 5 in Appendix B.
[Table]
D ISCUSSION AND F UTURE W ORK
We revisited DCT-based indirect encodings of neural network weight matrices in a character level language modelling task.
Our experiments seem to indicate that at least in the case of language modelling, large architectures with DCT-compressed weights don’t improve much over small networks with a similar number of parameters.
However, we also showed that the DCT-based encoding allows for controlling the number of parameters in each layer ﬂexibly and independently of hyperparameters in other layers.
This property may be of practical interest in some applications.
What is the best way of encoding neural network weights?
This fundamental question remains unanswered, but we hope our work will help trigger a discussion of this topic.
Our fast weight variants can be extended in various ways to build highly adaptive language models (Lazaridou et al., 2021; Irie et al., 2018).
In particular, the weight generator networks could be augmented by additional inputs for error signals (Schmidhuber, 1993b; Hochreiter et al., 2001; Rocki, 2016) to facilitate the weight update generation.
Future work will evaluate our models on tasks which explicitly require context-adaptive behaviour, e.g., texts containing multiple domain shifts. A CKNOWLEDGEMENTS
We thank Sjoerd van Steenkiste and Imanol Schlag for valuable comments and suggestions on the ﬁrst version of the manuscript.
This research was partially funded by ERC Advanced grant no: 742870, project AlgoRNN, and by Swiss National Science Foundation grant no: 200021 192356, project NEUSYM.
We thank NVIDIA Corporation for donating several DGX machines, and IBM for donating a Minsky machine.
A E XPERIMENTAL DETAILS
In this section, we provide experimental details.
As a baseline architecture to evaluate DCTparameterised NNs (Table 1), we follow the base conﬁguration of Merity et al.
(2018) which has three LSTM-RNN layers with 1840 nodes for the two ﬁrst layers and 400 for the third one, with an embedding layer of size 400 whose parameters are shared for input and output embeddings.
For all small models with less than 6 M parameters in all tables, dropout is completely disabled to avoid an extra factor for comparison.
For larger models in Table 1, the dropout conﬁguration from Merity et al.
(2018) is applied: 0.1 to the feed-forward hidden layers between the LSTM layers, 0.2 to the recurrent weights, and 0.4 to the output of the last LSTM layer.
Otherwise all models in all tables are trained with the same training conﬁguration: we use Adam optimiser with a learning rate of 0.001, with a batch size of 128, and a backpropagation span of 200 characters.
In the DCT models, we encode each weight matrix (separately for each gate) of an LSTM layer by
DCT as described in Sec. 2.
We apply the same compression rate for all weight matrices.
To initialise the model parameters which are DCT-coefﬁcient vectors, we ﬁrst generate and initialise a weight matrix in time domain, and apply the DCT compression to obtain the corresponding initialisation, which we found to work well in practice.
Implementation notes.
Like any fast weight variants which generate the entire weight matrix at each time step, implementation of DCT based fast weight RNNs requires a custom operation for backpropagation.
In fact, a naive implementation would store weights for each time step for the backward pass, which can quickly result in prohibitive space requirement.
Instead, in a custom implementation we only store the compressed weights for each time step during the forward pass, and the actual weights are recomputed via decompression for each backward step.
An obvious drawback is that the decompression operations need to be called at each backward step, which can be slow.
The corresponding code is available at https://github.com/kazuki-irie/dct-fast-weights .
B A BLATION STUDIES AND EXTRA EXPERIMENTAL RESULTS
Here we present extra experimental results which could not be included in the main text because of the space limitation.
B. C HOICE OF HIGH VS LOW FREQUENCY SPARSITY PATTERNS
The choice of high frequency sparsity pattern in the frequency matrix is arbitrary for our language models, because the textual data does not contain any a priori frequency property, and the input character embeddings are jointly learned with the rest of model parameters, given the pre-determined architectural choice of high or low-frequency coefﬁcients in the model.
Our choice is thus only supported by the empirical results, presented in Table 4, which show that the DCT models with high-frequency non-zero coefﬁcients tend to outperform those using low-frequency coefﬁcients.
[Table]
A LTERNATIVE FAST WEIGHT PARAMETERISATIONS
In Eqs.
7-11, we used two separate slow LSTM RNNs to parameterise DCT coefﬁcients of the two weight matrices of the fast RNN.
An alternaitve approach is to parameterise the two DCT vectors as the output of a single slow LSTM, such that it is aware of the states of all weights of the fast RNN to generate new weights.
The corresponding equations are as follows:
[Equation]
However, as shown in Table
[Table]
we experimentally found that the variant with two separate LSTMs (which we denote as Twin approach) perform better than the variant using only one LSTM (denoted as Single ).
Thus, in Table 3, we only reported results for twin model variants
[Table]
[Table]