# FlowLMModel Architecture: How Pocket‑TTS Generates Latent Audio Representations

> Discover the FlowLMModel architecture learn how Pocket-TTS generates latent audio codes in real-time with Lagrangian Self-Distillation for streaming text-to-speech

- Repository: [kyutai/pocket-tts](https://github.com/kyutai-labs/pocket-tts)
- Tags: architecture
- Published: 2026-07-11

---

**The FlowLMModel is a transformer-based flow language model that converts text and voice prompts into latent audio codes using Lagrangian Self-Distillation (LSD), enabling real-time streaming text-to-speech synthesis.**

The **FlowLMModel** serves as the core generative engine in the `kyutai-labs/pocket-tts` repository. This architecture uniquely combines causal transformer processing with flow-based generation to produce latent audio representations that the Mimi codec converts into speech waveforms. Unlike standard autoregressive token prediction models, FlowLMModel employs continuous flow matching to generate audio latents efficiently in just one or few steps.

## Core Architecture Components

The FlowLMModel architecture consists of five interconnected modules defined across specific source files that process text conditioning and generate latent audio streams.

### LUTConditioner for Text Encoding

Located in [`pocket_tts/conditioners/text.py`](https://github.com/kyutai-labs/pocket-tts/blob/main/pocket_tts/conditioners/text.py), the **LUTConditioner** tokenizes input text using SentencePiece and embeds tokens into the transformer's dimension (`d_model`). This module handles the initial conversion of raw text into vector representations suitable for the transformer backbone.

### SimpleMLPAdaLN Flow Network

The **SimpleMLPAdaLN** (implemented in [`pocket_tts/modules/mlp.py`](https://github.com/kyutai-labs/pocket-tts/blob/main/pocket_tts/modules/mlp.py)) functions as the lightweight flow network. This MLP predicts the flow direction `u_t` required for the Lagrangian Self-Distillation sampler, conditioning the flow on the transformer's output states.

### StreamingTransformer Backbone

Defined in [`pocket_tts/modules/transformer.py`](https://github.com/kyutai-labs/pocket-tts/blob/main/pocket_tts/modules/transformer.py), the **StreamingTransformer** is a causal transformer that processes concatenated sequences of text embeddings and previous latent embeddings. It maintains a KV-cache for streaming generation, enabling the model to attend to both modalities while preserving computational efficiency during autoregressive decoding.

### Latent Projection and BOS Handling

The model projects latent vectors from dimension `ldim` (default 64) to the transformer dimension using a linear projection layer. When voice prompts are provided, the model inserts a learned begin-of-sequence (BOS) embedding before the voice conditioning, replacing NaN placeholders in the input sequence with `self.bos_emb`.

### LSD Decoder Integration

The **LSD decoder** (implemented in [`pocket_tts/models/flow_lm.py`](https://github.com/kyutai-labs/pocket-tts/blob/main/pocket_tts/models/flow_lm.py)) reconstructs latent samples from a known prior (`x₀`) by repeatedly applying the learned flow `v_t`. This component implements the Lagrangian Self-Distillation method, integrating flow directions over discrete steps to produce the final latent representation of shape `[B, 1, ldim]`.

## Forward Pass and Latent Generation

The `forward()` method in `FlowLMModel` orchestrates the conversion of text and previous latents into new audio representations through a multi-step pipeline.

### Input Processing and Projection

The forward pass begins by handling special tokens and dimensionality alignment:

```python
def forward(self, sequence, text_embeddings, model_state,
            lsd_decode_steps, temp, noise_clamp, eos_threshold):
    # Replace NaN placeholders (BOS slots) with learned BOS embedding

    sequence = torch.where(torch.isnan(sequence), self.bos_emb, sequence)
    
    # Linear projection from latent dim → transformer dim

    input_ = self.input_linear(sequence)

```

### Transformer Conditioning and Flow Integration

The backbone concatenates text embeddings with projected latents and processes them through the StreamingTransformer. The model then samples Gaussian noise (optionally truncated by `noise_clamp` and scaled by `temp`) and conditions the flow network on transformer outputs:

```python
    # Run streaming transformer with text conditioning

    transformer_out = self.backbone(input_, text_embeddings,
                                    sequence, model_state=model_state)
    transformer_out = transformer_out.to(torch.float32)
    
    # Sample Gaussian noise with temperature scaling

    noise = torch.empty(..., dtype=transformer_out.dtype, device=transformer_out.device)
    torch.nn.init.normal_(noise, mean=0.0, std=temp**0.5)
    
    # Condition flow net on transformer output

    conditioned_flow = partial(self.flow_net, transformer_out)

```

### LSD Decoding and Output

The `lsd_decode` function integrates the flow for `lsd_decode_steps` (default 1) to produce the final latent, while the EOS head predicts sequence termination:

```python
    # LSD decode: integrate flow to get latent

    latent = lsd_decode(conditioned_flow, noise, lsd_decode_steps)
    
    # Predict end-of-sequence from last transformer output

    out_eos = self.out_eos(transformer_out[:, -1]) > eos_threshold
    return latent, out_eos

```

The returned latent tensor has shape `[B, 1, ldim]` and is subsequently fed into the Mimi codec via `TTSModel._decode_audio_worker` to produce a waveform.

## Streaming Generation Implementation

FlowLMModel enables real-time speech synthesis through autoregressive streaming with KV-cache management.

### Autoregressive Sampling with _sample_next_latent

During generation, the `_sample_next_latent` method receives the current latent sequence, text embeddings, and model state (KV-cache). It calls the forward method with `lsd_decode_steps > 0` to obtain the next latent and an end-of-sequence flag, which the surrounding `TTSModel` uses to manage the generation loop.

### Real-Time Decoding Pipeline

As soon as the model produces a latent representation, the `TTSModel` enqueues it for immediate decoding by the Mimi codec while the next latent is being generated. This pipelined architecture minimizes latency between text input and audio output, enabling conversational applications.

## Configuration and Model Instantiation

All architectural dimensions and hyperparameters are defined through a type-safe configuration system.

### FlowLMConfig Parameters

The `FlowLMConfig` class in [`pocket_tts/utils/config.py`](https://github.com/kyutai-labs/pocket-tts/blob/main/pocket_tts/utils/config.py) specifies:
- `d_model`: Transformer hidden dimension
- Number of transformer layers
- `ldim`: Latent dimension (typically 64)

### Factory Method Initialization

The model is instantiated via `FlowLMModel.from_pydantic_config`, which wires together the LUTConditioner, SimpleMLPAdaLN flow network, and StreamingTransformer according to the configuration values.

## Practical Implementation Examples

### Full TTS Pipeline Usage

```python
from pocket_tts import TTSModel

# Load default English model

tts = TTSModel.load_model()

# Create voice state from reference audio

voice_state = tts.get_state_for_audio_prompt(
    "hf://kyutai/tts-voices/alba-mackenna/casual.wav"
)

# Generate speech

audio = tts.generate_audio(voice_state, "Hello world! This is pocket‑tts.")
print(f"Audio shape: {audio.shape}, duration: {audio.shape[-1]/tts.sample_rate:.2f}s")

```

### Direct FlowLMModel Usage (Advanced)

```python
from pocket_tts.models.flow_lm import FlowLMModel
from pocket_tts.utils.config import load_config, CONFIGS_DIR

# Load configuration

cfg = load_config(CONFIGS_DIR / "english.yaml")
flow_lm = FlowLMModel.from_pydantic_config(
    cfg.flow_lm,
    latent_dim=cfg.mimi.quantizer.dimension,
    insert_bos_before_voice=False,
)

# Prepare dummy inputs

import torch
seq_len = 1
batch = torch.full((1, seq_len, flow_lm.ldim), float("nan"))  # BOS placeholder

text_emb = torch.zeros((1, 0, flow_lm.dim))  # No text conditioning

model_state = {}  # Empty KV cache

# Generate latent

latent, eos = flow_lm(
    sequence=batch,
    text_embeddings=text_emb,
    model_state=model_state,
    lsd_decode_steps=1,
    temp=1.0,
    noise_clamp=None,
    eos_threshold=0.5,
)
print("Generated latent shape:", latent.shape)  # → [1, 1, ldim]

```

### Streaming Audio Generation

```python
for chunk in tts.generate_audio_stream(voice_state, "A long paragraph for streaming..."):
    # chunk is a 1-D tensor of audio samples ready for playback

    play(chunk)  # Replace with audio output routine

```

## Summary

- **FlowLMModel** combines a transformer backbone with flow-based generation to produce latent audio codes in the `kyutai-labs/pocket-tts` architecture.
- **Core components** include the LUTConditioner (text), StreamingTransformer (backbone), SimpleMLPAdaLN (flow net), and LSD decoder.
- **Generation process** involves projecting latents from `ldim` to `d_model`, transformer conditioning, and LSD integration to produce outputs of shape `[B, 1, ldim]`.
- **Streaming capability** is achieved through KV-cache management in `_sample_next_latent` and parallel audio decoding via the Mimi codec.
- **Configuration** is handled via Pydantic models in `FlowLMConfig`, with instantiation through `from_pydantic_config`.

## Frequently Asked Questions

### What makes FlowLMModel different from standard autoregressive TTS models?

Unlike autoregressive models that predict discrete tokens sequentially, FlowLMModel uses **Lagrangian Self-Distillation (LSD)** to generate continuous latent representations through flow matching. This approach allows the model to generate high-quality audio latents in fewer steps (typically 1) while maintaining the parallelizable benefits of transformer architectures.

### How does the LSD decoder generate latent audio representations?

The LSD decoder reconstructs latent samples from a Gaussian prior by integrating learned flow directions `v_t` over `lsd_decode_steps`. The flow network (`SimpleMLPAdaLN`) predicts these directions conditioned on transformer outputs, and the `lsd_decode` function in [`pocket_tts/models/flow_lm.py`](https://github.com/kyutai-labs/pocket-tts/blob/main/pocket_tts/models/flow_lm.py) iteratively applies these flows to transform random noise into structured latent audio codes suitable for the Mimi codec.

### What is the role of the StreamingTransformer in the FlowLMModel architecture?

The **StreamingTransformer** serves as the causal attention mechanism that processes concatenated text and latent embeddings. Located in [`pocket_tts/modules/transformer.py`](https://github.com/kyutai-labs/pocket-tts/blob/main/pocket_tts/modules/transformer.py), it maintains a KV-cache to enable efficient autoregressive generation without recomputing attention for previous tokens. This component allows FlowLMModel to handle streaming generation where latents are produced one timestep at a time while attending to both historical audio context and text conditioning.

### Can FlowLMModel generate audio in real-time?

Yes, the architecture supports real-time generation through the `_sample_next_latent` method and streaming pipeline. The model produces latents of dimension `ldim` (default 64) incrementally, while a background thread decodes these latents via the Mimi codec. This design enables continuous audio playback with minimal latency between text input and audio output.