|
|
--- |
|
|
{} |
|
|
--- |
|
|
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. |
