keys, queries, values, and dot product attention?
“You need keys and queries for tokens to talk to each other,” says the popular Internet analogy. But why? There’s a lot of great analogies for how they work, but a lot less material about why we truly need them. Are there any alternatives or are these abstract concepts inevitable?
This might seem like a silly question given the utter success of the Transformer architecture in 2026. But if history is any indication, Transformers will eventually be replaced by something better. The more we understand why they work, the faster we can move beyond them.
Incredibly, we can show that the general shape of the Transformer is hard to avoid! Start with a few key design pressures, and the alphabet soup of matrices begins to look much less arbitrary. Without relying on any “token asks a question” analogies, we’ll see that queries arise from a symmetry problem, while values and attention heads appear when we replace an unwieldy dynamic weight matrix with a small set of reusable transformations.
And as we reconcile our toy model with the original (autoregressive) Transformer, we’ll end with one last fascinating connection: the MLP—the often-overlooked feedforward block—can be viewed as a key-value store of its own.
Ready to invent the Transformer for ourselves? First, we need to travel back to 2014 and ask why recurrent neural networks aren’t enough.
Why Fixed Memory Fails
There are many issues with standard recurrent neural networks (RNNs.) Some issues, like the “gradient vanishing” issue, were solved by the extremely popular LSTM (Long Short-Term Memory network) by Hochreiter and Schmidhuber [1] which is a more advanced flavor of RNN. But one core issue remains for all flavors: RNNs “squish” past inputs together into fixed memory.
To see why, recall that an RNN has a memory state that is “written to” by the current input and the previous state. This connection between states in time is where the “recurrence” comes from. Consider the following two layer RNN unrolled in time:
Here the inputs are green, the outputs are blue, and the two intermediate layers of the network that compute standard neural network activations are gray. All of the components are vectors, and the directed arrows are matrix weights that multiply their inputs. To see the network at one point in time, simply look at a single vertical slice of the diagram; in that slice, there are two fixed gray units that collectively represent the “memory” of the network at that point in time.
Let’s use a crude analogy to show why this “fixed memory” is an issue. Say you encode the sentence “I have five dollars” into the RNN’s memory, which we’ll represent with the grey rectangle below:
Now let’s extend that sentence to “I have five dollars and forty cents in my pocket”:
Yikes. After adding more information into finite memory, there is greater “competition” for real estate and some memory is overwritten. This is disastrous when you need to recall specific facts or follow very specific instructions.
You might be thinking, “Why not use dynamic memory that grows with the input sequence?”
Great intuition! This is exactly what Bahdanau et al. [2] tried in 2014 when they popularized the idea of “attention” within the RNN (yes, RNNs used attention before Transformers did!)
The idea, at a high level, is to keep the entire previous history of the RNN states as our “expanding memory”, which naturally has the property of growing with the input sequence. Consider the following updated diagram (this is not the exact architecture in the Bahdanau paper, but an analogy):
Note the new connections in red; these connect the gray states in the RNN layers to every previous input in time. Before adding these connections, each state was forced to compress all of the historical information and pass it along in the left-to-right “recurrent” connections. But this compression is no longer needed since every state now has direct access to the entire history of inputs, aka our growing memory!
There is one big challenge that remains: training speed. At training time, we have each input sequence available up front. But generating the final Nth output requires N sequential steps in time given the recurrent dependencies; with long input sequences, we have long sequential computations that cannot be parallelized by GPUs.
The idea in the landmark paper by Vaswani et al. [3] is this: What if we can remove those recurrent left-to-right connections? What if the red connections are all you need? See the following diagram that represents this idea at a high level:
Let’s stack the recurrent and non-recurrent architectures side by side, and compare their total number of compute steps by putting ordered labels for each step:
Note how the non-recurrent model needed only 2 compute steps, since it can compute each layer entirely in parallel once the previous layer is computed. On the other hand, the recurrent model needed 5 compute steps due to the recurrent dependencies within each layer. As the sequence gets longer, the non-recurrent model would stay at 2 steps while the recurrent model’s steps would grow forever with the sequence. It’s not looking good for recurrence!
Let’s pivot to the non-recurrent model on the right. Now, we run into our next challenge: How on Earth should we pick the weights for these red connections?
Transformers and Dynamic Weights
If you look at the diagram of our non-recurrent network, it looks just like an ordinary neural network with two layers and four units per layer. With such a network, we might ask: Why not learn fixed weights like we do with any other network?
But unlike this fixed diagram, sequences are not fixed during training or prediction. The network could encounter an input sequence of size 2 or it could encounter an input sequence of size 2000, and our diagram could grow indefinitely to the right with ever more units and weights. So how do we set those incoming new weights?
What we need is a function to generate new weights on the fly, with parameters that we can set during training. We need dynamic weights!
To identify a good function, let’s zoom into one particular unit with a length 3 sequence:
First, an important note on a new diagram addition: To make the diagram complete, we needed to add skip connections. Through these, the inputs x1, x2, and x3 are added back to the outputs of O1, O2, and O3 respectively (and this process repeats again at the next layer when O1, O2, and O3 become inputs themselves.) These “skips” free intermediate transformations from having to preserve the inputs and let them focus on the much easier task of additively adjusting those inputs. Skip connections are a critical performance optimization that were pioneered in the famous ResNet architecture by He et al. [4].
Now back to our function to generate weights dynamically. First, we need to decide what the weights should be a function of.
We could start by defining a given weight as a function of the output of the weight’s source unit as well as the position of that source unit in the left-right sequence. To make things simple, I’ll combine “input” and “position” together and just say “input” going forward (this merging can be implemented by encoding position into the input directly; to see how this might be done, read about sinusoidal encodings in the original Transformer architecture.) Then, our diagram looks something like this:
There’s one obvious downside here: If x1 needs to be “important” to O3 via a “large” weight, then it is forced to also be important to O2 and O1 since they all share the same value for their x1 weights. This means that nearby blue units within a layer will compute very similar things, defeating the flexibility of this architecture to model unique concepts. To fix this, we need to break the “symmetry” and make each of the weights of O1, O2, and O3 unique from one another.
We could break the symmetry by making a given weight a function of both the source unit’s output and the end unit’s output; however, using the end unit’s output directly is circular since we first need the weight to compute the end output.
To get around this circularity, note that each unit’s unique purpose is to modify its input stream; for example, O3 is uniquely responsible for modifying x3 via the skip connection. Since x3 and O3 are uniquely intertwined, x3 is a natural candidate for the second, symmetry-breaking argument to generate O3’s weights!
When we update the diagram it is clear that all weights are now unique:
If you have a good eye, you might start to spot the “key” and “query” of the Transformer architecture already! But if not, no worries; we’ll build up to that more formally.
One last call out: We could also break symmetry by making the weights zero for every input except the one that a unit modifies, but then you get the following diagram that shows this is effectively a state-less network:
In other words, symmetry breaking is a necessary but not sufficient condition. We also need some non-zero interactivity between units and other units from different time steps.
Keys, Queries, and Values Emerge Naturally
Now that we’ve decided to generate each weight via a function of two arguments, we need to decide what this function actually looks like.
Remember that each weight is a matrix that multiplies input vectors into output vectors, so our function actually needs to be matrix-valued. But that introduces a new problem: How do we make these matrices dynamic without blowing up the number of parameters in our model?
To see why parameter explosion is real, let’s write our function in matrix form, using the weight between x1 and O3 as an example:
Here, each position (i, j) of the (d x d) matrix has its own function. If we parametrize each function separately, that’s d-squared separate sets of parameters that need to be fit. When the dimension d is in the range of hundreds, that’s >100K sets of parameters that have to be learned. Hard pass!
We could vastly reduce that number by only fitting a diagonal matrix:
However, this is still hundreds of sets of function parameters to fit, and we haven’t even gotten into the size of those sets. Also, note that diagonal matrices perform element-wise multiplication on their inputs; if parts of the final vector need to be zero for reasons like sparsity, then functions in those diagonal elements need to be very close to zero or zero exactly—which puts a lot of burden on the functions themselves.
To see an alternative, let’s rewrite that diagonal matrix as a linear combination of one-hot matrices:
What if instead of using those one-hot matrices in the sum, we use any matrices we wanted? We could make each matrix a parameter to be learned, and then pick the final number of learnable matrices so that the total nested parameter count is reasonable. Then we could rewrite our weight between x1 and O3 as a new matrix sum:
Awesome; we just replaced a painful dynamic matrix with a small number of static matrices and dynamic coefficients! Here, the V’s are the static matrices that we learn as free parameters and the functions in the sum are our dynamic scalar-valued coefficients. Keep the V’s in mind; they’ll show up in the attention “value” calculations later.
Finally, we need to define the scalar-valued functions inside that sum above. To do that, let’s additively decompose one of the functions into two single argument “non-interaction” functions sandwiching a pure interaction function:
The downside to having non-interaction functions is the same downside we encountered with weight symmetries from the last section: If the left term is large, then it’ll be large for all units connected to x1 since x1 is the only dependency in the term. We want to break this symmetry so let’s keep only the middle interaction function v, which I’ll call the “attention” function from now on.
One candidate for the attention function was introduced in Bahdanau et al. [2]:
You may recognize this as a simple one layer neural network with hyperbolic tangent activation, where the final result is reduced to a scalar value via a dot product.
There is one downside to using the tanh function here: It is only weakly interacting through its non-linearity. The function’s contours get squashed non-linearly, but their basic structure still looks like a linear sum. To see this visually, compare the similarity in contour plots for tanh of x1 plus x3 vs. the straight sum of x1 and x3 (where we’ll make x1 and x3 1D inputs for ease of visualization):
Let’s really highlight the issue with an example: Suppose all the units in our network use one of the functions above, and one unit has inputs x1 and x3.
If this unit needs to output a large positive value, then there’s only one region that satisfies this on either plot: the top right half. But if another unit shares the same x1 input and requires x1 to be on the left side (i.e. negative) to achieve its goal, then you’re stuck on the first unit’s requirement unless x3 is a huge positive number to compensate. Both units’ requirements are, to a degree, incompatible. This network is less flexible in what its different units can model when they share some inputs.
On the other hand, look at the contour plot of the product function below:
Now, there are two distinct regions where the final output is large, not just one! If another unit shares the same x1 input and needs x1 to be negative, not a problem; the first unit can still output a large positive value if x3 is negative. Both units’ requirements are no longer as incompatible as before. This is a subtle mathematical argument for how “interactivity” (for example, via a product) matters in allowing units to model unique things despite having overlap in their inputs.
You might ask, “Why not use a deeper neural network to model a more interactive attention function?” This is quite costly in practice because the number of attention calculations scales quadratically with sequence length. And consider that we don’t make each layer of a traditional neural network complicated either; often we use simple non-linear functions like ReLU (rectified linear units) and let additional layers iteratively build up complexity.
Could we do something similar here, i.e. go with a simple yet interactive attention function and build up complexity over layers? If so, the product function is a great candidate—and GPUs love it too!
There is one tweak we need to make since our attention inputs are multi-dimensional vectors, not 1D inputs. Instead of a scalar product, we need a dot product. More generally, you might want to compute the dot product in a particular sub-space, which results in a bilinear form that looks like this:
Note that this reduces to a plain dot product when A is the identity matrix, so this form is more general.
There is another consideration that we need to start discussing here: the computation “cache” (what you may have heard as the “key-value cache” in modern architectures.) We can avoid a significant number of attention computations by caching matrix multiplications (e.g. A times x1 above) for past inputs, since we reuse these past calculations for every new time step at every layer.
However, the size of this cache becomes a pretty big pain point. Example: If we have a sequence length of 5,000 (quite a bit smaller than many standard LLM conversations), 50 layers, 20 different V matrices –> 20 different dot product functions (see our matrix sum definition earlier), cached matrix-vector products of dimension 1000, and 2 bytes per floating point number, then the final size of our cache is: 5000 x 50 x 20 x 1000 x 2 = 10 GB. Ouch.
Larger caches mean more GPU cost and memory overhead. One optimization is to simply reduce the dimension of the cached products from 1000 to a lower dimension r, as long as it doesn’t significantly impact model accuracy. But this requires the final dot product to now be in an r-dimensional space rather than a 1000-dimensional one.
We can lower that dot product dimension by factorizing the matrix A into two matrices that “project” the vectors x3 and x1 into the r-dimensional space:
Here, Wq and Wk are (r x 1000) matrices, and the right-most expression is the desired dot product between two r-dimensional vectors. If r is only 200, then the final cache size goes from 10 GB to 2 GB—a 5x reduction!
Now for the grand reveal: The left and right terms in the new dot product are nothing but the “query” and “key” in the Transformer architecture, and the projection matrices are the same learnable matrices in the original paper! (One caveat is that the Transformer architecture adds scaling for computational stability, hence the term “scaled dot product attention”. But for the rest of this article, I’ll focus on the shape of the architecture rather than on training optimizations like scaling.)
We can now substitute this attention function into the coefficients of our earlier matrix sum, with different parameters for each j:
Quick summary of where we are: The function “f” on the left spits out the weight matrix between input x1 and unit O3; this matrix is equal to a sum involving H scalar attention functions as coefficients, and based on our discussion, we’ve chosen to use dot product attention for each function. Then we can represent the sum total of all weighted inputs into O3 with a new sum s3:
The term WQ * x3 is the query for x3, the term WK * xi is the key for xi, and the term V * xi is the value for xi. The subscript j on the matrices denotes a particular attention “head” among the H heads; each head has a unique learnable query, key, and value projection matrix. This is starting to look just like a Transformer!
From Our Attention to Transformer Attention
In our earlier example, the outer sum is over a paltry sequence of size 3. But this sum blows up in size when the sequence gets really long.
One way to address this is a standard machine learning trick called normalization: Take the existing coefficients in the weighted sum and transform them so that the transformed coefficients sum to 1 regardless of the previous number and size of coefficients.
We might also want sparsity in the transformed coefficients; in the domain of language, there are hundreds of irrelevant words that can add up weight quickly, so we want to aggressively squash everything but a few coefficients.
To see how we should add normalization to our formula, note that we need to do this normalization in the sum over the sequence length, not in the sum over the heads. In that case, we need to flip the order of the sums as follows:
Now, the term inside the outer parentheses looks like a sum over the sequence length L with coefficient weights (key-query dot products) on the final value vectors—just as we wanted. We then apply a normalization transformation to those dot product coefficients. The natural choice of transformation for achieving soft sparsity is the softmax transformation, which exponentiates each coefficient and divides that result by the sum of all exponentiated coefficients:
This is called the “softmax” because it aggressively pushes the largest (max) coefficient to 1 and squashes all other coefficients towards 0. The tau parameter in the exponent controls how aggressive this transformation is.
Okay, all of this is great—but where are the matrices Q, K, and V that the article title promised us?
To match our expression with the iconic Transformer equation, let’s temporarily ignore the outer sum over heads by picking a particular head j and rewrite the normalized sum at position L for that head:
Here, q is the query vector for the Lth unit, and the rows of matrices K and V are the key and value vectors respectively. The subscript 1:L indicates that we only use the keys and values for positions 1 through L, which reflects the fact that the Lth unit is only connected to previous units in time. The product between q and K^T creates a vector containing every dot product between q and a key in K, and the softmax on top normalizes the final dot product scores. Afterwards, we do a weighted sum of value vectors in V using the resulting softmax scores.
We can extend this into a single consolidated matrix expression for the sum at every position, not just position L:
This is exactly the iconic Transformer equation (minus dimensional scaling) but for an autoregressive Transformer! Here, each row of S is the final sum at that position, and each row of Q, K, and V are the respective query, key, and value vectors at that position. Since we are using the full matrices Q, K, and V, we need a masking matrix M to zero out invalid dot products (such as the dot product between a query at position p and a key at a later position) by adding in negative infinity to those entries; the softmax then squashes those results to zero.
Why write this with matrices instead of our original sum if the two expressions are identical? The answer is simple: GPUs love matrix multiplication. It is often more efficient to rewrite a computation as a matrix multiplication, even if that means doing some unnecessary work like computing dot products which will be masked out anyway. That being said, I find the sum notation easier to use for explanations, so I’ll stick with it for the rest of the article.
Let’s close out the section by returning to the sum over attention heads. You may have noticed that while we sum over the H attention heads, Vaswani et al. [3] and implementations like PyTorch’s “MultiheadAttention” concatenate the outputs across each head instead—and multiply the concatenation by a “mixing matrix” Wo that collapses the concatenation back to the previous model dimension.
This might look like a cosmetic difference; why concatenate if you are just going to collapse back down?
The trick is caching, but now with a focus on value vectors: We need the mixing matrix to allow us to compress those value vectors (V * x’s) so our cache doesn’t explode. Let’s see what happens to our double sum expression if we concatenate the output of each head j into a block vector instead of summing over the heads:
For the sake of readability, I’ve substituted in alphas for the previous normalized coefficients so there aren’t a bunch of exponentials everywhere.
Now let’s left-multiply this long block vector by the mixing matrix Wo. Let’s first re-write Wo into a matching block matrix form and then perform the multiplication:
Note how similar the final result is to our original double sum! The only difference is the block matrix left multiplying V; when each block matrix is the identity matrix, the new expression is totally identical to our original expression.
You may ask, “Why not just absorb the Wo block matrices into the V matrices?” At first glance, the two matrices seem redundant if you can collapse them into a single one. But the reason we’d want to keep them separate is the same reason we didn’t use a single matrix A for our attention dot product: lowering our cache size.
Let’s say we did absorb everything into a single matrix V. If our model dimension is 1000 and we use the same conditions from our key-based cache example earlier, then the cumulative cache size for our value vectors is an unwieldy 10 GB. But if V has the shape (r x 1000) where r < 1000, then the final dimension of our value vectors will be r < 1000 and our cache size can be considerably smaller.
Herein lies the need for the Wo matrices; after we’re done having our fun in lower dimensions, we need to multiply the final values by a matrix of shape (1000 x r) to “up project” back to the larger model dimension of 1000.
The Transformer MLP: The Other Key-Value Store
Everything we did earlier was to define the weights in our network. If we isolate one input xi in the sum of weighted inputs from our last expression, then its weight into unit L is equal to:
Here, each alpha coefficient is unit L’s softmax normalized scalar for input i and a given head j.
Once you have all the weights, all you need to do is to sum up the weighted inputs into each unit, compute the position-wise non-linearity, and add back to the residual stream via the skip connection to update the positional representation. Voila, there’s your neural network!
These steps can be visually represented by the following recap diagram, using the same length 3 sequence for simplicity:
You may have noticed that Vaswani et al. [3] introduce two new weight matrices (W1 and W2) that sandwich the non-linearity in the blue units. At first glance, these additions feel like arbitrary add-ons. If we’ve already done all this work to define dynamic weights via attention, then do we really need two new matrices around the non-linearity?
It turns out that W2 isn’t there just for the sake of it. Geva et al. [5] give us important perspective for understanding why: The Transformer MLP (feed-forward block) acts like a key-value store in its own right.
To see this, let’s actually show where the “keys” and “values” would live in an example MLP network:
In this network, the middle layer has red, green, and blue units that each model different patterns. The incoming (first layer) weights of the same color represent the “key” vectors for each pattern, and the outgoing (last layer) weights of the same color represent the “value” vectors for each pattern. Geva et al. [5] demonstrate what some of these “patterns” could look like; for the model they trained, one example of a learned pattern was whenever the input sentence ended in the word “substitutes”, and another was whenever the input sentence contained a “part of” relation.
Now let’s walk through what the “keys” and “values” actually do here. When the input aligns closely with the key vector (incoming weights) for a given pattern, the dot product between the key and input is high, and the unit associated with that pattern outputs a large number. That number then multiplies the value vector (outgoing weights) for the pattern and this weighted value vector is finally added to the output layer.
So back to our extra parameters: What does this tell us about why we need W2?
Recall that W2 is just the matrix multiplied to the output of the non-linearity in the original Transformer. In our MLP example, W2 is exactly equal to the matrix of weights in the final layer, i.e. a matrix where each column is one value vector.
If we get rid of W2, our diagram would instead look something like this:
Each unit in this network is writing to a single coordinate of the residual stream (since we add the output layer directly back to the residual stream.) But there are two reasons why this is bad:
- If we want our units to learn arbitrary patterns, then each of those arbitrary patterns might not map cleanly to single coordinates of the “semantic space” at that layer. In that case, it is far more expressive to allow each unit to write arbitrary vectors (values) into the residual stream rather than coupling them to individual output coordinates.
- We want our MLPs to model as many patterns as possible, which pushes us to have more units in the MLP network relative to the model dimension in the residual stream (where the latter is kept smaller to avoid gnarly attention computations that don’t benefit as much from high dimensions.) In that case, a 1:1 correspondence of coordinates isn’t even possible; we need a matrix to “down project” back to the smaller model dimension anyways.
We’ve spent all this time discussing W2; what about W1, the extra matrix multiplying inside the non-linearity?
This may actually be the more arbitrary matrix. In our setup, W1 could indeed be absorbed into our dynamic weight matrices (specifically, Wo from the previous section) and dropped as an extra parameter. But there is a reason this can’t be done in the original Transformer: In that setup, there are two intermediate steps between the attention output and the non-linearity that prevent consolidation of matrices.
First, the Transformer splits our single big skip connection into a skip connection between the input and the sum output as well as a second skip between the sum output and the non-linearity output. See the modified diagram with the split skip connection:
Once you have this additional skip connection, you need a separate weight matrix applied on the sums of the output from the attention (capital sigma) blocks and the residuals from the first skips. The weight matrix applied to that sum of outputs can no longer be cleanly absorbed, so you need a separate parameter.
Secondly, there is an additional scaling transformation between the sum of attention outputs and the non-linearity that prevents absorption. While these scaling transformations are important, they are training optimizations, and so I’ve omitted them to focus purely on architectural shape here.
Recapping the Logic
Pat yourself on the back; you just invented the Transformer! Here’s a recap of the steps we took to get here:
- We needed to access every past state of the sequence without compression loss. This required direct connections to past states, aka “attention”.
- We needed parallelism for fast GPU training. This required dropping recurrence.
- We needed a way to set new weights for future inputs without introducing weight symmetries across the network. This required making the weights a function of the source unit and the end unit (both unique per weight), where we used the residual stream’s value as a stand-in for the “end unit”. These two function arguments are the “key” and “query” prior to projection.
- We needed to make the weight-generating function avoid symmetries by encoding pure interactions between the arguments. A GPU-efficient choice for this interaction function was the dot product between the key and query arguments.
- We needed to reduce our cache size for the reusable matrix-vector multiplies in the dot product, which required computing the dot product in lower dimensions. This required projection matrices Wk and Wq to project the key and query into that lower dimensional space.
- We needed to reduce our search over the entire d^2 space of linear operations (matrices) into a search over a smaller number H of possible operations that we can select from via our dot product weights. This is represented by H different value projections V representing our H attention heads.
- We needed to make the final weights on the value projections sparse so noise doesn’t blow up over larger sequences. This required softmax normalization of the dot product coefficients over sequence length, per head.
- We needed to reduce our cache size for the reusable matrix-vector multiplies between the value projections V and the inputs. This required the Vs to “down project” the final value vectors into a lower dimensional space, and a separate mixing matrix Wo to “up project” back into the model dimension.
- We needed a larger “feature space” for our non-linearities to learn lots of patterns, while allowing these features to write arbitrary vectors to the lower dimensional residual stream. This required a new matrix W2 that multiplies the output of the non-linearity.
Why Transformers Aren’t Inevitable
The only thing inevitable in AI/ML is one architecture replacing another, and Transformers are no exception to the rule.
But why will Transformers be replaced if they work so well today?
For starters, Transformers have one massive downside: Computations scale quadratically with sequence length. You could implement sliding window attention to try and get around this, but then you lose the ability to recall pin-point facts or instructions hidden within longer context.
You may be wondering, “But most of the attention scores are basically zero due to the softmax normalization. Do we need to calculate all n^2 scores?”
That’s a great insight, and attempts at sparse attention mechanisms have indeed been made. However, despite the major reduction in calculations, these mechanisms counterintuitively run slower than full quadratic attention due to GPU memory bottlenecks becoming the bigger pain point, as discussed by Dao et al. [6].
This brings us to another point: Transformers are benefitting enormously from being tailor-made for GPUs. Many of their design decisions revolve around the quirky memory and compute patterns of GPUs, and anything that could dethrone Transformers has to ultimately pass the “GPU smell test”. In other words, we are trapped in a hardware local minima and are missing out on much better algorithms because of it. Human intelligence doesn’t run on GPUs, so why should we settle for Transformers?
References
[1] S. Hochreiter and J. Schmidhuber, “Long Short-Term Memory,” Neural Computation, vol. 9, no. 8, pp. 1735-1780, 1997.
[2] D. Bahdanau, K. Cho, and Y. Bengio, “Neural Machine Translation by Jointly Learning to Align and Translate,” International Conference on Learning Representations (ICLR), 2015.
[3] A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, Ł. Kaiser, and I. Polosukhin, “Attention Is All You Need,” Advances in Neural Information Processing Systems 30 (NIPS), pp. 5998-6008, 2017.
[4] K. He, X. Zhang, S. Ren, and J. Sun, “Deep Residual Learning for Image Recognition,” IEEE Conference on Computer Vision and Pattern Recognition (CVPR), pp. 770-778, 2016.
[5] M. Geva, R. Schuster, J. Berant, and O. Levy, “Transformer Feed-Forward Layers Are Key-Value Memories,” Conference on Empirical Methods in Natural Language Processing (EMNLP), pp. 5484-5495, 2021.
[6] T. Dao, D. Y. Fu, S. Ermon, A. Rudra, and C. Ré, “FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness,” Advances in Neural Information Processing Systems 35 (NeurIPS), pp. 16344-16359, 2022.