large concept models

TLLM Architecture - Transformer Language Model for SONAR Embeddings

Model Configuration

• Input: SONAR embeddings (1024 dimensions)
• Hidden dimension: 512 (n_embd)
• Sequence length: 128 (block_size)
• Transformer layers: 8 (n_layer)
• Attention heads: 8 (n_head)
• Dropout: 0.2
• Architecture: Pre-norm GPT-style transformer

Mermaid Architecture Diagram

graph TD
    A[Input: SONAR Embeddings<br/>Shape: batch_size × seq_len × 1024] --> B[PreNet]
    
    subgraph PreNet["PreNet Module"]
        B1[LayerNorm<br/>Normalize SONAR embeddings] --> B2[Dense Layer<br/>1024 → 512 dimensions<br/>Project to n_embd]
    end
    
    B --> C[Add Positional Embeddings<br/>Learned embeddings<br/>Shape: seq_len × 512]
    C --> D[Dropout<br/>rate=0.2<br/>Applied if training]
    
    D --> E[Transformer Block 1]
    
    subgraph TB1["Transformer Block (×8 layers)"]
        E1[Input<br/>Shape: batch_size × seq_len × 512] --> E2[LayerNorm 1<br/>Pre-norm before attention<br/>epsilon=1e-5]
        E2 --> E3[Causal Self-Attention]
        
        subgraph CSA["Causal Self-Attention"]
            E3a[QKV Projection<br/>512 → 3×512<br/>Split into Q, K, V] --> E3b[Multi-Head Attention<br/>8 heads, head_dim=64<br/>Causal masking applied]
            E3b --> E3c[Attention Dropout<br/>rate=0.2 if training]
            E3c --> E3d[Output Projection<br/>512 → 512]
            E3d --> E3e[Attention Dropout<br/>rate=0.2 if training]
        end
        
        E1 --> E4[Residual Connection 1<br/>Add input to attention output]
        E3 --> E4
        
        E4 --> E5[LayerNorm 2<br/>Pre-norm before MLP<br/>epsilon=1e-5]
        E5 --> E6[MLP Block]
        
        subgraph MLP["MLP Block"]
            E6a[Dense Layer 1<br/>512 → 2048<br/>4 × n_embd expansion] --> E6b[GELU Activation<br/>Gaussian Error Linear Unit]
            E6b --> E6c[Dense Layer 2<br/>2048 → 512<br/>Project back to n_embd]
            E6c --> E6d[MLP Dropout<br/>rate=0.2 if training]
        end
        
        E4 --> E7[Residual Connection 2<br/>Add previous output to MLP output]
        E6 --> E7
    end
    
    E --> F[Transformer Block 2]
    F --> G[...]
    G --> H[Transformer Block 8]
    
    H --> I[Final LayerNorm<br/>Post-processing normalization<br/>epsilon=1e-5]
    
    I --> J[PostNet]
    
    subgraph PostNet["PostNet Module"]
        J1[Dense Layer<br/>512 → 1024 dimensions<br/>Project back to SONAR space]
    end
    
    J --> K[Output: Predicted SONAR Embeddings<br/>Shape: batch_size × seq_len × 1024]
    
    style A fill:#e1f5fe
    style K fill:#e8f5e8
    style PreNet fill:#fff3e0
    style PostNet fill:#fff3e0
    style TB1 fill:#f3e5f5
    style CSA fill:#e3f2fd
    style MLP fill:#fce4ec

Architecture Details

Layer Types and Components:

1. Input Processing:

   • SONAR embeddings (1024-dim) → PreNet

2. PreNet (Normalization + Projection):

   • LayerNorm: Normalizes input SONAR embeddings
   • Dense: Projects from 1024 → 512 dimensions

3. Positional Encoding:

   • Learned positional embeddings (max length: 128)
   • Added to projected embeddings

4. Transformer Blocks (8 layers):

   • Pre-LayerNorm: Applied before each sublayer
   • Causal Self-Attention:
     • 8 attention heads (64 dimensions each)
     • Causal masking for autoregressive generation
     • QKV projection + attention + output projection
     • Dropout after attention weights and output
   • MLP Block:
     • Expansion: 512 → 2048 → 512
     • GELU activation
     • Dropout after final projection
   • Residual Connections: Around both attention and MLP

5. Output Processing:

   • Final LayerNorm after all transformer blocks
   • PostNet: Projects back to SONAR space (512 → 1024)

Normalization Strategy:

• Pre-norm architecture: LayerNorm applied before each sublayer
• 3 LayerNorm operations per block:
  • Before self-attention
  • Before MLP
  • Final normalization (after all blocks)
• PreNet normalization: Input SONAR embeddings

Training Objective:

• Combined Loss Function:
  • MSE loss (α=0.7 weight)
  • Cosine similarity loss (1-α=0.3 weight)
  • Norm constraint (0.5 weight)
  • Total: 0.7*MSE + 0.3*cosine_loss + 0.5*norm_constraint