ESA & Quaternion Transformers: A Research Repository

This repository provides implementations of Eigensingular Attention (ESA) and Quaternionized Self-Attention, two efficient Transformer architectures designed for long-context NLP and parameter-efficient reasoning. It is designed to be a usable, swappable research tool for benchmarking and experimentation.

What this transformer is actually useful for (in practice)

• Long-context NLP/RAG: ESA’s landmark + spectral compression cuts quadratic complexity without the usual accuracy cliff. Use it for documents >8k tokens, minutes-to-hours transcripts, and contractual text.
• Math/code reasoning: Quaternion projections give rotation-aware mixing with ~4× fewer projection params. This buys either smaller models at constant quality or more heads at the same budget—useful for pass@k on GSM8K/HumanEval.
• Edge/low-VRAM: QESA (ESA + Quaternion) attacks both axes—fewer sequence tokens and slimmer projection stacks—so it’s friendly to L4/24GB, T4, and even 3090s.

API Contract

The core attention mechanisms expose a forward method with the following signature:

forward(
    x: torch.Tensor,
    attn_mask: Optional[torch.Tensor] = None,
    position_ids: Optional[torch.Tensor] = None,
    past_key_values: Optional[List[torch.Tensor]] = None,
    use_cache: bool = False,
    **kwargs
) -> (torch.Tensor, Optional[List[torch.Tensor]]):

Tensor Shapes:

• x: (B, N, d_model) - Input tensor.
• attn_mask: (B, 1, N, N) - Attention mask.
• position_ids: (B, N) - Position IDs.
• past_key_values: List[(K, V)] - A list of tuples, where each tuple contains the key and value tensors for a single attention head.
• y: (B, N, d_model) - Output tensor.
• new_past: List[(K, V)] - Updated past key values.

Constraints:

• For Quaternion attention, d_model must be divisible by 4 * n_heads.

Complexity & Memory Math

Eigensingular Attention (ESA):

• Compute: O(B * N * d * r + B * r^2 * k) where r, k « N. This provides near-linear behavior.
• Memory: O(B * N * d)

Quaternion Attention:

• Parameter Count: ≈ 0.25 * (3 * d * d) compared to vanilla QKV projections with quaternion tying.

Minimal Runnable Example

import torch
from models.transformer_qesa import QESATransformer
from utils.common import count_params
import yaml

# Load config from YAML
with open("configs/small_qesa.yaml", "r") as f:
    config = yaml.safe_load(f)

# Instantiate the model
model = QESATransformer(config).cuda().eval()
print("Params (M):", count_params(model) / 1e6)

# Create a random input tensor
x = torch.randn(2, 1024, config['model']['d_model'], device="cuda", dtype=torch.float16)

# Run a forward pass
with torch.no_grad():
    y = model(x)[0]

print("Out:", y.shape)

Hugging Face Integration

To use the models with the Hugging Face transformers library, first register them:

from transformers import AutoConfig, AutoModel
from models.qesa.configuration_qesa import QesaConfig
from models.qesa.modeling_qesa import QesaModel

AutoConfig.register("qesa", QesaConfig)
AutoModel.register(QesaConfig, QesaModel)

Then, you can use them like any other Hugging Face model:

pip install -e .
python - <<'PY'
from transformers import AutoConfig, AutoModel
cfg = AutoConfig.from_pretrained("qesa-small", trust_remote_code=True)
m = AutoModel.from_config(cfg, trust_remote_code=True)
print(sum(p.numel() for p in m.parameters())/1e6, "M params")
PY

Results

PG-19 (Perplexity)

Model	Params	Context	PPL (PG19)

WMT14 (BLEU)

Model	Params	Context	BLEU (WMT14)

SQuAD (F1)

Model	Params	Context	F1 (SQuAD)

Throughput (tokens/s)

Model	Params	Context	toks/s

Peak Memory (GB)

Model	Params	Context	peak GB

Inference Backends & Quantization

• FlashAttention-2/3: Use the --flash argument in bench.py to select the FlashAttention version.
• Quantization: The linear layers in the models are compatible with bitsandbytes and AWQ.
• vLLM: A CPU fallback is available. A full attention_backend='qesa' shim is not yet implemented.
• GGUF: Not yet supported.