Module neuralops.attentions.spectral_mamba
Spectral-Mamba Attention Module: Hybrid Attention-State Space Module with Spectral Channel Modulation – Production-Hardened.
Combines pretrained multi-head attention with a Spectral Mamba SSM branch via learned adaptive gating.
Supports: - Multi-architecture copy_weights (GPT-2, RoBERTa/BERT nested, Qwen/LLaMA, DistilBERT) - Trainable LoRA adapters with frozen base weights - GQA with repeat_kv - RoPE position embeddings - Causal masking for decoders - Nested encoder bypass - Frobenius norm verification (>= 1e-4) - Safe injection returning None on failure
Functions
def apply_rotary_pos_emb(q: torch.Tensor,
k: torch.Tensor,
cos: torch.Tensor,
sin: torch.Tensor,
position_ids: torch.Tensor | None = None) ‑> tuple[torch.Tensor, torch.Tensor]-
Expand source code
def apply_rotary_pos_emb( q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, position_ids: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: if cos is None or sin is None: return q, k if cos.dim() == 2: cos = cos.unsqueeze(0).unsqueeze(0) sin = sin.unsqueeze(0).unsqueeze(0) elif cos.dim() == 3: cos = cos.unsqueeze(1) sin = sin.unsqueeze(1) rotary_dim = cos.shape[-1] rotary_dim = rotary_dim - (rotary_dim % 2) q_rot, q_pass = q[..., :rotary_dim], q[..., rotary_dim:] k_rot, k_pass = k[..., :rotary_dim], k[..., rotary_dim:] cos = cos[..., :rotary_dim] sin = sin[..., :rotary_dim] q_rot_embed = (q_rot * cos) + (_rotate_half(q_rot) * sin) k_rot_embed = (k_rot * cos) + (_rotate_half(k_rot) * sin) q = torch.cat([q_rot_embed, q_pass], dim=-1) if q_pass.shape[-1] > 0 else q_rot_embed k = torch.cat([k_rot_embed, k_pass], dim=-1) if k_pass.shape[-1] > 0 else k_rot_embed return q, k def inject_spectral_mamba(layer: Any = None,
model: Any = None,
model_config: Any = None,
comp_config: Any = None,
layer_idx: int = 0,
total_layers: int = 12,
old_attn: nn.Module | None = None,
**kwargs: Any) ‑> Any | None-
Expand source code
@ATTENTION_REGISTRY.register("spectral_mamba_inject") @ATTENTION_REGISTRY.register("inject_spectral_mamba") def inject_spectral_mamba( layer: Any = None, model: Any = None, model_config: Any = None, comp_config: Any = None, layer_idx: int = 0, total_layers: int = 12, old_attn: nn.Module | None = None, **kwargs: Any, ) -> Any | None: """Safely replace attention in model or layer with SpectralMambaAttention.""" # Support both `layer=` and `model=` call conventions if layer is None and model is not None: layer = model cfg = model_config or getattr(layer, "config", None) # 1. Whole model traversal support if layer is not None and ( hasattr(layer, "transformer") or hasattr(layer, "distilbert") or hasattr(layer, "roberta") or hasattr(layer, "bert") or (hasattr(layer, "model") and not hasattr(layer, "forward")) ): layer_list = None if hasattr(layer, "distilbert") and hasattr(layer.distilbert, "transformer"): layer_list = layer.distilbert.transformer.layer elif hasattr(layer, "roberta") and hasattr(layer.roberta, "encoder"): layer_list = layer.roberta.encoder.layer elif hasattr(layer, "bert") and hasattr(layer.bert, "encoder"): layer_list = layer.bert.encoder.layer elif hasattr(layer, "transformer") and hasattr(layer.transformer, "h"): layer_list = layer.transformer.h elif hasattr(layer, "model") and hasattr(layer.model, "layers"): layer_list = layer.model.layers elif hasattr(layer, "encoder") and hasattr(layer.encoder, "layer"): layer_list = layer.encoder.layer if layer_list is not None: tot = len(layer_list) for idx, lyr in enumerate(layer_list): inject_spectral_mamba( layer=lyr, model_config=cfg, comp_config=comp_config, layer_idx=idx, total_layers=tot, **kwargs, ) return layer # 2. Named modules whole-model fallback if ( old_attn is None and layer is not None and hasattr(layer, "named_modules") and not ( hasattr(layer, "attention") or hasattr(layer, "self") or hasattr(layer, "attn") or hasattr(layer, "q_proj") or hasattr(layer, "q_lin") or hasattr(layer, "query") ) ): try: replaced = 0 for name, module in layer.named_modules(): if ( hasattr(module, "q_lin") or hasattr(module, "query") or hasattr(module, "q_proj") or hasattr(module, "c_attn") ): parent_name = name.rsplit(".", 1)[0] if "." in name else "" child_name = name.rsplit(".", 1)[-1] if "." in name else name parent = layer.get_submodule(parent_name) if parent_name else layer new_attn = inject_spectral_mamba( layer=None, model_config=cfg, comp_config=comp_config, old_attn=module, **kwargs, ) if new_attn is not None: setattr(parent, child_name, new_attn) replaced += 1 if replaced > 0: return layer except Exception: return None src_attn = old_attn if src_attn is None and layer is not None: src_attn = getattr(layer, "attention", getattr(layer, "self", getattr(layer, "attn", None))) hidden_size = ( kwargs.get("hidden_size") or getattr(src_attn, "hidden_size", None) or getattr(src_attn, "embed_dim", None) or ( src_attn.q_proj.in_features if hasattr(src_attn, "q_proj") and hasattr(src_attn.q_proj, "in_features") else None ) or ( src_attn.query.in_features if hasattr(src_attn, "query") and hasattr(src_attn.query, "in_features") else None ) or ( src_attn.q_lin.in_features if hasattr(src_attn, "q_lin") and hasattr(src_attn.q_lin, "in_features") else None ) or (layer.in_features if hasattr(layer, "in_features") else None) or _cfg_get(cfg, "hidden_size", None) or _cfg_get(cfg, "d_model", 64) ) num_heads = ( kwargs.get("num_heads") or kwargs.get("num_attention_heads") or getattr(src_attn, "num_heads", None) or getattr(src_attn, "num_attention_heads", None) or _cfg_get(cfg, "num_attention_heads", None) or _cfg_get(cfg, "num_heads", 4) ) try: new_attn = SpectralMambaAttention( config=cfg, pretrained_attn=src_attn, hidden_size=hidden_size, num_heads=num_heads, num_attention_heads=num_heads, comp_config=comp_config, **kwargs, ) if layer is not None: if hasattr(layer, "attention"): if hasattr(layer.attention, "self"): layer.attention.self = new_attn new_attn.is_nested_encoder = True else: layer.attention = new_attn elif hasattr(layer, "self"): layer.self = new_attn new_attn.is_nested_encoder = True elif hasattr(layer, "attn"): layer.attn = new_attn elif hasattr(layer, "self_attn"): layer.self_attn = new_attn else: if hasattr(layer, "attention") and hasattr(layer.attention, "self"): layer.attention.self = new_attn new_attn.is_nested_encoder = True else: layer.attention = new_attn return new_attn except Exception: return NoneSafely replace attention in model or layer with SpectralMambaAttention.
Classes
class LoRALinear (in_features: int,
out_features: int,
bias: bool = True,
lora_rank: int = 16,
lora_alpha: float = 32.0,
freeze_base: bool = False,
device: torch.device | None = None,
dtype: torch.dtype | None = None)-
Expand source code
class LoRALinear(nn.Module): def __init__( self, in_features: int, out_features: int, bias: bool = True, lora_rank: int = 16, lora_alpha: float = 32.0, freeze_base: bool = False, device: torch.device | None = None, dtype: torch.dtype | None = None, ) -> None: super().__init__() factory_kwargs = {"device": device, "dtype": dtype} self.in_features = in_features self.out_features = out_features self.lora_rank = lora_rank self.lora_alpha = lora_alpha self.lora_scaling = (lora_alpha / lora_rank) if lora_rank > 0 else 1.0 self.freeze_base = freeze_base self.weight = nn.Parameter(torch.empty((out_features, in_features), **factory_kwargs)) if bias: self.bias = nn.Parameter(torch.empty(out_features, **factory_kwargs)) else: self.register_parameter("bias", None) if self.lora_rank > 0: self.lora_A = nn.Parameter(torch.empty((lora_rank, in_features), **factory_kwargs)) self.lora_B = nn.Parameter(torch.empty((out_features, lora_rank), **factory_kwargs)) else: self.register_parameter("lora_A", None) self.register_parameter("lora_B", None) self.reset_parameters() if freeze_base: self.freeze_base_weights() def reset_parameters(self) -> None: nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5)) if self.bias is not None: fan_in, _ = nn.init._calculate_fan_in_and_fan_out(self.weight) bound = 1 / math.sqrt(fan_in) if fan_in > 0 else 0 nn.init.uniform_(self.bias, -bound, bound) if self.lora_rank > 0 and self.lora_A is not None: nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5)) self.lora_A.requires_grad = True if self.lora_rank > 0 and self.lora_B is not None: nn.init.zeros_(self.lora_B) self.lora_B.requires_grad = True def freeze_base_weights(self) -> None: self.weight.requires_grad = False if self.bias is not None: self.bias.requires_grad = False if self.lora_rank > 0: if self.lora_A is not None: self.lora_A.requires_grad = True if self.lora_B is not None: self.lora_B.requires_grad = True def forward(self, x: torch.Tensor) -> torch.Tensor: out = F.linear(x, self.weight, self.bias) if self.lora_rank > 0 and self.lora_A is not None and self.lora_B is not None: lora_out = F.linear(F.linear(x, self.lora_A), self.lora_B) * self.lora_scaling out = out + lora_out return outBase class for all neural network modules.
Your models should also subclass this class.
Modules can also contain other Modules, allowing them to be nested in a tree structure. You can assign the submodules as regular attributes::
import torch.nn as nn import torch.nn.functional as F class Model(nn.Module): def __init__(self) -> None: super().__init__() self.conv1 = nn.Conv2d(1, 20, 5) self.conv2 = nn.Conv2d(20, 20, 5) def forward(self, x): x = F.relu(self.conv1(x)) return F.relu(self.conv2(x))Submodules assigned in this way will be registered, and will also have their parameters converted when you call :meth:
to, etc.Note
As per the example above, an
__init__()call to the parent class must be made before assignment on the child.:ivar training: Boolean represents whether this module is in training or evaluation mode. :vartype training: bool
Initialize internal Module state, shared by both nn.Module and ScriptModule.
Ancestors
- torch.nn.modules.module.Module
Methods
def forward(self, x: torch.Tensor) ‑> torch.Tensor-
Expand source code
def forward(self, x: torch.Tensor) -> torch.Tensor: out = F.linear(x, self.weight, self.bias) if self.lora_rank > 0 and self.lora_A is not None and self.lora_B is not None: lora_out = F.linear(F.linear(x, self.lora_A), self.lora_B) * self.lora_scaling out = out + lora_out return outDefine the computation performed at every call.
Should be overridden by all subclasses.
Note
Although the recipe for forward pass needs to be defined within this function, one should call the :class:
Moduleinstance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them. def freeze_base_weights(self) ‑> None-
Expand source code
def freeze_base_weights(self) -> None: self.weight.requires_grad = False if self.bias is not None: self.bias.requires_grad = False if self.lora_rank > 0: if self.lora_A is not None: self.lora_A.requires_grad = True if self.lora_B is not None: self.lora_B.requires_grad = True def reset_parameters(self) ‑> None-
Expand source code
def reset_parameters(self) -> None: nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5)) if self.bias is not None: fan_in, _ = nn.init._calculate_fan_in_and_fan_out(self.weight) bound = 1 / math.sqrt(fan_in) if fan_in > 0 else 0 nn.init.uniform_(self.bias, -bound, bound) if self.lora_rank > 0 and self.lora_A is not None: nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5)) self.lora_A.requires_grad = True if self.lora_rank > 0 and self.lora_B is not None: nn.init.zeros_(self.lora_B) self.lora_B.requires_grad = True
class Mamba (d_model: int = 768, **kwargs: Any)-
Expand source code
class Mamba(nn.Module): def __init__(self, d_model: int = 768, **kwargs: Any) -> None: super().__init__() self.in_proj = nn.Linear(d_model, 2 * d_model) self.out_proj = nn.Linear(d_model, d_model) nn.init.zeros_(self.out_proj.weight) if self.out_proj.bias is not None: nn.init.zeros_(self.out_proj.bias) def forward(self, x: torch.Tensor, **kwargs: Any) -> torch.Tensor: x1, x2 = self.in_proj(x).chunk(2, dim=-1) y = F.silu(x1) * x2 return self.out_proj(y)Base class for all neural network modules.
Your models should also subclass this class.
Modules can also contain other Modules, allowing them to be nested in a tree structure. You can assign the submodules as regular attributes::
import torch.nn as nn import torch.nn.functional as F class Model(nn.Module): def __init__(self) -> None: super().__init__() self.conv1 = nn.Conv2d(1, 20, 5) self.conv2 = nn.Conv2d(20, 20, 5) def forward(self, x): x = F.relu(self.conv1(x)) return F.relu(self.conv2(x))Submodules assigned in this way will be registered, and will also have their parameters converted when you call :meth:
to, etc.Note
As per the example above, an
__init__()call to the parent class must be made before assignment on the child.:ivar training: Boolean represents whether this module is in training or evaluation mode. :vartype training: bool
Initialize internal Module state, shared by both nn.Module and ScriptModule.
Ancestors
- torch.nn.modules.module.Module
Methods
def forward(self, x: torch.Tensor, **kwargs: Any) ‑> torch.Tensor-
Expand source code
def forward(self, x: torch.Tensor, **kwargs: Any) -> torch.Tensor: x1, x2 = self.in_proj(x).chunk(2, dim=-1) y = F.silu(x1) * x2 return self.out_proj(y)Define the computation performed at every call.
Should be overridden by all subclasses.
Note
Although the recipe for forward pass needs to be defined within this function, one should call the :class:
Moduleinstance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.
class SpectralMamba (d_model: int = 768, d_state: int = 16, **kwargs: Any)-
Expand source code
class SpectralMamba(nn.Module): """ Spectral modulation with actual frequency filtering and low-rank state projection. Uses a learned frequency mask over FFT bins followed by low-rank adaptation. """ def __init__(self, d_model: int = 768, d_state: int = 16, **kwargs: Any) -> None: super().__init__() self.d_model = d_model # Use a fixed resolution for frequency weights and interpolate to sequence length self.register_buffer( "freq_weights", torch.linspace(0.1, 1.0, 256).unsqueeze(1).expand(256, d_model).clone(), ) self.register_buffer("dt_scale", torch.zeros(d_model)) # Low-rank adaptation projection self.lora_down = nn.Parameter(torch.randn(d_model, d_state) * 0.01) self.lora_up = nn.Parameter(torch.zeros(d_state, d_model)) def forward(self, x: torch.Tensor) -> torch.Tensor: """ Spectral modulation with actual frequency filtering. """ B, L, D = x.shape dtype = x.dtype # FFT along sequence axis x_float = x.float() x_fft = torch.fft.rfft(x_float, dim=1) # (B, L//2+1, D) freq_weights = self.freq_weights if freq_weights.shape[0] != x_fft.shape[1]: freq_weights = ( F.interpolate( freq_weights.transpose(0, 1).unsqueeze(0), size=x_fft.shape[1], mode="linear", align_corners=False, ) .squeeze(0) .transpose(0, 1) ) x_fft_filtered = x_fft * freq_weights.unsqueeze(0) # Back to time domain x_filtered = torch.fft.irfft(x_fft_filtered, n=L, dim=1).to(dtype) # Low-rank state projection state_act = F.linear( F.linear(x_filtered, self.lora_down.t()), self.lora_up.t(), ) # Learned channel gate gate = torch.sigmoid(self.dt_scale).unsqueeze(0).unsqueeze(0) return (x_filtered + state_act) * gateSpectral modulation with actual frequency filtering and low-rank state projection.
Uses a learned frequency mask over FFT bins followed by low-rank adaptation.
Initialize internal Module state, shared by both nn.Module and ScriptModule.
Ancestors
- torch.nn.modules.module.Module
Methods
def forward(self, x: torch.Tensor) ‑> torch.Tensor-
Expand source code
def forward(self, x: torch.Tensor) -> torch.Tensor: """ Spectral modulation with actual frequency filtering. """ B, L, D = x.shape dtype = x.dtype # FFT along sequence axis x_float = x.float() x_fft = torch.fft.rfft(x_float, dim=1) # (B, L//2+1, D) freq_weights = self.freq_weights if freq_weights.shape[0] != x_fft.shape[1]: freq_weights = ( F.interpolate( freq_weights.transpose(0, 1).unsqueeze(0), size=x_fft.shape[1], mode="linear", align_corners=False, ) .squeeze(0) .transpose(0, 1) ) x_fft_filtered = x_fft * freq_weights.unsqueeze(0) # Back to time domain x_filtered = torch.fft.irfft(x_fft_filtered, n=L, dim=1).to(dtype) # Low-rank state projection state_act = F.linear( F.linear(x_filtered, self.lora_down.t()), self.lora_up.t(), ) # Learned channel gate gate = torch.sigmoid(self.dt_scale).unsqueeze(0).unsqueeze(0) return (x_filtered + state_act) * gateSpectral modulation with actual frequency filtering.
class SpectralMambaAttention (config: Any | None = None,
rank: int | None = 64,
pretrained_attn: nn.Module | None = None,
**kwargs: Any)-
Expand source code
@ATTENTION_REGISTRY.register("spectral_mamba") class SpectralMambaAttention(BaseAttention): is_custom_attention: bool = True def __init__( self, config: Any | None = None, rank: int | None = 64, pretrained_attn: nn.Module | None = None, **kwargs: Any, ) -> None: super().__init__(config, rank or 64) self.config = config comp_config = _cfg_get(config, "compression_config", kwargs.get("comp_config")) self.hidden_size: int = kwargs.get( "hidden_size", _cfg_get( config, "hidden_size", _cfg_get(config, "dim", _cfg_get(config, "d_model", 768)), ), ) self.dim = self.hidden_size self.num_heads: int = kwargs.get( "num_heads", kwargs.get( "num_attention_heads", _cfg_get( config, "num_attention_heads", _cfg_get(config, "num_heads", _cfg_get(config, "n_heads", 12)), ), ), ) self.num_kv_heads: int = kwargs.get( "num_kv_heads", kwargs.get( "num_key_value_heads", _cfg_get( config, "num_key_value_heads", _cfg_get(config, "num_kv_heads", self.num_heads), ), ), ) self.head_dim: int = kwargs.get( "head_dim", _cfg_get( config, "head_dim", self.dim // self.num_heads if self.num_heads > 0 else 64, ), ) self.all_head_size: int = self.num_heads * self.head_dim self.all_kv_head_size: int = self.num_kv_heads * self.head_dim is_dec = _cfg_get(config, "is_decoder", kwargs.get("is_decoder", False)) model_type = str(_cfg_get(config, "model_type", "") or "").lower() decoder_prefixes = ( "gpt", "qwen", "llama", "mistral", "gemma", "phi", "opt", "bloom", "falcon", "olmo", "mpt", "stablelm", ) if ( any(model_type.startswith(p) for p in decoder_prefixes) or hasattr(config, "architectures") and any( "GPT" in str(a) or "CausalLM" in str(a) for a in (getattr(config, "architectures", []) or []) ) ): is_dec = True if kwargs.get("is_causal", False): is_dec = True self.is_decoder = bool(is_dec) self.is_causal = bool(is_dec) self.is_nested_encoder = False self.lora_rank = _cfg_get( config, "lora_rank", _cfg_get(comp_config, "lora_rank", kwargs.get("lora_rank", 16)), ) self.lora_alpha = _cfg_get( config, "lora_alpha", _cfg_get(comp_config, "lora_alpha", kwargs.get("lora_alpha", 32.0)), ) self.freeze_base = self.lora_rank > 0 # Projections self.q_lin = LoRALinear( self.dim, self.all_head_size, bias=True, lora_rank=self.lora_rank, lora_alpha=self.lora_alpha, freeze_base=self.freeze_base, ) self.k_lin = LoRALinear( self.dim, self.all_kv_head_size, bias=True, lora_rank=self.lora_rank, lora_alpha=self.lora_alpha, freeze_base=self.freeze_base, ) self.v_lin = LoRALinear( self.dim, self.all_kv_head_size, bias=True, lora_rank=self.lora_rank, lora_alpha=self.lora_alpha, freeze_base=self.freeze_base, ) self.out_lin = LoRALinear( self.all_head_size, self.dim, bias=True, lora_rank=self.lora_rank, lora_alpha=self.lora_alpha, freeze_base=self.freeze_base, ) self.q_proj = self.q_lin self.k_proj = self.k_lin self.v_proj = self.v_lin self.out_proj = self.out_lin self.query = self.q_lin self.key = self.k_lin self.value = self.v_lin d_state = _cfg_get(comp_config, "d_state", _cfg_get(config, "d_state", 16)) self.spectral_mamba = SpectralMamba(self.hidden_size, d_state=d_state) # Gate implemented with frozen buffer weight/bias + LoRA adaptation so zero base parameters are added # Gate initialization: # output = gate * attn + (1 - gate) * mamba # With gate_bias=10.0, sigmoid(10.0) ≈ 0.99995, so the model starts as ~pure attention (identical to pretrained). self.register_buffer("gate_weight", torch.zeros(1, self.hidden_size)) self.register_buffer("gate_bias", torch.tensor([10.0])) self.gate_lora_A = nn.Parameter(torch.randn(4, self.hidden_size) * 0.01) self.gate_lora_B = nn.Parameter(torch.zeros(1, 4)) dropout_p = _cfg_get( config, "dropout", _cfg_get(config, "attention_probs_dropout_prob", 0.1) ) self.dropout = nn.Dropout(dropout_p) source_attn = pretrained_attn if pretrained_attn is not None else kwargs.get("source_attn") if source_attn is not None: self.copy_weights(source_attn) def copy_weights(self, src: Any) -> None: if src is None: return cls_name = src.__class__.__name__ has_out_lin = hasattr(src, "out_lin") or hasattr(src, "c_proj") or hasattr(src, "o_proj") is_nested = False if ( ( hasattr(src, "self") and hasattr(src.self, "query") and hasattr(src, "output") and hasattr(src.output, "dense") ) or ( cls_name in [ "BertSelfAttention", "RobertaSelfAttention", "XLMRobertaSelfAttention", "ElectraSelfAttention", "DebertaSelfAttention", ] or (cls_name.endswith("SelfAttention") and "distil" not in cls_name.lower()) ) and not has_out_lin ): is_nested = True self.is_nested_encoder = bool(is_nested) extracted: dict[str, tuple[torch.Tensor | None, torch.Tensor | None]] = { "q": (None, None), "k": (None, None), "v": (None, None), "o": (None, None), } # 1. GPT-2 fused if hasattr(src, "c_attn") and src.c_attn is not None: c_attn_mod = src.c_attn base_c = getattr(c_attn_mod, "base_layer", c_attn_mod) c_w = getattr(base_c, "weight", getattr(c_attn_mod, "weight", None)) if c_w is not None: c_w = c_w.data.clone() if ( (c_attn_mod.__class__.__name__ == "Conv1D") or (base_c.__class__.__name__ == "Conv1D") or (c_w.ndim == 2 and c_w.shape[0] < c_w.shape[1]) ): c_w = c_w.t().contiguous() split_size = c_w.size(0) // 3 bias_mod = getattr(base_c, "bias", getattr(c_attn_mod, "bias", None)) c_b = bias_mod.data.clone() if bias_mod is not None else None extracted["q"] = ( c_w[:split_size, :], c_b[:split_size] if c_b is not None else None, ) extracted["k"] = ( c_w[split_size : 2 * split_size, :], c_b[split_size : 2 * split_size] if c_b is not None else None, ) extracted["v"] = ( c_w[2 * split_size :, :], c_b[2 * split_size :] if c_b is not None else None, ) if hasattr(src, "c_proj") and src.c_proj is not None: c_proj_mod = src.c_proj base_o = getattr(c_proj_mod, "base_layer", c_proj_mod) o_w = getattr(base_o, "weight", getattr(c_proj_mod, "weight", None)) if o_w is not None: o_w = o_w.data.clone() if ( (c_proj_mod.__class__.__name__ == "Conv1D") or (base_o.__class__.__name__ == "Conv1D") or (o_w.ndim == 2 and o_w.shape[0] < o_w.shape[1]) ): o_w = o_w.t().contiguous() bias_o = getattr(base_o, "bias", getattr(c_proj_mod, "bias", None)) extracted["o"] = ( o_w, bias_o.data.clone() if bias_o is not None else None, ) # 2. RoBERTa nested elif hasattr(src, "self") and hasattr(src.self, "query"): s = src.self for proj_key, sub_name in [("q", "query"), ("k", "key"), ("v", "value")]: mod = getattr(s, sub_name, None) w, b = _unpack_weight_bias(mod) if w is not None: extracted[proj_key] = (w, b) if hasattr(src, "output") and hasattr(src.output, "dense"): out_mod = src.output.dense w, b = _unpack_weight_bias(out_mod) if w is not None: extracted["o"] = (w, b) # 3. Flat modules else: proj_candidates = { "q": ["q_lin", "q_proj", "query"], "k": ["k_lin", "k_proj", "key"], "v": ["v_lin", "v_proj", "value"], "o": ["out_lin", "out_proj", "dense", "o_proj"], } for proj_key, names in proj_candidates.items(): for name in names: if hasattr(src, name): mod = getattr(src, name) w, b = _unpack_weight_bias(mod) if w is not None: extracted[proj_key] = (w, b) break missing = [ k for k, (w, _) in extracted.items() if w is None and (k != "o" or not self.is_nested_encoder) ] if missing: try: from neuralops.core.adapters import UniversalAdapter adapter = UniversalAdapter() for k in missing: w, b = adapter.get_weight_bias(src, k) if w is not None: extracted[k] = (w.clone(), b.clone() if b is not None else None) except Exception as e: logger.debug("UniversalAdapter fallback failed: %s", e) missing_after = [ k for k, (w, _) in extracted.items() if w is None and (k != "o" or not self.is_nested_encoder) ] if missing_after: logger.warning( f"SpectralMambaAttention copy_weights failed: missing mandatory projections {missing_after} from {src.__class__.__name__}." ) return target_map = { "q": self.q_lin, "k": self.k_lin, "v": self.v_lin, "o": self.out_lin, } with torch.no_grad(): for proj_key, target_mod in target_map.items(): w, b = extracted[proj_key] if w is None: if proj_key == "o" and self.is_nested_encoder: nn.init.eye_(target_mod.weight) if target_mod.bias is not None: nn.init.zeros_(target_mod.bias) continue else: logger.warning(f"Missing required weights for {proj_key}") return if w.shape == target_mod.weight.shape: target_mod.weight.copy_(w) elif w.t().contiguous().shape == target_mod.weight.shape: target_mod.weight.copy_(w.t().contiguous()) elif ( w.shape[0] < target_mod.weight.shape[0] and target_mod.weight.shape[0] % w.shape[0] == 0 ): repeats = target_mod.weight.shape[0] // w.shape[0] # For GQA expansion or repeating target_mod.weight.copy_(w.repeat(repeats, 1)) elif ( w.shape[0] > target_mod.weight.shape[0] and w.shape[0] % target_mod.weight.shape[0] == 0 ): repeats = w.shape[0] // target_mod.weight.shape[0] # For GQA: mean-average the head groups from source w_grouped = w.view(target_mod.weight.shape[0], repeats, w.shape[1]).mean(dim=1) target_mod.weight.copy_(w_grouped) else: min_0 = min(target_mod.weight.shape[0], w.shape[0]) min_1 = min(target_mod.weight.shape[1], w.shape[1]) target_mod.weight.zero_() target_mod.weight[:min_0, :min_1].copy_(w[:min_0, :min_1]) if b is not None and target_mod.bias is not None: if b.shape == target_mod.bias.shape: target_mod.bias.copy_(b) elif ( b.shape[0] < target_mod.bias.shape[0] and target_mod.bias.shape[0] % b.shape[0] == 0 ): repeats = target_mod.bias.shape[0] // b.shape[0] target_mod.bias.copy_(b.repeat(repeats)) elif ( b.shape[0] > target_mod.bias.shape[0] and b.shape[0] % target_mod.bias.shape[0] == 0 ): repeats = b.shape[0] // target_mod.bias.shape[0] b_grouped = b.view(target_mod.bias.shape[0], repeats).mean(dim=1) target_mod.bias.copy_(b_grouped) else: min_b = min(target_mod.bias.shape[0], b.shape[0]) target_mod.bias.zero_() target_mod.bias[:min_b].copy_(b[:min_b]) elif target_mod.bias is not None: nn.init.zeros_(target_mod.bias) source_norm = torch.linalg.norm(w.float()).item() fnorm = torch.linalg.norm(target_mod.weight.float()).item() if fnorm < 1e-4: if source_norm > 1e-3: raise RuntimeError( f"SpectralMambaAttention: weight copy for {proj_key} produced near-zero norm " f"({fnorm:.2e}) from source norm ({source_norm:.2e}). Aborting." ) logger.warning(f"Weight copy notice for {proj_key}: norm {fnorm:.6f} < 1e-4") if self.freeze_base: for mod in [self.q_lin, self.k_lin, self.v_lin, self.out_lin]: mod.freeze_base_weights() _copy_pretrained = copy_weights @classmethod def from_distilbert_pretrained( cls, old_attn: nn.Module, model_config: Any | None = None, **kwargs: Any ) -> SpectralMambaAttention: return cls(config=model_config, pretrained_attn=old_attn, **kwargs) def forward( self, hidden_states: torch.Tensor | None = None, attention_mask: torch.Tensor | None = None, head_mask: torch.Tensor | None = None, encoder_hidden_states: torch.Tensor | None = None, encoder_attention_mask: torch.Tensor | None = None, past_key_value: Any | None = None, output_attentions: bool | None = False, output_hidden_states: bool | None = False, return_dict: bool | None = True, position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None, *args: Any, **kwargs: Any, ) -> tuple[torch.Tensor, torch.Tensor | None]: if hidden_states is None: if "hidden_states" in kwargs and kwargs["hidden_states"] is not None: hidden_states = kwargs["hidden_states"] elif "x" in kwargs and kwargs["x"] is not None: hidden_states = kwargs["x"] elif "input_ids" in kwargs and kwargs["input_ids"] is not None: hidden_states = kwargs["input_ids"] elif len(args) > 0 and args[0] is not None: hidden_states = args[0] if hidden_states is None: raise ValueError( "SpectralMambaAttention forward requires hidden_states " "(pass as positional or via kwargs['hidden_states'])" ) if attention_mask is None: attention_mask = kwargs.get("attention_mask") if hidden_states.dim() == 2: hidden_states = hidden_states.unsqueeze(0) B, N, D = hidden_states.shape q = self.q_lin(hidden_states).view(B, N, self.num_heads, self.head_dim).permute(0, 2, 1, 3) k = ( self.k_lin(hidden_states) .view(B, N, self.num_kv_heads, self.head_dim) .permute(0, 2, 1, 3) ) v = ( self.v_lin(hidden_states) .view(B, N, self.num_kv_heads, self.head_dim) .permute(0, 2, 1, 3) ) # Apply RoPE if provided if position_embeddings is None: position_embeddings = kwargs.get("position_embeddings") if position_embeddings is not None: cos, sin = position_embeddings q, k = apply_rotary_pos_emb(q, k, cos, sin) # Repeat KV for GQA if self.num_heads != self.num_kv_heads: repeats = self.num_heads // self.num_kv_heads k = repeat_kv(k, repeats) v = repeat_kv(v, repeats) scale = 1.0 / math.sqrt(self.head_dim) # Memory-efficient: use SDPA when attention weights are not requested attn_out = None attn_weights = None if not output_attentions and head_mask is None: try: _qs = q.half() if q.dtype == torch.float32 else q _ks = k.half() if k.dtype == torch.float32 else k _vs = v.half() if v.dtype == torch.float32 else v _sdpa_out = torch.nn.functional.scaled_dot_product_attention( _qs, _ks, _vs, is_causal=getattr(self, "is_decoder", False) and N > 1, dropout_p=self.dropout.p if self.training else 0.0, ) _sdpa_out = ( _sdpa_out.to(hidden_states.dtype) .transpose(1, 2) .contiguous() .view(B, N, self.all_head_size) ) _sdpa_out = torch.nan_to_num(_sdpa_out, nan=0.0, posinf=1e4, neginf=-1e4) attn_out = ( self.out_proj(_sdpa_out) if not getattr(self, "is_nested_encoder", False) else _sdpa_out ) except Exception as e: logger.debug("SDPA failed, falling back to manual attention: %s", e) attn_out = None if attn_out is None: scores = torch.matmul(q.float(), k.float().transpose(-2, -1).contiguous()) * scale # Causal masking is_causal = self.is_decoder or kwargs.get("is_causal", False) if is_causal: causal_mask = torch.triu(torch.full((N, N), -1e4, device=scores.device), diagonal=1) scores = scores + causal_mask.unsqueeze(0).unsqueeze(0) # Attention mask if attention_mask is not None: scores = apply_attention_mask(scores, attention_mask) attn_weights = F.softmax(scores, dim=-1).to(hidden_states.dtype) attn_weights = torch.nan_to_num(attn_weights, nan=0.0) attn_weights = self.dropout(attn_weights) # head_mask handling if head_mask is not None: if head_mask.dim() == 1 and head_mask.size(0) == self.num_heads: attn_weights = attn_weights * head_mask.view(1, self.num_heads, 1, 1).to( attn_weights.dtype ) elif head_mask.dim() == 2 and head_mask.size(1) == self.num_heads: attn_weights = attn_weights * head_mask.view( head_mask.size(0), self.num_heads, 1, 1 ).to(attn_weights.dtype) elif head_mask.dim() == 4: attn_weights = attn_weights * head_mask.to(attn_weights.dtype) attn_ctx = ( torch.matmul(attn_weights, v) .permute(0, 2, 1, 3) .contiguous() .view(B, N, self.all_head_size) ) if self.is_nested_encoder: attn_out = attn_ctx else: attn_out = self.out_lin(attn_ctx) # Spectral Mamba stream mamba_out = self.spectral_mamba(hidden_states) while isinstance(mamba_out, (tuple, list)): mamba_out = mamba_out[0] if mamba_out.shape[-1] != attn_out.shape[-1]: if mamba_out.shape[-1] < attn_out.shape[-1]: mamba_out = F.pad(mamba_out, (0, attn_out.shape[-1] - mamba_out.shape[-1])) else: mamba_out = mamba_out[..., : attn_out.shape[-1]] # Adaptive gating: base linear + lora adapter gate_logits = F.linear(hidden_states, self.gate_weight, self.gate_bias) gate_logits = ( gate_logits + F.linear(F.linear(hidden_states, self.gate_lora_A), self.gate_lora_B) * 2.0 ) g = torch.sigmoid(gate_logits) # (B, N, 1) out = g * attn_out + (1.0 - g) * mamba_out if out.shape[-1] != hidden_states.shape[-1]: if out.shape[-1] < hidden_states.shape[-1]: out = F.pad(out, (0, hidden_states.shape[-1] - out.shape[-1])) else: out = out[..., : hidden_states.shape[-1]] weights_to_return = None if output_attentions: weights_to_return = locals().get("attn_weights", None) return out, weights_to_returnEvery custom attention module must inherit from this.
Initialize internal Module state, shared by both nn.Module and ScriptModule.
Ancestors
- BaseAttention
- torch.nn.modules.module.Module
Class variables
var is_custom_attention : bool-
The type of the None singleton.
Static methods
def from_distilbert_pretrained(old_attn: nn.Module, model_config: Any | None = None, **kwargs: Any) ‑> SpectralMambaAttention
Inherited members
class SpectralMambaConfig (hidden_size: int = 768,
num_attention_heads: int = 12,
head_dim: int = 64,
d_state: int = 16,
d_conv: int = 4,
expand: int = 2,
dropout: float = 0.1,
gate_init_bias: float = 1.8)-
Expand source code
@dataclass class SpectralMambaConfig: hidden_size: int = 768 num_attention_heads: int = 12 head_dim: int = 64 d_state: int = 16 d_conv: int = 4 expand: int = 2 dropout: float = 0.1 gate_init_bias: float = 1.8SpectralMambaConfig(hidden_size: 'int' = 768, num_attention_heads: 'int' = 12, head_dim: 'int' = 64, d_state: 'int' = 16, d_conv: 'int' = 4, expand: 'int' = 2, dropout: 'float' = 0.1, gate_init_bias: 'float' = 1.8)
Instance variables
var d_conv : int-
The type of the None singleton.
var d_state : int-
The type of the None singleton.
var dropout : float-
The type of the None singleton.
var expand : int-
The type of the None singleton.
var gate_init_bias : float-
The type of the None singleton.
var head_dim : int-
The type of the None singleton.
-
The type of the None singleton.
var num_attention_heads : int-
The type of the None singleton.