Attention maps reveal stimulus-dependent retinal population codes

Understanding how deep learning models map neural population activity to stimuli requires both high predictive accuracy and interpretable internal mechanisms.

In this work, we employ the POYO framework, a scalable transformer architecture based on spike tokenization and latent modeling, to decode large-scale retinal ganglion cell recordings. We ask whether the model's attention mechanisms can provide biologically meaningful insight by evaluating two contrasting conditions: a uniform flash stimulus and a spatiotemporally structured moving ball stimulus.

We show that the model decodes both stimuli reliably and adapts rapidly to new preparations via fine-tuning, suggesting the capture of transferable population codes. We then analyze the model's internal organization, revealing that encoder attention patterns adapt to stimulus complexity: attention heads appear synchronized and broadly distributed for the flash stimulus, whereas they exhibit heterogeneous, specialized allocation strategies for the moving ball. By aggregating attention weights to identify the most relevant neurons for each task, we demonstrate that these high-attention units possess distinct physiological signatures-concentrating sustained, high-firing rates responses for the flash vs. diverse kinetics for the structured input. We confirm the causal validity of these findings via attention-guided ablations, where the progressive removal of these top-ranked units yields systematic losses in decoding performance. Furthermore, we expand the analysis to the decoder's attention, uncovering stimulus-specific retrieval strategies where individual heads exhibit distinct directional tuning preferences.

We conclude that generic attention mechanisms can spontaneously recover biological coding strategies, identifying functionally distinct neural subpopulations without supervision, thus validating the utility of transformer-based architectures for neuroscientific discovery.