assets/architecture.md

# SmolVLA Architecture Diagrams

Reference diagrams for understanding the SmolVLA model structure and what
the `--model-health` report measures.

---

## 1. High-level data flow

```
  Camera Images (512x512)       Task String            Robot State     Noisy Actions
  +--------+ +--------+         "pick up the           (joint          (from flow
  |  side  | |   up   |          red block"             positions)       matching)
  +---+----+ +---+----+              |                     |               |
      |          |                   |                     |               |
      v          v                   v                     v               v
  +------------------+     +---------------+     +---------------+  +---------------+
  | SigLIP Vision    |     | Tokenizer     |     | state_proj    |  | action_in_proj|
  | Encoder          |     |               |     | (Linear)      |  | (Linear)      |
  | (frozen)         |     |               |     |               |  | + time MLP    |
  | 12L, 12H         |     |               |     | trainable     |  | trainable     |
  +--------+---------+     +-------+-------+     +-------+-------+  +-------+-------+
           |                       |                     |                  |
     1024 patches            token embeds            1 token          10 tokens
           |                       |                     |                  |
           v                       |                     |                  |
  +-------------------+            |                     |                  |
  | Connector         |            |                     |                  |
  | (pixel shuffle)   |            |                     |                  |
  | trainable         |            |                     |                  |
  +--------+----------+            |                     |                  |
           |                       |                     |                  |
     64 tokens                     |                     |                  |
           |                       |                     |                  |
           +-----------+-----------+                     |                  |
                       |                                 |                  |
                       v                                 |                  |
           +-----------------------+                     |                  |
           | PREFIX SEQUENCE       |                     |                  |
           | [vision] [language]   +<--------------------+                  |
           | [state]               |     VLM dim (960)                     |
           +-----------+-----------+                                       |
                       |                                                   |
                       |                                 +-----------------+
                       |                                 |
                       v                                 v
           +---------------------------------------------------+
           |         SmolVLMWithExpert (16 layers)               |
           |                                                     |
           |   PREFIX  ------------------->  VLM Text Model      |
           |   (vision+lang+state)            (frozen)           |
           |                                  960-dim, 15 heads  |
           |                                                     |
           |   SUFFIX  ------------------->  Action Expert       |
           |   (noisy actions+time)           (trainable)        |
           |                                  480-dim, 8 heads   |
           |                                                     |
           |         (see attention diagrams below)              |
           +--------------------------+--------------------------|
                                      |
                                      |  Expert output (last 10 tokens)
                                      v
                          +-------------------+
                          | action_out_proj   |
                          | (Linear)          |
                          | trainable         |
                          +---------+---------+
                                    |
                                    v
                          Predicted Actions
                          (10 steps x 6 DOF)
```

---

## 2. The three attention components

These are the three distinct attention operations that the health report
measures for entropy and head redundancy.

### A. SigLIP Vision (12 layers, 12 heads)

Standard self-attention inside the vision encoder. Happens before anything
enters SmolVLMWithExpert. Image patches attend to other image patches.

```
  patch_1  patch_2  patch_3  ...  patch_1024
     |        |        |              |
     v        v        v              v
  +----------------------------------------------+
  | SigLIP Self-Attention                         |
  |                                               |
  | Each patch can see all patches.               |
  |                                               |
  | Q <-- patches        K <-- patches            |
  | (1024 x 1024 attention matrix)                |
  +----------------------------------------------+

  Report label: "SigLIP Vision (12L, 12H)"
```

### B. VLM+Expert Joint Self-Attention (16 layers, 15 heads)

Runs during prefill (initial encoding). VLM and Expert tokens are
concatenated into one sequence and attend to each other in a single
attention call.

```
  VLM tokens (prefix)           Expert tokens (suffix)
  [vision][lang][state]          [action+time]
        |                              |
        v                              v
  +-------------+               +-------------+
  | VLM layer   |               | Expert layer|
  | q/k/v proj  |               | q/k/v proj  |
  +------+------+               +------+------+
         |                             |
         +-------------+---------------+
                       |
                       v
                torch.cat(dim=seq)
                       |
                       v
  +----------------------------------------------+
  | JOINT Self-Attention                          |
  |                                               |
  | All tokens see all tokens:                    |
  |   vision  <-->  vision                        |
  |   vision  <-->  language                      |
  |   vision  <-->  action                        |
  |   action  <-->  language                      |
  |   action  <-->  action                        |
  |   ... etc                                     |
  |                                               |
  | Q <-- [VLM+Expert]     K <-- [VLM+Expert]     |
  | Q_len == K_len  (self-attention)              |
  +----------------------------------------------+

  Report label: "VLM+Expert Joint Self-Attn (16L, 15H)"
```

### C. Expert-to-VLM Cross-Attention (16 layers, 8 heads)

Runs during generation (autoregressive action decoding). The Expert
generates queries from its own hidden states, but keys and values come
from the VLM's cached representations, re-projected through the Expert's
k_proj and v_proj. Repeats 10 times (one per action token).

```
  VLM KV Cache                       Expert hidden states
  (built during prefill)             (current action token)
        |                                    |
        |                                    v
        |                            +---------------+
        |                            | Expert q_proj |
        |                            +-------+-------+
        |                                    |
        |     re-projected through           |
        |     Expert k_proj / v_proj         |
        v                                    |
  +-------------+                            |
  | Expert K, V |<-- VLM cache               |
  | (from VLM)  |   re-projected             |
  +------+------+                            |
         |                                   |
         +---------------+-------------------+
                         |
                         v
  +----------------------------------------------+
  | CROSS-Attention                               |
  |                                               |
  | Expert reads VLM:                             |
  |   action --> vision   (what to grab?)         |
  |   action --> language  (what was asked?)      |
  |   action --> state     (current pose?)        |
  |                                               |
  | VLM does NOT read Expert here.                |
  |                                               |
  | Q <-- Expert tokens (few)                     |
  | K <-- VLM prefix (many)                       |
  | Q_len != K_len  (cross-attention)             |
  +----------------------------------------------+

  Report label: "Expert-to-VLM Cross-Attn (16L, 8H)"
```

---

## 3. Execution timeline

```
  IMAGE ENCODING            PREFILL                   GENERATION
  +--------------+    +------------------+    +------+------+     +------+
  | SigLIP       |    | Joint Self-Attn  |    | XA   | XA   |    | XA   |
  | 12 layers    +--->| 16 layers        +--->| #1   | #2   |...>| #10  |
  | (A)          |    | (B)              |    |      |      |    |      |
  +--------------+    +------------------+    +------+------+    +------+
                                              <--- 10 action steps --->

  Report:              Report:                 Report:
  "SigLIP Vision"      "VLM+Expert             "Expert-to-VLM
                        Joint Self-Attn"        Cross-Attn"
  12L x 12H            16L x 15H               16L x 8H
                                                (x10 steps, averaged)
```

---

## 4. What the health report measures

### Section 1: Weight Spectral Analysis (alpha)

Analyzes weight matrices directly (no data needed). Each transformer layer
contains multiple weight matrices (q_proj, k_proj, v_proj, o_proj,
gate_proj, up_proj, down_proj). A power-law is fit to each matrix's
singular values; the exponent (alpha) indicates training quality.

```
  +----------------------------------+-------------+---------------------+
  | Component                        | Wt Matrices | What it is          |
  +----------------------------------+-------------+---------------------+
  | Expert (trainable)               | 112 = 16x7  | Action decoder      |
  | VLM Text Model (frozen)          | 113 = 16x7+ | Language model      |
  | Vision Encoder (frozen)          |  74 = 12x~6 | SigLIP              |
  | Connector (trainable)            |   1          | Pixel shuffle       |
  | Projections (trainable)          | too small    | state/action linear |
  +----------------------------------+-------------+---------------------+

  Healthy: alpha 2-4     Undertrained: 4-6     Severe: >6
```

### Section 2: Attention Entropy

Requires a forward pass with real data. Measures how spread out each
attention head's focus is (normalized by log of sequence length).

```
  Collapsed: <0.10     Healthy: 0.10-0.80     Unfocused: >0.80     Dead: >0.95
```

### Section 3: Head Redundancy

Same forward pass data. Measures pairwise cosine similarity between
flattened head attention patterns within each layer. Heads should learn
different patterns.

```
  Diverse: <0.70       High redundancy: >0.70       Collapsed: >0.90
```

---

## 5. Trainable vs frozen components

```
  FROZEN (pretrained, not updated during fine-tuning)
  +----------------------------------------------------+
  | SigLIP Vision Encoder                               |
  | 12 layers, 12 heads, 74 weight matrices             |
  +----------------------------------------------------+
  | VLM Text Model (SmolLM2)                            |
  | 16 layers, 15 heads, 113 weight matrices            |
  +----------------------------------------------------+

  TRAINABLE (learned during fine-tuning)
  +----------------------------------------------------+
  | Action Expert                                       |
  | 16 layers, 8 heads, 112 weight matrices             |
  +----------------------------------------------------+
  | Connector (pixel shuffle)                           |
  | 1 weight matrix                                     |
  +----------------------------------------------------+
  | state_proj, action_in_proj, action_out_proj         |
  | (too small for spectral analysis)                   |
  +----------------------------------------------------+
```