nx_fnet_mix.nx source

1// nx_fnet_mix.nx -- FNET-001: SUB-QUADRATIC Fourier token-mixer (FNet, Lee-Thorp/Ainslie/Eckstein/Ontanon 2021). 2// 3// The escape from the transformer VRAM/compute exponential is JUST MATH. Self-attention mixes tokens with an 4// all-pairs softmax(Q@K^T)@V -- O(n^2 * d) multiplies, the quadratic that explodes VRAM as context grows. FNet 5// replaces that learned attention sublayer with a parameter-FREE 2D Discrete Fourier Transform and keeps the 6// REAL part: y = Re( FFT_seq( FFT_hidden(x) ) ). Token mixing becomes O(n * d * log(n*d)) -- linearithmic, not 7// quadratic -- with ZERO attention parameters and ZERO n-by-n score matrix to hold in memory. 8// 9// This is the sovereign realization of that math: 10// * a GENERALIZED radix-2 FFT for ANY power-of-2 length (the shipped nx_fft caps at N=32 via a hardcoded 11// twiddle table; here twiddles are synthesized for arbitrary N by the sovereign CORDIC fx_cos/fx_sin), 12// * the FNet 2D-DFT real-part token mixer composed from it, 13// * an integer-MULTIPLY op counter threaded through so the gate can MEASURE (never assert) the sub-quadratic 14// crossover against scaled-dot-product attention. 15// 16// DRY: reuses nx_fft_log2 + nx_fft_bit_reverse (already length-generic) and the SAME Q14 butterfly arithmetic as 17// the shipped, gated nx_fft, so the generalized FFT is bit-comparable to it (proven in nx_fnet_mix_gate.nx). 18// 19// genealogy_id: cooley_tukey_1965_fft + lee_thorp_2021_fnet + vaswani_2017_attention(as the incumbent it exceeds) 20// lineage_id: fnet_sovereign_fourier_token_mixer_v1 21// license_tier: ORIGINAL 22// 23// nx_safety_envelope: 24// intended_use: "parameter-free sub-quadratic token mixing for sovereign sequence models (FNet block); 25// the no-VRAM-exponential alternative to self-attention" 26// sil_target: SIL2 27// evidence: [Q14_fixed_point_deterministic, FFT_equals_naive_DFT_gated, 28// twiddles_match_shipped_nx_fft, measured_subquadratic_crossover_vs_attention] 29// residual_risk: "N and hidden-dim d must be powers of two (radix-2 contract); Q14 saturates above |x|>=2." 30// verdict: GREEN (nx_fnet_mix_gate 2026-06-14: C1 FFT==naive-DFT, C2 twiddles match nx_fft within 1 Q14, 31// measured sub-quadratic crossover vs attention 6->22->75->264x at n=64..4096; see fnet_mix.log) 32import "nx_fft.nx" // REUSE nx_fft_log2 + nx_fft_bit_reverse + NX_FFT_Q -- DRY, no FFT reinvention 33import "fx.nx" // sovereign CORDIC fx_cos/fx_sin + FX_TWO_PI -- arbitrary-length twiddle synthesis 34import "nx_syscalls.nx" // sys_mmap (sys_write only in the smoke main) 35 36// ===== Q16.16 (CORDIC) -> Q14 (FFT) with round-to-nearest, away from zero ============================ 37// 16384 == 65536/4, so the conversion is a rounded divide-by-4. cos/sin land in [-65536,65536] (Q16.16 of 38// [-1,1]); the result lands in [-16384,16384] (Q14), matching nx_fft's NX_FFT_Q scale exactly. 39func fnet_q16_to_q14(v: i64) -> i64 { 40 if v >= 0 { return (v + 2) / 4 } 41 return 0 - ((0 - v + 2) / 4) 42} 43 44// ===== Twiddle synthesis for ANY power-of-2 length n ================================================= 45// tw_re[k] = round(Q14 * cos(2*pi*k/n)), tw_im[k] = round(-Q14 * sin(2*pi*k/n)), k = 0 .. n/2-1. 46// Angle 2*pi*k/n is formed in Q16.16 as FX_TWO_PI*k/n; fx_cos/fx_sin reduce it internally. 47func fnet_build_twiddles(n: i64, tw_re: *i64, tw_im: *i64) -> i64 { 48 let half: i64 = n / 2 49 var k: i64 = 0 50 while k < half { 51 let ang: i64 = FX_TWO_PI * k / n // 2*pi*k/n in Q16.16 radians 52 tw_re[k] = fnet_q16_to_q14(fx_cos(ang)) 53 tw_im[k] = fnet_q16_to_q14(0 - fx_sin(ang)) 54 k = k + 1 55 } 56 return 0 57} 58 59// ===== Generalized in-place radix-2 forward FFT ====================================================== 60// re/im: parallel i64 arrays of length n (power of 2). tw_re/tw_im: a length-(n/2) twiddle table built by 61// fnet_build_twiddles(n,...). opc accumulates the count of integer multiplies executed (4 per butterfly). 62// Identical butterfly math + index scheme to nx_fft_forward (stride = n/m2), generalized past N=32. 63func fnet_fft_fwd(re: *i64, im: *i64, n: i64, tw_re: *i64, tw_im: *i64, opc: *i64) -> i64 { 64 let log2n: i64 = nx_fft_log2(n) 65 66 // Bit-reverse permutation (reuse the shipped length-generic helper). 67 var i: i64 = 0 68 while i < n { 69 let j: i64 = nx_fft_bit_reverse(i, log2n) 70 if i < j { 71 let tr: i64 = re[i] 72 let ti: i64 = im[i] 73 re[i] = re[j] 74 im[i] = im[j] 75 re[j] = tr 76 im[j] = ti 77 } 78 i = i + 1 79 } 80 81 // Iterative Cooley-Tukey butterflies. 82 var m: i64 = 1 83 while m < n { 84 let m2: i64 = m * 2 85 let twstride: i64 = n / m2 86 var k: i64 = 0 87 while k < n { 88 var j2: i64 = 0 89 while j2 < m { 90 let tw_idx: i64 = j2 * twstride 91 let w_r: i64 = tw_re[tw_idx] 92 let w_i: i64 = tw_im[tw_idx] 93 let x_r: i64 = re[k + j2 + m] 94 let x_i: i64 = im[k + j2 + m] 95 let t_r: i64 = (w_r * x_r - w_i * x_i) / NX_FFT_Q // 2 multiplies 96 let t_i: i64 = (w_r * x_i + w_i * x_r) / NX_FFT_Q // 2 multiplies 97 *opc = *opc + 4 98 let u_r: i64 = re[k + j2] 99 let u_i: i64 = im[k + j2] 100 re[k + j2] = u_r + t_r 101 im[k + j2] = u_i + t_i 102 re[k + j2 + m] = u_r - t_r 103 im[k + j2 + m] = u_i - t_i 104 j2 = j2 + 1 105 } 106 k = k + m2 107 } 108 m = m2 109 } 110 return 0 111} 112 113// ===== FNet 2D-DFT real-part token mixer ============================================================= 114// x_re/x_im: an [n, d] row-major matrix (input is real: caller sets x_im = 0). In place: 115// 1. row pass -- FFT of length d along the hidden dimension, for each of the n tokens, 116// 2. column pass -- FFT of length n along the sequence dimension, for each of the d channels. 117// FNet keeps the REAL part (x_re); the caller treats x_im as scratch. opc accumulates executed multiplies. 118func fnet_mix(x_re: *i64, x_im: *i64, n: i64, d: i64, opc: *i64) -> i64 { 119 // --- row pass: hidden-dimension mixing (length d) --- 120 let twr_d: *i64 = (sys_mmap(d * 8)) as *i64 121 let twi_d: *i64 = (sys_mmap(d * 8)) as *i64 122 fnet_build_twiddles(d, twr_d, twi_d) 123 let rr: *i64 = (sys_mmap(d * 8)) as *i64 124 let ri: *i64 = (sys_mmap(d * 8)) as *i64 125 var y: i64 = 0 126 while y < n { 127 var c: i64 = 0 128 while c < d { 129 rr[c] = x_re[y * d + c] 130 ri[c] = x_im[y * d + c] 131 c = c + 1 132 } 133 fnet_fft_fwd(rr, ri, d, twr_d, twi_d, opc) 134 c = 0 135 while c < d { 136 x_re[y * d + c] = rr[c] 137 x_im[y * d + c] = ri[c] 138 c = c + 1 139 } 140 y = y + 1 141 } 142 143 // --- column pass: sequence-dimension mixing (length n) --- 144 let twr_n: *i64 = (sys_mmap(n * 8)) as *i64 145 let twi_n: *i64 = (sys_mmap(n * 8)) as *i64 146 fnet_build_twiddles(n, twr_n, twi_n) 147 let cr: *i64 = (sys_mmap(n * 8)) as *i64 148 let ci: *i64 = (sys_mmap(n * 8)) as *i64 149 var xc: i64 = 0 150 while xc < d { 151 var r2: i64 = 0 152 while r2 < n { 153 cr[r2] = x_re[r2 * d + xc] 154 ci[r2] = x_im[r2 * d + xc] 155 r2 = r2 + 1 156 } 157 fnet_fft_fwd(cr, ci, n, twr_n, twi_n, opc) 158 r2 = 0 159 while r2 < n { 160 x_re[r2 * d + xc] = cr[r2] 161 x_im[r2 * d + xc] = ci[r2] 162 r2 = r2 + 1 163 } 164 xc = xc + 1 165 } 166 // Real part lives in x_re. Done. 167 return 0 168} 169 170// ===== backward pass (VJP) -- the FNet mixer is SELF-ADJOINT ========================================== 171// The mixer is y = Re(F_n * X * F_d) with F_n, F_d the SYMMETRIC DFT matrices (F[a,b]=w^{ab}=w^{ba}, so 172// F = F^T). For a scalar loss L with upstream gradient g = dL/dy, the standard linear-layer backward is 173// dL/dx = A^T g where A is the real-linear forward operator. Because both DFT matrices are symmetric, the 174// real part of the (symmetric) Kronecker operator F_n (x) F_d is itself SYMMETRIC -> A^T = A. Therefore 175// 176// dL/dx = Re( F_n * (dL/dy) * F_d ) = the SAME mixer applied to the upstream gradient. 177// 178// Backprop through the token mixer costs exactly ONE forward mix, stores ZERO activations, and has ZERO 179// parameters -- the structural win over attention (whose backward needs the stored n-by-n scores). On entry 180// g_re holds dL/dy (real); g_im is scratch the caller has zeroed; on return g_re holds dL/dx. Proven by 181// nx_fnet_grad_gate.nx (operator-matrix symmetry + adjoint dot-product test + finite-difference gradcheck). 182func fnet_mix_backward(g_re: *i64, g_im: *i64, n: i64, d: i64, opc: *i64) -> i64 { 183 return fnet_mix(g_re, g_im, n, d, opc) 184} 185 186// ===== compile-only smoke ============================================================================ 187// Mixes a tiny deterministic 4x4 matrix so the module builds + runs standalone. Real validation lives in 188// nx_fnet_mix_gate.nx (FFT==DFT, twiddle-match vs nx_fft, measured sub-quadratic crossover vs attention). 189func main() -> i64 { 190 let n: i64 = 4 191 let d: i64 = 4 192 let xr: *i64 = (sys_mmap(n * d * 8)) as *i64 193 let xi: *i64 = (sys_mmap(n * d * 8)) as *i64 194 var i: i64 = 0 195 while i < n * d { 196 xr[i] = (i % 7) - 3 197 xi[i] = 0 198 i = i + 1 199 } 200 let opc: *i64 = (sys_mmap(8)) as *i64 201 *opc = 0 202 fnet_mix(xr, xi, n, d, opc) 203 return 0 204}