% Generated by roxygen2: do not edit by hand
% Please edit documentation in R/ag_layers.R
\name{ag_multihead_attention}
\alias{ag_multihead_attention}
\title{Create a Multi-Head Attention layer}
\usage{
ag_multihead_attention(d_model, n_heads, dropout = 0, bias = TRUE)
}
\arguments{
\item{d_model}{Model (embedding) dimension}

\item{n_heads}{Number of attention heads. \code{d_model} must be divisible
by \code{n_heads}.}

\item{dropout}{Attention dropout probability (default 0, applied in
training mode only)}

\item{bias}{Logical: add bias to output projection (default TRUE)}
}
\value{
An \code{ag_multihead_attention} environment with
  \code{$forward(q, k, v, causal_mask)} and \code{$parameters()}
}
\description{
Implements scaled dot-product multi-head attention as in "Attention Is All
You Need" (Vaswani et al., 2017).
}
\details{
Calling convention (mirrors PyTorch \code{nn.MultiheadAttention}):
\itemize{
  \item \code{layer$forward(q)} — self-attention (\code{k = v = q})
  \item \code{layer$forward(q, k, v)} — cross-attention
}

Tensor layout: \strong{\code{[d_model, seq_len]}} — columns are tokens,
consistent with the rest of the ag_* API.

Forward pass:
\preformatted{
  Q = W_q \%*\% q            [d_k * n_heads, seq_len]
  K = W_k \%*\% k            [d_k * n_heads, seq_len]
  V = W_v \%*\% v            [d_v * n_heads, seq_len]

  for each head h:
    q_h = Q[h*d_k+1 : (h+1)*d_k, ]    [d_k, seq_len]
    k_h = K[h*d_k+1 : (h+1)*d_k, ]    [d_k, seq_len]
    v_h = V[h*d_v+1 : (h+1)*d_v, ]    [d_v, seq_len]
    A_h = softmax(t(q_h) \%*\% k_h / sqrt(d_k))   [seq_len, seq_len]
    if causal_mask: A_h[i,j] = 0 for j > i
    head_h = v_h \%*\% A_h             [d_v, seq_len]

  concat = rbind(head_1, ..., head_H)  [d_v*n_heads, seq_len]
  out    = W_o \%*\% concat + b_o      [d_model, seq_len]
}
}
\examples{
\donttest{
# Self-attention
mha <- ag_multihead_attention(64L, 8L)
x   <- ag_tensor(matrix(rnorm(64 * 10), 64, 10))  # [d_model=64, seq_len=10]
out <- mha$forward(x)                              # [64, 10]

# Cross-attention
q   <- ag_tensor(matrix(rnorm(64 * 10), 64, 10))
kv  <- ag_tensor(matrix(rnorm(64 * 15), 64, 15))
out <- mha$forward(q, kv, kv)

# Causal (GPT-style)
out <- mha$forward(x, causal_mask = TRUE)
}
}
