# Block Causal Transformer (BCT) Architecture

**Implementation Reference:** `sunoGPT/utils/bct.py`

## Overview

The Block Causal Transformer (BCT) is a data structure framework for organizing training data for transformer-based audio generation models. BCT provides a unified way to represent sequences of different data types (text, audio, embeddings) as "blocks" that are processed left-to-right by a transformer.

**Key Point:** BCT is primarily a data organization layer, not a model architecture specification. It defines how to structure and pack training data, not how the model itself works.

## What BCT Actually Is

BCT provides:
1. **Data structures** for organizing training sequences
2. **Type definitions** for blocks with different properties
3. **Packing utilities** for efficient batching
4. **Serialization** for saving/loading sequences
5. **Statistics tracking** for monitoring training data distribution

BCT does **NOT** provide:
- Model architecture (just uses standard transformers)
- Attention mask computation (handled elsewhere)
- Loss functions (model-specific)
- Training loops

## Core Data Structures

### 1. BlockType - Specification

**See:** [bct.py](bct.py)

```python
@dataclass
class BlockType:
    """Specifies a type of block like text or audio"""
    name: str                      # e.g., "text", "audio", "ditto"
    is_causal: bool                # Causality within the block
    description: Optional[str] = None
```

**Properties:**
- `name`: Identifier for the block type
- `is_causal`: Whether tokens within this block attend causally (unidirectional) or bidirectionally
- Has serialization methods: `to_dict()`, `from_dict()`
- Has colored string representation for debugging

**Example BlockTypes:**
```python
text_block = BlockType(name="text", is_causal=False)      # Bidirectional
audio_block = BlockType(name="audio", is_causal=True)     # Autoregressive
ditto_block = BlockType(name="ditto", is_causal=False)    # Bidirectional embedding
```

### 2. Block - Instance with Data

**See:** [bct.py](bct.py)

```python
@dataclass
class Block:
    """An instance of a block with data"""
    spec: BlockType
    inputs: dict[str, torch.Tensor]   # (T, D) - named input tensors
    targets: dict[str, torch.Tensor]  # (T, D) - named target tensors
    debug_text: Optional[str] = None
```

**Properties:**
- `spec`: The BlockType specification
- `inputs`: Dictionary of named input tensors, shape `(T, D)` where T is time steps, D is dimensions
- `targets`: Dictionary of named target tensors for training
- `debug_text`: Optional text for debugging (e.g., actual lyrics for text blocks)

**Key Methods:**
- `__len__()`: Returns number of tokens (T dimension of first input tensor)
- `shift_left()`: Static method for shifting tensors left (used for autoregressive targets)
- `__str__()`: Pretty-printed representation showing tensor shapes and values

**Example:**
```python
# Text block with 20 tokens
text_block = Block(
    spec=BlockType(name="text", is_causal=False),
    inputs={"tokens": torch.tensor([...])},  # shape: (20, 768)
    targets={},
    debug_text="verse about mountains"
)
```

### 3. BlockSequence - Single Training Sample

**See:** [bct.py](bct.py)

```python
@dataclass
class BlockSequence:
    """A sequence of blocks representing a single sample"""
    blocks: list[Block]
```

**Properties:**
- `blocks`: Ordered list of Block instances
- `n_tokens`: Total token count across all blocks

**Key Methods:**
- `__len__()`: Number of blocks in sequence
- `__getitem__()`, `__setitem__()`: Index into blocks list
- `__iter__()`: Iterate over blocks
- `__add__()`: Concatenate two sequences
- `crop_to_max_tokens(max_tokens)`: Truncate sequence to fit token budget
- `save(path)`, `load(path)`: Serialize to/from disk

**Example Structure:**
```python
sequence = BlockSequence([
    Block(spec=text_type, inputs={...}, targets={}),           # Text conditioning
    Block(spec=ditto_type, inputs={...}, targets={}),          # Style embedding
    Block(spec=audio_type, inputs={...}, targets={...}),       # Audio output
])
```

### 4. PackedBlockSequence - Batched Samples

**See:** [bct.py](bct.py)

```python
@dataclass
class PackedBlockSequence:
    """Multiple block sequences packed together, representing a single batch"""
    block_sequences: list[BlockSequence]
```

**Purpose:** Efficiently batch multiple training samples with different lengths by packing them together.

**Properties:**
- `block_sequences`: List of BlockSequence samples
- `n_tokens`: Total token count across all sequences

**Key Methods:**
- `__len__()`: Number of sequences in pack
- `__getitem__()`, `__setitem__()`: Index into sequences
- `__iter__()`: Iterate over sequences
- `append(block_sequence)`: Add a sequence to the pack
- `crop_to_max_tokens(max_tokens)`: Truncate pack to fit token budget
- `save(path)`, `load(path)`: Serialize to/from disk

**Usage:**
```python
batch = PackedBlockSequence([
    BlockSequence([text1, ditto1, audio1]),  # Sample 1
    BlockSequence([text2, audio2]),           # Sample 2
    BlockSequence([text3, ditto3, audio3]),  # Sample 3
])
```

### 5. BlockUsageStatistics - Training Monitoring

**See:** [bct.py](bct.py)

```python
class BlockUsageStatistics:
    """Statistics about training data using a rolling window"""
    def __init__(self, window_size: int = 100)
```

**Purpose:** Track distribution of block types and their losses during training.

**Tracks:**
- Block frequency per sample
- Token distribution across block types
- Block pair co-occurrence
- Average loss per block type (optional)

**Key Methods:**
- `update(packed_block_sequence, per_block_losses)`: Add sample to rolling window
- `__str__()`: Human-readable summary
- `get_wandb_metrics()`: Dictionary of metrics for Weights & Biases logging

**Output Example:**
```
Blocks/sample: audio=1.000 | text=0.890 | ditto=0.340 | playlist=0.120
% tokens: audio=0.812 | text=0.156 | ditto=0.024 | playlist=0.008
Avg loss: audio=2.3451 | text=1.8923 | ditto=1.2341
```

## Block Causality

The `is_causal` property defines attention within a block:

**Causal Blocks (`is_causal=True`):**
- Token at position `t` can only attend to positions `0` to `t`
- Used for autoregressive generation (GPT-style)
- Examples: audio output tokens, semantic tokens

**Non-causal Blocks (`is_causal=False`):**
- Token at position `t` can attend to all positions in the block
- Used for bidirectional understanding
- Examples: text prompts, embeddings, conditioning audio

**Inter-block Attention:**
All tokens in block N can attend to all tokens in blocks 1 through N-1, regardless of causality. The `is_causal` flag only affects attention *within* a block.

## Practical Usage Patterns

### Basic Text-to-Audio

```python
BlockSequence([
    Block(spec=BlockType("text", is_causal=False), ...),     # Prompt
    Block(spec=BlockType("audio", is_causal=True), ...),     # Output
])
```

### With Style Conditioning

```python
BlockSequence([
    Block(spec=BlockType("text", is_causal=False), ...),     # Lyrics
    Block(spec=BlockType("ditto", is_causal=False), ...),    # Style embedding
    Block(spec=BlockType("audio", is_causal=True), ...),     # Output
])
```

### Cover Task

```python
BlockSequence([
    Block(spec=BlockType("text", is_causal=False), ...),     # Lyrics
    Block(spec=BlockType("cover", is_causal=False), ...),    # Cover audio
    Block(spec=BlockType("audio", is_causal=True), ...),     # Original audio
])
```

### Infill with Context

```python
BlockSequence([
    Block(spec=BlockType("text", is_causal=False), ...),     # Lyrics
    Block(spec=BlockType("past", is_causal=False), ...),     # Past audio context
    Block(spec=BlockType("future", is_causal=False), ...),   # Future audio context
    Block(spec=BlockType("audio", is_causal=True), ...),     # Infill output
])
```

### Multiple Conditioning

```python
BlockSequence([
    Block(spec=BlockType("text", is_causal=False), ...),     # Lyrics
    Block(spec=BlockType("playlist", is_causal=False), ...), # Reference track 1
    Block(spec=BlockType("playlist", is_causal=False), ...), # Reference track 2
    Block(spec=BlockType("ditto", is_causal=False), ...),    # Style embedding
    Block(spec=BlockType("audio", is_causal=True), ...),     # Output
])
```

## Named Inputs and Targets

Each block can have multiple named tensors in both `inputs` and `targets` dictionaries. This allows for:

1. **Multi-head inputs**: Different input modalities (e.g., tokens + positions)
2. **Multi-head outputs**: Different prediction targets (e.g., semantic + acoustic)
3. **Flexible data**: Custom data per block type

**Example:**
```python
audio_block = Block(
    spec=BlockType("audio", is_causal=True),
    inputs={
        "semantic_tokens": torch.tensor(...),  # (T, 1)
        "acoustic_tokens": torch.tensor(...),  # (T, 8)
    },
    targets={
        "semantic_tokens": torch.tensor(...),  # (T, 1)
        "acoustic_tokens": torch.tensor(...),  # (T, 8)
    }
)
```

## Token Budgets and Cropping

**See:** [bct.py](bct.py)

Both `BlockSequence` and `PackedBlockSequence` support cropping to maximum token counts:

```python
# Crop sequence to 2048 tokens
sequence = sequence.crop_to_max_tokens(2048)

# Crop packed batch to 8192 tokens
batch = batch.crop_to_max_tokens(8192)
```

**Behavior:**
- Keeps blocks in order
- Truncates the last block if needed
- Truncates all tensor dimensions equally

## Serialization

**See:** [bct.py](bct.py)

All data structures support save/load:

```python
# Save
sequence.save("sample.pt")
batch.save("batch.pt")

# Load
sequence = BlockSequence.load("sample.pt")
batch = PackedBlockSequence.load("batch.pt")
```

**Format:**
- Single PyTorch file per sequence/batch
- Includes metadata (block specs) and tensors
- Uses unique keys for all tensors
- CPU-mapped on load

## Composable Conditioning (Product Context)

From a product perspective, BCT enables **composable conditioning** - the ability to combine multiple conditioning signals:

**Available Block Types (in practice):**
- Text blocks: Lyrics and style tags
- Cover blocks: Cover version audio
- Artist blocks: Audio from same artist
- Playlist blocks: Audio from thematically similar songs
- Overpaint blocks: Instrumental audio (to preserve during vocal generation)
- Underpaint blocks: Vocal audio (to preserve during instrumental generation)
- Stem blocks: Individual instrument tracks
- Sample blocks: Short audio samples to build around
- Ditto blocks: Learned style embeddings
- Past/Future blocks: Temporal context for infilling

**Key Product Insight:**
You can combine any number of these blocks in a sequence. Want to:
- Add new vocals to a section while preserving instrumentals? → Past + Future + Artist + Overpaint
- Make a cover in a specific style? → Text + Cover + Ditto
- Sample-based generation with playlist vibe? → Sample + Playlist + Text

Each new block type multiplies with existing blocks, creating exponentially growing capabilities.

## Training Integration

**Typical Training Flow:**

1. **Dataloader** creates `PackedBlockSequence` batches
2. **Model** processes sequences (computes attention masks, forward pass)
3. **Loss function** computes per-block losses using `targets` dictionaries
4. **Statistics** tracked via `BlockUsageStatistics.update()`
5. **Logging** via `get_wandb_metrics()` for monitoring

**Example Statistics Usage:**

```python
stats = BlockUsageStatistics(window_size=100)

for batch in dataloader:
    # Forward pass and compute losses
    per_block_losses = model(batch)

    # Update statistics
    stats.update(batch, per_block_losses)

    # Log to wandb every N steps
    if step % log_interval == 0:
        wandb.log(stats.get_wandb_metrics())
        print(stats)
```

## Implementation Details

### Color-Coded Debugging

**See:** [bct.py](bct.py)

`BlockType.__str__()` uses MD5 hashing to assign consistent colors to block names for terminal output, making it easy to visually distinguish block types in logs.

### Lazy Statistics

**See:** [bct.py](bct.py)

`BlockUsageStatistics` uses lazy recomputation - statistics are only recalculated when accessed after updates, not on every `update()` call.

### Rolling Window

**See:** [bct.py](bct.py)

Statistics use a fixed-size rolling window (default 100 samples) to provide recent statistics without unbounded memory growth.

### Tensor Shape Convention

All tensors use `(T, D)` shape:
- `T`: Time dimension (number of tokens/frames)
- `D`: Feature dimension (embedding size, number of channels, etc.)

Note: This is opposite of the common `(B, T, D)` batch convention. Batching happens at the `PackedBlockSequence` level, not in individual tensors.

## Relationship to Model Architecture

BCT describes data organization, not model architecture. The HTML docs mention various architectures:

**How BCT data maps to architectures:**

- **GPT**: All blocks are causal, processed left-to-right
- **PrefixLM**: Conditioning blocks are non-causal, output is causal
- **Diffusion**: Most blocks are non-causal, used as conditioning
- **Chunked Diffusion**: Alternating causal semantic + non-causal audio blocks

The model's attention mechanism determines how blocks interact. BCT just labels the data.

## Common Block Types in Practice

Based on the composable conditioning context:

| Block Name | Typical Causality | Purpose |
|------------|------------------|---------|
| `text` | Non-causal | Lyrics and style tags |
| `audio` | Causal | Output audio generation |
| `semantic` | Causal | Semantic token prediction |
| `cover` | Non-causal | Cover audio conditioning |
| `artist` | Non-causal | Artist identity from reference track |
| `playlist` | Non-causal | Thematic/genre conditioning |
| `overpaint` | Non-causal | Instrumental preservation |
| `underpaint` | Non-causal | Vocal preservation |
| `stem` | Non-causal | Individual instrument conditioning |
| `sample` | Non-causal | Audio sample to build around |
| `ditto` | Non-causal | Learned style embedding |
| `past` | Non-causal | Past temporal context |
| `future` | Non-causal | Future temporal context |

## Design Philosophy

BCT treats training data like **Lego blocks**:
- Each block is a self-contained piece
- Blocks snap together in sequences
- Combining blocks creates new capabilities
- The dataloader is the "builder"

The framework is intentionally minimal - it provides structure without constraining what you build.

## Summary

**BCT is:**
- A data structure framework for organizing transformer training data
- A way to label sequences with block types and causality
- A set of utilities for packing, cropping, and serializing
- A statistics tracker for monitoring data distribution

**BCT is not:**
- A model architecture
- An attention mask implementation
- A training framework
- A loss function specification

**Key takeaway:** BCT makes it easy to experiment with different combinations of conditioning signals by providing a clean data abstraction. The real power comes from composing blocks in creative ways.

## Code Reference

**Full implementation:** [bct.py](bct.py) (~450 lines)

**Key components:**
- 4 core dataclasses (BlockType, Block, BlockSequence, PackedBlockSequence)
- 1 statistics class (BlockUsageStatistics)
- Serialization, cropping, and utility methods