Diagnostic for learned-projection attention heads

Extend the participation-ratio memorization diagnostic to attention heads with learned projections by identifying the analogue of the condensation variable x when the logits are no longer direct overlaps between queries and keys.

Background

The paper analyzes memorization in attention heads whose keys are noisy training examples and whose logits are direct overlaps between a query and those keys. In this setting, the within-class logit fluctuations are characterized by the variable x=β√{N(1-r²)}, and the participation ratio of the softmax weights serves as a diagnostic for condensation and memorization through a random-energy-model description.

A learned-projection attention head transforms queries and keys before forming logits, so its logits are not simply the raw overlaps treated in Lemma 'logit statistics'. The paper therefore leaves unresolved how to define the corresponding scale or variable controlling softmax condensation, and consequently how to transfer the participation-ratio diagnostic to trained attention architectures.