相対位置エンコーディングの理論 — Shaw et al.からTransformer-XLまで、相対位置の系譜

import numpy as np
import matplotlib.pyplot as plt

np.random.seed(42)

def shaw_relative_attention(X, W_Q, W_K, W_V, rel_pos_K, rel_pos_V, max_rel_dist):
    """
    Shaw et al. (2018) の相対位置Attention

    Parameters:
    -----------
    X : (seq_len, d_model) 入力
    W_Q, W_K, W_V : 射影行列
    rel_pos_K : (2*max_rel_dist+1, d_k) 相対位置ベクトル (Key側)
    rel_pos_V : (2*max_rel_dist+1, d_v) 相対位置ベクトル (Value側)
    max_rel_dist : クリッピングの最大距離
    """
    seq_len = X.shape[0]
    d_k = W_Q.shape[1]

    # Query, Key, Value の計算
    Q = X @ W_Q  # (seq_len, d_k)
    K = X @ W_K  # (seq_len, d_k)
    V = X @ W_V  # (seq_len, d_v)

    # 標準のAttentionスコア: Q @ K^T
    content_scores = Q @ K.T  # (seq_len, seq_len)

    # 相対位置スコアの計算
    position_scores = np.zeros((seq_len, seq_len))
    for i in range(seq_len):
        for j in range(seq_len):
            # クリッピング: 相対距離を[-max_rel_dist, max_rel_dist]に制限
            rel_dist = np.clip(i - j, -max_rel_dist, max_rel_dist)
            # インデックスを [0, 2*max_rel_dist] に変換
            idx = rel_dist + max_rel_dist
            position_scores[i, j] = Q[i] @ rel_pos_K[idx]

    # スコアの合計とスケーリング
    scores = (content_scores + position_scores) / np.sqrt(d_k)

    # Softmax
    scores_exp = np.exp(scores - np.max(scores, axis=-1, keepdims=True))
    attn_weights = scores_exp / np.sum(scores_exp, axis=-1, keepdims=True)

    # Value側の相対位置も考慮した出力
    output = np.zeros_like(V)
    for i in range(seq_len):
        for j in range(seq_len):
            rel_dist = np.clip(i - j, -max_rel_dist, max_rel_dist)
            idx = rel_dist + max_rel_dist
            output[i] += attn_weights[i, j] * (V[j] + rel_pos_V[idx])

    return output, attn_weights, content_scores / np.sqrt(d_k), position_scores / np.sqrt(d_k)

# パラメータ設定
seq_len = 12
d_model = 16
d_k = d_v = 8
max_rel_dist = 4

# 入力とパラメータの初期化
X = np.random.randn(seq_len, d_model) * 0.1
W_Q = np.random.randn(d_model, d_k) * 0.1
W_K = np.random.randn(d_model, d_k) * 0.1
W_V = np.random.randn(d_model, d_v) * 0.1

# 相対位置ベクトル (2*max_rel_dist+1 = 9個)
rel_pos_K = np.random.randn(2 * max_rel_dist + 1, d_k) * 0.5
rel_pos_V = np.random.randn(2 * max_rel_dist + 1, d_v) * 0.5

# 実行
output, attn_weights, content_scores, position_scores = shaw_relative_attention(
    X, W_Q, W_K, W_V, rel_pos_K, rel_pos_V, max_rel_dist
)

# 可視化: 内容スコア vs 位置スコア vs 合計
fig, axes = plt.subplots(1, 3, figsize=(18, 5))

im0 = axes[0].imshow(content_scores, cmap='RdBu_r', aspect='auto')
axes[0].set_title('Content Scores (q_i^T k_j)', fontsize=12)
axes[0].set_xlabel('Key position j')
axes[0].set_ylabel('Query position i')
plt.colorbar(im0, ax=axes[0], shrink=0.8)

im1 = axes[1].imshow(position_scores, cmap='RdBu_r', aspect='auto')
axes[1].set_title('Relative Position Scores (q_i^T a_K)', fontsize=12)
axes[1].set_xlabel('Key position j')
axes[1].set_ylabel('Query position i')
plt.colorbar(im1, ax=axes[1], shrink=0.8)

im2 = axes[2].imshow(attn_weights, cmap='viridis', aspect='auto')
axes[2].set_title('Final Attention Weights', fontsize=12)
axes[2].set_xlabel('Key position j')
axes[2].set_ylabel('Query position i')
plt.colorbar(im2, ax=axes[2], shrink=0.8)

plt.tight_layout()
plt.savefig('shaw_attention.png', dpi=150, bbox_inches='tight')
plt.show()

import numpy as np
import matplotlib.pyplot as plt

def t5_relative_position_bucket(relative_position, bidirectional=True,
                                 num_buckets=32, max_distance=128):
    """
    T5の相対位置バケット化関数

    Parameters:
    -----------
    relative_position : int or ndarray 相対距離 (i - j)
    bidirectional : bool 双方向かどうか
    num_buckets : int バケット数
    max_distance : int 最大距離

    Returns:
    --------
    バケット番号
    """
    relative_buckets = 0

    if bidirectional:
        num_buckets //= 2
        # 正方向と負方向を分離
        relative_buckets += (relative_position > 0).astype(np.int32) * num_buckets
        relative_position = np.abs(relative_position)
    else:
        relative_position = -np.minimum(relative_position, 0)

    # 近距離: そのままの距離をバケット番号に
    max_exact = num_buckets // 2
    is_small = relative_position < max_exact

    # 遠距離: 対数スケールで圧縮
    relative_position_if_large = max_exact + (
        np.log(relative_position.astype(np.float32) / max_exact)
        / np.log(max_distance / max_exact)
        * (num_buckets - max_exact)
    ).astype(np.int32)

    # 最大バケット番号でクリップ
    relative_position_if_large = np.minimum(
        relative_position_if_large, num_buckets - 1
    )

    relative_buckets += np.where(is_small, relative_position, relative_position_if_large)

    return relative_buckets

# バケット化の可視化
distances = np.arange(-150, 151)
buckets = t5_relative_position_bucket(distances, bidirectional=True)

fig, axes = plt.subplots(1, 2, figsize=(16, 5))

# 距離 → バケット番号のマッピング
axes[0].plot(distances, buckets, 'b-', linewidth=1.5)
axes[0].set_xlabel('Relative Distance (i - j)', fontsize=12)
axes[0].set_ylabel('Bucket Index', fontsize=12)
axes[0].set_title('T5 Relative Position Bucket Mapping', fontsize=13)
axes[0].axhline(y=16, color='gray', linestyle='--', alpha=0.5, label='Boundary (bucket 16)')
axes[0].axvline(x=0, color='gray', linestyle='--', alpha=0.5)
axes[0].legend()
axes[0].grid(True, alpha=0.3)

# 各バケットに含まれる距離の範囲
bucket_counts = np.zeros(32)
for d in range(-128, 129):
    b = t5_relative_position_bucket(np.array([d]), bidirectional=True)[0]
    bucket_counts[b] += 1

axes[1].bar(range(32), bucket_counts, color='steelblue', alpha=0.8)
axes[1].set_xlabel('Bucket Index', fontsize=12)
axes[1].set_ylabel('Number of Distances in Bucket', fontsize=12)
axes[1].set_title('Distance Count per Bucket', fontsize=13)
axes[1].grid(True, alpha=0.3, axis='y')

plt.tight_layout()
plt.savefig('t5_bucket.png', dpi=150, bbox_inches='tight')
plt.show()

import numpy as np
import matplotlib.pyplot as plt

np.random.seed(42)

def t5_relative_position_bucket(relative_position, bidirectional=True,
                                 num_buckets=32, max_distance=128):
    """T5の相対位置バケット化関数"""
    relative_buckets = 0
    if bidirectional:
        num_buckets //= 2
        relative_buckets += (relative_position > 0).astype(np.int32) * num_buckets
        relative_position = np.abs(relative_position)
    else:
        relative_position = -np.minimum(relative_position, 0)

    max_exact = num_buckets // 2
    is_small = relative_position < max_exact

    relative_position_if_large = max_exact + (
        np.log(relative_position.astype(np.float32).clip(min=1) / max_exact)
        / np.log(max_distance / max_exact)
        * (num_buckets - max_exact)
    ).astype(np.int32)
    relative_position_if_large = np.minimum(relative_position_if_large, num_buckets - 1)

    relative_buckets += np.where(is_small, relative_position, relative_position_if_large)
    return relative_buckets

def t5_attention(X, W_Q, W_K, W_V, bias_table, num_buckets=32, max_distance=128):
    """
    T5方式の相対位置バイアス付きAttention

    Parameters:
    -----------
    X : (seq_len, d_model) 入力
    W_Q, W_K, W_V : 射影行列
    bias_table : (num_buckets,) 各バケットに対応するバイアス値
    """
    seq_len = X.shape[0]
    d_k = W_Q.shape[1]

    Q = X @ W_Q
    K = X @ W_K
    V = X @ W_V

    # 内容ベースのスコア
    content_scores = Q @ K.T / np.sqrt(d_k)

    # 相対位置バイアスの計算
    positions = np.arange(seq_len)
    relative_positions = positions[:, None] - positions[None, :]
    buckets = t5_relative_position_bucket(relative_positions, bidirectional=True,
                                          num_buckets=num_buckets, max_distance=max_distance)
    position_bias = bias_table[buckets]

    # スコアにバイアスを加算
    scores = content_scores + position_bias

    # Softmax
    scores_exp = np.exp(scores - np.max(scores, axis=-1, keepdims=True))
    attn_weights = scores_exp / np.sum(scores_exp, axis=-1, keepdims=True)

    output = attn_weights @ V
    return output, attn_weights, content_scores, position_bias

# パラメータ設定
seq_len = 20
d_model = 16
d_k = d_v = 8
num_buckets = 32

X = np.random.randn(seq_len, d_model) * 0.1
W_Q = np.random.randn(d_model, d_k) * 0.1
W_K = np.random.randn(d_model, d_k) * 0.1
W_V = np.random.randn(d_model, d_v) * 0.1

# バイアステーブル: 近い距離に大きなバイアス、遠い距離に小さなバイアスを模擬
bias_table = np.zeros(num_buckets)
for b in range(num_buckets):
    if b < 16:  # 負方向 (j > i)
        dist_approx = b if b < 8 else 8 * 2**((b - 8) / 8 * np.log2(128 / 8))
        bias_table[b] = -0.3 * np.log1p(dist_approx)
    else:  # 正方向 (i > j)
        dist_approx = (b - 16) if (b - 16) < 8 else 8 * 2**(((b - 16) - 8) / 8 * np.log2(128 / 8))
        bias_table[b] = -0.3 * np.log1p(dist_approx)
# 距離0のバイアスを最大にする
bias_table[0] = 0.5
bias_table[16] = 0.5

output, attn_weights, content_scores, position_bias = t5_attention(
    X, W_Q, W_K, W_V, bias_table, num_buckets
)

fig, axes = plt.subplots(1, 3, figsize=(18, 5))

im0 = axes[0].imshow(content_scores, cmap='RdBu_r', aspect='auto')
axes[0].set_title('Content Scores', fontsize=12)
axes[0].set_xlabel('Key position j')
axes[0].set_ylabel('Query position i')
plt.colorbar(im0, ax=axes[0], shrink=0.8)

im1 = axes[1].imshow(position_bias, cmap='RdBu_r', aspect='auto')
axes[1].set_title('T5 Position Bias b(i-j)', fontsize=12)
axes[1].set_xlabel('Key position j')
axes[1].set_ylabel('Query position i')
plt.colorbar(im1, ax=axes[1], shrink=0.8)

im2 = axes[2].imshow(attn_weights, cmap='viridis', aspect='auto')
axes[2].set_title('Final Attention Weights', fontsize=12)
axes[2].set_xlabel('Key position j')
axes[2].set_ylabel('Query position i')
plt.colorbar(im2, ax=axes[2], shrink=0.8)

plt.tight_layout()
plt.savefig('t5_attention.png', dpi=150, bbox_inches='tight')
plt.show()

import numpy as np
import matplotlib.pyplot as plt

def shaw_distance_profile(distances, max_rel_dist=16):
    """Shaw方式: クリッピング後の有効距離（概念的プロファイル）"""
    clipped = np.clip(distances, -max_rel_dist, max_rel_dist)
    # クリッピングされた距離を正規化して返す（学習前の概念的なプロファイル）
    return -np.abs(clipped) / max_rel_dist

def t5_distance_profile(distances, num_buckets=32, max_distance=128):
    """T5方式: バケット番号に基づくプロファイル"""
    buckets = t5_relative_position_bucket(distances.astype(np.int64),
                                           bidirectional=True,
                                           num_buckets=num_buckets,
                                           max_distance=max_distance)
    # バケット番号を正規化（概念的なプロファイル）
    center = num_buckets // 2
    return -np.abs(buckets - center) / center

def alibi_distance_profile(distances, slope=0.125):
    """ALiBi方式: 線形減衰"""
    return -slope * np.abs(distances)

distances = np.arange(-64, 65)

fig, axes = plt.subplots(2, 2, figsize=(14, 10))

# Shaw方式
profile_shaw = shaw_distance_profile(distances, max_rel_dist=16)
axes[0, 0].plot(distances, profile_shaw, 'b-', linewidth=2)
axes[0, 0].set_title('Shaw et al. (Clipping, k=16)', fontsize=12)
axes[0, 0].set_xlabel('Relative Distance')
axes[0, 0].set_ylabel('Effective Bias (conceptual)')
axes[0, 0].axvline(x=-16, color='red', linestyle='--', alpha=0.5, label='Clip boundary')
axes[0, 0].axvline(x=16, color='red', linestyle='--', alpha=0.5)
axes[0, 0].legend()
axes[0, 0].grid(True, alpha=0.3)

# T5方式
profile_t5 = t5_distance_profile(distances)
axes[0, 1].plot(distances, profile_t5, 'g-', linewidth=2)
axes[0, 1].set_title('T5 (Log Bucketing)', fontsize=12)
axes[0, 1].set_xlabel('Relative Distance')
axes[0, 1].set_ylabel('Effective Bias (conceptual)')
axes[0, 1].grid(True, alpha=0.3)

# ALiBi方式（異なるスロープ）
slopes = [0.5, 0.25, 0.125, 0.0625]
colors = ['#e74c3c', '#e67e22', '#2ecc71', '#3498db']
for slope, color in zip(slopes, colors):
    profile = alibi_distance_profile(distances, slope)
    axes[1, 0].plot(distances, profile, color=color, linewidth=1.5, label=f'm={slope}')
axes[1, 0].set_title('ALiBi (Linear Decay, multiple slopes)', fontsize=12)
axes[1, 0].set_xlabel('Relative Distance')
axes[1, 0].set_ylabel('Bias Value')
axes[1, 0].legend()
axes[1, 0].grid(True, alpha=0.3)
axes[1, 0].set_ylim(-10, 1)

# 全手法の比較（正の距離のみ）
pos_dist = np.arange(0, 65)
axes[1, 1].plot(pos_dist, shaw_distance_profile(pos_dist, 16), 'b-',
                linewidth=2, label='Shaw (clip=16)')
axes[1, 1].plot(pos_dist, t5_distance_profile(pos_dist), 'g-',
                linewidth=2, label='T5 (log bucket)')
axes[1, 1].plot(pos_dist, alibi_distance_profile(pos_dist, 0.125), 'r-',
                linewidth=2, label='ALiBi (m=0.125)')
axes[1, 1].set_title('Comparison (Positive Direction)', fontsize=12)
axes[1, 1].set_xlabel('Relative Distance')
axes[1, 1].set_ylabel('Effective Bias (normalized)')
axes[1, 1].legend()
axes[1, 1].grid(True, alpha=0.3)

plt.tight_layout()
plt.savefig('position_encoding_comparison.png', dpi=150, bbox_inches='tight')
plt.show()

import numpy as np
import matplotlib.pyplot as plt

np.random.seed(123)

def softmax(x, axis=-1):
    """数値安定なSoftmax"""
    x_max = np.max(x, axis=axis, keepdims=True)
    e_x = np.exp(x - x_max)
    return e_x / np.sum(e_x, axis=axis, keepdims=True)

def t5_relative_position_bucket(relative_position, bidirectional=True,
                                 num_buckets=32, max_distance=128):
    """T5バケット化関数"""
    relative_buckets = 0
    if bidirectional:
        num_buckets //= 2
        relative_buckets += (relative_position > 0).astype(np.int32) * num_buckets
        relative_position = np.abs(relative_position)
    else:
        relative_position = -np.minimum(relative_position, 0)
    max_exact = num_buckets // 2
    is_small = relative_position < max_exact
    relative_position_if_large = max_exact + (
        np.log(relative_position.astype(np.float32).clip(min=1) / max_exact)
        / np.log(max_distance / max_exact)
        * (num_buckets - max_exact)
    ).astype(np.int32)
    relative_position_if_large = np.minimum(relative_position_if_large, num_buckets - 1)
    relative_buckets += np.where(is_small, relative_position, relative_position_if_large)
    return relative_buckets

seq_len = 24
d_model = 32
d_k = 16

# 共通の入力
X = np.random.randn(seq_len, d_model) * 0.1
W_Q = np.random.randn(d_model, d_k) * 0.1
W_K = np.random.randn(d_model, d_k) * 0.1

Q = X @ W_Q
K = X @ W_K
content_scores = Q @ K.T / np.sqrt(d_k)

# --- Shaw方式のAttention ---
max_rel_dist = 6
rel_pos_K_shaw = np.random.randn(2 * max_rel_dist + 1, d_k) * 0.3
# 距離0に近いほど大きなスコアになるよう調整
for r in range(2 * max_rel_dist + 1):
    dist = abs(r - max_rel_dist)
    rel_pos_K_shaw[r] *= np.exp(-0.2 * dist)

shaw_pos_scores = np.zeros((seq_len, seq_len))
for i in range(seq_len):
    for j in range(seq_len):
        rel_dist = np.clip(i - j, -max_rel_dist, max_rel_dist) + max_rel_dist
        shaw_pos_scores[i, j] = Q[i] @ rel_pos_K_shaw[rel_dist]

shaw_scores = content_scores + shaw_pos_scores
shaw_attn = softmax(shaw_scores)

# --- T5方式のAttention ---
num_buckets = 32
bias_table_t5 = np.zeros(num_buckets)
for b in range(num_buckets):
    center = 0 if b < 16 else 16
    dist_idx = b - center
    bias_table_t5[b] = -0.5 * dist_idx  # 距離が大きいほど負のバイアス

positions = np.arange(seq_len)
rel_positions = positions[:, None] - positions[None, :]
buckets = t5_relative_position_bucket(rel_positions.astype(np.int64), bidirectional=True)
t5_bias = bias_table_t5[buckets]

t5_scores = content_scores + t5_bias
t5_attn = softmax(t5_scores)

# --- 標準Attention（位置バイアスなし）---
standard_attn = softmax(content_scores)

# 可視化
fig, axes = plt.subplots(1, 3, figsize=(18, 5))

im0 = axes[0].imshow(standard_attn, cmap='viridis', aspect='auto', vmin=0, vmax=0.2)
axes[0].set_title('Standard Attention (No Position)', fontsize=12)
axes[0].set_xlabel('Key position j')
axes[0].set_ylabel('Query position i')
plt.colorbar(im0, ax=axes[0], shrink=0.8)

im1 = axes[1].imshow(shaw_attn, cmap='viridis', aspect='auto', vmin=0, vmax=0.2)
axes[1].set_title('Shaw Relative Attention (k=6)', fontsize=12)
axes[1].set_xlabel('Key position j')
axes[1].set_ylabel('Query position i')
plt.colorbar(im1, ax=axes[1], shrink=0.8)

im2 = axes[2].imshow(t5_attn, cmap='viridis', aspect='auto', vmin=0, vmax=0.2)
axes[2].set_title('T5 Relative Bias Attention', fontsize=12)
axes[2].set_xlabel('Key position j')
axes[2].set_ylabel('Query position i')
plt.colorbar(im2, ax=axes[2], shrink=0.8)

plt.tight_layout()
plt.savefig('attention_pattern_comparison.png', dpi=150, bbox_inches='tight')
plt.show()

# 特定の位置のAttention分布を比較
fig, axes = plt.subplots(1, 2, figsize=(14, 5))

query_pos = 12
axes[0].bar(range(seq_len), standard_attn[query_pos], alpha=0.4, label='Standard', color='gray')
axes[0].bar(range(seq_len), shaw_attn[query_pos], alpha=0.6, label='Shaw', color='blue')
axes[0].bar(range(seq_len), t5_attn[query_pos], alpha=0.4, label='T5', color='green')
axes[0].axvline(x=query_pos, color='red', linestyle='--', alpha=0.5, label=f'Query pos={query_pos}')
axes[0].set_xlabel('Key Position j', fontsize=12)
axes[0].set_ylabel('Attention Weight', fontsize=12)
axes[0].set_title(f'Attention Distribution from Position {query_pos}', fontsize=12)
axes[0].legend()
axes[0].grid(True, alpha=0.3, axis='y')

# 対角線からの距離に対するAttention重みの平均
max_d = seq_len - 1
avg_attn_standard = np.zeros(max_d + 1)
avg_attn_shaw = np.zeros(max_d + 1)
avg_attn_t5 = np.zeros(max_d + 1)
counts = np.zeros(max_d + 1)

for i in range(seq_len):
    for j in range(seq_len):
        d = abs(i - j)
        avg_attn_standard[d] += standard_attn[i, j]
        avg_attn_shaw[d] += shaw_attn[i, j]
        avg_attn_t5[d] += t5_attn[i, j]
        counts[d] += 1

avg_attn_standard /= counts
avg_attn_shaw /= counts
avg_attn_t5 /= counts

axes[1].plot(range(max_d + 1), avg_attn_standard, 'gray', linewidth=2,
             marker='o', markersize=3, label='Standard')
axes[1].plot(range(max_d + 1), avg_attn_shaw, 'b-', linewidth=2,
             marker='s', markersize=3, label='Shaw')
axes[1].plot(range(max_d + 1), avg_attn_t5, 'g-', linewidth=2,
             marker='^', markersize=3, label='T5')
axes[1].set_xlabel('Absolute Distance |i - j|', fontsize=12)
axes[1].set_ylabel('Average Attention Weight', fontsize=12)
axes[1].set_title('Average Attention Weight vs Distance', fontsize=12)
axes[1].legend()
axes[1].grid(True, alpha=0.3)

plt.tight_layout()
plt.savefig('attention_distance_decay.png', dpi=150, bbox_inches='tight')
plt.show()