coconut-pytorch

The Coconut class implements the architecture proposed in the paper "Training Large Language Models to Reason in a Continuous Latent Space".

Initialization

Initialize the model by specifying the number of reasoning steps, the number of latents per step, and a dictionary defining the transformer architecture parameters (num_tokens, dim, and depth).

Training

Pass a prompt and an answer tensor to the model instance to compute the loss. The model supports standard PyTorch backpropagation via loss.backward().

Generation

After training, use the .generate() method to produce an answer from a given prompt by specifying max_length.

import torch
from coconut_pytorch import Coconut

model = Coconut(
    num_reasoning_steps = 3,
    num_latents_per_step = 1,
    transformer = dict(
        num_tokens = 256,
        dim = 512,
        depth = 6
    )
)

prompt = torch.randint(0, 256, (2, 1024))
answer = torch.randint(0, 256, (2, 64))

# Training step
loss = model(prompt, answer)
loss.backward()

# Generation step
answer = model.generate(prompt, max_length = 64) # (2, 64)

When num_hypothesis > 1 is set in the Coconut class, the model implements a multi-stream reasoning process:

1. Expansion: After the initial prompt processing, the latent tokens are projected into multiple streams using to_streams.
2. Recurrent Reasoning: Each stream evolves independently through the reasoning steps. If synthesize_hypothesis_per_step is True, the streams interact via a synthesis layer during each step.
3. Diversity: A hyp_diff_loss can be applied to encourage the different hypothesis streams to represent different latent states (using cosine similarity).
4. Merging: At the end of the reasoning phase, the streams are merged back into a single representation using merge_streams before the final answer generation.
5. KV Cache Handling: The key/value cache is expanded to accommodate all hypothesis streams during reasoning and then averaged/rearranged so the final answer tokens can attend to the combined information.

The Coconut class is the primary interface for the library.

__call__(prompt, answer)

Computes the loss for a given prompt and target answer.

• Parameters:
  • prompt: Tensor of shape (batch, prompt_length) containing token indices.
  • answer: Tensor of shape (batch, answer_length) containing target token indices.
• Returns: A scalar loss tensor.

generate(prompt, max_length)

Generates an answer sequence from the provided prompt.

• Parameters:
  • prompt: Tensor of shape (batch, prompt_length).
  • max_length: Integer specifying the maximum number of tokens to generate.
• Returns: A tensor of shape (batch, max_length) containing the generated token indices.

# Example of the API usage
model = Coconut(num_reasoning_steps=3, num_latents_per_step=1, transformer=dict(num_tokens=256, dim=512, depth=6))
loss = model(prompt, answer)
answer = model.generate(prompt, max_length=64)

The Transformer class is the backbone of the Coconut model. It supports multiple residual streams and a specialized attention mechanism that allows for value residual mixing.

Key parameters:

• num_tokens: Vocabulary size.
• dim: Model dimension.
• depth: Number of layers.
• num_residual_streams: Number of parallel residual streams (uses hyper-connections to expand/reduce).
• accept_value_residual: If True, enables the Attention mechanism to mix in a value_residual from the previous layer, which is critical for the Coconut reasoning mechanism.

Inputs can be a single Tensor of integer tokens [b n] or a list of Tensors (e.g., combining discrete tokens and continuous embeddings).

from coconut_pytorch.coconut import Transformer
import torch

model = Transformer(
    num_tokens=5000,
    dim=512,
    depth=6,
    num_residual_streams=4,
    accept_value_residual=True
)

# Input can be integer tokens
x = torch.randint(0, 5000, (2, 10))
logits = model(x)

The Coconut class implements a model architecture that performs recurrent latent reasoning. It uses a Transformer backbone to process a prompt, then enters a reasoning phase where it iteratively updates latent tokens (thoughts) before generating an answer.

Key features include:

• Latent Reasoning: Iterative updates of num_reasoning_steps.
• Multi-stream Hypothesis: By setting num_hypothesis > 1, the model can maintain multiple independent latent streams (hypotheses) that are merged at the end of the reasoning phase.
• Thought Tokens: Uses learnable begin_of_thought and end_of_thought tokens to demarcate the reasoning phase.
• Losses: Supports begin_thought_loss (to teach the model when to start thinking) and hyp_diff_loss (to encourage diversity among multiple hypotheses).

from coconut_pytorch.multi_stream_coconut import Coconut, Transformer

# Define transformer configuration
transformer_config = {
    'num_tokens': 50000,
    'dim': 512,
    'depth': 12,
    'num_residual_streams': 4
}

# Initialize Coconut model
model = Coconut(
    num_reasoning_steps=5,
    transformer=transformer_config,
    num_latents_per_step=1,
    num_hypothesis=3,  # Enables multiple latent streams
    learn_begin_of_thought=True,
    hyp_diff_loss_weight=0.1
)

# Forward pass with prompt and answer
# prompt: Tensor of shape [batch, seq_len]
# answer: Tensor of shape [batch, seq_len]
loss, (loss_breakdown, prompt_logits, latent_tokens, answer_logits, cached_kv) = model(prompt, answer=answer)

# Backward pass
loss.backward()

The Coconut class is a wrapper around a Transformer that implements recurrent latent reasoning. It uses special <bot> (begin of thought) and <eot> (end of thought) tokens to facilitate reasoning steps in a continuous latent space.

Key features:

• Recurrent Latent Reasoning: Performs a fixed number of reasoning steps (num_reasoning_steps) where the last hidden state of one step is passed as input to the next, while keeping the prompt's KV cache constant.
• Customizable Latents: Allows specifying num_latents_per_step to extend the paper's implementation.
• Training Options: Supports learning the <bot> token via learn_begin_of_thought and uses gradient checkpointing to save memory during long reasoning chains.
• Generation: Includes a .generate() method for autoregressive decoding after the reasoning phase.

from coconut_pytorch.coconut import Coconut, Transformer
import torch

# Define transformer configuration
transformer_config = {
    'num_tokens': 5000,
    'dim': 512,
    'depth': 6,
}

# Initialize Coconut wrapper
model = Coconut(
    num_reasoning_steps=4,
    transformer=transformer_config,
    num_latents_per_step=1,
    learn_begin_of_thought=True,
    checkpoint=True
)

# Example forward pass with prompt and answer
prompt = torch.randint(0, 5000, (2, 10))
answer = torch.randint(0, 5000, (2, 15))

total_loss, (loss_breakdown, prompt_logits, latent_tokens, answer_logits, cached_kv) = model(
    prompt=prompt,
    answer=answer,
    return_loss=True
)