Polynomial kernel attention of odd degree

[sunnysfwong]

Hi. I recently trying to replace exponential with polynomial in softmax attention and achieve validation performance (val_bpb) on par with, or slightly better than, softmax attention across model scales (d=6,8,10,12,14). The formulation looks like the following:

import torch
import torch.nn.functional as F

def poly_attn_naive(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor:
    """
    Input/Output shapes: (B, H, T, D)
    """
    q = F.normalize(q, dim=-1)
    k = F.normalize(k, dim=-1)
    dot = (q * 2) @ k.transpose(2, 3)
    att = torch.tril(dot.pow(9) + 0.01 * dot)
    x = att @ v
    return F.rms_norm(x, (x.size(-1),))

Experiments were conducted on an A100 GPU (Colab Runtime Version 2026.07) with a vocabulary size of 32768. Below is the final validation bits-per-byte (val_bpb, lower is better) across depths:

Depth (d)	Softmax (val_bpb)	Polynomial (val_bpb)	$\Delta$
6	1.0223	1.0198	-0.0025
8	0.9458	0.9421	-0.0037
10*	0.8928	0.8925	-0.0003
12*	0.8465	0.8475	+0.0010
14*	0.8115	0.8107	-0.0008

*Results for depths 10, 12, and 14 were obtained using an earlier revision of the kernel and a different runtime version.

The code can be found in https://github.com/sunnysfwong/nanochat_poly/tree/poly_attn

I current don't have enough compute units in Colab so I will rerun depth=10,12,14 later.

I hope you find this attention mechanism interesting.