Learning positional encodings in transformers depends on initialization
Keywords: transformers, representation learning, deep learning, reasoning, transformer
TL;DR: Learning accurate and interpretable positional encodings that improve generalization depends on their initialization
Abstract: The attention mechanism is central to the transformer's ability to capture complex dependencies between tokens of an input sequence.
Key to the successful application of the attention mechanism in transformers is its choice of positional encoding (PE).
The PE provides essential information that distinguishes the position and order amongst tokens in a sequence.
Most prior investigations of PE effects on generalization were tailored to 1D input sequences, such as those presented in natural language, where adjacent tokens (e.g., words) are highly related.
In contrast, many real world tasks involve datasets with highly non-trivial positional arrangements, such as datasets organized in multiple spatial dimensions, or datasets for which ground truth positions are not known, such as in biological data.
Here we study the importance of learning accurate PE for problems which rely on a non-trivial arrangement of input tokens.
Critically, we find that the choice of initialization of a learnable PE greatly influences its ability to learn accurate PEs that lead to enhanced generalization.
We empirically demonstrate our findings in a 2D relational reasoning task and a real world 3D neuroscience dataset, applying interpretability analyses to verify the learning of accurate PEs.
Overall, we find that a learned PE initialized from a small-norm distribution can 1) uncover interpretable PEs that mirror ground truth positions, 2) learn non-trivial and modular PEs in a real-world neuroscience dataset, and 3) lead to improved downstream generalization in both datasets.
Importantly, choosing an ill-suited PE can be detrimental to both model interpretability and generalization.
Together, our results illustrate the feasibility of learning identifiable and interpretable PEs for enhanced generalization.
Supplementary Material: zip
Primary Area: interpretability and explainable AI
Code Of Ethics: I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics.
Submission Guidelines: I certify that this submission complies with the submission instructions as described on https://iclr.cc/Conferences/2025/AuthorGuide.
Anonymous Url: I certify that there is no URL (e.g., github page) that could be used to find authors’ identity.
No Acknowledgement Section: I certify that there is no acknowledgement section in this submission for double blind review.
Submission Number: 4613
Loading