In this section, we highlight some simple ways to further accelerate the parallel generation algorithm. For example, note that one can reduce the attention's memory cost by computing the diagonal attention simply as elementwise (rather than pairwise) dot-products. That is, given two input embedding sequences of shapes \((b,t,l,d)\) and \((b,1,l,d)\) where \(t\) is the number of timesteps ahead, \(b\) is batch size, \(l\) is sequence length, and \(d\) is embedding dimension, we do not need to compute their pairwise attention of shape \((b,t,l,l)\), we only need to compute the attention for the paired elements along the diagonal of shape \((b,t,l)\). Additionally, to avoid generating continuations for all of the tokens (for example, if one wanted to apply a value function to determine where thoughts would be most useful), one can index into this generated attention mask. Notably, however, this also requires manipulation of the other inputs during the forward pass such as the positional embeddings.