A Modern Self-Referential Weight Matrix That Learns to Modify Itself
Abstract
The weight matrix (WM) of a neural network (NN) is its program.
The programs of many traditional NNs are learned through gradient descent in some error function, then remain ﬁxed.
The WM or program of a self-referential NN, however, can keep rapidly modifying all of itself during runtime.
In principle, such NNs can meta-learn to learn, and meta-meta-learn to meta-learn to learn, and so on, in the sense of recursive self-improvement.
Here we revisit such NNs, building upon recent successes of fast weight programmers (FWPs) and closely related linear
Transformers.
We propose a scalable self-referential WM (SRWM) that uses selfgenerated training patterns, outer products and the delta update rule to modify itself.
We evaluate our SRWM in a multi-task reinforcement learning setting with procedurally generated ProcGen game environments.
Our experiments demonstrate both practical applicability and competitive performance of the SRWM. Our code is public.
Introduction
The program of a neural network (NN) is its weight matrix (WM) [1].
For example, with prediction tasks, starting from random values, an NN training procedure based on gradient descent might update the WM to minimize an error function that favors compression of given input-output observations [2].
The WM becomes permanent once training ends and its usefulness is evaluated with respect to its generalisation capability on yet unseen data.
Many environments, however, continue to evolve after training has halted (e g., [3, 4]) and the test setting may deviate from training in ways that exceed the NN’s generalisation capability.
Then human intervention might be required to re-train or ﬁne-tune the model.
To minimize such intervention, we consider NNs that can learn to update their own programs in the light of new experience.
Especially in multi-task learning and meta-learning (learning to learn [5]), it may be useful to learn how to keep changing and ﬁne-tuning the WM in a way that quickly adapts to new challenges [6].
In principle, a WM could learn by itself a way of executing rapid WM adaptations in task-dependent and context-dependent fashion through a generic mechanism for self-modiﬁcation.
Various selfmodifying NNs have been proposed (see Sec. 5).
We revisit the self-referential WM [7, 8, 9, 10] from the ’90s in the light of modern techniques for updating and generating weights.
In particular, we leverage mechanisms which are now well established in the context of Fast Weight Programmers (FWPs; reviewed in Sec. 2) [11, 12, 13].
FWPs have recently seen advancements in terms of performance and scalability, inspired by their formal equivalence [14] to linear variants [15, 16, 17] of the popular Transformer [18].
Here we derive a new type of self-modifying WM which naturally emerges as an extension to recent works on FWPs.
Using ProcGen [19], we evaluate it in a multi-task reinforcement learning (RL) setting with procedurally generated game environments.
We demonstrate both practical applicability and competitive performance of the proposed method.
Background on Fast Weight Programmers
Here we brieﬂy review the essential components of fast weight programmers (FWPs) [11, 12, 13] which our model is built upon (Sec. 3).
FWPs have a slow NN which can rapidly modify weights of another fast NN.
The concept has seen a recent revival, in particular in light of its direct formal connection [14] to linear variants [15, 16, 17, 20] of the popular Transformer [18] when the weight generation is based on outer products between keys and values generated by the slow NN [11].
Recent work augmented the basic FWPs [11, 12] with an improved elementary programming instruction or update rule invoked by the slow NN to reprogram the fast NN, called delta update rule (akin to the delta rule by Widrow and Hoff [21]).
The resulting DeltaNet [14] is a general purpose auto-regressive
NN with linear complexity w.r.t.
input sequence length, which transforms the input x t ∈ R d in to the output
y t ∈ R d out as follows:
[Equation]
[Equation] where ⊗ denotes the outer product, σ is a sigmoid function, and φ is an element-wise activation function whose output elements are all positive and sum up to one (e.g. softmax).
In Eq.
1, the input x t is ﬁrst projected to key k t ∈ R d key , value v
t ∈ R d out , query q t ∈ R d key vectors and a scalar β t ∈ R using a trainable weight matrix W slow ∈ R ( d value +2 ∗ d key +1)
× d in .
The generated key vector k t and a learning rate β t (generated by the slow NN) are used to update the fast weight matrix W t − 1 using the delta rule expressed in Eqs.
The fast weight matrix is typically initialized to zero i.e. W 0 = 0 .
The ﬁnal output y t is obtained by querying the updated fast weight matrix W t using a generated query vector q t (Eq. 4).
We note that the use of φ function for both writing to (Eq. 3) and reading from (Eqs. 2 and 4) the fast weights is crucial for stability when the delta rule is used [14].
In practice, we use the multi-head version [18] of the computation above, that is, after the projection (Eq. 1), the vectors k t , v t , q t are split into equally sized H sub-vectors, and the operations in Eqs.
2-4 are conducted by H computation heads independently.
So the slow NN or the programmer (here a one-layer feedforward NN; Eq. 1) with slow weights
W slow learns by gradient descent to continuously modify or program the fast NN (here also a onelayer feedforward NN; Eq. 4) with fast weights W t as it continually receives a stream of inputs.
For further extensions of this concept to more complex slow and fast NN architectures such as recurrent
NNs, we refer the readers to another recent study [20].
We note that FWPs are also of interest from the perspective of context-sensitive systems, since the fast WM is completely context-dependent: while processing some sequence, a continually changing custom fast NN is built on the ﬂy.
Here we leverage this mechanism to design a new kind of FWP which programs itself.
It can be naturally derived from the operations described above, resulting in a modern version of the selfreferential weight matrix [7, 8, 9, 10] of the ’90s.
A Modern Self-Referential Weight Matrix
Our modern self-referential weight matrix (SRWM) learns to train itself through self-invented key/value “training” patterns and learning rates, invoking sequences of elementary programming instructions based on outer products and the delta update rule, as in the recently proposed variants [14] of FWPs (Sec. 2).
Given an input x t ∈ R d in at time t , our SRWM W t − 1 ∈ R ( d out +2 ∗ d in +1) ×
d in produces four variables [ y t , q t , k t , β t ] where y t ∈ R d out is the output of this layer at the current time step, q t ∈ R d in and k t ∈ R d in are query and key vectors, and β t ∈ R is the self-invented learning rate to be used by the delta rule.
In analogy to the terminology introduced by the original SRWM papers [7, 8, 9, 10], k t ∈ R d in is the modiﬁer -key vector, representing the key whose current value in the SRWM has to be modiﬁed, and q t ∈ R d in is the analyser -query which is again fed to the SRWM to retrieve a new “value” vector to be associated with the modiﬁer-key.
The overall dynamics can be expressed as simply as follows:
[Equation]
where the value vectors have dimensions: v t , ¯ v t ∈ R ( d out +2 ∗ d in +1) .
Figure 1 illustrates the model.
Importantly, the initial values of the SRWM W 0 are the only parameters in this layer which are trained by gradient descent.
In practice, we extend the output dimension of the matrix from “3D+1” ( d out + 2 ∗ d in + 1) to “3D+4” ( d out + 2 ∗ d in + 4) to generate four different learning rates β t ∈ R 4 to be used in Eq.
8 for the four sub-matrices of W t − 1 =
[ W yt − 1 , W qt − 1 , W kt − 1 , W βt − 1 ] used to produce y t , q t , k t , and β t in Eq.
For efﬁcient computation, we also make use of multi-head computation as is done in regular Transformers [18, 15].
Please refer to Appendix A for the full description.
The SRWM described above can potentially be used to replace any regular WM.
Here we replace the WM of a simple feedforward layer.
We consider two models which we can describe using the
DeltaNet as a base model.
Model (1) consists of the DeltaNet in which we replace Eqs.
1-4 by the corresponding SRWM equations Eqs.
In model (2), we replace the slow weight matrix (Eq. 1) in the DeltaNet by an SRWM (Eqs.
In the experimental section, we refer to the model (1) simply as “SRWM” and to the model (2) as “SR-Delta”.
The effective weight matrix is a function of task/episode speciﬁc input
Experiments 4.1 Experimental Settings
We evaluate the proposed model in a multi-task reinforcement learning setting using procedurally generated game environments of ProcGen [19].
The corresponding setting is illustrated in Figure 2.
Unlike other popular game environments such as Atari [22], ProcGen provides various procedurally different levels.
This allows for creating clean train/test splits.
Working with diverse levels is especially relevant to our setting, as we wish to build models which are adaptive to changes across game types, as well as diversity within the same game.
General Settings.
In our main experiment, we jointly train on 6 environments, namely Bigﬁsh ,
Fruitbot , Maze , Leaper , Plunder , and Starpilot in the easy distribution.
We conduct distributed training using the standard IMPALA [23] architecture implemented in Torchbeast [24].
We useactors (i.e. 8 actors per environment).
All our models use the common large architecture of Espeholt [Table] et al.
[23] which consists of a 15-layer residual convolutional vision model.
They differ from each other by the “memory” module inserted between the vision stem and the output layer.
In addition to our SRWM model
[Table]
we train the baseline IMPALA feed-forward and LSTM [25, 26] models, as well as two additional baselines: the DeltaNet [14] and a “Fake SR” model which has exactly the same architecture as the SRWM model but from which we removed the self-modiﬁcation mechanism.
(i.e. we only keep the “y”-part in Eq. 5).
We set the backpropagation span of 50 steps to train the self-modiﬁcation, as well as LSTM and DeltaNet baselines.
The LSTM model has 1 layer withnodes as in the IMPALA baseline.
Both DeltaNet and SRWM have two layers with a hidden size of 128, the same setting as the DeltaNet used for Atari in the previous work [20].
These 6 environments are known for not explicitly requiring “memory” to perform the task (we also conﬁrm this trend by comparing our baseline feed-forward and LSTM RNN models).
In principle, this allows us to evaluate the effect of self-modiﬁcations in isolation (even if it is difﬁcult to completely dissociate self-modiﬁcation from the concept of “memory”).
In addition to the above setting usingenvironments in the easy distribution, we also conduct an extra experiment using 4 environments in the memory distribution ( Dodgeball , Heist , Maze , Miner ) to evaluate our models also in partially observable settings.
In ProcGen memory distributions, the size of the world is increased and the observation is restricted to a small patch of space around the agent [19].
The backpropagation span is increased to 100 time steps for the memory distribution.
In all cases, the memory states (including the weight changes for the SRWM) are only reset at episode boundaries for all stateful models (LSTM,
DeltaNet, and SRWM).
We train for a total of 300 M steps (ca.
50 M per environment) for joint training on 6 environments in the easy distribution, and 800 M steps (200 M steps per environment) for 4 environments in the memory distribution.
Train/Test split.
Following Cobbe et al.
[19], we use 200 levels (level ID 0 to 199) to train in the easy distribution.
For evaluation, instead of randomly sampling the test levels as is commonly done for ProcGen, we consistently use the same set of 3 distinct test splits for all models.
Each of our test splits contains 200 levels, respectively including levels 1000 to 1199, 1200 to 1399, and 1400 to 1599.
We opt for using 3 test splits and report an average score to take into account the performance variability across the choice of test levels.
The performance in the training set is computed using all 200 training levels.
We train each model three times, and thus report training performance averaged over three runs, and test performance over 9 data points (3 test splits for 3 training runs).
In the memory distribution setting, since there is no standard convention [19], we opt for training ontraining levels (level ID 0 to 499) as recommended for the hard distribution.
For testing, we use the exact same setting as in our easy distribution setting described above.
4.2 Results
Overall performance.
Table 1 shows the aggregated normalised scores.
First of all, comparing the LSTM and feed-forward baselines, we conﬁrm Cobbe et al.
[19]’s ﬁnding that the LSTM layer does not provide any improvements regarding the test performance (while some improvements are obtained on the train set).
The two fast weight models, DeltaNet and SRWM, clearly outperform the feedforward and LSTM baselines.
The SRWM achieves a slightly better test score than the DeltaNet.
The trend is slightly different in the memory distribution setting.
While having a similar parameter count, the SRWM variant achieves a better training score than the DeltaNet baseline, while the test scores are rather close.
Overall, the proposed SRWM based model variants achieve very competitive performance.
Comparison to expert models.
We observed that the performance gains achieved by our SRWM over the baselines in the easy distribution are particularly large for two of the environments, Bigﬁsh and Starpilot .
Here we study these two cases in isolation.
We compare the performance of multi-task agents presented above with expert agents trained speciﬁcally on one environment for 50 M steps.
As shown in Table 2, the performance gap between the two agents tends to be larger in case of the model without self-modiﬁcation, especially on Bigﬁsh .
[Table]
Ablation on State Reset.
The SRWM models presented above are trained by carrying over the weight modiﬁcations across entire episodes whose lengths are variable—often episodes are getting longer during training as the agent becomes better at the task.
We were initially uncertain about the empirical stability of the dynamics described by Eqs.
5-8 in such a scenario.
As an ablation study, we also trained and evaluated an SRWM agent by resetting the weight update every ﬁxed time span whose length is the backpropagation span.
We found such a model to fail in leveraging the SRWM mechanism.
It obtained scores of 28.5 (1.2) and 16.1 (2.2) on the train and test splits respectively, which are similar to that of the baseline without self-modiﬁcation (Table 1).
4.3 Discussion
Interpretability.
As is often the case with NNs, interpreting the model behavior is not straightforward.
Looking into the values of β t in Eq. 8 which intuitively deﬁne the strength of the weight modiﬁcation, we observe that the value of all four components of β t varies between 0.50 to 0. depending on the input, rather than the full range of sigmoid values between 0 and 1.
We ﬁnd it difﬁcult to derive any further interpretation beyond those statistics.
Also β t alone does not fully describe the behavior of Eq. 8, as it also depends on the actual values of key and query.
We leave further analysis to future work on certain simpler supervised learning tasks.
Implementation/Limitation.
Similar to recent works on fast weight programmers (FWPs) [14, 15, 20], our SRWM is implemented as a custom CUDA kernel.
While this approach yields competitive computation time and memory-efﬁcient custom backpropagation, its ﬂexibility is limited.
For instance, if we want to replace all weight matrices in the agent’s vision module, a custom implementation for convolution would be required, even if we claim that in theory, the SRWM presented above could replace any regular weight matrix.
Regarding speed, the feedforward and LSTM baselines process about 3,500 steps per second, while DeltaNet and SRWM do 2,300 and 1,700 steps per second respectively on a single P100 GPU.
Related Work
Original Self-Referential Weight Matrix.
The original SRWM was proposed in the ’90s as a framework for self-modifying recurrent NNs [7, 8, 9, 10].
Such an RNN has special output and input units to directly address and read and modify any of its own current weights through an index for each weight of its weight matrix (i.e., for a weight matrix with an input/output dimension N , the weight index ranges from 0 to N 2 − 1 which is encoded as a binary vector).
In contrast, our self-modiﬁcation is based on key/value associations , i.e., to encode a WM modiﬁcation, our NN generates a key vector, a value vector, and a temporary learning rate which allows for the rapid modiﬁcation of an entire rank at a time [11, 13].
This design is reinforced by the recent success of linear Transformers and fast weight programmers [15, 14, 20].
In this sense, our SRWM is a modern approach to self-modiﬁcation, even if the use of outer products to parameterise fast weight generation itself is not (e.g. [11, 13]).
Other Self-Modifying Neural Networks.
There are also more recent works on self-modifying
NNs.
Neuromodulated plasticity is a Hebbian-style self-modiﬁcation [27, 28, 29, 30] which also makes use of outer products to generate a modulation term which is added to the base weights.
The corresponding computations can also be interpreted as key/value/query association operations.
However, the key, query, and value patterns are hard-coded to be one of the input/output pairs of the corresponding layer at each time step.
While this circumvents the necessity to allocate parameters for generating those vectors, it is known that the resulting program can be simply expressed as a regular attention [13] over the past outputs [31].
In contrast, in our model, all these patterns are arbitrary as they are generated from learned transformations whose parameters are themselves self-modifying.
Hierarchical Fast Weight Programmers.
As reviewed in Sec. 2, an FWP is an NN which learns to generate, update, and maintain weights of another NN.
However, a typical FWP has a slow NN with a weight matrix that remains ﬁxed after training.
Previous work [20] has proposed to go one step further by parameterising the slow weights in the Delta Net with another FWP to obtain the
DeltaDelta Net.
However, such a hierarchy has no end, as the highest level programmer would still have a ﬁxed weight matrix.
In this work, we follow the spirit of early work [7, 8, 9, 10, 13] and collapse these potentially hierarchical meta-levels into one single self-referential weight matrix.
Fixed Weight Meta-RNNs.
Learning learning dynamics using a ﬁxed-weight NN (typically an
RNN) has become a common approach [32, 33, 34, 35, 36, 37, 38, 39].
A truly self-referential weight matrix, however, would allow for modifying all of its own components.
The only thing that’s trained by gradient descent are the SWRM’s initial weights at the beginning of each episode—all of them, however, may rapidly change during sequence processing, in a way that’s driven by the SRWM itself.
Recursive Self-Improvements.
Beyond the scope of NNs, the concept of self-modiﬁcation is of general interest when considering autonomous, self-improving machines [5, 40, 41, 42, 43, 44].
In future work, we intend to use the proposed SRWM as a backbone for exploiting the beneﬁts of recursive self-improvement.
Conclusion
We proposed a new type of self-referential weight matrix (SRWM) with a modern mechanism for self-modiﬁcation.
Our self-improving NNs learn to generate patterns of keys and values and learning rates, translating these patterns into rapid changes of their own weight matrix through sequential outer products and invocations of the delta update rule.
We demonstrated that our SRWM is practical and performs well in a multi-task reinforcement learning setting using game environments procedurally generated by ProcGen.
In future work, we plan to evaluate the SRWM also on supervised learning tasks such as few shot learning. Acknowledgments
We would like to thank Karl Cobbe for answering some practical questions about ProcGen.
Kazuki
Irie wishes to thank Anand Gopalakrishnan for letting him know about ProcGen.
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 Model Details
Equations for the four-learning rate case.
In Sec. 3, for the purpose of clarity, we presented the equations for our SRWM model in the case where we only have a single learning rate β t .
Here we provide a complete description of an SRWM with a separate self-inveted learning rate for each component.
As we noted in Sec. 3, the SRWM can be split into sub-matrices: W t − 1 =
[ W yt − 1 , W qt − 1 , W kt − 1 , W βt − 1 ] according to the sub-components used to produce y t , q t , k t , and β t in Eq.
In case where we use separate learning rates, we need separate equations to describe the update of each sub-matrices.
For example, for the “y”-part W yt − 1 , while keeping the same equation for the ﬁrst projection (Eq. 5), the rest becomes:
[Equation]
) where y kt and y qt are the “y”-part of ¯ v t and v t in Eq.
6 and 7 respectively, and β y,t ∈ R is one of four learning rate dedicated to the “y”-part.
The equations for other sub-matrices W qt − 1 , W kt − 1 , W βt −are analogous.
Use of multiple heads.
In this work, the SRWM was inserted between other layers with learned parameters with conﬁgurable dimensionalities.
This allows us for an efﬁcient computation using multiple heads for computation described as follows.
Given a number of heads H which we use in the SRWM layer, we conﬁgured the model dimensions such that the input dimension to an SRWM layer d in is divisible by H .
The input is then split into H equally sized components, and each head executes separate SRWM operations (Eqs.
5-8) on one of the input components.
In consequence, the
SRWM with the same model hyper-parameters as the DeltaNet has less parameters than the DeltaNet.
For example, if d in = d key , the common head dimension is d = d in //H , the parameter shape of key projection in the SRWM is ( H , d , d ) while it is ( d in , d in ) =
( H ∗ d , H ∗ d ) for the DeltaNet.
In case the input size of the SRWM layer is not conﬁgurable, this option has to be disabled and a single head version should be used.
B Extra Result Tables
[Table]
[Table]
Plunder
[Table]
Starpilot
[Table]
[Table]