SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Evaluation points Our results as presented are for specific classes of evaluation points. A natural question to ask is if our results can be extended to more general set of evaluation points. It turns out that our results for  $p$ -variate Monarch matrices can be extended to a wider class of evaluation points. Specifically, for each  $0 \leq a < p$ , let  $S_{a} \subset \mathbb{C}$  with  $|S_{a}| = \sqrt[p]{N}$ . Then our results in this sub-section hold if we replace the evaluation points from  $A^{p}$  to  $\times_{a=0}^{p-1} S_{a}$ . The only thing that changes in our proofs is that in Theorem 16, we replace  $\mod \left(T_{\overline{p}\sqrt{N}}(X_{0}), \ldots, T_{\overline{p}\sqrt{N}}(X_{p-1})\right)$  by  $\mod \left(q_{S_{0}}(X_{0}), \ldots, q_{S_{p-1}}(X_{p-1})\right)$ , where  $q_{A}(Z)$  is as defined in Equation (81). This result can then be propagated throughout the rest of our proofs.
On the other hand, our results in Appendix D.3 and Appendix D.5 do exploit specific properties of the evaluation points (specifically  $\left(\omega_N^i\right)^{\sqrt{N}} = \omega_{\sqrt{N}}^{i_0}$  for Appendix D.3 and  $T_{\sqrt{N}}(\omega_{N,i}) = (-1)^{i_1}\omega_{\sqrt{N},i_0}$  for Appendix D.5). To generalize these results to other sets of evaluation points, we need the existence of degree  $\sqrt{N}$  polynomial that maps (in a  $\sqrt{N}$ -to-1 fashion)  $A$  to a set of  $\sqrt{N}$  elements. Another interesting open question is to avoid the blowup  $n\rightarrow 2^p\cdot n$  in Theorem 17 and ideally only pay a blowup  $n\rightarrow 2n$  for every  $p\geq 2$  as we were able to do in Appendix D.3 and Appendix D.5 (with  $p = 2$ ).
1 from einops import rearrange   
2 import torch   
3 from torch import nn   
4   
5 def blockdiag_matmul(x, w): return torch.einsum( "bnm,...bm->...bn",w,x.view(*x.shape[-1],w.shape[0],w. shape[-1]) ).reshape(*x.shape)   
9   
10 class MonarchMatrix(nnModule):   
11   
12 def __init__(self,sqrt_n:int): super().__init_(） self_sqrt_n  $=$  sqrt_n self.L  $\equiv$  nn.Parameters(torch.randint((sqrt_n，sqrt_n，sqrt_n))) self.R  $\equiv$  nn.Parameters(torch.randint((sqrt_n，sqrt_n，sqrt_n)))   
17   
18 def forward(self,x): x  $=$  rearrange(x，"..（mn）->...（n m)，n=self_sqrt_n) x  $=$  blockdiag_matmul(x,self.L) x  $=$  rearrange(x，"..（mn）->...（n m)，n=self_sqrt_n) x  $=$  blockdiag_matmul(x,self.R) return rearrange(x，"..（mn）->...（n m)，n=self_sqrt_n)   
24   
25 class MonarchMixerLayer(nnModule):   
26 def __init__(self,sqrt_n:int,sqrt_d:int): super().__init_(） self.m1  $=$  MonarchMatrix(square_n) self.m2  $=$  MonarchMatrix(square_n) self.m3  $=$  MonarchMatrix(square_d) self.m4  $=$  MonarchMatrix(square_d)   
32   
33 self.n_kernel = nn_PARAMETER(torch.randint(square_d**2,sqrt_n \*\*2)) self.d_kernel = nn_PARAMETER(torch.randint(1,sqrt_d**2)) self(layer_norm  $\equiv$  nn.layersNorm(square_d\*\*2)   
34   
35 def forward(self,x:torch.Tensor): #x.shape  $\equiv$  (b,n,d) x_tilde  $\equiv$  self.m2(self.n_kernel\*self.m1(x.transpose(-1,-2)) ).transpose(-1,-2)# mix sequence y  $=$  self.m4(torchrelu(self.d_kernel\*self.m3(x_tilde))) # mix features   
40 return self(layer_norm(y+x_tilde)# skip connection Listing 1:A basic implementation of the M2 layer.