I noticed a discrepancy in how attention weights are summed in pyramidkv/pyramidkv_utils.py:
PyramidKV implementation (Line 231):
attn_weights_sum = attn_weights[:, :, -self.window_size:, :-self.window_size].sum(dim=-2)
Other baselines (e.g., Line 554):
attn_weights_sum = attn_weights[:, :, :, :-self.window_size].sum(dim=-2)
The key difference is that PyramidKV only considers the last self.window_size positions in the third dimension (-self.window_size:), while other baselines use all positions (:).
Issue: This approach seems problematic when important instructions appear at the beginning of the sequence, as PyramidKV would ignore attention weights from early positions. This could explain why performance degrades significantly when instructions are placed at the start rather than the end.
Question: Is this selective attention weight calculation intentional, or could this be causing the performance drop when instructions appear early in the sequence? Should PyramidKV also consider attention weights from all positions like the other baselines?
I noticed a discrepancy in how attention weights are summed in
pyramidkv/pyramidkv_utils.py:PyramidKV implementation (Line 231):
Other baselines (e.g., Line 554):
The key difference is that PyramidKV only considers the last
self.window_sizepositions in the third dimension (-self.window_size:), while other baselines use all positions (:).Issue: This approach seems problematic when important instructions appear at the beginning of the sequence, as PyramidKV would ignore attention weights from early positions. This could explain why performance degrades significantly when instructions are placed at the start rather than the end.
Question: Is this selective attention weight calculation intentional, or could this be causing the performance drop when instructions appear early in the sequence? Should PyramidKV also consider attention weights from all positions like the other baselines?