New Record: Paired Head Attention (-3.5s, -65 steps)#191
Conversation
|
I think it'll make more sense to calculate relative improvements in % instead of absolute seconds for the track? |
Do you mean in the readme show percent drop from the last record? |
|
dont get distracted by the recursion- that's a new experiment. class Attention(nn.Module):
def __init__(self, config):
super().__init__()
self.n_embd = config.n_embd
self.n_branch = config.n_branch
self.block_size = config.block_size
# Projections: Q and V produce NB branches, K is shared
self.q_proj = nn.Linear(config.n_embd, config.n_embd * self.n_branch, bias=config.bias)
self.k_proj = nn.Linear(config.n_embd, config.n_embd, bias=config.bias)
self.v_proj = nn.Linear(config.n_embd, config.n_embd * self.n_branch, bias=config.bias)
self.o_proj = nn.Linear(config.n_embd, config.n_embd, bias=config.bias)
self.register_buffer("mask", torch.tril(torch.ones(config.block_size, config.block_size))
.view(1, 1, config.block_size, config.block_size))
# RoPE should handle the embedding dimension directly now
self.rope = RoPE(self.n_embd, max_len=config.block_size)
self.attn_drop = nn.Dropout(config.dropout)
def forward(self, a, x):
B, T, C = x.shape
NB = self.n_branch
# Project and reshape: (B, T, NB, C) -> (B, NB, T, C)
q = self.q_proj(a).view(B, T, NB, C).transpose(1, 2)
v = self.v_proj(a).view(B, T, NB, C).transpose(1, 2)
# K has no branch dimension initially: (B, T, C) -> (B, 1, T, C)
k = self.k_proj(x).view(B, T, 1, C).transpose(1, 2)
# Apply RoPE
q, k = self.rope(q), self.rope(k)
# Raw scores: (B, NB, T, C) @ (B, 1, C, T) -> (B, NB, T, T)
att = (q @ k.transpose(-2, -1)) / math.sqrt(C)
# Causal Masking
mask = self.mask[:, :, :T, :T]
att = att.masked_fill(mask == 0, float("-inf"))
# Branch Routing Logic (Softmax across the NB dimension)
soft_probs = F.softmax(att, dim=1)
soft_probs = torch.nan_to_num(soft_probs, nan=0.0)
# Straight-Through Estimator (STE)
with torch.no_grad():
max_val = soft_probs.max(dim=1, keepdim=True)[0]
hard_mask = (soft_probs == max_val).float()
route_mask = (hard_mask - soft_probs).detach() + soft_probs
# Temporal Attention (Softmax across the T dimension)
# Find global importance by taking max over branches
att_max, _ = att.max(dim=1)
attn_probs = F.softmax(att_max, dim=-1)
attn_probs = self.attn_drop(attn_probs)
# Composition: Combine temporal weights and branch router
# (B, 1, T, T) * (B, NB, T, T) -> (B, NB, T, T)
combined_weights = attn_probs.unsqueeze(1) * route_mask
# Weighted sum of values and sum across branches: (B, NB, T, T) @ (B, NB, T, C)
y = (combined_weights @ v).sum(dim=1) # (B, T, C)
out = self.o_proj(y)
return out |
|
Observe how i'm routing V to select individual V from individual branch where individual Q -> K. |
|
I overlooked something when I refactored RoPE. It used to be [32 rotating, 32 stationary, 32 rotating, 32 stationary], and partial_key_offset is set to only apply to stationary dims. Now its [64 rotating, 64 stationary], but partial key offset is still applied to dims [32:64] and [96:128]. When I made this update I didn't notice a meaningful loss shift. So maybe it doesn't matter, but this may be performing slightly worse than if partial key offset is only applied to stationary dims. (In the prior PR I remember ablating only stationary vs only rotating dims and only stationary performed much better. But maybe a mix is fine too) |
|
Nice work — really clean trick. One thing I’m curious about from a downstream perspective: I’m wondering whether some attention-level efficiency tricks also shift how “certain” the model behaves, not just how fast it runs. |
for example here you calculated the relative change as -3.5s and used it to lock in in readme, but it would make more sense to calculate the relative change as -3.0876% and lock in readme time using that... but i guess with readme times being in minutes this won't matter anytime soon.. |
4 participants
Updates in PR:
The main idea behind paired head attention is to let queries attend to multiple representations of every position, and have multiple logits for the same position going into the same softmax. Implemented by interleaving the k, q, and v for pairs of heads to form twice as long sequences. EG [k1_h1, k2_h1, k3_h1], [k1_h2, k2_h2, k3_h2] -> [k1_h1, k1_h2, k2_h1, k2_h2, k3_h1, k3_h2], repeat for q and v
I update rotary() to run in a single line. Interestingly, pytorch is capable of computing x_flip without performing any data copying.
I implemented a seperate custom rotary method for PairedHeadAttention, that enables the attention forward pass to save an extra parameter reshape copy:
I staggered the custom rotary for PairedHeadAttention such that head1 rotates at an offset from head 2. From limited testing this seemed better. Note that because the data is stored sequentially and attention applies a causal mask, query 4 from head 1 can only attend to keys 1-3 for head 2, whereas query 4 from head 2 can attend to keys 1-4 for head 1.
I double the max doc length and seqlens to account for concatenating two heads:
I chose to not alter the window size. This means that window size of 128 is really only looking at 64 positions back (64 from head 1 and 64 from head2)
I chose to apply this to some of the short window layers, since I wasn't sure how it would interact with partial key offset on the longer windows. Kept adding to more layers and it kept getting better. I stopped after 4 (maybe more is better).
Seperately, I tested adding a mantissa to Adam, which seemed to give slight improvement. Leaving off this PR to keep the scope focused.
Timing and Validation
retiming prior record: 112.262 [112.187, 112.278, 112.321]
If no changes, will merge at 109.2s to be consistent with 3.5s improvement.