JustANormalTinkerer/

hayai-ocr-v2-alpha

1. Vision Encoder

• Backbone: google/siglip2-base-patch16-naflex (SigLIP2).
• Naflex Batching: It natively supports images of varying aspect ratios and sizes without requiring uniform resizing. The encoder outputs visual features alongside spatial shape and attention mask tensors.
• Feature Dimension: The vision encoder produces features with a dimension of 768.

2. Vision-Language Bridge (Projector)

• Architecture: A Multi-Layer Perceptron (MLPProjector) consisting of two linear layers with a GELU activation function in between.
• Function: Maps the 768-dimensional visual representation into the 512-dimensional hidden space utilized by the decoder.

3. Causal Decoder

• Layer Count: 12 decoder layers.
• Hidden Dimension: 512.
• Feed-Forward Network Dimension: 2048.
• Normalization: RMSNorm is used instead of LayerNorm, applied at the input of the attention block, feed-forward block, and before the final output head.
• Grouped-Query Attention (GQA):
  • Configured with 8 Query heads and 2 Key/Value heads. This reduces the memory footprint of the KV cache during generation.
  • Uses a head dimension of 64.
  • Applies RMSNorm directly to the Projected Query and Key states (q_norm and k_norm) before computing attention.
• SwiGLU with Clamping:
  • The feed-forward network uses a SwiGLU activation variant.
  • To prevent instability, the intermediate activations are clamped: the gate tensor is capped at a maximum of 10.0, and the up-projection tensor is clamped between -10.0 and 10.0.
• Positional Embeddings: Rotary Position Embeddings (RoPE) are applied to the Query and Key states in the attention blocks.
• Attention Masking:
  • Uses a block-causal mask.
  • The visual features act as a prefix context where all visual tokens can attend to one another, but cannot attend to future text tokens.
  • The text tokens are causally masked, meaning they can attend to all visual features and prior text tokens, but not future text tokens.

4. Tokenizer & Output Vocab

• Tokenizer: A custom UTF-8 byte-level tokenizer (ByteTokenizer).
• Vocabulary Size: 260.
  • Bytes 0–255 represent standard UTF-8 raw bytes.
  • 4 special tokens: BOS (256), EOS (257), PAD (258), and UNK (259).
• Weight Tying: The input token embedding matrix weight is shared with the final linear output projection head (output_head).