Developing an intuition for how to replace RNNs with Self-Attention

actuallyaswin: <p>Because self-attention is capable of learning long-term relationships in sequential data, I was wondering if there was an intuitive way in PyTorch to substitute attention mechanisms in place of recurrent networks. Suppose I devise a task in which a RNN must predict the future value of a time-series input (e.g., a 5 Hz sinusoid).</p>
<pre><code class="lang-auto">torch.manual_seed(0)
class Temp(torch.nn.Module):
    def __init__(self, hidden_size):
        super().__init__()
        self.layer1 = torch.nn.GRU(1, hidden_size, 2, batch_first=True)
        self.layer2 = torch.nn.Linear(hidden_size, 1)
        self.tanh = torch.nn.Tanh()
    def forward(self, x):
        x = self.layer1(x)[0]
        x = self.tanh(self.layer2(x))
        return x
    
net = Temp(4).cuda()
n_pred = 1  # learn to predict one timestep forward
x = torch.sin(torch.linspace(0, 2*np.pi*5, 200+n_pred))  # make a sine wave
y = x[n_pred:].unsqueeze(0).unsqueeze(-1).cuda()  # target, shape=(B, L, 1)
x = x[:-n_pred].unsqueeze(0).unsqueeze(-1).cuda()  # source, shape=(B, L, 1)
criterion = torch.nn.MSELoss()
optimizer = torch.optim.Adam(net.parameters(), lr=1e-2)

n_iterations = 300
loss_vals = []
for _ in range(n_iterations):
    y_hat = net(x)
    loss = criterion(y_hat, y)
    loss_vals.append(loss.item())
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

# net is now fully trained (loss reaches about 0.002)
</code></pre>
<p>For the working snippet above, and in particular, by modifying the <strong>Temp</strong> model, how would one leverage <code>nn.MultiheadAttention</code> or <code>nn.Transformer</code> to accomplish the same task?</p>
<p>With this problem, the “input dimension” is simply 1 (a floating point <em>y</em> value on the sine wave), so I’m a bit confused how this input size can be split evenly over <em>N</em> heads in the case of multi-head attention; both <code>nn.MultiheadAttention</code> and <code>nn.Transformer</code> require the input dimension to be divisible by the number of heads. Also, when using an RNN, I can arbitrarily specify a latent dimension size of 4, but I’m not sure how that can also be done with <code>nn.MultiheadAttention</code> or <code>nn.Transformer</code>. Changing the values of “vdim” and “kdim” doesn’t make sense since usually people apply attention where the Query=Key=Value=the_input_sequence, so therefore “vdim” and “kdim” must also be 1.</p>
<p>All the examples I’ve seen on online blogs relate to neural translation, so the input sequence is made up of word embeddings; in this case, I want to understand how attention can be applied directly with real-valued time-series data. Here I set the sequence length to about 200 values, but once I get a working attention-based model, I want to try smaller or larger sequences (e.g., as low as 60 but maybe up to 300, let’s say).</p>

J_Johnson: <p>A positional encoder actually adds a sinusoidal “bias” to the embeddings. But for all intents and purposes, embeddings are n number dims. For a sinusoidal wave, you only have 2 dims, x and y.</p>
<p>So in the simple case, you’re just feeding a sequence with 2 embeddings that are however many t-steps you want in length.</p>
<pre><code class="lang-auto">alpha=10
batch_size=32
seq_len=100
beta=torch.rand((batch_size, 1, 1))*6.283
t=torch.arange((1, seq_len, 1)).repeat(batch_size,1).unsqueeze(2)
emb1 = torch.cos(t/alpha+beta)
emb2 = torch.sin(t/alpha+beta)
data=torch.cat([emb1, emb2], dim=2)
</code></pre>
<p>Where alpha is some constant that establishes the frequency. Beta would be a random vector between 0 and 2π that offsets the beginning of the sinusoidal wave, for each sample in the batch. And then to create a sequence of, say, 100 length, you’d just have <code>t</code> be a range between 0 and 99.</p>