Adaptive Locality Attention (ALA)

Injecting and Controlling Locality Priors in Transformer Self-Attention

Results • Training • Analysis • Installation • Quick Start

---

Overview

Standard Transformer self-attention treats all positions equally, requiring models to learn locality from scratch. This is sample-inefficient and leads to "attention noise"—models wasting probability mass on irrelevant distant tokens.

Adaptive Locality Attention (ALA) addresses this by injecting a learnable locality prior directly into attention logits, with a per-token gate that allows the model to dynamically switch between local and global attention modes.

Key Contributions

• Locality Prior: A learnable distance penalty that encourages attention to nearby tokens
• Adaptive Gating: Per-token gates that allow the model to bypass locality when global context is needed
• Mechanistic Analysis: Beyond perplexity—I analyze attention distance, local mass concentration, and gate behavior
• Reproducible Framework: Multi-seed experiments with statistical significance testing

---

Key Results

Experiments on WikiText-103 (~51M parameters, 10K training steps, 3 seeds):

Performance Summary

Variant	Val Perplexity ↓	Val Loss ↓	Top-1 Acc ↑	Top-5 Acc ↑
Baseline	36.20 ± 0.44	3.589 ± 0.012	38.56%	59.43%
DistBias	35.54 ± 0.27	3.571 ± 0.008	38.79%	59.74%
ALA (Gated)	35.71 ± 0.33	3.575 ± 0.009	38.62%	59.59%

Key Findings

• ~0.65 PPL improvement from injecting locality prior
• >50% reduction in mean attention distance (92 → 43-50 tokens)
• Adaptive gating works: g ≈ 0.75 shows the model learns a nuanced policy
• Minimal overhead: ~12-17% throughput reduction with unchanged O(T²) complexity

---

Performance Metrics

Validation Perplexity

Both locality-biased variants achieve lower perplexity than baseline. DistBias shows the strongest improvement, while ALA (gated) trades a small amount of PPL for adaptive control.

---

Training Dynamics

Loss Curves

All variants converge stably. The locality-biased models (distbias, gated) show consistently lower validation loss throughout training.

Perplexity Over Training

Bits Per Token

Top-1 Accuracy

The biased variants show a slight but consistent edge in validation accuracy, confirming that the locality prior helps the model make better predictions.

Top-5 Accuracy

---

Mechanistic Analysis

Beyond perplexity, I analyze how the attention mechanism changes internally.

Attention Distance

| Variant | Mean Expected Distance E[|i-j|] ↓ | |:--------|:---------------------------------:| | Baseline | 92.67 ± 2.14 | | DistBias | 43.70 ± 2.83 | | ALA (Gated) | 50.54 ± 1.15 |

Key insight: The baseline model spreads attention broadly (~93 tokens back on average). The locality prior drastically reduces this to ~44-50 tokens—a >50% reduction in attention distance.

Local Attention Mass

Variant	LocalMass@32 ↑	LocalMass@64 ↑
Baseline	33.9%	52.0%
DistBias	55.6%	77.1%
ALA (Gated)	53.0%	73.7%

Key insight: In the baseline, nearly half of attention is cast outside the 64 most recent tokens. With locality bias, over 70% of attention stays local—a significant sharpening of focus.

Attention Entropy

Lower entropy indicates more concentrated (less uniform) attention distributions. The biased variants show reduced entropy, confirming sharper attention patterns.

Adaptive Gate Values (ALA Only)

Average Gate Value: 0.7504 ± 0.0014

Key insight: The model does NOT collapse to g=1 (pure local) or g=0 (pure global). Instead, it learns a nuanced policy—applying ~75% of the maximum locality penalty on average, while retaining the ability to reduce it for tokens that need long-range context. This is clear evidence of adaptive control.

---

Speed & Efficiency

Variant	Throughput (tok/s)	Overhead	Parameters
Baseline	150,547 ± 193	—	51.20M
DistBias	131,584 ± 74	-12.6%	51.20M
ALA (Gated)	124,396 ± 192	-17.3%	51.23M

• Complexity: O(T²) unchanged from baseline
• Memory: Nearly identical footprint
• Parameter overhead: ALA adds only ~33K parameters (the gating projection)

The throughput reduction comes from extra element-wise operations (distance matrix computation and gating), but remains practical for research and production use.

---

Methodology

The Core Idea

Problem: Standard attention computes similarity without positional preference:

$$S_{i,j} = \frac{Q_i \cdot K_j^T}{\sqrt{d_k}}$$

Solution: Add a learnable distance penalty:

$$S'_{i,j} = \frac{Q_i \cdot K_j^T}{\sqrt{d_k}} - \alpha_h \cdot |i-j|$$

Innovation (ALA): Make the penalty conditional via a per-token gate:

$$S''_{i,j} = \frac{Q_i \cdot K_j^T}{\sqrt{d_k}} - \alpha_h \cdot g_i^{(h)} \cdot |i-j|$$

$$g_i^{(h)} = \sigma(W_g \cdot x_i + b_g) \in (0,1)$$

Interpretation

• When g ≈ 1: Head enforces strong locality (syntax, morphology)
• When g ≈ 0: Head ignores distance (coreference, retrieval)
• The model learns when to be local vs. global

Three Variants Compared

Variant	Description	Added Parameters
baseline	Standard scaled dot-product attention	None
distbias	+ Learnable per-head distance penalty (αₕ)	H scalars
gated_distbias (ALA)	+ Per-token gate to modulate penalty	+ Linear(d→H)

---

Installation

Requirements

• Python 3.9+
• PyTorch 2.2+
• CUDA-capable GPU (recommended)

Setup

# Clone the repository
git clone https://github.com/yourusername/adaptive-locality-attention.git
cd adaptive-locality-attention

# Create virtual environment
python -m venv .venv
source .venv/bin/activate  # Windows: .venv\Scripts\activate

# Install dependencies
pip install -r requirements.txt

Dependencies

torch>=2.2
transformers>=4.40
datasets>=2.18
tqdm>=4.66
numpy>=1.24
pandas>=2.0
matplotlib>=3.8
pyyaml>=6.0
tabulate>=0.9

---

Quick Start

Run the Full Experiment Suite

This trains all 3 variants with 3 seeds each and generates all results:

bash scripts/run_wt103_512_suite.sh

Outputs:

• results/wt103_bs512_aggregate.md — Aggregated metrics (mean ± std)
• results/wt103_bs512_runs.md — Per-run detailed results
• results/wt103_bs512_*.png — All training curves and bar charts

Train a Single Model

python src/train.py --config configs/wt103_512_gated.yaml

Analyze Attention Patterns

python src/analyze_attention.py --checkpoint runs/your_run/checkpoint.pt

---

Project Structure

adaptive_locality_attention/
├── configs/                    # YAML experiment configurations
│   ├── wt103_512_baseline.yaml
│   ├── wt103_512_distbias.yaml
│   └── wt103_512_gated.yaml
├── src/                        # Source code
│   ├── attention.py            # ALA implementation
│   ├── model.py                # Transformer architecture
│   ├── train.py                # Training loop
│   ├── eval.py                 # Evaluation utilities
│   ├── analyze_attention.py    # Mechanistic analysis
│   ├── aggregate_results.py    # Multirun aggregation
│   └── plot_metrics.py         # Visualization
├── scripts/                    # Automation scripts
│   └── run_wt103_512_suite.sh  # Full experiment suite
├── results/                    # Generated plots and tables
│   ├── wt103_bs512_*.png       # All visualizations
│   ├── wt103_bs512_aggregate.md
│   └── wt103_bs512_runs.md
├── reports/                    # Documentation
│   ├── ALA_Report.pdf          # Full technical report
├── presentation/               # Presentation slides
├── requirements.txt
└── README.md

---

Configuration

Experiments are controlled via YAML configs:

# configs/wt103_512_gated.yaml
seed: 1337

data:
  dataset: wikitext103
  tokenizer_name: gpt2
  block_size: 512
  batch_size: 16

model:
  n_layers: 8
  n_heads: 8
  d_model: 512
  d_ff: 2048
  dropout: 0.1
  
  # ALA-specific
  attn_variant: gated_distbias  # baseline | distbias | gated_distbias
  alpha_init: 0.01              # Initial locality strength
  alpha_learnable: true         # Learn α during training
  gate_enabled: true            # Enable per-token gating

train:
  max_steps: 10000
  lr: 0.0002
  warmup_steps: 400
  weight_decay: 0.1

---

Understanding the Code

Core Implementation (src/attention.py)

The key innovation is in the attention score computation:

# Standard attention logits
att = (q @ k.transpose(-2, -1)) / math.sqrt(self.head_dim)

# Add locality bias
if self.use_distbias:
    alpha = F.softplus(self.raw_alpha)  # Ensure α > 0
    dist = self.dist[:T, :T]            # |i-j| matrix
    
    if self.gate_enabled:
        gate = torch.sigmoid(self.gate_proj(x))  # Per-token gate
        bias = -alpha * gate * dist
    else:
        bias = -alpha * dist
    
    att = att + bias

---

Detailed Results

Per-Run Results (All Seeds)

| Run | Variant | Val PPL | Val Acc1 | E[|i-j|] | LocalMass@64 | Gate | |:----|:--------|:-------:|:--------:|:--------:|:------------:|:----:| | seed1337 | baseline | 36.63 | 38.54% | 94.91 | 50.97% | — | | seed42 | baseline | 35.60 | 38.68% | 89.79 | 53.02% | — | | seed2025 | baseline | 36.37 | 38.46% | 93.31 | 52.10% | — | | seed1337 | distbias | 35.68 | 38.75% | 47.57 | 74.71% | — | | seed42 | distbias | 35.16 | 38.62% | 40.86 | 79.13% | — | | seed2025 | distbias | 35.76 | 39.01% | 42.68 | 77.52% | — | | seed1337 | gated | 35.31 | 38.41% | 52.17 | 72.59% | 0.749 | | seed42 | gated | 36.13 | 39.05% | 49.81 | 74.42% | 0.751 | | seed2025 | gated | 35.69 | 38.41% | 49.65 | 74.18% | 0.752 |

---

Reports & Documentation

Document	Description
reports/ALA_Report.pdf	Full technical report
presentation/	Presentation slides
results/wt103_bs512_aggregate.md	Aggregated results

---

Reproducibility

All experiments use:

• 3 random seeds: 1337, 42, 2025
• Fixed hyperparameters across variants
• Same data splits and preprocessing

Results are reported as mean ± standard deviation across seeds.

# Reproduce all results
bash scripts/run_wt103_512_suite.sh

---

Future Directions

1. Sparse Attention: Use learned locality to prune distant keys → O(T log T) or O(T)
2. Multi-scale Gating: Different α values for different distance ranges
3. Cross-attention: Apply locality priors to encoder-decoder attention
4. Larger Scale: Validate on billion-parameter models

---

Citation

@misc{agnihotri2025ala,
  author       = {Agnihotri, Aman},
  title        = {Adaptive Locality Attention: Injecting and Controlling 
                  Locality Priors in Transformer Self-Attention},
  year         = {2025},
  howpublished = {Independent Research Project},
  institution  = {Delhi Skill & Entrepreneurship University}
}

---

References

• Vaswani, A., et al. (2017). Attention Is All You Need. NeurIPS.
• Merity, S., et al. (2017). Pointer Sentinel Mixture Models. ICLR.

---

Aman Agnihotri
B.Sc. Data Analytics, Delhi Skill & Entrepreneurship University
Independent Research Project • December 2025

A solo research project exploring interpretable and controllable attention mechanisms in Transformers.