formula/FORMULA_MODEL_SUMMARY.md

# PP-FormulaNet Model Summary

## Overview

**Model**: PP-FormulaNet (PaddlePaddle Formula Recognition)
**Purpose**: Mathematical formula recognition from images
**Type**: Encoder-Decoder Vision-Language Model
**Model Size**: 694 MB (PyTorch safetensors)
**Total Parameters**: 181,896,960 (~182M)
**Total Weight Tensors**: 398
**Total Module Instances**: 279
**Unique Module Types**: 20

---

## Model Architecture

### High-Level Structure

```
Input Image (3, 768, 768)
    ↓
[Image Preprocessing] → Resize, Normalize, Pad
    ↓
[Vision Encoder] → Patch Embedding + Transformer Layers
    ↓
[Neck + Multi-Modal Projector] → Feature projection to text space
    ↓
[Text Decoder] → Transformer decoder with cross-attention
    ↓
[Language Modeling Head] → Token logits (vocab_size=50000)
    ↓
Output: LaTeX formula tokens
```

---

## Vision Encoder (Image Understanding)

### Configuration
| Parameter | Value |
|-----------|-------|
| Image Size | 768 x 768 |
| Patch Size | 16 x 16 |
| Number of Patches | (768/16)² = 2304 patches |
| Input Channels | 3 (RGB) |
| Hidden Size | 768 |
| Output Channels | 256 (post-conv) |
| Number of Layers | 12 |
| Attention Heads | 12 |
| MLP Dimension | 3072 |
| Global Attention Indexes | [2, 5, 8, 11] |
| Window Size | 14 |
| Position Encoding | Absolute + Relative |
| QKV Bias | True |

### Components

1. **Patch Embedding**
   - `model.encoder.patch_embed.projection.weight`: (768, 3, 16, 16)
   - Conv2d layer converting patches to embeddings

2. **Position Embedding**
   - `model.encoder.pos_embed`: Learnable position embeddings

3. **Neck (Post-Conv)**
   - `model.encoder.neck.conv1.weight`: First conv layer
   - `model.encoder.neck.conv2.weight`: Second conv layer
   - `model.encoder.neck.layer_norm1.*`: First layer norm
   - `model.encoder.neck.layer_norm2.*`: Second layer norm
   - Output channels: 256 → 512 → 1024

4. **Multi-Modal Projector**
   - `model.encoder.multi_modal_projector.linear_1.*`: (1024, 1024)
   - `model.encoder.multi_modal_projector.linear_2.*`: Projects to text dimension (512)
   - `model.encoder.multi_modal_projector.conv1.weight`: Conv projection
   - `model.encoder.multi_modal_projector.conv2.weight`: Conv projection

5. **Transformer Layers (12 layers)**
   Each layer contains:
   - `attn.qkv.weight/bias`: Fused QKV projection (2304 = 3*768)
   - `attn.proj.weight/bias`: Output projection
   - `attn.rel_pos_h/w`: Relative position encodings
   - `mlp.lin1.weight/bias`: First linear (3072)
   - `mlp.lin2.weight/bias`: Second linear (768)
   - `layer_norm1.*`: Pre-attention norm
   - `layer_norm2.*`: Pre-MLP norm

---

## Text Decoder (LaTeX Generation)

### Configuration
| Parameter | Value |
|-----------|-------|
| Model Type | Transformer Decoder |
| Hidden Size | 512 |
| Vocab Size | 50000 |
| Number of Layers | 8 |
| Attention Heads | 16 |
| FFN Dimension | 2048 |
| Max Position Embeddings | 2560 |
| Dropout | 0.1 |
| Activation | GELU |
| Scale Embedding | True |
| Tie Word Embeddings | False |

### Special Tokens
| Token ID | Token | Description |
|----------|-------|-------------|
| 0 | `<s>` | BOS (Begin of Sequence) |
| 1 | `<pad>` | Padding |
| 2 | `</s>` | EOS (End of Sequence) |
| 3 | `<unk>` | Unknown |
| 4 | `[START_REF]` | Start reference |
| 5 | `[END_REF]` | End reference |
| 6 | `[IMAGE]` | Image token |
| 7 | `<fragments>` | Start fragments |
| 8 | `</fragments>` | End fragments |
| 9 | `<work>` | Start work |
| 10 | `</work>` | End work |
| 11 | `[START_SUP]` | Start superscript |
| 12 | `[END_SUP]` | End superscript |
| ... | ... | More special tokens for LaTeX structure |

### Components

1. **Token Embedding**
   - `model.decoder.embed_tokens.weight`: (50000, 512)

2. **Position Embedding**
   - `model.decoder.embed_positions.weight`: Positional encodings

3. **Layer Norm Embedding**
   - `model.decoder.layernorm_embedding.*`: Input embedding normalization

4. **Decoder Layers (8 layers)**
   Each layer contains:
   - `self_attn.q_proj/k_proj/v_proj.*`: Self-attention projections
   - `self_attn.out_proj.*`: Self-attention output
   - `self_attn_layer_norm.*`: Pre-self-attention norm
   - `encoder_attn.q_proj/k_proj/v_proj.*`: Cross-attention (to encoder output)
   - `encoder_attn.out_proj.*`: Cross-attention output
   - `encoder_attn_layer_norm.*`: Pre-cross-attention norm
   - `fc1.*`: First FFN linear (2048)
   - `fc2.*`: Second FFN linear (512)
   - `final_layer_norm.*`: Pre-FFN norm

5. **Final Layer Norm**
   - `model.decoder.layer_norm.*`: Output normalization

6. **Language Modeling Head**
   - `lm_head.weight`: (50000, 512) - Projects to vocabulary logits

---

## Image Preprocessing

### Configuration (from `processor_config.json`)
| Parameter | Value |
|-----------|-------|
| Image Size | 768 x 768 |
| Resample Method | Bicubic (2) |
| Do Resize | True |
| Do Rescale | True |
| Do Normalize | True |
| Do Pad | True |
| Do Crop Margin | True |
| Do Align Long Axis | False |
| Do Thumbnail | True |
| Image Mean | [0.7931, 0.7931, 0.7931] |
| Image Std | [0.1738, 0.1738, 0.1738] |

### Preprocessing Pipeline
1. **Thumbnail**: Resize maintaining aspect ratio
2. **Crop Margin**: Remove white margins around formula
3. **Resize**: Resize to 768 x 768
4. **Rescale**: Scale pixel values to [0, 1]
5. **Normalize**: Apply mean/std normalization
6. **Pad**: Pad to target size if needed

---

## Key Tensor Shapes

| Component | Tensor | Shape |
|-----------|--------|-------|
| Patch Embedding | `model.encoder.patch_embed.projection.weight` | (768, 3, 16, 16) |
| QKV Projection | `model.encoder.layers.0.attn.qkv.weight` | (2304, 768) |
| Token Embedding | `model.decoder.embed_tokens.weight` | (50000, 512) |
| Projector | `model.encoder.multi_modal_projector.linear_1.weight` | (1024, 1024) |
| LM Head | `lm_head.weight` | (50000, 512) |

---

## Model Tree Structure

```
PPFormulaNetForConditionalGeneration
├── model
│   ├── encoder (Vision Encoder)
│   │   ├── patch_embed
│   │   │   └── projection (Conv2d)
│   │   ├── pos_embed (Learnable)
│   │   ├── neck
│   │   │   ├── conv1
│   │   │   ├── conv2
│   │   │   ├── layer_norm1
│   │   │   └── layer_norm2
│   │   ├── multi_modal_projector
│   │   │   ├── linear_1
│   │   │   ├── linear_2
│   │   │   ├── conv1
│   │   │   └── conv2
│   │   └── layers (12 Transformer layers)
│   │       └── [0-11]
│   │           ├── attn
│   │           │   ├── qkv
│   │           │   ├── proj
│   │           │   ├── rel_pos_h
│   │           │   └── rel_pos_w
│   │           ├── mlp
│   │           │   ├── lin1
│   │           │   └── lin2
│   │           ├── layer_norm1
│   │           └── layer_norm2
│   └── decoder (Text Decoder)
│       ├── embed_tokens
│       ├── embed_positions
│       ├── layernorm_embedding
│       ├── layers (8 Transformer layers)
│       │   └── [0-7]
│       │       ├── self_attn
│       │       │   ├── q_proj
│       │       │   ├── k_proj
│       │       │   ├── v_proj
│       │       │   └── out_proj
│       │       ├── self_attn_layer_norm
│       │       ├── encoder_attn
│       │       │   ├── q_proj
│       │       │   ├── k_proj
│       │       │   ├── v_proj
│       │       │   └── out_proj
│       │       ├── encoder_attn_layer_norm
│       │       ├── fc1
│       │       ├── fc2
│       │       └── final_layer_norm
│       └── layer_norm
└── lm_head (Linear)
```

---

## Inference Usage

### HuggingFace Transformers

```python
from PIL import Image
from transformers import AutoProcessor, PPFormulaNetForConditionalGeneration

# Load model and processor
model_path = "model/formula"
model = PPFormulaNetForConditionalGeneration.from_pretrained(model_path)
processor = AutoProcessor.from_pretrained(model_path)

# Preprocess image
image = Image.open("formula.png").convert("RGB")
inputs = processor(images=image, return_tensors="pt")

# Generate LaTeX
outputs = model.generate(**inputs)
result = processor.post_process(outputs)
print(result)  # LaTeX string
```

### Key Points
- Use `model.generate()` for autoregressive decoding (inference mode)
- Use `model.forward()` only for training (requires decoder input_ids)
- The processor handles image preprocessing and text post-processing

### Decoding Strategy
**Default: Greedy Decoding** (beam_size=1)
- `num_beams`: 1 (default)
- `do_sample`: False (default)
- `max_length`: 1537 tokens

To enable beam search for better quality:
```python
outputs = model.generate(**inputs, num_beams=5, early_stopping=True)
```

Common generation parameters:
- `num_beams`: Number of beams for beam search (default=1)
- `max_length`: Maximum sequence length (default=1537)
- `early_stopping`: Stop when all beams reach EOS (default=False)
- `length_penalty`: Penalty for longer sequences (default=1.0)
- `do_sample`: Enable sampling (default=False)
- `temperature`: Sampling temperature (default=1.0)
- `top_k`: Top-k sampling (default=50)
- `top_p`: Nucleus sampling (default=1.0)

---

## Porting Considerations for MLX

### Key Challenges

1. **Relative Position Encoding**: The vision encoder uses both absolute and relative position encodings (`rel_pos_h`, `rel_pos_w`). MLX's attention mechanism may need custom implementation.

2. **Global Attention**: Layers at indexes [2, 5, 8, 11] use global attention. This may require special handling.

3. **Window Attention**: Window size of 14 suggests window-based attention (similar to Swin Transformer).

4. **Cross-Attention**: The decoder has cross-attention to encoder outputs, requiring careful memory management.

5. **Autoregressive Generation**: The decoder generates tokens one at a time, which may be slow on MLX without optimization.

6. **Large Vocabulary**: 50,000 vocab size means the LM head is large (50000 x 512).

### MLX Module Structure

The MLX implementation would need:
- `PPFormulaNetVisionEncoder`: Vision encoder with patch embedding and transformer layers
- `PPFormulaNetNeck`: Post-conv layers for feature refinement
- `PPFormulaNetMultiModalProjector`: Projects vision features to text space
- `PPFormulaNetDecoder`: Transformer decoder with self and cross attention
- `PPFormulaNetLMHead`: Linear layer for vocabulary projection
- `PPFormulaNetModel`: Combined model
- `PPFormulaNetForConditionalGeneration`: Root model with generation support

---

## Module Types Breakdown

| Module Type | Count | Description |
|-------------|-------|-------------|
| Linear | 131 | Linear projection layers |
| LayerNorm | 50 | Layer normalization layers |
| GELUActivation | 20 | GELU activation functions |
| PPFormulaNetAttention | 16 | Attention mechanisms (12 vision + 4 decoder cross-attention) |
| PPFormulaNetVisionLayer | 12 | Vision transformer encoder layers |
| PPFormulaNetVisionAttention | 12 | Vision-specific attention with relative position encoding |
| PPFormulaNetMLPBlock | 12 | MLP blocks in vision encoder |
| PPFormulaNetDecoderLayer | 8 | Text decoder layers |
| Conv2d | 5 | Convolutional layers (patch embed + neck + projector) |
| ModuleList | 2 | Module lists for repeated layers |
| PPFormulaNetLayerNorm | 2 | Custom layer norm implementations |
| PPFormulaNetForConditionalGeneration | 1 | Root model class |
| PPFormulaNetModel | 1 | Core model (encoder + decoder) |
| PPFormulaNetTextModel | 1 | Text decoder wrapper |
| PPFormulaNetScaledWordEmbedding | 1 | Scaled word embedding layer |
| PPFormulaNetLearnedPositionalEmbedding | 1 | Learnable position embeddings |
| PPFormulaNetVisionModel | 1 | Vision encoder wrapper |
| PPFormulaNetPatchEmbeddings | 1 | Patch embedding layer |
| PPFormulaNetVisionNeck | 1 | Vision neck (post-conv) |
| PPFormulaNetMultiModalProjector | 1 | Multi-modal projection layer |

---

## Files in `model/formula/`

| File | Description |
|------|-------------|
| `model.safetensors` | PyTorch weights (694 MB, 398 tensors) |
| `config.json` | Model architecture configuration |
| `processor_config.json` | Image preprocessing configuration |
| `tokenizer_config.json` | Tokenizer special tokens and settings |
| `tokenizer.json` | Tokenizer vocabulary |
| `generation_config.json` | Generation parameters |
| `inference.yml` | Inference configuration |
| `README.md` | Model documentation (Apache 2.0 license) |
| `formula_model_tree.json` | Generated model tree structure |

---

## References

- HuggingFace Model: `PaddlePaddle/PP-FormulaNet_plus-L_safetensors`
- Original Paper: PP-FormulaNet (PaddlePaddle)
- Transformers Class: `PPFormulaNetForConditionalGeneration`
- Processor Class: `PPFormulaNetProcessor`