Hooshaai/gpt2-adasplash-attention

GPT-2 with AdaSplash Attention via NeuralOps

Model Details

• Model Identifier: Hooshaai/gpt2-adasplash-attention
• Base Architecture: GPT-2 Base (gpt2, 124M parameters, 12 layers, 12 heads, hidden size 768)
• Attention Mechanism: adasplash_attention (Adaptive Causal Splash Attention with Splash Radius $W=32$ and $\alpha=1.5$ Entmax)
• Attention Sparsity: 41% exact mathematical zero weights ($A_{ij} \equiv 0$) across autoregressive causal generation
• Parameters: 125.6M (+0.96% parameter count via AdaSplashLinear residual LoRA adapters)
• Framework: neuralops v1.4.3
• Organization: Hooshaai

Architecture Overview

Autoregressive causal generation in GPT-2 requires tokens at time $t$ to evaluate attention over all preceding tokens $j \le t$. As generation length extends to thousands of tokens, caching and attending over full-length Key-Value tensors consumes extensive memory bandwidth while attending to distant, irrelevant context.

AdaSplash Attention on GPT-2 combines two complementary principles:

1. Selective Sliding-Window Splash Masking ($W=32$): Enforces a local causal receptive field that concentrates dense generation compute on the most recent 32 tokens while preserving attention sinks at token position 0 (<|endoftext|>).
2. Tsallis 1.5-Entmax Simplex Projection: Replaces dense Softmax with $\alpha=1.5$ entmax, projecting attention affinities onto the probability simplex and assigning identically zero attention weight ($A_{ij} \equiv 0$) to 41% of entries within the active window.
3. AdaSplashLinear Adapters: Augments pre-trained $W_{\text{attn}}$ and $W_{\text{proj}}$ projections with low-rank adapters, maintaining rapid convergence and 87.84% SST-2 classification fidelity.

Autoregressive Input X [B, L, 768]
       │
       ├── AdaSplashLinear Projections: Q, K, V via (W_base + B · A)
       ├── Causal Splash Masking:
       │     Restricts attention to j ∈ {0} ∪ [max(1, i - 32), i]
       ├── Exact Tsallis 1.5-Entmax Simplex Projection:
       │     Solves: ∑ max(0, (S_{i, j} - τ)/2)^2 = 1 via bisection search
       │     Produces A_{i, j} with 41% exact zeros
       ├── Sparse Value Aggregation: O_i = ∑_{j: A_{i, j} > 0} A_{i, j} V_j
       └── Output Projection: Y = AdaSplashLinear_proj(O)  [B, L, 768]

Mathematical Formulation & Derivation

1. Causal Splash Masking with Sink Preservation

For query position $i \in {1, \dots, L}$ and key position $j \in {1, \dots, i}$, the active set $\mathcal{A}_i$ preserves both the sequence initial sink and the local splash radius $W = 32$: [ \mathcal{A}_i = {0} \cup \left{ j \in {1, \dots, i} : i - j \le W \right} ] Tokens outside $\mathcal{A}i$ are set to $-\infty$ before simplex projection: [ \widehat{S}{i, j} = \begin{cases} \frac{Q_i K_j^\top}{\sqrt{d_k}} & \text{if } j \in \mathcal{A}_i \ -\infty & \text{otherwise} \end{cases} ]

2. Tsallis 1.5-Entmax on Causal Logits

The attention weights over active tokens are: [ A_{i, j} = \left[ \frac{\widehat{S}{i, j} - \tau_i}{2} \right]+^2 = \max\left(0, ; \frac{\widehat{S}{i, j} - \tau_i}{2}\right)^2 ] where $\tau_i$ is solved via bisection search satisfying: [ \sum{j \in \mathcal{A}i} \left[ \frac{\widehat{S}{i, j} - \tau_i}{2} \right]_+^2 = 1 ] Tokens with logits below $\tau_i$ receive mathematically exact zero weight, preventing attention dilution during long-context causal generation.

Empirical Benchmark Table

Evaluated on GLUE SST-2 and WikiText-2 generation on NVIDIA GeForce RTX 4070 (SM 8.9) with sequence length $L = 128$:

Metric	Dense GPT-2 Baseline	Zero-Shot AdaSplash	Post-Distillation (v1.4.3)	Benefit vs. Baseline
SST-2 Classification	89.33%	79.39%	87.84%	Minimal delta (-1.49%)
WikiText-2 Perplexity (PPL)	39.33	44.20	41.10	High generation fluency
Attention Sparsity (% Zeros)	0.0% (Softmax)	41.0%	41.0%	Exact Mathematical Zero Weights
Active Parameters	124.4M	125.6M	125.6M	+0.96% LoRA adapters
Triton Kernel Latency	4.20 ms	1.54 ms	1.54 ms	2.73x Kernel Speedup
Inference Step Time	20.5 ms	14.4 ms	14.4 ms	1.42x Speedup
Peak VRAM Footprint	793 MB	391 MB	391 MB	-50.7% Reduction

Quickstart Code

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from neuralops import replace_attention

# 1. Load tokenizer and model from Hugging Face Hub
tokenizer = AutoTokenizer.from_pretrained("gpt2")
model = AutoModelForCausalLM.from_pretrained("Hooshaai/gpt2-adasplash-attention")

# 2. Or surgically inject AdaSplash into any GPT-2 in 1 line
base_model = AutoModelForCausalLM.from_pretrained("gpt2")
adasplash_model = replace_attention(base_model, "adasplash_attention")

# 3. Generate text
prompt = "Sparse attention mechanisms with exact mathematical zero weights"
inputs = tokenizer(prompt, return_tensors="pt")

with torch.no_grad():
    outputs = model.generate(
        **inputs,
        max_new_tokens=48,
        temperature=0.8,
        do_sample=True,
    )

print(tokenizer.decode(outputs[0], skip_special_tokens=True))

Hardware & Software Requirements

• PyTorch: 2.11+
• CUDA: 12.8+
• Triton: 3.6+
• Accelerators: NVIDIA SM 8.0+ (Ampere, Ada Lovelace, Hopper, Blackwell)
• Universal Fallback: Automatic Level 3 pure PyTorch vectorization available on CPU and Apple Silicon (MPS).

Citation

@software{neuralops2026,
  author = {Majlesi, Taha and contributors},
  title = {NeuralOps: Unified Platform for Subquadratic Attention and Deep Model Compression},
  year = {2026},
  publisher = {GitHub},
  howpublished = {\url{https://github.com/tahamajs/NeuralOps}},
  version = {1.4.3}
}