Query, Key, Value › Putting It All Together

class SelfAttention:
    """
    Complete self-attention layer with QKV projections.
    """

    def __init__(self, embed_dim, d_k, d_v, seed=None):
        """
        Initialize projection matrices.

        Args:
            embed_dim: Dimension of input embeddings
            d_k: Dimension of queries and keys
            d_v: Dimension of values
            seed: Random seed for reproducibility
        """
        if seed is not None:
            np.random.seed(seed)

        # Xavier/Glorot initialization for stable gradients
        scale_qk = np.sqrt(2.0 / (embed_dim + d_k))
        scale_v = np.sqrt(2.0 / (embed_dim + d_v))

        self.W_q = np.random.randn(embed_dim, d_k) * scale_qk
        self.W_k = np.random.randn(embed_dim, d_k) * scale_qk
        self.W_v = np.random.randn(embed_dim, d_v) * scale_v
        self.d_k = d_k

    def forward(self, X):
        """
        Compute self-attention output.

        Args:
            X: Input embeddings of shape (seq_len, embed_dim)

        Returns:
            output: Attention output of shape (seq_len, d_v)
            attention_weights: Weights of shape (seq_len, seq_len)
        """
        # Project to Q, K, V
        Q = X @ self.W_q
        K = X @ self.W_k
        V = X @ self.W_v

        # Scaled dot-product attention
        scores = Q @ K.T / np.sqrt(self.d_k)

        # Softmax
        scores_stable = scores - scores.max(axis=1, keepdims=True)
        exp_scores = np.exp(scores_stable)
        attention_weights = exp_scores / exp_scores.sum(axis=1, keepdims=True)

        # Aggregate values
        output = attention_weights @ V

        return output, attention_weights

# Test the complete layer
embed_dim = 8
d_k = d_v = 6
seq_len = 5

np.random.seed(42)
X = np.random.randn(seq_len, embed_dim)

# Create and apply self-attention layer
attention = SelfAttention(embed_dim, d_k, d_v, seed=123)
output, weights = attention.forward(X)