How did you come up with this update rule?

[kaikai23]

Hi, thanks for sharing this impressive project.

I wonder how did you come up with the update rule? Specifically, why the average over the attention scores before softmax is suitable for being the beta_t? Is it derived theoretically or randomly tried out with a lot of experiments? May you can share some insights to take away? thanks!

[rover-xingyu]

Hi @kaikai23, thanks for your interest!

The update rule is explained in Section 3.3 of the paper. Taking the average over the observations reflects the alignment confidence between the query state and the key observations, and this idea is inspired by the findings from Easi3R.

I believe your question is really about why we take the average before applying softmax, rather than using the sum or some other operation. The answer lies in the mathematical similarity between the softmax and sigmoid functions:
softmax = exp(x_i) / Σ_j exp(x_j)
sigmoid = exp(x_1) / (exp(0) + exp(x_1))
As we use a sigmoid to modulate the learning rate, taking the average before softmax becomes a natural choice.

[chrisoffner3d]

I have a question along similar lines, except w.r.t. the update direction rather than the learning rate $\beta_t.$

Specifically, how exactly is the objective function $\min \left\lVert \mathbf{S}_{t-1} \mathbf{K}_{\mathbf{X}_t} - \mathbf{V}_{\mathbf{X_t}}\right\rVert^2$ motivated?

It is not clear to me why the memory/state $\mathbf{S}_{t-1}$ should be optimized so that its product with the key matrix $\mathbf{K}_{\mathbf{X}_t}$ of the current observation is as close as possible to the value matrix of the current observation $\mathbf{V}_{\mathbf{X}_t}.$

I suppose I'm missing an intuition for the semantics of the key and value matrix here, and why aiming for $\mathbf{S}_{t-1}$ to be such that it transforms $\mathbf{K}_{\mathbf{X}_t}$ to (more closely) match $\mathbf{V}_{\mathbf{X}_t}$ is the right thing to do here.

Unfortunately formulations like

encourage the network to associate the current observation with the state

and

optimizing the state $\mathbf{S}$ to accurately reconstruct the observation value $\mathbf{V}_{\mathbf{X}_t}$

from the paper didn't help me understand why this is what we want.

[rover-xingyu]

Hi @chrisoffner3d! Here’s how I understand it:

The observation key ($\mathbf{K}_{\mathbf{X}_t}$) works like an index into the previous state ($\mathbf S_{t-1}$). It selects exactly which slots/entries of the state should be inspected (and potentially overwritten).

The optimization objective therefore drives the model to store the current observation value ($\mathbf V_{\mathbf X_t}$) at the positions of state that specified by that key ($\mathbf{K}_{\mathbf{X}_t}$).

The optimization objective its to enhance the product ($\mathbf S_{t-1} \mathbf K_{\mathbf X_t}$) to reproduce ($\mathbf V_{\mathbf X_t}$) as closely as possible, so the state/memory really holds the information we have just observed.

To summarize concisely:
• The key indicates where in the state to write.
• The value contains the content to be written.
• The learning rate controls how strongly the new content is written into memory—how much of it is actually retained.

https://arxiv.org/pdf/2501.12352 gives a good discussion of this test-time regression problem.