Span Corruption — T5方式の事前学習を理論と実装で完全理解する

import numpy as np
from typing import List, Tuple

def sample_span_lengths(
    num_tokens: int,
    corruption_rate: float = 0.15,
    mean_span_length: float = 3.0,
    rng: np.random.Generator = None
) -> List[int]:
    """幾何分布からスパン長をサンプリングする"""
    if rng is None:
        rng = np.random.default_rng(42)

    # マスクする総トークン数
    num_masked = int(np.round(num_tokens * corruption_rate))
    if num_masked == 0:
        return []

    # 幾何分布のパラメータ p = 1/μ
    p = 1.0 / mean_span_length

    # スパン長を繰り返しサンプリング
    span_lengths = []
    total = 0
    while total < num_masked:
        # 幾何分布からサンプリング（最低1）
        length = rng.geometric(p)
        # 残りのトークン数を超えないようにクリップ
        length = min(length, num_masked - total)
        span_lengths.append(length)
        total += length

    return span_lengths


def create_span_corruption(
    token_ids: List[int],
    sentinel_start_id: int = 32000,
    corruption_rate: float = 0.15,
    mean_span_length: float = 3.0,
    seed: int = 42
) -> Tuple[List[int], List[int]]:
    """Span Corruptionによる入力系列とターゲット系列を生成する"""
    rng = np.random.default_rng(seed)
    n = len(token_ids)

    # スパン長をサンプリング
    span_lengths = sample_span_lengths(
        n, corruption_rate, mean_span_length, rng
    )
    k = len(span_lengths)  # スパン数
    if k == 0:
        return token_ids[:], []

    num_masked = sum(span_lengths)
    num_unmasked = n - num_masked

    # 非マスクトークンをk+1個のセグメントにランダム分割
    # （多項分布を使って均一にランダム配置）
    segment_lengths = np.zeros(k + 1, dtype=int)
    positions = rng.choice(num_unmasked + k, k, replace=False)
    positions.sort()
    # スターズ・アンド・バーズ法で分割
    positions_shifted = positions - np.arange(k)
    segment_lengths[0] = positions_shifted[0]
    for i in range(1, k):
        segment_lengths[i] = positions_shifted[i] - positions_shifted[i-1] - 1
    segment_lengths[k] = num_unmasked - positions_shifted[-1] - 1

    # マスク配列を構成
    mask = np.zeros(n, dtype=bool)
    pos = 0
    span_starts = []
    for i in range(k):
        pos += segment_lengths[i]  # 非マスクセグメント
        span_starts.append(pos)
        for j in range(span_lengths[i]):
            if pos + j < n:
                mask[pos + j] = True
        pos += span_lengths[i]

    # 入力系列の構成（スパンをSentinelで置換）
    input_ids = []
    sentinel_idx = 0
    i = 0
    while i < n:
        if mask[i]:
            # Sentinelトークンを挿入
            input_ids.append(sentinel_start_id + sentinel_idx)
            sentinel_idx += 1
            # スパンの終端までスキップ
            while i < n and mask[i]:
                i += 1
        else:
            input_ids.append(token_ids[i])
            i += 1

    # ターゲット系列の構成
    target_ids = []
    sentinel_idx = 0
    i = 0
    while i < n:
        if mask[i]:
            target_ids.append(sentinel_start_id + sentinel_idx)
            sentinel_idx += 1
            while i < n and mask[i]:
                target_ids.append(token_ids[i])
                i += 1
        else:
            i += 1

    return input_ids, target_ids

# デモ: テキストをトークンIDに見立てて実行
np.random.seed(42)
# 簡易トークン列（0〜19の整数をトークンIDに見立てる）
tokens = list(range(20))
print(f"元の系列:       {tokens}")
print(f"系列長:         {len(tokens)}")

input_ids, target_ids = create_span_corruption(
    tokens, sentinel_start_id=100, corruption_rate=0.15, mean_span_length=3.0
)
print(f"入力系列:       {input_ids}")
print(f"ターゲット系列: {target_ids}")
print(f"入力長: {len(input_ids)}, ターゲット長: {len(target_ids)}")

import numpy as np
import matplotlib.pyplot as plt

# スパン長のサンプリングを大量に繰り返す
rng = np.random.default_rng(42)
mean_span_length = 3.0
p = 1.0 / mean_span_length
num_samples = 50000

sampled_lengths = rng.geometric(p, size=num_samples)

# 理論分布
max_length = 15
l_values = np.arange(1, max_length + 1)
theoretical_probs = (1 - p) ** (l_values - 1) * p

# 経験分布
empirical_probs = np.array([
    np.sum(sampled_lengths == l) / num_samples for l in l_values
])

# 可視化
fig, ax = plt.subplots(figsize=(10, 6))
width = 0.35
ax.bar(l_values - width/2, theoretical_probs, width, label='Geometric (theory)',
       color='#00bcd4', alpha=0.8, edgecolor='white')
ax.bar(l_values + width/2, empirical_probs, width, label='Sampled (empirical)',
       color='#ff9800', alpha=0.8, edgecolor='white')
ax.set_xlabel('Span Length', fontsize=13)
ax.set_ylabel('Probability', fontsize=13)
ax.set_title(f'Span Length Distribution (mean={mean_span_length})', fontsize=14)
ax.set_xticks(l_values)
ax.legend(fontsize=12)
ax.grid(axis='y', alpha=0.3)
plt.tight_layout()
plt.savefig('span_length_distribution.png', dpi=150, bbox_inches='tight')
plt.show()

print(f"理論平均: {1/p:.2f}")
print(f"経験平均: {np.mean(sampled_lengths):.2f}")
print(f"理論分散: {(1-p)/p**2:.2f}")
print(f"経験分散: {np.var(sampled_lengths):.2f}")

import numpy as np
import matplotlib.pyplot as plt

fig, axes = plt.subplots(1, 3, figsize=(15, 5))
mean_lengths = [2, 3, 5]
colors = ['#e91e63', '#00bcd4', '#4caf50']

for ax, mu, color in zip(axes, mean_lengths, colors):
    p = 1.0 / mu
    l_values = np.arange(1, 16)
    probs = (1 - p) ** (l_values - 1) * p

    ax.bar(l_values, probs, color=color, alpha=0.8, edgecolor='white')
    ax.set_xlabel('Span Length', fontsize=12)
    ax.set_ylabel('Probability', fontsize=12)
    ax.set_title(f'μ = {mu} (p = {p:.3f})', fontsize=13)
    ax.set_xticks(l_values[::2])
    ax.set_ylim(0, 0.55)
    ax.grid(axis='y', alpha=0.3)
    # 期待値を縦線で表示
    ax.axvline(x=mu, color='red', linestyle='--', alpha=0.7, label=f'E[L]={mu}')
    ax.legend(fontsize=11)

plt.suptitle('Span Length Distribution for Different Mean Lengths',
             fontsize=14, y=1.02)
plt.tight_layout()
plt.savefig('span_length_comparison.png', dpi=150, bbox_inches='tight')
plt.show()

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches

def visualize_span_corruption(
    words: list,
    corruption_rate: float = 0.15,
    mean_span_length: float = 3.0,
    seed: int = 42
):
    """テキストに対するSpan Corruptionを可視化する"""
    n = len(words)
    token_ids = list(range(n))

    input_ids, target_ids = create_span_corruption(
        token_ids, sentinel_start_id=1000,
        corruption_rate=corruption_rate,
        mean_span_length=mean_span_length,
        seed=seed
    )

    # マスク位置を特定
    masked_positions = set()
    sentinel_spans = {}
    sentinel_count = 0

    for i, tid in enumerate(input_ids):
        if tid >= 1000:
            sentinel_spans[tid] = sentinel_count
            sentinel_count += 1

    # 元のtoken_idから復元
    mask_array = np.zeros(n, dtype=int)  # 0: unmasked, 1+: sentinel_id
    pos = 0
    sid = 0
    for iid in input_ids:
        if iid >= 1000:
            # このsentinelに対応する元のトークンを特定
            span_tokens = []
            in_target = False
            current_sentinel = iid
            for tid in target_ids:
                if tid == current_sentinel:
                    in_target = True
                    continue
                elif tid >= 1000:
                    if in_target:
                        break
                    continue
                elif in_target:
                    span_tokens.append(tid)
            for t in span_tokens:
                mask_array[t] = sid + 1
            sid += 1
        else:
            pos += 1

    # 可視化
    fig, ax = plt.subplots(figsize=(16, 4))
    colors_map = plt.cm.Set2(np.linspace(0, 1, max(sid, 1) + 1))

    for i, word in enumerate(words):
        if mask_array[i] > 0:
            bg_color = colors_map[mask_array[i] - 1]
            label = f'<extra_id_{mask_array[i]-1}>'
        else:
            bg_color = '#2a2a3e'
            label = None

        rect = mpatches.FancyBboxPatch(
            (i * 1.6, 0.5), 1.4, 0.8,
            boxstyle="round,pad=0.1",
            facecolor=bg_color,
            edgecolor='white' if mask_array[i] == 0 else 'red',
            linewidth=1.5
        )
        ax.add_patch(rect)
        ax.text(i * 1.6 + 0.7, 0.9, word,
                ha='center', va='center', fontsize=9,
                color='white' if mask_array[i] == 0 else 'black',
                fontweight='bold')
        if mask_array[i] > 0:
            ax.text(i * 1.6 + 0.7, 0.6, f'<id_{mask_array[i]-1}>',
                    ha='center', va='center', fontsize=7,
                    color='red', fontstyle='italic')

    ax.set_xlim(-0.2, len(words) * 1.6 + 0.2)
    ax.set_ylim(0, 2)
    ax.set_aspect('equal')
    ax.axis('off')
    ax.set_title('Span Corruption Visualization', fontsize=14, pad=10)

    # 凡例
    legend_elements = [
        mpatches.Patch(facecolor='#2a2a3e', edgecolor='white', label='Unmasked'),
        mpatches.Patch(facecolor=colors_map[0], edgecolor='red', label='Masked Span'),
    ]
    ax.legend(handles=legend_elements, loc='upper right', fontsize=10)

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

# サンプルテキスト
words = ["The", "quick", "brown", "fox", "jumps", "over",
         "the", "lazy", "dog", "and", "runs", "into",
         "the", "dark", "forest", "at", "night"]
visualize_span_corruption(words, corruption_rate=0.15, mean_span_length=3.0, seed=42)

import numpy as np
from typing import List, Tuple

def create_mlm_masking(
    token_ids: List[int],
    mask_token_id: int = 99,
    corruption_rate: float = 0.15,
    seed: int = 42
) -> Tuple[List[int], List[int], List[int]]:
    """BERTスタイルのMLMマスキングを実行する"""
    rng = np.random.default_rng(seed)
    n = len(token_ids)
    num_masked = int(np.round(n * corruption_rate))

    # ランダムにマスク位置を選択
    mask_positions = sorted(rng.choice(n, num_masked, replace=False))

    # 入力系列（マスクトークンで置換）
    input_ids = token_ids[:]
    for pos in mask_positions:
        input_ids[pos] = mask_token_id

    # ターゲット: マスク位置の正解トークン
    target_tokens = [token_ids[pos] for pos in mask_positions]

    return input_ids, target_tokens, mask_positions

import numpy as np
import matplotlib.pyplot as plt

def compare_sequence_lengths(
    seq_lengths: list,
    corruption_rate: float = 0.15,
    mean_span_length: float = 3.0,
    num_trials: int = 100
):
    """複数の系列長でMLMとSpan Corruptionの出力長を比較する"""
    results = {
        'seq_len': [],
        'mlm_input_len': [], 'mlm_target_len': [],
        'sc_input_len': [], 'sc_target_len': []
    }

    for n in seq_lengths:
        mlm_inputs, mlm_targets = [], []
        sc_inputs, sc_targets = [], []

        for trial in range(num_trials):
            tokens = list(range(n))

            # MLM
            mlm_in, mlm_tgt, _ = create_mlm_masking(
                tokens, mask_token_id=99,
                corruption_rate=corruption_rate, seed=trial
            )
            mlm_inputs.append(len(mlm_in))
            mlm_targets.append(len(mlm_tgt))

            # Span Corruption
            sc_in, sc_tgt = create_span_corruption(
                tokens, sentinel_start_id=100,
                corruption_rate=corruption_rate,
                mean_span_length=mean_span_length, seed=trial
            )
            sc_inputs.append(len(sc_in))
            sc_targets.append(len(sc_tgt))

        results['seq_len'].append(n)
        results['mlm_input_len'].append(np.mean(mlm_inputs))
        results['mlm_target_len'].append(np.mean(mlm_targets))
        results['sc_input_len'].append(np.mean(sc_inputs))
        results['sc_target_len'].append(np.mean(sc_targets))

    return results

seq_lengths = [32, 64, 128, 256, 512]
results = compare_sequence_lengths(seq_lengths)

# 可視化
fig, axes = plt.subplots(1, 2, figsize=(14, 6))

# 入力系列長
ax = axes[0]
ax.plot(results['seq_len'], results['mlm_input_len'],
        'o-', color='#e91e63', linewidth=2, markersize=8, label='MLM Input')
ax.plot(results['seq_len'], results['sc_input_len'],
        's-', color='#00bcd4', linewidth=2, markersize=8, label='Span Corruption Input')
ax.plot(results['seq_len'], results['seq_len'],
        '--', color='gray', alpha=0.5, label='Original Length')
ax.set_xlabel('Original Sequence Length', fontsize=13)
ax.set_ylabel('Input Sequence Length', fontsize=13)
ax.set_title('Input Sequence Length Comparison', fontsize=14)
ax.legend(fontsize=11)
ax.grid(alpha=0.3)

# ターゲット系列長
ax = axes[1]
ax.plot(results['seq_len'], results['mlm_target_len'],
        'o-', color='#e91e63', linewidth=2, markersize=8, label='MLM Target')
ax.plot(results['seq_len'], results['sc_target_len'],
        's-', color='#00bcd4', linewidth=2, markersize=8, label='Span Corruption Target')
ax.plot(results['seq_len'], results['seq_len'],
        '--', color='gray', alpha=0.5, label='Original Length')
ax.set_xlabel('Original Sequence Length', fontsize=13)
ax.set_ylabel('Target Sequence Length', fontsize=13)
ax.set_title('Target Sequence Length Comparison', fontsize=14)
ax.legend(fontsize=11)
ax.grid(alpha=0.3)

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

# 数値比較
print("=" * 60)
print(f"{'SeqLen':>8} | {'MLM Input':>10} {'MLM Tgt':>10} | "
      f"{'SC Input':>10} {'SC Tgt':>10} | {'Tgt Ratio':>10}")
print("-" * 60)
for i, n in enumerate(results['seq_len']):
    ratio = results['sc_target_len'][i] / results['mlm_target_len'][i]
    print(f"{n:>8} | {results['mlm_input_len'][i]:>10.1f} "
          f"{results['mlm_target_len'][i]:>10.1f} | "
          f"{results['sc_input_len'][i]:>10.1f} "
          f"{results['sc_target_len'][i]:>10.1f} | {ratio:>10.2f}")

import numpy as np
import matplotlib.pyplot as plt

def compute_mask_diversity(
    num_tokens: int = 100,
    corruption_rate: float = 0.15,
    mean_span_length: float = 3.0,
    num_trials: int = 500
):
    """マスクパターンの多様性を測定する"""
    tokens = list(range(num_tokens))

    # 各トークン位置がマスクされる頻度を集計
    mlm_freq = np.zeros(num_tokens)
    sc_freq = np.zeros(num_tokens)

    for trial in range(num_trials):
        # MLM
        _, _, mask_positions = create_mlm_masking(
            tokens, corruption_rate=corruption_rate, seed=trial
        )
        for pos in mask_positions:
            mlm_freq[pos] += 1

        # Span Corruption
        sc_in, _ = create_span_corruption(
            tokens, sentinel_start_id=1000,
            corruption_rate=corruption_rate,
            mean_span_length=mean_span_length, seed=trial
        )
        # 入力系列に含まれないトークンIDを特定
        remaining_ids = set(tid for tid in sc_in if tid < 1000)
        for i in range(num_tokens):
            if i not in remaining_ids:
                sc_freq[i] += 1

    mlm_freq /= num_trials
    sc_freq /= num_trials

    return mlm_freq, sc_freq

mlm_freq, sc_freq = compute_mask_diversity(num_tokens=100, num_trials=500)

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

# MLMのマスク頻度
axes[0].bar(range(100), mlm_freq, color='#e91e63', alpha=0.7, edgecolor='none')
axes[0].axhline(y=0.15, color='black', linestyle='--', alpha=0.5, label='Expected Rate (15%)')
axes[0].set_xlabel('Token Position', fontsize=12)
axes[0].set_ylabel('Masking Frequency', fontsize=12)
axes[0].set_title('MLM: Per-Position Masking Frequency (500 trials)', fontsize=13)
axes[0].set_ylim(0, 0.3)
axes[0].legend(fontsize=11)
axes[0].grid(axis='y', alpha=0.3)

# Span Corruptionのマスク頻度
axes[1].bar(range(100), sc_freq, color='#00bcd4', alpha=0.7, edgecolor='none')
axes[1].axhline(y=0.15, color='black', linestyle='--', alpha=0.5, label='Expected Rate (15%)')
axes[1].set_xlabel('Token Position', fontsize=12)
axes[1].set_ylabel('Masking Frequency', fontsize=12)
axes[1].set_title('Span Corruption: Per-Position Masking Frequency (500 trials)', fontsize=13)
axes[1].set_ylim(0, 0.3)
axes[1].legend(fontsize=11)
axes[1].grid(axis='y', alpha=0.3)

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

# 統計量の比較
print(f"MLM       - 平均マスク率: {np.mean(mlm_freq):.4f}, "
      f"分散: {np.var(mlm_freq):.6f}")
print(f"Span Corr - 平均マスク率: {np.mean(sc_freq):.4f}, "
      f"分散: {np.var(sc_freq):.6f}")

import numpy as np
import matplotlib.pyplot as plt

corruption_rates = [0.10, 0.15, 0.20, 0.25, 0.30, 0.50]
seq_len = 256
num_trials = 200

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

# ターゲット系列長
mean_target_lens = []
std_target_lens = []
mean_input_lens = []
mean_num_spans = []

for rate in corruption_rates:
    target_lens = []
    input_lens = []
    num_spans = []
    for trial in range(num_trials):
        tokens = list(range(seq_len))
        sc_in, sc_tgt = create_span_corruption(
            tokens, sentinel_start_id=1000,
            corruption_rate=rate, mean_span_length=3.0, seed=trial
        )
        target_lens.append(len(sc_tgt))
        input_lens.append(len(sc_in))
        # Sentinel数を数える
        n_sentinels = sum(1 for t in sc_in if t >= 1000)
        num_spans.append(n_sentinels)

    mean_target_lens.append(np.mean(target_lens))
    std_target_lens.append(np.std(target_lens))
    mean_input_lens.append(np.mean(input_lens))
    mean_num_spans.append(np.mean(num_spans))

# 左: ターゲット系列長
ax = axes[0]
ax.errorbar(corruption_rates, mean_target_lens, yerr=std_target_lens,
            fmt='o-', color='#00bcd4', linewidth=2, markersize=8,
            capsize=5, label='Target Length')
ax.plot(corruption_rates, mean_input_lens,
        's-', color='#ff9800', linewidth=2, markersize=8, label='Input Length')
ax.axhline(y=seq_len, color='gray', linestyle='--', alpha=0.5,
           label=f'Original ({seq_len})')
ax.set_xlabel('Corruption Rate', fontsize=13)
ax.set_ylabel('Sequence Length', fontsize=13)
ax.set_title('Effect of Corruption Rate on Sequence Length', fontsize=14)
ax.legend(fontsize=11)
ax.grid(alpha=0.3)

# 右: スパン数
ax = axes[1]
theoretical_spans = [rate * seq_len / 3.0 for rate in corruption_rates]
ax.plot(corruption_rates, mean_num_spans,
        'o-', color='#4caf50', linewidth=2, markersize=8, label='Empirical')
ax.plot(corruption_rates, theoretical_spans,
        '--', color='red', linewidth=2, alpha=0.7, label='Theoretical (rn/μ)')
ax.set_xlabel('Corruption Rate', fontsize=13)
ax.set_ylabel('Number of Spans', fontsize=13)
ax.set_title('Number of Spans vs Corruption Rate', fontsize=14)
ax.legend(fontsize=11)
ax.grid(alpha=0.3)

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

# 数値の表示
print(f"{'Rate':>6} | {'Input Len':>10} {'Target Len':>12} | "
      f"{'#Spans':>8} {'Theory':>8} | {'Target/Orig':>12}")
print("-" * 65)
for i, rate in enumerate(corruption_rates):
    print(f"{rate:>6.2f} | {mean_input_lens[i]:>10.1f} "
          f"{mean_target_lens[i]:>12.1f} | "
          f"{mean_num_spans[i]:>8.1f} {theoretical_spans[i]:>8.1f} | "
          f"{mean_target_lens[i]/seq_len:>12.2%}")