Google

In this model, the variables $$\mu_i$$ and $$\nu_i$$ serve as representations of the sequence history up to position $$i$$. They are calculated using recurrent forms, effectively summarizing past data: $$\mu_i = \mu_{i-1} + \mathbf{k'}_{i}^{\mathrm{T}} \mathbf{v}_{i}$$ and $$\nu_i = \nu_{i-1} + \mathbf{k'}_{i}^{\mathrm{T}}$$.

Recurrent Computation of $$\mu_i$$ and $$\nu_i$$ in Linear Attention

A primary benefit of the recurrent model utilizing $$\mu_i$$ and $$\nu_i$$ is that it eliminates the need to retain all past queries and values. By relying exclusively on the latest representations, $$\mu_i$$ and $$\nu_i$$, the computational cost of each individual step remains constant. Consequently, this allows the model to be easily extended to handle very long sequences.

Computational and Memory Efficiency of Linear Attention's Recurrent Method

A system processes a sequence of inputs by maintaining two state variables, μ and ν. These variables are updated at each step `i` according to the following recurrent formulas, using the state from the previous step (`i-1`) and the current inputs (`k'_i` and `v_i`). Given the formulas and a two-step input sequence, calculate the final values of the state variables.

Recurrent State Update Calculation

A system maintains two state variables, μ and ν, which are updated at each step `i` according to the following rules:

μ_i = μ_{i-1} + k'_iᵀ v_i
ν_i = ν_{i-1} + k'_iᵀ

Assuming the initial states μ₀ and ν₀ are both zero vectors, write the expanded mathematical expression for μ₃ solely in terms of the input vectors k' and v from steps 1, 2, and 3.

Unrolling a Recurrent State Update

In this variant of linear attention, the final output is calculated by combining the current transformed query vector $$\mathbf{q'}_i$$ with the accumulated state variables $$\mu_i$$ and $$\nu_i$$. The numerator is the product of the query and the key-value state $$\mu_i$$, while the denominator is the product of the query and the key state $$\nu_i$$, serving as a normalization term. The formula is: $$\text{Att}_{\text{linear}}(\mathbf{q}_i, \mathbf{K}_{\le i}, \mathbf{V}_{\le i}) = \frac{\mathbf{q'}_i \mu_i}{\mathbf{q'}_i \nu_i}$$ This approach replaces the standard Softmax operation with simpler matrix-vector products, leading to computational savings.

Linear Attention Output Calculation

A sequential model updates two history-representing variables, μ and ν, at each step `i` using the following rules: 

μ_i = μ_{i-1} + k'_i^T * v_i 
ν_i = ν_{i-1} + k'_i^T 

Consider the update at a single step `i`. If the input value vector `v_i` is a zero vector (a vector of all zeros), but the input key vector `k'_i` is a non-zero vector, what is the outcome of the update from step `i-1` to step `i`?

Learn Before

Related