Standard attention often reuses a token’s own value vector, even though the residual stream already carries that information. XSA removes that redundant component with two lines of code.
An interactive marimo notebook that reproduces and extends Exclusive Self Attention by Shuangfei Zhai (Apple, 2026).
Built for the alphaXiv × marimo Notebook Competition.
This notebook is designed to make Exclusive Self Attention (XSA) intuitive, observable, and interactive.
Rather than aiming for a full-scale reproduction of the paper, it focuses on the paper’s core mechanism:
- showing how standard self‑attention develops a bias toward a token’s own value vector
- explaining why that component can be redundant
- demonstrating how XSA removes it through a simple projection
- testing how this affects training, bias, and sequence length behavior
- exploring a new extension: partial self‑projection
The result is a notebook that helps readers understand not only what XSA does, but also why it works and how to experiment with it themselves on CPU.
| Part | Focus | What you explore |
|---|---|---|
| 1. The problem | Self‑alignment in attention | Live measurement of attention output similarity to the token’s own value |
| 2. The fix | XSA mechanism | Geometric intuition and a direct code comparison between SA and XSA |
| 3. Live training | SA vs XSA | Side‑by‑side training runs from identical initialization |
| 4. Sequence length | Scaling behavior | How the XSA advantage changes as context length increases |
| 5. Mechanism validation | Pre vs post projection | XSA develops strong raw self‑alignment internally, then removes it at the output |
| 6. Partial XSA | Notebook extension | A tunable λ sweep showing that partial removal can outperform full removal |
Standard self‑attention computes:
In practice, the attention output
But the residual connection already carries token‑local information forward unchanged. That means part of the attention output may be spent repeating information the model already has, instead of integrating context from other tokens.
Exclusive Self Attention (XSA) removes this self‑value‑aligned component from the attention output.
In code:
Vn = F.normalize(V, dim=-1) # unit self‑value
Z = Y - (Y * Vn).sum(-1, keepdim=True) * Vn # remove projection onto VBy construction, the final output
This adds:
- no new parameters
- almost no extra compute
- a clean way to separate self‑reuse from context integration
In addition to reproducing the core XSA idea, this notebook adds several interactive and interpretive components:
- live bias visualization, so the self‑alignment problem can be observed directly
- side‑by‑side training from identical weights, for cleaner SA vs XSA comparison
- mechanism‑level validation, showing that XSA does not stop self‑alignment from forming internally, but removes it at the output
- partial self‑projection (λ‑XSA), a notebook extension that explores whether full removal is always optimal
This last extension is especially interesting:
In small interactive runs, full projection removes self‑alignment completely, but intermediate values such as λ = 0.5 can produce better validation loss.
That suggests a trade‑off between eliminating redundancy and preserving useful signal.
XSA is appealing because it is both simple and revealing.
It changes only the geometry of the attention output, yet it exposes an important question:
How much of attention is actually integrating context, and how much is just reusing the token’s own representation?
This notebook makes that question visible.
Instead of treating XSA as a black‑box trick, it turns the method into something you can inspect, train, compare, and modify.
# Clone the repository
git clone https://github.com/turancannb02/xsa-marimo-notebook
cd xsa-marimo-notebook
# Create and activate a virtual environment
python -m venv .venv
source .venv/bin/activate
# Install dependencies
pip install -r requirements.txt
# Set HuggingFace token for TinyStories
export HF_TOKEN="your_token_here"
# Run the notebook
marimo run notebook.pyThen open:
http://localhost:2718
For the full experience, work through the notebook from top to bottom. Parts 3, 4, and 5 include run buttons for interactive experiments.
.
├── notebook.py # main marimo notebook
├── train_utils.py # training loop, data loading, bias metrics
├── requirements.txt # dependencies
├── EXA_repo/ # model implementation generated with PaperCoder
│ ├── model.py # GPT with standard attention and XSA
│ ├── config.py # model configuration
│ └── config.yaml # architecture specs from paper Table 1
└── cache/ # generated at runtime (checkpoints, experiment results)
All experiments are intentionally small and CPU‑friendly so the mechanism can be explored interactively.
Trained on TinyStories with a 2‑layer, 128‑dim model for 300 steps:
| Model | Final validation loss |
|---|---|
| Standard SA | 3.8972 |
| XSA | 3.8678 |
Relative improvement: +0.8%
Toy‑scale sequence length sweep:
| Sequence length | XSA improvement |
|---|---|
| 64 | +1.03% |
| 128 | +0.60% |
| 256 | +0.02% |
Pre‑ and post‑projection cosine similarity with the token’s own value vector:
| Layer | SA ⟨Y, v⟩ | XSA ⟨Y, v⟩ (pre) | XSA ⟨Z, v⟩ (post) |
|---|---|---|---|
| Layer 0 | 0.290 | 0.811 | 0.001 |
| Layer 1 | 0.413 | 0.675 | 0.000 |
A key takeaway:
XSA models can develop stronger raw self‑value alignment than standard attention before projection, then remove it almost completely in the final output.
So XSA does not prevent the bias from arising. It cleans it up.
To explore whether full projection is always optimal, this notebook introduces a scaling factor λ:
- λ = 0.0 → standard attention
- λ = 1.0 → full XSA
- 0 < λ < 1 → partial removal of the self‑value projection
| λ | Paired SA baseline | XSA loss | Improvement |
|---|---|---|---|
| 0.00 | 3.9017 | 3.9017 | +0.00% |
| 0.25 | 3.8822 | 3.8706 | +0.30% |
| 0.50 | 3.9036 | 3.8762 | +0.70% |
| 0.75 | 3.8943 | 3.8735 | +0.53% |
| 1.00 | 3.8993 | 3.8897 | +0.24% |
This suggests that the loss‑optimal amount of projection may not always be full removal.
At least in this interactive setting, partial projection can preserve useful signal while still reducing redundancy.
That makes λ‑XSA a simple but interesting extension of the original idea, and one that can be explored directly inside the notebook.
This notebook shows that:
- standard attention can spend part of its output reusing token‑local information
- XSA removes that component with a simple projection
- even small CPU‑scale experiments make the mechanism visible
- XSA may benefit from partial, not just full, self‑projection
More broadly, the notebook turns a concise research idea into something interactive, inspectable, and extendable.
@article{zhai2026xsa,
title = {Exclusive Self Attention},
author = {Zhai, Shuangfei},
journal = {arXiv preprint arXiv:2603.09078},
year = {2026}
}Built with marimo for the alphaXiv × marimo Notebook Competition.