Fine-grained Artificial Neurons in Audio-transformers for Disentangling Neural Auditory Encoding (2023)

Abstract: The Wav2Vec and its variants have achieved unprecedented success in computational auditory and speech processing. Meanwhile, neural encoding studies that integrate the superb representation capability of Wav2Vec and link those representations to brain activities have provided novel insights into a fundamental question of how auditory and speech processing unfold in the human brain. Without an explicit definition, most existing studies treat each transformer encoding layer in Wav2Vec as a single artificial neuron (AN). That is, the layer-level embeddings are used to predict neural responses. However, the comprehensive layer-level embedding aggregates multiple types of contextual attention captured by multi-head self-attention (MSA) modules. Thus, the layer-level ANs lack fine-granularity for neural encoding. To address this limitation, we define the elementary units, i.e., each hidden dimension, as neuron-level ANs in Wav2Vec2.0, quantify their temporal responses, and couple those ANs with their biological-neuron (BN) counterparts in the human brain. Our experimental results demonstrated that: 1) The proposed neuron-level ANs carry meaningful neurolinguistic information; 2) Those ANs anchor to their BN signatures; 3) The AN-BN anchoring patterns are interpretable from a neurolinguistic perspective. More importantly, our results suggest an intermediate stage in both the computational representation in Wav2Vec2.0 and the cortical representation in the brain. Our study validates the fine-grained ANs in Wav2Vec2.0, which may serve as a novel and general strategy to link transformer-based deep learning models to neural responses for probing the sensory processing in the brain.

## introduction the wav2vec model and its variants @xcite have achieved superb performance in learning acoustic information representations and on a variety of downstream * the corresponding author tasks such as automatic speech recognition. meanwhile, recent studies that link the computational representations in wav2vec to neural responses recorded by functional brain imaging techniques have provided novel insights into the model's interpretability and neural sensory perception of acoustic information @xcite @xcite . such studies can be formulated as a general framework of brain encoding and decoding @xcite @xcite . in brief, a predictive model is trained to build a mapping between the computational feature representation (the feature space, referred to artificial neurons, ans) of the input stimuli and the brain activities (the brain activity space, referred to biological neurons, bns) evoked by the same set of stimuli. the fitness of the predictive model, also known as the "brain score", is used to infer the correspondence between specific features and the underlying brain regions. in most existing studies that link audiotransformers to brain responses, the layer-level contextual embeddings in the transformer encoding layers are used as the feature space @xcite @xcite . the layer-level representations aggregate multiple types of attentional relationships among the input sequences captured by multi-head self-attention (msa) modules @xcite . the aggregation operation results in comprehensive representations. however, these representations lack specificity. thus, treating each encoding layer as a single an is relatively coarse and consequently degenerates the capability of audio-transformers in brain encoding and decoding studies. multi-level visualizations of transformer attentions @xcite @xcite may provide some inspirations to address this problem. for example, bertviz visualizes the attention at the model-level, head-level and neuronlevel @xcite . more specifically, the neuronlevel visualization factorizes the attention score matrix in each head into a set of element-wise product matrices corresponding to the hidden dimensions. the neuron-level visualization enables computational interpretation of transformers with fine granularity. however, whether each hidden dimension can be defined as a fine-grained an for neural encoding and decoding study is not clear. do those ans carry meaningful linguistic information? do those ans anchor to their bn signatures in the human brain? are the coupled an-bn pairs interpretable from a neurolinguistic perspective? we sought to answer these questions in this study. to this end, we propose a general framework for coupling the fine-grained ans in wav2vec2.0 @xcite and the bns in the human brain. we adopt the pre-trained wav2vec2.0 to embed the spoken story stimuli in the narratives functional magnetic resonance imaging (fmri) dataset @xcite . the temporal response of an an is then quantified according to the element-wise product of the queries and keys. functional brain networks (fbns) are identified from the fmri data and each fbn is regarded as a single bn. afterwards, the coupling relationship between ans and bns are built by maximizing the synchronizations between their temporal responses. our experimental results show that those finegrained ans carry meaningful linguistic information and well synchronize to their bn signatures, and the anchored an-bn pairs are interpretable. more importantly, our results suggest an intermediate stage in both the computational representation in wav2vec2.0 and the cortical representation in the brain. the proposed fine-grained ans may also serve as a general strategy to link transformerbased deep learning models to neural responses for probing the sensory processing in the brain. ## related works features from computational models have long been used to model the feature space for exploring the auditory neural encoding. conventional hand-crafted features that capture low-level acoustic properties (e.g., sound intensity, timbre, rhythm, pitch, and spectrograms) have been found to be closely correlated to brain responses @xcite @xcite @xcite @xcite @xcite @xcite . some studies replicate similar findings for the combinations of those low-level features optimized for specific tasks such as auditory attention @xcite and melodic pitch expectations @xcite . the deep neural networks (dnns) developed for auditory and speech processing bring new opportunities to model the feature space. the model architecture and the training objective are two basic ingredients of dnns. existing studies have investigated the similarity between brain responses and dnns in different architectures including convolutional neural network (cnn) @xcite @xcite @xcite , convolutional auto-encoder (cae) @xcite , generative adversarial network (gan) @xcite , cnn followed by recurrent neural network (rnn) @xcite @xcite , spiking neural networks @xcite , and transformers @xcite @xcite . the training objectives include unsupervised, self-supervised and supervised by various tasks such as musical genre prediction, acoustic scene classification, and speech recognition. these studies have provided fruitful insights into neural auditory encoding, model interpretation, and brain-like model development. for example, correlating the hierarchical representations derived from cnn-based models for automatic music tagging have revealed the representational gradients in the superior temporal gyrus (stg). the anterior stg (astg) and posterior stg (pstg) have been shown to be more sensitive to low-level and highlevel features encoded in shallow and deep layers, respectively @xcite . by optimizing a cnn-based model for dual-task of word and music genre classification, kell et al. showed that the bestperforming network may resemble the hierarchical organization of the human auditory cortex. that is, brain responses in the primary and non-primary auditory cortices are most well predicted by middle and late cnn layers, respectively @xcite . by modeling the feature space via cnn-rnn-based deepspeech2 @xcite op-timized for acoustic scene classification and speechto-text with different types of inputs (i.e., english, dutch and bengali), millet et al. replicated such a hierarchy and suggested that the brain utilizes sound-generic representations in the first processing stage of its hierarchy, and then builds speechspecific representations in higher-level processing stages @xcite . more recently, the transformer based on multihead self-attention (msa) has emerged as a powerful dnn architecture to learn comprehensive contextual representations @xcite . in this context, audio-transformers such as wav2vec 2.0 have also been used to model the feature space @xcite @xcite . for example, millet et al. compared wav2vec 2.0 to neural activities in a large cohort, and found that the representational hierarchy of wav2vec 2.0 aligns with the cortical hierarchy of speech processing. more specifically, wav2vec2.0 learns sound-generic, speech-specific and language-specific representations that are analogous to those of the temporal and prefrontal cortices @xcite . li et al. compared the representational similarity of hubert @xcite , wav2vec 2.0 @xcite and deepspeech2 @xcite with different training objectives to the human auditory pathway. they showed that the representational hierarchy in the dnns correlates well to the ascending auditory pathway, and unsupervised models achieve optimal neural correlations @xcite . tuckute et al. examined brain-dnn similarities within the auditory cortex for a large set of models based on various architectures and trained on different tasks. they found that most dnns predicted brain responses in the auditory cortex better than the filterbank models as a baseline, and the models trained on multiple tasks produced the best overall predictions. more importantly, they showed that most of the dnns exhibited a correspondence between model stages and brain regions, for example, the neural responses in lateral, anterior and posterior non-primary auditory cortices were better predicted by deeper layers @xcite . despite those fruitful findings, the feature space defined in existing studies that assess the representational similarity between wav2vec2.0 and brain responses relies on layer-level embeddings. that is, these studies implicitly treat each layer as a single artificial neuron. considering the heterogeneity of the attentional heads, this operation may lose the specificity of each head, which is designed to capture different types of contextual attention. as argued in the field of natural language processing (nlp), a fine decomposition of a model's components into elementary units is among the keys for mapping computational models to their neurobiological counterparts @xcite . this demand also applies to audiotransformers. meanwhile, our previous study has shown the validity of fine-grained ans defined as the hidden dimensions of the pre-trained bert model @xcite . however, whether those fine-grained ans hold similar premises in audiotransformers is unknown. thus, the key objective of this work is to validate those fine-grained ans in wav2vec 2.0 for neural encoding studies. ## discussion and conclusion we proposed to define fine-grained artificial neurons (ans) in the audio-transformer wav2vec2.0 and map them to their neurobiological counterparts. our experimental results showed that the fine-grained ans carried meaningful linguistic information and well synchronized to their bn signatures. moreover, the anchored an-bn pairs are partly interpretable in a neurolinguistic view. although a comprehensive mapping of the cortical speech hierarchy is out of the scope of this study, we observed some interesting results, facilitated by the fine-grained ans. first, the alignment between the computational hierarchy in wav2vec2.0 and the cortical speech hierarchy is largely in line with existing studies @xcite @xcite . second, and more importantly, we preliminarily discovered an intermediate stage in both the computational representation in wav2vec2.0 and the cortical representation in the brain. it gradually integrates global phonetic localization encoded in the early stages of neural speech hierarchy through the functional interactions between the working memory and language networks to form local language localization. in comparison, a good predictive performance from computational representation in audio-transformers to brain activities has rarely been reported previously. third, we observed phoneme-selective neural-level ans in wav2vec2.0, and the associated bns are partly in line with existing studies @xcite . thus, the fine-grained ans defined here may potentially provide an alternative approach to explore whether there are phonemeselective neural activities in the brain. the fine-grained ans defined in this study may also serve as the brain-based test-bed to evaluate and interpret audio-transformers, and provide neurolinguistic support for better understanding of the role of self-attention for efficient speech computation. for example, after interpreting distinctive attentional profiles in different attention heads, shim et al. applied a layer-wise attention map reuse strategy to improve model performance @xcite . a similar but with more fine-grained strategy could further improve model performance. in conclusion, we defined and validated neuronlevel ans in wav2vec2.0. this definition may serve as a general strategy to link transformer-based deep learning models to neural responses for probing the sensory processing in the brain.