Skip to content

machinelearning:: MSA treats all elements as equally important, but in our case we would like to promote a recency bias. Gated attention allows us to apply a time-discounted weight to each entry in the sequence

Summary

  • creates some random features in order to learn a linear-time approximation of attention-softmax, turning it from \(O(n^2)\) to \(O(n)\).
  • Calculating attention for a single query takes \(O(N)\) time for a sequence of length N. In MSA, we do so for every element in the sequence, resulting in an \(O(N^2)\) computation when the computation cannot be paralellized, like in autoregressive decoding
  • Gating enables us to decide what information in a sequence to propagate, and with what weight.
    • without gating, attention weights are solely a result of how similar a query and key are, irrespective of positional importance.
      • positional encodings allow for positional awareness, but not much positional understanding.
  • The gating mechanism used was $$ g_t = sigmoid(w_g \cdot x_t)$$ where \(g_t \in [0, 1]\) is a scalar indicating the importance of timepoint t, \(x_t\) is the \(t\)'th element in the sequence, and \(w_g\) is a learned parameter with the same shape as \(x_t\).

Inuitively, this means we learn a vector \(w_g\) where entry \(i\) is a scalar indicating the relative importance of feature \(i\)

We apply this gating in an autoregressive manner to obtain output \(S_t\) from \(S_{t-1}\) $$ S_t = g_t S_{t-1} + (1 - g_t)Attn(x_t) $$

so here \(g_t\) indicates the importance of the past vs current state.

I could apply this gated mechanism either at the feature level (directly on x_t), or after the attention weight is obtained.

Ideas

  • my big idea: apply a gating mechanism alongside a time-discounting factor.
    • we let the gating mechanism learn how to scale the time-discounting factor for each feature
  • time discounting factor: let the time discounting factor \(\delta(t) = log(t+c)\) where \(t\) is the time from the first biopsy and \(c\) is some scalar.
    • then we can weight the feature vector from timepoint \(t\) by \(\delta\) such that \(\bar x_t = \delta (t) x_t\)
      • or maybe it could be the time from the current timepoint T. so \(distance = T - t\)
    • Furthermore, we can apply the gating idea to learn a feature-importance vector \(v \in [0,1]^d\) (or maybe no cap on size) indicating how relevant each feature is over time, just like in gated attention
    • apply that on top of the time-discounting factor so that \(\bar{x} = \delta(t) v x_t\)
  • see Gated Recurrent Unit for specifics. This paper is a nice proof-of-concept that we can merge gating with attention

Abstract

pdf (openreview.net)

Transformers are state-of-the-art models for a variety of sequence modeling tasks. At their core is an attention function which models pairwise interactions between the inputs at every timestep. While attention is powerful, it does not scale efficiently to long sequences due to its quadratic time and space complexity in the sequence length. We propose RFA, a linear time and space attention that uses random feature methods to approximate the softmax function, and explore its application in transformers. RFA can be used as a drop-in replacement for conventional softmax attention and offers a straightforward way of learning with recency bias through an optional gating mechanism. Experiments on language modeling and machine translation demonstrate that RFA achieves similar or better performance compared to strong transformer baselines. In the machine translation experiment, RFA decodes twice as fast as a vanilla transformer. Compared to existing efficient transformer variants, RFA is competitive in terms of both accuracy and efficiency on three long text classification datasets. Our analysis shows that RFA’s efficiency gains are especially notable on long sequences, suggesting that RFA will be particularly useful in tasks that require working with large inputs, fast decoding speed, or low memory footprints