goat-attention 0.1.1

Generalized Optimal Transport Attention with Trainable Priors (PyTorch).

GOAT Attention

Generalized Optimal Transport Attention with Trainable Priors (GOAT), available as a PyTorch multi-head attention module.

Install name: goat-attention (PyPI) · Import name: goat

Installation

• From PyPI (recommended):

uv add goat-attention

• pip:

pip install goat-attention

• From source (editable):

uv pip install -e .

• From source (editable, pip):

pip install -e .

Quickstart

import torch
from goat import GoatAttention

B, L, S, E, H = 2, 5, 7, 64, 8
xq = torch.randn(B, L, E)
xk = torch.randn(B, S, E)
xv = torch.randn(B, S, E)

attn = GoatAttention(
    embed_dim=E,
    num_heads=H,
    batch_first=True,
    pos_rank=2,
    abs_rank=4,
    enable_key_bias=True,
)

out, weights = attn(xq, xk, xv, is_causal=False, need_weights=True)
print(out.shape, None if weights is None else weights.shape)

CLI

After installation:

goat info
goat smoke

Documentation

See docs/:

• docs/index.md
• docs/usage.md
• docs/api.md
• docs/development.md

Development

uv pip install -e ".[dev]"
pytest

License

MIT (see LICENSE).

Citation

If you find GOAT useful, please cite:

@misc{goat,
  title         = {You Need Better Attention Priors},
  author        = {Litman, Elon and ...},
  year          = {2026},
  eprint        = {XXXX.XXXXX},
  archivePrefix = {arXiv},
  primaryClass  = {cs.LG}
}