Enhancing Length Extrapolation in Sequential Models with Pointer-Augmented Neural Memory

TL;DR

PANM introduces a memory module with explicit physical addresses and a Pointer Unit that learns to manipulate pointers for symbol processing, enabling robust length extrapolation across diverse sequential tasks. By isolating pointer manipulation from input data and coupling it with an address bank, PANM supports Mode-1 and Mode-2 accesses that empower fundamental models like Transformer and LSTM to perform complex symbolic operations without task-specific architecture tweaks. Across algorithmic reasoning, Dyck language recognition, SCAN compositional learning, and practical NLP tasks, PANM consistently improves length generalization and reduces overfitting while maintaining compatibility with common backbones. The work demonstrates that explicit pointer-based memory and data-symbol separation provide a principled route to systematic generalization with potential applicability to real-world reasoning and language tasks.

Abstract

We propose Pointer-Augmented Neural Memory (PANM) to help neural networks understand and apply symbol processing to new, longer sequences of data. PANM integrates an external neural memory that uses novel physical addresses and pointer manipulation techniques to mimic human and computer symbol processing abilities. PANM facilitates pointer assignment, dereference, and arithmetic by explicitly using physical pointers to access memory content. Remarkably, it can learn to perform these operations through end-to-end training on sequence data, powering various sequential models. Our experiments demonstrate PANM's exceptional length extrapolating capabilities and improved performance in tasks that require symbol processing, such as algorithmic reasoning and Dyck language recognition. PANM helps Transformer achieve up to 100% generalization accuracy in compositional learning tasks and significantly better results in mathematical reasoning, question answering and machine translation tasks.

Figures (8)

• Figure 1: PANM architecture. (a) The data memory contains the encoded input sequence (b) The address bank contains physical addresses associated with data memory slots. The base and end addresses ($p_{B},p_{E}$) define the address range of the input sequence. (c) The Pointer Unit takes $p_{B},p_{E}$, recurrently generates the current pointer $p_{t}^{a}$ and gets its value $^{*}p_{t}^{a}$ via Mode-1 (red)/2 (green) Access. (d) The Controller takes pointer information, decoding input ($z_{t}=y_{t-}$), and produce the $t$-th output token $\hat{y_{t}}$.
• Figure 2: Exemplar results on 2 algorithms. (a, b) Test accuracy (mean $\pm$ std) over 5 runs on Copy and ID Sort on each length test, respectively. Random predictor would reach around 10% accuracy. (c,d) Visualization of data and pointer's slots for Copy and ID Sort, respectively.
• Figure 3: (a) Dyck: mean $\pm$ std. accuracy over 5 runs with different testing lengths. (b) Machine translation task: Perplexity on Multi30K dataset (the lower the better). We sort the sequences in the data by length and create 2 settings using train/test split of 0.8 and 0.5, respectively. The baselines are Transformer and PANM. Left: The best test perplexity over 2 settings for different number of Transformer's layers (1 to 3 layers). Right: an example of testing perplexity curves over training epochs for the case of 0.5 train/test split (2 layers) where we run 3 times and report the mean$\pm$std. The y-axis is visualized using log scale.
• Figure 4: PANM as a plug-and-play architecture. The encoder and decoder can be any model (LSTM, Transformer or BERT). PANM Controller can be used as the last layer of the Decoder to access the memory during decoding. To reduce the number of parameters of the augmented architecture, the decoder's number of layers can be decreased.
• Figure 5: Dyck (Left): mean $\pm$ std. accuracy over 5 runs with different testing lengths. bAbI QA (Right): mean $\pm$ std. testing accuracy and cross-entropy loss across 100 training epochs over 5 runs.