Repository: huggingface/transformers
Pull request: [EfficientLoFTR] Fixed bugs and added original implementation bug

# What does this PR do?

Fixes a couple of bugs introduced in the implementation of EfficientLoFTR.
Adds a cross attention bug introduced by the authors which makes the weights dependant on (see https://github.com/zju3dv/MatchAnything/issues/26)

## Who can review?

@qubvel 


Changed files:
filename: src/transformers/models/efficientloftr/modeling_efficientloftr.py
patch: @@ -322,7 +322,6 @@ def eager_attention_forward(
 class EfficientLoFTRAttention(nn.Module):
     """Multi-headed attention from 'Attention Is All You Need' paper"""
 
-    # Copied from transformers.models.llama.modeling_llama.LlamaAttention.__init__ with Llama->EfficientLoFTR
     def __init__(self, config: EfficientLoFTRConfig, layer_idx: int):
         super().__init__()
         self.config = config
@@ -331,7 +330,7 @@ def __init__(self, config: EfficientLoFTRConfig, layer_idx: int):
         self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads
         self.scaling = self.head_dim**-0.5
         self.attention_dropout = config.attention_dropout
-        self.is_causal = True
+        self.is_causal = False
 
         self.q_proj = nn.Linear(
             config.hidden_size, config.num_attention_heads * self.head_dim, bias=config.attention_bias
@@ -478,12 +477,16 @@ def forward(
         hidden_states = hidden_states.reshape(-1, embed_dim, height, width)
         hidden_states = self.self_attention(hidden_states, position_embeddings=position_embeddings, **kwargs)
 
-        encoder_hidden_states = hidden_states.reshape(-1, 2, embed_dim, height, width)
-        encoder_hidden_states = encoder_hidden_states.flip(1)
-        encoder_hidden_states = encoder_hidden_states.reshape(-1, embed_dim, height, width)
-
-        hidden_states = self.cross_attention(hidden_states, encoder_hidden_states, **kwargs)
-        hidden_states = hidden_states.reshape(batch_size, -1, embed_dim, height, width)
+        ###
+        # Implementation of a bug in the original implementation regarding the cross-attention
+        # See : https://github.com/zju3dv/MatchAnything/issues/26
+        hidden_states = hidden_states.reshape(-1, 2, embed_dim, height, width)
+        features_0 = hidden_states[:, 0]
+        features_1 = hidden_states[:, 1]
+        features_0 = self.cross_attention(features_0, features_1, **kwargs)
+        features_1 = self.cross_attention(features_1, features_0, **kwargs)
+        hidden_states = torch.stack((features_0, features_1), dim=1)
+        ###
 
         return hidden_states
 
@@ -750,8 +753,15 @@ def mask_border(tensor: torch.Tensor, border_margin: int, value: Union[bool, flo
     if border_margin <= 0:
         return tensor
 
-    tensor[:, :border_margin, :border_margin, :border_margin, :border_margin] = value
-    tensor[:, -border_margin:, -border_margin:, -border_margin:, -border_margin:] = value
+    tensor[:, :border_margin] = value
+    tensor[:, :, :border_margin] = value
+    tensor[:, :, :, :border_margin] = value
+    tensor[:, :, :, :, :border_margin] = value
+    tensor[:, -border_margin:] = value
+    tensor[:, :, -border_margin:] = value
+    tensor[:, :, :, -border_margin:] = value
+    tensor[:, :, :, :, -border_margin:] = value
+
     return tensor
 
 
@@ -1276,6 +1286,7 @@ def forward(
 
         # 3. Fine-level refinement
         residual_features = features[1:]
+        coarse_features = coarse_features / self.config.hidden_size**0.5
         fine_features_0, fine_features_1 = self.refinement_layer(coarse_features, residual_features)
 
         # Filter fine features with coarse matches indices

filename: tests/models/efficientloftr/test_modeling_efficientloftr.py
patch: @@ -436,10 +436,10 @@ def test_inference(self):
         expected_matching_scores_shape = torch.Size((len(images), 2, expected_number_of_matches))
 
         expected_top10_matches_indices = torch.tensor(
-            [3145, 3065, 3143, 3066, 3144, 1397, 1705, 3151, 2342, 2422], dtype=torch.int64, device=torch_device
+            [3145, 3065, 3143, 3144, 1397, 1705, 3151, 2422, 3066, 2342], dtype=torch.int64, device=torch_device
         )
         expected_top10_matching_scores = torch.tensor(
-            [0.9997, 0.9996, 0.9996, 0.9995, 0.9995, 0.9995, 0.9994, 0.9994, 0.9994, 0.9994], device=torch_device
+            [0.9998, 0.9997, 0.9997, 0.9996, 0.9996, 0.9996, 0.9996, 0.9995, 0.9995, 0.9995], device=torch_device
         )
 
         self.assertEqual(outputs.matches.shape, expected_matches_shape)