Modelpublic

nkiyohara/lra-image-s4d-token-mixer-rope

sourceHugging Faceapache-2.0updated 5mo agoView on Hugging Face

0likes

Model Card

LRA-Image S4D Token Mixer + RoPE

S4D Token Mixer + RoPE model trained on the Long Range Arena (LRA) sCIFAR-10 benchmark.

Model Details

—Architecture: S4D Token Mixer + RoPE
—Task: Sequential CIFAR-10 (grayscale, 1024 tokens)
—Parameters: ~4.1M
—Position encoding: RoPE
—Causal: False
—Test Accuracy: 75.19%

Training

—Protocol: S4 paper LRA-Image (200 epochs, lr=1e-3, batch=64, warmup=18k steps)
—Backbone: Llama-style (6 layers, d=512, 8 heads, RMSNorm, SwiGLU)
—Seed: 2222
—WandB run: harrisonzhu/InterdomainAttention/ptjwayj3

Usage

python

import torch
from model import LlamaLRAImage  # requires interdomain-attention repo

state_dict = torch.load("lra_image_best.pt", weights_only=True)
model = LlamaLRAImage(...)  # match config
model.load_state_dict(state_dict["model"])

Citation

bibtex

@article{interdomain2026,
  title={Interdomain Attention},
  author={...},
  year={2026}
}

License

Apache 2.0