latent-thoughts-tuning

Enable models to reason in continuous latent space via context-prediction fusion, combining hidden state context with vocabulary embeddings to maintain scaling across model sizes.

Zur Installation springen

Repository: ADu2021/skillXiv
Letzte Quellaktivität: 24. März 2026 um 19:42
Erkannte Sprache von SKILL.md: Englisch
Sterne: 6
Forks: 0

Installationsoptionen

Standardmäßig ist der Prompt ausgewählt, der zuerst die Quelle prüft. Sie können zu einem direkten Befehl wechseln oder eine lokale Kopie herunterladen.

Quelldateien prüfen

Lesen Sie SKILL.md und alle von SkillsMP angezeigten Begleitdateien, bevor Sie sich für eine Installation entscheiden.

Zurück zur Suche Zurück zur Installation

name	latent-thoughts-tuning
title	Latent Thoughts Tuning: Bridging Context and Reasoning with Fused Information
version	0.0.2
engine	skillxiv-v0.0.2-claude-opus-4.6
license	MIT
url	https://arxiv.org/abs/2602.10229
keywords	["Latent Reasoning","Chain-of-Thought","Hidden State Fusion","Model Scaling","Adaptive Reasoning Allocation"]
description	Enable models to reason in continuous latent space via context-prediction fusion, combining hidden state context with vocabulary embeddings to maintain scaling across model sizes.

Latent Thoughts Tuning: Bridging Context and Reasoning with Fused Information

Problem Context

Prior latent reasoning approaches face critical limitations: directly reusing hidden states as input embeddings causes distribution mismatch and feature collapse, especially in larger models. Existing methods like Coconut degrade severely with scale (50.3% → 41.5% accuracy), while fixed reasoning schedules ignore varying problem difficulty.

Core Concept

Latent Thoughts Tuning (LT-Tuning) enables reasoning in continuous latent space through Context-Prediction Fusion—combining contextual information from hidden states with predictive semantic guidance from vocabulary embeddings. This mitigates feature collapse while enabling sample-specific reasoning diversity.

Architecture Overview

Three-Stage Curriculum: Standard CoT fine-tuning → confidence-driven token insertion → fusion mechanism training
Confidence-Driven Insertion: Dynamically insert "" tokens at uncertain positions based on prediction confidence
Fusion Mechanism: Blend hidden state context with probability-weighted embeddings: e_fusion = α·h_{t-1,I} + (1-α)·e_pred
Variable-Depth Reasoning: Adapt reasoning depth based on problem difficulty

Implementation

Phase 1: Standard CoT Fine-Tuning

def stage1_cot_finetuning(model, training_data, num_epochs=2):
    """Establish reasoning foundation with standard CoT"""

    for epoch in range(num_epochs):
        for example in training_data:
            input_text = example['input']
            reasoning = example['reasoning']
            output = example['output']

            # Format: input + thinking + output
            full_text = f"{input_text}\n<thinking>\n{reasoning}\n</thinking>\n{output}"

            logits = model.forward(full_text)
            loss = cross_entropy(logits, full_text)
            loss.backward()
            optimizer.step()

def stage2_confidence_insertion(model, validation_data, threshold_τ=0.7):
    """
    Learn where to insert <thinking> tokens based on confidence.
    Inference: dynamically insert tokens at uncertain positions.
    """

    # Analyze confidence on validation data
    confidence_stats = []

    for example in validation_data:
        input_text = example['input']

        # Get token-by-token confidences
        logits = model.forward(input_text, return_per_token=True)
        probs = softmax(logits, dim=-1)
        confidences = np.max(probs, axis=-1)

        confidence_stats.append(confidences)

    # Identify confident vs. uncertain tokens
    global_confidence = np.concatenate(confidence_stats)
    confidence_threshold = np.percentile(global_confidence, (1 - threshold_τ) * 100)

    return confidence_threshold

def inference_with_dynamic_insertion(model, input_text, confidence_threshold):
    """
    During inference, insert <thinking> tokens at uncertain positions.
    """

    tokens = tokenize(input_text)
    output_tokens = []
    current_confidence = 1.0

    for i, token in enumerate(tokens):
        # Predict confidence for next token
        partial_text = detokenize(tokens[:i])
        logits = model.forward(partial_text, return_per_token=True)
        probs = softmax(logits[-1], dim=-1)
        current_confidence = np.max(probs)

        # Decide whether to insert thinking token
        if current_confidence < confidence_threshold:
            output_tokens.append('<thinking>')
            # Generate reasoning step
            reasoning = model.sample(partial_text + ' <thinking>')
            output_tokens.extend(tokenize(reasoning))
            output_tokens.append('</thinking>')

        output_tokens.append(token)

    return detokenize(output_tokens)

class ContextPredictionFusion(nn.Module):
    """Fuses hidden state context with prediction embeddings"""

    def __init__(self, d_model):
        super().__init__()
        self.fusion_weight = nn.Parameter(torch.tensor(0.5))  # Learnable α

    def forward(self, h_prev, e_pred):
        """
        h_prev: hidden state from previous step
        e_pred: prediction embedding (from vocabulary)
        """
        # Constrain α to [0, 1]
        alpha = torch.sigmoid(self.fusion_weight)

        # Fuse: e_fusion = α·h + (1-α)·e_pred
        e_fused = alpha * h_prev + (1 - alpha) * e_pred

        return e_fused

def stage3_fusion_training(model, training_data, num_epochs=3):
    """
    Train fusion mechanism to maintain informative representations.
    """

    fusion_layer = ContextPredictionFusion(d_model=model.d_model)
    optimizer = Adam([fusion_layer.fusion_weight], lr=0.001)

    for epoch in range(num_epochs):
        for example in training_data:
            input_text = example['input']
            reasoning = example['reasoning']
            output = example['output']

            # Extract hidden states
            with torch.no_grad():
                h_states = model.get_hidden_states(
                    input_text + '\n<thinking>\n' + reasoning
                )

            # For each reasoning position, train fusion
            for t in range(len(h_states) - 1):
                h_prev = h_states[t]  # Previous hidden state

                # Prediction embedding: next token's embedding
                next_token = reasoning[t]
                e_pred = model.embedding(next_token)

                # Fuse
                e_fused = fusion_layer(h_prev, e_pred)

                # Predict next token from fused embedding
                logits = model.lm_head(e_fused)
                loss = cross_entropy(logits, reasoning[t+1])

                loss.backward()
                optimizer.step()

                # Monitor fusion weight
                alpha = torch.sigmoid(fusion_layer.fusion_weight)
                if t % 10 == 0:
                    print(f"Fusion α = {alpha.item():.3f}")

def full_lt_tuning_pipeline(model, training_data, validation_data):
    """
    Complete three-stage training.
    """

    print("Stage 1: CoT Finetuning...")
    stage1_cot_finetuning(model, training_data)

    print("Stage 2: Confidence Analysis...")
    confidence_threshold = stage2_confidence_insertion(
        model, validation_data, threshold_τ=0.7
    )

    print("Stage 3: Fusion Training...")
    stage3_fusion_training(model, training_data)

    return model, confidence_threshold

def infer_with_lt_tuning(model, input_text, confidence_threshold):
    """
    Inference combining dynamic insertion and fusion.
    """

    # Use confidence-driven insertion
    output = inference_with_dynamic_insertion(
        model, input_text, confidence_threshold
    )

    return output