1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42

import numpy as np
 
def demonstrate_variance_growth():
    """Show how dot product variance grows with dimension."""
    np.random.seed(42)
    
    dimensions = [1, 4, 16, 64, 256, 1024]
    n_samples = 10000
    
    print("Dot Product Variance vs. Dimension")
    print("=" * 50)
    print(f"{'d_k':<8} {'Theoretical Var':<18} {'Empirical Var':<18} {'Std Dev':<12}")
    print("-" * 50)
    
    for d_k in dimensions:
        # Sample random Q and K vectors
        Q = np.random.randn(n_samples, d_k)
        K = np.random.randn(n_samples, d_k)
        
        # Compute dot products
        dot_products = np.sum(Q * K, axis=1)
        
        empirical_var = np.var(dot_products)
        empirical_std = np.std(dot_products)
        
        print(f"{d_k:<8} {d_k:<18.2f} {empirical_var:<18.2f} {empirical_std:<12.2f}")
    
    return dot_products
 
# Run demonstration
dots = demonstrate_variance_growth()
 
# Show distribution for a specific d_k
d_k = 512
Q = np.random.randn(10000, d_k)
K = np.random.randn(10000, d_k)
dot_products = np.sum(Q * K, axis=1)
 
print(f"\nFor d_k = {d_k}:")
print(f"  Mean: {dot_products.mean():.3f} (expected: 0)")
print(f"  Std:  {dot_products.std():.3f} (expected: {np.sqrt(d_k):.3f})")
print(f"  Range: [{dot_products.min():.1f}, {dot_products.max():.1f}]")

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48

import numpy as np
from scipy.special import softmax
 
def compare_scaled_unscaled(d_k: int, n: int = 10):
    """Compare attention behavior with and without scaling."""
    np.random.seed(42)
    
    # Generate random Q and K
    Q = np.random.randn(n, d_k)
    K = np.random.randn(n, d_k)
    
    # Compute scores
    scores_unscaled = Q @ K.T
    scores_scaled = scores_unscaled / np.sqrt(d_k)
    
    # Apply softmax
    attn_unscaled = softmax(scores_unscaled, axis=-1)
    attn_scaled = softmax(scores_scaled, axis=-1)
    
    # Analyze distributions
    print(f"Dimension d_k = {d_k}")
    print("=" * 50)
    
    print(f"\nScore Statistics:")
    print(f"  Unscaled - Var: {scores_unscaled.var():.2f}, Range: [{scores_unscaled.min():.1f}, {scores_unscaled.max():.1f}]")
    print(f"  Scaled   - Var: {scores_scaled.var():.2f}, Range: [{scores_scaled.min():.1f}, {scores_scaled.max():.1f}]")
    
    # Entropy measures how "spread out" the attention is
    def entropy(p):
        return -np.sum(p * np.log(p + 1e-10), axis=-1).mean()
    
    max_entropy = np.log(n)  # Uniform distribution entropy
    
    print(f"\nAttention Entropy (max = {max_entropy:.2f}):")
    print(f"  Unscaled: {entropy(attn_unscaled):.3f} ({100*entropy(attn_unscaled)/max_entropy:.1f}% of max)")
    print(f"  Scaled:   {entropy(attn_scaled):.3f} ({100*entropy(attn_scaled)/max_entropy:.1f}% of max)")
    
    # Max attention weight (peakedness)
    print(f"\nMax Attention Weight (per row):")
    print(f"  Unscaled: {attn_unscaled.max(axis=-1).mean():.4f}")
    print(f"  Scaled:   {attn_scaled.max(axis=-1).mean():.4f}")
    
    return attn_scaled, attn_unscaled
 
# Test with different dimensions
for d_k in [4, 64, 512]:
    compare_scaled_unscaled(d_k)
    print()

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70

import numpy as np
from scipy.special import softmax
 
def analyze_softmax_gradients(scores: np.ndarray, label: str):
    """Analyze gradient magnitudes for softmax on given scores."""
    s = softmax(scores)
    
    # Compute diagonal of Jacobian: d(s_i)/d(z_i) = s_i * (1 - s_i)
    diag_grad = s * (1 - s)
    
    # Off-diagonal gradients: d(s_i)/d(z_j) = -s_i * s_j
    # Average magnitude of off-diagonal
    jacobian = np.diag(diag_grad) - np.outer(s, s)
    off_diag = jacobian - np.diag(np.diag(jacobian))
    
    print(f"\n{label}")
    print("-" * 40)
    print(f"Score range: [{scores.min():.1f}, {scores.max():.1f}]")
    print(f"Softmax output: {np.round(s, 4)}")
    print(f"Max softmax: {s.max():.6f}")
    print(f"Entropy: {-np.sum(s * np.log(s + 1e-10)):.4f}")
    print(f"Diagonal gradient magnitudes: {np.round(diag_grad, 6)}")
    print(f"Max gradient magnitude: {diag_grad.max():.6f}")
    print(f"Mean gradient magnitude: {diag_grad.mean():.6f}")
 
# Scenario: d_k = 512 dimensions
d_k = 512
np.random.seed(42)
 
# Sample a single query-key dot product scenario
Q = np.random.randn(d_k)
K = np.random.randn(4, d_k)  # 4 keys
 
# Unscaled scores
scores_unscaled = K @ Q  # Shape (4,)
scores_scaled = scores_unscaled / np.sqrt(d_k)
 
analyze_softmax_gradients(scores_unscaled, "UNSCALED (d_k=512)")
analyze_softmax_gradients(scores_scaled, "SCALED (d_k=512)")
 
# Extreme case: even larger unscaled
print("\n" + "=" * 50)
print("Impact on training:")
print("=" * 50)
 
# Simulate gradient accumulation over many examples
n_examples = 1000
unscaled_grads = []
scaled_grads = []
 
for _ in range(n_examples):
    Q = np.random.randn(d_k)
    K = np.random.randn(4, d_k)
    
    scores_u = K @ Q
    scores_s = scores_u / np.sqrt(d_k)
    
    s_u = softmax(scores_u)
    s_s = softmax(scores_s)
    
    grad_u = s_u * (1 - s_u)
    grad_s = s_s * (1 - s_s)
    
    unscaled_grads.append(grad_u.mean())
    scaled_grads.append(grad_s.mean())
 
print(f"\nMean gradient magnitude over {n_examples} samples:")
print(f"  Unscaled: {np.mean(unscaled_grads):.6f}")
print(f"  Scaled:   {np.mean(scaled_grads):.6f}")
print(f"  Ratio (scaled/unscaled): {np.mean(scaled_grads)/np.mean(unscaled_grads):.1f}x")

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42

import numpy as np
from scipy.special import softmax
 
def compare_scaling_strategies(d_k: int = 64, n: int = 8):
    """Compare different scaling strategies."""
    np.random.seed(42)
    
    Q = np.random.randn(n, d_k)
    K = np.random.randn(n, d_k)
    scores = Q @ K.T
    
    strategies = {
        'No scaling': 1,
        'Scale by √d_k': np.sqrt(d_k),
        'Scale by d_k': d_k,
        'Scale by d_k^(1/4)': d_k ** 0.25,
        'Scale by 2√d_k': 2 * np.sqrt(d_k),
    }
    
    print(f"Scaling Strategy Comparison (d_k = {d_k})")
    print("=" * 70)
    print(f"{'Strategy':<20} {'Scale':<10} {'Score Var':<12} {'Entropy':<12} {'Max Attn':<10}")
    print("-" * 70)
    
    max_entropy = np.log(n)
    
    for name, scale in strategies.items():
        scaled_scores = scores / scale
        attn = softmax(scaled_scores, axis=-1)
        
        score_var = scaled_scores.var()
        entropy = -np.sum(attn * np.log(attn + 1e-10), axis=-1).mean()
        max_attn = attn.max(axis=-1).mean()
        
        print(f"{name:<20} {scale:<10.2f} {score_var:<12.2f} {entropy:<12.3f} {max_attn:<10.4f}")
    
    print(f"\n(Reference: Max entropy = {max_entropy:.3f} for uniform distribution)")
    print("√d_k scaling achieves unit variance and balanced entropy.")
 
compare_scaling_strategies(d_k=64)
print()
compare_scaling_strategies(d_k=512)

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58

import numpy as np
from scipy.special import softmax
 
def analyze_initialization_effects(init_std: float, d_model: int = 64, d_k: int = 16):
    """Analyze how initialization affects attention score distribution."""
    np.random.seed(42)
    n = 10
    
    # Input with unit variance
    X = np.random.randn(n, d_model)
    
    # Initialize projections
    W_q = np.random.randn(d_model, d_k) * init_std
    W_k = np.random.randn(d_model, d_k) * init_std
    
    # Compute Q, K
    Q = X @ W_q
    K = X @ W_k
    
    # Compute attention scores
    scores = Q @ K.T
    scores_scaled = scores / np.sqrt(d_k)
    
    # Analyze
    print(f"Init std = {init_std}")
    print(f"  Q variance: {Q.var():.4f}")
    print(f"  Score variance (unscaled): {scores.var():.4f}")
    print(f"  Score variance (√d_k scaled): {scores_scaled.var():.4f}")
    
    # Compute optimal scale
    optimal_scale = np.sqrt(scores.var())
    scores_optimal = scores / optimal_scale
    
    print(f"  Optimal scale: {optimal_scale:.2f} (vs √d_k = {np.sqrt(d_k):.2f})")
    print(f"  Optimally scaled variance: {scores_optimal.var():.4f}")
 
print("Effect of Initialization Standard Deviation")
print("=" * 60)
print(f"d_model = 64, d_k = 16, √d_k = 4.0")
print()
 
# Standard initialization: ~1/sqrt(d_model)
std_standard = 1 / np.sqrt(64)
print("Standard initialization (1/√d_model):")
analyze_initialization_effects(std_standard)
 
# Xavier initialization
std_xavier = np.sqrt(2 / (64 + 16))
print("\nXavier initialization:")
analyze_initialization_effects(std_xavier)
 
# Small initialization
print("\nSmall initialization (0.01):")
analyze_initialization_effects(0.01)
 
# Large initialization
print("\nLarge initialization (0.5):")
analyze_initialization_effects(0.5)

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96

import numpy as np
from scipy.special import softmax
 
def scaled_dot_product_attention(
    Q: np.ndarray,           # (n_q, d_k)
    K: np.ndarray,           # (n_k, d_k)
    V: np.ndarray,           # (n_k, d_v)
    mask: np.ndarray = None, # (n_q, n_k) or broadcastable
    return_weights: bool = False
) -> np.ndarray:
    """
    Scaled dot-product attention as defined in 'Attention is All You Need'.
    
    Args:
        Q: Query matrix with shape (n_q, d_k)
        K: Key matrix with shape (n_k, d_k)
        V: Value matrix with shape (n_k, d_v)
        mask: Optional mask where -inf values prevent attention
        return_weights: If True, also return attention weights
    
    Returns:
        Output with shape (n_q, d_v)
        Optionally, attention weights with shape (n_q, n_k)
    """
    d_k = Q.shape[-1]
    
    # Step 1: Compute QK^T
    scores = Q @ K.T  # (n_q, n_k)
    
    # Step 2: Scale by sqrt(d_k)
    scores = scores / np.sqrt(d_k)
    
    # Step 3: Apply mask (if provided)
    if mask is not None:
        scores = scores + mask
    
    # Step 4: Softmax normalization
    attention_weights = softmax(scores, axis=-1)  # (n_q, n_k)
    
    # Step 5: Weighted aggregation
    output = attention_weights @ V  # (n_q, d_v)
    
    if return_weights:
        return output, attention_weights
    return output
 
 
def create_causal_mask(n: int) -> np.ndarray:
    """Create causal (lower triangular) mask."""
    mask = np.triu(np.ones((n, n)), k=1) * -1e9
    return mask
 
 
def create_padding_mask(lengths: np.ndarray, max_len: int) -> np.ndarray:
    """Create padding mask from sequence lengths."""
    # lengths: (batch_size,)
    # Output: (batch_size, max_len) - True where should be masked
    mask = np.arange(max_len)[None, :] >= lengths[:, None]
    return mask.astype(float) * -1e9
 
 
# Demonstration
np.random.seed(42)
n_q, n_k, d_k, d_v = 3, 5, 8, 16
 
Q = np.random.randn(n_q, d_k)
K = np.random.randn(n_k, d_k)
V = np.random.randn(n_k, d_v)
 
# Standard attention
output, weights = scaled_dot_product_attention(Q, K, V, return_weights=True)
 
print("Scaled Dot-Product Attention")
print("=" * 50)
print(f"Q shape: {Q.shape}")
print(f"K shape: {K.shape}")
print(f"V shape: {V.shape}")
print(f"Output shape: {output.shape}")
print(f"Weights shape: {weights.shape}")
print(f"\nAttention weights (each row sums to 1):")
print(np.round(weights, 3))
print(f"Row sums: {weights.sum(axis=-1)}")
 
# With causal mask (self-attention case where n_q = n_k)
Q_self = np.random.randn(4, d_k)
K_self = np.random.randn(4, d_k)
V_self = np.random.randn(4, d_v)
causal_mask = create_causal_mask(4)
 
output_causal, weights_causal = scaled_dot_product_attention(
    Q_self, K_self, V_self, mask=causal_mask, return_weights=True
)
 
print("\nWith Causal Mask:")
print(np.round(weights_causal, 3))
print("Notice: Upper triangle is 0 (can't attend to future)")

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59

import numpy as np
import time
 
def benchmark_attention(sequence_lengths: list, d_model: int = 512, num_runs: int = 5):
    """Benchmark attention computation for different sequence lengths."""
    from scipy.special import softmax
    
    d_k = d_v = d_model // 8  # Typical head dimension
    
    print("Attention Complexity Benchmark")
    print("=" * 60)
    print(f"d_model = {d_model}, d_k = d_v = {d_k}")
    print(f"{'Seq Len':<10} {'Time (ms)':<15} {'Memory (MB)':<15} {'n² scaling':<15}")
    print("-" * 60)
    
    reference_time = None
    
    for n in sequence_lengths:
        Q = np.random.randn(n, d_k).astype(np.float32)
        K = np.random.randn(n, d_k).astype(np.float32)
        V = np.random.randn(n, d_v).astype(np.float32)
        
        # Warmup
        _ = softmax(Q @ K.T / np.sqrt(d_k), axis=-1) @ V
        
        # Benchmark
        times = []
        for _ in range(num_runs):
            start = time.perf_counter()
            scores = Q @ K.T / np.sqrt(d_k)
            attn = softmax(scores, axis=-1)
            output = attn @ V
            times.append(time.perf_counter() - start)
        
        avg_time = np.mean(times) * 1000  # ms
        
        # Memory estimate (attention matrix is the bottleneck)
        attn_memory_mb = (n * n * 4) / (1024 * 1024)  # float32
        
        # n² scaling factor
        if reference_time is None:
            reference_time = avg_time
            reference_n = n
            scaling = 1.0
        else:
            expected_scaling = (n / reference_n) ** 2
            actual_scaling = avg_time / reference_time
            scaling = actual_scaling
        
        print(f"{n:<10} {avg_time:<15.3f} {attn_memory_mb:<15.2f} {scaling:<15.2f}")
    
    return None
 
# Run benchmark
sequence_lengths = [128, 256, 512, 1024, 2048]
benchmark_attention(sequence_lengths)
 
print("\nNote: Time should scale ~quadratically with sequence length")
print("Memory for attention matrix scales exactly quadratically")

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132

import torch
import torch.nn as nn
import torch.nn.functional as F
import math
from typing import Optional, Tuple
 
def scaled_dot_product_attention(
    query: torch.Tensor,     # (batch, seq_q, d_k)
    key: torch.Tensor,       # (batch, seq_k, d_k)
    value: torch.Tensor,     # (batch, seq_k, d_v)
    attn_mask: Optional[torch.Tensor] = None,  # (batch, seq_q, seq_k) or broadcastable
    dropout_p: float = 0.0,
    is_causal: bool = False,
    training: bool = True
) -> Tuple[torch.Tensor, torch.Tensor]:
    """
    Scaled dot-product attention with all standard features.
    
    Args:
        query, key, value: Query, key, and value tensors
        attn_mask: Additive mask (0 for attend, -inf for don't attend)
        dropout_p: Dropout probability on attention weights
        is_causal: If True, apply causal masking
        training: Whether in training mode (affects dropout)
    
    Returns:
        output: (batch, seq_q, d_v)
        attn_weights: (batch, seq_q, seq_k)
    """
    d_k = query.size(-1)
    
    # Compute attention scores
    scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(d_k)
    
    # Apply causal mask if requested
    if is_causal:
        seq_q, seq_k = query.size(-2), key.size(-2)
        causal_mask = torch.triu(
            torch.ones(seq_q, seq_k, device=query.device, dtype=torch.bool),
            diagonal=1
        )
        scores.masked_fill_(causal_mask, float('-inf'))
    
    # Apply explicit mask if provided
    if attn_mask is not None:
        scores = scores + attn_mask
    
    # Softmax normalization
    attn_weights = F.softmax(scores, dim=-1)
    
    # Apply dropout
    if training and dropout_p > 0.0:
        attn_weights = F.dropout(attn_weights, p=dropout_p)
    
    # Compute output
    output = torch.matmul(attn_weights, value)
    
    return output, attn_weights
 
 
class SelfAttention(nn.Module):
    """Self-attention module with QKV projections."""
    
    def __init__(self, d_model: int, d_k: int, d_v: int, dropout: float = 0.1):
        super().__init__()
        self.d_k = d_k
        
        # QKV projections
        self.W_q = nn.Linear(d_model, d_k, bias=False)
        self.W_k = nn.Linear(d_model, d_k, bias=False)
        self.W_v = nn.Linear(d_model, d_v, bias=False)
        
        # Output projection
        self.W_o = nn.Linear(d_v, d_model, bias=False)
        
        self.dropout = dropout
        
        # Initialize with Xavier
        for module in [self.W_q, self.W_k, self.W_v, self.W_o]:
            nn.init.xavier_uniform_(module.weight)
    
    def forward(
        self,
        x: torch.Tensor,
        attn_mask: Optional[torch.Tensor] = None,
        is_causal: bool = False
    ) -> torch.Tensor:
        # Project to Q, K, V
        Q = self.W_q(x)
        K = self.W_k(x)
        V = self.W_v(x)
        
        # Apply scaled dot-product attention
        attn_output, _ = scaled_dot_product_attention(
            Q, K, V,
            attn_mask=attn_mask,
            dropout_p=self.dropout if self.training else 0.0,
            is_causal=is_causal,
            training=self.training
        )
        
        # Output projection
        output = self.W_o(attn_output)
        
        return output
 
 
# Usage example
if __name__ == "__main__":
    batch_size, seq_len, d_model = 2, 10, 64
    d_k = d_v = 16
    
    x = torch.randn(batch_size, seq_len, d_model)
    
    # Create self-attention layer
    attention = SelfAttention(d_model, d_k, d_v, dropout=0.1)
    
    # Forward pass (causal for decoder-style)
    output = attention(x, is_causal=True)
    
    print(f"Input shape: {x.shape}")
    print(f"Output shape: {output.shape}")
    
    # Verify causal masking
    attention.eval()  # Disable dropout for checking
    Q = attention.W_q(x)
    K = attention.W_k(x)
    V = attention.W_v(x)
    
    _, weights = scaled_dot_product_attention(Q, K, V, is_causal=True, training=False)
    print(f"\nCausal attention weights (should be lower triangular):")
    print(weights[0].round(decimals=2))