Attention Mechanisms and Transformers

The attention mechanism addresses a fundamental challenge in deep learning: transforming variable-dimensional inputs into fixed-dimensional outputs through a weighted aggregation process. This capability proves essential when dealing with sequences or sets of varying sizes, where traditional fixed-parameter approaches struggle to maintain consistent output dimensions regardless of input cardinality.

Core Attention Framework

Consider the mathematical formulation of attention with three fundamental components: a query vector q ∈ ℝQ×1, a set of key-value pairs {(km, vm)} where m = 1..M, and corresponding projection matrices WQ ∈ ℝD×Q, WK ∈ ℝD×K, and WV ∈ ℝH×V. The attention operation computes a weighted sum of value projections, where weights derive from the similarity between query and key projections:

Attention(q, k, v) = W_V · v · softmax((W_K · k)^T · W_Q · q / √D) ∈ ℝ^(H×1)

</div>This mechanism transforms an arbitrary number M of key-value pairs into a fixed H-dimensional feature vector. The physical interpretation involves computing weighted sums over M H-dimensional value projections, with softmax-normalized similarity scores determining each value's contribution to the final output.

Dictionary Analogy
------------------

To build intuition, consider the Python dictionary analogy that mirrors attention's behavior. A dictionary stores key-value associations, and retrieval involves matching a query against available keys to retrieve corresponding values. Unlike exact dictionary lookup, attention performs soft matching—returning weighted combinations based on similarity rather than requiring exact key matches.

Suppose we have a dictionary mapping semantic features to values:

feature_map = {'color': 'blue', 'age': 22, 'category': 'pickup'} target_query = 'color' matched_value = feature_map[target_query]

For queries without exact matches, soft attention would return weighted combinations. A query for 'species' might retrieve 'pickup' with highest weight if semantic similarity exists between 'species' and 'category'.

Ganeralized Attention Formulation
---------------------------------

Extending to multiple queries and key-value pairs, let **Q** ∈ ℝ<sup>n×q</sup> represent n queries, **K** ∈ ℝ<sup>m×k</sup> represent m keys, and **V** ∈ ℝ<sup>m×v</sup> represent m corresponding values. The attention mechanism proceeds through three stages:

1. **Attention Scoring**: Compute similarity scores a(**Q**, **K**) ∈ ℝ<sup>n×m</sup>
2. **Weight Normalization**: Apply softmax to obtain probability distributions α(**Q**, **K**) = softmax(a(**Q**, **K**))
3. **Weighted Aggregation**: Compute output f(**Q**, **K**, **V**) = α(**Q**, **K**)<sup>T</sup> · **V** ∈ ℝ<sup>n×v</sup>

Additive Attention
------------------

Additive attention employs learned scoring functions with hidden layers. Given query **q** ∈ ℝ<sup>1×q and key **k** ∈ ℝ<sup>1×k, the scoring mechanism involves:</sup></sup>

<div style="background-color: #f5f5f5; padding: 15px; border-radius: 5px; margin: 15px 0;">```

a(q, k) = w_v^T · tanh(W_q · q^T + W_k · k^T) ∈ ℝ

Scaled Dot-Product Attention

Scaled dot-product attention offers a computationally efficient alternative using vector dot products. For normalized query and key dimensions d:

a(q, k) = (1/√d) · q · k^T ∈ ℝ

</div>Batched computation across multiple queries and key-value pairs yields:

<div style="background-color: #f5f5f5; padding: 15px; border-radius: 5px; margin: 15px 0;">```

Attention(Q, K, V) = softmax((1/√d) · Q · K^T) · V ∈ ℝ^(n×v)

This formulation provides a crucial property: the output dimension remains independent of the key-value pair count m, enabling handling of variable-sized inputs while producing consistent output dimensions.

Multi-Head Attention

Multi-head attention extends single attention by computing multiple parallel attention transformations and combining their outputs. This approach allows the model to jointly attend to information from different representation subspaces at different positions.

For sequence length L, embedding dimension dmodel, and h attention heads with subspace dimensions pq, pk, pv:

Each attention head performs linear projections and scaled dot-product attention:

H_i = Attention(Q · W_i^(q), K · W_i^(k), V · W_i^(v)) ∈ ℝ^(L×p_v)

</div>where W<sub>i</sub><sup>(q)</sup> ∈ ℝ<sup>p\_q×d\_q</sup>, W<sub>i</sub><sup>(k)</sup> ∈ ℝ<sup>p\_k×d\_k</sup>, and W<sub>i</sub><sup>(v)</sup> ∈ ℝ<sup>p\_v×d\_v</sup> are projection matrices for head i.

Concatenating all head outputs and applying a final linear projection produces the multi-head attantion output:

<div style="background-color: #f5f5f5; padding: 15px; border-radius: 5px; margin: 15px 0;">```

H = [H_1, ..., H_h] ∈ ℝ^(L×h·p_v)
Output = H · W_o^T ∈ ℝ^(L×p_o)

A common configuration sets po = h·pq = h·pk = h·pv to maintain computational efficiency through parallel computation across all heads while preserving the original embedding dimension in the output.

Self-Attention

Self-attention applies attention where queries, keys, and values derive from the same sequence. Given input sequence X ∈ ℝn×d with row vectors xi representing each position:

Q = X, K = X, V = X Y = softmax((1/√d) · X · X^T) · X ∈ ℝ^(n×d)

</div>The output y<sub>i</sub> for each position aggregates information from all positions through learned attention weights, enabling direct modeling of long-range dependencies regardless of distance in the sequence.

When queries and keys derive from different sources (different numbers of positions), learned linear transformations project inputs to consistent dimensions:

<div style="background-color: #f5f5f5; padding: 15px; border-radius: 5px; margin: 15px 0;">```

Q = tanh(W_q · X_q) ∈ ℝ^(n×d_att)
K = tanh(W_k · X_k) ∈ ℝ^(m×d_att)
V = tanh(W_v · X_v) ∈ ℝ^(m×d_att)

Practical Implementation

The following PyTorch implementation demonstrates a generalized attention layer suitable for multi-agent scenarios:

class AttentionBlock(nn.Module):
    def __init__(self, input_dim, latent_dim, device):
        super(AttentionBlock, self).__init__()
        self.input_features = input_dim
        self.device = device
        self.query_proj = self._init_linear(nn.Linear(self.input_features, latent_dim))
        self.key_proj = self._init_linear(nn.Linear(self.input_features, latent_dim))
        self.value_proj = self._init_linear(nn.Linear(self.input_features, latent_dim))
        self.latent_dim = latent_dim

    def _init_linear(self, layer):
        nn.init.xavier_uniform_(layer.weight)
        return layer

    def forward(self, features):
        """
        Input: [num_agents, num_targets, input_dim]
        Output: [num_agents, num_targets, latent_dim], [num_agents, latent_dim]
        """
        projected_q = torch.tanh(self.query_proj(features))
        projected_k = torch.tanh(self.key_proj(features))
        projected_v = torch.tanh(self.value_proj(features))

        attention_scores = torch.bmm(projected_q, projected_k.permute(0, 2, 1))
        normalized_weights = F.softmax(attention_scores, dim=2)
        aggregated = torch.bmm(normalized_weights, projected_v)

        global_repr = aggregated.sum(dim=1)
        return aggregated, global_repr

Transformer Ancoder Architecture

The original Transformer implementation employs multi-head self-attention with consistent dimensionality across all components. Given embedding dimension E = 128, attention subspace dimension d = 128, and 8 attention heads:

class MultiHeadAttention(nn.Module):
    def __init__(self, model_dim, num_heads, dropout_rate=0.1):
        super(MultiHeadAttention, self).__init__()
        assert model_dim % num_heads == 0
        self.head_dim = model_dim // num_heads
        self.num_heads = num_heads

        self.query_linear = nn.Linear(model_dim, model_dim)
        self.key_linear = nn.Linear(model_dim, model_dim)
        self.value_linear = nn.Linear(model_dim, model_dim)
        self.output_linear = nn.Linear(model_dim, model_dim)
        self.dropout = nn.Dropout(dropout_rate)

    def forward(self, query_seq, key_seq, value_seq):
        batch_size = query_seq.size(0)
        
        # Linear projections and reshape for multi-head processing
        query_proj = self.query_linear(query_seq)
        key_proj = self.key_linear(key_seq)
        value_proj = self.value_linear(value_seq)

        query_heads = query_proj.view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2)
        key_heads = key_proj.view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2)
        value_heads = value_proj.view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2)

        # Scaled dot-product attention with attention weight capture
        attention_output, attention_weights = self._scaled_dot_attention(
            query_heads, key_heads, value_heads
        )

        # Concatenate heads and apply final linear transformation
        concatenated = attention_output.transpose(1, 2).contiguous()
        concatenated = concatenated.view(batch_size, -1, self.num_heads * self.head_dim)

        output = self.output_linear(concatenated)
        output = self.dropout(output)
        return output, attention_weights

    def _scaled_dot_attention(self, q, k, v):
        scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.head_dim)
        attention_weights = F.softmax(scores, dim=-1)
        output = torch.matmul(attention_weights, v)
        return output, attention_weights

Applications

Attention mechanisms and Transformers have demonstrated remarkable success across diverse domains. In neural information processing, attention enables multi-agent coordination by allowing dynamic information sharing based on relevance rather than fixed connectivity patterns. Wireless communication systems leverage attention for semantic encoding, extracting task-relevant features from high-dimensional signals. The flexibility of attention's variable-to-fixed transformation makes it particularly valuable for reinforcement learning with varying state dimensions and for processing sequences of arbitrary length in natural language processing and computer vision applications.