Explanatory (non-vectorized) code for how attention works

almost-attention.py

# This code doesn't work, and isn't intended to.
# The goal of this code is to explain how attention mechansisms work, in code.
# It is deliberately not vectorized to make it clearer.

def attention(self, X_in:List[Tensor]):
    # For every token transform previous layer's out
    for i in range(self.sequence_length):
        query[i] = self.Q * X_in[i]
        key[i]   = self.K * X_in[i]
        value[i] = self.V * X_in[i]

    # Compute output values, one at a time
    for i in range(self.sequence_length):
        this_query = query[i]
        # how relevant is each input to this out?
        for j in range(self.sequence_length):
            relevance[j] = this_query * key[j]
        # normalize relevance scores to sum to 1
        relevance = scaled_softmax(relevance)

        # compute a weighted sum of the values
        out[i] = 0 
        for j in range(self.sequence_length):
            out[i] += relevance[j] * value[j]

    return out

Q,K,V are matrices to be learned.