codi-gpt2-prontoqa-latent

CODI — GPT-2 ProntoQA (Latent Reasoning)

A CODI (Chain-of-thought Distillation) model trained on ProntoQA for latent chain-of-thought reasoning. The model wraps GPT-2 with LoRA adapters and a distillation objective that compresses explicit chain-of-thought steps into latent embeddings.

Quick Start

pip install transformers peft huggingface_hub torch

from huggingface_hub import snapshot_download
import sys

# Download the model
local_dir = snapshot_download("simon-pltk/codi-gpt2-prontoqa-latent")

# Load it
sys.path.insert(0, local_dir)
from load_model import load_codi_model

model = load_codi_model(local_dir, device="cuda")

Architecture

CODI is a custom torch.nn.Module that:

1. Wraps a base GPT-2 model loaded via AutoModelForCausalLM
2. Applies LoRA adapters (rank=128, alpha=16) for parameter-efficient tuning
3. Generates latent embeddings that replace explicit chain-of-thought tokens
4. Uses a layer-wise distillation loss (SmoothL1) to align the student (latent) representations with a teacher (explicit CoT) across all layers

Training Details

Final Metrics

Files