On Coding Your First Attention

<p>

While it’s not necessarily the case that you must code the attention block of a transformer from scratch to understand how it works, yet it sure is the closest you can get to having a first-principles understanding of why/how transformers behave the way they do.</p>
<p><a href="https://cdn-uploads.huggingface.co/production/uploads/6438a9027de34e8ea7e4b257/4MMtJDefZBU8dpmHana6B.png" rel="nofollow"><img alt="image/png" src="https://cdn-uploads.huggingface.co/production/uploads/6438a9027de34e8ea7e4b257/4MMtJDefZBU8dpmHana6B.png"/></a></p>
<p>@karpathy covered attention in detail in his nanoGPT video (strongly recommend watching). Now I would like to share some thoughts and experience in writing my first attention.</p>
<p>First let’s zoom out quickly and explain what attention is in transformers:
Attention in transformers is a communication mechanism that allows the model to focus on different parts of the input sequence when making predictions. </p>
<p>It assigns weights to each input token based on its relevance to the current context, enabling the model to weigh information selectively. This mechanism helps transformers capture long-range dependencies and contextual information effectively.</p>
<p>The  official AIAN paper introduced two commonly used forms of attentions: Scaled Dot-Product Attention (also known as Self-Attention) and a stack of self-attention blocks known as Multi-Head Attention.</p>
<h1 class="relative group flex items-center">
<a class="block pr-1.5 text-lg md:absolute md:p-1.5 md:opacity-0 md:group-hover:opacity-100 md:right-full" href="#the-code" id="the-code" rel="nofollow">
<span class="header-link"><svg aria-hidden="true" class="text-gray-500 hover:text-black dark:hover:text-gray-200 w-4" height="1em" preserveaspectratio="xMidYMid meet" role="img" viewbox="0 0 256 256" width="1em" xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink"><path d="M167.594 88.393a8.001 8.001 0 0 1 0 11.314l-67.882 67.882a8 8 0 1 1-11.314-11.315l67.882-67.881a8.003 8.003 0 0 1 11.314 0zm-28.287 84.86l-28.284 28.284a40 40 0 0 1-56.567-56.567l28.284-28.284a8 8 0 0 0-11.315-11.315l-28.284 28.284a56 56 0 0 0 79.196 79.197l28.285-28.285a8 8 0 1 0-11.315-11.314zM212.852 43.14a56.002 56.002 0 0 0-79.196 0l-28.284 28.284a8 8 0 1 0 11.314 11.314l28.284-28.284a40 40 0 0 1 56.568 56.567l-28.285 28.285a8 8 0 0 0 11.315 11.314l28.284-28.284a56.065 56.065 0 0 0 0-79.196z" fill="currentColor"></path></svg></span>
</a>
<span>
		The Code
	</span>
</h1>
<p>Now, attention as for most deep learning algorithms boils down to a math equation. So writing the code can get really trivial especially with a deep learning framework like PyTorch. Below is what's called a Single Head Attention:</p>
<p><a href="https://cdn-uploads.huggingface.co/production/uploads/6438a9027de34e8ea7e4b257/c-RzcFcoyRVFqYCSxgvsS.png" rel="nofollow"><img alt="image/png" src="https://cdn-uploads.huggingface.co/production/uploads/6438a9027de34e8ea7e4b257/c-RzcFcoyRVFqYCSxgvsS.png"/></a></p>
<p>The code defines single-head attention in PyTorch - it transforms input vectors, computes attention scores and weights, and then calculates the weighted sum of values based on these weights (as per the attention equation)</p>
<p>When you have multiple of those stacked in parallel, you get what's called Multi-Head Attention. This gives a much simpler code if you are inheriting from the SingleHeadAttention class:</p>
<p><a href="https://cdn-uploads.huggingface.co/production/uploads/6438a9027de34e8ea7e4b257/QfUeWOfSU1J64OR7yIYUn.png" rel="nofollow"><img alt="image/png" src="https://cdn-uploads.huggingface.co/production/uploads/6438a9027de34e8ea7e4b257/QfUeWOfSU1J64OR7yIYUn.png"/></a></p>
<p>This one creates multiple attentions (inheriting from SingleHeadAttention class) and stacks them in parallel. During the forward pass, it applies each head to the input tensors Q, K, and V, concatenates the outputs, and linearly transforms them to produce the final output.</p>
<h1 class="relative group flex items-center">
<a class="block pr-1.5 text-lg md:absolute md:p-1.5 md:opacity-0 md:group-hover:opacity-100 md:right-full" href="#end-note" id="end-note" rel="nofollow">
<span class="header-link"><svg aria-hidden="true" class="text-gray-500 hover:text-black dark:hover:text-gray-200 w-4" height="1em" preserveaspectratio="xMidYMid meet" role="img" viewbox="0 0 256 256" width="1em" xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink"><path d="M167.594 88.393a8.001 8.001 0 0 1 0 11.314l-67.882 67.882a8 8 0 1 1-11.314-11.315l67.882-67.881a8.003 8.003 0 0 1 11.314 0zm-28.287 84.86l-28.284 28.284a40 40 0 0 1-56.567-56.567l28.284-28.284a8 8 0 0 0-11.315-11.315l-28.284 28.284a56 56 0 0 0 79.196 79.197l28.285-28.285a8 8 0 1 0-11.315-11.314zM212.852 43.14a56.002 56.002 0 0 0-79.196 0l-28.284 28.284a8 8 0 1 0 11.314 11.314l28.284-28.284a40 40 0 0 1 56.568 56.567l-28.285 28.285a8 8 0 0 0 11.315 11.314l28.284-28.284a56.065 56.065 0 0 0 0-79.196z" fill="currentColor"></path></svg></span>
</a>
<span>
		End Note
	</span>
</h1>
<p>So, in essence, having to code this made me revisit some key underlying concepts (pytorch's matmul, softmax, dropout, even backprop) that helped clear some stuff up. This along with Karpathy's nanoGPT video shaped my understanding - which is still a work in progress as new forms of the transformer architecture emerge.</p>
