Understanding “You Solely Cache As soon as” | by Matthew Gunton | Jun, 2024

To grasp the modifications made right here, we first want to debate the Key-Worth Cache. Within the transformer we’ve 3 vectors which might be vital for consideration to work — key, worth, and question. From a excessive degree, consideration is how we cross alongside vital details about the earlier tokens to the present token in order that it might predict the subsequent token. Within the instance of self-attention with one head, we multiply the question vector on the present token with the important thing vectors from the earlier tokens after which normalize the ensuing matrix (the ensuing matrix we name the eye sample). We now multiply the worth vectors with the eye sample to get the updates to every token. This information is then added to the present tokens embedding in order that it now has the context to find out what comes subsequent.

We create the eye sample for each single new token we create, so whereas the queries have a tendency to alter, the keys and the values are fixed. Consequently, the present architectures attempt to cut back compute time by caching the important thing and worth vectors as they’re generated by every successive spherical of consideration. This cache is known as the Key-Worth Cache.

Whereas architectures like encoder-only and encoder-decoder transformer fashions have had success, the authors posit that the autoregression proven above, and the velocity it permits its fashions, is the rationale why decoder-only fashions are probably the most generally used right now.

To grasp the YOCO structure, we’ve to begin out by understanding the way it units out its layers.

For one half of the mannequin, we use one kind of consideration to generate the vectors wanted to fill the KV Cache. As soon as it crosses into the second half, it’ll use the KV Cache completely for the important thing and worth vectors respectively, now producing the output token embeddings.

This new structure requires two varieties of consideration — environment friendly self-attention and cross-attention. We’ll go into every beneath.

Environment friendly Self-Consideration (ESA) is designed to realize a relentless inference reminiscence. Put in a different way we wish the cache complexity to rely not on the enter size however on the variety of layers in our block. Within the beneath equation, the authors abstracted ESA, however the the rest of the self-decoder is constant as proven beneath.

Let’s undergo the equation step-by-step. X^l is our token embedding and Y^l is an middleman variable used to generate the subsequent token embedding X^l+1. Within the equation, ESA is Environment friendly Self-Consideration, LN is the layer normalization operate — which right here was at all times Root Imply Sq. Norm (RMSNorm ), and eventually SwiGLU. SwiGLU is outlined by the beneath:

Right here swish = x*sigmoid (Wg * x), the place Wg is a trainable parameter. We then discover the element-wise product (Hadamard Product) between that consequence and X*W1 earlier than then multiplying that entire product by W2. The aim with SwiGLU is to get an activation operate that may conditionally cross by way of completely different quantities of knowledge by way of the layer to the subsequent token.

Now that we see how the self-decoder works, let’s go into the 2 methods the authors thought of implementing ESA.

First, they thought of what is known as Gated Retention. Retention and self-attention are admittedly very related, with the authors of the “Retentive Community: A Successor to Transformer for Massive Language Fashions” paper saying that the important thing distinction lies within the activation operate — retention removes softmax permitting for a recurrent formulation. They use this recurrent formulation together with the parallelizability to drive reminiscence efficiencies.

To dive into the mathematical particulars:

We have now our typical matrices of Q, Okay, and V — every of that are multiplied by the learnable weights related to every matrix. We then discover the Hadamard product between the weighted matrices and the scalar Θ. The aim in utilizing Θ is to create exponential decay, whereas we then use the D matrix to assist with informal masking (stopping future tokens from interacting with present tokens) and activation.

Gated Retention is distinct from retention by way of the γ worth. Right here the matrix Wγ is used to permit our ESA to be data-driven.

Sliding Window ESA introduces the concept of limiting what number of tokens the eye window ought to take note of. Whereas in common self-attention all earlier tokens are attended to ultimately (even when their worth is 0), in sliding window ESA, we select some fixed worth C that limits the scale of those matrices. Which means throughout inference time the KV cache may be lowered to a relentless complexity.

To once more dive into the maths:

We have now our matrices being scaled by their corresponding weights. Subsequent, we compute the top just like how multi-head consideration is computed, the place B acts each as a causal map and in addition to verify solely the tokens C again are attended to.