OzTianlu commited on
Commit
9d0ba4a
Β·
verified Β·
1 Parent(s): c8de6b4

Update README.md

Browse files
Files changed (1) hide show
  1. README.md +0 -138
README.md CHANGED
@@ -31,144 +31,6 @@ A 1.3B parameter language model that replaces softmax attention with **causal mo
31
 
32
  ![Core Mechanism: The Monoid Recurrence](ARCH.png)
33
 
34
- ## Architecture Overview
35
-
36
- ```
37
- ╔═══════════════════════════════════════════════════════════════════════════╗
38
- β•‘ MonoidForCausalLM (1.34B) β•‘
39
- ╠═══════════════════════════════════════════════════════════════════════════╣
40
- β•‘ β•‘
41
- β•‘ token_ids ──> [ embed_tokens 128256 Γ— 2048 ] ──> x_0 β•‘
42
- β•‘ β•‘
43
- β•‘ β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β•‘
44
- β•‘ β”‚ MonoidDecoderLayer Γ— 16 β”‚ ◄── see detail below β•‘
45
- β•‘ β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β•‘
46
- β•‘ β”‚ β•‘
47
- β•‘ [ RMSNorm ] β•‘
48
- β•‘ β”‚ β•‘
49
- β•‘ [ lm_head 2048 Γ— 128256 ] ──> logits β•‘
50
- β•‘ (tied with embed_tokens) β•‘
51
- β•šβ•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•
52
-
53
-
54
- ╔═══════════════════════════════════════════════════════════════════════════╗
55
- β•‘ MonoidDecoderLayer (Γ— 16 layers) β•‘
56
- ╠═══════════════════════════════════════════════════════════════════════════╣
57
- β•‘ β•‘
58
- β•‘ x ─────────────────────────────────────────┐ (residual) β•‘
59
- β•‘ β”‚ β”‚ β•‘
60
- β•‘ [ input_layernorm RMSNorm ] β”‚ β•‘
61
- β•‘ β”‚ β”‚ β•‘
62
- β•‘ [ MonoidAttention ] ◄── see detail below β”‚ β•‘
63
- β•‘ β”‚ β”‚ β•‘
64
- β•‘ + <β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β•‘
65
- β•‘ β”‚ β•‘
66
- β•‘ x ─────────────────────────────────────────┐ (residual) β•‘
67
- β•‘ β”‚ β”‚ β•‘
68
- β•‘ [ post_attention_layernorm RMSNorm ] β”‚ β•‘
69
- β•‘ β”‚ β”‚ β•‘
70
- β•‘ [ LlamaMLP 2048 β†’ 8192 β†’ 2048 ] β”‚ β•‘
71
- β•‘ β”‚ gate_proj ─┐ β”‚ β•‘
72
- β•‘ β”‚ up_proj ─────> SiLU(gate) βŠ™ up β”‚ β•‘
73
- β•‘ β”‚ └──> down_proj ──> out β”‚ β•‘
74
- β•‘ β”‚ β”‚ β•‘
75
- β•‘ + <β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β•‘
76
- β•‘ β”‚ β•‘
77
- β•‘ out β•‘
78
- β•šβ•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•
79
-
80
-
81
- ╔═══════════════════════════════════════════════��═══════════════════════════╗
82
- β•‘ MonoidAttention (32 heads, d=64 per head) β•‘
83
- ╠═══════════════════════════════════════════════════════════════════════════╣
84
- β•‘ β•‘
85
- β•‘ x_t ∈ R^{2048} β•‘
86
- β•‘ β”‚ β•‘
87
- β•‘ β”œβ”€β”€> q_proj ──> [B,H,T,d] ──> RMSNorm ──> Γ—(1/√d) ──────> q_t β•‘
88
- β•‘ β”‚ β•‘
89
- β•‘ β”œβ”€β”€> k_proj ──> [B,H,T,d] ──> RMSNorm ──> SiLU ──────────> k_t β‰₯0 β•‘
90
- β•‘ β”‚ β•‘
91
- β•‘ β”œβ”€β”€> v_proj ──> [B,H,T,d] ────────────────────────────────> v_t β•‘
92
- β•‘ β”‚ β•‘
93
- β•‘ β”œβ”€β”€> decay_proj ──> Sigmoid ──> Ξ±_t ∈ (0,1)^d (vector decay gate) β•‘
94
- β•‘ β”‚ bias init = 3.0 β•‘
95
- β•‘ β”‚ β†’ Οƒ(3) β‰ˆ 0.95 at start β•‘
96
- β•‘ β”‚ β•‘
97
- β•‘ └──> gate_proj ──> SiLU ──────> g_t ∈ R^{H*d} (output gate) β•‘
98
- β•‘ β•‘
99
- β•‘ β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„ β•‘
100
- β•‘ Monoid Recurrence (training: parallel prefix scan, decode: O(1)) β•‘
101
- β•‘ β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„ β•‘
102
- β•‘ β•‘
103
- β•‘ k_t βŠ— v_t ──────────────┐ β•‘
104
- β•‘ [dΓ—d] v β•‘
105
- β•‘ β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β•‘
106
- β•‘ S_{t-1} ────> β”‚ S_t = diag(Ξ±_t)Β·S_{t-1}β”‚ β•‘
107
- β•‘ [dΓ—d] β”‚ + k_t βŠ— v_t │──> S_t β•‘
108
- β•‘ β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ [dΓ—d] β•‘
109
- β•‘ "compressed causal history" β•‘
110
- β•‘ β•‘
111
- β•‘ h0 (learnable, zero-init) ──> S_0 at sequence start β•‘
112
- β•‘ β•‘
113
- β•‘ β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„ β•‘
114
- β•‘ Readout + Output Projection β•‘
115
- β•‘ β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„β”„ β•‘
116
- β•‘ β•‘
117
- β•‘ q_t ──> einsum(q, S_t) ──> o_t ──> RMSNorm ──┐ β•‘
118
- β•‘ (o_norm) β”‚ β•‘
119
- β•‘ v β•‘
120
- β•‘ g_t ──────────────────────────────────> g_t βŠ™ o_t ──> o_proj ──> out β•‘
121
- β•‘ β•‘
122
- β•šβ•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•
123
-
124
-
125
- ╔═══════════════════════════════════════════════════════════════════════════╗
126
- β•‘ MonoidCache β€” O(1) State (replaces O(T) KV-Cache) β•‘
127
- ╠═══════════════════��═══════════════════════════════════════════════════════╣
128
- β•‘ β•‘
129
- β•‘ Transformer KV-Cache: Monoid State Cache: β•‘
130
- β•‘ β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β•‘
131
- β•‘ β”‚ K: [B,H,T,d] β”‚ β”‚ S: [B,H,d,d] β”‚ ← fixed size β•‘
132
- β•‘ β”‚ V: [B,H,T,d] β”‚ β”‚ Ξ±_acc: [B,H,d] β”‚ β•‘
133
- β•‘ β”‚ grows with T ↑↑↑ β”‚ β”‚ per layer β”‚ β•‘
134
- β•‘ β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β•‘
135
- β•‘ Memory: O(TΒ·HΒ·d) Memory: O(HΒ·dΒ²) β•‘
136
- β•‘ 1000 tok β†’ 2M floats/layer ANY length β†’ 131K floats/layer β•‘
137
- β•‘ β•‘
138
- β•‘ Decode step: Decode step: β•‘
139
- β•‘ o = softmax(qΒ·K^T)Β·V S_t = Ξ±_tΒ·S_{t-1} + k_tβŠ—v_t β•‘
140
- β•‘ scan T keys ↑ o_t = q_t Β· S_t β•‘
141
- β•‘ Time: O(TΒ·d) Time: O(dΒ²) ← constant! β•‘
142
- β•šβ•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•
143
-
144
-
145
- ╔═══════════════════════════════════════════════════════════════════════════╗
146
- β•‘ Weight Transfer from Llama-3.2-1B-Instruct β•‘
147
- ╠═══════════════════════════════════════════════════════════════════════════╣
148
- β•‘ β•‘
149
- β•‘ Reused directly (frozen-compatible): β•‘
150
- β•‘ β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β•‘
151
- β•‘ β”‚ embed_tokens 128256 Γ— 2048 β”‚ β•‘
152
- β•‘ β”‚ lm_head 2048 Γ— 128256 (tied) β”‚ β•‘
153
- β•‘ β”‚ LlamaMLP Γ— 16 gate/up/down_proj β”‚ β•‘
154
- β•‘ β”‚ LlamaRMSNorm Γ— 33 input/post_attn/final β”‚ β•‘
155
- β•‘ β”‚ q_proj Γ— 16 2048 β†’ 2048 β”‚ β•‘
156
- β•‘ β”‚ k_proj Γ— 16 2048 β†’ 2048 (tiled 8β†’32 heads from GQA) β”‚ β•‘
157
- β•‘ β”‚ v_proj Γ— 16 2048 β†’ 2048 (tiled 8β†’32 heads from GQA) β”‚ β•‘
158
- β•‘ β”‚ o_proj Γ— 16 2048 β†’ 2048 β”‚ β•‘
159
- β•‘ β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β•‘
160
- β•‘ β•‘
161
- β•‘ Novel (randomly initialized): β•‘
162
- β•‘ β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” β•‘
163
- β•‘ β”‚ decay_proj Γ— 16 2048 β†’ 2048 (bias=3.0) β”‚ β•‘
164
- β•‘ β”‚ gate_proj Γ— 16 2048 β†’ 2048 (std=0.01) β”‚ β•‘
165
- β•‘ β”‚ q_norm Γ— 16 RMSNorm(64) β”‚ β•‘
166
- β•‘ β”‚ k_norm Γ— 16 RMSNorm(64) β”‚ β•‘
167
- β•‘ β”‚ o_norm Γ— 16 RMSNorm(64) (weight=1) β”‚ β•‘
168
- β•‘ β”‚ h0 Γ— 16 [1,32,64,64] (zeros) β”‚ β•‘
169
- β•‘ β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ β•‘
170
- β•šβ•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•β•
171
- ```
172
 
173
  ## Key Properties
174
 
 
31
 
32
  ![Core Mechanism: The Monoid Recurrence](ARCH.png)
33
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
34
 
35
  ## Key Properties
36