# Model Architecture

This document describes the architecture and training methodology for models fine-tuned on Zebra-CoT.

## Mixture-of-Transformer-Experts (MoT)

The architecture adopts a **Mixture-of-Transformer-Experts (MoT)** design to maximize the model's capacity to learn from richly diverse multimodal information.

### Key Design Principles

1. **Capacity Maximization**: MoT enables the model to handle the diversity of visual reasoning tasks across scientific, 2D, 3D, and logic/game domains.

2. **Expert Specialization**: Different experts can specialize in different types of reasoning patterns (geometric, spatial, strategic, etc.).

## Dual Encoder System

Following the capacity maximization principle, the architecture utilizes **two separate encoders**:

| Encoder | Purpose | Features Captured |
|---------|---------|-------------------|
| **Pixel-Level Encoder** | Low-level visual processing | Edges, textures, fine details |
| **Semantic-Level Encoder** | High-level understanding | Objects, relationships, concepts |

```
┌─────────────────────────────────────────────────────────┐
│                      Input Image                         │
└─────────────────────────────────────────────────────────┘
                            │
              ┌─────────────┴─────────────┐
              ▼                           ▼
┌─────────────────────────┐   ┌─────────────────────────┐
│   Pixel-Level Encoder   │   │  Semantic-Level Encoder │
│   (Fine visual details) │   │  (High-level concepts)  │
└─────────────────────────┘   └─────────────────────────┘
              │                           │
              └─────────────┬─────────────┘
                            ▼
┌─────────────────────────────────────────────────────────┐
│              Mixture-of-Transformer-Experts              │
│                        (MoT)                             │
└─────────────────────────────────────────────────────────┘
                            │
                            ▼
┌─────────────────────────────────────────────────────────┐
│           Next Group of Token Prediction                 │
└─────────────────────────────────────────────────────────┘
```

## Next Group of Token Prediction (NGTP)

The training paradigm follows **Next Group of Token Prediction**, where the model predicts the next group of language or visual tokens as a compression target.

### Advantages

- **Interleaved Generation**: Naturally supports generating interleaved text-image reasoning traces
- **Efficient Compression**: Groups of tokens provide better information density
- **Flexible Modality**: Can predict either language tokens or visual tokens based on context

### Token Groups

```
Input:  [Question] [Problem Image]
          │
          ▼
Step 1: [THOUGHT 1] ────────────────► Language tokens
          │
          ▼  
Step 2: [REASONING IMAGE 1] ────────► Visual tokens
          │
          ▼
Step 3: [THOUGHT 2] ────────────────► Language tokens
          │
          ▼
Output: [FINAL ANSWER]
```

## Performance Results

### In-Distribution Test Accuracy

| Model | Before Fine-tuning | After Fine-tuning | Improvement |
|-------|-------------------|-------------------|-------------|
| Anole-7B | 4.2% | 16.9% | **+12.7%** |
| Bagel-7B | — | High-quality interleaved chains | Qualitative |

### Benchmark Improvements

Fine-tuning on Zebra-CoT yields up to **+13%** performance gain on standard VLM benchmarks:

- Enhanced visual reasoning capabilities
- Improved chain-of-thought generation
- Better intermediate step visualization

## Category-Specific Performance

| Category | Description | Key Reasoning Skills |
|----------|-------------|---------------------|
| Scientific | Geometry, Physics, Algorithms | Diagram construction, step-by-step derivation |
| 2D Visual | Visual search, Jigsaw | Pattern recognition, spatial arrangement |
| 3D Visual | Multi-hop inference, Embodied planning | Depth perception, navigation |
| Logic/Games | Chess, Visual logic | Strategic thinking, rule application |

## References

- [Zebra-CoT Paper](https://arxiv.org/abs/2507.16746)
- [Anole-Zebra-CoT Model](https://huggingface.co/multimodal-reasoning-lab/Anole-Zebra-CoT)
- [Bagel-Zebra-CoT Model](https://huggingface.co/multimodal-reasoning-lab/Bagel-Zebra-CoT)