fishy.models.deep.augformer

AugFormer: Multi-View Self-Attention over Augmented Spectral Views.

The core insight: instead of splitting the spectrum into spatial patches (which requires large datasets for attention to learn), generate K augmented views of the same spectrum and treat each view as a token. The transformer can then attend across views and learn which features are consistent across noise, scale, and shift — i.e., the features that are invariant to plausible measurement variation, which are exactly the robust, discriminative features for classification.

Classes

class fishy.models.deep.augformer.AugFormer(input_dim: int, output_dim: int, hidden_dim: int = 512, num_layers: int = 4, num_heads: int = 8, dropout: float = 0.3, num_views: int = 6, use_xsa: bool = False, use_mla: bool = False, use_mhc: bool = False, use_engram: bool = False, engram_slots: int = 128, latent_dim: int = 64, **kwargs)

Bases: Module, TTTMixin

Augmentation-as-Sequence Transformer for REIMS spectral classification. Iteration 1 Baseline: Anchor-centric, Deep Spectral Gating, no tokenizer.

__init__(input_dim: int, output_dim: int, hidden_dim: int = 512, num_layers: int = 4, num_heads: int = 8, dropout: float = 0.3, num_views: int = 6, use_xsa: bool = False, use_mla: bool = False, use_mhc: bool = False, use_engram: bool = False, engram_slots: int = 128, latent_dim: int = 64, **kwargs) → None

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(x: Tensor, return_attention: bool = False, *args, **kwargs) → Tensor | Tuple[Tensor, List]

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 Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class fishy.models.deep.augformer.EngramMemory(dim: int, num_slots: int = 128)

Bases: Module

Engram Memory bank inspired by DeepSeek-V4. Decouples static spectral knowledge (peak signatures) from dynamic reasoning. Tokens query a learnable memory bank to retrieve “known” spectral patterns.

__init__(dim: int, num_slots: int = 128)

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(x: Tensor) → Tensor

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 Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class fishy.models.deep.augformer.ManifoldProjection(dim: int)

Bases: Module

Manifold-Constrained Hyper-Connection (mHC) integrator. Ensures that features from all depths are preserved and projected onto the optimal representation manifold.

__init__(dim: int)

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(x: Tensor, stream: Tensor) → Tensor

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 Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class fishy.models.deep.augformer.MultiHeadAttention(embed_dim: int, num_heads: int, use_xsa: bool = False)

Bases: Module

__init__(embed_dim: int, num_heads: int, use_xsa: bool = False)

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(x: Tensor) → Tensor

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 Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class fishy.models.deep.augformer.MultiHeadLatentAttention(embed_dim: int, num_heads: int, latent_dim: int = 64, kv_lora_rank: int = 64, qk_lora_rank: int = 64, use_xsa: bool = False)

Bases: Module

Multi-Head Latent Attention (MLA) inspired by DeepSeek-V3/V4. Compresses KV and Q into low-rank latent spaces to improve parameter efficiency and force more robust feature extraction from spectral tokens.

__init__(embed_dim: int, num_heads: int, latent_dim: int = 64, kv_lora_rank: int = 64, qk_lora_rank: int = 64, use_xsa: bool = False)

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(x: Tensor) → Tensor

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 Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class fishy.models.deep.augformer.RMSNorm(dim: int, eps: float = 1e-06)

Bases: Module

__init__(dim: int, eps: float = 1e-06)

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(x: Tensor) → Tensor

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 Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class fishy.models.deep.augformer.SpectralGating(dim: int, hidden_dim: int, dropout: float = 0.3)

Bases: Module

Coral-style gating applied spectral-wise to each view independently.

__init__(dim: int, hidden_dim: int, dropout: float = 0.3)

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(x: Tensor) → Tensor

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 Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class fishy.models.deep.augformer.TransformerBlock(embed_dim: int, num_heads: int, mlp_ratio: int = 2, dropout: float = 0.1, use_xsa: bool = False, use_mla: bool = False, latent_dim: int = 64)

Bases: Module

Pre-norm block: RMSNorm → Attention → RMSNorm → SwiGLU FFN.

__init__(embed_dim: int, num_heads: int, mlp_ratio: int = 2, dropout: float = 0.1, use_xsa: bool = False, use_mla: bool = False, latent_dim: int = 64)

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(x: Tensor) → Tensor

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 Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

s