notebook.py

import marimo

__generated_with = "0.23.1"
app = marimo.App(width="medium")


@app.cell
def _():
    import marimo as mo

    return (mo,)


@app.cell
def _(mo):
    mo.md(r"""
    # Exclusive Self Attention (XSA)
    ### *Standard attention wastes capacity talking to itself. XSA fixes this with two lines of code.*

    ---

    **Paper:** [Exclusive Self Attention - Shuangfei Zhai, Apple (2026)](https://arxiv.org/abs/2603.09078)
     · 
    **Competition:** alphaXiv × marimo Notebook Competition

    This notebook walks through the core idea of XSA from first principles:
    diagnose the problem, understand the fix, compare training dynamics, and explore
    both sequence-length scaling and a tunable partial-XSA variant.
    """)
    return


@app.cell
def _():
    import sys
    import os
    from pathlib import Path

    NOTEBOOK_DIR = Path(__file__).parent
    REPO_PATH    = str(NOTEBOOK_DIR / "EXA_repo")
    CACHE_DIR    = NOTEBOOK_DIR / "cache"
    CACHE_DIR.mkdir(exist_ok=True)

    if REPO_PATH not in sys.path:
        sys.path.insert(0, REPO_PATH)

    import torch
    import torch.nn.functional as F
    import numpy as np
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt
    import plotly.graph_objects as go
    from plotly.subplots import make_subplots
    import plotly.io as pio
    pio.templates.default = "plotly_white"

    SA_COL  = "#E74C3C"
    XSA_COL = "#2471A3"
    C1      = "#E74C3C"
    C2      = "#E67E22"
    C3      = "#7D3C98"

    if torch.backends.mps.is_available():
        DEVICE = torch.device("mps")
        DEVICE_NAME = "Apple MPS"
    elif torch.cuda.is_available():
        DEVICE = torch.device("cuda")
        DEVICE_NAME = "CUDA"
    else:
        DEVICE = torch.device("cpu")
        DEVICE_NAME = "CPU"
    return (
        C1,
        C2,
        C3,
        CACHE_DIR,
        DEVICE,
        DEVICE_NAME,
        REPO_PATH,
        SA_COL,
        XSA_COL,
        go,
        make_subplots,
        np,
        plt,
        torch,
    )


@app.cell
def _(DEVICE_NAME, REPO_PATH, mo):
    mo.md(f"""
    **Device:** `{DEVICE_NAME}`  ·  **Repo:** `{REPO_PATH}`
    """)
    return


@app.cell
def _(mo):
    mo.md(r"""
    ---
    ## Part 1 - The Problem: Attention Similarity Bias

    A transformer attention layer computes:

    $$y_i = \sum_j a_{ij} v_j$$

    where $a_{ij}$ are attention weights and $v_j$ are value vectors.

    **The hidden problem:** value vectors within a sequence tend to be positively
    correlated. Combined with non-trivial self-attention weights $a_{ii}$, the
    output $y_i$ ends up highly aligned with the *self-value* $v_i$ -
    the token's own value vector.

    This means attention is partly duplicating work the FFN already does via
    the residual path. The paper calls this the **attention similarity bias**.

    The chart below mirrors the paper's Figure 1 diagnostic at toy scale: three
    quantities measured layer by layer on an untrained 2-layer model.
    """)
    return


@app.cell
def _(C1, C2, C3, DEVICE, REPO_PATH, go, make_subplots, mo):
    from train_utils import (
        make_model, load_tinystories, make_loaders, compute_bias_metrics
    )

    _probe   = make_model("standard", REPO_PATH)
    _ds0     = load_tinystories(seq_len=64, max_tokens=80_000)
    _, _vl0  = make_loaders(_ds0, batch_size=16)
    _m0      = compute_bias_metrics(_probe, _vl0, DEVICE, n_batches=8)

    _layers  = _m0.layer_indices
    _specs   = [
        ("Avg value-vector cosine sim  (vᵢ, vⱼ)  i≠j",
         _m0.value_cosine_sim,  C1),
        ("Avg diagonal attention weight  aᵢᵢ",
         _m0.self_attn_weight,  C2),
        ("Avg output–self-value cosine sim  ⟨yᵢ, vᵢ⟩  ← the bias",
         _m0.preproj_output_self_cosim, C3),
    ]

    _fig1 = make_subplots(
        rows=1, cols=3,
        subplot_titles=[s[0] for s in _specs],
        horizontal_spacing=0.10,
    )

    for _ci, (_title, _vals, _col) in enumerate(_specs, start=1):
        _fig1.add_trace(go.Scatter(
            x=_layers, y=_vals,
            mode="lines+markers+text",
            line=dict(color=_col, width=3),
            marker=dict(size=11, color=_col),
            text=[f"{v:.3f}" for v in _vals],
            textposition="top center",
            textfont=dict(color=_col, size=11),
            showlegend=False,
            hovertemplate="Layer %{x}<br>%{y:.4f}<extra></extra>",
        ), row=1, col=_ci)
        _fig1.update_xaxes(tickvals=_layers, title_text="Layer index",
                           row=1, col=_ci)
        _fig1.update_yaxes(rangemode="tozero", gridcolor="#eee", row=1, col=_ci)

    _fig1.update_layout(
        title=dict(
            text="Reproducing Paper Figure 1 - Attention Similarity Bias (untrained 2-layer model)",
            font=dict(size=13), x=0.5,
        ),
        height=360,
        margin=dict(t=90, b=50, l=50, r=30),
        plot_bgcolor="white",
    )

    bias_fig_untrained = _fig1
    mo.ui.plotly(bias_fig_untrained)
    return compute_bias_metrics, load_tinystories, make_loaders, make_model


@app.cell
def _(mo):
    mo.md(r"""
    Even in an *untrained* model the output $y_i$ already has non-trivial
    cosine similarity with $v_i$ - and it grows with depth.
    After training this gets significantly worse (we confirm in Part 5).

    **Key intuition:** the residual connection already carries $v_i$ to the FFN
    unchanged. If attention output also aligns with $v_i$, two paths say the
    same thing - one is wasted.
    """)
    return


@app.cell
def _(mo):
    mo.md(r"""
    ---
    ## Part 2 - The Fix: Two Lines of Code

    XSA removes the self-value component from the attention output before
    the output projection.

    ### Standard Self Attention
    ```python
    Y   = scaled_dot_product_attention(Q, K, V)   # (B, H, T, Dh)
    out = Y @ Wo
    ```

    ### Exclusive Self Attention (XSA)
    ```python
    Y   = scaled_dot_product_attention(Q, K, V)
    Vn  = F.normalize(V, dim=-1)                   # ← line 1
    Z   = Y - (Y * Vn).sum(-1, keepdim=True) * Vn  # ← line 2
    out = Z @ Wo
    ```

    $Z$ is the component of $Y$ **orthogonal to $V$**. The diagram below shows
    what this means geometrically.

    XSA does not remove the token's own information from the network. It only
    removes the component of the attention output aligned with the self-value;
    the residual path still carries the token representation forward.
    """)
    return


@app.cell
def _(mo):
    lam_slider = mo.ui.slider(
        start=0.0, stop=1.0, value=1.0, step=0.05,
        label="\u03bb (projection strength)"
    )
    lam_slider
    return (lam_slider,)


@app.cell
def _(lam_slider, mo, np, plt):
    _fig, _axes = plt.subplots(1, 2, figsize=(11, 4.2))

    def _draw(ax, title, lam=1.0):
        ax.set_xlim(-0.2, 1.7); ax.set_ylim(-0.25, 1.35)
        ax.set_aspect("equal"); ax.axis("off")
        ax.set_title(title, fontsize=12, fontweight="bold", pad=10)

        v    = np.array([1.0, 0.0])
        y    = np.array([0.6, 0.9])
        proj = np.dot(y, v) * v

        ax.annotate("", xy=v, xytext=[0, 0],
                    arrowprops=dict(arrowstyle="->", color="#e74c3c", lw=2.5))
        ax.text(1.06, -0.13, "$V$  (self-value)", color="#e74c3c", fontsize=11)

        ax.annotate("", xy=y, xytext=[0, 0],
                    arrowprops=dict(arrowstyle="->", color="#2980b9", lw=2.5))
        ax.text(y[0] + 0.05, y[1] + 0.04, "$Y$  (attn output)",
                color="#2980b9", fontsize=11)

        ax.plot([proj[0], proj[0]], [0, proj[1]], "--", color="#bbb", lw=1.2)
        ax.plot([y[0], proj[0]], [y[1], 0], "--", color="#bbb", lw=1.2)
        ax.annotate("", xy=proj, xytext=[0, 0],
                    arrowprops=dict(arrowstyle="->", color="#e74c3c", lw=1.4, alpha=0.45))

        if lam > 0:
            removed = lam * proj
            ax.annotate("", xy=removed, xytext=[0, 0],
                        arrowprops=dict(arrowstyle="->", color="#e74c3c", lw=1.4, alpha=0.45))
            ax.text(proj[0] / 2 - 0.2, -0.19,
                    f"$\\lambda \\cdot \\alpha V_n$  (\u03bb={lam:.2f})",
                    color="#e74c3c", fontsize=9, alpha=0.7)

            z = y - lam * proj
            ax.annotate("", xy=z, xytext=[0, 0],
                        arrowprops=dict(arrowstyle="->", color="#27ae60", lw=2.5))
            orth_label = "orthogonal!" if abs(lam - 1.0) < 0.01 else f"(\u03bb={lam:.2f})"
            ax.text(z[0] - 0.38, z[1] + 0.04,
                    f"$Z_{{\\lambda}} = Y - \\lambda \\alpha V_n$\n{orth_label}",
                    color="#27ae60", fontsize=11)
            s = 0.07
            c = proj + s * z / np.linalg.norm(z)
            ax.plot([c[0], c[0] - s * v[0]], [c[1], c[1] - s * v[1]],
                    color="#27ae60", lw=1.2)
            ax.plot([c[0] - s * v[0], proj[0] - s * v[0]],
                    [c[1] - s * v[1], -s * v[1]], color="#27ae60", lw=1.2)

    _draw(_axes[0], "Standard SA \u2014 $Y$ contains $V$-component", lam=0.0)
    _draw(_axes[1], "XSA \u2014 $Z_\\lambda$ with adjustable $\\lambda$", lam=lam_slider.value)
    _fig.tight_layout()

    mo.md(f"$Z_{{\\lambda}} = Y - \\lambda \\cdot \\alpha V_n$ &nbsp; with $\\lambda$ = **{lam_slider.value:.2f}**")
    _fig
    return


@app.cell
def _(mo):
    mo.md(r"""
    At $\lambda = 1$, $Z$ is orthogonal to $V$ by construction, always, not just
    approximately. At intermediate values, partial removal still reduces the
    redundant component. Use the slider above to see how $\lambda$ controls the
    projection strength.

    **No extra parameters. Negligible overhead. Just a projection subtraction.**
    """)
    return


@app.cell
def _(mo):
    mo.md(r"""
    ---
    ## Part 3 - Live Training: SA vs XSA

    Both models train side-by-side starting from **identical weights**.
    Any difference in loss is purely due to the XSA mechanism.

    - **Model:** 2-layer transformer, 128-dim, 2 heads
    - **Data:** TinyStories (real English text, GPT-2 tokeniser)
    - **Steps:** adjustable below  (~2–3 min for 300 steps on Apple Silicon MPS)
    """)
    return


@app.cell
def _(mo):
    steps_slider = mo.ui.slider(start=100, stop=500, value=300, step=50,
                                 label="Training steps")
    train_button = mo.ui.run_button(label="▶  Train SA vs XSA")
    mo.vstack([steps_slider, train_button])
    return steps_slider, train_button


@app.cell
def _(
    CACHE_DIR,
    DEVICE,
    REPO_PATH,
    SA_COL,
    XSA_COL,
    go,
    load_tinystories,
    make_loaders,
    make_subplots,
    mo,
    steps_slider,
    torch,
    train_button,
):
    from train_utils import train_pair, make_model_pair

    mo.stop(not train_button.value,
            mo.md("*Click the button above to start training.*"))

    _dataset = load_tinystories(seq_len=64, max_tokens=500_000)
    _train_loader, _val_loader = make_loaders(_dataset, batch_size=16)

    _model_sa, _model_xsa, _ = make_model_pair(REPO_PATH, seq_len=64, proj_strength_xsa=1.0)

    _n_params = sum(p.numel() for p in _model_sa.parameters())

    _result = train_pair(
        _model_sa, _model_xsa,
        _train_loader, _val_loader,
        n_steps=steps_slider.value, lr=1e-3, eval_every=20,
    )

    torch.save(_model_sa.state_dict(),  CACHE_DIR / f"model_sa_steps{steps_slider.value}.pt")
    torch.save(_model_xsa.state_dict(), CACHE_DIR / f"model_xsa_steps{steps_slider.value}.pt")

    import pickle
    with open(CACHE_DIR / f"train_result_steps{steps_slider.value}.pkl", "wb") as _f:
        pickle.dump(_result, _f)

    _fig3 = make_subplots(
        rows=1, cols=2,
        subplot_titles=["Training Loss", "Validation Loss"],
        horizontal_spacing=0.10,
    )

    for _ci3, (_sa, _xs) in enumerate([
        (_result.train_losses_sa, _result.train_losses_xsa),
        (_result.val_losses_sa,   _result.val_losses_xsa),
    ], start=1):
        _show3 = _ci3 == 1

        _fig3.add_trace(go.Scatter(
            x=_result.steps, y=_sa,
            name="Standard SA", legendgroup="sa",
            line=dict(color=SA_COL, width=2.5),
            showlegend=_show3,
            hovertemplate="Step %{x}<br>Loss: %{y:.4f}<extra>Standard SA</extra>",
        ), row=1, col=_ci3)

        _fig3.add_trace(go.Scatter(
            x=_result.steps, y=_xs,
            name="XSA", legendgroup="xsa",
            line=dict(color=XSA_COL, width=2.5),
            showlegend=_show3,
            hovertemplate="Step %{x}<br>Loss: %{y:.4f}<extra>XSA</extra>",
        ), row=1, col=_ci3)

        _fig3.add_trace(go.Scatter(
            x=_result.steps + _result.steps[::-1],
            y=_sa + _xs[::-1],
            fill="toself", fillcolor=XSA_COL,
            opacity=0.08, line=dict(width=0),
            name="XSA advantage", legendgroup="adv",
            showlegend=_show3, hoverinfo="skip",
        ), row=1, col=_ci3)

        _fig3.update_xaxes(title_text="Step", row=1, col=_ci3)
        _fig3.update_yaxes(title_text="Cross-Entropy Loss",
                           gridcolor="#eee", row=1, col=_ci3)

    _fig3.update_layout(
        height=420,
        legend=dict(orientation="h", yanchor="bottom", y=1.02,
                    xanchor="right", x=1),
        margin=dict(t=70, b=50, l=60, r=30),
        plot_bgcolor="white",
        hovermode="x unified",
    )

    _fsa = _result.val_losses_sa[-1]
    _fxs = _result.val_losses_xsa[-1]
    _imp = (_fsa - _fxs) / _fsa * 100

    mo.vstack([
        mo.ui.plotly(_fig3),
        mo.md(f"""
    **Results** ({_result.elapsed_sec:.0f}s on `{DEVICE}`, {_result.steps[-1]} steps, {_n_params/1e6:.1f}M params)

    | | Standard SA | XSA | Improvement |
    |---|---|---|---|
    | Final val loss | `{_fsa:.4f}` | `{_fxs:.4f}` | **+{_imp:.1f}%** |

    In this run, XSA converges to lower loss with identical architecture and identical starting
    weights. The shaded region shows the growing XSA advantage - matching the
    qualitative result of **Paper Figure 3**.
        """),
    ])
    return make_model_pair, pickle, train_pair


@app.cell
def _(mo):
    mo.md(r"""
    ---
    ## Part 4 - Extension: The Benefit Grows with Sequence Length

    **Paper Figure 5** shows XSA's advantage increases with longer sequences.
    Longer contexts mean more tokens, so the self-attention bias accumulates
    more and XSA's constraint becomes more valuable.

    We reproduce this trend at small scale (64 / 128 / 256 tokens).
    Results are cached after the first run.
    """)
    return


@app.cell
def _(mo):
    seq_button = mo.ui.run_button(
        label="▶  Run sequence-length experiment  (64 / 128 / 256)"
    )
    seq_button
    return (seq_button,)


@app.cell
def _(
    CACHE_DIR,
    REPO_PATH,
    SA_COL,
    XSA_COL,
    go,
    load_tinystories,
    make_loaders,
    make_model_pair,
    make_subplots,
    mo,
    pickle,
    seq_button,
    train_pair,
):
    mo.stop(not seq_button.value,
            mo.md("*Click above to run. Results are cached - re-running is instant.*"))

    _cache = CACHE_DIR / "seq_len_results.pkl"

    if _cache.exists():
        with open(_cache, "rb") as _f:
            _sr = pickle.load(_f)
        print("Loaded cached seq-len results")
    else:
        _sr = {}
        for _sl in [64, 128, 256]:
            print(f"  seq_len={_sl} ...")
            _msa, _mxs, _ = make_model_pair(REPO_PATH)
            _ds = load_tinystories(seq_len=_sl, max_tokens=150_000)
            _tl, _vl = make_loaders(_ds, batch_size=8)
            _r = train_pair(_msa, _mxs, _tl, _vl,
                            n_steps=500, lr=1e-3, eval_every=300)
            _sr[_sl] = {"sa": _r.val_losses_sa[-1], "xsa": _r.val_losses_xsa[-1]}
        with open(_cache, "wb") as _f:
            pickle.dump(_sr, _f)

    _sls  = sorted(_sr.keys())
    _sa_l = [_sr[s]["sa"]  for s in _sls]
    _xs_l = [_sr[s]["xsa"] for s in _sls]
    _imps = [(_sr[s]["sa"] - _sr[s]["xsa"]) / _sr[s]["sa"] * 100 for s in _sls]

    _fig4 = make_subplots(
        rows=1, cols=2,
        subplot_titles=[
            "Validation Loss vs Sequence Length  (Paper Figure 5)",
            "XSA Improvement  (%)",
        ],
        horizontal_spacing=0.12,
    )

    _fig4.add_trace(go.Scatter(
        x=_sls, y=_sa_l, name="Standard SA",
        mode="lines+markers",
        line=dict(color=SA_COL, width=2.5),
        marker=dict(size=10),
        hovertemplate="seq=%{x}<br>loss=%{y:.4f}<extra>Standard SA</extra>",
    ), row=1, col=1)

    _fig4.add_trace(go.Scatter(
        x=_sls, y=_xs_l, name="XSA",
        mode="lines+markers",
        line=dict(color=XSA_COL, width=2.5),
        marker=dict(size=10),
        hovertemplate="seq=%{x}<br>loss=%{y:.4f}<extra>XSA</extra>",
    ), row=1, col=1)

    _fig4.add_trace(go.Scatter(
        x=_sls + _sls[::-1],
        y=_sa_l + _xs_l[::-1],
        fill="toself", fillcolor=XSA_COL,
        opacity=0.08, line=dict(width=0),
        showlegend=False, hoverinfo="skip",
    ), row=1, col=1)

    _bar_cols = [XSA_COL if v >= 0 else SA_COL for v in _imps]
    _fig4.add_trace(go.Bar(
        x=[str(s) for s in _sls],
        y=_imps,
        marker_color=_bar_cols,
        marker_line_width=0,
        text=[f"{v:+.2f}%" for v in _imps],
        textposition="outside",
        textfont=dict(size=12, color=_bar_cols),
        showlegend=False,
        hovertemplate="seq=%{x}<br>improvement=%{y:.2f}%<extra></extra>",
    ), row=1, col=2)

    _fig4.update_xaxes(title_text="Sequence length (tokens)",
                       tickvals=_sls, row=1, col=1)
    _fig4.update_xaxes(title_text="Sequence length (tokens)", row=1, col=2)
    _fig4.update_yaxes(title_text="Validation loss",
                       gridcolor="#eee", row=1, col=1)
    _fig4.update_yaxes(title_text="Improvement (%)", gridcolor="#eee",
                       zeroline=True, zerolinecolor="#aaa", zerolinewidth=1,
                       row=1, col=2)

    _fig4.update_layout(
        height=420,
        legend=dict(orientation="h", yanchor="bottom", y=1.05,
                    xanchor="left", x=0),
        margin=dict(t=80, b=60, l=60, r=30),
        plot_bgcolor="white",
    )

    _rows = "\n".join(
        f"| {s} | `{_sr[s]['sa']:.4f}` | `{_sr[s]['xsa']:.4f}` | **{i:+.2f}%** |"
        for s, i in zip(_sls, _imps)
    )

    mo.vstack([
        mo.ui.plotly(_fig4),
        mo.md(f"""
    | Seq len | SA loss | XSA loss | Improvement |
    |---|---|---|---|
    {_rows}

    In this toy-scale run, XSA improves validation loss across the tested sequence
    lengths, though the pattern is much noisier than in the large-scale result
    reported in the paper. The paper demonstrates a clear monotonic trend at 1.3B
    parameters (Figure 5), where longer contexts give XSA more opportunity to
    eliminate the self-value bias.
    """),
    ])
    return


@app.cell
def _(mo):
    mo.md(r"""
    ---
    ## Part 5 - Before vs After: Bias Elimination

    Let's return to Part 1's diagnostic and compare three states:
    **untrained** (gray), **trained SA** (red), and **trained XSA** (blue).

    The untrained baseline shows that some output–self alignment already exists at
    initialization. Training then amplifies this tendency, especially in standard
    self-attention. Interestingly, the trained XSA model can develop an even
    stronger raw output–self alignment before the projection step. The key
    effect of XSA is not to prevent that alignment from arising inside $Y$, but
    to remove it from the final attention output $Z$.

    We track **both** pre-projection alignment (standard attention output Y vs V)
    and post-projection alignment (XSA output Z vs V) to show how the projection
    removes the redundant self-value direction.
    """)
    return


@app.cell
def _(mo):
    bias_button = mo.ui.run_button(label="▶  Compute bias metrics on trained models")
    bias_button
    return (bias_button,)


@app.cell
def _(
    CACHE_DIR,
    DEVICE,
    REPO_PATH,
    SA_COL,
    XSA_COL,
    bias_button,
    compute_bias_metrics,
    go,
    load_tinystories,
    make_loaders,
    make_model,
    make_subplots,
    mo,
    torch,
):
    mo.stop(not bias_button.value,
            mo.md("*Run Part 3 first, then click above.*"))

    _ckpt_sa  = CACHE_DIR / "model_sa.pt"
    _ckpt_xsa = CACHE_DIR / "model_xsa.pt"

    # Fallback: also check step-suffixed names from newer runs
    if not (_ckpt_sa.exists() and _ckpt_xsa.exists()):
        _candidates = sorted(CACHE_DIR.glob("model_sa_steps*.pt"), reverse=True)
        if _candidates:
            _ckpt_sa = _candidates[0]
            _ckpt_xsa = CACHE_DIR / _ckpt_sa.name.replace("model_sa_", "model_xsa_")

    mo.stop(
        not (_ckpt_sa.exists() and _ckpt_xsa.exists()),
        mo.md("No checkpoints found - please run Part 3 first.")
    )

    _msa = make_model("standard", REPO_PATH)
    _mxs = make_model("xsa",      REPO_PATH)
    _msa.load_state_dict(torch.load(_ckpt_sa,  map_location=DEVICE))
    _mxs.load_state_dict(torch.load(_ckpt_xsa, map_location=DEVICE))

    _ds5     = load_tinystories(seq_len=64, max_tokens=80_000)
    _, _vl5  = make_loaders(_ds5, batch_size=16)

    _met_sa  = compute_bias_metrics(_msa, _vl5, DEVICE, n_batches=8)
    _met_xsa = compute_bias_metrics(_mxs, _vl5, DEVICE, n_batches=8)

    # Untrained baseline: shows bias exists at init and grows with training
    _mu = make_model("standard", REPO_PATH)
    _met_un = compute_bias_metrics(_mu, _vl5, DEVICE, n_batches=8)

    _GRAY = "#999999"
    _layers5 = _met_sa.layer_indices
    _triples = [
        ("Avg value-vector cosine sim  (vᵢ, vⱼ)  i≠j",
         _met_sa.value_cosine_sim,  _met_xsa.value_cosine_sim,
         _met_un.value_cosine_sim),
        ("Avg diagonal attention weight  aᵢᵢ",
         _met_sa.self_attn_weight,  _met_xsa.self_attn_weight,
         _met_un.self_attn_weight),
        ("Pre-projection ⟨Yᵢ, vᵢ⟩  (raw attention bias)",
         _met_sa.preproj_output_self_cosim, _met_xsa.preproj_output_self_cosim,
         _met_un.preproj_output_self_cosim),
        ("Post-projection ⟨Zᵢ, vᵢ⟩  (final output bias)  ← key metric",
         _met_sa.postproj_output_self_cosim, _met_xsa.postproj_output_self_cosim,
         _met_un.postproj_output_self_cosim),
    ]

    _fig5 = make_subplots(
        rows=2, cols=2,
        subplot_titles=[t for t, *_ in _triples],
        horizontal_spacing=0.12,
        vertical_spacing=0.22,
    )

    for _idx5, (_title5, _dsa, _dxs, _dun) in enumerate(_triples):
        _row5 = _idx5 // 2 + 1
        _col5 = _idx5 % 2 + 1
        _show5 = _idx5 == 0

        _fig5.add_trace(go.Scatter(
            x=_layers5, y=_dsa,
            name="SA (trained)", legendgroup="sa5",
            mode="lines+markers+text",
            line=dict(color=SA_COL, width=2.5),
            marker=dict(size=10, symbol="circle"),
            text=[f"{v:.3f}" for v in _dsa],
            textposition="top right",
            textfont=dict(color=SA_COL, size=10),
            showlegend=_show5,
            hovertemplate="Layer %{x}<br>%{y:.4f}<extra>SA</extra>",
        ), row=_row5, col=_col5)

        _fig5.add_trace(go.Scatter(
            x=_layers5, y=_dxs,
            name="XSA (trained)", legendgroup="xsa5",
            mode="lines+markers+text",
            line=dict(color=XSA_COL, width=2.5, dash="dash"),
            marker=dict(size=10, symbol="square"),
            text=[f"{v:.3f}" for v in _dxs],
            textposition="bottom right",
            textfont=dict(color=XSA_COL, size=10),
            showlegend=_show5,
            hovertemplate="Layer %{x}<br>%{y:.4f}<extra>XSA</extra>",
        ), row=_row5, col=_col5)

        _fig5.add_trace(go.Scatter(
            x=_layers5, y=_dun,
            name="Untrained", legendgroup="un5",
            mode="lines+markers",
            line=dict(color=_GRAY, width=1.8, dash="dot"),
            marker=dict(size=8, symbol="diamond", color=_GRAY),
            showlegend=_show5,
            hovertemplate="Layer %{x}<br>%{y:.4f}<extra>Untrained</extra>",
        ), row=_row5, col=_col5)

        _fig5.update_xaxes(tickvals=_layers5, title_text="Layer index",
                           row=_row5, col=_col5)
        _fig5.update_yaxes(rangemode="tozero", gridcolor="#eee",
                           row=_row5, col=_col5)

    _fig5.update_layout(
        title=dict(
            text="SA (red) vs XSA (blue) vs Untrained (gray) | Pre & Post Projection Bias",
            font=dict(size=12), x=0.5,
        ),
        height=680,
        legend=dict(orientation="h", yanchor="bottom", y=1.04,
                    xanchor="right", x=1),
        margin=dict(t=100, b=60, l=50, r=30),
        plot_bgcolor="white",
    )

    _sa_pre  = _met_sa.preproj_output_self_cosim
    _xs_pre  = _met_xsa.preproj_output_self_cosim
    _xs_post = _met_xsa.postproj_output_self_cosim
    _n5      = len(_layers5)

    mo.vstack([
        mo.ui.plotly(_fig5),
        mo.md(
            "**Pre vs Post projection: XSA can develop stronger raw bias, then remove it:**\n\n"
            "| Layer | SA \u27e8Y,v\u27e9 | XSA \u27e8Y,v\u27e9 (pre) | XSA \u27e8Z,v\u27e9 (post) | Bias removed |\n"
            "|---|---|---|---|---|\n"
            + "\n".join(
                f"| Layer {i} | `{_sa_pre[i]:.4f}` | `{_xs_pre[i]:.4f}`"
                f" | `{_xs_post[i]:.4f}` | **{(_xs_pre[i] - _xs_post[i]) / _xs_pre[i] * 100:.0f}%** |"
                for i in range(_n5)
            )
            + "\n\nInterestingly, the trained XSA model can develop **even stronger** raw "
              "output\u2013self alignment before projection (the pre-projection values are higher "
              "than standard SA). The key effect of XSA is not to prevent that alignment from "
              "arising inside $Y$, but to remove it from the final attention output $Z$. "
              "The gray dotted line (untrained) shows the bias exists at initialization and "
              "is amplified by training."
        ),
    ])
    return


@app.cell
def _(mo):
    mo.md(r"""
    ---
    ## Part 6 - Extension: How Much Projection Is Needed?

    XSA removes the self-value projection with strength $\lambda$:

    $$Z_\lambda = Y - \lambda \cdot \frac{(Y \cdot V_n)}{\|V_n\|^2} V_n$$

    - $\lambda = 0$: standard SA (no removal)
    - $\lambda = 1$: full XSA (complete removal)
    - $0 < \lambda < 1$: **Partial XSA**, tunable trade-off

    We sweep five values of $\lambda$ and plot the resulting loss landscape.
    """)
    return


@app.cell
def _(mo):
    lambda_button = mo.ui.run_button(
        label="▶  Run \u03bb-sweep  (0.0, 0.25, 0.5, 0.75, 1.0)  (~10 min)"
    )
    lambda_button
    return (lambda_button,)


@app.cell
def _(
    CACHE_DIR,
    DEVICE,
    REPO_PATH,
    XSA_COL,
    compute_bias_metrics,
    go,
    lambda_button,
    load_tinystories,
    make_loaders,
    make_model,
    make_subplots,
    mo,
    pickle,
    torch,
    train_pair,
):
    mo.stop(not lambda_button.value,
            mo.md("*Click the button above to run the \u03bb-sweep experiment.*"))

    _cache6 = CACHE_DIR / "lambda_sweep_results.pkl"

    if _cache6.exists():
        with open(_cache6, "rb") as _f:
            _sr6 = pickle.load(_f)
        print("Loaded cached \u03bb-sweep results")
    else:
        _sr6 = {}
        _lambdas = [0.0, 0.25, 0.5, 0.75, 1.0]
        _ds6 = load_tinystories(seq_len=64, max_tokens=500_000)
        _tl6, _vl6 = make_loaders(_ds6, batch_size=16)

        # Single fixed init across all lambda values
        _base_sa = make_model("standard", REPO_PATH)
        _base_sd = _base_sa.state_dict()

        for _lam in _lambdas:
            print(f"  \u03bb={_lam} ...")
            # Each lambda gets its own SA baseline + XSA model, all from same init
            _msa = make_model("standard", REPO_PATH)
            _mxs = make_model("xsa", REPO_PATH, proj_strength=_lam)
            _msa.load_state_dict(_base_sd)
            _mxs.load_state_dict(_base_sd)

            _r6 = train_pair(
                _msa, _mxs, _tl6, _vl6,
                n_steps=300, lr=1e-3, eval_every=20,
            )
            _sr6[_lam] = {
                "final_val_sa": _r6.val_losses_sa[-1],
                "final_val_xsa": _r6.val_losses_xsa[-1],
            }
            torch.save(_mxs.state_dict(), CACHE_DIR / f"model_xsa_lam{_lam}.pt")
            torch.save(_msa.state_dict(), CACHE_DIR / f"model_sa_lam{_lam}.pt")

        with open(_cache6, "wb") as _f:
            pickle.dump(_sr6, _f)

    _lams = sorted(_sr6.keys())

    # Per-lambda paired improvement (each lambda vs its own SA baseline)
    _paired_sa  = [_sr6[l]["final_val_sa"]  for l in _lams]
    _paired_xsa = [_sr6[l]["final_val_xsa"] for l in _lams]
    _paired_imp = [(s - x) / s * 100 for s, x in zip(_paired_sa, _paired_xsa)]

    # Bias metrics for each lambda
    _ds6b = load_tinystories(seq_len=64, max_tokens=80_000)
    _, _vl6b = make_loaders(_ds6b, batch_size=16)
    _bias_post_avg = []
    for _lam in _lams:
        _ck = CACHE_DIR / f"model_xsa_lam{_lam}.pt"
        if _ck.exists():
            _m = make_model("xsa", REPO_PATH, proj_strength=_lam)
            _m.load_state_dict(torch.load(_ck, map_location=DEVICE))
            _bm = compute_bias_metrics(_m, _vl6b, DEVICE, n_batches=3)
            _bias_post_avg.append(float(sum(_bm.postproj_output_self_cosim) / max(len(_bm.postproj_output_self_cosim), 1)))
        else:
            _bias_post_avg.append(float("nan"))

    _fig6 = make_subplots(
        rows=1, cols=2,
        subplot_titles=[
            "Paired Improvement over SA  (%)",
            "Post-projection Bias  (avg \u27e8Z,v\u27e9 across layers)",
        ],
        horizontal_spacing=0.12,
    )

    _fig6.add_trace(go.Scatter(
        x=_lams, y=_paired_imp,
        mode="lines+markers+text",
        line=dict(color=XSA_COL, width=2.5),
        marker=dict(size=10),
        text=[f"{v:+.2f}%" for v in _paired_imp],
        textposition="top center",
        textfont=dict(color=XSA_COL, size=11),
        showlegend=False,
        hovertemplate="\u03bb=%{x}<br>improvement=%{y:.2f}%<extra></extra>",
    ), row=1, col=1)

    _fig6.add_trace(go.Scatter(
        x=_lams, y=_bias_post_avg,
        mode="lines+markers+text",
        line=dict(color="#7D3C98", width=2.5),
        marker=dict(size=10, color="#7D3C98"),
        text=[f"{v:.4f}" for v in _bias_post_avg],
        textposition="top center",
        textfont=dict(color="#7D3C98", size=11),
        showlegend=False,
        hovertemplate="\u03bb=%{x}<br>bias=%{y:.4f}<extra></extra>",
    ), row=1, col=2)

    _fig6.update_xaxes(title_text="\u03bb (projection strength)",
                       tickvals=_lams, row=1, col=1)
    _fig6.update_xaxes(title_text="\u03bb (projection strength)",
                       tickvals=_lams, row=1, col=2)
    _fig6.update_yaxes(title_text="Improvement over paired SA (%)", gridcolor="#eee",
                       zeroline=True, zerolinecolor="#aaa", zerolinewidth=1,
                       row=1, col=1)
    _fig6.update_yaxes(title_text="Avg post-proj cosine sim", gridcolor="#eee",
                       row=1, col=2)

    _fig6.update_layout(
        height=420,
        legend=dict(orientation="h", yanchor="bottom", y=1.05,
                    xanchor="left", x=0),
        margin=dict(t=80, b=60, l=60, r=30),
        plot_bgcolor="white",
    )

    _rows6 = "\n".join(
        f"| \u03bb = {l} | `{_sr6[l]['final_val_sa']:.4f}` | `{_sr6[l]['final_val_xsa']:.4f}`"
        f" | **{imp:+.2f}%** |"
        for l, imp in zip(_lams, _paired_imp)
    )

    mo.vstack([
        mo.ui.plotly(_fig6),
        mo.md(
            "| \u03bb (proj strength) | Paired SA baseline | XSA loss | Improvement |\n"
            "|---|---|---|---|\n"
            + _rows6
            + "\n\nEach \u03bb uses its own paired SA baseline from the *same initialization*, "
            "so the comparison is fair. In this run, partial projection performs best, "
            "with \u03bb = 0.5 giving the largest improvement over its paired SA baseline. "
            "This suggests that at toy scale, the loss-optimal amount of self-value removal "
            "may be less than full projection, even though full XSA drives the "
            "post-projection bias closest to zero. "
            "The right panel shows how post-projection bias decreases monotonically with "
            "increasing \u03bb \u2014 the loss-optimal \u03bb may trade some residual bias for "
            "preserved signal."
        ),
    ])
    return


@app.cell
def _(mo):
    mo.md(r"""
    ---
    ## Summary

    | Part | Paper figure | Finding |
    |---|---|---|
    | **1. The problem** | Figure 1 | SA output cosine-sim with self-value grows with depth |
    | **2. The fix** | Algorithm 1 | Two lines remove the projection, negligible overhead |
    | **3. Live training** | Figure 3 | XSA converges to lower loss with identical weights |
    | **4. Scaling** | Figure 5 | XSA advantage grows with sequence length |
    | **5. Validation** | Figure 1 | XSA can develop stronger raw bias than SA, but removes it from the final output $Z$ |
    | **6. Partial XSA** | Novel extension | Loss-optimal $\lambda$ may be less than 1; partial removal trades bias for signal |

    ### Why it works

    The residual connection already carries $v_i$ to the FFN unchanged.
    Making the attention output orthogonal to $v_i$ enforces a clean
    **division of labour**: attention handles context, FFN handles the token.
    This is why the gain is free - it removes redundancy rather than adding capacity.

    The key effect of XSA is not to prevent self-value alignment from arising
    inside the raw attention output $Y$, but to remove it from the final attention
    output $Z$. The model can still learn strong internal representations; XSA
    just ensures the attention layer contributes information orthogonal to what
    the residual path already carries.

    ### The two lines

    ```python
    Vn = F.normalize(V, dim=-1)                        # unit self-value
    Z  = Y - (Y * Vn).sum(-1, keepdim=True) * Vn      # remove its projection
    ```

    ---

    **Paper:** Shuangfei Zhai, *Exclusive Self Attention*, Apple, 2026.
    [arxiv.org/abs/2603.09078](https://arxiv.org/abs/2603.09078)

    *Notebook built for the alphaXiv × marimo Notebook Competition.*
    """)
    return


if __name__ == "__main__":
    app.run()