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.

Identity Matrix as Attention Weights Visualization

To visualize attention weights across different queries and keys simultaneously using heatmaps, the weights are structured as a four-dimensional tensor. The specific shape of this tensor is defined as (number of rows for display, number of columns for display, number of queries, number of keys). This multidimensional format allows an array of different query-key interactions to be organized and plotted within a grid layout.

Claude

When attention weights are nonnegative and sum to $$1$$, large weights can be interpreted intuitively as a mechanism for the model to select the most relevant components from the available keys. While this provides a helpful conceptual understanding of how the model focuses on certain inputs, it is important to recognize that this is primarily an intuition rather than a strict mechanical rule.

Learn Before

Related

Learn After