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 串行架构的先天上限。但它的价值不在训练速度,在:

1. 推理永远 O(1)
2. 长序列显存不爆
3. 总成本在推理部署场景下胜出
4. 架构完全可控

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

附录: FRSM V6 Dense MoE 完整代码

文件:frsm_v6_moe/frsm_v6a_dense_moe.py

""" FRSM V6a Dense MoE — 全部专家用堆叠 einsum（无 gather/chunk/检查点） """importmath,torch,torch.nnasnn,torch.nn.functionalasFclassFRSM_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//4self.embed=nn.Embedding(vocab_size,D);self.input_proj=nn.Linear(D,D)fornin['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))ifn_shared>0:fornin['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):foreinrange(p.size(0)):forsinrange(self.num_scales):nn.init.kaiming_uniform_(p[e,s],a=math.sqrt(5))forpnin['W_forget','W_input','W_cand','gate_W1','gate_W2']:_k(getattr(self,pn))foreinrange(self.n_experts):nn.init.kaiming_uniform_(self.fusion_W[e],a=math.sqrt(5))ifself.n_shared>0:forpnin['W_forget','W_input','W_cand','gate_W1','gate_W2']:_k(getattr(self,pn+'_sh'))foreinrange(self.n_shared):nn.init.kaiming_uniform_(getattr(self,'fusion_W_sh')[e],a=math.sqrt(5))forn,pinself.named_parameters():if'bias'inn: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)ifself.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)returnHn,fuseddef_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)ifself.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,0probs=self._route(inp)combined=((probs.t().unsqueeze(-1)*fused).sum(dim=0))+sfreturnHn,Hsn,combined,probsdef_route(self,inp):l=self.router(inp)ifself.trainingandself.router_noise>0:l=l+torch.randn_like(l)*self.router_noisereturnF.softmax(l,dim=-1)defforward(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)ifh_previsNone: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)ifself.n_shared>0elseNoneelse: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_sizeifself.chunk_size>0elsemax(1,int(math.sqrt(T)))fortsinrange(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)ifHsisnotNoneelseNoneHnf,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)ifself.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,0ifsfisnotNone: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,:,:]ifHsnfisnotNoneelseNonetpe=probs.mean(0);aux=aux+E*torch.sum(tpe*probs.mean(0))self.aux_loss=aux/max(1,(T+C-1)//C)ifreturn_state:returnlogits,(H,Hs)returnlogits@torch.no_grad()defgenerate_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)ifself.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,0probs=self._route(inp)comb=((probs.t().unsqueeze(-1)*fu).sum(dim=0))+sfreturnself.output_proj(self.output_norm(comb)),(Hn,Hsn)