FurkanNar/noodle_baseline

Noodle Baseline

An alternative latent reasoning architecture for math reasoning tasks using latent chain-of-thought reasoning.

Architecture

The model consists of three main components:

1. Prompt Encoder

• Encodes input tokens into latent vectors
• Uses GPT-2 small weights (768 hidden, 12 layers, 12 heads)
• Projects to latent dimension (128) for reasoning

2. Latent Reasoning

• Custom reasoning module with learnable latent tokens
• Perceiver-style cross-attention resampling
• Iterative self-attention for multi-step reasoning
• Kaiming initialization for stable training

3. Decoder

• Decodes latent tokens to output tokens
• Uses GPT-2 small weights
• Autoregressive generation capability

Model Configuration

d_model = 768           # GPT-2 small hidden size
n_layers = 12           # Number of transformer layers
n_heads = 12            # Attention heads
latent_dim = 128        # Latent space dimension
latent_heads = 4        # Latent reasoner attention heads
n_latent_tokens = 16    # Number of latent reasoning tokens
max_latent_steps = 4    # Maximum reasoning steps
vocab_size = 50257      # GPT-2 vocabulary

Training Setup

• Dataset: GSM8K (1000 samples)
• Training steps: 200
• Warmup steps: 20
• Batch size: 16
• Learning rate: 2e-4 with cosine scheduler
• Mixed precision: FP16 enabled for CUDA
• Gradient clipping: 1.0

Training

The training script:

• Loads GPT-2 weights for encoder and decoder
• Validates model architecture before training
• Saves checkpoints to checkpoint.safetensors every 50 steps
• Uses mixed precision training for speed

Key Features

• Pre-trained weights: GPT-2 small for encoder/decoder
• Custom reasoning: LatentReasoner with iterative reasoning
• Fast training: Mixed precision, optimized batch size
• Validation: Pre-training shape validation to catch errors
• Checkpoints: Automatic saving every 50 steps

Notes

• The Latent Reasoner is trained from scratch with Kaiming initialization
• GPT-2 weights are adapted for the custom architecture (QKV splitting, transposition)
• Training is optimized for speed with reduced steps and mixed precision
• This is a custom latent reasoning architecture