Attention Mechanisms

5.1 The Problem Attention Solves

In RNNs, all information must flow through a single hidden state vector. When translating a long sentence, the encoder must compress the entire input into one vector. Attention lets each output position directly access relevant parts of the input.

Example: translating "The quick brown fox jumps over the lazy dog" — each output word can attend (focus on) whichever input words are relevant, rather than pulling everything from a single summary vector.

5.2 Scaled Dot-Product Attention

The Formula

Intuition

• Compute similarity between each query and all keys (Q K^T)
• Scale by 1/√d_k to prevent softmax saturation at large dimensions
• Softmax to get attention weights (probability distribution over positions)
• Weight values by attention weights to produce output

-- Scaled dot-product attention
attention
  :: Matrix Double      -- Q: queries  (seq_len_query, d_k)
  -> Matrix Double      -- K: keys     (seq_len_key, d_k)
  -> Matrix Double      -- V: values   (seq_len_value, d_v)
  -> Matrix Double      -- Output      (seq_len_query, d_v)
attention queries keys values =
  let d_k          = fromIntegral (cols keys)
      scores       = matmul queries (transpose keys)  -- (seq_q, seq_k)
      scaledScores = map (/ sqrt d_k) scores
      weights      = softmaxRows scaledScores         -- rows sum to 1
  in matmul weights values

-- Softmax per row for numerical stability
softmaxVector :: Vector Double -> Vector Double
softmaxVector vec =
  let maxVal  = maximum vec
      shifted = map (\x -> x - maxVal) vec  -- stability: subtract max first
      exps    = map exp shifted
      sumExps = sum exps
  in map (/ sumExps) exps

5.3 Multi-Head Attention

A single attention head can only attend to one type of relationship at a time. With h heads, the model can simultaneously attend to different aspects of the input (e.g., syntactic vs. semantic relationships).

data MultiHeadAttention = MultiHeadAttention
  { numHeads          :: Int
  , d_model           :: Int
  , queryProjections  :: [Matrix Double]
  , keyProjections    :: [Matrix Double]
  , valueProjections  :: [Matrix Double]
  , outputProjection  :: Matrix Double
  }

multiHeadAttention
  :: MultiHeadAttention
  -> Matrix Double   -- Query
  -> Matrix Double   -- Key
  -> Matrix Double   -- Value
  -> Matrix Double
multiHeadAttention mha queries keys values =
  let queryHeads    = [matmul queries (queryProjections mha !! i) | i <- [0..numHeads mha - 1]]
      keyHeads      = [matmul keys   (keyProjections   mha !! i) | i <- [0..numHeads mha - 1]]
      valueHeads    = [matmul values (valueProjections mha !! i) | i <- [0..numHeads mha - 1]]
      headOutputs   = zipWith3 attention queryHeads keyHeads valueHeads
      concatenated  = horizontalConcat headOutputs
  in matmul concatenated (outputProjection mha)

5.4 Masked Attention

During generation, a model must not attend to future tokens (they don't exist yet). We apply a causal mask before the softmax — future positions are set to a large negative number so they receive near-zero attention weight.

-- Mask future positions before softmax
maskedAttention
  :: Matrix Double   -- queries
  -> Matrix Double   -- keys
  -> Matrix Double   -- values
  -> Matrix Double   -- causal mask (1 = attend, 0 = block)
  -> Matrix Double
maskedAttention queries keys values mask =
  let d_k          = fromIntegral (cols keys)
      scores       = matmul queries (transpose keys)
      scaledScores = map (/ sqrt d_k) scores
      -- Set future positions to -1e9 before softmax
      maskedScores = elementwiseAdd
                       (elementwiseMultiply scaledScores mask)
                       (scalarMultiply (-1.0e9) (oneMinus mask))
      weights      = softmaxRows maskedScores
  in matmul weights values