Upload folder using huggingface_hub

2c93889 verified 2 months ago

5.66 kB

"""LinearNO attention block (Stage 2 / Gate B).
Reimplemented from the equations of:
    "Transolver is a Linear Transformer: Revisiting Physics-Attention through the Lens of
     Linear Attention" — Hu, Liu, Qiao, Sun, Dou (NUDT), AAAI 2026, arXiv:2511.06294.
Reconstructed from the equations BEFORE the official code was released (github.com/HiPRL/LinearNO,
Jan 2026); we did not consult it. A post-hoc check confirms structural agreement (softmax axes,
associativity). Differences from the official Elasticity block: it adds a learnable per-head softmax
temperature (init 0.5, clamped [0.01,1]) that we omit, and uses one shared in_project_x lift + small
per-head q/k/v maps vs our full C->inner projections. Only the attention block differs from Transolver.
Core idea: Physics-Attention is the special case of linear attention
``Attention(Q,K,V) ~ phi(Q) (psi^T(K) V)`` in which (a) phi and psi come from the SAME linear
layer (differing only by normalization) and (b) there is an extra slice self-attention step.
LinearNO removes BOTH constraints:
  1. learn Q and K projections independently (break weight sharing), and
  2. drop the slice self-attention (identity).
  LinearNO(H) = phi(Q) @ ( psi(K)^T @ V )
  Q = H Wq ; K = H Wk ; V = Linear_V(H)
  phi(Q) = softmax_over_M( Linear_Q(Q) )   # (N, M)  rows sum to 1  <- softmax along M
  psi(K) = softmax_over_N( Linear_K(K) )   # (N, M)  cols sum to 1  <- softmax along N
### THE make-or-break detail (paper Table 6, Elasticity):
    phi softmax over M (slices), psi softmax over N (points) -> 0.0050  (correct)
    swapping these dims -> 0.0081..0.0112 (wrong). If LinearNO lands at 0.008-0.011, the
    softmax dimensions are almost certainly swapped.
Associativity: compute ``psi^T V`` first (M x d), then ``phi @ (...)`` -> O(N*M*d), linear in N.
Two variants (see docs/RECONCILIATION.md for the parameter-count tension):
  - "independent" (paper skeleton, default): separate Wq, Wk, Wv as full dim->inner
    projections. Most literal to the plan's reference code. ~0.85M params (> baseline).
  - "shared_qk": share one dim->inner base for the slice projections (phi, psi), keep a
    separate dim->inner V; breaks slice-weight sharing via two separate small slice layers.
    ~0.72M params (~= baseline). Use this to satisfy the Gate-B "<= baseline params" constraint.
An optional ``project_out=False`` drops the output projection (folds the per-head concat
directly), reaching ~0.59M (the plan's stated LinearNO size).
"""
from __future__ import annotations

import torch
import torch.nn as nn
from einops import rearrange


class LinearNO(nn.Module):
    def __init__(
        self,
        dim,
        heads=8,
        dim_head=16,
        slice_num=64,
        dropout=0.0,
        variant: str = "independent",
        project_out: bool = True,
        temperature: bool = False,
    ):
        super().__init__()
        inner = heads * dim_head
        self.h, self.m, self.dh = heads, slice_num, dim_head
        self.variant = variant
        self.project_out = project_out
        self.temperature = temperature
        if temperature:
            # Matches the official LinearNO `temp` block: learnable per-head temperature on both
            # softmaxes, init 0.5, clamped to [0.01, 1] (github.com/HiPRL/LinearNO).
            self.temp_q = nn.Parameter(torch.ones(1, heads, 1, 1) * 0.5)
            self.temp_k = nn.Parameter(torch.ones(1, heads, 1, 1) * 0.5)

        if variant == "independent":
            # Independent Q, K, V projections (paper modification 1, literal).
            self.to_q = nn.Linear(dim, inner, bias=False)
            self.to_k = nn.Linear(dim, inner, bias=False)
            self.to_v = nn.Linear(dim, inner, bias=False)
        elif variant == "shared_qk":
            # Share one base projection for the slice space (Q==K base), keep V separate.
            self.to_qk = nn.Linear(dim, inner, bias=False)
            self.to_v = nn.Linear(dim, inner, bias=False)
        else:
            raise ValueError(f"unknown variant {variant!r} (expected 'independent' or 'shared_qk')")

        # Asymmetric slice projections: query->slices and key->slices are SEPARATE weights.
        self.lin_q = nn.Linear(dim_head, slice_num)
        self.lin_k = nn.Linear(dim_head, slice_num)

        if project_out:
            self.to_out = nn.Sequential(nn.Linear(inner, dim), nn.Dropout(dropout))
        else:
            self.to_out = nn.Dropout(dropout)

    def _heads(self, t, B, N):
        return t.reshape(B, N, self.h, self.dh).permute(0, 2, 1, 3).contiguous()  # (B,H,N,dh)

    def forward(self, x):  # x: (B, N, C)
        B, N, C = x.shape
        if self.variant == "independent":
            q = self._heads(self.to_q(x), B, N)  # (B,H,N,dh)
            k = self._heads(self.to_k(x), B, N)
            v = self._heads(self.to_v(x), B, N)
        else:  # shared_qk
            base = self._heads(self.to_qk(x), B, N)
            q = base
            k = base
            v = self._heads(self.to_v(x), B, N)

        sq = self.lin_q(q)  # (B,H,N,M)
        sk = self.lin_k(k)  # (B,H,N,M)
        if self.temperature:
            sq = sq / self.temp_q.clamp(0.01, 1.0)
            sk = sk / self.temp_k.clamp(0.01, 1.0)
        phi = sq.softmax(dim=-1)  # softmax OVER M (slices)  <- rows sum to 1
        psi = sk.softmax(dim=-2)  # softmax OVER N (points)  <- cols sum to 1
        kv = torch.einsum("bhnm,bhnd->bhmd", psi, v)  # (B,H,M,dh)  cheap inner product first
        out = torch.einsum("bhnm,bhmd->bhnd", phi, kv)  # (B,H,N,dh)  linear in N
        out = rearrange(out, "b h n d -> b n (h d)")
        return self.to_out(out)

	"""LinearNO attention block (Stage 2 / Gate B).

	Reimplemented from the equations of:
	"Transolver is a Linear Transformer: Revisiting Physics-Attention through the Lens of
	Linear Attention" — Hu, Liu, Qiao, Sun, Dou (NUDT), AAAI 2026, arXiv:2511.06294.
	Reconstructed from the equations BEFORE the official code was released (github.com/HiPRL/LinearNO,
	Jan 2026); we did not consult it. A post-hoc check confirms structural agreement (softmax axes,
	associativity). Differences from the official Elasticity block: it adds a learnable per-head softmax
	temperature (init 0.5, clamped [0.01,1]) that we omit, and uses one shared in_project_x lift + small
	per-head q/k/v maps vs our full C->inner projections. Only the attention block differs from Transolver.

	Core idea: Physics-Attention is the special case of linear attention
	``Attention(Q,K,V) ~ phi(Q) (psi^T(K) V)`` in which (a) phi and psi come from the SAME linear
	layer (differing only by normalization) and (b) there is an extra slice self-attention step.
	LinearNO removes BOTH constraints:
	1. learn Q and K projections independently (break weight sharing), and
	2. drop the slice self-attention (identity).

	LinearNO(H) = phi(Q) @ ( psi(K)^T @ V )
	Q = H Wq ; K = H Wk ; V = Linear_V(H)
	phi(Q) = softmax_over_M( Linear_Q(Q) ) # (N, M) rows sum to 1 <- softmax along M
	psi(K) = softmax_over_N( Linear_K(K) ) # (N, M) cols sum to 1 <- softmax along N

	### THE make-or-break detail (paper Table 6, Elasticity):
	phi softmax over M (slices), psi softmax over N (points) -> 0.0050 (correct)
	swapping these dims -> 0.0081..0.0112 (wrong). If LinearNO lands at 0.008-0.011, the
	softmax dimensions are almost certainly swapped.

	Associativity: compute ``psi^T V`` first (M x d), then ``phi @ (...)`` -> O(NMd), linear in N.

	Two variants (see docs/RECONCILIATION.md for the parameter-count tension):
	- "independent" (paper skeleton, default): separate Wq, Wk, Wv as full dim->inner
	projections. Most literal to the plan's reference code. ~0.85M params (> baseline).
	- "shared_qk": share one dim->inner base for the slice projections (phi, psi), keep a
	separate dim->inner V; breaks slice-weight sharing via two separate small slice layers.
	~0.72M params (~= baseline). Use this to satisfy the Gate-B "<= baseline params" constraint.
	An optional ``project_out=False`` drops the output projection (folds the per-head concat
	directly), reaching ~0.59M (the plan's stated LinearNO size).
	"""
	from __future__ import annotations

	import torch
	import torch.nn as nn
	from einops import rearrange


	class LinearNO(nn.Module):
	def __init__(
	self,
	dim,
	heads=8,
	dim_head=16,
	slice_num=64,
	dropout=0.0,
	variant: str = "independent",
	project_out: bool = True,
	temperature: bool = False,
	):
	super().__init__()
	inner = heads * dim_head
	self.h, self.m, self.dh = heads, slice_num, dim_head
	self.variant = variant
	self.project_out = project_out
	self.temperature = temperature
	if temperature:
	# Matches the official LinearNO `temp` block: learnable per-head temperature on both
	# softmaxes, init 0.5, clamped to [0.01, 1] (github.com/HiPRL/LinearNO).
	self.temp_q = nn.Parameter(torch.ones(1, heads, 1, 1) * 0.5)
	self.temp_k = nn.Parameter(torch.ones(1, heads, 1, 1) * 0.5)

	if variant == "independent":
	# Independent Q, K, V projections (paper modification 1, literal).
	self.to_q = nn.Linear(dim, inner, bias=False)
	self.to_k = nn.Linear(dim, inner, bias=False)
	self.to_v = nn.Linear(dim, inner, bias=False)
	elif variant == "shared_qk":
	# Share one base projection for the slice space (Q==K base), keep V separate.
	self.to_qk = nn.Linear(dim, inner, bias=False)
	self.to_v = nn.Linear(dim, inner, bias=False)
	else:
	raise ValueError(f"unknown variant {variant!r} (expected 'independent' or 'shared_qk')")

	# Asymmetric slice projections: query->slices and key->slices are SEPARATE weights.
	self.lin_q = nn.Linear(dim_head, slice_num)
	self.lin_k = nn.Linear(dim_head, slice_num)

	if project_out:
	self.to_out = nn.Sequential(nn.Linear(inner, dim), nn.Dropout(dropout))
	else:
	self.to_out = nn.Dropout(dropout)

	def _heads(self, t, B, N):
	return t.reshape(B, N, self.h, self.dh).permute(0, 2, 1, 3).contiguous() # (B,H,N,dh)

	def forward(self, x): # x: (B, N, C)
	B, N, C = x.shape
	if self.variant == "independent":
	q = self._heads(self.to_q(x), B, N) # (B,H,N,dh)
	k = self._heads(self.to_k(x), B, N)
	v = self._heads(self.to_v(x), B, N)
	else: # shared_qk
	base = self._heads(self.to_qk(x), B, N)
	q = base
	k = base
	v = self._heads(self.to_v(x), B, N)

	sq = self.lin_q(q) # (B,H,N,M)
	sk = self.lin_k(k) # (B,H,N,M)
	if self.temperature:
	sq = sq / self.temp_q.clamp(0.01, 1.0)
	sk = sk / self.temp_k.clamp(0.01, 1.0)
	phi = sq.softmax(dim=-1) # softmax OVER M (slices) <- rows sum to 1
	psi = sk.softmax(dim=-2) # softmax OVER N (points) <- cols sum to 1
	kv = torch.einsum("bhnm,bhnd->bhmd", psi, v) # (B,H,M,dh) cheap inner product first
	out = torch.einsum("bhnm,bhmd->bhnd", phi, kv) # (B,H,N,dh) linear in N
	out = rearrange(out, "b h n d -> b n (h d)")
	return self.to_out(out)