File size: 2,432 Bytes
5aeae2e 771a34c 40e27a9 771a34c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 |
---
{}
---
1-layer simple transformer described in [A Mathematical Framework for Transformer Circuits](https://transformer-circuits.pub/2021/framework/index.html).
Load with
```python
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. |