Butanium/simple-stories-one-layer-simple-transformer

1-layer simple transformer described in A Mathematical Framework for Transformer Circuits. Load with

class OneLayerTransformer(PreTrainedModel):
    config_class = LlamaConfig
    
    def __init__(self, config: LlamaConfig):
        super().__init__(config)
        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)

        # Single self-attention layer
        self.self_attn = nn.MultiheadAttention(
            embed_dim=config.hidden_size,
            num_heads=config.num_attention_heads,
            dropout=getattr(config, 'attention_dropout', 0.0),
            batch_first=True,
        )

        # Output head
        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)

    def forward(self, input_ids=None, attention_mask=None, labels=None, **kwargs):
        batch_size, seq_len = input_ids.shape

        # Embeddings
        hidden_states = self.embed_tokens(input_ids)
        assert hidden_states.shape == (batch_size, seq_len, self.config.hidden_size)

        # Create causal mask for self-attention
        causal_mask = torch.triu(torch.ones(seq_len, seq_len), diagonal=1).bool()
        causal_mask = causal_mask.to(hidden_states.device)

        # Self-attention with residual connection
        attn_output, _ = self.self_attn(
            hidden_states,
            hidden_states,
            hidden_states,
            attn_mask=causal_mask,
            key_padding_mask=None if attention_mask is None else ~attention_mask.bool(),
        )
        hidden_states = hidden_states + attn_output
        assert hidden_states.shape == (batch_size, seq_len, self.config.hidden_size)

        # Output projection
        logits = self.lm_head(hidden_states)
        assert logits.shape == (batch_size, seq_len, self.config.vocab_size)

        loss = None
        if labels is not None:
            shift_logits = logits[..., :-1, :].contiguous()
            shift_labels = labels[..., 1:].contiguous()
            loss_fct = nn.CrossEntropyLoss()
            loss = loss_fct(
                shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1)
            )

        return {"loss": loss, "logits": logits}



model = OneLayerTransformer.from_pretrained('Butanium/simple-stories-one-layer-simple-transformer')

The model is trained on the SimpleStories dataset.