metadata

license: mit

Gated State Space

This repo contains pretrain model for the gated state space paper. The model has been trained on C4 dataset. I have used Lucidrains' implementation (commit) for the model. I think the main benefit of this model is the ability to scale beyond the training context length. As authors noted in the paper, they trained the model on 4k sequence length but it generalized beyond that length. I have written a blog post on how I started the training here.

Here are the training logs (report) on W&B. Since the training is still going on, this repo currently contains the checkpoint @ ~200k step. The report contains both the loss and the output results. While the generated text somewhat makes sense, given the loss is still ~4.2, it is not useful out-of-the-box for most tasks. So need to fine-tune it before using it for anything else.

How to use this.

Since it is not based on transformers library, it is a bit tricky to use the model out of the box. Here are the general steps:

pip install gated-state-spaces-pytorch
Download the model weights from here.
Download the config from here.
Following code to patch the original model:

    model = AutoregressiveWrapper(
        GatedStateSpacesLM(
            **config
        ),
    )
    model.net.to_logits = nn.Sequential(
        nn.LayerNorm(f_emb),
        model.net.to_logits,
    )

Load the state dict: model.load_state_dict(torch.load('model.pt'))
If you want to fine-tune the model, you can freeze the embeddings:

    model.net.token_emb.weight.requires_grad_(False)
    model.net.token_emb.weight.copy_(emb)

    model.net.to_logits[1].weight.requires_grad_(False)
    model.net.to_logits[1].weight.copy_(emb)

Training proceedure

Primarily it has been trained by language model objective. However; I added a few tricks to further optimize the training. The main trick of using pretrained embeddings is explained in the Devlog blog post linked above. The batch size is 8 with a sequence length of 128, the optimizer is AdamW with a learning rate of 2e-5. However; it is trained using gradient accumulation of 4 so the effective batch size is 32. Training happens with two types of losses, one is simple cross entropy for the next token prediction and other is distillation loss from GPT2-xl. During training each loss is alternated. Gradient norm are also clipped at 1.0.