question

[banxianerr]

Very good job. Congratulations! And thank you for sharing codes.
I have some questions, Which variables represent search and retrieval? Are qk_dim and nrules?

[sarthmit]

Hello @banxianerr
The variable nheads represents searches, and nrules represents the number of values. So for each search in nheads, a value from nrules is chosen through the secondary attention mechanism, which is parameterized by the qk_dim. Does that answer your question?

As an example, nheads=2 and nrules=3 means 2 searches, and 3 retrievals per search. Now if qk_dim = 32, that means the dimensionality of secondary query-key attention is 32, which is used to soft-select retrievals per search.