sanskxr02/temporal-self-attention

TemporalSelfAttention - A Time-Biased Attention Module

Give Transformers a sense of time - not by scaling, but by structure.

---

Why?

Standard attention treats all tokens equally in time.
This works for syntax, but breaks for:

• Temporal event ordering
• Causal reasoning
• Timeline consistency
• Long-range narrative coherence

💡 Insight: These models simulate time via token position. We inject it structurally with a tiny inductive bias.

---

Core Equation

The time-aware attention score is computed as:

$${\text{score}}_{ij}=\frac{Q_i\cdot K_j^{\mathrm{\top }}}{\sqrt{d_k}}+\gamma \cdot f(t_j-t_i)\backslash text\{score\}\_\{ij\}\; =\; \backslash frac\{Q\_i\; \backslash cdot\; K\_j^\backslash top\}\{\backslash sqrt\{d\_k\}\}\; +\; \backslash gamma\; \backslash cdot\; f(t\_j\; -\; t\_i)$$

Notation

Symbol	Description
${\text{score}}_{ij}\backslash text\{score\}\_\{ij\}$	Attention score between query at position $ii$ and key at position $jj$
$Q_{i}Q\_i$	Query vector for position $ii$
$K_{j}K\_j$	Key vector for position $jj$
$d_{k}d\_k$	Dimension of key vectors
$\gamma \backslash gamma$	Learnable time bias strength
$f(\cdot )f(\backslash cdot)$	Time difference function
$t_j-t_{i}t\_j\; -\; t\_i$	Relative time difference

How To Use

from temporal_attention import TemporalSelfAttention

model = TemporalSelfAttention(
    embed_dim=64,
    num_heads=1,
    bias_type="linear",  # or 'gaussian'
    gamma=1.0,
    causal=False
)

# x: (B, T, D), timestamps: (B, T)
output, weights = model(x, timestamps)