Unify FFN and TrajMixer model architectures
This commit is contained in:
203
backbones.py
203
backbones.py
@@ -4,6 +4,12 @@ import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
from model_architectures import (
|
||||
TRAJ_MIXER_ARCHITECTURE,
|
||||
TRANSFORMER_FFN_ARCHITECTURE,
|
||||
resolve_model_architecture,
|
||||
)
|
||||
|
||||
|
||||
class TimeRoPE(nn.Module):
|
||||
def __init__(self, dim: int, base: float = 10000.0):
|
||||
@@ -213,7 +219,133 @@ class SwiGLU(nn.Module):
|
||||
return self.drop(self.w3(F.silu(self.w1(x)) * self.w2(x)))
|
||||
|
||||
|
||||
class GPTBlock(nn.Module):
|
||||
class TrajMixer(nn.Module):
|
||||
"""PreNorm gated mixing within and across latent trajectory groups.
|
||||
|
||||
The groups are contiguous partitions of the post-``W_O`` residual
|
||||
representation. They are deliberately not treated as attention heads.
|
||||
All operations are position-wise, so the sequence dimension remains fully
|
||||
parallel and no temporal information can leak between positions here.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
n_embd: int,
|
||||
n_head: int = 10,
|
||||
dropout: float = 0.0,
|
||||
):
|
||||
super().__init__()
|
||||
if n_embd <= 0:
|
||||
raise ValueError(f"n_embd must be > 0, got {n_embd}")
|
||||
if n_head <= 0:
|
||||
raise ValueError(f"n_head must be > 0, got {n_head}")
|
||||
if n_embd % n_head != 0:
|
||||
raise ValueError(
|
||||
f"n_embd must be divisible by n_head, got {n_embd} and {n_head}"
|
||||
)
|
||||
self.n_embd = n_embd
|
||||
self.n_group = n_head
|
||||
self.d_group = n_embd // n_head
|
||||
self.intra_hidden = 4 * self.d_group
|
||||
self.hidden_group = 4 * n_head
|
||||
|
||||
self.norm = nn.LayerNorm(self.n_embd)
|
||||
|
||||
self.intra_gate_proj = nn.Parameter(
|
||||
torch.empty(self.n_group, self.d_group, self.intra_hidden)
|
||||
)
|
||||
self.intra_value_proj = nn.Parameter(
|
||||
torch.empty(self.n_group, self.d_group, self.intra_hidden)
|
||||
)
|
||||
self.intra_output_proj = nn.Parameter(
|
||||
torch.empty(self.n_group, self.intra_hidden, self.d_group)
|
||||
)
|
||||
self.intra_gate_logits = nn.Parameter(
|
||||
torch.empty(self.n_group, self.d_group)
|
||||
)
|
||||
|
||||
self.gate_proj = nn.Parameter(
|
||||
torch.empty(self.d_group, self.n_group, self.hidden_group)
|
||||
)
|
||||
self.value_proj = nn.Parameter(
|
||||
torch.empty(self.d_group, self.n_group, self.hidden_group)
|
||||
)
|
||||
self.output_proj = nn.Parameter(
|
||||
torch.empty(self.d_group, self.hidden_group, self.n_group)
|
||||
)
|
||||
self.drop = nn.Dropout(dropout)
|
||||
self.reset_parameters()
|
||||
|
||||
def reset_parameters(self) -> None:
|
||||
for group_idx in range(self.n_group):
|
||||
nn.init.xavier_uniform_(self.intra_gate_proj[group_idx])
|
||||
nn.init.xavier_uniform_(self.intra_value_proj[group_idx])
|
||||
nn.init.xavier_uniform_(self.intra_output_proj[group_idx])
|
||||
nn.init.constant_(
|
||||
self.intra_gate_logits,
|
||||
math.log(0.1 / 0.9),
|
||||
)
|
||||
|
||||
for feature_idx in range(self.d_group):
|
||||
nn.init.xavier_uniform_(self.gate_proj[feature_idx])
|
||||
nn.init.xavier_uniform_(self.value_proj[feature_idx])
|
||||
nn.init.normal_(self.output_proj, mean=0.0, std=1e-3)
|
||||
|
||||
def _intra_mix(self, grouped: torch.Tensor) -> torch.Tensor:
|
||||
"""Mix features independently inside each residual-space group."""
|
||||
intra_gate = torch.einsum(
|
||||
"blgd,gdh->blgh", grouped, self.intra_gate_proj
|
||||
)
|
||||
intra_value = torch.einsum(
|
||||
"blgd,gdh->blgh", grouped, self.intra_value_proj
|
||||
)
|
||||
intra_hidden = F.silu(intra_gate) * intra_value
|
||||
return torch.einsum(
|
||||
"blgh,ghd->blgd", intra_hidden, self.intra_output_proj
|
||||
)
|
||||
|
||||
def _cross_mix(self, grouped: torch.Tensor) -> torch.Tensor:
|
||||
"""Mix groups independently for each within-group coordinate."""
|
||||
gate = torch.einsum(
|
||||
"blgr,rgh->blhr", grouped, self.gate_proj
|
||||
)
|
||||
value = torch.einsum(
|
||||
"blgr,rgh->blhr", grouped, self.value_proj
|
||||
)
|
||||
hidden = F.silu(gate) * value
|
||||
return torch.einsum(
|
||||
"blhr,rhg->blgr", hidden, self.output_proj
|
||||
)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
"""Apply one full-width PreNorm and one outer residual update."""
|
||||
if x.ndim != 3:
|
||||
raise ValueError(
|
||||
f"TrajMixer expects a 3D tensor, got shape {tuple(x.shape)}"
|
||||
)
|
||||
if x.size(-1) != self.n_embd:
|
||||
raise ValueError(
|
||||
f"Expected hidden size {self.n_embd}, got {x.size(-1)}"
|
||||
)
|
||||
|
||||
batch_size, seq_len, _ = x.shape
|
||||
grouped = self.norm(x).reshape(
|
||||
batch_size, seq_len, self.n_group, self.d_group
|
||||
)
|
||||
|
||||
intra_output = self._intra_mix(grouped)
|
||||
intra_gate = torch.sigmoid(self.intra_gate_logits).view(
|
||||
1, 1, self.n_group, self.d_group
|
||||
)
|
||||
mixed_input = grouped + intra_gate * intra_output
|
||||
|
||||
update = self._cross_mix(mixed_input).reshape(
|
||||
batch_size, seq_len, self.n_embd
|
||||
)
|
||||
return x + self.drop(update)
|
||||
|
||||
|
||||
class TransformerFFNBlock(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
n_embd: int,
|
||||
@@ -250,6 +382,75 @@ class GPTBlock(nn.Module):
|
||||
return x
|
||||
|
||||
|
||||
class TrajMixerBlock(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
n_embd: int,
|
||||
n_head: int,
|
||||
attn_dropout: float = 0.0,
|
||||
mlp_dropout: float = 0.0,
|
||||
use_time_rope: bool = False,
|
||||
use_rbf_bias: bool = False,
|
||||
n_rbf_bases: int = 16,
|
||||
):
|
||||
super().__init__()
|
||||
self.attn = TemporalAttention(
|
||||
n_embd=n_embd,
|
||||
n_head=n_head,
|
||||
n_rbf_bases=n_rbf_bases,
|
||||
dropout=attn_dropout,
|
||||
use_time_rope=use_time_rope,
|
||||
use_rbf_bias=use_rbf_bias,
|
||||
)
|
||||
self.mlp = TrajMixer(
|
||||
n_embd=n_embd,
|
||||
n_head=n_head,
|
||||
dropout=mlp_dropout,
|
||||
)
|
||||
self.ln1 = nn.LayerNorm(n_embd)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
x: torch.Tensor,
|
||||
rope_cache: tuple[torch.Tensor, torch.Tensor] | None = None,
|
||||
rbf_cache: torch.Tensor | None = None,
|
||||
attn_mask: torch.Tensor | None = None,
|
||||
) -> torch.Tensor:
|
||||
x = x + self.attn(self.ln1(x), rope_cache, rbf_cache, attn_mask)
|
||||
return self.mlp(x)
|
||||
|
||||
|
||||
def build_backbone_block(
|
||||
model_architecture: str,
|
||||
*,
|
||||
n_embd: int,
|
||||
n_head: int,
|
||||
attn_dropout: float = 0.0,
|
||||
mlp_dropout: float = 0.0,
|
||||
use_time_rope: bool = False,
|
||||
use_rbf_bias: bool = False,
|
||||
n_rbf_bases: int = 16,
|
||||
) -> nn.Module:
|
||||
"""Build one history block for a supported model architecture."""
|
||||
architecture = resolve_model_architecture(model_architecture)
|
||||
block_class: type[nn.Module]
|
||||
if architecture == TRANSFORMER_FFN_ARCHITECTURE:
|
||||
block_class = TransformerFFNBlock
|
||||
elif architecture == TRAJ_MIXER_ARCHITECTURE:
|
||||
block_class = TrajMixerBlock
|
||||
else: # pragma: no cover - guarded by resolve_model_architecture.
|
||||
raise ValueError(f"Unsupported model architecture: {architecture!r}")
|
||||
return block_class(
|
||||
n_embd=n_embd,
|
||||
n_head=n_head,
|
||||
attn_dropout=attn_dropout,
|
||||
mlp_dropout=mlp_dropout,
|
||||
use_time_rope=use_time_rope,
|
||||
use_rbf_bias=use_rbf_bias,
|
||||
n_rbf_bases=n_rbf_bases,
|
||||
)
|
||||
|
||||
|
||||
class TokenAutoDiscretization(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
|
||||
Reference in New Issue
Block a user