# Jamba

> Hybrid SSM-Transformer-MoE — interleaves Mamba, attention, and MoE blocks in one stack (AI21 2024)

Pick to see three paradigms in one compact graph: state-space, attention, and mixture-of-experts. Good for studying hybrid long-context designs.

- Category: NLP/LLM
- Layers: 53
- Parameters: 13.03B
- Input shape (batchless): 1 × 4096
- Output shape: 1 × 4096 × 65536
- Verifier verdict: warn
- Graph JSON: https://neurarch.com/templates/jamba/model.json
- Open on the canvas: https://neurarch.com/?template=jamba

## Structure

| # | Layer | Type | Parameters | Output shape |
| --- | --- | --- | --- | --- |
| 1 | tokens | Input | shape=[1, 4096] | 1 × 4096 |
| 2 | token_embed | Embedding |  | 1 × 4096 × 4096 |
| 3 | mixer_norm_1 | RMSNorm | normalizedShape=4096 | 1 × 4096 × 4096 |
| 4 | mamba_1 | Mamba (SSM) |  | 1 × 4096 × 4096 |
| 5 | mixer_residual_1 | Add |  | 1 × 4096 × 4096 |
| 6 | ffn_norm_1 | RMSNorm | normalizedShape=4096 | 1 × 4096 × 4096 |
| 7 | mlp_1 | Feed Forward | embedDim=4096, ffDim=14336 | 1 × 4096 × 4096 |
| 8 | ffn_residual_1 | Add |  | 1 × 4096 × 4096 |
| 9 | mixer_norm_2 | RMSNorm | normalizedShape=4096 | 1 × 4096 × 4096 |
| 10 | mamba_2 | Mamba (SSM) |  | 1 × 4096 × 4096 |
| 11 | mixer_residual_2 | Add |  | 1 × 4096 × 4096 |
| 12 | ffn_norm_2 | RMSNorm | normalizedShape=4096 | 1 × 4096 × 4096 |
| 13 | moe_2 | MoE Layer | embedDim=4096, numExperts=16, topK=2 | 1 × 4096 × 4096 |
| 14 | ffn_residual_2 | Add |  | 1 × 4096 × 4096 |
| 15 | mixer_norm_3 | RMSNorm | normalizedShape=4096 | 1 × 4096 × 4096 |
| 16 | mamba_3 | Mamba (SSM) |  | 1 × 4096 × 4096 |
| 17 | mixer_residual_3 | Add |  | 1 × 4096 × 4096 |
| 18 | ffn_norm_3 | RMSNorm | normalizedShape=4096 | 1 × 4096 × 4096 |
| 19 | mlp_3 | Feed Forward | embedDim=4096, ffDim=14336 | 1 × 4096 × 4096 |
| 20 | ffn_residual_3 | Add |  | 1 × 4096 × 4096 |
| 21 | mixer_norm_4 | RMSNorm | normalizedShape=4096 | 1 × 4096 × 4096 |
| 22 | mamba_4 | Mamba (SSM) |  | 1 × 4096 × 4096 |
| 23 | mixer_residual_4 | Add |  | 1 × 4096 × 4096 |
| 24 | ffn_norm_4 | RMSNorm | normalizedShape=4096 | 1 × 4096 × 4096 |
| 25 | moe_4 | MoE Layer | embedDim=4096, numExperts=16, topK=2 | 1 × 4096 × 4096 |
| 26 | ffn_residual_4 | Add |  | 1 × 4096 × 4096 |
| 27 | mixer_norm_5 | RMSNorm | normalizedShape=4096 | 1 × 4096 × 4096 |
| 28 | attn_5 | Grouped Query Attn | embedDim=4096, numHeads=32, numKVHeads=8 | 1 × 4096 × 4096 |
| 29 | mixer_residual_5 | Add |  | 1 × 4096 × 4096 |
| 30 | ffn_norm_5 | RMSNorm | normalizedShape=4096 | 1 × 4096 × 4096 |
| 31 | mlp_5 | Feed Forward | embedDim=4096, ffDim=14336 | 1 × 4096 × 4096 |
| 32 | ffn_residual_5 | Add |  | 1 × 4096 × 4096 |
| 33 | mixer_norm_6 | RMSNorm | normalizedShape=4096 | 1 × 4096 × 4096 |
| 34 | mamba_6 | Mamba (SSM) |  | 1 × 4096 × 4096 |
| 35 | mixer_residual_6 | Add |  | 1 × 4096 × 4096 |
| 36 | ffn_norm_6 | RMSNorm | normalizedShape=4096 | 1 × 4096 × 4096 |
| 37 | moe_6 | MoE Layer | embedDim=4096, numExperts=16, topK=2 | 1 × 4096 × 4096 |
| 38 | ffn_residual_6 | Add |  | 1 × 4096 × 4096 |
| 39 | mixer_norm_7 | RMSNorm | normalizedShape=4096 | 1 × 4096 × 4096 |
| 40 | mamba_7 | Mamba (SSM) |  | 1 × 4096 × 4096 |
| 41 | mixer_residual_7 | Add |  | 1 × 4096 × 4096 |
| 42 | ffn_norm_7 | RMSNorm | normalizedShape=4096 | 1 × 4096 × 4096 |
| 43 | mlp_7 | Feed Forward | embedDim=4096, ffDim=14336 | 1 × 4096 × 4096 |
| 44 | ffn_residual_7 | Add |  | 1 × 4096 × 4096 |
| 45 | mixer_norm_8 | RMSNorm | normalizedShape=4096 | 1 × 4096 × 4096 |
| 46 | mamba_8 | Mamba (SSM) |  | 1 × 4096 × 4096 |
| 47 | mixer_residual_8 | Add |  | 1 × 4096 × 4096 |
| 48 | ffn_norm_8 | RMSNorm | normalizedShape=4096 | 1 × 4096 × 4096 |
| 49 | moe_8 | MoE Layer | embedDim=4096, numExperts=16, topK=2 | 1 × 4096 × 4096 |
| 50 | ffn_residual_8 | Add |  | 1 × 4096 × 4096 |
| 51 | final_norm | RMSNorm | normalizedShape=4096 | 1 × 4096 × 4096 |
| 52 | lm_head | Linear | outFeatures=65536, inFeatures=4096 | 1 × 4096 × 65536 |
| 53 | logits | Output |  | 1 × 4096 × 65536 |

## Verifier findings

- **warn** `attention-no-pe` at `attn_5`: 1 attention layer(s) present but no positional encoding found. Attention is permutation-invariant, without position information the model cannot distinguish token order. Fix: Add a PositionalEncoding (sinusoidal) or RoPE layer before the first attention layer.
- **info** `moe-no-aux-loss` at `moe_2`: MoE layers require an auxiliary router z-loss + load-balance loss during training to prevent expert collapse. This is not visible in the architecture diagram but must be in the training loop. Applies to all 4: moe_2, moe_4, moe_6, moe_8. Fix: Add a note on these layers. Typical aux_loss coefficient: 1e-2 (Mixtral/Switch Transformer).

## Exported PyTorch (first 46 lines)

```python
# Architecture designed with Neurarch: https://neurarch.com
# PyTorch: compatible with Python 3.8+ and torch>=1.12
# Colab: pip install torch torchvision  (usually pre-installed)

import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Tuple

class Jamba(nn.Module):
    def __init__(self):
        super().__init__()

        self.embedding_1 = nn.Embedding(65536, 4096)
        self.rmsNorm_1 = nn.RMSNorm(4096)
        self.mamba_1 = nn.Identity()  # Mamba(d_model=4096, d_state=16, d_conv=4, expand=2), pip install mamba-ssm and swap in
        self.rmsNorm_2 = nn.RMSNorm(4096)
        self.feedForward_1 = nn.Sequential(
            nn.Linear(4096, 14336),
            nn.ReLU(),
            nn.Linear(14336, 4096)
        )
        self.rmsNorm_3 = nn.RMSNorm(4096)
        self.mamba_2 = nn.Identity()  # Mamba(d_model=4096, d_state=16, d_conv=4, expand=2), pip install mamba-ssm and swap in
        self.rmsNorm_4 = nn.RMSNorm(4096)
        self.moeLayer_1 = nn.ModuleDict({
            'router': nn.Linear(4096, 16, bias=False),
            'experts': nn.ModuleList([
                nn.Sequential(
                    nn.Linear(4096, 14336, bias=False), nn.SiLU(),
                    nn.Linear(14336, 4096, bias=False),
                ) for _ in range(16)
            ]),
        })  # MoE top-2
        self.rmsNorm_5 = nn.RMSNorm(4096)
        self.mamba_3 = nn.Identity()  # Mamba(d_model=4096, d_state=16, d_conv=4, expand=2), pip install mamba-ssm and swap in
        self.rmsNorm_6 = nn.RMSNorm(4096)
        self.feedForward_2 = nn.Sequential(
            nn.Linear(4096, 14336),
            nn.ReLU(),
            nn.Linear(14336, 4096)
        )
        self.rmsNorm_7 = nn.RMSNorm(4096)
        self.mamba_4 = nn.Identity()  # Mamba(d_model=4096, d_state=16, d_conv=4, expand=2), pip install mamba-ssm and swap in
        self.rmsNorm_8 = nn.RMSNorm(4096)
        self.moeLayer_2 = nn.ModuleDict({
```

## Machine access

- Every architecture: https://neurarch.com/a/index.json
- Verify a graph of your own: `POST https://www.neurarch.com/api/v1/check` (see https://neurarch.com/developer.html)
- MCP server, so an agent edits the graph with the checks in the loop: https://neurarch.com/docs/mcp.md