awf / functional-transformer Goto Github PK

Contributors

Stargazers

Forkers

jamesthesnake arthurperret ml-lab alexm-gc anas-zafar kaushalya

# Multi-head self-attention
for head in layer.heads:

    # Project into this head's query/key space
    query = linear(head.query, t1)                  # L x Dk
    key = linear(head.key, t1)                      # L x Dk

    # Compute L x L attention matrix
    score = query @ key.T + mask                    # L x L
    attn = jax.nn.softmax(cfg.tau * score, axis=1)  # L x L

    value = linear(head.value, t1)                  # L x Dm
    self_attn = attn @ value                        # L x Dm

    # Add this head's contribution into embeddings
    embeddings += self_attn                         # L x Dm     <---- sum instead of concatenate
    
# <-- after concatenating all attention heads there should be another linear layer here.