DOCSEARCHbyAlgolia
Submit Docs

H-Transformer-1D

repository·main·Indexed 41 days ago

https://github.com/lucidrains/h-transformer-1d

An implementation of the H-Transformer-1D architecture using hierarchical attention to achieve subquadratic costs for long-range sequence modeling. The HTransformer1D class supports both autoregressive (causal) and non-autoregressive (encoder) modes and automatically handles padding to the nearest power of 2.

Tokens
854
Snippets
2
Records
3
Agent score
1%

What's inside H-Transformer-1D

  2. Configure HTransformer1D parameters

    main

    When initializing HTransformer1D, you can configure the following parameters:

    ParameterDescription
    num_tokensThe size of the vocabulary (number of tokens).
    dimThe embedding dimension.
    depthThe number of layers in the transformer.
    causalBoolean; set to True for autoregressive modeling, False for non-autoregressive (encoder) modeling.
    max_seq_lenThe maximum sequence length supported.
    headsThe number of attention heads.
    dim_headThe dimension per attention head.
    block_sizeThe size of the blocks used for hierarchical attention.
    reversibleBoolean; if True, uses reversibility to save memory at the cost of increased depth.
    shift_tokensBoolean; (experimental) whether to shift half the feature space by one along the sequence dimension for faster convergence.

  3. Use the HTransformer1D model

    main

    The HTransformer1D class implements a Transformer using hierarchical attention for sequence learning with subquadratic costs. It supports both autoregressive (causal) and non-autoregressive (encoder) modes. The network automatically handles padding to the nearest power of 2 and performs hierarchical attention internally.

    When passing inputs, you can provide a mask to handle variable sequence lengths.

    import torch
from h_transformer_1d import HTransformer1D

model = HTransformer1D(
    num_tokens = 256,          # number of tokens
    dim = 512,                 # dimension
    depth = 12,                # depth
    causal = False,            # autoregressive or not
    max_seq_len = 8192,        # maximum sequence length
    heads = 8,                 # heads
    dim_head = 64,             # dimension per head
    block_size = 128,          # block size
    reversible = True,         # use reversibility, to save on memory with increased depth
    shift_tokens = True        # whether to shift half the feature space by one along the sequence dimension, for faster convergence (experimental feature)
)

x = torch.randint(0, 256, (1, 8000))   # variable sequence length
mask = torch.ones((1, 8000)).bool()    # variable mask length

# network will automatically pad to power of 2, do hierarchical attention, etc

logits = model(x, mask = mask) # (1, 8000, 256)