| { |
| "name": "OMNI-LITE (Unified Sparse-Multimodal Transformer)", |
| "layers": [ |
| { |
| "type": "Conv2d", |
| "params": { |
| "in_channels": 3, |
| "out_channels": 1024, |
| "kernel_size": 14, |
| "stride": 14, |
| "note": "Vision Patch Embedding for ViT encoder" |
| } |
| }, |
| { |
| "type": "TransformerBlock", |
| "params": { |
| "embed_dim": 1024, |
| "num_heads": 16, |
| "ff_dim": 4096, |
| "depth": 12, |
| "note": "Lightweight Vision Transformer (ViT) Backbone" |
| } |
| }, |
| { |
| "type": "TransformerBlock", |
| "params": { |
| "type": "PerceiverResampler", |
| "num_latents": 64, |
| "embed_dim": 2048, |
| "note": "Maps visual features to text latent space" |
| } |
| }, |
| { |
| "type": "Linear", |
| "params": { |
| "in_features": 32000, |
| "out_features": 2048, |
| "note": "Text Token Embedding layer" |
| } |
| }, |
| { |
| "type": "TransformerBlock", |
| "params": { |
| "type": "GQA_MoE_Layer", |
| "repeat": 24, |
| "num_experts": 16, |
| "top_k": 2, |
| "hidden_dim": 2048, |
| "num_heads": 32, |
| "num_kv_heads": 8, |
| "rope_dim": 64, |
| "note": "Shared Backbone: 480M active parameters per token" |
| } |
| }, |
| { |
| "type": "Linear", |
| "params": { |
| "in_features": 2048, |
| "out_features": 32000, |
| "note": "Causal Language Modeling (CLM) Head" |
| } |
| }, |
| { |
| "type": "Linear", |
| "params": { |
| "in_features": 2048, |
| "out_features": 64, |
| "note": "Rectified Flow-Matching (RFM) Head for DiT Latents" |
| } |
| }, |
| { |
| "type": "Conv2d", |
| "params": { |
| "in_channels": 4, |
| "out_channels": 3, |
| "kernel_size": 3, |
| "stride": 1, |
| "note": "VQ-VAE Decoder for 8x8 Latent Reconstruction" |
| } |
| } |
| ], |
| "explanation": "OMNI-LITE utilizes a Sparse MoE backbone to minimize active compute (480M params) while maintaining 2.5B knowledge capacity, fitting within the 6GB VRAM limit when quantized via NF4/AWQ. Grouped-Query Attention (GQA) significantly reduces the KV-cache footprint for edge deployment. The Perceiver Resampler allows the model to treat visual inputs as a fixed set of tokens within the Causal Transformer's context window. For generation, the dual-head design supports standard autoregressive text while a separate Flow-Matching head handles Diffusion Transformer (DiT) logic within the same latent space, ensuring hardware-agnostic efficiency through RFM's low-step count." |
| } |
