How were the precomputed tensors (.pth) generated and what is the implementation of lookup_func in Mamba-2?

[xuboshen]

Hello, this is a very impressive work! I’m currently trying to conduct visualization analysis on models such as Mamba-2, and I’d like to build everything from scratch. I noticed entries like

"global_channel_indexes": "./remapping_configs/precomputed_tensors/mamba2-1.3b_global_channel_indexes.pth",
"per_channel_decays": "./remapping_configs/precomputed_tensors/mamba2-1.3b_per_channel_decays.pth"

Could you please share how these .pth files were computed? I couldn’t find the corresponding scripts—would it be possible to open-source them?

Additionally, I’m very curious about the lookup_func functions. I couldn’t locate their implementations either, and I see that in the mamba2.py
file, the default values for several related parameters are set to None.

Thank you very much!

[xuboshen]

In my understanding, the per_channel_decays tensor corresponds to the precomputed constant factor defined in Equation (6), i.e., $\Pi_{k=1}^{L_0} \overline{A}_k$ for each layer and each head.

After inspecting the global_channel_indexes tensors, I found it contains boolean values. Could you please clarify what these represent? It appears to randomly select from the 64 heads to calculate attentions, but according to the paper, the term “across all global channels” suggests that all channels should be selected.

Please correct me if my understanding is mistaken.