Selective-Update RNNs (`suRNN`/`suRNNCell`)

[MartinuzziFrancesco]

Paper

https://arxiv.org/abs/2603.02226

Equations

$$\begin{align} \mathbf{D}_t &\triangleq \mathrm{diag}(\mathbf{g}_t), \qquad \mathbf{g}_t \in \{0,1\}^H, \qquad \mathbf{h}_t \in \mathbb{R}^H \\[2pt] \Delta \mathbf{h}_t &\triangleq f_\theta(\mathbf{h}_{t-1}, \mathbf{x}_t) - \mathbf{h}_{t-1} \\[2pt] \mathbf{h}_t &= \mathbf{h}_{t-1} + \mathbf{D}_t \, \Delta \mathbf{h}_t \\[-2pt] &= \mathbf{h}_{t-1} + \mathbf{D}_t \big(f_\theta(\mathbf{h}_{t-1}, \mathbf{x}_t) - \mathbf{h}_{t-1}\big) \\[-2pt] &= (\mathbf{I}-\mathbf{D}_t)\mathbf{h}_{t-1} + \mathbf{D}_t f_\theta(\mathbf{h}_{t-1}, \mathbf{x}_t) \\[6pt] \mathbf{J}_t &\triangleq \frac{\partial \mathbf{h}_t}{\partial \mathbf{h}_{t-1}} = \mathbf{I} + \mathbf{D}_t\big(\mathbf{J}^{(f)}_t - \mathbf{I}\big), \qquad \mathbf{J}^{(f)}_t \triangleq \left.\frac{\partial f_\theta}{\partial \mathbf{h}_{t-1}}\right|_{(\mathbf{h}_{t-1},\mathbf{x}_t)} \\[6pt] a_{t,i} &= b_i + \sum_{k=1}^{K} \alpha_{i,k}\,\sin(\omega_k t + \phi_{i,k}), \qquad g_{t,i} = H(a_{t,i}) \quad (i=1,\dots,H) \end{align}$$

Official implementation

This is linked in the paper but no code is there yet
https://anonymous.4open.science/r/suGRU-EB5C/README.md

[P-r-e-m-i-u-m]

"Hi @MartinuzziFrancesco, I'd like to work on this! I'll implement suRNNCell and suRNN following the same pattern as the MiRU implementation."

[MartinuzziFrancesco]

sounds good to me!