Simple Hierarchical Transformer

The model uses hierarchies to define compression levels and window_sizes to control the receptive field of attention at each level:

1. Hierarchies: A value of 1 represents a regular transformer. Higher values (e.g., 2, 8) represent compressed hierarchical tokens. As you compress up the hierarchies, it is recommended to increase model dimensions to maintain capacity.
2. Window Sizes: This controls local attention. Higher hierarchies can pass distant information to the local (fine) hierarchy.
   • None: Full receptive field.
   • Integer: A local attention window of that size.
   • 0: Turns off attention at that hierarchy (though token shift remains in effect).
3. Hierarchical Stride: Controls the stride during compression and the repetition of tokens when concatenating hierarchical tokens back to fine tokens. This maintains causality using a shift trick similar to hourglass transformers.

The HierarchicalTransformer class implements a model that induces multiple hierarchical predictive coding models within a GPT. It uses compressed hierarchical tokens that undergo their own transformer blocks, with information pooled into one hierarchy at each layer.

To use it for training, pass return_loss = True to the forward call. For inference, call the model with token IDs to get logits.

import torch
from simple_hierarchical_transformer import HierarchicalTransformer

model = HierarchicalTransformer(
    num_tokens = 20000,                # number of tokens
    dim = 512,                         # model dimensions
    depth = 6,                         # depth
    dim_head = 64,                     # dimension per attention head
    heads = 8,                         # attention heads
    seq_len = 2048,                    # sequence lengths
    hierarchies = (1, 2, 8),           # hierarchies (e.g., 1x, 2x, and 8x compressed tokens)
    window_sizes = (32, 64, None)      # local attention window sizes (None = full receptive field)
)

ids = torch.randint(0, 20000, (1, 2048))

# Training mode
loss, _ = model(ids, return_loss = True)
loss.backward()

# Inference mode
logits = model(ids)

The HierarchicalTransformer constructor accepts several parameters to define the hierarchy and attention behavior:

• num_tokens: Integer. Vocabulary size.
• dim: Integer or Tuple. Model dimensions. If a tuple, it defines dimensions for each hierarchy level.
• depth: Integer. Number of transformer blocks.
• dim_head: Integer or Tuple. Dimension per attention head.
• heads: Integer or Tuple. Number of attention heads.
• seq_len: Integer. Sequence length.
• hierarchies: Integer or Tuple. Defines the compression levels (e.g., (1, 2, 4, 16)). Defaults to 1.
• window_sizes: Tuple of Integers/None. Local attention window sizes for each hierarchy. None represents a full receptive field. Defaults to None (full sequence length).
• hierarchical_stride: Tuple of Integers. Determines the stride when compressing. When concatenating hierarchical tokens to fine tokens, past tokens are repeated this amount of times. Recommended to stick with 1 except for highly compressed hierarchies.
• use_flash_attn: Boolean. Whether to use Flash Attention.
• ff_mult: Tuple of Floats. Feed-forward multiplier for each hierarchy level.

The HierarchicalTransformer is the main class for building hierarchical predictive coding models. It uses multiple levels of compression (hierarchies) to process sequences. You can customize each hierarchy's dimension, head count, and window size.

Key parameters:

• num_tokens: Vocabulary size.
• dim: A tuple of dimensions for each hierarchy level.
• depth: Number of transformer layers per hierarchy.
• hierarchies: A tuple of compression factors (must be powers of two, e.g., (1, 2, 4)).
• window_sizes: A tuple of attention window sizes for each hierarchy level.
• hierarchical_stride: A tuple of strides for each hierarchy level.
• predict_hierarchy: The hierarchy level used for the final next-token prediction (defaults to the minimum compression factor).
• predict_use_all_hierarchy: If True, the final prediction uses concatenated information from all hierarchies instead of just the selected predict_hierarchy.

from simple_hierarchical_transformer.simple_hierarchical_transformer import HierarchicalTransformer
import torch

model = HierarchicalTransformer(
    num_tokens = 5000,
    dim = (128, 256, 512),
    depth = 4,
    seq_len = 1024,
    hierarchies = (1, 2, 4),
    window_sizes = (None, 32, 64),
    hierarchical_stride = (1, 1, 1),
    predict_hierarchy = 1
)

ids = torch.randint(0, 5000, (8, 1024))
loss, (ce_loss, recon_loss, ar_loss, sigreg_loss) = model(ids, return_loss=True)