Stochastic Parroting in Temporal Attention -- Regulating the Diagonal Sink

The gist The authors provide a theoretical characterization of the Jacobian of Temporal Attention to show that temporal attention layers suffer from a diagonal attention sink, which leads to stochastic parroting in TA and showing that this intensifies with increasing sequence length

Introduction

Spatio-temporal deep learning models aim to predict how systems evolve over time and across multiple locations. A common challenge in spatio-temporal forecasting is to exchange information effectively between both nearby and distant nodes with two well-known information degeneration issues: over-squashing, which is the insensitivity of a node’s features to information contained at distant nodes, and over-smoothing, where node representations become increasingly similar. Most spatio-temporal models either mix temporal and spatial processing at every layer, i.e. Time&Space (T&S), or do all temporal processing first and all spatial processing afterwards, i.e. Timethen-Space (TTS). Both T&S and TTS are affected by spatio-temporal oversquashing because their gradients still multiply across layers in the same way.

Previous Work on Over-squashing and Over-smoothing

Di Giovanni et al. laid the groundwork for subsequent research on oversquashing by introducing sensitivity bounds on the Jacobian, which measure how sensitive the final representation of a node v at layer L, h v(L), is to the initial representation of anoter node u, h u(0). Arroyo et al. further showed that oversmoothing is a representational collapse to a zero fixed point, driven by the same contractive Jacobians responsible for vanishing gradients. For GNNs this has been solved recently and in Temporal Convolutional Networks (TCNs) temporal causal convolutions loose focus on recent information, while giving disproportionate weight to inputs from the distant past, creating a “primacy bias”. A similar phenomenon can be seen in causal Transformers, which only use preceding context via unidirectional attention. In [7], the emergence of an ”attention sink” in the first token is empirically demonstrated for auto-regressive Large Language Transformers. Barbero et al. examine the Jacobian of auto-regressive decoders, showing that for a number of layers L → ∞ y n is only sensitive to the first token. In [9] the authors examine full attention and show that the residual of tokens shrink cubically per layer with a contraction factor proportional to √4β dq k, where β is the upper bound of the norm of weight matrices. As a consequence residual (skip) connections and appropriate choices of network dimensionalities are important to avoid rank collapse.

Deriving Sensitivity Bounds of Temporal Attention

The overall Jacobian without residual is thus: ∂h i/∂x j = α ijWVz> Value Path + X T m=1 [vm (X T k=1 α im(δmk − αik)z] Softmax Gradient · 1/√dk (δijk T WQz Query Path + δkjq T i W Kz Key Path] Score Gradient. The overall Jacobian is given by: ∂h i/∂x j = α ijWVz> Value Path + X T m=1 [vm (X T k=1 α im(δmk − αik)z] Softmax Gradient · 1/√dk (δijk T WQz Query Path + δkjq T i W Kz Key Path] Score Gradient. We derive a uniform expectation of the Jacobian norm over the sequence length j ∈ I for the Value Path, Key Path, and Query Path separately. The expected norm for off-diagonal terms (i ≠ j) is bounded by 1/T X T j=1 α ij kWV k = kWV k / T.

Analysis and Regularization of Diagonal Attention Sinks

The 1/T-term in the upper part of Eq. 8 shows that the non-local signal strength (i ≠ j) decays as the sequences are longer (O(1/T)). The lower part of Eq. 8 for the diagonal (i = j) has a larger bound since it additionally has the query term’s bound CQ/√dk. To reduce the diagonal attention sink, we evaluated three different regularization methods: a diagonal mask with diagonal entries −∞ and others zero, similar to the approach in SparseBert [12], standard dropout on the diagonal of the attention matrix, and a negative scalar penalty added to the raw attention scores.

Empirical Evaluation

Table 1 shows the results over four runs with different seeds for a) without residual and without regularization, b) with residual and without regularization, c) with residual and with diagonal mask, d) with residual and with diagonal dropout, and e) with residual and diagonal penalty. Using a diagonal mask (c) gives similar results as the variant without regularization (b), while dropout and penalty show a significant improvement of around 2.5 percent. In (d) and (e) we see clear temporal patterns with specific keys attending to specific queries.

Conclusion

We prove the existence of diagonal attention sinks, induced by the residual and the weight matrix of the query. While full diagonal masks suppress the Query Path of the softmax attention, resulting in attention emphasizing specific keys for the whole sequence, diagonal penalty and diagonal dropout significantly improve the performance of temporal attention layers. The larger T is, the smaller off-diagonal entries become. Overall, the diagonal penalty gives best results for longer horizons.

How it works

The authors derive sensitivity bounds on the Jacobian of temporal attention layers to analyze information degeneration. They show that TA layers can collapse temporal information, which leads to stochastic parroting in TA and showing that this intensifies with increasing sequence length. The derivation involves examining the influence of x j ∈ R dx for an input x i ∈ R dx on the hidden state h i ∈ R dv after one layer by deriving the Jacobian.

Key Findings and Regularization Effects

1. The expected norm for off-diagonal terms (i ≠ j) decays as the sequences are longer (O(1/T)).

2. The lower part of Eq. 8 for the diagonal (i = j) has a larger bound since it additionally has the query term’s bound CQ/√dk.

3. Regularization methods include a diagonal mask with diagonal entries −∞ and others zero, standard dropout on the diagonal of the attention matrix, and a negative scalar penalty added to the raw attention scores.

4. The experiments suggest that the partial control of the diagonal using dropout or penalty is superior to a full diagonal mask.

5. With a full diagonal mask (c) we see similar results as without regularization (b).

Evaluation on Long Horizon Forecasting

The evaluation extended to long-horizon forecasting, predicting 36 future steps (i.e., 3 hours) <ref:2602.

1. textbfProcessing temporal attention collapse via diagonal regularization: The improved system will mitigate stochastic parroting by applying a diagonal mask with diagonal entries −∞ and others zero, which is suggested as effective in mitigating the diagonal attention sink. This allows the model to allocate learning capacity towards the off-diagonal entries, leading to more expressive attention patterns without sacrificing temporal stability, as shown in Table 1 and Figure 1(e).

2. textbfIntegrating residual connections: The system should utilize residual connections because they remain universally expressible even when the diagonal is completely removed, which suggests that the diagonal provides the least important information. This feature ensures that self-information remains stable regardless of sequence length (O(1)) and significantly improves performance, as evidenced by the substantial error reduction when comparing models with and without residual connections.

3. textbfEmploying diagonal penalties for long-horizon forecasting: For long-term predictions, the model should implement a negative scalar penalty added to the raw attention scores. This method is found to be superior to masking or dropout, as it allows the model to use the diagonal entries when beneficial, which is crucial for capturing more heterogeneous patterns in longer sequences.

4. textbfEnhanced spatio-temporal information exchange: By deriving sensitivity bounds, the system can theoretically understand how off-diagonal attention scores depend on the sequence length, enabling informed architectural adjustments that address over-squashing and prevent the degradation of distant node information when integrating temporal and spatial processing.