GridPE: Unifying Positional Encoding in Transformers with a Grid Cell-Inspired Framework

TL;DR

This paper introduces GridPE, a novel positional encoding inspired by grid cells and Fourier analysis, which improves transformer performance by effectively encoding spatial information in high-dimensional spaces.

Contribution

We propose a neuroscience-inspired positional encoding scheme, GridPE, that unifies spatial encoding in transformers and demonstrates improved performance in high-dimensional tasks.

Findings

GridPE enhances transformer accuracy on spatial tasks.

Theoretical analysis confirms translational invariance of GridPE.

Optimal grid scale ratios improve encoding efficiency.

Abstract

Understanding spatial location and relationships is a fundamental capability for modern artificial intelligence systems. Insights from human spatial cognition provide valuable guidance in this domain. Neuroscientific discoveries have highlighted the role of grid cells as a fundamental neural component for spatial representation, including distance computation, path integration, and scale discernment. In this paper, we introduce a novel positional encoding scheme inspired by Fourier analysis and the latest findings in computational neuroscience regarding grid cells. Assuming that grid cells encode spatial position through a summation of Fourier basis functions, we demonstrate the translational invariance of the grid representation during inner product calculations. Additionally, we derive an optimal grid scale ratio for multi-dimensional Euclidean spaces based on principles of biological efficiency. Utilizing these computational principles, we have developed a Grid-cell inspired Positional Encoding technique, termed GridPE, for encoding locations within high-dimensional spaces. We integrated GridPE into the Pyramid Vision Transformer architecture. Our theoretical analysis shows that GridPE provides a unifying framework for positional encoding in arbitrary high-dimensional spaces. Experimental results demonstrate that GridPE significantly enhances the performance of transformers, underscoring the importance of incorporating neuroscientific insights into the design of artificial intelligence systems.