Simplified the model by always computing batch-first

modeling_norbert.py

@@ -101,23 +101,6 @@ class FeedForward(nn.Module):
         return self.mlp(x)
 
 
-class MaskedSoftmax(torch.autograd.Function):
-    @staticmethod
-    def forward(self, x, mask, dim):
-        self.dim = dim
-        x.masked_fill_(mask, float('-inf'))
-        x = torch.softmax(x, self.dim)
-        x.masked_fill_(mask, 0.0)
-        self.save_for_backward(x)
-        return x
-
-    @staticmethod
-    def backward(self, grad_output):
-        output, = self.saved_tensors
-        input_grad = softmax_backward_data(self, grad_output, output, self.dim, output)
-        return input_grad, None, None
-
-
 class Attention(nn.Module):
     def __init__(self, config):
         super().__init__()
@@ -142,7 +125,7 @@ class Attention(nn.Module):
             - torch.arange(config.max_position_embeddings, dtype=torch.long).unsqueeze(0)
         position_indices = self.make_log_bucket_position(position_indices, config.position_bucket_size, config.max_position_embeddings)
         position_indices = config.position_bucket_size - 1 + position_indices
-        self.register_buffer("position_indices", position_indices, persistent=False)
+        self.register_buffer("position_indices", position_indices.contiguous(), persistent=False)
 
         self.dropout = nn.Dropout(config.attention_probs_dropout_prob)
         self.scale = 1.0 / math.sqrt(3 * self.head_size)
@@ -155,10 +138,11 @@ class Attention(nn.Module):
         bucket_pos = torch.where(abs_pos <= mid, relative_pos, log_pos * sign).long()
         return bucket_pos
 
-    def compute_attention_scores(self, hidden_states, relative_embedding):
-        key_len, batch_size, _ = hidden_states.size()
+    def forward(self, hidden_states, attention_mask, relative_embedding):
+        batch_size, key_len, _ = hidden_states.size()
         query_len = key_len
 
+        # Recompute position_indices if sequence length exceeds the precomputed size
         if self.position_indices.size(0) < query_len:
             position_indices = torch.arange(query_len, dtype=torch.long).unsqueeze(1) \
                 - torch.arange(query_len, dtype=torch.long).unsqueeze(0)
@@ -166,48 +150,52 @@ class Attention(nn.Module):
             position_indices = self.config.position_bucket_size - 1 + position_indices
             self.position_indices = position_indices.to(hidden_states.device)
 
-        hidden_states = self.pre_layer_norm(hidden_states)
-
-        query, key = self.in_proj_qk(hidden_states).chunk(2, dim=2)  # shape: [T, B, D]
-        value = self.in_proj_v(hidden_states)  # shape: [T, B, D]
+        # Pre-LN and project query/key/value.
+        hidden_states = self.pre_layer_norm(hidden_states)  # shape: [B, T, D]
+        query, key = self.in_proj_qk(hidden_states).chunk(2, dim=-1)  # shape: [B, T, D]
+        value = self.in_proj_v(hidden_states)  # shape: [B, T, D]
 
-        query = query.reshape(query_len, batch_size * self.num_heads, self.head_size).transpose(0, 1)
-        key = key.reshape(key_len, batch_size * self.num_heads, self.head_size).transpose(0, 1)
-        value = value.view(key_len, batch_size * self.num_heads, self.head_size).transpose(0, 1)
+        # Reshape to [B, num_heads, T, head_size]
+        query = query.view(batch_size, query_len, self.num_heads, self.head_size).transpose(1, 2)  # shape: [B, num_heads, T_q, head_size]
+        key = key.view(batch_size, key_len, self.num_heads, self.head_size).permute(0, 2, 3, 1)  # shape: [B, num_heads, head_size, T_k]
+        value = value.view(batch_size, key_len, self.num_heads, self.head_size).transpose(1, 2)  # shape: [B, num_heads, T_k, head_size]
 
-        attention_scores = torch.bmm(query, key.transpose(1, 2) * self.scale)
+        # Compute relative positional contributions
+        pos = self.in_proj_qk(self.dropout(relative_embedding))  # shape: [2*position_bucket_size - 1, 2D]
+        query_pos, key_pos = pos.view(-1, self.num_heads, 2*self.head_size).chunk(2, dim=2)  # shape: [2*position_bucket_size - 1, num_heads, head_size]
+        query_pos = query_pos.transpose(0, 1)  # shape: [num_heads, 2*position_bucket_size - 1, head_size]
+        key_pos = key_pos.permute(1, 2, 0)  # shape: [num_heads, head_size, 2*position_bucket_size - 1]
 
-        pos = self.in_proj_qk(self.dropout(relative_embedding))  # shape: [2T-1, 2D]
-        query_pos, key_pos = pos.view(-1, self.num_heads, 2*self.head_size).chunk(2, dim=2)
-        query = query.view(batch_size, self.num_heads, query_len, self.head_size)
-        key = key.view(batch_size, self.num_heads, query_len, self.head_size)
+        # Scale the keys
+        key = key * self.scale
+        key_pos = key_pos * self.scale
 
-        attention_c_p = torch.einsum("bhqd,khd->bhqk", query, key_pos.squeeze(1) * self.scale)
-        attention_p_c = torch.einsum("bhkd,qhd->bhqk", key * self.scale, query_pos.squeeze(1))
+        # Compute standard content-to-content attention scores
+        attention_c_to_c = torch.matmul(query, key)  # shape: [B, num_heads, T_q, T_k]
 
-        position_indices = self.position_indices[:query_len, :key_len].expand(batch_size, self.num_heads, -1, -1)
-        attention_c_p = attention_c_p.gather(3, position_indices)
-        attention_p_c = attention_p_c.gather(2, position_indices)
+        # Compute content-to-position and position-to-content attention scores
+        position_indices = self.position_indices[:query_len, :key_len].expand(batch_size, self.num_heads, -1, -1)  # shape: [B, num_heads, T_q, T_k]
+        attention_c_to_p = torch.matmul(query, key_pos.unsqueeze(0))  # shape: [B, num_heads, T_q, 2*position_bucket_size - 1]
+        attention_p_to_c = torch.matmul(query_pos.unsqueeze(0), key)  # shape: [B, num_heads, 2*position_bucket_size - 1, T_k]
+        attention_c_to_p = attention_c_to_p.gather(3, position_indices)  # shape: [B, num_heads, T_q, T_k]
+        attention_p_to_c = attention_p_to_c.gather(2, position_indices)  # shape: [B, num_heads, T_q, T_k]
 
-        attention_scores = attention_scores.view(batch_size, self.num_heads, query_len, key_len)
-        attention_scores.add_(attention_c_p)
-        attention_scores.add_(attention_p_c)
+        # Full attention score
+        attention_scores = attention_c_to_c + attention_c_to_p + attention_p_to_c  # shape: [B, num_heads, T_q, T_k]
 
-        return attention_scores, value
+        # Masked softmax
+        attention_scores = attention_scores.masked_fill(attention_mask, float('-inf'))  # shape: [B, num_heads, T_q, T_k]
+        attention_probs = F.softmax(attention_scores, dim=-1)  # shape: [B, num_heads, T_q, T_k]
 
-    def compute_output(self, attention_probs, value):
-        attention_probs = self.dropout(attention_probs)
-        context = torch.bmm(attention_probs.flatten(0, 1), value)  # shape: [B*H, Q, D]
-        context = context.transpose(0, 1).reshape(context.size(1), -1, self.hidden_size)  # shape: [Q, B, H*D]
-        context = self.out_proj(context)
-        context = self.post_layer_norm(context)
-        context = self.dropout(context)
-        return context
+        # Collect the weighted-averaged values
+        attention_probs = self.dropout(attention_probs)  # shape: [B, num_heads, T_q, T_k]
+        output = torch.matmul(attention_probs, value)  # shape: [B, num_heads, T_q, head_size]
+        output = output.transpose(1, 2).flatten(2, 3)  # shape: [B, T_q, D]
+        output = self.out_proj(output)
+        output = self.post_layer_norm(output)
+        output = self.dropout(output)
 
-    def forward(self, hidden_states, attention_mask, relative_embedding):
-        attention_scores, value = self.compute_attention_scores(hidden_states, relative_embedding)
-        attention_probs = MaskedSoftmax.apply(attention_scores, attention_mask, -1)
-        return self.compute_output(attention_probs, value), attention_probs.detach()
+        return output, attention_probs.detach()
 
 
 class Embedding(nn.Module):
@@ -290,9 +278,8 @@ class NorbertModel(NorbertPreTrainedModel):
             attention_mask = ~attention_mask.bool()
         attention_mask = attention_mask.unsqueeze(1).unsqueeze(2)
  
-        static_embeddings, relative_embedding = self.embedding(input_ids.t())
+        static_embeddings, relative_embedding = self.embedding(input_ids)
         contextualized_embeddings, attention_probs = self.transformer(static_embeddings, attention_mask, relative_embedding)
-        contextualized_embeddings = [e.transpose(0, 1) for e in contextualized_embeddings]
         last_layer = contextualized_embeddings[-1]
         contextualized_embeddings = [contextualized_embeddings[0]] + [
             contextualized_embeddings[i] - contextualized_embeddings[i - 1]