Sim Cot Supervised Implicit Chain Of Thought

Stabilize Implicit Reasoning with Step-Level Supervision

Build chain-of-thought reasoning into latent token representations rather than explicit text tokens, maintaining inference efficiency while achieving explicit CoT accuracy.

The Problem: Latent Instability in Implicit CoT

Implicit chain-of-thought methods compress reasoning into hidden states rather than generated text tokens, achieving 2.3× speedup over explicit CoT. However, this efficiency comes at a cost: training becomes unstable as you scale reasoning complexity. When you increase implicit reasoning tokens from 3 to 5, models experience training collapse where latent representations become nearly identical, information dissolves, and the model fails to learn meaningful reasoning.

The root cause: current implicit CoT methods provide only answer-level supervision. The model receives feedback only on whether the final answer matches, but nothing guides the intermediate latent states. Without step-level guidance, these representations drift into semantic homogeneity. The latent space geometrically deteriorates—distances between reasoning steps shrink while representations drift from the vocabulary embedding space.

Core Concept: Supervised Implicit Chain-of-Thought

SIM-CoT solves latent instability through auxiliary step-level supervision during training. An auxiliary decoder aligns each implicit latent state with its corresponding explicit reasoning step. This dual-objective training distributes learning signals across the entire reasoning chain rather than concentrating them only at the final answer.

The key insight: you can supervise intermediate representations during training without paying an inference cost. The auxiliary decoder exists only during training—at inference time, you remove it and get pure latent reasoning with all the efficiency benefits.

This enables:

• Stable training with 8-16 implicit tokens (prior methods collapsed at 5)
• Per-step interpretability by projecting latent tokens onto reasoning vocabulary
• Performance gains of +3 to +8 percentage points over prior implicit methods
• 2.3× token efficiency maintained at inference

Architecture Overview

Two-Phase Inference Structure:

• Implicit phase: Model generates K fixed latent reasoning steps, extracting the hidden state at each step as a latent vector z_k
• Explicit phase: After K implicit steps, model switches to standard vocabulary decoding for the final answer

Training-Time Architecture:

• Main model: Transformer LLM with embedded implicit reasoning capability
• Auxiliary decoder: Lightweight module that maps each latent z_k to vocabulary logits for reasoning step s_k
• Dual loss: Step-level loss (auxiliary decoder outputs vs. explicit reasoning steps) + answer-level loss (final generated answer vs. ground truth)

Latent Representation Characteristics:

• Each z_k is the hidden state extracted after step k of implicit reasoning
• Representations remain in hidden space (typically 768d for GPT-2, up to 8192d for LLaMA 8B)
• No vocabulary projection during inference—pure latent computation preserves speedup

Implementation Strategy

Step 1: Design the Implicit Reasoning Curriculum

Before building the architecture, establish how many implicit steps you need and design training progression. Implicit reasoning works best for tasks with clear step-by-step solutions (math, logic, symbolic reasoning).

# pseudo code - curriculum learning schedule
curriculum_steps = [
    {"stage": 1, "implicit_tokens": 1, "explicit_percentage": 1.0, "epochs": 5},
    {"stage": 2, "implicit_tokens": 2, "explicit_percentage": 0.8, "epochs": 5},
    {"stage": 3, "implicit_tokens": 3, "explicit_percentage": 0.6, "epochs": 5},
    {"stage": 4, "implicit_tokens": 5, "explicit_percentage": 0.4, "epochs": 10},
    {"stage": 5, "implicit_tokens": 8, "explicit_percentage": 0.2, "epochs": 15},
]

# Start with mostly explicit reasoning, gradually shift to implicit
# Each stage uses curriculum learning to swap explicit steps with implicit latents
# The model learns to pack meaning into fixed-length latent vectors over time

Curriculum learning is essential. Don't attempt 8 implicit tokens from epoch 1—the model will collapse immediately. Start with 1-2 implicit tokens and explicit reasoning for most steps, then progressively increase implicit capacity. This gives the latent space time to develop meaningful semantic structure.

Step 2: Implement the Auxiliary Decoder Module

The auxiliary decoder is a lightweight projection layer that maps latent vectors to vocabulary logits. It exists only during training and enables step-level supervision.

import torch
import torch.nn as nn

class AuxiliaryDecoder(nn.Module):
    """
    Maps latent representations to vocabulary logits for training supervision.
    Removed at inference time to preserve computational efficiency.
    """
    def __init__(self, hidden_size, vocab_size, num_layers=1):
        super().__init__()
        self.hidden_size = hidden_size
        self.vocab_size = vocab_size

        # Simple linear projection or light MLP
        if num_layers == 1:
            self.decoder = nn.Linear(hidden_size, vocab_size)
        else:
            layers = []
            layers.append(nn.Linear(hidden_size, hidden_size))
            layers.append(nn.ReLU())
            layers.append(nn.Linear(hidden_size, vocab_size))
            self.decoder = nn.Sequential(*layers)

    def forward(self, latent_vectors):
        """
        Args:
            latent_vectors: (batch_size, num_steps, hidden_size)

        Returns:
            logits: (batch_size, num_steps, vocab_size)
        """
        return self.decoder(latent_vectors)

# Usage during training:
# aux_decoder = AuxiliaryDecoder(hidden_size=768, vocab_size=50257)
# step_logits = aux_decoder(latent_states)  # supervise against explicit steps
# At inference: remove auxiliary decoder entirely

The auxiliary decoder should be minimal—a single linear layer for most cases, or a 2-3 layer MLP for larger models. Heavy architectures waste training compute since you discard them anyway.

Step 3: Extract Latent Representations During Forward Pass

Modify your model's forward pass to capture hidden states at specific token positions where implicit reasoning steps occur.

class ImplicitCoTModel(nn.Module):
    """
    Language model extended with implicit chain-of-thought capability.
    """
    def __init__(self, base_model, hidden_size, vocab_size, num_implicit_steps):
        super().__init__()
        self.model = base_model  # GPT-2, LLaMA, etc.
        self.num_implicit_steps = num_implicit_steps
        self.hidden_size = hidden_size

        # Track which positions correspond to implicit steps
        # For example, if prompt is 50 tokens, implicit steps at 51-58
        self.implicit_step_positions = None

        # Auxiliary decoder (training only)
        self.aux_decoder = AuxiliaryDecoder(hidden_size, vocab_size)

    def forward(self, input_ids, labels=None, return_latents=False, training=True):
        """
        Args:
            input_ids: (batch_size, seq_len)
            labels: (batch_size, seq_len) - for standard LM loss
            return_latents: whether to extract and return latent vectors
            training: whether in training mode

        Returns:
            loss: combined answer-level + step-level loss (if training)
            logits: final output logits for answer generation
            latent_states: extracted latents (if return_latents=True)
        """
        # Forward through base model, capturing all hidden states
        outputs = self.model(input_ids, output_hidden_states=True)
        hidden_states = outputs.hidden_states[-1]  # (batch, seq_len, hidden_size)
        logits = outputs.logits

        # Extract latent vectors at implicit step positions
        latent_states = hidden_states[:, -self.num_implicit_steps:, :]

        # Answer-level loss: standard language modeling loss
        answer_loss = None
        if labels is not None:
            # Standard cross-entropy on final answer tokens
            answer_loss = torch.nn.functional.cross_entropy(
                logits.view(-1, logits.size(-1)),
                labels.view(-1),
                reduction='mean'
            )

        # Step-level loss: auxiliary decoder supervision (training only)
        step_loss = None
        if training and labels is not None:
            step_logits = self.aux_decoder(latent_states)
            # Create targets from explicit reasoning steps in the input sequence
            # For example, if implicit steps 3-5 correspond to explicit steps 20-22
            step_targets = labels[:, -self.num_implicit_steps:]  # simplified mapping
            step_loss = torch.nn.functional.cross_entropy(
                step_logits.view(-1, step_logits.size(-1)),
                step_targets.view(-1),
                reduction='mean'
            )

        # Combine losses
        total_loss = answer_loss
        if step_loss is not None:
            # Weight the step-level supervision (typically 0.5-1.0 weight)
            total_loss = answer_loss + 0.7 * step_loss

        outputs_dict = {
            'loss': total_loss,
            'logits': logits,
            'answer_loss': answer_loss,
            'step_loss': step_loss,
        }

        if return_latents:
            outputs_dict['latent_states'] = latent_states

        return outputs_dict

# Usage:
# model = ImplicitCoTModel(base_model, hidden_size=768, vocab_size=50257, num_implicit_steps=5)
# outputs = model(input_ids, labels=targets, training=True)
# loss = outputs['loss']
# loss.backward()

Step 4: Implement Curriculum Learning Training Loop

Progressively transition from explicit to implicit reasoning to stabilize training.

import torch.optim as optim
from torch.utils.data import DataLoader

def train_with_curriculum(model, train_dataset, curriculum_schedule, device):
    """
    Train model using curriculum learning to gradually introduce implicit reasoning.

    Args:
        model: ImplicitCoTModel instance
        train_dataset: Dataset with (input_ids, labels, explicit_steps) tuples
        curriculum_schedule: List of dicts with stage, implicit_tokens, explicit_percentage, epochs
        device: torch device
    """
    optimizer = optim.AdamW(model.parameters(), lr=2e-4)

    for stage in curriculum_schedule:
        print(f"\nStage {stage['stage']}: {stage['implicit_tokens']} implicit tokens, "
              f"{stage['explicit_percentage']*100:.0f}% explicit")

        # Set number of implicit steps for this stage
        model.num_implicit_steps = stage['implicit_tokens']

        # Create dataloader with curriculum applied
        train_loader = DataLoader(
            train_dataset,
            batch_size=128,
            shuffle=True,
            collate_fn=lambda x: apply_curriculum_collate(
                x,
                implicit_tokens=stage['implicit_tokens'],
                explicit_percentage=stage['explicit_percentage']
            )
        )

        # Train for specified epochs in this stage
        for epoch in range(stage['epochs']):
            total_loss = 0.0
            num_batches = 0

            for batch in train_loader:
                input_ids = batch['input_ids'].to(device)
                labels = batch['labels'].to(device)

                optimizer.zero_grad()
                outputs = model(input_ids, labels=labels, training=True)
                loss = outputs['loss']

                # Monitor loss components
                if outputs['step_loss'] is not None:
                    print(f"  Batch: answer_loss={outputs['answer_loss']:.4f}, "
                          f"step_loss={outputs['step_loss']:.4f}")

                loss.backward()
                torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
                optimizer.step()

                total_loss += loss.item()
                num_batches += 1

            avg_loss = total_loss / num_batches
            print(f"  Epoch {epoch+1}/{stage['epochs']}: loss={avg_loss:.4f}")

def apply_curriculum_collate(batch, implicit_tokens, explicit_percentage):
    """
    Apply curriculum learning: replace some explicit steps with implicit latents.

    For a sample with 8 reasoning steps and explicit_percentage=0.6:
    - Keep 5 explicit steps
    - Replace 3 steps with implicit latents
    """
    # Simplified: in practice, you'd manipulate token sequences
    # to mark which positions should be implicit vs explicit
    input_ids = torch.stack([item['input_ids'] for item in batch])
    labels = torch.stack([item['labels'] for item in batch])

    return {
        'input_ids': input_ids,
        'labels': labels,
        'implicit_tokens': implicit_tokens,
        'explicit_percentage': explicit_percentage,
    }

# Training call:
# curriculum = [
#     {"stage": 1, "implicit_tokens": 1, "explicit_percentage": 1.0, "epochs": 5},
#     {"stage": 2, "implicit_tokens": 2, "explicit_percentage": 0.8, "epochs": 5},
#     {"stage": 3, "implicit_tokens": 5, "explicit_percentage": 0.6, "epochs": 10},
# ]
# train_with_curriculum(model, train_dataset, curriculum, device='cuda')

Step 5: Monitor Latent Space Health with Geometric Diagnostics

During training, measure whether latent representations are developing proper semantic structure. Failed models show geometric collapse; healthy models maintain diverse, dispersed latent vectors.

import torch
import numpy as np

class LatentSpaceMonitor:
    """
    Track geometric properties of latent representations to detect training failure.
    Healthy models: inter-latent distances ~32.81
    Failed models: inter-latent distances ~4.21 (collapse)
    """

    def __init__(self):
        self.metrics = {
            'inter_latent_distances': [],
            'drift_from_embeddings': [],
            'variance_per_dimension': [],
        }

    def compute_metrics(self, latent_states, embedding_matrix):
        """
        Args:
            latent_states: (batch_size, num_steps, hidden_size)
            embedding_matrix: (vocab_size, hidden_size)

        Returns:
            metrics: dict with health indicators
        """
        batch_size, num_steps, hidden_size = latent_states.shape

        # 1. Inter-latent distances: average L2 distance between different steps
        latent_flat = latent_states.view(-1, hidden_size)  # (batch*steps, hidden_size)
        distances = torch.cdist(latent_flat, latent_flat, p=2)
        # Extract off-diagonal (distances between different steps)
        inter_latent_distance = distances[distances > 0].mean().item()

        # 2. Drift from embeddings: mean distance to vocabulary embedding center
        embedding_center = embedding_matrix.mean(dim=0)
        drift = torch.norm(latent_flat - embedding_center, dim=1).mean().item()

        # 3. Per-dimension variance: should be diverse, not collapsed
        variance = latent_states.var(dim=(0, 1)).mean().item()

        metrics = {
            'inter_latent_distance': inter_latent_distance,
            'drift_from_embeddings': drift,
            'dimension_variance': variance,
            'healthy': inter_latent_distance > 20 and variance > 0.1,
        }

        return metrics

    def check_collapse(self, metrics, tolerance=10.0):
        """
        Determine if latent space is collapsing.
        Return True if intervention needed.
        """
        if metrics['inter_latent_distance'] < tolerance:
            return True  # Collapse detected
        if metrics['dimension_variance'] < 0.05:
            return True  # Variance too low
        return False

# Usage during training:
# monitor = LatentSpaceMonitor()
# outputs = model(input_ids, labels=labels, training=True, return_latents=True)
# latents = outputs['latent_states']
# metrics = monitor.compute_metrics(latents, model.model.embeddings.word_embeddings.weight)
#
# if monitor.check_collapse(metrics):
#     print(f"WARNING: Latent collapse detected. Distance: {metrics['inter_latent_distance']:.2f}")
#     # Reduce learning rate, increase step-level loss weight, or backtrack curriculum

Step 6: Remove Auxiliary Decoder at Inference

Once training completes, the auxiliary decoder is no longer needed. Remove it before deployment.

def prepare_for_inference(model):
    """
    Remove training-only components and optimize for inference.
    """
    # Delete auxiliary decoder
    if hasattr(model, 'aux_decoder'):
        del model.aux_decoder

    # Set to eval mode
    model.eval()

    # Optional: quantize for faster inference
    # model = torch.quantization.quantize_dynamic(model)

    return model

# Usage:
# model = prepare_for_inference(model)
#
# # Inference: get answer directly from latent reasoning
# with torch.no_grad():
#     outputs = model(input_ids, training=False)
#     answer_logits = outputs['logits'][:, -1, :]  # take final token
#     answer_token = answer_logits.argmax(dim=-1)

Practical Guidance

Hyperparameters and Configuration

When to Use SIM-CoT

Use SIM-CoT when you have:

1. Long reasoning chains: Tasks requiring 5+ distinct reasoning steps (math, logic, planning)
2. Latency constraints: You need faster inference than explicit CoT (2.3× speedup is substantial)
3. Limited compute for inference: Implicit reasoning uses fewer tokens at inference time
4. Step-level supervision data: Explicit intermediate reasoning steps are available in training data
5. Stability over pure implicit methods: You need reliable training without collapse
6. Interpretability requirements: You can visualize implicit steps via the auxiliary decoder during debugging

When NOT to Use SIM-CoT

Do not use SIM-CoT for:

1. Simple tasks: One-step or two-step reasoning doesn't benefit from implicit methods. Use standard language modeling.
2. Unstructured reasoning: Tasks without clear intermediate steps (narrative generation, translation). Implicit structure assumes step-wise computation.
3. Streaming generation: Implicit tokens are fixed-length and pre-computed, not adaptive to input. Explicit CoT is more flexible.
4. Absence of explicit supervision: Without ground-truth intermediate steps during training, you lose the key advantage of step-level supervision.
5. Model serving constraints: Some inference frameworks don't support custom hidden state extraction. Verify before deployment.
6. Cold training start: Don't skip curriculum learning. Jumping directly to many implicit tokens causes immediate collapse.
7. Very large language models: Auxiliary decoder training overhead increases with hidden size. For 70B+ models, ROI may not justify added complexity.

Common Pitfalls and Fixes

Pitfall: Immediate training collapse

• Root cause: Insufficient curriculum. Attempting too many implicit tokens too early.
• Fix: Start with 1-2 implicit tokens for 10 epochs, then gradually increase. Monitor latent space metrics.

Pitfall: Step-level loss dominates, answer accuracy drops

• Root cause: Step-level loss weight too high.
• Fix: Reduce step loss weight from 1.0 to 0.5-0.7. Answer-level loss should drive final performance.

Pitfall: Latent space diverges, distances grow without bound

• Root cause: Learning rate too high or gradient clipping disabled.
• Fix: Lower learning rate to 1e-4, enable gradient clipping (max_norm=1.0).

Pitfall: Auxiliary decoder improves during training but answer accuracy stalls

• Root cause: Auxiliary decoder and final answer are decoupling. Step supervision isn't helping end-to-end.
• Fix: Increase implicit tokens more slowly in curriculum. Use lower step-level loss weight to avoid distraction.

Pitfall: Inference speed doesn't match expected 2.3× speedup

• Root cause: Auxiliary decoder wasn't fully removed or other overheads present.
• Fix: Verify auxiliary decoder is deleted. Profile latent vs. token generation separately. Check batch size (smaller batches reduce overhead visibility).

Reference

Paper: SIM-CoT: Supervised Implicit Chain-of-Thought arXiv: https://arxiv.org/abs/2509.20317 Key contributions: Step-level supervision for latent reasoning stability, 2.3× inference speedup, interpretability via auxiliary decoder projection, strong performance on math reasoning tasks (GSM8K, SVAMP, GSM-Hard, MultiArith).