Create focus.py

focus.py

@@ -0,0 +1,257 @@
+class Adaptivefocus(nn.Module):
+    def __init__(self, base, dims, head, max_dist, sharpen, win_size, max_span, temp_scale=0.01, num_iterations=3):
+        super().__init__()
+        self.max_dist = max_dist
+        self.win_size = win_size
+        self.max_span = max_span
+        self.temp_scale = temp_scale
+        self.multihead_attn = MultiheadAttention(base=base, dims=dims, head=head, max_dist=max_dist)
+        self.span_scale = nn.Parameter(torch.tensor(1.0))
+        self.sharpen = sharpen
+        self.num_iterations = num_iterations
+        self.base_threshold = 1e-4
+        self.scaling_factor = 0.1
+
+    def _focus(self, query, key, value, span_scale):
+        max_iterations = self.num_iterations
+        iteration = 0
+        prev_attn_out = torch.zeros_like(query)
+
+        while iteration < max_iterations:
+            span_len = int(self.max_span * span_scale.mean().item())
+            span_len = min(span_len, query.shape[1], key.shape[1], value.shape[1])
+            eff_span = min(span_len, self.max_dist)
+
+            if eff_span == 0:
+                break
+
+            q_span = query[:, :eff_span, :]
+            k_span = key[:, :eff_span, :]
+            v_span = value[:, :eff_span, :]
+
+            batch_size, seq_len, dims = q_span.size()
+            scale = (dims // self.multihead_attn.head) ** -0.25
+
+            q = q_span.view(q_span.shape[0], q_span.shape[1], self.multihead_attn.head, -1).permute(0, 2, 1, 3)
+            k = k_span.view(k_span.shape[0], k_span.shape[1], self.multihead_attn.head, -1).permute(0, 2, 1, 3)
+            v = v_span.view(v_span.shape[0], v_span.shape[1], self.multihead_attn.head, -1).permute(0, 2, 1, 3)
+
+            if self.sharpen:
+                temperature = 1.0 + self.temp_scale * (1.0 - span_scale.mean().item())
+            else:
+                temperature = 0.5 + self.temp_scale * span_scale.mean().item()
+
+            attn_scores = torch.matmul(q, k.transpose(-2, -1))
+            attn_weights = torch.softmax((attn_scores / temperature) * scale, dim=-1)
+            attn_out = torch.matmul(attn_weights, v)
+            attn_out = attn_out.permute(0, 2, 1, 3).flatten(start_dim=2)
+            attn_out = attn_out.contiguous().view(batch_size, eff_span, dims)
+
+            diff = torch.abs(attn_out - prev_attn_out).mean()
+
+            dynamic_threshold = self.base_threshold + self.scaling_factor * diff
+
+            if diff < dynamic_threshold:
+                break
+
+            prev_attn_out = attn_out
+            query = query + attn_out  
+            iteration += 1  
+
+        return attn_out, attn_weights
+
+    def forward(self, query, key, value, span_scale):
+        return self._focus(query, key, value, span_scale)
+
+
+class SpanPredictor(nn.Module):
+    def __init__(self, dims):
+        super().__init__()
+        self.linear = nn.Linear(in_features=dims, out_features=1)
+
+    def forward(self, global_out):
+        scale = torch.sigmoid(self.linear(global_out))
+        return scale
+        
+class FocusedAttention(nn.Module):
+    def __init__(self, base, dims, head, max_dist, sharpen, win_size=32, max_span=32, slid_win=32, temp_scale=0.01, num_iterations=3):
+        super().__init__()
+        self.max_dist = max_dist
+        self.win_size = win_size
+        self.max_span = max_span
+        self.slid_win = slid_win
+
+        self.span_pred = SpanPredictor(dims=dims)
+        self.dist_local = max_dist
+        self.dist_global = max_dist
+
+        self.attn_local = Adaptivefocus(base=base, dims=dims, head=head, max_dist=max_dist, sharpen=sharpen, win_size=win_size, max_span=max_span, temp_scale=temp_scale, num_iterations=num_iterations)
+        self.attn_global = MultiheadAttention(base=base, dims=dims, head=head, max_dist=self.dist_global)
+        self.ln_local = LayerNorm(normalized_shape=dims)
+        self.ln_global = LayerNorm(normalized_shape=dims)
+        self.projection = Linear(in_features=2 * dims, out_features=dims)
+
+    def forward(self, x, new_dist=None, new_base=None, xa=None, mask=None, kv_cache=None):
+        local = self.ln_local(x)
+        globe = self.ln_global(x)
+
+        globe_out, _ = self.attn_global(globe, globe, globe)
+
+        span_scale = self.span_pred(globe_out.mean(dim=1))
+
+        win_size = max(1, int(self.slid_win * span_scale.mean().item()))
+        span_len = max(1, int(self.max_span * span_scale.mean().item()))
+
+        effective_max_dist = min(self.max_dist, local.size(1))
+        local_max_dist = min(self.dist_local, span_len, win_size)
+        globe_max_dist = effective_max_dist
+
+        self.attn_local.max_dist = local_max_dist
+        self.attn_global.max_dist = globe_max_dist
+
+        local_out = self.slide_win(x=local, win_size=win_size, span_len=span_len, span_scale=span_scale)
+
+        combined = torch.cat(tensors=[local_out, globe_out], dim=-1)
+        x = self.projection(combined)
+
+        return x
+
+    def slide_win(self, x, win_size, span_len, span_scale):
+        batch_size, seq_len, dims = x.size()
+        out = torch.zeros_like(x, device=x.device)
+
+        for i in range(0, seq_len, win_size):
+            end = min(i + win_size, seq_len)
+            query = x[:, i:end, :]
+
+            start = max(0, i - span_len + win_size)
+            key = x[:, start:i + span_len, :]
+            value = x[:, start:i + span_len, :]
+            attn_out, _ = self.attn_local(query, key, value, span_scale)
+            out[:, i:end, :] = attn_out
+
+        return out
+        
+## different version
+# class FocusedAttention(nn.Module):
+#     def __init__(self, base, dims, head, max_dist, sharpen, win_size=32, max_span=32, slid_win=32, temp_scale=0.01):
+#         super().__init__()
+#         self.base = base
+#         self.dims = dims
+#         self.head = head
+#         self.max_dist = max_dist
+#         self.sharpen = sharpen
+#         self.win_size = win_size
+#         self.max_span = max_span
+#         self.slid_win = slid_win
+#         self.temp_scale = temp_scale
+
+#         self.span_scale_param = nn.Parameter(torch.tensor(1.0))
+#         self.span_predictor = nn.Linear(in_features=dims, out_features=1)
+
+#         self.multihead_attn_local = MultiheadAttention(base=base, dims=dims, head=head, max_dist=max_dist)
+#         self.multihead_attn_global = MultiheadAttention(base=base, dims=dims, head=head, max_dist=max_dist)
+
+#         self.ln_local = LayerNorm(normalized_shape=dims)
+#         self.ln_global = LayerNorm(normalized_shape=dims)
+#         self.projection = Linear(in_features=2 * dims, out_features=dims)
+
+#     def forward(self, x):
+
+#         local = self.ln_local(x)
+#         global_ = self.ln_global(x)
+
+#         globe_out, _ = self.multihead_attn_global(global_, global_, global_)
+
+#         span_scale = torch.sigmoid(self.span_predictor(globe_out.mean(dim=1)))
+
+#         win_size = max(1, int(self.slid_win * span_scale.mean().item()))
+#         span_len = max(1, int(self.max_span * span_scale.mean().item()))
+
+#         effective_max_dist = min(self.max_dist, local.size(1))
+#         local_max_dist = min(self.max_dist, span_len, win_size)
+#         globe_max_dist = effective_max_dist
+
+#         self.multihead_attn_local.max_dist = local_max_dist
+#         self.multihead_attn_global.max_dist = globe_max_dist
+
+#         local_out = self._window(local, win_size, span_len, span_scale)
+
+#         combined = torch.cat([local_out, globe_out], dim=-1)
+#         x = self.projection(combined)
+
+#         return x
+
+#     def _window(self, x, win_size, span_len, span_scale):
+#         batch_size, seq_len, dims = x.size()
+#         output = torch.zeros_like(x, device=x.device)
+
+#         for i in range(0, seq_len, win_size):
+#             end = min(i + win_size, seq_len)
+#             query = x[:, i:end, :]
+
+#             start = max(0, i - span_len + win_size)
+#             key = x[:, start:i + span_len, :]
+#             value = x[:, start:i + span_len, :]
+
+#             attn_out, _ = self._focus(query, key, value, span_scale)
+#             output[:, i:end, :] = attn_out
+
+#         return output
+
+#     def _focus(self, query, key, value, span_scale):
+#         span_len = int(self.max_span * span_scale.mean().item())
+#         span_len = min(span_len, query.size(1), key.size(1), value.size(1))
+#         eff_span = min(span_len, self.max_dist)
+
+#         q_span = query[:, :eff_span, :]
+#         k_span = key[:, :eff_span, :]
+#         v_span = value[:, :eff_span, :]
+
+#         batch_size, seq_len, dims = q_span.size()
+#         scale_factor = (dims // self.head) ** -0.25
+
+#         q = q_span.view(batch_size, seq_len, self.head, -1).permute(0, 2, 1, 3)
+#         k = k_span.view(batch_size, seq_len, self.head, -1).permute(0, 2, 1, 3)
+#         v = v_span.view(batch_size, seq_len, self.head, -1).permute(0, 2, 1, 3)
+
+#         if self.sharpen:
+#             temperature = 1.0 + self.temp_scale * (1.0 - span_scale.mean().item())
+#         else:
+#             temperature = 0.5 + self.temp_scale * span_scale.mean().item()
+
+#         attn_scores = torch.matmul(q, k.transpose(-2, -1))
+#         attn_weights = torch.softmax((attn_scores / temperature) * scale_factor, dim=-1)
+#         attn_out = torch.matmul(attn_weights, v)
+
+#         attn_out = attn_out.permute(0, 2, 1, 3).contiguous().view(batch_size, seq_len, -1)
+
+#         return attn_out, attn_weights
+
+
+# #Batch: 
+
+#     def _window(self, x, win_size, span_len, span_scale):
+#         batch_size, seq_len, dims = x.size()
+#         num_windows = (seq_len + win_size - 1) // win_size  # Calculate the number of windows
+    
+#         # Create tensors to store the outputs
+#         output = torch.zeros_like(x, device=x.device)
+    
+#         # Iterate over the windows in a more efficient manner
+#         for i in range(num_windows):
+#             start_idx = i * win_size
+#             end_idx = min((i + 1) * win_size, seq_len)
+#             query = x[:, start_idx:end_idx, :]
+    
+#             # Define the range of keys and values
+#             key_start = max(0, start_idx - span_len + win_size)
+#             key_end = min(start_idx + span_len, seq_len)
+#             key = x[:, key_start:key_end, :]
+#             value = x[:, key_start:key_end, :]
+    
+#             attn_out, _ = self._focus(query, key, value, span_scale)
+#             output[:, start_idx:end_idx, :] = attn_out
+    
+#         return output
+