Block Causal Transformer

At Suno we experiment with many transformer variants such as MusicFM, GPT, diffusion, and chunked diffusion. We also support an ever growing list of tasks such as cover, artist, infill, upsample, stem, and dry. Model architectures are split into multiple codebases and branches. This is an attempt to unify our architectures and tasks into a single framework: the Block Causal Transformer.

What is a Block Causal Transformer (BCT)?

A block causal transformer processes information in blocks of tokens from left to right. Each block has context of all the previous blocks, and has its own properties such as causality and discreteness. Each block can have different causality and types of io. This architecture encompasses all the architectures we've experimented with at Suno. For example, lets look at a possible BCT:

Here we have 3 blocks - text, semantic, and audio. Our text prompt is noncausal for better text understanding. Semantic is predicted autoregressively like GPT, and audio predicted with diffusion with text and semantic as context.

We can describe all our current models in this format: Notation: underline indicates a causal block, normal text is noncausal

You can see that all transformer models we use can be represented as BCTs. A natural question to ask is if future models will fit into this framework, or if we've overfit to past architectures. The core assumption of a BCT is that we input various types of data sequentially into a transformer. This is opposed to the encoder decoder architecture which has fallen out of favor. It doesn't prescribe what happens within a block, just that blocks come in some order. This is general enough that I believe it encompasses almost all transformers so far.

As an example of the generality of BCT's, we can add one property to the block specification that extends the power of BCT's to cover architectures like delay patterns, shallow diffusion, early exit, and hydra heads. This is layerwise IO. When defining our blocks, we can specify at which layers inputs and outputs should be placed. By defining composable block properties such as causality, discreteness, and layerwise IO we can cover most of the space of efficient architectures.

BCT's have some very nice practical properties. Thanks to FlexAttention, training is simple and efficient. Packing allows us to efficiently train on long sequence lengths. Inference is also efficient since causality allows the use of a KV cache. I suspect that any model that utilizes a KV cache can be represented as a BCT. flex attention is the backbone of the implementation. The pytorch team is planning to make flex attention a core feature with similar performance as flash 3

Specification

Now lets look at the actual specification of a BCT.

@dataclass
class InputType:
	"""Specifies a single input head"""
	name: str
	dimensions: int
	is_discrete: bool
	norm: bool = True
	layer: int = 0

First we define the IO heads. These can be discrete or continuous. The model will create embedding or linear heads depending on the type.

@dataclass
class OutputType:
	"""Specifies a single output head"""
	name: str
	dimensions: int
	is_discrete: bool
	layer: int = -1

Blocks can have multiple inputs and outputs, like in the case of musicgen with many inputs. is_causal specifies causality within the block.

@dataclass
class BlockType:
	"""Specifies a type of block like text or audio"""
	name: str
	is_causal: bool
	input_types: list[InputType]
	output_types: list[OutputType]

@dataclass
class Block:
	"""An instance of a block with data"""
	spec: BlockType
	inputs: list[Tensor] # (B, D, T)
	targets: list[Optional[Tensor]] = None # (B, D, T)

	def __len__(self):
		return in.shape[-1]

Finally Block is an instance of a block prepared in the dataloader. A list of Blocks will be fed into the model during training.

For each batch of data, we will compute the appropriate attention mask for flex attention given the input blocks and packing.

Summary

The Block Causal Transformer (BCT) framework provides a flexible way to implement these different architectures with appropriate attention masking. It handles the complexity of managing different types of blocks (text, audio, etc.) and their causal relationships, making it easier to experiment with hybrid approaches like chunked diffusion.

Concretely, BCT defines the types of data a transformer model supports, and provides data structures to simplify data loading. It computes the attention mask for the Block's specified. It does not implement the loss function or prepare the input and target tensors.

Refactoring sunoGPT to use the BCT framework shouldn't be a large change, and it shouldn't change functionality at all. (why do I keep saying things like this)