ReFNet changes: Use encoders for refiner pooler inputs and fix edge case for MS loss calculation when no negative pairs found. (#1160)

Summary:
Pull Request resolved: #1160

Loss value dict needs to be initialized and set to default values when no negatives per-batch (for MS loss calculation). Also, the refiner class should take in encoded layers rather than sequence output.

Reviewed By: ebsmothers

Differential Revision: D32696676

fbshipit-source-id: b52d8532445141499152353c7893fb83ef6142c4

Author: Sethu Sankaran

mmf/models/transformers/heads/refiner.py

@@ -120,10 +120,10 @@ def forward(
             end_token[modality] = start_token[modality] + sz[1] - 1
             prev_end_token = end_token[modality] + 1
 
-        attention_mask = torch.cat(masks, dim=1)
+        pad_mask = torch.cat(masks, dim=1)
         processed_sample_list["refiner_outputs"] = {}
         processed_sample_list["refiner_outputs"]["fused_embedding"] = self.pooler(
-            sequence_output, attention_mask
+            encoded_layers, pad_mask
         )
         processed_sample_list["refiner_targets"] = {}
         for modality in self.modalities:
@@ -132,7 +132,7 @@ def forward(
             tk_end = end_token[modality]
             for enc_layers in encoded_layers:
                 modality_encodings.append(enc_layers[:, tk_start : tk_end + 1, :])
-            modality_mask_encodings = attention_mask[:, tk_start : tk_end + 1]
+            modality_mask_encodings = pad_mask[:, tk_start : tk_end + 1]
             processed_sample_list["refiner_targets"][modality] = self.pooler(
                 modality_encodings, modality_mask_encodings
             )
@@ -169,7 +169,7 @@ def forward(
                 ]
                 refiner_modal_outputs = {}
                 refiner_modal_outputs["scores"] = refiner_reconstruct[modality]
-                loss = self.refinerloss(modality_targets, refiner_modal_outputs)
+                loss = self.refiner_loss(modality_targets, refiner_modal_outputs)
 
             else:
                 loss = self.weights[modality] * self.refiner_loss(

mmf/models/transformers/heads/refnet_classifier.py

@@ -58,6 +58,8 @@ def forward(
             targets_subset = {}
             targets_subset["targets"] = processed_sample_list["target_key"]["targets"]
             targets_subset["targets"] = targets_subset["targets"][:score_max]
+            if "losses" not in output_dict.keys():
+                output_dict["losses"] = {}
             output_dict["losses"][self.loss_name] = self.loss_fn(
                 targets_subset, scores_subset
             )

mmf/modules/losses.py

@@ -933,8 +933,10 @@ def forward(self, sample_list, model_output):
             pos_loss = calc_ms_loss(pos_pair, self.base, self.beta, -1)
             neg_loss = calc_ms_loss(neg_pairs, self.base, self.alpha, 1)
             loss.append(pos_loss + neg_loss)
-
-        loss = sum(loss) / n
+        if n > 0:
+            loss = sum(loss) / n
+        else:
+            loss = inputs.new_zeros(1, requires_grad=True)
         return loss
 
 

projects/mmbt/configs/mmimdb/paper_ablations_reducedlabel.yaml

@@ -36,23 +36,30 @@ model_config:
           hidden_size: 768
           vocab_size: 30522
           loss_type: "cosine"
+          refiner_target_pooler: "average_k_from_last"
+          refiner_target_layer_depth: 2
           modalities:
             - "text"
             - "image"
           weights:
             - 0.0
             - 0.0
-        mlp_config:
-          type: mlp
-          freeze: false
-          num_labels: 24
-          lr_multiplier: 1.0
-          hidden_size: 768
-          vocab_size: 30522
+        mlp_loss_config:
+          config:
+            type: mlp
+            num_labels: 24
+            hidden_size: 768
+            hidden_dropout_prob: 0.1
+            layer_norm_eps: 0.000001
+            hidden_act: gelu
+            pooler_name: bert_pooler
+            num_layers: 1
+          loss_name: classification_loss
           loss: logit_bce
-          max_sample_size: 32
-        ms_loss_weight: 0.05
+          max_sample_size: 33
+        ms_loss_weight: 0.0
         use_msloss: true
+        num_labels: 24
     self_weight_decay: 0.997
 dataset_config:
   mmimdb: