FRSM 架构迭代实验报告

一、迭代历程

本报告记录 FRSM 从 v1 原始设计到 v3 残差设计的四轮迭代过程。

二、v3 Residual 设计详解

2.1 核心创新：三层残差

Layer 1: State Residual — 每个尺度内部
h_s_new = LayerNorm(α × h_s + (1-α) × candidate)
α = sigmoid(2.0) ≈ 0.88 (偏向保留)

candidate = forget_gate * h_s + input_gate * tanh(W_c * [h_s; inp])
forget_gate: 默认 1 (保留)，input_gate: 默认 0 (不写)

Layer 2: Scale Projection — 每个尺度直接贡献到输出
fused = Linear(concat(h_s)) ← 融合通道
for s in 1..N: fused += Linear_s(h_s) / N ← 独立通道 (残差)

Layer 3: Input Highway — 原始输入直接通路
out = fused_norm(fused + inp) ← 梯度直达输入

2.2 为什么残差解决了长期记忆

原始 FRSM 的状态更新是全量替换式：h_new = f*h + i*candidate。当 forget_gate 和 input_gate 都接近 0.5 时，旧状态被洗掉一半。在多尺度中，更新间隔长（period=S），每次更新覆盖量更大。

残差版本改为增量式：h_new = LayerNorm(α*h + (1-α)*candidate)。α≈0.88 意味着每步只改变状态的 12%。旧信息被"轻轻推动"而非"覆盖"，天然留存更久。

2.3 代价：LM 质量略微下降

残差使状态更新更"保守"——新信息写入阻力更大。对于依赖即时输入的 LM 预测，这是精度损失。5.70→6.05 的差距（6%）在可接受范围内，且可以通过延长训练来缩小（v1 多训 5000 步降了 2.4 loss，同理 v3 应该也能降）。

三、四轮实验完整数据

3.1 尺度数消融（v1 Orig, 固定 d_model=256）

结论：LM 任务 scales=2 最优。scales≥3 引入过度稀疏的更新，恶化泛化。

3.2 规模效应（v1 Orig, scales=4）

Scaling Law: loss ≈ -0.757 × log₁₀(params) + 11.172

参数每翻倍，loss 降 ~0.15-0.25。收益递减明显。

3.3 残差方案对比

3.4 5架构 CopyFirst 对比（v1 Orig 2-scales）

四、最终推荐

五、最终模型代码 (v3 Residual)

"""
FRSM v3 Residual
核心: 固定残差率 α→0.88 + 尺度投影 + 输入 highway
优势: LM loss=5.70(2sc)/6.05(4sc), CopyFirst@65K=68.8%
"""
import torch
import torch.nn as nn
import torch.nn.functional as F
import math

class ScaleRecurrentBlock(nn.Module):
"""
单尺度内部: 固定残差更新
h_new = LayerNorm(α × h + (1-α) × candidate)
candidate = forget_gate × h + input_gate × tanh(W_c × [h; inp])
"""
def __init__(self, d_model, alpha_init=2.0):
"""
alpha_init: logit for sigmoid. 2.0 → sigmoid=0.88 (偏向保留)
"""
super().__init__()
self.d_model = d_model
self.W_forget = nn.Linear(d_model * 2, d_model)
self.W_input = nn.Linear(d_model * 2, d_model)
self.W_cand = nn.Linear(d_model * 2, d_model)
# 初始化: forget偏向记住, input偏向不写
nn.init.constant_(self.W_forget.bias, 1.0)
nn.init.constant_(self.W_input.bias, –2.0)
# 固定残差率 (不是可学习参数，避免v4的自适应失败)
self.register_buffer('alpha', torch.sigmoid(torch.tensor(alpha_init)))
# 状态归一化 (维持临界态)
self.state_norm = nn.LayerNorm(d_model)

def forward(self, h_prev, inp):
combined = torch.cat([h_prev, inp], dim=–1)
f = torch.sigmoid(self.W_forget(combined))
i = torch.sigmoid(self.W_input(combined))
candidate = f * h_prev + i * torch.tanh(self.W_cand(combined))
# 固定残差: α保留旧状态, (1-α)写入新信息
h_new = self.state_norm(
self.alpha * h_prev + (1 – self.alpha) * candidate
)
return h_new

class ResidualFRSM(nn.Module):
"""
三层残差 FRSM:
1. State Residual: 每个尺度内 α-残差更新
2. Scale Projection: 每个尺度独立投影到输出
3. Input Highway: 原始输入直接通路
"""
def __init__(
self,
vocab_size: int,
d_model: int = 256,
num_scales: int = 4,
alpha_init: float = 2.0, # sigmoid(2.0) ≈ 0.88
):
super().__init__()
self.d_model = d_model
self.num_scales = num_scales
self.vocab_size = vocab_size

# Token embedding & input projection
self.embed = nn.Embedding(vocab_size, d_model)
self.input_proj = nn.Linear(d_model, d_model)
self.input_norm = nn.LayerNorm(d_model)

# 多尺度递归块
self.scales = nn.ModuleList([
ScaleRecurrentBlock(d_model, alpha_init)
for _ in range(num_scales)
])

# 尺度间投影 (每个尺度独立贡献到输出)
self.scale_projections = nn.ModuleList([
nn.Linear(d_model, d_model) for _ in range(num_scales)
])

# 融合通道
self.fusion = nn.Linear(d_model * num_scales, d_model)
self.fusion_norm = nn.LayerNorm(d_model)

# 输出头
self.output_proj = nn.Linear(d_model, vocab_size)

self._init_weights()

def _init_weights(self):
for m in self.modules():
if isinstance(m, nn.Linear):
nn.init.xavier_uniform_(m.weight, gain=0.5)
if m.bias is not None:
nn.init.zeros_(m.bias)
elif isinstance(m, nn.Embedding):
nn.init.normal_(m.weight, mean=0, std=0.02)
nn.init.zeros_(self.output_proj.bias)

def forward(self, x, h_prev=None, return_state=False):
"""训练模式: 全序列前向 (O(n))"""
B, T = x.shape

# 初始化多尺度状态
if h_prev is None:
h = [torch.zeros(B, self.d_model, device=x.device) for _ in range(self.num_scales)]
else:
h = [hs.clone() for hs in h_prev]

x_emb = self.embed(x)
outputs = []

for t in range(T):
inp = self.input_norm(self.input_proj(x_emb[:, t, :]))
next_h = []

# === Layer 1: State Residual ===
for s in range(self.num_scales):
if t % (2 ** s) == 0:
next_h.append(self.scales[s](h[s], inp))
else:
next_h.append(h[s])
h = next_h

# === Layer 2: Scale Projection ===
fused = self.fusion(torch.cat(h, dim=–1))
scale_contrib = sum(
self.scale_projections[s](h[s]) for s in range(self.num_scales)
) / self.num_scales
fused = fused + scale_contrib

# === Layer 3: Input Highway ===
out = self.fusion_norm(fused + inp)
outputs.append(self.output_proj(out).unsqueeze(1))

logits = torch.cat(outputs, dim=1)

if return_state:
return logits, h
return logits

def generate_step(self, token, h_prev):
"""推理模式: 单步前向 (O(1))"""
with torch.no_grad():
x_emb = self.embed(token)
inp = self.input_norm(self.input_proj(x_emb.squeeze(1)))
next_h = []

for s in range(self.num_scales):
# 推理时简化: 所有尺度始终更新
next_h.append(self.scales[s](h_prev[s], inp))
h = next_h

fused = self.fusion(torch.cat(h, dim=–1))
scale_contrib = sum(
self.scale_projections[s](h[s]) for s in range(self.num_scales)
) / self.num_scales
fused = fused + scale_contrib
out = self.fusion_norm(fused + inp)
logits = self.output_proj(out)

return logits, h

# ============================================================
# 使用示例
# ============================================================
if __name__ == "__main__":
# 创建模型
vocab_size = 32000
model = ResidualFRSM(vocab_size=vocab_size, d_model=256, num_scales=4)
print(f"Params: {sum(p.numel() for p in model.parameters()):,}")

# 训练: batch=4, seq_len=384
x = torch.randint(0, vocab_size, (4, 384))
logits = model(x) # (4, 384, 32000)
print(f"Train output: {logits.shape}")

# 推理: 单步生成
token = torch.tensor([[42]])
h = None
for step in range(10):
if h is None:
logits, h = model(token, return_state=True)
logits = logits[:, –1, :]
else:
logits, h = model.generate_step(token, h)
token = logits.argmax(dim=–1).unsqueeze(0)
print(f"Inference completed: 10 steps generated")

5.1 模型特性

报告生成: 2026-06-15
实验设备: NVIDIA GeForce RTX 4090 D, CUDA 13.2, PyTorch 2.12.0