how to show the heatmap on Figure 10 in the paper?

[DavideHe]

I see the heatmap in this paper Figure 10 and RMT Figure 6
I have some question:

1. many LLMs are casual. attention heatmap is bidirectional attention in RMT's Figure 6. But bidirectional attention generally is symmetric matrix . In the RMT's Figure 6, it show the casual feature.

2. what information you use to show that heatmap ? how to use these information . If these are some code or implementation principle，thank you very much for providing these .

[FYYFU]

Hi,
I'm also interest in how to obtain the Figure 9 of the paper. How can we change the insert position? (Like the original Needle-in-a-haystack).

[booydar]

@FYYFU The main dataset is already pre-generated, so the fact position is fixed. Hovewer, you can generate the dataset yourself, e.g. using the QA1 example notebook. To constrain the relative position of facts in the sample, you can use the task_start_pct and task_end_pct arguments of the NoiseInjectionDataset. Please make sure that both these values are between 0 and 1, and the sample_size is large enough to fit all the facts in this range. Don't forget to set noise_dataset to PG-19 to match the original BABILong.

[FYYFU]

@FYYFU The main dataset is already pre-generated, so the fact position is fixed. Hovewer, you can generate the dataset yourself, e.g. using the QA1 example notebook. To constrain the relative position of facts in the sample, you can use the task_start_pct and task_end_pct arguments of the NoiseInjectionDataset. Please make sure that both these values are between 0 and 1, and the sample_size is large enough to fit all the facts in this range. Don't forget to set noise_dataset to PG-19 to match the original BABILong.

Thank you very much! I'll have a try to use this code.

[booydar]

@DavideHe Thanks for your questions.

For Figure 6 of the Scaling Transformer to 1M paper, we used RMT based on encoder-only BERT, which does not have a causal attention mask. The assymetry is caused by writing or reading operations of important tokens to memory.

Regarding Figure 10 of the In Search of Needles paper, plots (c) and (d) were obtained using selected regions (crops) of attention maps of RMT based on a decoder-only GPT-2. Plots (a) and (b) are different, they represent pairwise distances between memory states of different segments. We collect the memory tokens after processing all segments, and then calculate the cosine distance between each two of them.

Hope this makes these figures more clear!