Could you provide the code for visualizing attention in Figure 2, or help us identify if there are any issues with our approach?

[yuhkalhic]

Thank you for your excellent work. We have a question regarding the attention visualization in Figure 2 of your paper.

We attempted to reproduce the visualization using the following approach:

• Taking the last transformer layer
• Summing across all attention heads

And an example of our method is below. However, our results differ significantly from yours. Even when I set my code to be the first head of the first transformer layer, I can't get such a high score and such a significant distribution pattern like yours.

Could you help us understand, if there might be any specific preprocess or normalization steps we're missing?

To help us better understand and reproduce your results, would it be possible to share the visualization code you used? This would be incredibly helpful for our research.

Thank you for your time and assistance.

[lonleyodd]

@yuhkalhic hi, here is my exp code ,and can generate fig like follow, maybe can help you`

prompt1='The following are multiple choice questions (with answers) about machine learning.\n\n Which of the following is a clustering algorithm in machine learning?,Expectation Maximization,CART,Gaussian Naïve Bayes,Apriori'
prompt2='The following are multiple choice questions (with answers) about machine learning \n\n Which of the following is a clustering algorithm in machine learning?,Expectation Maximization,CART,Gaussian Naïve Bayes,Apriori'
prompt3='The following are multiple choice questions。\n Which of the following is a clustering algorithm in machine learning... \n Expectation Maximization\n,CART.... \n,Gaussian Naïve Bayes\n,Apriori'
prompt4='\n\n The following are multiple choice questions (with answers) about machine learning Which of the following is a clustering algorithm in machine learning?,Expectation Maximization,CART,Gaussian Naïve Bayes,Apriori'
prompt5='The following are multiple choice questions (with answers) about machine learning Which of the following is a clustering algorithm in machine learning? Expectation Maximization,CART,Gaussian Naïve Bayes,Apriori'
prompt = prompt1
inputs=tokenizer(prompt,return_tensors='pt').input_ids

token_q = inputs.squeeze().tolist()
token_q = tokenizer.convert_ids_to_tokens(token_q)
# token_q = tokenizer.convert_ids_to_tokens(token_ids.squeeze().tolist())
if prefix:
    bos_token = tokenizer.bos_token
    prefix_token = tokenizer.tokenize('.\n') + [bos_token]
    token_k = prefix_token + token_q
else:
    token_k = token_q

plt.figure(figsize=(40, 30))
# mask = np.triu(np.ones_like(data, dtype=bool),k=1)
# mask[0:3,0:3] = False
mask = (causal_mask<0).cpu().numpy()
data[mask] = np.nan

sns.heatmap(data, xticklabels=token_k, yticklabels=token_q,cmap='RdBu_r',mask=mask)

plt.xticks(fontsize=30)
plt.yticks(fontsize=30)
plt.title('Attention Scores with Masked Upper Triangle',fontsize=30)
plt.xlabel('Key',fontsize=30)
plt.ylabel('Query',fontsize=30)

save_path="./exp/exp_softmax/llama2_7B_after_softmax_with_prompt1_apply_prefix_visualization_v2"
os.makedirs(save_path,exist_ok=True)
plt.savefig(os.path.join(save_path,file_name+".png"))
plt.clf()`

[yuhkalhic]

Thank you so much