KRAFTON

/

Raon-1-mini-preview

Like 189

Follow

KRAFTON 340

Thanks for the question! RDA and GDN-2 are closely related: both decouple the erase and write strengths that KDA ties to a single gate ${\beta }_t\backslash beta\_t$. The difference is granularity.

• GDN-2 makes both gates channel-wise: an erase gate $b_t\in {\mathbb{R}}^{d_k}b\_t\; \backslash in\; \backslash mathbb\{R\}^\{d\_k\}$ on the key and a write gate $w_t\in {\mathbb{R}}^{d_v}w\_t\; \backslash in\; \backslash mathbb\{R\}^\{d\_v\}$ on the value.
• RDA keeps both gates as per-head scalars, $e_t,w_t\in \mathbb{R}e\_t,\; w\_t\; \backslash in\; \backslash mathbb\{R\}$, so erase and write are learned independently with almost no extra parameters. It starts exactly as KDA at initialization and learns to separate the two during training.

More precisely, the three state updates are:

$$\mathbf{\text{KDA:}}{S}_t=(I-{\beta }_t{k}_t{k}_t^{\mathrm{\top }})\mathrm{D}\mathrm{i}\mathrm{a}\mathrm{g}({\alpha }_t)S_{t-1}+{\beta }_t{k}_t{v}_t^{\mathrm{\top }}\backslash textbf\{KDA:\}\backslash quad\; S\_t\; =\; (I\; -\; \backslash beta\_t\backslash ,\; k\_t\; k\_t^\backslash top)\backslash ,\backslash mathrm\{Diag\}(\backslash alpha\_t)\backslash ,S\_\{t-1\}\; +\; \backslash beta\_t\backslash ,\; k\_t\; v\_t^\backslash top$$

$$\mathbf{\text{RDA:}}{S}_t=(I-e_t{k}_t{k}_t^{\mathrm{\top }})\mathrm{D}\mathrm{i}\mathrm{a}\mathrm{g}({\alpha }_t)S_{t-1}+w_t{k}_t{v}_t^{\mathrm{\top }}\backslash textbf\{RDA:\}\backslash quad\; S\_t\; =\; (I\; -\; e\_t\backslash ,\; k\_t\; k\_t^\backslash top)\backslash ,\backslash mathrm\{Diag\}(\backslash alpha\_t)\backslash ,S\_\{t-1\}\; +\; w\_t\backslash ,\; k\_t\; v\_t^\backslash top$$

$$\mathbf{\text{GDN-2:}}{S}_t=(I-k_t(b_t\odot k_t{)}^{\mathrm{\top }})\mathrm{D}\mathrm{i}\mathrm{a}\mathrm{g}({\alpha }_t)S_{t-1}+k_t(w_t\odot v_t{)}^{\mathrm{\top }}\backslash textbf\{GDN-2:\}\backslash quad\; S\_t\; =\; (I\; -\; k\_t\; (b\_t\; \backslash odot\; k\_t)^\backslash top)\backslash ,\backslash mathrm\{Diag\}(\backslash alpha\_t)\backslash ,S\_\{t-1\}\; +\; k\_t\; (w\_t\; \backslash odot\; v\_t)^\backslash top$$

All three share KDA's channel-wise decay ${\alpha }_t\backslash alpha\_t$. RDA reduces to KDA when $e_t=w_{t}e\_t\; =\; w\_t$.
In our 1.3B ablations under the GDN-2 paper setup (recurrent-only models), RDA showed better long-context retrieval (NIAH) than both KDA and GDN-2.

Is this right?

Thanks for the nice summary!
Yes, that's right. In terms of gate granularity, RDA sits between KDA and GDN-2: it decouples erase from write like GDN-2, but keeps per-head scalar gates like KDA.

Ah ok that makes sense. Thank you very much. Althoug i find the better performance in ablations quite surprising, considering, at least intuitively, that a more fine grained edit should allow the model to better control the information in its state. Will there be a detailed tech report (including ablations)?

Thanks! We were also surprised at first. As a simple ablation, we tried to reproduce the MQAR experiment from the Kimi Linear paper under a fair setting for RDA, KDA and GDN-2. The code for that experiment was not released, so we wrote our own codebase and swept the same five learning rates for all three models. In our runs at L=2048 with 32 KV pairs, RDA reached ~97% accuracy while GDN-2 reached ~25% and KDA stayed near 0% at every learning rate. This does not match the trend in the Kimi Linear paper, especially for KDA, possibly because the original setup is not fully specified.

Like GDN-2, we think decoupling erase and write is the key idea. But in our experiments the granularity of KDA was already enough.

Yes, we also plan to release a technical report.