1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
| import torch
import torch.nn as nn
class MoTLayer(nn.Module):
"""
Mixture-of-Transformers 层
论文核心: 每层都有 Reasoner 和 Generator 两个专家
通过双流注意力(Dual-Stream Attention)联合
与传统 DiT-on-VLM 的区别:
- 传统: DiT 只在最后 cross-attend VLM 的输出一次
- Cosmos 3: 每层都联合, 生成始终基于 Reasoner 的推理结果
"""
def __init__(self, dim=4096, num_heads=32):
super().__init__()
self.reasoner_attn = nn.MultiheadAttention(
dim, num_heads, batch_first=True
)
self.reasoner_ffn = nn.Sequential(
nn.Linear(dim, dim*4), nn.GELU(), nn.Linear(dim*4, dim)
)
self.reasoner_norm1 = nn.LayerNorm(dim)
self.reasoner_norm2 = nn.LayerNorm(dim)
self.generator_attn = nn.MultiheadAttention(
dim, num_heads, batch_first=True
)
self.generator_ffn = nn.Sequential(
nn.Linear(dim, dim*4), nn.GELU(), nn.Linear(dim*4, dim)
)
self.generator_norm1 = nn.LayerNorm(dim)
self.generator_norm2 = nn.LayerNorm(dim)
self.cross_attn = nn.MultiheadAttention(
dim, num_heads, batch_first=True
)
self.cross_norm = nn.LayerNorm(dim)
def forward(self, ar_tokens, dm_tokens):
"""
Args:
ar_tokens: (B, L_ar, D) 自回归区 token (文本+视觉理解)
dm_tokens: (B, L_dm, D) 扩散区 token (视频+音频+动作)
Attention Mask:
- AR queries: 因果注意力, 只看前面的 AR token
- DM queries: 全注意力, 看所有 AR + DM token
"""
ar_normed = self.reasoner_norm1(ar_tokens)
ar_attn_out, _ = self.reasoner_attn(
ar_normed, ar_normed, ar_normed,
is_causal=True
)
ar_tokens = ar_tokens + ar_attn_out
ar_tokens = ar_tokens + self.reasoner_ffn(
self.reasoner_norm2(ar_tokens)
)
dm_normed = self.generator_norm1(dm_tokens)
dm_attn_out, _ = self.generator_attn(
dm_normed, dm_normed, dm_normed
)
dm_tokens = dm_tokens + dm_attn_out
cross_normed = self.cross_norm(ar_tokens)
cross_out, _ = self.cross_attn(
dm_tokens,
cross_normed,
cross_normed
)
dm_tokens = dm_tokens + cross_out
dm_tokens = dm_tokens + self.generator_ffn(
self.generator_norm2(dm_tokens)
)
return ar_tokens, dm_tokens
if __name__ == "__main__":
layer = MoTLayer(dim=512, num_heads=8)
B, L_ar, L_dm, D = 2, 100, 50, 512
ar = torch.randn(B, L_ar, D)
dm = torch.randn(B, L_dm, D)
ar_out, dm_out = layer(ar, dm)
print(f"AR output: {ar_out.shape}")
print(f"DM output: {dm_out.shape}")
print(f"Params: {sum(p.numel() for p in layer.parameters()):,}")
|
