SRU++

SRU++ is a self-attentive recurrent unit that combines fast recurrence and attention for sequence modeling, extending the SRU unit. The key modification of SRU++ is to incorporate more expressive non-linear operations into the recurrent network. Specifically, given the input sequence represented as a matrix 𝐗 ∈ℝ^(L ×d), the attention component computes the query, key and value representations using the following multiplications,

𝐐 =𝐖^q 𝐗^⊤

𝐊 =𝐖ᵏ 𝐐

𝐕 =𝐖ᵛ 𝐐

where 𝐖^q ∈ℝ^(d^' ×d), 𝐖ᵏ, 𝐖ᵛ ∈ℝ^(d^' ×d^') are model parameters. d^' is the attention dimension that is typically much smaller than d. Note that the keys 𝐊 and values 𝐕 are computed using 𝐐 instead of 𝐗 such that the weight matrices 𝐖ᵏ and 𝐖ᵛ are significantly smaller.

Next, we compute a weighted average output 𝐀 ∈ℝ^(d^' ×L) using scaled dot-product attention:

𝐀^⊤=softmax((𝐐^⊤ 𝐊)/(√(d^'))) 𝐕^⊤

The final output U required by the elementwise recurrence is obtained by another linear projection,

𝐔^⊤=𝐖ᵒ(𝐐+α·𝐀)

where α∈ℝ is a learned scalar and 𝐖ₒ ∈ℝ^(3 d ×d^') is a parameter matrix. 𝐐+α·𝐀 is a residual connection which improves gradient propagation and stabilizes training. We initialize α to zero and as a result,

𝐔^⊤=𝐖ᵒ 𝐐=(𝐖ᵒ 𝐖^q) 𝐗^⊤

initially falls back to a linear transformation of the input X skipping the attention transformation. Intuitively, skipping attention encourages leveraging recurrence to capture sequential patterns during early stage of training. As |α| grows, the attention mechanism can learn long-range dependencies for the model. In addition, 𝐖ᵒ 𝐖^q can be interpreted as applying a matrix factorization trick with a small inner dimension d^'<d, reducing the total number of parameters. The Figure compares the differences of SRU, SRU with this factorization trick (but without attention), and SRU++.

The last modification is adding layer normalization to each SRU++ layer. We apply normalization after the attention operation and before the matrix multiplication with 𝐖ᵒ

𝐔^⊤=𝐖ᵒ layernorm(𝐐+α·𝐀)

This implementation is post-layer normalization in which the normalization is added after the residual connection.