FRSM V6 Dense MoE vs Transformer — 全维度技术报告

核心结论

FRSM V6 Dense MoE 训练速度慢于 Transformer(同结构下 3.6x),但推理 O(1)、长序列显存恒定、总成本在推理部署场景下更优。它不是 Transformer 的替代品,而是特定场景下的更好选择。

一、FRSM 架构概述

FRSM(Fast Recurrent State Machine)是一个多尺度内容门控状态机——RNN 的现代化变体:

• 每个专家有 num_scales 个并行的时间尺度,各自维护一个状态向量
• 内容门控网络动态决定每个尺度的写入强度
• Dense MoE 版本:16 个路由专家 + 1 个共享专家,全部通过堆叠 einsum 计算
• 路由器产生软权重,专家输出按权重混合
• 共享专家始终激活,捕获通用知识

关键改进:去掉了 Sparse MoE 的 token-to-expert gather 参数拷贝(原占总步时间 77%),改为 Dense/Soft MoE 的全专家堆叠 einsum + chunk 并行。

二、性能数据(实测于 RTX 4090 D, 24GB, T=512)

2.1 公平对比:两边都是 Dense MoE 结构

公平对比:Transformer 也用上同样的 16 专家 Dense MoE,保证 FLOPs/参数量可比。

结论:同样 MoE 结构下,FRSM 比 Transformer 慢 3.6 倍。 这是 RNN 串行 vs Transformer 并行的架构性差距。

2.2 FRSM 在不同 chunk 下的训练速度

chunk 将差距从 114x 缩到 3.6x。 C=512 时 FRSM 和 Transformer 一样一次性处理全部 token,但 FRSM 的 16 专家 × 4 尺度 × 3 门控 = 192 个独立 matmul 无法融合成一个大 matmul,GPU 利用率先天不足。

2.3 推理速度(生成)

FRSM 的 generate_step 永远常数时间,Transformer 的注意力成本随序列增长。在生成长度 >1000 时 FRSM 推理反超。

2.4 序列长度与显存

FRSM 的显存与 T 弱相关(仅受 B×T logits 影响),Transformer 受 O(T²) 注意力矩阵拖累。在 T>2K 时 Transformer 先 OOM。

三、总成本分析

以训练一个 45M 模型 + 长期推理部署(1B token 生成)为例:

对于推理部署为主的场景,FRSM 的总成本比 Transformer 低 73%。 训练端的 3.6x 差距被推理端的 16x 优势轻松覆盖。

四、技术总结

五、最终结论

FRSM V6 Dense MoE 训练速度追不上 Transformer——3.6x 是 RNN 串行架构的先天上限。但它的价值不在训练速度,在:

如果你的场景以推理部署为主(对话、生成、Agent),FRSM 的长期总成本远低于 Transformer。如果追求极致训练速度,Transformer 是正确选择。

附录: FRSM V6 Dense MoE 完整代码

文件: frsm_v6_moe/frsm_v6a_dense_moe.py

"""
FRSM V6a Dense MoE — 全部专家用堆叠 einsum（无 gather/chunk/检查点）
"""
import math, torch, torch.nn as nn, torch.nn.functional as F

class FRSM_V6_DenseMoE(nn.Module):
def __init__(self, vocab_size, d_model=256, num_scales=4,
n_experts=16, n_shared=1, router_noise=1.0,
aux_loss_weight=0.01, chunk_size=0):
super().__init__()
self.d_model=d_model; self.num_scales=num_scales
self.n_experts=n_experts; self.n_shared=n_shared; self.router_noise=router_noise
self.aux_loss_weight=aux_loss_weight; self.chunk_size=chunk_size
self.aux_loss=torch.tensor(0.0)
E,S,D=n_experts,num_scales,d_model; dh=D//4

self.embed=nn.Embedding(vocab_size,D); self.input_proj=nn.Linear(D,D)
for n in ['W_forget','W_input','W_cand']:
setattr(self,n,nn.Parameter(torch.empty(E,S,D,2*D)))
setattr(self,'b_'+n[2:],nn.Parameter(torch.empty(E,S,D)))
self.gate_W1=nn.Parameter(torch.empty(E,S,dh,2*D))
self.gate_b1=nn.Parameter(torch.empty(E,S,dh))
self.gate_W2=nn.Parameter(torch.empty(E,S,1,dh))
self.gate_b2=nn.Parameter(torch.empty(E,S,1))
self.fusion_W=nn.Parameter(torch.empty(E,S*D,D))
self.fusion_b=nn.Parameter(torch.empty(E,D))

if n_shared>0:
for n in ['W_forget','W_input','W_cand']:
setattr(self,n+'_sh',nn.Parameter(torch.empty(n_shared,S,D,2*D)))
setattr(self,'b_'+n.split('_')[1]+'_sh',nn.Parameter(torch.empty(n_shared,S,D)))
self.gate_W1_sh=nn.Parameter(torch.empty(n_shared,S,dh,2*D))
self.gate_b1_sh=nn.Parameter(torch.empty(n_shared,S,dh))
self.gate_W2_sh=nn.Parameter(torch.empty(n_shared,S,1,dh))
self.gate_b2_sh=nn.Parameter(torch.empty(n_shared,S,1))
self.fusion_W_sh=nn.Parameter(torch.empty(n_shared,S*D,D))
self.fusion_b_sh=nn.Parameter(torch.empty(n_shared,D))

self.router=nn.Linear(D,E)
self.output_norm=nn.LayerNorm(D); self.output_proj=nn.Linear(D,vocab_size)
self._init_w()

def _init_w(self):
def _k(p):
for e in range(p.size(0)):
for s in range(self.num_scales):
nn.init.kaiming_uniform_(p[e,s],a=math.sqrt(5))
for pn in ['W_forget','W_input','W_cand','gate_W1','gate_W2']:
_k(getattr(self,pn))
for e in range(self.n_experts):
nn.init.kaiming_uniform_(self.fusion_W[e],a=math.sqrt(5))
if self.n_shared>0:
for pn in ['W_forget','W_input','W_cand','gate_W1','gate_W2']:
_k(getattr(self,pn+'_sh'))
for e in range(self.n_shared):
nn.init.kaiming_uniform_(getattr(self,'fusion_W_sh')[e],a=math.sqrt(5))
for n,p in self.named_parameters():
if 'bias' in n: nn.init.zeros_(p)
nn.init.zeros_(self.b_cand); nn.init.zeros_(self.gate_b1); nn.init.zeros_(self.gate_b2); nn.init.zeros_(self.fusion_b)
nn.init.constant_(self.b_forget,1.0); nn.init.constant_(self.b_input,–2.0)
if self.n_shared>0:
nn.init.zeros_(self.b_cand_sh); nn.init.zeros_(self.gate_b1_sh); nn.init.zeros_(self.gate_b2_sh); nn.init.zeros_(self.fusion_b_sh)
nn.init.constant_(self.b_forget_sh,1.0); nn.init.constant_(self.b_input_sh,–2.0)
nn.init.normal_(self.router.weight,0,0.02); nn.init.normal_(self.embed.weight,0,0.02)
nn.init.kaiming_uniform_(self.input_proj.weight,a=math.sqrt(5))
nn.init.kaiming_uniform_(self.output_proj.weight,a=math.sqrt(5))

def _estep(self, H, inp, Wf, Wi, Wc, bf, bi, bc, gW1, gb1, gW2, gb2, fW, fb):
E,B=H.shape[:2]; S,D=self.num_scales,self.d_model
inp=inp.reshape(–1,D) # (B_actual, D)
ie=inp.unsqueeze(0).unsqueeze(2).expand(E,B,S,D)
g=torch.cat([H,ie],dim=–1)
f=torch.sigmoid(torch.einsum('ebsj,esij->ebsi',g,Wf)+bf.unsqueeze(1))
i=torch.sigmoid(torch.einsum('ebsj,esij->ebsi',g,Wi)+bi.unsqueeze(1))
c=torch.tanh(torch.einsum('ebsj,esij->ebsi',g,Wc)+bc.unsqueeze(1))
cand=f*H+i*c
h1=F.gelu(torch.einsum('ebsj,esij->ebsi',g,gW1)+gb1.unsqueeze(1))
st=torch.sigmoid(torch.einsum('ebsi,esoi->ebso',h1,gW2)+gb2.unsqueeze(1))
Hn=st*cand+(1–st)*H
fused=torch.einsum('ebk,eki->ebi',Hn.reshape(E,B,S*D),fW)+fb.unsqueeze(1)
return Hn, fused

def _step(self, H, Hs, inp):
Hn, fused = self._estep(H, inp, self.W_forget,self.W_input,self.W_cand,
self.b_forget,self.b_input,self.b_cand,self.gate_W1,self.gate_b1,
self.gate_W2,self.gate_b2,self.fusion_W,self.fusion_b)
if self.n_shared>0:
Hsn, sf = self._estep(Hs, inp, self.W_forget_sh,self.W_input_sh,self.W_cand_sh,
self.b_forget_sh,self.b_input_sh,self.b_cand_sh,self.gate_W1_sh,self.gate_b1_sh,
self.gate_W2_sh,self.gate_b2_sh,self.fusion_W_sh,self.fusion_b_sh)
sf = sf.sum(dim=0) # (NS,B,D) -> (B,D)
else: Hsn, sf = None, 0
probs = self._route(inp)
combined = ((probs.t().unsqueeze(–1)*fused).sum(dim=0)) + sf
return Hn, Hsn, combined, probs

def _route(self, inp):
l=self.router(inp)
if self.training and self.router_noise>0: l=l+torch.randn_like(l)*self.router_noise
return F.softmax(l,dim=–1)

def forward(self,x,h_prev=None,return_state=False):
B,T=x.shape; E,S,D=self.n_experts,self.num_scales,self.d_model
xe=self.embed(x); iseq=self.input_proj(xe)
if h_prev is None:
H=torch.zeros(E,B,S,D,device=x.device,dtype=iseq.dtype)
Hs=torch.zeros(self.n_shared,B,S,D,device=x.device,dtype=iseq.dtype) if self.n_shared>0 else None
else: H,Hs=h_prev
logits=torch.zeros(B,T,self.output_proj.out_features,device=x.device,dtype=iseq.dtype)
aux=torch.zeros((),device=x.device,dtype=torch.float32)
C=self.chunk_size if self.chunk_size>0 else max(1,int(math.sqrt(T)))

for ts in range(0,T,C):
te=min(ts+C,T); ch=te–ts
ic=iseq[:,ts:te,:]
bch=B*ch
inf=ic.reshape(bch,D)
Hf=H.unsqueeze(2).expand(E,B,ch,S,D).reshape(E,bch,S,D)
Hsf=Hs.unsqueeze(2).expand(self.n_shared,B,ch,S,D).reshape(self.n_shared,bch,S,D) if Hs is not None else None

Hnf,fused_f=self._estep(Hf,inf,self.W_forget,self.W_input,self.W_cand,
self.b_forget,self.b_input,self.b_cand,self.gate_W1,self.gate_b1,
self.gate_W2,self.gate_b2,self.fusion_W,self.fusion_b)
if self.n_shared>0:
Hsnf,sf=self._estep(Hsf,inf,self.W_forget_sh,self.W_input_sh,self.W_cand_sh,
self.b_forget_sh,self.b_input_sh,self.b_cand_sh,self.gate_W1_sh,self.gate_b1_sh,
self.gate_W2_sh,self.gate_b2_sh,self.fusion_W_sh,self.fusion_b_sh)
else: Hsnf,sf=None,0
if sf is not None: sf = sf.sum(dim=0) # (NS,bch,D)->(bch,D)

probs=self._route(ic[:,0,:])
pbf=probs.unsqueeze(1).expand(B,ch,E).reshape(bch,E)
comb_f=((pbf.t().unsqueeze(–1)*fused_f).sum(dim=0))+sf
comb=comb_f.reshape(B,ch,D)
logits[:,ts:te,:]=self.output_proj(self.output_norm(comb))

li=torch.arange(B,device=x.device)*ch+(ch–1)
H=Hnf[:,li,:,:]
Hs=Hsnf[:,li,:,:] if Hsnf is not None else None
tpe=probs.mean(0); aux=aux+E*torch.sum(tpe*probs.mean(0))

self.aux_loss=aux/max(1,(T+C–1)//C)
if return_state: return logits,(H,Hs)
return logits

@torch.no_grad()
def generate_step(self,token,h_prev):
H,Hs=h_prev; B=token.size(0)
xe=self.embed(token).squeeze(1); inp=self.input_proj(xe)
Hn,fu=self._estep(H,inp,self.W_forget,self.W_input,self.W_cand,
self.b_forget,self.b_input,self.b_cand,self.gate_W1,self.gate_b1,
self.gate_W2,self.gate_b2,self.fusion_W,self.fusion_b)
if self.n_shared>0:
Hsn,sf=self._estep(Hs,inp,self.W_forget_sh,self.W_input_sh,self.W_cand_sh,
self.b_forget_sh,self.b_input_sh,self.b_cand_sh,self.gate_W1_sh,self.gate_b1_sh,
self.gate_W2_sh,self.gate_b2_sh,self.fusion_W_sh,self.fusion_b_sh)
sf = sf.sum(dim=0)
else: Hsn,sf=None,0
probs=self._route(inp)
comb=((probs.t().unsqueeze(–1)*fu).sum(dim=0))+sf
return self.output_proj(self.output_norm(comb)),(Hn,Hsn)