Gated recurrent neural networks discover attention

We thank the reviewer for their on-point feedback.

• Here is the precise description of the values the different weight matrices take. The key-values are stored in the first $d^2$ recurrent neurons and the queries in the last $d$ ones (indices $d^2+1$ to $d^2 + d$.

  • For the input gating: $W\_\mathrm{x}^\mathrm{in}$ and $W^\mathrm{in}\_\mathrm{g}$ are matrices of size $(d^2 + d) \times (d+1)$ with

  (W\_\mathrm{x}^\mathrm{in})\_{i,j} = \left \\{ \begin{array}{ll} {(W_V)}\_{i/d,j} &\text{ if } j \leq d \text{ and } i \leq d^2 \\\\ {(W\_Q)}\_{i-d^2,j} & \text{ if } j \leq d \text{ and } i > d^2 \\\\ 0 & \text{ otherwise} \end{array} \right .

  {(W\_\mathrm{g}^\mathrm{in})}\_{i,j} = \left \\{ \begin{array}{ll} {(W_K)}\_{i\,\mathrm{mod}\,d,j} &\text{ if } j \leq d \text{ and } i \leq d^2 \\\\ 1 & \text{ if } j = d+1 \text{ and } i > d^2\\\\ 0 & \text{ otherwise} \end{array} \right .

  • For the recurrent neurons: $\lambda$ is a vector of size $d^2 + d$ with

  \lambda\_i = \left \\{ \begin{array}{ll} 1 & \text{ if } i \leq d^2\\\\ 0 & \text{ otherwise } \end{array} \right .

  • For the output gating: $W^\mathrm{out}\_\mathrm{x}$ and $W^\mathrm{out}\_\mathrm{g}$ are matrices of size $d^2 \times (d^2 + d)$ with
  0 & \text{ otherwise} \end{array} \right . $$ $$ {(W\_\mathrm{g}^\mathrm{out})}\_{i,j} = \left \\{ \begin{array}{ll} 1 & \text{ if } j > d^2 \text{ and } i = j \\, \mathrm{mod} \\, d \\ 0 & \text{ otherwise} \end{array} \right . $$ For the readout matrix: $D$ is of size $d \times d^2$ with $$ D\_{i,j} = \left \\{ \begin{array}{ll} 1 & \text{ if } i = j / d\\\\ 0 & \text{ otherwise } \end{array} \right . $$ We will include this in the appendix of the next version of the paper.

• The associative recall task is a great suggestion. We included promising preliminary results on this task in the global response.

• For the first point, we refer the reviewer to our global response. Regarding the second point, we agree with the reviewer that this is a very intriguing result. We have been investigating the reasons behind this performance but so far remained unsuccessful. We believe that this is an interesting direction for future work.

We believe the new experimental results we provide following the reviewer advice greatly improve the paper and we hope that they may convince the reviewer to reconsider their score.

We appreciate the reviewer’s valuable time in reviewing our paper and would like to nuance their point of view on the limitations imposed on the two architectures we consider.

• As argued in the global response, there is a growing literature showing how extensions of linear self-attention layers can reach, if not surpass, Transformers-level performance. The interest of linear self-attention thus goes beyond its amenability to theoretical analysis.
• While we provide a theoretical construction for the linear case in the main text, we show how this still holds when the gating is nonlinear in Appendix A.1. We agree with the reviewer that it makes learning more difficult, as shown in Figure 4.A. However, we think that it still helps understand the inductive bias behind architectures such as LRUs, by showing that they are more akin to reproducing linear self-attention like behaviors than other types of RNNs.

We made clearer the reasons why we think our paper can have a great impact in the global response. We encourage the reviewer to consider those arguments and hope they may convince them to change their score.

We thank the reviewer for their positive feedback on our work and for pointing out sources of confusion. We provide some clarifications below that we hope will be helpful:

(a)

1. At the end of learning, not all recurrent neurons have their $\lambda$ being 0 or 1 (c.f. figure in the global answer). However, inspection of the input and output gating weights reveals that those neurons will always receive 0 as input and that they are ignored by the output gating. This is what we mean by “we observe that only perfect memory neurons (λ = 1) and perfect forget neurons (λ = 0) influence the network output”. Therefore, we can safely prune those neurons (and their corresponding input and output weights) from the network without changing its behavior. We plot the resulting network in Figure 2.A.
2. In our construction the input and output gating matrices have a characteristic 2x2 block structure. For example, if we look at the input gating, one of the matrix has top right and bottom left blocks that are equal to 0. We recover the same structure in $W_x^\mathrm{in}$ in Figure 2.A. The same thing holds for the 4 other weights of the gating matrix.
3. We will make this caption clearer. If this may help the reviewer, here is a more detailed explanation. In Figure 2.A, the output gating weights do not fully match the structure of the construction (that is one matrix with non zero values on the left side, and one with non-zero values on the right side). However, it turns out that the last 3 neurons in between output gating and the $D$ matrix (rows 11, 12 and 13 in the output gating matrices, and the corresponding columns in $D$) are functionally equivalent to one. The input weights of this neurons are described in the right part of Figure 2.B (the row and columns vectors in the right part of Fig.2.B). The nice thing is that the corresponding weights for this neuron match the expected structure: the row vector only has zero entries in the first 10 columns, and the column vector has zero entries in the last 4 rows. As a consequence, it is possible to rewrite the output gating weights (and $D$) in a way that matches the construction, while not changing at all the behavior of the network. For the mathematical details, see Appendix B.2.
4. Indeed, the caption was not sufficiently explaining the table. The input at time $t$ consists in the concatenated vectors $x_t=(x_{t,1},x_{t,2},x_{t,3})^\top$ and $y_t=(y_{t,1},y_{t,2},y_{t,3})^\top$. The instantaneous function for each output neuron can implement a polynomial of degree 4 in these terms. The table shows the coefficient associated to the terms of the polynomial implemented by the first neuron after training. Interestingly, the only terms without a negligible coefficients are $x_{t,1}^2y_{t,1}, x_{t,2}^2y_{t,1}$ and $x_{t,3}^2y_{t,1}$, where $x^2$ refer to $xx$. This is virtually identical to the polynomial implemented by a single step gradient descent algorithm (GD), as shown in the table.

(b) In all our experiments, we vary the number of neurons and as, a consequence, the number of parameters. The reviewer might be interested in Figure 5 in the appendix, in which we look at how the loss behaves as a function of the number of parameters of the different models. We are happy to provide more information if needed.

(c) We thank the reviewer for this great suggestion. We report the result of this experiment in the global answer to all reviewers.

We thank the reviewer for their valuable time reviewing our paper. We have answered part of their questions in the global response and we provide additional details in the following:

• Invertibility of the value matrix. We require the matrix $W_V$ to be invertible to decrease the number of neurons required by our construction, which of course depends on the self-attention layer we want to replicate. See the global discussion for a more comprehensive discussion on this topic.

• Section 4.1. In this experiment, we observe three things when looking at the weights of the network (we provide a visualization of those weights in the global answer) at the end of learning:

  • Most of the weights in the input gating are 0s.
  • There are some $\lambda$ values in the recurrent that are not 0 or 1.
  • Most of the weights in the output gating are 0s.

  It turns out that recurrent neurons whose $\lambda$ is not equal to 0 or 1 have corresponding weights in the input and output gating to be 0. As a consequence, they are effectively ignored by the network. If we remove those neurons, we end up with as many neurons as in the construction. The weights of the network after this pruning are displayed in Figure 2.

• Section 4.2. As highlighted in the previous point, not all recurrent neurons end up being used at the end of learning. However, having those extra neurons is important for learning as some neurons will get “killed” during learning. This is why overparametrization (in the sense of more recurrent neurons than the construction) is needed and we empirically observe that having twice as many neurons as in the construction is enough to reach 0 training loss. Yet, it turns out that only the neurons from the construction are actually used.