As a fundamental sanity check for attention visualizations, an identity matrix can be used to represent the attention weights. In this scenario, the attention weight is exactly $1$ only when the query and the key correspond to the same index, representing a perfect one-to-one focus. This idealized case serves to verify that the visualization mechanism correctly maps the query-key pairs to the expected heatmap layout.