Reproducing Figure 1 in the paper

[freesky01]

Thank you for the great work, quite interesting view!

I'm trying to use DCT to analyze the power spectrum of input hidden states of currently popular LLMs, like you have shown in Figure 1 in the paper. But I'm not sure about the details to plot it (e.g. w/o normalization and how are amplitudes averaged over the entire validation set, sample-level or token level?). Most of the amplitudes of tokens by my side is about zero, which does not follow the pattern of Figure 1, so I think I am very likely to mess up some settings.

Could you provide the snippet of plotting Figure 1 or give me some further hints about it? I would really appreciate it if you would help me.