Time-MoE — 24億パラメータのMoE時系列基盤モデル【ICLR 2025】

import numpy as np
import matplotlib.pyplot as plt

np.random.seed(42)

# --- 合成時系列データ（複数パターン混在）---
def generate_mixed_series(length=2000):
    """複数パターンが混在する時系列"""
    t = np.arange(length, dtype=float)
    series = np.zeros(length)
    pattern_labels = np.zeros(length, dtype=int)

    for i in range(0, length, 200):
        seg_type = np.random.choice([0, 1, 2, 3])
        end = min(i + 200, length)
        seg_t = t[i:end] - t[i]
        if seg_type == 0:  # 定常
            series[i:end] = 0.5 + np.random.randn(end-i) * 0.1
        elif seg_type == 1:  # 周期
            series[i:end] = np.sin(2 * np.pi * seg_t / 50) + np.random.randn(end-i) * 0.1
        elif seg_type == 2:  # トレンド
            series[i:end] = np.linspace(-1, 1, end-i) + np.random.randn(end-i) * 0.1
        else:  # パルス（異常）
            series[i:end] = np.random.randn(end-i) * 0.1
            for _ in range(3):
                pos = np.random.randint(0, end-i-5)
                series[i+pos:i+pos+5] += np.random.choice([-3, 3])
        pattern_labels[i:end] = seg_type
    return series, pattern_labels

series, labels = generate_mixed_series(length=2000)

# --- パッチ化 ---
patch_size = 32
stride = 16
patches = []
patch_labels = []
for i in range(0, len(series) - patch_size, stride):
    patches.append(series[i:i+patch_size])
    # パッチの主パターン（最頻値）
    patch_labels.append(np.bincount(labels[i:i+patch_size]).argmax())
patches = np.array(patches)
patch_labels = np.array(patch_labels)

# --- 簡易MoEモデル ---
class SimpleMoETimeSeries:
    def __init__(self, input_dim=32, hidden_dim=16, n_experts=4, top_k=2, lr=0.01):
        self.n_experts = n_experts
        self.top_k = top_k
        self.lr = lr

        # ゲーティングネットワーク
        self.W_gate = np.random.randn(input_dim, n_experts) * 0.1

        # エキスパート（各エキスパートは独立なFFN）
        self.experts_W1 = [np.random.randn(input_dim, hidden_dim) * 0.1 for _ in range(n_experts)]
        self.experts_b1 = [np.zeros(hidden_dim) for _ in range(n_experts)]
        self.experts_W2 = [np.random.randn(hidden_dim, input_dim) * 0.1 for _ in range(n_experts)]
        self.experts_b2 = [np.zeros(input_dim) for _ in range(n_experts)]

    def gate(self, x):
        logits = x @ self.W_gate
        probs = np.exp(logits - logits.max(axis=1, keepdims=True))
        probs /= probs.sum(axis=1, keepdims=True)
        return probs

    def expert_forward(self, x, expert_idx):
        h = np.maximum(0, x @ self.experts_W1[expert_idx] + self.experts_b1[expert_idx])
        return h @ self.experts_W2[expert_idx] + self.experts_b2[expert_idx]

    def forward(self, x):
        gate_probs = self.gate(x)
        batch_size = len(x)
        output = np.zeros_like(x)
        routing_info = []

        for b in range(batch_size):
            top_k_idx = np.argsort(gate_probs[b])[-self.top_k:]
            top_k_probs = gate_probs[b, top_k_idx]
            top_k_probs /= top_k_probs.sum()

            sample_output = np.zeros(x.shape[1])
            for k, idx in enumerate(top_k_idx):
                expert_out = self.expert_forward(x[b:b+1], idx)[0]
                sample_output += top_k_probs[k] * expert_out
            output[b] = sample_output
            routing_info.append(top_k_idx.tolist())

        return output, gate_probs, routing_info

    def train_step(self, x_input, x_target):
        output, gate_probs, routing_info = self.forward(x_input)
        loss = np.mean((output - x_target) ** 2)

        # 負荷分散損失
        expert_counts = np.zeros(self.n_experts)
        for ri in routing_info:
            for idx in ri:
                expert_counts[idx] += 1
        expert_frac = expert_counts / len(x_input)
        expert_prob_mean = gate_probs.mean(axis=0)
        balance_loss = self.n_experts * np.sum(expert_frac * expert_prob_mean)

        return loss, balance_loss, routing_info

# --- 学習 ---
model = SimpleMoETimeSeries(input_dim=32, hidden_dim=16, n_experts=4, top_k=2, lr=0.01)

# 次パッチ予測タスク
X_input = patches[:-1]
X_target = patches[1:]
labels_input = patch_labels[:-1]

n_epochs = 50
batch_size = 64
all_routing = []

for epoch in range(n_epochs):
    perm = np.random.permutation(len(X_input))
    epoch_routing = []
    for i in range(0, len(X_input), batch_size):
        idx = perm[i:i+batch_size]
        loss, bl, ri = model.train_step(X_input[idx], X_target[idx])
        epoch_routing.extend(zip(labels_input[idx], ri))
    all_routing.append(epoch_routing)

# --- ルーティング分析の可視化 ---
last_epoch = all_routing[-1]
pattern_names = ['Stationary', 'Periodic', 'Trend', 'Anomaly']
expert_pattern_matrix = np.zeros((4, 4))  # [pattern, expert]

for label, experts in last_epoch:
    for e in experts:
        expert_pattern_matrix[label, e] += 1

# 行ごとに正規化
row_sums = expert_pattern_matrix.sum(axis=1, keepdims=True)
expert_pattern_norm = expert_pattern_matrix / (row_sums + 1e-8)

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

im = axes[0].imshow(expert_pattern_norm, cmap='YlOrRd', aspect='auto')
axes[0].set_xticks(range(4))
axes[0].set_xticklabels([f'Expert {i}' for i in range(4)])
axes[0].set_yticks(range(4))
axes[0].set_yticklabels(pattern_names)
axes[0].set_title('Expert Routing by Pattern Type (normalized)')
for i in range(4):
    for j in range(4):
        axes[0].text(j, i, f'{expert_pattern_norm[i,j]:.2f}',
                     ha='center', va='center', fontsize=10,
                     color='white' if expert_pattern_norm[i,j] > 0.4 else 'black')
plt.colorbar(im, ax=axes[0])

# 各エキスパートの総使用回数
expert_total = expert_pattern_matrix.sum(axis=0)
colors = ['#00d4ff', '#ff6b6b', '#ffd93d', '#6bff6b']
bars = axes[1].bar(range(4), expert_total, color=colors)
axes[1].set_xticks(range(4))
axes[1].set_xticklabels([f'Expert {i}' for i in range(4)])
axes[1].set_ylabel('Total routing count')
axes[1].set_title('Expert Load Distribution')
axes[1].grid(True, alpha=0.3, axis='y')

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

# --- スケーリング則のシミュレーション ---
model_sizes = [4, 8, 16, 32, 64]  # エキスパート内の隠れ層次元
losses_by_size = []

for hidden_dim in model_sizes:
    model_test = SimpleMoETimeSeries(
        input_dim=32, hidden_dim=hidden_dim, n_experts=4, top_k=2, lr=0.01
    )
    # 50エポック学習
    final_losses = []
    for epoch in range(50):
        perm = np.random.permutation(len(X_input))
        epoch_loss = 0
        n_batches = 0
        for i in range(0, len(X_input), 64):
            idx = perm[i:i+64]
            loss, _, _ = model_test.train_step(X_input[idx], X_target[idx])
            epoch_loss += loss
            n_batches += 1
        final_losses.append(epoch_loss / n_batches)
    losses_by_size.append(final_losses[-1])

# パラメータ数の概算: 4 experts × (input_dim × hidden_dim + hidden_dim × input_dim)
param_counts = [4 * (32 * h + h * 32) + 32 * 4 for h in model_sizes]

fig, ax = plt.subplots(1, 1, figsize=(8, 5))
ax.loglog(param_counts, losses_by_size, 'o-', color='#00d4ff', markersize=8, linewidth=2)
ax.set_xlabel('Total Parameters (log scale)')
ax.set_ylabel('Prediction Loss (log scale)')
ax.set_title('Scaling Law: Loss vs Model Size')
ax.grid(True, alpha=0.3, which='both')

# 冪乗則フィット
log_params = np.log(param_counts)
log_losses = np.log(losses_by_size)
coeffs = np.polyfit(log_params, log_losses, 1)
fit_line = np.exp(coeffs[1]) * np.array(param_counts) ** coeffs[0]
ax.loglog(param_counts, fit_line, '--', color='#ff6b6b',
          label=f'Power law fit: L ∝ N^{{{coeffs[0]:.3f}}}')
ax.legend()

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

print(f"スケーリング指数 α = {coeffs[0]:.4f}")
print(f"パラメータ数が2倍になると損失は {2**coeffs[0]:.4f} 倍")