Brain-Like Language Processing via a Shallow Untrained Multihead Attention Network (2406.15109v1)

Abstract: LLMs have been shown to be effective models of the human language system, with some models predicting most explainable variance of brain activity in current datasets. Even in untrained models, the representations induced by architectural priors can exhibit reasonable alignment to brain data. In this work, we investigate the key architectural components driving the surprising alignment of untrained models. To estimate LLM-to-brain similarity, we first select language-selective units within an LLM, similar to how neuroscientists identify the language network in the human brain. We then benchmark the brain alignment of these LLM units across five different brain recording datasets. By isolating critical components of the Transformer architecture, we identify tokenization strategy and multihead attention as the two major components driving brain alignment. A simple form of recurrence further improves alignment. We further demonstrate this quantitative brain alignment of our model by reproducing landmark studies in the language neuroscience field, showing that localized model units -- just like language voxels measured empirically in the human brain -- discriminate more reliably between lexical than syntactic differences, and exhibit similar response profiles under the same experimental conditions. Finally, we demonstrate the utility of our model's representations for LLMing, achieving improved sample and parameter efficiency over comparable architectures. Our model's estimates of surprisal sets a new state-of-the-art in the behavioral alignment to human reading times. Taken together, we propose a highly brain- and behaviorally-aligned model that conceptualizes the human language system as an untrained shallow feature encoder, with structural priors, combined with a trained decoder to achieve efficient and performant language processing.

Brain-Like Language Processing via a Shallow Untrained Multihead Attention Network

The paper "Brain-Like Language Processing via a Shallow Untrained Multihead Attention Network" offers a rigorously detailed investigation into the architectural components that enable high alignment between LLMs and human brain activity during language processing. The authors present a streamlined model that achieves significant brain alignment and offers promising implications for both cognitive neuroscience and artificial intelligence.

Architectural Components and Brain Alignment

LLMs, even when untrained, exhibit internal representations that align with human brain data. This paper investigates the specific architectural components responsible for this surprising alignment. The authors employ a methodology akin to neuroscientific approaches, such as functional localization, to identify language-selective units within LLMs. The paper draws insights from transformer architectures, focusing particularly on tokenization strategies and multihead attention mechanisms.

Key Findings

1. Token Aggregation: A significant component driving the alignment is the model's method of aggregating tokens. The paper identifies that utilizing a Byte Pair Encoding (BPE) tokenizer and aggregating tokens through multihead attention mechanisms critically enhances the model's alignment with brain data. Interestingly, even simple mean pooling of tokens contributes to a high degree of alignment.
2. Attention Mechanisms: Increasing the diversity of token aggregation through multihead attention was found to improve brain alignment further. This is due to the diverse context-dependent associations that multihead attention mechanisms can encode.
3. Recurrent Processing: The paper highlights the benefits of recurrently applying shared weights to enhance brain alignment. This form of recurrence, analogous to repeated processing in neural circuits, yielded substantial improvements.

Proposed Model: \ourmodel

The culmination of the paper's insights is a model named \ourmodel, composed of a shallow, untrained multihead attention mechanism. This simplified model, with structural priors, captures most of the variance in current brain recording benchmarks and achieves competitive alignment scores efficiently.

Efficacy in Brain-alignment Benchmarks

Across five diverse brain recording datasets, the model exhibits robust alignment. Quantitative evaluations demonstrate that \ourmodel can explain a significant portion of the variance in brain activity, which is notable given its untrained status. The use of a BPE tokenizer and token aggregation proved crucial in replicating the nuanced response profiles observed in human neuroimaging studies.

Replicating Landmark Neuroscience Studies

The paper undertakes a thorough validation by replicating landmark studies in language neuroscience. Findings show that localized units in \ourmodel, akin to language voxels in the brain, are more sensitive to lexical than syntactic differences. These results highlight the model's ability to capture essential properties of the human language system.

Implications for LLMing

The model's utility extends beyond brain alignment. By integrating a trainable decoder module with the untrained \ourmodel, the paper demonstrates enhanced sample efficiency and LLMing performance. This combined architecture achieves state-of-the-art behavioral alignment in predicting human reading times, underscoring the potential practical applications of the model in language technology.

Discussion and Future Directions

The findings prompt a reconsideration of the structural simplicity underlying the human language system. The model suggests that effective language representations can be derived from simple, untrained architectures with structural priors. This supports a conceptual framework where the human language system functions as a hierarchical encoder feeding into a downstream decoder.

However, the paper acknowledges the need for improved brain benchmarks with higher signal-to-noise ratios and consistency across different metrics and datasets. Future work should address these limitations by fostering more refined datasets and evaluating models under diverse linguistic conditions.

Conclusion

This paper establishes that a shallow, untrained multihead attention network can achieve significant alignment with human brain activity during language processing. This work forwards our understanding of both machine learning and cognitive neuroscience, proposing a simpler yet effective framework for modeling the human language system. Future research will benefit from these insights by exploring more human-like processing models and improving cross-disciplinary evaluation methods.