One-sentence Summary

Researchers from University of California, Irvine and University of Massachusetts Amherst inject tunable noise into the hidden representations of Transformer-based language models and train them with a hybrid objective that constrains total encoding precision, simulating resource-rational working memory allocation and yielding improved alignment with human reading times as well as more compressed and categorical context representations.

Key Contributions

• Working-memory constraints are added to Transformer-based language models by injecting tunable noise into hidden representations and training with a hybrid objective that balances next-word prediction against an explicit limit on encoding precision.
• These constraints improve alignment with human reading times and reshape context representations to be more compressed and categorical, supporting a resource-rational account of memory allocation.
• Moderate constraints increase locality bias in attention weights, while stricter constraints weaken it, revealing a dissociation between retrieval mechanisms and encoded representations and linking working-memory constraints to quantization.

Introduction

Human sentence processing relies on working memory because earlier words must influence later interpretation after they are no longer directly available. Since human working memory is highly limited, understanding how these constraints shape language representations is important for both psycholinguistics and for building language models that behave more like humans. Prior computational work has mainly simulated working memory limits through ad hoc choices, such as smaller models, shorter context windows, biased local attention, or binary erasure of past words, but these approaches do not directly explain how limited resources affect the quality of information stored in memory. The authors instead treat working memory constraints as limits on representational precision and show that noisier, less precise memory encodings can emerge as a resource rational response to linguistic prediction demands.

Method

The authors model working memory constraints in Transformer-based language models by making the internal memory representation deliberately noisy. The core idea is that linguistic inputs are not encoded perfectly in working memory, and the fidelity of that encoding depends on the amount of working memory resources available to the model. In the proposed framework, working memory capacity is operationalized through encoding precision, and precision is manipulated by injecting Gaussian noise into the value vectors used by self-attention.

The method builds on the standard self-attention mechanism. Normally, a query vector interacts with key vectors to determine attention weights, and those weights are used to aggregate value vectors from earlier positions. To model imperfect memory, the authors corrupt each historical value vector before aggregation. Given a value vector $\mathbf{v}_{qk}$vqk​ corresponding to a previous key position $k$k relative to the current query position $q$q, the noisy value $\tilde{\mathbf{v}}_{qk}$v~qk​ is sampled from a Gaussian distribution:

$$\tilde{\mathbf{v}}_{qk} \sim \mathcal{N} \left(\mathbf{v}_{qk}, \tau^{-1} \mathbf{I}\right).$$v~qk​∼N(vqk​,τ−1I).

Here, $\tau$τ denotes the encoding precision, defined as the inverse of the variance. A higher precision corresponds to lower noise and therefore a more faithful memory representation. This connects the computational model to an information-theoretic interpretation: the noisy attention pathway behaves like a communication channel with limited signal-to-noise ratio. Models with greater working memory resources can maintain higher precision, while models with tighter constraints transmit information under stronger noise.

Precision is not fixed across all past tokens. Instead, it decays exponentially as a token becomes more distant from the current query position. The precision for a key token at position $k$k relative to query position $q$q is given by:

$$\tau_{qk} = \tau_0 e^{-\theta (q-k)}.$$τqk​=τ0​e−θ(q−k).

The parameter $\theta$θ controls the rate of decay, and $\tau_0$τ0​ is the highest achievable precision, corresponding to the current token at query position where $q-k=0$q−k=0. Even this base precision is finite, reflecting the assumption that encoding is never perfectly noiseless. This design allows the model to represent a graded memory effect: recent information is encoded with higher precision, while more remote information becomes progressively noisier.

The output of an attention head is then formed by aggregating the noisy value vectors using the standard softmax attention weights:

$$\mathbf{z}_q = \sum_k a_{qk} \tilde{\mathbf{v}}_{qk},$$zq​=k∑​aqk​v~qk​,

with attention weights

$$\mathbf{a}_q = \operatorname{softmax} \left(\frac{\mathbf{K}^{\top} \mathbf{q}}{\sqrt{d}}\right).$$aq​=softmax(d​K⊤q​).

To optimize this stochastic attention mechanism efficiently, the authors apply a reparameterization trick. The attention output is decomposed into a deterministic component and a noise component:

$$\mathbf{z}_q = \mathbf{z}_q^{\mathrm{det}} + \tilde{\mathbf{z}}_q.$$zq​=zqdet​+z~q​.

The deterministic part $\mathbf{z}_q^{\mathrm{det}}$zqdet​ is computed as in standard self-attention, as if no noise were present. The noisy part is expressed as a scaled unit Gaussian:

$$\tilde{\mathbf{z}}_q = \sigma_q \boldsymbol{\epsilon}, \quad \boldsymbol{\epsilon} \sim \mathcal{N}(0, I).$$z~q​=σq​ϵ,ϵ∼N(0,I).

Assuming independence of the noise contributions from different previous tokens, the effective variance after attention aggregation is

$$\sigma_q^2 = \sum_k \left(a_{qk} \sigma_{qk}\right)^2 = \sum_k a_{qk}^2 \tau_{qk}^{-1}.$$σq2​=k∑​(aqk​σqk​)2=k∑​aqk2​τqk−1​.

This formulation preserves parallelized attention computation while introducing position-dependent memory noise. The authors inject noise specifically into self-attention rather than at the output embeddings because self-attention most closely mirrors the working memory process of retrieving and combining representations from different points in context. It also makes it natural to control noise as a function of the relative distance between the stored token and the current query position.

Working memory constraints are then manipulated during training through a hybrid loss function. The model is not optimized solely for language modeling accuracy. Instead, it must balance predictive performance against the cost of maintaining precise memory representations. The training objective is

$$\mathcal{L} = \underbrace{- \sum_{\mathbf{y} \in \mathbf{S}} \sum_{t=1}^{T} \log p \left(y_t \mid \mathbf{y}_{< t}\right)}_{\text{cross-entropy loss}} + \lambda \underbrace{\sum_{h,l,t,k} \tau_{h,l,t,k} z^{-1}}_{\text{memory capacity}},$$L=cross-entropy loss−y∈S∑​t=1∑T​logp(yt​∣y<t​)​​+λmemory capacityh,l,t,k∑​τh,l,t,k​z−1​​,

where

$$z = L H T \frac{T-1}{2}.$$z=LHT2T−1​.

The first term is the standard next-word prediction loss over training examples. The second term penalizes the average encoding precision of past tokens across attention heads and layers. This term represents the amount of working memory resources used by the model. The scalar $\lambda$λ controls the strength of the working memory constraint. A larger $\lambda$λ imposes a stronger penalty on high precision, encouraging the model to operate with fewer working memory resources. As a result, the model learns a trade-off between linguistic prediction performance and the cost of encoding information in memory.

This design yields a unified method for studying working memory limitations in Transformers. Gaussian noise injection creates a continuous, distributional representation of memory imperfection, while the precision decay function introduces distance-dependent degradation. The hybrid training objective then explicitly shapes how much precise memory the model is allowed to allocate, making working memory capacity a controllable property of the language model architecture.

Experiment

The experiments evaluate GPT-2-style language models trained with varying working-memory constraints, examining learned precision decay, alignment with human reading times, and the geometry of next-word prediction spaces. As the memory constraint increases, models show asymmetric precision decay across layers and generate predictions that align better with human reading behavior at moderate constraint levels, despite higher perplexity and reduced extreme surprisal estimates. Stronger constraints also compress the prediction space and promote more categorical encoding, indicating a shift from fine-grained distinctions toward broader semantic generalization.