diff --git "a/trainer_state.json" "b/trainer_state.json"
new file mode 100644--- /dev/null
+++ "b/trainer_state.json"
@@ -0,0 +1,97630 @@
+{
+  "best_metric": null,
+  "best_model_checkpoint": null,
+  "epoch": 1.0,
+  "eval_steps": 500,
+  "global_step": 9760,
+  "is_hyper_param_search": false,
+  "is_local_process_zero": true,
+  "is_world_process_zero": true,
+  "log_history": [
+    {
+      "epoch": 0.0,
+      "gumbel_tau": 1.0,
+      "image_masks_loss": 9.375,
+      "llm_learning_rate": 1.7064846416382254e-08,
+      "loss": 10.4927,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 1.7064846416382253e-06,
+      "step": 1
+    },
+    {
+      "epoch": 0.0,
+      "gumbel_tau": 0.9998,
+      "image_masks_loss": 9.1797,
+      "llm_learning_rate": 3.412969283276451e-08,
+      "loss": 10.4797,
+      "output_text_mask_loss": 0.0984,
+      "predictor_learning_rate": 3.4129692832764506e-06,
+      "step": 2
+    },
+    {
+      "epoch": 0.0,
+      "gumbel_tau": 0.9995,
+      "image_masks_loss": 8.3496,
+      "llm_learning_rate": 5.119453924914676e-08,
+      "loss": 9.6042,
+      "output_text_mask_loss": 0.0275,
+      "predictor_learning_rate": 5.119453924914675e-06,
+      "step": 3
+    },
+    {
+      "epoch": 0.0,
+      "gumbel_tau": 0.9993,
+      "image_masks_loss": 8.6914,
+      "llm_learning_rate": 6.825938566552902e-08,
+      "loss": 10.0235,
+      "output_text_mask_loss": 0.0391,
+      "predictor_learning_rate": 6.825938566552901e-06,
+      "step": 4
+    },
+    {
+      "epoch": 0.0,
+      "gumbel_tau": 0.9991,
+      "image_masks_loss": 8.4961,
+      "llm_learning_rate": 8.532423208191127e-08,
+      "loss": 9.7061,
+      "output_text_mask_loss": 0.033,
+      "predictor_learning_rate": 8.532423208191126e-06,
+      "step": 5
+    },
+    {
+      "epoch": 0.0,
+      "gumbel_tau": 0.9988,
+      "image_masks_loss": 8.9844,
+      "llm_learning_rate": 1.0238907849829352e-07,
+      "loss": 10.2919,
+      "output_text_mask_loss": 0.0433,
+      "predictor_learning_rate": 1.023890784982935e-05,
+      "step": 6
+    },
+    {
+      "epoch": 0.0,
+      "gumbel_tau": 0.9986,
+      "image_masks_loss": 9.7656,
+      "llm_learning_rate": 1.1945392491467578e-07,
+      "loss": 10.4853,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 1.1945392491467578e-05,
+      "step": 7
+    },
+    {
+      "epoch": 0.0,
+      "gumbel_tau": 0.9983,
+      "image_masks_loss": 8.8867,
+      "llm_learning_rate": 1.3651877133105803e-07,
+      "loss": 10.1011,
+      "output_text_mask_loss": 0.0259,
+      "predictor_learning_rate": 1.3651877133105803e-05,
+      "step": 8
+    },
+    {
+      "epoch": 0.0,
+      "gumbel_tau": 0.9981,
+      "image_masks_loss": 8.252,
+      "llm_learning_rate": 1.535836177474403e-07,
+      "loss": 9.3879,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 1.5358361774744027e-05,
+      "step": 9
+    },
+    {
+      "epoch": 0.0,
+      "gumbel_tau": 0.9979,
+      "image_masks_loss": 8.1543,
+      "llm_learning_rate": 1.7064846416382255e-07,
+      "loss": 9.652,
+      "output_text_mask_loss": 0.0565,
+      "predictor_learning_rate": 1.7064846416382253e-05,
+      "step": 10
+    },
+    {
+      "epoch": 0.0,
+      "gumbel_tau": 0.9976,
+      "image_masks_loss": 8.3496,
+      "llm_learning_rate": 1.877133105802048e-07,
+      "loss": 9.2976,
+      "output_text_mask_loss": 0.0282,
+      "predictor_learning_rate": 1.877133105802048e-05,
+      "step": 11
+    },
+    {
+      "epoch": 0.0,
+      "gumbel_tau": 0.9974,
+      "image_masks_loss": 7.959,
+      "llm_learning_rate": 2.0477815699658704e-07,
+      "loss": 9.0414,
+      "output_text_mask_loss": 0.0629,
+      "predictor_learning_rate": 2.04778156996587e-05,
+      "step": 12
+    },
+    {
+      "epoch": 0.0,
+      "gumbel_tau": 0.9972,
+      "image_masks_loss": 8.1055,
+      "llm_learning_rate": 2.218430034129693e-07,
+      "loss": 9.2871,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 2.218430034129693e-05,
+      "step": 13
+    },
+    {
+      "epoch": 0.0,
+      "gumbel_tau": 0.9969,
+      "image_masks_loss": 6.7383,
+      "llm_learning_rate": 2.3890784982935155e-07,
+      "loss": 8.0306,
+      "output_text_mask_loss": 0.0278,
+      "predictor_learning_rate": 2.3890784982935157e-05,
+      "step": 14
+    },
+    {
+      "epoch": 0.0,
+      "gumbel_tau": 0.9967,
+      "image_masks_loss": 7.0312,
+      "llm_learning_rate": 2.559726962457338e-07,
+      "loss": 8.1173,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 2.559726962457338e-05,
+      "step": 15
+    },
+    {
+      "epoch": 0.0,
+      "gumbel_tau": 0.9965,
+      "image_masks_loss": 6.9824,
+      "llm_learning_rate": 2.7303754266211607e-07,
+      "loss": 8.2247,
+      "output_text_mask_loss": 0.0023,
+      "predictor_learning_rate": 2.7303754266211605e-05,
+      "step": 16
+    },
+    {
+      "epoch": 0.0,
+      "gumbel_tau": 0.9962,
+      "image_masks_loss": 6.5918,
+      "llm_learning_rate": 2.901023890784983e-07,
+      "loss": 7.8378,
+      "output_text_mask_loss": 0.0614,
+      "predictor_learning_rate": 2.901023890784983e-05,
+      "step": 17
+    },
+    {
+      "epoch": 0.0,
+      "gumbel_tau": 0.996,
+      "image_masks_loss": 5.9082,
+      "llm_learning_rate": 3.071672354948806e-07,
+      "loss": 7.2037,
+      "output_text_mask_loss": 0.0916,
+      "predictor_learning_rate": 3.0716723549488054e-05,
+      "step": 18
+    },
+    {
+      "epoch": 0.0,
+      "gumbel_tau": 0.9958,
+      "image_masks_loss": 5.542,
+      "llm_learning_rate": 3.242320819112628e-07,
+      "loss": 6.8942,
+      "output_text_mask_loss": 0.0429,
+      "predictor_learning_rate": 3.2423208191126276e-05,
+      "step": 19
+    },
+    {
+      "epoch": 0.0,
+      "gumbel_tau": 0.9955,
+      "image_masks_loss": 5.2734,
+      "llm_learning_rate": 3.412969283276451e-07,
+      "loss": 6.5651,
+      "output_text_mask_loss": 0.0278,
+      "predictor_learning_rate": 3.4129692832764505e-05,
+      "step": 20
+    },
+    {
+      "epoch": 0.0,
+      "gumbel_tau": 0.9953,
+      "image_masks_loss": 4.2725,
+      "llm_learning_rate": 3.583617747440273e-07,
+      "loss": 5.473,
+      "output_text_mask_loss": 0.1503,
+      "predictor_learning_rate": 3.583617747440273e-05,
+      "step": 21
+    },
+    {
+      "epoch": 0.0,
+      "gumbel_tau": 0.9951,
+      "image_masks_loss": 4.3945,
+      "llm_learning_rate": 3.754266211604096e-07,
+      "loss": 5.5969,
+      "output_text_mask_loss": 0.2563,
+      "predictor_learning_rate": 3.754266211604096e-05,
+      "step": 22
+    },
+    {
+      "epoch": 0.0,
+      "gumbel_tau": 0.9948,
+      "image_masks_loss": 3.8818,
+      "llm_learning_rate": 3.9249146757679187e-07,
+      "loss": 5.2452,
+      "output_text_mask_loss": 0.0343,
+      "predictor_learning_rate": 3.924914675767918e-05,
+      "step": 23
+    },
+    {
+      "epoch": 0.0,
+      "gumbel_tau": 0.9946,
+      "image_masks_loss": 2.3315,
+      "llm_learning_rate": 4.0955631399317407e-07,
+      "loss": 3.371,
+      "output_text_mask_loss": 0.0101,
+      "predictor_learning_rate": 4.09556313993174e-05,
+      "step": 24
+    },
+    {
+      "epoch": 0.0,
+      "gumbel_tau": 0.9944,
+      "image_masks_loss": 2.4658,
+      "llm_learning_rate": 4.266211604095564e-07,
+      "loss": 3.5179,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 4.266211604095564e-05,
+      "step": 25
+    },
+    {
+      "epoch": 0.0,
+      "gumbel_tau": 0.9941,
+      "image_masks_loss": 1.5991,
+      "llm_learning_rate": 4.436860068259386e-07,
+      "loss": 2.9296,
+      "output_text_mask_loss": 0.0452,
+      "predictor_learning_rate": 4.436860068259386e-05,
+      "step": 26
+    },
+    {
+      "epoch": 0.0,
+      "gumbel_tau": 0.9939,
+      "image_masks_loss": 0.9949,
+      "llm_learning_rate": 4.6075085324232084e-07,
+      "loss": 2.1732,
+      "output_text_mask_loss": 0.0095,
+      "predictor_learning_rate": 4.6075085324232084e-05,
+      "step": 27
+    },
+    {
+      "epoch": 0.0,
+      "gumbel_tau": 0.9937,
+      "image_masks_loss": 0.5463,
+      "llm_learning_rate": 4.778156996587031e-07,
+      "loss": 1.8103,
+      "output_text_mask_loss": 0.1404,
+      "predictor_learning_rate": 4.778156996587031e-05,
+      "step": 28
+    },
+    {
+      "epoch": 0.0,
+      "gumbel_tau": 0.9934,
+      "image_masks_loss": 0.2258,
+      "llm_learning_rate": 4.948805460750854e-07,
+      "loss": 1.4283,
+      "output_text_mask_loss": 0.0832,
+      "predictor_learning_rate": 4.9488054607508536e-05,
+      "step": 29
+    },
+    {
+      "epoch": 0.0,
+      "gumbel_tau": 0.9932,
+      "image_masks_loss": 0.1038,
+      "llm_learning_rate": 5.119453924914676e-07,
+      "loss": 1.414,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 5.119453924914676e-05,
+      "step": 30
+    },
+    {
+      "epoch": 0.0,
+      "gumbel_tau": 0.9929,
+      "image_masks_loss": 0.0345,
+      "llm_learning_rate": 5.290102389078499e-07,
+      "loss": 1.109,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 5.290102389078499e-05,
+      "step": 31
+    },
+    {
+      "epoch": 0.0,
+      "gumbel_tau": 0.9927,
+      "image_masks_loss": 0.1999,
+      "llm_learning_rate": 5.460750853242321e-07,
+      "loss": 1.4674,
+      "output_text_mask_loss": 0.1068,
+      "predictor_learning_rate": 5.460750853242321e-05,
+      "step": 32
+    },
+    {
+      "epoch": 0.0,
+      "gumbel_tau": 0.9925,
+      "image_masks_loss": 0.415,
+      "llm_learning_rate": 5.631399317406143e-07,
+      "loss": 1.5806,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 5.631399317406143e-05,
+      "step": 33
+    },
+    {
+      "epoch": 0.0,
+      "gumbel_tau": 0.9922,
+      "image_masks_loss": 0.3235,
+      "llm_learning_rate": 5.802047781569966e-07,
+      "loss": 1.6224,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 5.802047781569966e-05,
+      "step": 34
+    },
+    {
+      "epoch": 0.0,
+      "gumbel_tau": 0.992,
+      "image_masks_loss": 0.531,
+      "llm_learning_rate": 5.972696245733788e-07,
+      "loss": 1.8502,
+      "output_text_mask_loss": 0.0052,
+      "predictor_learning_rate": 5.9726962457337885e-05,
+      "step": 35
+    },
+    {
+      "epoch": 0.0,
+      "gumbel_tau": 0.9918,
+      "image_masks_loss": 0.5188,
+      "llm_learning_rate": 6.143344709897612e-07,
+      "loss": 1.4884,
+      "output_text_mask_loss": 0.0235,
+      "predictor_learning_rate": 6.143344709897611e-05,
+      "step": 36
+    },
+    {
+      "epoch": 0.0,
+      "gumbel_tau": 0.9915,
+      "image_masks_loss": 0.589,
+      "llm_learning_rate": 6.313993174061434e-07,
+      "loss": 1.914,
+      "output_text_mask_loss": 0.0322,
+      "predictor_learning_rate": 6.313993174061433e-05,
+      "step": 37
+    },
+    {
+      "epoch": 0.0,
+      "gumbel_tau": 0.9913,
+      "image_masks_loss": 0.4242,
+      "llm_learning_rate": 6.484641638225256e-07,
+      "loss": 1.6582,
+      "output_text_mask_loss": 0.0504,
+      "predictor_learning_rate": 6.484641638225255e-05,
+      "step": 38
+    },
+    {
+      "epoch": 0.0,
+      "gumbel_tau": 0.9911,
+      "image_masks_loss": 0.2838,
+      "llm_learning_rate": 6.65529010238908e-07,
+      "loss": 1.5518,
+      "output_text_mask_loss": 0.0111,
+      "predictor_learning_rate": 6.655290102389079e-05,
+      "step": 39
+    },
+    {
+      "epoch": 0.0,
+      "gumbel_tau": 0.9908,
+      "image_masks_loss": 0.1503,
+      "llm_learning_rate": 6.825938566552902e-07,
+      "loss": 1.3173,
+      "output_text_mask_loss": 0.1717,
+      "predictor_learning_rate": 6.825938566552901e-05,
+      "step": 40
+    },
+    {
+      "epoch": 0.0,
+      "gumbel_tau": 0.9906,
+      "image_masks_loss": 0.0399,
+      "llm_learning_rate": 6.996587030716724e-07,
+      "loss": 1.3421,
+      "output_text_mask_loss": 0.1495,
+      "predictor_learning_rate": 6.996587030716723e-05,
+      "step": 41
+    },
+    {
+      "epoch": 0.0,
+      "gumbel_tau": 0.9904,
+      "image_masks_loss": 0.0244,
+      "llm_learning_rate": 7.167235494880546e-07,
+      "loss": 1.0825,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 7.167235494880546e-05,
+      "step": 42
+    },
+    {
+      "epoch": 0.0,
+      "gumbel_tau": 0.9901,
+      "image_masks_loss": 0.0313,
+      "llm_learning_rate": 7.337883959044369e-07,
+      "loss": 1.3682,
+      "output_text_mask_loss": 0.174,
+      "predictor_learning_rate": 7.337883959044368e-05,
+      "step": 43
+    },
+    {
+      "epoch": 0.0,
+      "gumbel_tau": 0.9899,
+      "image_masks_loss": 0.0835,
+      "llm_learning_rate": 7.508532423208192e-07,
+      "loss": 1.2334,
+      "output_text_mask_loss": 0.0618,
+      "predictor_learning_rate": 7.508532423208191e-05,
+      "step": 44
+    },
+    {
+      "epoch": 0.0,
+      "gumbel_tau": 0.9897,
+      "image_masks_loss": 0.1244,
+      "llm_learning_rate": 7.679180887372014e-07,
+      "loss": 1.3364,
+      "output_text_mask_loss": 0.1213,
+      "predictor_learning_rate": 7.679180887372014e-05,
+      "step": 45
+    },
+    {
+      "epoch": 0.0,
+      "gumbel_tau": 0.9894,
+      "image_masks_loss": 0.2106,
+      "llm_learning_rate": 7.849829351535837e-07,
+      "loss": 1.4084,
+      "output_text_mask_loss": 0.0012,
+      "predictor_learning_rate": 7.849829351535836e-05,
+      "step": 46
+    },
+    {
+      "epoch": 0.0,
+      "gumbel_tau": 0.9892,
+      "image_masks_loss": 0.3799,
+      "llm_learning_rate": 8.020477815699659e-07,
+      "loss": 1.6601,
+      "output_text_mask_loss": 0.0322,
+      "predictor_learning_rate": 8.020477815699658e-05,
+      "step": 47
+    },
+    {
+      "epoch": 0.0,
+      "gumbel_tau": 0.989,
+      "image_masks_loss": 0.3967,
+      "llm_learning_rate": 8.191126279863481e-07,
+      "loss": 1.6461,
+      "output_text_mask_loss": 0.1907,
+      "predictor_learning_rate": 8.19112627986348e-05,
+      "step": 48
+    },
+    {
+      "epoch": 0.01,
+      "gumbel_tau": 0.9887,
+      "image_masks_loss": 0.3403,
+      "llm_learning_rate": 8.361774744027303e-07,
+      "loss": 1.6222,
+      "output_text_mask_loss": 0.3571,
+      "predictor_learning_rate": 8.361774744027303e-05,
+      "step": 49
+    },
+    {
+      "epoch": 0.01,
+      "gumbel_tau": 0.9885,
+      "image_masks_loss": 0.293,
+      "llm_learning_rate": 8.532423208191128e-07,
+      "loss": 1.5358,
+      "output_text_mask_loss": 0.1602,
+      "predictor_learning_rate": 8.532423208191128e-05,
+      "step": 50
+    },
+    {
+      "epoch": 0.01,
+      "gumbel_tau": 0.9883,
+      "image_masks_loss": 0.2975,
+      "llm_learning_rate": 8.70307167235495e-07,
+      "loss": 1.4384,
+      "output_text_mask_loss": 0.058,
+      "predictor_learning_rate": 8.70307167235495e-05,
+      "step": 51
+    },
+    {
+      "epoch": 0.01,
+      "gumbel_tau": 0.988,
+      "image_masks_loss": 0.1099,
+      "llm_learning_rate": 8.873720136518772e-07,
+      "loss": 1.4917,
+      "output_text_mask_loss": 0.0429,
+      "predictor_learning_rate": 8.873720136518772e-05,
+      "step": 52
+    },
+    {
+      "epoch": 0.01,
+      "gumbel_tau": 0.9878,
+      "image_masks_loss": 0.1518,
+      "llm_learning_rate": 9.044368600682595e-07,
+      "loss": 1.2001,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 9.044368600682594e-05,
+      "step": 53
+    },
+    {
+      "epoch": 0.01,
+      "gumbel_tau": 0.9876,
+      "image_masks_loss": 0.0572,
+      "llm_learning_rate": 9.215017064846417e-07,
+      "loss": 1.435,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 9.215017064846417e-05,
+      "step": 54
+    },
+    {
+      "epoch": 0.01,
+      "gumbel_tau": 0.9873,
+      "image_masks_loss": 0.0448,
+      "llm_learning_rate": 9.385665529010239e-07,
+      "loss": 1.2901,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 9.385665529010239e-05,
+      "step": 55
+    },
+    {
+      "epoch": 0.01,
+      "gumbel_tau": 0.9871,
+      "image_masks_loss": 0.0607,
+      "llm_learning_rate": 9.556313993174062e-07,
+      "loss": 1.1596,
+      "output_text_mask_loss": 0.0618,
+      "predictor_learning_rate": 9.556313993174063e-05,
+      "step": 56
+    },
+    {
+      "epoch": 0.01,
+      "gumbel_tau": 0.9869,
+      "image_masks_loss": 0.0326,
+      "llm_learning_rate": 9.726962457337885e-07,
+      "loss": 1.2636,
+      "output_text_mask_loss": 0.2472,
+      "predictor_learning_rate": 9.726962457337885e-05,
+      "step": 57
+    },
+    {
+      "epoch": 0.01,
+      "gumbel_tau": 0.9866,
+      "image_masks_loss": 0.0341,
+      "llm_learning_rate": 9.897610921501708e-07,
+      "loss": 1.3986,
+      "output_text_mask_loss": 0.1007,
+      "predictor_learning_rate": 9.897610921501707e-05,
+      "step": 58
+    },
+    {
+      "epoch": 0.01,
+      "gumbel_tau": 0.9864,
+      "image_masks_loss": 0.0759,
+      "llm_learning_rate": 1.006825938566553e-06,
+      "loss": 1.3569,
+      "output_text_mask_loss": 0.4211,
+      "predictor_learning_rate": 0.0001006825938566553,
+      "step": 59
+    },
+    {
+      "epoch": 0.01,
+      "gumbel_tau": 0.9862,
+      "image_masks_loss": 0.177,
+      "llm_learning_rate": 1.0238907849829352e-06,
+      "loss": 1.4202,
+      "output_text_mask_loss": 0.4089,
+      "predictor_learning_rate": 0.00010238907849829352,
+      "step": 60
+    },
+    {
+      "epoch": 0.01,
+      "gumbel_tau": 0.9859,
+      "image_masks_loss": 0.0263,
+      "llm_learning_rate": 1.0409556313993173e-06,
+      "loss": 1.007,
+      "output_text_mask_loss": 0.1099,
+      "predictor_learning_rate": 0.00010409556313993174,
+      "step": 61
+    },
+    {
+      "epoch": 0.01,
+      "gumbel_tau": 0.9857,
+      "image_masks_loss": 0.0969,
+      "llm_learning_rate": 1.0580204778156999e-06,
+      "loss": 1.1982,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00010580204778156998,
+      "step": 62
+    },
+    {
+      "epoch": 0.01,
+      "gumbel_tau": 0.9855,
+      "image_masks_loss": 0.074,
+      "llm_learning_rate": 1.075085324232082e-06,
+      "loss": 1.2425,
+      "output_text_mask_loss": 0.4578,
+      "predictor_learning_rate": 0.0001075085324232082,
+      "step": 63
+    },
+    {
+      "epoch": 0.01,
+      "gumbel_tau": 0.9852,
+      "image_masks_loss": 0.0858,
+      "llm_learning_rate": 1.0921501706484643e-06,
+      "loss": 1.3013,
+      "output_text_mask_loss": 1.0498,
+      "predictor_learning_rate": 0.00010921501706484642,
+      "step": 64
+    },
+    {
+      "epoch": 0.01,
+      "gumbel_tau": 0.985,
+      "image_masks_loss": 0.0694,
+      "llm_learning_rate": 1.1092150170648466e-06,
+      "loss": 1.6468,
+      "output_text_mask_loss": 0.5066,
+      "predictor_learning_rate": 0.00011092150170648464,
+      "step": 65
+    },
+    {
+      "epoch": 0.01,
+      "gumbel_tau": 0.9848,
+      "image_masks_loss": 0.0313,
+      "llm_learning_rate": 1.1262798634812287e-06,
+      "loss": 1.1948,
+      "output_text_mask_loss": 0.1862,
+      "predictor_learning_rate": 0.00011262798634812287,
+      "step": 66
+    },
+    {
+      "epoch": 0.01,
+      "gumbel_tau": 0.9845,
+      "image_masks_loss": 0.0196,
+      "llm_learning_rate": 1.143344709897611e-06,
+      "loss": 1.2201,
+      "output_text_mask_loss": 0.3204,
+      "predictor_learning_rate": 0.00011433447098976109,
+      "step": 67
+    },
+    {
+      "epoch": 0.01,
+      "gumbel_tau": 0.9843,
+      "image_masks_loss": 0.1053,
+      "llm_learning_rate": 1.1604095563139933e-06,
+      "loss": 1.8991,
+      "output_text_mask_loss": 0.1076,
+      "predictor_learning_rate": 0.00011604095563139932,
+      "step": 68
+    },
+    {
+      "epoch": 0.01,
+      "gumbel_tau": 0.9841,
+      "image_masks_loss": 0.0984,
+      "llm_learning_rate": 1.1774744027303756e-06,
+      "loss": 1.6352,
+      "output_text_mask_loss": 0.3693,
+      "predictor_learning_rate": 0.00011774744027303755,
+      "step": 69
+    },
+    {
+      "epoch": 0.01,
+      "gumbel_tau": 0.9839,
+      "image_masks_loss": 0.1335,
+      "llm_learning_rate": 1.1945392491467577e-06,
+      "loss": 1.1763,
+      "output_text_mask_loss": 0.8484,
+      "predictor_learning_rate": 0.00011945392491467577,
+      "step": 70
+    },
+    {
+      "epoch": 0.01,
+      "gumbel_tau": 0.9836,
+      "image_masks_loss": 0.1915,
+      "llm_learning_rate": 1.21160409556314e-06,
+      "loss": 1.1943,
+      "output_text_mask_loss": 0.1305,
+      "predictor_learning_rate": 0.00012116040955631399,
+      "step": 71
+    },
+    {
+      "epoch": 0.01,
+      "gumbel_tau": 0.9834,
+      "image_masks_loss": 0.1564,
+      "llm_learning_rate": 1.2286689419795223e-06,
+      "loss": 1.228,
+      "output_text_mask_loss": 0.2579,
+      "predictor_learning_rate": 0.00012286689419795221,
+      "step": 72
+    },
+    {
+      "epoch": 0.01,
+      "gumbel_tau": 0.9832,
+      "image_masks_loss": 0.0511,
+      "llm_learning_rate": 1.2457337883959044e-06,
+      "loss": 1.3548,
+      "output_text_mask_loss": 0.7935,
+      "predictor_learning_rate": 0.00012457337883959045,
+      "step": 73
+    },
+    {
+      "epoch": 0.01,
+      "gumbel_tau": 0.9829,
+      "image_masks_loss": 0.053,
+      "llm_learning_rate": 1.2627986348122867e-06,
+      "loss": 1.4312,
+      "output_text_mask_loss": 0.8179,
+      "predictor_learning_rate": 0.00012627986348122866,
+      "step": 74
+    },
+    {
+      "epoch": 0.01,
+      "gumbel_tau": 0.9827,
+      "image_masks_loss": 0.0232,
+      "llm_learning_rate": 1.279863481228669e-06,
+      "loss": 1.6028,
+      "output_text_mask_loss": 0.2609,
+      "predictor_learning_rate": 0.0001279863481228669,
+      "step": 75
+    },
+    {
+      "epoch": 0.01,
+      "gumbel_tau": 0.9825,
+      "image_masks_loss": 0.0054,
+      "llm_learning_rate": 1.2969283276450511e-06,
+      "loss": 1.2457,
+      "output_text_mask_loss": 0.322,
+      "predictor_learning_rate": 0.0001296928327645051,
+      "step": 76
+    },
+    {
+      "epoch": 0.01,
+      "gumbel_tau": 0.9822,
+      "image_masks_loss": 0.0565,
+      "llm_learning_rate": 1.3139931740614337e-06,
+      "loss": 1.4503,
+      "output_text_mask_loss": 0.4822,
+      "predictor_learning_rate": 0.00013139931740614337,
+      "step": 77
+    },
+    {
+      "epoch": 0.01,
+      "gumbel_tau": 0.982,
+      "image_masks_loss": 0.0332,
+      "llm_learning_rate": 1.331058020477816e-06,
+      "loss": 1.6198,
+      "output_text_mask_loss": 0.7111,
+      "predictor_learning_rate": 0.00013310580204778158,
+      "step": 78
+    },
+    {
+      "epoch": 0.01,
+      "gumbel_tau": 0.9818,
+      "image_masks_loss": 0.0504,
+      "llm_learning_rate": 1.348122866894198e-06,
+      "loss": 1.4199,
+      "output_text_mask_loss": 1.1108,
+      "predictor_learning_rate": 0.0001348122866894198,
+      "step": 79
+    },
+    {
+      "epoch": 0.01,
+      "gumbel_tau": 0.9815,
+      "image_masks_loss": 0.032,
+      "llm_learning_rate": 1.3651877133105804e-06,
+      "loss": 2.066,
+      "output_text_mask_loss": 1.2695,
+      "predictor_learning_rate": 0.00013651877133105802,
+      "step": 80
+    },
+    {
+      "epoch": 0.01,
+      "gumbel_tau": 0.9813,
+      "image_masks_loss": 0.0357,
+      "llm_learning_rate": 1.3822525597269627e-06,
+      "loss": 1.9931,
+      "output_text_mask_loss": 2.1606,
+      "predictor_learning_rate": 0.00013822525597269626,
+      "step": 81
+    },
+    {
+      "epoch": 0.01,
+      "gumbel_tau": 0.9811,
+      "image_masks_loss": 0.0406,
+      "llm_learning_rate": 1.3993174061433448e-06,
+      "loss": 1.614,
+      "output_text_mask_loss": 1.0254,
+      "predictor_learning_rate": 0.00013993174061433447,
+      "step": 82
+    },
+    {
+      "epoch": 0.01,
+      "gumbel_tau": 0.9808,
+      "image_masks_loss": 0.0828,
+      "llm_learning_rate": 1.416382252559727e-06,
+      "loss": 2.0747,
+      "output_text_mask_loss": 0.9033,
+      "predictor_learning_rate": 0.0001416382252559727,
+      "step": 83
+    },
+    {
+      "epoch": 0.01,
+      "gumbel_tau": 0.9806,
+      "image_masks_loss": 0.0145,
+      "llm_learning_rate": 1.4334470989761092e-06,
+      "loss": 1.9274,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0001433447098976109,
+      "step": 84
+    },
+    {
+      "epoch": 0.01,
+      "gumbel_tau": 0.9804,
+      "image_masks_loss": 0.1251,
+      "llm_learning_rate": 1.4505119453924915e-06,
+      "loss": 2.3284,
+      "output_text_mask_loss": 1.1963,
+      "predictor_learning_rate": 0.00014505119453924915,
+      "step": 85
+    },
+    {
+      "epoch": 0.01,
+      "gumbel_tau": 0.9801,
+      "image_masks_loss": 0.0353,
+      "llm_learning_rate": 1.4675767918088738e-06,
+      "loss": 1.7805,
+      "output_text_mask_loss": 1.7334,
+      "predictor_learning_rate": 0.00014675767918088736,
+      "step": 86
+    },
+    {
+      "epoch": 0.01,
+      "gumbel_tau": 0.9799,
+      "image_masks_loss": 0.037,
+      "llm_learning_rate": 1.484641638225256e-06,
+      "loss": 2.01,
+      "output_text_mask_loss": 0.7202,
+      "predictor_learning_rate": 0.0001484641638225256,
+      "step": 87
+    },
+    {
+      "epoch": 0.01,
+      "gumbel_tau": 0.9797,
+      "image_masks_loss": 0.0584,
+      "llm_learning_rate": 1.5017064846416384e-06,
+      "loss": 1.4137,
+      "output_text_mask_loss": 0.5524,
+      "predictor_learning_rate": 0.00015017064846416383,
+      "step": 88
+    },
+    {
+      "epoch": 0.01,
+      "gumbel_tau": 0.9795,
+      "image_masks_loss": 0.0402,
+      "llm_learning_rate": 1.5187713310580207e-06,
+      "loss": 1.7328,
+      "output_text_mask_loss": 0.4395,
+      "predictor_learning_rate": 0.00015187713310580207,
+      "step": 89
+    },
+    {
+      "epoch": 0.01,
+      "gumbel_tau": 0.9792,
+      "image_masks_loss": 0.0263,
+      "llm_learning_rate": 1.5358361774744028e-06,
+      "loss": 1.7785,
+      "output_text_mask_loss": 1.77,
+      "predictor_learning_rate": 0.00015358361774744027,
+      "step": 90
+    },
+    {
+      "epoch": 0.01,
+      "gumbel_tau": 0.979,
+      "image_masks_loss": 0.0439,
+      "llm_learning_rate": 1.5529010238907852e-06,
+      "loss": 1.559,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0001552901023890785,
+      "step": 91
+    },
+    {
+      "epoch": 0.01,
+      "gumbel_tau": 0.9788,
+      "image_masks_loss": 0.0782,
+      "llm_learning_rate": 1.5699658703071675e-06,
+      "loss": 1.8771,
+      "output_text_mask_loss": 0.1602,
+      "predictor_learning_rate": 0.00015699658703071672,
+      "step": 92
+    },
+    {
+      "epoch": 0.01,
+      "gumbel_tau": 0.9785,
+      "image_masks_loss": 0.0462,
+      "llm_learning_rate": 1.5870307167235496e-06,
+      "loss": 1.7002,
+      "output_text_mask_loss": 1.178,
+      "predictor_learning_rate": 0.00015870307167235496,
+      "step": 93
+    },
+    {
+      "epoch": 0.01,
+      "gumbel_tau": 0.9783,
+      "image_masks_loss": 0.0584,
+      "llm_learning_rate": 1.6040955631399319e-06,
+      "loss": 1.8069,
+      "output_text_mask_loss": 0.5798,
+      "predictor_learning_rate": 0.00016040955631399316,
+      "step": 94
+    },
+    {
+      "epoch": 0.01,
+      "gumbel_tau": 0.9781,
+      "image_masks_loss": 0.0496,
+      "llm_learning_rate": 1.6211604095563142e-06,
+      "loss": 1.6974,
+      "output_text_mask_loss": 0.3357,
+      "predictor_learning_rate": 0.0001621160409556314,
+      "step": 95
+    },
+    {
+      "epoch": 0.01,
+      "gumbel_tau": 0.9778,
+      "image_masks_loss": 0.0211,
+      "llm_learning_rate": 1.6382252559726963e-06,
+      "loss": 1.9288,
+      "output_text_mask_loss": 0.9155,
+      "predictor_learning_rate": 0.0001638225255972696,
+      "step": 96
+    },
+    {
+      "epoch": 0.01,
+      "gumbel_tau": 0.9776,
+      "image_masks_loss": 0.0261,
+      "llm_learning_rate": 1.6552901023890786e-06,
+      "loss": 1.4767,
+      "output_text_mask_loss": 0.9888,
+      "predictor_learning_rate": 0.00016552901023890785,
+      "step": 97
+    },
+    {
+      "epoch": 0.01,
+      "gumbel_tau": 0.9774,
+      "image_masks_loss": 0.05,
+      "llm_learning_rate": 1.6723549488054607e-06,
+      "loss": 1.6855,
+      "output_text_mask_loss": 0.9216,
+      "predictor_learning_rate": 0.00016723549488054606,
+      "step": 98
+    },
+    {
+      "epoch": 0.01,
+      "gumbel_tau": 0.9771,
+      "image_masks_loss": 0.0523,
+      "llm_learning_rate": 1.689419795221843e-06,
+      "loss": 1.6608,
+      "output_text_mask_loss": 0.3052,
+      "predictor_learning_rate": 0.0001689419795221843,
+      "step": 99
+    },
+    {
+      "epoch": 0.01,
+      "gumbel_tau": 0.9769,
+      "image_masks_loss": 0.0584,
+      "llm_learning_rate": 1.7064846416382255e-06,
+      "loss": 1.5641,
+      "output_text_mask_loss": 0.264,
+      "predictor_learning_rate": 0.00017064846416382255,
+      "step": 100
+    },
+    {
+      "epoch": 0.01,
+      "gumbel_tau": 0.9767,
+      "image_masks_loss": 0.0338,
+      "llm_learning_rate": 1.7235494880546078e-06,
+      "loss": 1.3973,
+      "output_text_mask_loss": 0.4395,
+      "predictor_learning_rate": 0.00017235494880546076,
+      "step": 101
+    },
+    {
+      "epoch": 0.01,
+      "gumbel_tau": 0.9765,
+      "image_masks_loss": 0.1099,
+      "llm_learning_rate": 1.74061433447099e-06,
+      "loss": 1.5314,
+      "output_text_mask_loss": 1.3611,
+      "predictor_learning_rate": 0.000174061433447099,
+      "step": 102
+    },
+    {
+      "epoch": 0.01,
+      "gumbel_tau": 0.9762,
+      "image_masks_loss": 0.0418,
+      "llm_learning_rate": 1.7576791808873722e-06,
+      "loss": 1.359,
+      "output_text_mask_loss": 0.7019,
+      "predictor_learning_rate": 0.0001757679180887372,
+      "step": 103
+    },
+    {
+      "epoch": 0.01,
+      "gumbel_tau": 0.976,
+      "image_masks_loss": 0.106,
+      "llm_learning_rate": 1.7747440273037543e-06,
+      "loss": 1.4042,
+      "output_text_mask_loss": 0.6287,
+      "predictor_learning_rate": 0.00017747440273037544,
+      "step": 104
+    },
+    {
+      "epoch": 0.01,
+      "gumbel_tau": 0.9758,
+      "image_masks_loss": 0.0254,
+      "llm_learning_rate": 1.7918088737201367e-06,
+      "loss": 1.3158,
+      "output_text_mask_loss": 0.4028,
+      "predictor_learning_rate": 0.00017918088737201365,
+      "step": 105
+    },
+    {
+      "epoch": 0.01,
+      "gumbel_tau": 0.9755,
+      "image_masks_loss": 0.0353,
+      "llm_learning_rate": 1.808873720136519e-06,
+      "loss": 1.0601,
+      "output_text_mask_loss": 0.0839,
+      "predictor_learning_rate": 0.0001808873720136519,
+      "step": 106
+    },
+    {
+      "epoch": 0.01,
+      "gumbel_tau": 0.9753,
+      "image_masks_loss": 0.0187,
+      "llm_learning_rate": 1.825938566552901e-06,
+      "loss": 1.1416,
+      "output_text_mask_loss": 0.5981,
+      "predictor_learning_rate": 0.0001825938566552901,
+      "step": 107
+    },
+    {
+      "epoch": 0.01,
+      "gumbel_tau": 0.9751,
+      "image_masks_loss": 0.0523,
+      "llm_learning_rate": 1.8430034129692834e-06,
+      "loss": 1.1433,
+      "output_text_mask_loss": 0.1389,
+      "predictor_learning_rate": 0.00018430034129692833,
+      "step": 108
+    },
+    {
+      "epoch": 0.01,
+      "gumbel_tau": 0.9748,
+      "image_masks_loss": 0.1343,
+      "llm_learning_rate": 1.8600682593856657e-06,
+      "loss": 1.452,
+      "output_text_mask_loss": 0.3281,
+      "predictor_learning_rate": 0.00018600682593856654,
+      "step": 109
+    },
+    {
+      "epoch": 0.01,
+      "gumbel_tau": 0.9746,
+      "image_masks_loss": 0.0326,
+      "llm_learning_rate": 1.8771331058020478e-06,
+      "loss": 1.2239,
+      "output_text_mask_loss": 0.3204,
+      "predictor_learning_rate": 0.00018771331058020478,
+      "step": 110
+    },
+    {
+      "epoch": 0.01,
+      "gumbel_tau": 0.9744,
+      "image_masks_loss": 0.0237,
+      "llm_learning_rate": 1.89419795221843e-06,
+      "loss": 1.2525,
+      "output_text_mask_loss": 0.1953,
+      "predictor_learning_rate": 0.000189419795221843,
+      "step": 111
+    },
+    {
+      "epoch": 0.01,
+      "gumbel_tau": 0.9742,
+      "image_masks_loss": 0.0465,
+      "llm_learning_rate": 1.9112627986348124e-06,
+      "loss": 1.6726,
+      "output_text_mask_loss": 1.0498,
+      "predictor_learning_rate": 0.00019112627986348125,
+      "step": 112
+    },
+    {
+      "epoch": 0.01,
+      "gumbel_tau": 0.9739,
+      "image_masks_loss": 0.0483,
+      "llm_learning_rate": 1.9283276450511947e-06,
+      "loss": 1.3543,
+      "output_text_mask_loss": 0.2899,
+      "predictor_learning_rate": 0.00019283276450511946,
+      "step": 113
+    },
+    {
+      "epoch": 0.01,
+      "gumbel_tau": 0.9737,
+      "image_masks_loss": 0.042,
+      "llm_learning_rate": 1.945392491467577e-06,
+      "loss": 1.6207,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0001945392491467577,
+      "step": 114
+    },
+    {
+      "epoch": 0.01,
+      "gumbel_tau": 0.9735,
+      "image_masks_loss": 0.0938,
+      "llm_learning_rate": 1.9624573378839593e-06,
+      "loss": 2.0347,
+      "output_text_mask_loss": 0.2472,
+      "predictor_learning_rate": 0.0001962457337883959,
+      "step": 115
+    },
+    {
+      "epoch": 0.01,
+      "gumbel_tau": 0.9732,
+      "image_masks_loss": 0.0467,
+      "llm_learning_rate": 1.9795221843003416e-06,
+      "loss": 1.7988,
+      "output_text_mask_loss": 0.8789,
+      "predictor_learning_rate": 0.00019795221843003414,
+      "step": 116
+    },
+    {
+      "epoch": 0.01,
+      "gumbel_tau": 0.973,
+      "image_masks_loss": 0.0355,
+      "llm_learning_rate": 1.9965870307167235e-06,
+      "loss": 1.8511,
+      "output_text_mask_loss": 1.0132,
+      "predictor_learning_rate": 0.00019965870307167235,
+      "step": 117
+    },
+    {
+      "epoch": 0.01,
+      "gumbel_tau": 0.9728,
+      "image_masks_loss": 0.0334,
+      "llm_learning_rate": 2.013651877133106e-06,
+      "loss": 1.6861,
+      "output_text_mask_loss": 0.8057,
+      "predictor_learning_rate": 0.0002013651877133106,
+      "step": 118
+    },
+    {
+      "epoch": 0.01,
+      "gumbel_tau": 0.9725,
+      "image_masks_loss": 0.0843,
+      "llm_learning_rate": 2.030716723549488e-06,
+      "loss": 2.4694,
+      "output_text_mask_loss": 1.2451,
+      "predictor_learning_rate": 0.0002030716723549488,
+      "step": 119
+    },
+    {
+      "epoch": 0.01,
+      "gumbel_tau": 0.9723,
+      "image_masks_loss": 0.0755,
+      "llm_learning_rate": 2.0477815699658705e-06,
+      "loss": 2.9211,
+      "output_text_mask_loss": 1.6357,
+      "predictor_learning_rate": 0.00020477815699658703,
+      "step": 120
+    },
+    {
+      "epoch": 0.01,
+      "gumbel_tau": 0.9721,
+      "image_masks_loss": 0.05,
+      "llm_learning_rate": 2.0648464163822528e-06,
+      "loss": 2.1128,
+      "output_text_mask_loss": 1.5442,
+      "predictor_learning_rate": 0.00020648464163822524,
+      "step": 121
+    },
+    {
+      "epoch": 0.01,
+      "gumbel_tau": 0.9719,
+      "image_masks_loss": 0.0629,
+      "llm_learning_rate": 2.0819112627986347e-06,
+      "loss": 1.9704,
+      "output_text_mask_loss": 0.8789,
+      "predictor_learning_rate": 0.00020819112627986348,
+      "step": 122
+    },
+    {
+      "epoch": 0.01,
+      "gumbel_tau": 0.9716,
+      "image_masks_loss": 0.0782,
+      "llm_learning_rate": 2.098976109215017e-06,
+      "loss": 2.2882,
+      "output_text_mask_loss": 0.4761,
+      "predictor_learning_rate": 0.0002098976109215017,
+      "step": 123
+    },
+    {
+      "epoch": 0.01,
+      "gumbel_tau": 0.9714,
+      "image_masks_loss": 0.0271,
+      "llm_learning_rate": 2.1160409556313997e-06,
+      "loss": 2.6233,
+      "output_text_mask_loss": 2.4536,
+      "predictor_learning_rate": 0.00021160409556313995,
+      "step": 124
+    },
+    {
+      "epoch": 0.01,
+      "gumbel_tau": 0.9712,
+      "image_masks_loss": 0.0675,
+      "llm_learning_rate": 2.133105802047782e-06,
+      "loss": 1.0958,
+      "output_text_mask_loss": 0.4761,
+      "predictor_learning_rate": 0.00021331058020477816,
+      "step": 125
+    },
+    {
+      "epoch": 0.01,
+      "gumbel_tau": 0.9709,
+      "image_masks_loss": 0.0429,
+      "llm_learning_rate": 2.150170648464164e-06,
+      "loss": 2.0029,
+      "output_text_mask_loss": 0.5371,
+      "predictor_learning_rate": 0.0002150170648464164,
+      "step": 126
+    },
+    {
+      "epoch": 0.01,
+      "gumbel_tau": 0.9707,
+      "image_masks_loss": 0.0477,
+      "llm_learning_rate": 2.1672354948805462e-06,
+      "loss": 1.3788,
+      "output_text_mask_loss": 1.2024,
+      "predictor_learning_rate": 0.0002167235494880546,
+      "step": 127
+    },
+    {
+      "epoch": 0.01,
+      "gumbel_tau": 0.9705,
+      "image_masks_loss": 0.0427,
+      "llm_learning_rate": 2.1843003412969285e-06,
+      "loss": 1.4163,
+      "output_text_mask_loss": 0.2747,
+      "predictor_learning_rate": 0.00021843003412969284,
+      "step": 128
+    },
+    {
+      "epoch": 0.01,
+      "gumbel_tau": 0.9703,
+      "image_masks_loss": 0.0212,
+      "llm_learning_rate": 2.201365187713311e-06,
+      "loss": 1.5372,
+      "output_text_mask_loss": 1.0193,
+      "predictor_learning_rate": 0.00022013651877133105,
+      "step": 129
+    },
+    {
+      "epoch": 0.01,
+      "gumbel_tau": 0.97,
+      "image_masks_loss": 0.0343,
+      "llm_learning_rate": 2.218430034129693e-06,
+      "loss": 1.9073,
+      "output_text_mask_loss": 1.1597,
+      "predictor_learning_rate": 0.00022184300341296929,
+      "step": 130
+    },
+    {
+      "epoch": 0.01,
+      "gumbel_tau": 0.9698,
+      "image_masks_loss": 0.0473,
+      "llm_learning_rate": 2.235494880546075e-06,
+      "loss": 1.1556,
+      "output_text_mask_loss": 0.2472,
+      "predictor_learning_rate": 0.00022354948805460752,
+      "step": 131
+    },
+    {
+      "epoch": 0.01,
+      "gumbel_tau": 0.9696,
+      "image_masks_loss": 0.0437,
+      "llm_learning_rate": 2.2525597269624573e-06,
+      "loss": 1.3231,
+      "output_text_mask_loss": 0.6805,
+      "predictor_learning_rate": 0.00022525597269624573,
+      "step": 132
+    },
+    {
+      "epoch": 0.01,
+      "gumbel_tau": 0.9693,
+      "image_masks_loss": 0.045,
+      "llm_learning_rate": 2.2696245733788397e-06,
+      "loss": 1.2525,
+      "output_text_mask_loss": 0.1648,
+      "predictor_learning_rate": 0.00022696245733788397,
+      "step": 133
+    },
+    {
+      "epoch": 0.01,
+      "gumbel_tau": 0.9691,
+      "image_masks_loss": 0.0675,
+      "llm_learning_rate": 2.286689419795222e-06,
+      "loss": 0.9701,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00022866894197952218,
+      "step": 134
+    },
+    {
+      "epoch": 0.01,
+      "gumbel_tau": 0.9689,
+      "image_masks_loss": 0.0843,
+      "llm_learning_rate": 2.3037542662116043e-06,
+      "loss": 1.0154,
+      "output_text_mask_loss": 0.106,
+      "predictor_learning_rate": 0.00023037542662116044,
+      "step": 135
+    },
+    {
+      "epoch": 0.01,
+      "gumbel_tau": 0.9687,
+      "image_masks_loss": 0.0763,
+      "llm_learning_rate": 2.3208191126279866e-06,
+      "loss": 1.0776,
+      "output_text_mask_loss": 0.0238,
+      "predictor_learning_rate": 0.00023208191126279865,
+      "step": 136
+    },
+    {
+      "epoch": 0.01,
+      "gumbel_tau": 0.9684,
+      "image_masks_loss": 0.0675,
+      "llm_learning_rate": 2.337883959044369e-06,
+      "loss": 1.1668,
+      "output_text_mask_loss": 0.0084,
+      "predictor_learning_rate": 0.00023378839590443688,
+      "step": 137
+    },
+    {
+      "epoch": 0.01,
+      "gumbel_tau": 0.9682,
+      "image_masks_loss": 0.018,
+      "llm_learning_rate": 2.354948805460751e-06,
+      "loss": 1.2392,
+      "output_text_mask_loss": 0.5463,
+      "predictor_learning_rate": 0.0002354948805460751,
+      "step": 138
+    },
+    {
+      "epoch": 0.01,
+      "gumbel_tau": 0.968,
+      "image_masks_loss": 0.0303,
+      "llm_learning_rate": 2.3720136518771335e-06,
+      "loss": 1.1046,
+      "output_text_mask_loss": 0.0893,
+      "predictor_learning_rate": 0.00023720136518771333,
+      "step": 139
+    },
+    {
+      "epoch": 0.01,
+      "gumbel_tau": 0.9677,
+      "image_masks_loss": 0.0452,
+      "llm_learning_rate": 2.3890784982935154e-06,
+      "loss": 1.1573,
+      "output_text_mask_loss": 0.0858,
+      "predictor_learning_rate": 0.00023890784982935154,
+      "step": 140
+    },
+    {
+      "epoch": 0.01,
+      "gumbel_tau": 0.9675,
+      "image_masks_loss": 0.0648,
+      "llm_learning_rate": 2.4061433447098977e-06,
+      "loss": 1.0259,
+      "output_text_mask_loss": 0.1602,
+      "predictor_learning_rate": 0.00024061433447098977,
+      "step": 141
+    },
+    {
+      "epoch": 0.01,
+      "gumbel_tau": 0.9673,
+      "image_masks_loss": 0.0186,
+      "llm_learning_rate": 2.42320819112628e-06,
+      "loss": 1.0321,
+      "output_text_mask_loss": 0.0374,
+      "predictor_learning_rate": 0.00024232081911262798,
+      "step": 142
+    },
+    {
+      "epoch": 0.01,
+      "gumbel_tau": 0.9671,
+      "image_masks_loss": 0.066,
+      "llm_learning_rate": 2.4402730375426623e-06,
+      "loss": 1.2075,
+      "output_text_mask_loss": 0.2289,
+      "predictor_learning_rate": 0.00024402730375426622,
+      "step": 143
+    },
+    {
+      "epoch": 0.01,
+      "gumbel_tau": 0.9668,
+      "image_masks_loss": 0.0702,
+      "llm_learning_rate": 2.4573378839590446e-06,
+      "loss": 1.0752,
+      "output_text_mask_loss": 0.0618,
+      "predictor_learning_rate": 0.00024573378839590443,
+      "step": 144
+    },
+    {
+      "epoch": 0.01,
+      "gumbel_tau": 0.9666,
+      "image_masks_loss": 0.0414,
+      "llm_learning_rate": 2.4744027303754265e-06,
+      "loss": 1.6189,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00024744027303754264,
+      "step": 145
+    },
+    {
+      "epoch": 0.01,
+      "gumbel_tau": 0.9664,
+      "image_masks_loss": 0.0222,
+      "llm_learning_rate": 2.491467576791809e-06,
+      "loss": 1.2671,
+      "output_text_mask_loss": 0.3571,
+      "predictor_learning_rate": 0.0002491467576791809,
+      "step": 146
+    },
+    {
+      "epoch": 0.02,
+      "gumbel_tau": 0.9661,
+      "image_masks_loss": 0.053,
+      "llm_learning_rate": 2.5085324232081916e-06,
+      "loss": 1.3658,
+      "output_text_mask_loss": 0.5005,
+      "predictor_learning_rate": 0.0002508532423208191,
+      "step": 147
+    },
+    {
+      "epoch": 0.02,
+      "gumbel_tau": 0.9659,
+      "image_masks_loss": 0.0307,
+      "llm_learning_rate": 2.5255972696245735e-06,
+      "loss": 1.364,
+      "output_text_mask_loss": 1.0071,
+      "predictor_learning_rate": 0.0002525597269624573,
+      "step": 148
+    },
+    {
+      "epoch": 0.02,
+      "gumbel_tau": 0.9657,
+      "image_masks_loss": 0.0203,
+      "llm_learning_rate": 2.5426621160409558e-06,
+      "loss": 1.2633,
+      "output_text_mask_loss": 0.1007,
+      "predictor_learning_rate": 0.0002542662116040956,
+      "step": 149
+    },
+    {
+      "epoch": 0.02,
+      "gumbel_tau": 0.9655,
+      "image_masks_loss": 0.0172,
+      "llm_learning_rate": 2.559726962457338e-06,
+      "loss": 1.3092,
+      "output_text_mask_loss": 0.5981,
+      "predictor_learning_rate": 0.0002559726962457338,
+      "step": 150
+    },
+    {
+      "epoch": 0.02,
+      "gumbel_tau": 0.9652,
+      "image_masks_loss": 0.0387,
+      "llm_learning_rate": 2.5767918088737204e-06,
+      "loss": 1.0592,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 0.000257679180887372,
+      "step": 151
+    },
+    {
+      "epoch": 0.02,
+      "gumbel_tau": 0.965,
+      "image_masks_loss": 0.0357,
+      "llm_learning_rate": 2.5938566552901023e-06,
+      "loss": 1.0869,
+      "output_text_mask_loss": 0.3098,
+      "predictor_learning_rate": 0.0002593856655290102,
+      "step": 152
+    },
+    {
+      "epoch": 0.02,
+      "gumbel_tau": 0.9648,
+      "image_masks_loss": 0.0372,
+      "llm_learning_rate": 2.610921501706485e-06,
+      "loss": 1.1765,
+      "output_text_mask_loss": 0.0374,
+      "predictor_learning_rate": 0.00026109215017064847,
+      "step": 153
+    },
+    {
+      "epoch": 0.02,
+      "gumbel_tau": 0.9645,
+      "image_masks_loss": 0.0732,
+      "llm_learning_rate": 2.6279863481228673e-06,
+      "loss": 1.1275,
+      "output_text_mask_loss": 0.0687,
+      "predictor_learning_rate": 0.00026279863481228674,
+      "step": 154
+    },
+    {
+      "epoch": 0.02,
+      "gumbel_tau": 0.9643,
+      "image_masks_loss": 0.0572,
+      "llm_learning_rate": 2.645051194539249e-06,
+      "loss": 0.8769,
+      "output_text_mask_loss": 0.0549,
+      "predictor_learning_rate": 0.0002645051194539249,
+      "step": 155
+    },
+    {
+      "epoch": 0.02,
+      "gumbel_tau": 0.9641,
+      "image_masks_loss": 0.0168,
+      "llm_learning_rate": 2.662116040955632e-06,
+      "loss": 0.7999,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00026621160409556315,
+      "step": 156
+    },
+    {
+      "epoch": 0.02,
+      "gumbel_tau": 0.9639,
+      "image_masks_loss": 0.0813,
+      "llm_learning_rate": 2.679180887372014e-06,
+      "loss": 1.2546,
+      "output_text_mask_loss": 0.0973,
+      "predictor_learning_rate": 0.00026791808873720136,
+      "step": 157
+    },
+    {
+      "epoch": 0.02,
+      "gumbel_tau": 0.9636,
+      "image_masks_loss": 0.0113,
+      "llm_learning_rate": 2.696245733788396e-06,
+      "loss": 1.0249,
+      "output_text_mask_loss": 0.1541,
+      "predictor_learning_rate": 0.0002696245733788396,
+      "step": 158
+    },
+    {
+      "epoch": 0.02,
+      "gumbel_tau": 0.9634,
+      "image_masks_loss": 0.0206,
+      "llm_learning_rate": 2.713310580204778e-06,
+      "loss": 1.0851,
+      "output_text_mask_loss": 0.1244,
+      "predictor_learning_rate": 0.0002713310580204778,
+      "step": 159
+    },
+    {
+      "epoch": 0.02,
+      "gumbel_tau": 0.9632,
+      "image_masks_loss": 0.028,
+      "llm_learning_rate": 2.7303754266211608e-06,
+      "loss": 1.1269,
+      "output_text_mask_loss": 0.0299,
+      "predictor_learning_rate": 0.00027303754266211604,
+      "step": 160
+    },
+    {
+      "epoch": 0.02,
+      "gumbel_tau": 0.963,
+      "image_masks_loss": 0.0256,
+      "llm_learning_rate": 2.7474402730375426e-06,
+      "loss": 0.9415,
+      "output_text_mask_loss": 0.2472,
+      "predictor_learning_rate": 0.00027474402730375425,
+      "step": 161
+    },
+    {
+      "epoch": 0.02,
+      "gumbel_tau": 0.9627,
+      "image_masks_loss": 0.0626,
+      "llm_learning_rate": 2.7645051194539254e-06,
+      "loss": 1.0479,
+      "output_text_mask_loss": 0.0275,
+      "predictor_learning_rate": 0.0002764505119453925,
+      "step": 162
+    },
+    {
+      "epoch": 0.02,
+      "gumbel_tau": 0.9625,
+      "image_masks_loss": 0.0454,
+      "llm_learning_rate": 2.7815699658703073e-06,
+      "loss": 1.2268,
+      "output_text_mask_loss": 0.1762,
+      "predictor_learning_rate": 0.00027815699658703067,
+      "step": 163
+    },
+    {
+      "epoch": 0.02,
+      "gumbel_tau": 0.9623,
+      "image_masks_loss": 0.0404,
+      "llm_learning_rate": 2.7986348122866896e-06,
+      "loss": 1.1014,
+      "output_text_mask_loss": 0.6958,
+      "predictor_learning_rate": 0.00027986348122866893,
+      "step": 164
+    },
+    {
+      "epoch": 0.02,
+      "gumbel_tau": 0.962,
+      "image_masks_loss": 0.036,
+      "llm_learning_rate": 2.8156996587030723e-06,
+      "loss": 1.2246,
+      "output_text_mask_loss": 0.4517,
+      "predictor_learning_rate": 0.0002815699658703072,
+      "step": 165
+    },
+    {
+      "epoch": 0.02,
+      "gumbel_tau": 0.9618,
+      "image_masks_loss": 0.1007,
+      "llm_learning_rate": 2.832764505119454e-06,
+      "loss": 0.9837,
+      "output_text_mask_loss": 0.209,
+      "predictor_learning_rate": 0.0002832764505119454,
+      "step": 166
+    },
+    {
+      "epoch": 0.02,
+      "gumbel_tau": 0.9616,
+      "image_masks_loss": 0.0116,
+      "llm_learning_rate": 2.8498293515358365e-06,
+      "loss": 0.9578,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00028498293515358367,
+      "step": 167
+    },
+    {
+      "epoch": 0.02,
+      "gumbel_tau": 0.9614,
+      "image_masks_loss": 0.0343,
+      "llm_learning_rate": 2.8668941979522184e-06,
+      "loss": 1.051,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 0.0002866894197952218,
+      "step": 168
+    },
+    {
+      "epoch": 0.02,
+      "gumbel_tau": 0.9611,
+      "image_masks_loss": 0.0828,
+      "llm_learning_rate": 2.883959044368601e-06,
+      "loss": 1.0464,
+      "output_text_mask_loss": 0.1099,
+      "predictor_learning_rate": 0.0002883959044368601,
+      "step": 169
+    },
+    {
+      "epoch": 0.02,
+      "gumbel_tau": 0.9609,
+      "image_masks_loss": 0.1648,
+      "llm_learning_rate": 2.901023890784983e-06,
+      "loss": 1.1888,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0002901023890784983,
+      "step": 170
+    },
+    {
+      "epoch": 0.02,
+      "gumbel_tau": 0.9607,
+      "image_masks_loss": 0.0345,
+      "llm_learning_rate": 2.9180887372013657e-06,
+      "loss": 0.9621,
+      "output_text_mask_loss": 0.1099,
+      "predictor_learning_rate": 0.00029180887372013656,
+      "step": 171
+    },
+    {
+      "epoch": 0.02,
+      "gumbel_tau": 0.9605,
+      "image_masks_loss": 0.0471,
+      "llm_learning_rate": 2.9351535836177476e-06,
+      "loss": 1.0936,
+      "output_text_mask_loss": 0.0771,
+      "predictor_learning_rate": 0.0002935153583617747,
+      "step": 172
+    },
+    {
+      "epoch": 0.02,
+      "gumbel_tau": 0.9602,
+      "image_masks_loss": 0.0771,
+      "llm_learning_rate": 2.95221843003413e-06,
+      "loss": 1.0369,
+      "output_text_mask_loss": 0.0116,
+      "predictor_learning_rate": 0.000295221843003413,
+      "step": 173
+    },
+    {
+      "epoch": 0.02,
+      "gumbel_tau": 0.96,
+      "image_masks_loss": 0.1022,
+      "llm_learning_rate": 2.969283276450512e-06,
+      "loss": 0.9353,
+      "output_text_mask_loss": 0.0572,
+      "predictor_learning_rate": 0.0002969283276450512,
+      "step": 174
+    },
+    {
+      "epoch": 0.02,
+      "gumbel_tau": 0.9598,
+      "image_masks_loss": 0.0607,
+      "llm_learning_rate": 2.9863481228668946e-06,
+      "loss": 0.8633,
+      "output_text_mask_loss": 0.0208,
+      "predictor_learning_rate": 0.00029863481228668945,
+      "step": 175
+    },
+    {
+      "epoch": 0.02,
+      "gumbel_tau": 0.9596,
+      "image_masks_loss": 0.0324,
+      "llm_learning_rate": 3.003412969283277e-06,
+      "loss": 0.944,
+      "output_text_mask_loss": 0.0086,
+      "predictor_learning_rate": 0.00030034129692832766,
+      "step": 176
+    },
+    {
+      "epoch": 0.02,
+      "gumbel_tau": 0.9593,
+      "image_masks_loss": 0.0683,
+      "llm_learning_rate": 3.0204778156996588e-06,
+      "loss": 1.0827,
+      "output_text_mask_loss": 0.1968,
+      "predictor_learning_rate": 0.00030204778156996587,
+      "step": 177
+    },
+    {
+      "epoch": 0.02,
+      "gumbel_tau": 0.9591,
+      "image_masks_loss": 0.1175,
+      "llm_learning_rate": 3.0375426621160415e-06,
+      "loss": 0.9507,
+      "output_text_mask_loss": 0.0679,
+      "predictor_learning_rate": 0.00030375426621160413,
+      "step": 178
+    },
+    {
+      "epoch": 0.02,
+      "gumbel_tau": 0.9589,
+      "image_masks_loss": 0.0324,
+      "llm_learning_rate": 3.0546075085324234e-06,
+      "loss": 1.1019,
+      "output_text_mask_loss": 0.1038,
+      "predictor_learning_rate": 0.00030546075085324234,
+      "step": 179
+    },
+    {
+      "epoch": 0.02,
+      "gumbel_tau": 0.9586,
+      "image_masks_loss": 0.0507,
+      "llm_learning_rate": 3.0716723549488057e-06,
+      "loss": 0.8756,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00030716723549488055,
+      "step": 180
+    },
+    {
+      "epoch": 0.02,
+      "gumbel_tau": 0.9584,
+      "image_masks_loss": 0.0275,
+      "llm_learning_rate": 3.088737201365188e-06,
+      "loss": 0.8336,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 0.00030887372013651876,
+      "step": 181
+    },
+    {
+      "epoch": 0.02,
+      "gumbel_tau": 0.9582,
+      "image_masks_loss": 0.0576,
+      "llm_learning_rate": 3.1058020477815703e-06,
+      "loss": 1.1216,
+      "output_text_mask_loss": 0.119,
+      "predictor_learning_rate": 0.000310580204778157,
+      "step": 182
+    },
+    {
+      "epoch": 0.02,
+      "gumbel_tau": 0.958,
+      "image_masks_loss": 0.087,
+      "llm_learning_rate": 3.122866894197952e-06,
+      "loss": 0.8907,
+      "output_text_mask_loss": 0.164,
+      "predictor_learning_rate": 0.00031228668941979523,
+      "step": 183
+    },
+    {
+      "epoch": 0.02,
+      "gumbel_tau": 0.9577,
+      "image_masks_loss": 0.0328,
+      "llm_learning_rate": 3.139931740614335e-06,
+      "loss": 0.9657,
+      "output_text_mask_loss": 0.1495,
+      "predictor_learning_rate": 0.00031399317406143344,
+      "step": 184
+    },
+    {
+      "epoch": 0.02,
+      "gumbel_tau": 0.9575,
+      "image_masks_loss": 0.0319,
+      "llm_learning_rate": 3.156996587030717e-06,
+      "loss": 1.1309,
+      "output_text_mask_loss": 0.0633,
+      "predictor_learning_rate": 0.00031569965870307165,
+      "step": 185
+    },
+    {
+      "epoch": 0.02,
+      "gumbel_tau": 0.9573,
+      "image_masks_loss": 0.0383,
+      "llm_learning_rate": 3.174061433447099e-06,
+      "loss": 0.9624,
+      "output_text_mask_loss": 0.145,
+      "predictor_learning_rate": 0.0003174061433447099,
+      "step": 186
+    },
+    {
+      "epoch": 0.02,
+      "gumbel_tau": 0.9571,
+      "image_masks_loss": 0.05,
+      "llm_learning_rate": 3.191126279863481e-06,
+      "loss": 0.9577,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 0.0003191126279863481,
+      "step": 187
+    },
+    {
+      "epoch": 0.02,
+      "gumbel_tau": 0.9568,
+      "image_masks_loss": 0.071,
+      "llm_learning_rate": 3.2081911262798638e-06,
+      "loss": 1.0396,
+      "output_text_mask_loss": 0.0437,
+      "predictor_learning_rate": 0.00032081911262798633,
+      "step": 188
+    },
+    {
+      "epoch": 0.02,
+      "gumbel_tau": 0.9566,
+      "image_masks_loss": 0.0507,
+      "llm_learning_rate": 3.225255972696246e-06,
+      "loss": 0.9046,
+      "output_text_mask_loss": 0.1984,
+      "predictor_learning_rate": 0.0003225255972696246,
+      "step": 189
+    },
+    {
+      "epoch": 0.02,
+      "gumbel_tau": 0.9564,
+      "image_masks_loss": 0.023,
+      "llm_learning_rate": 3.2423208191126284e-06,
+      "loss": 0.9534,
+      "output_text_mask_loss": 0.037,
+      "predictor_learning_rate": 0.0003242320819112628,
+      "step": 190
+    },
+    {
+      "epoch": 0.02,
+      "gumbel_tau": 0.9562,
+      "image_masks_loss": 0.0374,
+      "llm_learning_rate": 3.2593856655290107e-06,
+      "loss": 1.0294,
+      "output_text_mask_loss": 0.0828,
+      "predictor_learning_rate": 0.00032593856655290107,
+      "step": 191
+    },
+    {
+      "epoch": 0.02,
+      "gumbel_tau": 0.9559,
+      "image_masks_loss": 0.0313,
+      "llm_learning_rate": 3.2764505119453926e-06,
+      "loss": 0.9395,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 0.0003276450511945392,
+      "step": 192
+    },
+    {
+      "epoch": 0.02,
+      "gumbel_tau": 0.9557,
+      "image_masks_loss": 0.0637,
+      "llm_learning_rate": 3.2935153583617753e-06,
+      "loss": 1.0068,
+      "output_text_mask_loss": 0.0671,
+      "predictor_learning_rate": 0.0003293515358361775,
+      "step": 193
+    },
+    {
+      "epoch": 0.02,
+      "gumbel_tau": 0.9555,
+      "image_masks_loss": 0.0088,
+      "llm_learning_rate": 3.310580204778157e-06,
+      "loss": 0.9319,
+      "output_text_mask_loss": 0.1747,
+      "predictor_learning_rate": 0.0003310580204778157,
+      "step": 194
+    },
+    {
+      "epoch": 0.02,
+      "gumbel_tau": 0.9553,
+      "image_masks_loss": 0.0721,
+      "llm_learning_rate": 3.3276450511945395e-06,
+      "loss": 0.9497,
+      "output_text_mask_loss": 0.0488,
+      "predictor_learning_rate": 0.00033276450511945396,
+      "step": 195
+    },
+    {
+      "epoch": 0.02,
+      "gumbel_tau": 0.955,
+      "image_masks_loss": 0.0195,
+      "llm_learning_rate": 3.3447098976109214e-06,
+      "loss": 1.0612,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 0.0003344709897610921,
+      "step": 196
+    },
+    {
+      "epoch": 0.02,
+      "gumbel_tau": 0.9548,
+      "image_masks_loss": 0.1007,
+      "llm_learning_rate": 3.361774744027304e-06,
+      "loss": 1.0582,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.0003361774744027304,
+      "step": 197
+    },
+    {
+      "epoch": 0.02,
+      "gumbel_tau": 0.9546,
+      "image_masks_loss": 0.0299,
+      "llm_learning_rate": 3.378839590443686e-06,
+      "loss": 0.8801,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 0.0003378839590443686,
+      "step": 198
+    },
+    {
+      "epoch": 0.02,
+      "gumbel_tau": 0.9544,
+      "image_masks_loss": 0.014,
+      "llm_learning_rate": 3.3959044368600687e-06,
+      "loss": 1.0591,
+      "output_text_mask_loss": 0.2502,
+      "predictor_learning_rate": 0.00033959044368600685,
+      "step": 199
+    },
+    {
+      "epoch": 0.02,
+      "gumbel_tau": 0.9541,
+      "image_masks_loss": 0.0326,
+      "llm_learning_rate": 3.412969283276451e-06,
+      "loss": 0.9768,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0003412969283276451,
+      "step": 200
+    },
+    {
+      "epoch": 0.02,
+      "gumbel_tau": 0.9539,
+      "image_masks_loss": 0.0315,
+      "llm_learning_rate": 3.430034129692833e-06,
+      "loss": 1.1212,
+      "output_text_mask_loss": 0.2884,
+      "predictor_learning_rate": 0.00034300341296928326,
+      "step": 201
+    },
+    {
+      "epoch": 0.02,
+      "gumbel_tau": 0.9537,
+      "image_masks_loss": 0.0322,
+      "llm_learning_rate": 3.4470989761092157e-06,
+      "loss": 0.9482,
+      "output_text_mask_loss": 0.1007,
+      "predictor_learning_rate": 0.0003447098976109215,
+      "step": 202
+    },
+    {
+      "epoch": 0.02,
+      "gumbel_tau": 0.9535,
+      "image_masks_loss": 0.0269,
+      "llm_learning_rate": 3.4641638225255976e-06,
+      "loss": 0.8352,
+      "output_text_mask_loss": 0.0618,
+      "predictor_learning_rate": 0.00034641638225255974,
+      "step": 203
+    },
+    {
+      "epoch": 0.02,
+      "gumbel_tau": 0.9532,
+      "image_masks_loss": 0.0315,
+      "llm_learning_rate": 3.48122866894198e-06,
+      "loss": 0.8881,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.000348122866894198,
+      "step": 204
+    },
+    {
+      "epoch": 0.02,
+      "gumbel_tau": 0.953,
+      "image_masks_loss": 0.1091,
+      "llm_learning_rate": 3.4982935153583618e-06,
+      "loss": 0.9645,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00034982935153583615,
+      "step": 205
+    },
+    {
+      "epoch": 0.02,
+      "gumbel_tau": 0.9528,
+      "image_masks_loss": 0.0278,
+      "llm_learning_rate": 3.5153583617747445e-06,
+      "loss": 0.9791,
+      "output_text_mask_loss": 0.3815,
+      "predictor_learning_rate": 0.0003515358361774744,
+      "step": 206
+    },
+    {
+      "epoch": 0.02,
+      "gumbel_tau": 0.9526,
+      "image_masks_loss": 0.033,
+      "llm_learning_rate": 3.5324232081911264e-06,
+      "loss": 1.0559,
+      "output_text_mask_loss": 0.1038,
+      "predictor_learning_rate": 0.0003532423208191126,
+      "step": 207
+    },
+    {
+      "epoch": 0.02,
+      "gumbel_tau": 0.9523,
+      "image_masks_loss": 0.0462,
+      "llm_learning_rate": 3.5494880546075087e-06,
+      "loss": 0.8559,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0003549488054607509,
+      "step": 208
+    },
+    {
+      "epoch": 0.02,
+      "gumbel_tau": 0.9521,
+      "image_masks_loss": 0.033,
+      "llm_learning_rate": 3.566552901023891e-06,
+      "loss": 0.8854,
+      "output_text_mask_loss": 0.0313,
+      "predictor_learning_rate": 0.00035665529010238904,
+      "step": 209
+    },
+    {
+      "epoch": 0.02,
+      "gumbel_tau": 0.9519,
+      "image_masks_loss": 0.0793,
+      "llm_learning_rate": 3.5836177474402733e-06,
+      "loss": 0.9323,
+      "output_text_mask_loss": 0.0916,
+      "predictor_learning_rate": 0.0003583617747440273,
+      "step": 210
+    },
+    {
+      "epoch": 0.02,
+      "gumbel_tau": 0.9517,
+      "image_masks_loss": 0.0113,
+      "llm_learning_rate": 3.600682593856656e-06,
+      "loss": 0.9161,
+      "output_text_mask_loss": 0.0549,
+      "predictor_learning_rate": 0.00036006825938566557,
+      "step": 211
+    },
+    {
+      "epoch": 0.02,
+      "gumbel_tau": 0.9514,
+      "image_masks_loss": 0.0462,
+      "llm_learning_rate": 3.617747440273038e-06,
+      "loss": 0.9211,
+      "output_text_mask_loss": 0.1007,
+      "predictor_learning_rate": 0.0003617747440273038,
+      "step": 212
+    },
+    {
+      "epoch": 0.02,
+      "gumbel_tau": 0.9512,
+      "image_masks_loss": 0.0252,
+      "llm_learning_rate": 3.6348122866894202e-06,
+      "loss": 0.9872,
+      "output_text_mask_loss": 0.0763,
+      "predictor_learning_rate": 0.000363481228668942,
+      "step": 213
+    },
+    {
+      "epoch": 0.02,
+      "gumbel_tau": 0.951,
+      "image_masks_loss": 0.0104,
+      "llm_learning_rate": 3.651877133105802e-06,
+      "loss": 0.8372,
+      "output_text_mask_loss": 0.0076,
+      "predictor_learning_rate": 0.0003651877133105802,
+      "step": 214
+    },
+    {
+      "epoch": 0.02,
+      "gumbel_tau": 0.9508,
+      "image_masks_loss": 0.0465,
+      "llm_learning_rate": 3.668941979522185e-06,
+      "loss": 0.9219,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.00036689419795221846,
+      "step": 215
+    },
+    {
+      "epoch": 0.02,
+      "gumbel_tau": 0.9505,
+      "image_masks_loss": 0.0534,
+      "llm_learning_rate": 3.6860068259385667e-06,
+      "loss": 0.8352,
+      "output_text_mask_loss": 0.2075,
+      "predictor_learning_rate": 0.00036860068259385667,
+      "step": 216
+    },
+    {
+      "epoch": 0.02,
+      "gumbel_tau": 0.9503,
+      "image_masks_loss": 0.0381,
+      "llm_learning_rate": 3.703071672354949e-06,
+      "loss": 0.7811,
+      "output_text_mask_loss": 0.0429,
+      "predictor_learning_rate": 0.0003703071672354949,
+      "step": 217
+    },
+    {
+      "epoch": 0.02,
+      "gumbel_tau": 0.9501,
+      "image_masks_loss": 0.0483,
+      "llm_learning_rate": 3.7201365187713314e-06,
+      "loss": 0.8824,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0003720136518771331,
+      "step": 218
+    },
+    {
+      "epoch": 0.02,
+      "gumbel_tau": 0.9499,
+      "image_masks_loss": 0.0196,
+      "llm_learning_rate": 3.7372013651877137e-06,
+      "loss": 0.9776,
+      "output_text_mask_loss": 0.1785,
+      "predictor_learning_rate": 0.00037372013651877135,
+      "step": 219
+    },
+    {
+      "epoch": 0.02,
+      "gumbel_tau": 0.9496,
+      "image_masks_loss": 0.0378,
+      "llm_learning_rate": 3.7542662116040956e-06,
+      "loss": 0.9586,
+      "output_text_mask_loss": 0.132,
+      "predictor_learning_rate": 0.00037542662116040956,
+      "step": 220
+    },
+    {
+      "epoch": 0.02,
+      "gumbel_tau": 0.9494,
+      "image_masks_loss": 0.0368,
+      "llm_learning_rate": 3.7713310580204783e-06,
+      "loss": 0.8727,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.00037713310580204777,
+      "step": 221
+    },
+    {
+      "epoch": 0.02,
+      "gumbel_tau": 0.9492,
+      "image_masks_loss": 0.018,
+      "llm_learning_rate": 3.78839590443686e-06,
+      "loss": 0.8275,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.000378839590443686,
+      "step": 222
+    },
+    {
+      "epoch": 0.02,
+      "gumbel_tau": 0.949,
+      "image_masks_loss": 0.0456,
+      "llm_learning_rate": 3.8054607508532425e-06,
+      "loss": 0.9244,
+      "output_text_mask_loss": 0.0275,
+      "predictor_learning_rate": 0.00038054607508532424,
+      "step": 223
+    },
+    {
+      "epoch": 0.02,
+      "gumbel_tau": 0.9487,
+      "image_masks_loss": 0.0347,
+      "llm_learning_rate": 3.822525597269625e-06,
+      "loss": 0.8746,
+      "output_text_mask_loss": 0.0923,
+      "predictor_learning_rate": 0.0003822525597269625,
+      "step": 224
+    },
+    {
+      "epoch": 0.02,
+      "gumbel_tau": 0.9485,
+      "image_masks_loss": 0.0641,
+      "llm_learning_rate": 3.839590443686007e-06,
+      "loss": 0.8642,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00038395904436860066,
+      "step": 225
+    },
+    {
+      "epoch": 0.02,
+      "gumbel_tau": 0.9483,
+      "image_masks_loss": 0.1038,
+      "llm_learning_rate": 3.8566552901023894e-06,
+      "loss": 1.0432,
+      "output_text_mask_loss": 0.1816,
+      "predictor_learning_rate": 0.0003856655290102389,
+      "step": 226
+    },
+    {
+      "epoch": 0.02,
+      "gumbel_tau": 0.9481,
+      "image_masks_loss": 0.0148,
+      "llm_learning_rate": 3.873720136518771e-06,
+      "loss": 0.9725,
+      "output_text_mask_loss": 0.0484,
+      "predictor_learning_rate": 0.00038737201365187713,
+      "step": 227
+    },
+    {
+      "epoch": 0.02,
+      "gumbel_tau": 0.9479,
+      "image_masks_loss": 0.0366,
+      "llm_learning_rate": 3.890784982935154e-06,
+      "loss": 0.8957,
+      "output_text_mask_loss": 0.0187,
+      "predictor_learning_rate": 0.0003890784982935154,
+      "step": 228
+    },
+    {
+      "epoch": 0.02,
+      "gumbel_tau": 0.9476,
+      "image_masks_loss": 0.0771,
+      "llm_learning_rate": 3.907849829351536e-06,
+      "loss": 0.9074,
+      "output_text_mask_loss": 0.1419,
+      "predictor_learning_rate": 0.0003907849829351536,
+      "step": 229
+    },
+    {
+      "epoch": 0.02,
+      "gumbel_tau": 0.9474,
+      "image_masks_loss": 0.0155,
+      "llm_learning_rate": 3.924914675767919e-06,
+      "loss": 0.9179,
+      "output_text_mask_loss": 0.0338,
+      "predictor_learning_rate": 0.0003924914675767918,
+      "step": 230
+    },
+    {
+      "epoch": 0.02,
+      "gumbel_tau": 0.9472,
+      "image_masks_loss": 0.0275,
+      "llm_learning_rate": 3.9419795221843006e-06,
+      "loss": 0.883,
+      "output_text_mask_loss": 0.0641,
+      "predictor_learning_rate": 0.00039419795221843,
+      "step": 231
+    },
+    {
+      "epoch": 0.02,
+      "gumbel_tau": 0.947,
+      "image_masks_loss": 0.0648,
+      "llm_learning_rate": 3.959044368600683e-06,
+      "loss": 0.9598,
+      "output_text_mask_loss": 0.0565,
+      "predictor_learning_rate": 0.0003959044368600683,
+      "step": 232
+    },
+    {
+      "epoch": 0.02,
+      "gumbel_tau": 0.9467,
+      "image_masks_loss": 0.0374,
+      "llm_learning_rate": 3.976109215017065e-06,
+      "loss": 0.856,
+      "output_text_mask_loss": 0.1305,
+      "predictor_learning_rate": 0.0003976109215017065,
+      "step": 233
+    },
+    {
+      "epoch": 0.02,
+      "gumbel_tau": 0.9465,
+      "image_masks_loss": 0.0519,
+      "llm_learning_rate": 3.993174061433447e-06,
+      "loss": 0.8509,
+      "output_text_mask_loss": 0.001,
+      "predictor_learning_rate": 0.0003993174061433447,
+      "step": 234
+    },
+    {
+      "epoch": 0.02,
+      "gumbel_tau": 0.9463,
+      "image_masks_loss": 0.0128,
+      "llm_learning_rate": 4.01023890784983e-06,
+      "loss": 0.7986,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.00040102389078498297,
+      "step": 235
+    },
+    {
+      "epoch": 0.02,
+      "gumbel_tau": 0.9461,
+      "image_masks_loss": 0.0782,
+      "llm_learning_rate": 4.027303754266212e-06,
+      "loss": 0.8143,
+      "output_text_mask_loss": 0.0198,
+      "predictor_learning_rate": 0.0004027303754266212,
+      "step": 236
+    },
+    {
+      "epoch": 0.02,
+      "gumbel_tau": 0.9458,
+      "image_masks_loss": 0.0195,
+      "llm_learning_rate": 4.044368600682594e-06,
+      "loss": 0.8543,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.00040443686006825944,
+      "step": 237
+    },
+    {
+      "epoch": 0.02,
+      "gumbel_tau": 0.9456,
+      "image_masks_loss": 0.0439,
+      "llm_learning_rate": 4.061433447098976e-06,
+      "loss": 0.899,
+      "output_text_mask_loss": 0.0221,
+      "predictor_learning_rate": 0.0004061433447098976,
+      "step": 238
+    },
+    {
+      "epoch": 0.02,
+      "gumbel_tau": 0.9454,
+      "image_masks_loss": 0.0839,
+      "llm_learning_rate": 4.078498293515359e-06,
+      "loss": 0.8814,
+      "output_text_mask_loss": 0.1007,
+      "predictor_learning_rate": 0.00040784982935153586,
+      "step": 239
+    },
+    {
+      "epoch": 0.02,
+      "gumbel_tau": 0.9452,
+      "image_masks_loss": 0.0166,
+      "llm_learning_rate": 4.095563139931741e-06,
+      "loss": 0.8872,
+      "output_text_mask_loss": 0.1205,
+      "predictor_learning_rate": 0.00040955631399317407,
+      "step": 240
+    },
+    {
+      "epoch": 0.02,
+      "gumbel_tau": 0.945,
+      "image_masks_loss": 0.0813,
+      "llm_learning_rate": 4.112627986348124e-06,
+      "loss": 0.9527,
+      "output_text_mask_loss": 0.1068,
+      "predictor_learning_rate": 0.00041126279863481233,
+      "step": 241
+    },
+    {
+      "epoch": 0.02,
+      "gumbel_tau": 0.9447,
+      "image_masks_loss": 0.0156,
+      "llm_learning_rate": 4.1296928327645055e-06,
+      "loss": 0.8074,
+      "output_text_mask_loss": 0.0065,
+      "predictor_learning_rate": 0.0004129692832764505,
+      "step": 242
+    },
+    {
+      "epoch": 0.02,
+      "gumbel_tau": 0.9445,
+      "image_masks_loss": 0.037,
+      "llm_learning_rate": 4.1467576791808874e-06,
+      "loss": 0.998,
+      "output_text_mask_loss": 0.1793,
+      "predictor_learning_rate": 0.00041467576791808875,
+      "step": 243
+    },
+    {
+      "epoch": 0.03,
+      "gumbel_tau": 0.9443,
+      "image_masks_loss": 0.0305,
+      "llm_learning_rate": 4.163822525597269e-06,
+      "loss": 0.8935,
+      "output_text_mask_loss": 0.0193,
+      "predictor_learning_rate": 0.00041638225255972696,
+      "step": 244
+    },
+    {
+      "epoch": 0.03,
+      "gumbel_tau": 0.9441,
+      "image_masks_loss": 0.0305,
+      "llm_learning_rate": 4.180887372013652e-06,
+      "loss": 0.8686,
+      "output_text_mask_loss": 0.0687,
+      "predictor_learning_rate": 0.0004180887372013652,
+      "step": 245
+    },
+    {
+      "epoch": 0.03,
+      "gumbel_tau": 0.9438,
+      "image_masks_loss": 0.042,
+      "llm_learning_rate": 4.197952218430034e-06,
+      "loss": 0.8741,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.0004197952218430034,
+      "step": 246
+    },
+    {
+      "epoch": 0.03,
+      "gumbel_tau": 0.9436,
+      "image_masks_loss": 0.0763,
+      "llm_learning_rate": 4.215017064846417e-06,
+      "loss": 0.9365,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 0.00042150170648464164,
+      "step": 247
+    },
+    {
+      "epoch": 0.03,
+      "gumbel_tau": 0.9434,
+      "image_masks_loss": 0.0351,
+      "llm_learning_rate": 4.232081911262799e-06,
+      "loss": 0.9097,
+      "output_text_mask_loss": 0.0275,
+      "predictor_learning_rate": 0.0004232081911262799,
+      "step": 248
+    },
+    {
+      "epoch": 0.03,
+      "gumbel_tau": 0.9432,
+      "image_masks_loss": 0.028,
+      "llm_learning_rate": 4.249146757679181e-06,
+      "loss": 0.8887,
+      "output_text_mask_loss": 0.0549,
+      "predictor_learning_rate": 0.0004249146757679181,
+      "step": 249
+    },
+    {
+      "epoch": 0.03,
+      "gumbel_tau": 0.9429,
+      "image_masks_loss": 0.0229,
+      "llm_learning_rate": 4.266211604095564e-06,
+      "loss": 0.9336,
+      "output_text_mask_loss": 0.01,
+      "predictor_learning_rate": 0.0004266211604095563,
+      "step": 250
+    },
+    {
+      "epoch": 0.03,
+      "gumbel_tau": 0.9427,
+      "image_masks_loss": 0.0595,
+      "llm_learning_rate": 4.283276450511946e-06,
+      "loss": 0.8548,
+      "output_text_mask_loss": 0.0488,
+      "predictor_learning_rate": 0.00042832764505119453,
+      "step": 251
+    },
+    {
+      "epoch": 0.03,
+      "gumbel_tau": 0.9425,
+      "image_masks_loss": 0.1007,
+      "llm_learning_rate": 4.300341296928328e-06,
+      "loss": 1.0235,
+      "output_text_mask_loss": 0.2258,
+      "predictor_learning_rate": 0.0004300341296928328,
+      "step": 252
+    },
+    {
+      "epoch": 0.03,
+      "gumbel_tau": 0.9423,
+      "image_masks_loss": 0.1007,
+      "llm_learning_rate": 4.31740614334471e-06,
+      "loss": 1.0307,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.000431740614334471,
+      "step": 253
+    },
+    {
+      "epoch": 0.03,
+      "gumbel_tau": 0.9421,
+      "image_masks_loss": 0.025,
+      "llm_learning_rate": 4.3344709897610924e-06,
+      "loss": 0.9835,
+      "output_text_mask_loss": 0.3113,
+      "predictor_learning_rate": 0.0004334470989761092,
+      "step": 254
+    },
+    {
+      "epoch": 0.03,
+      "gumbel_tau": 0.9418,
+      "image_masks_loss": 0.0418,
+      "llm_learning_rate": 4.351535836177474e-06,
+      "loss": 1.1469,
+      "output_text_mask_loss": 0.0885,
+      "predictor_learning_rate": 0.0004351535836177474,
+      "step": 255
+    },
+    {
+      "epoch": 0.03,
+      "gumbel_tau": 0.9416,
+      "image_masks_loss": 0.0257,
+      "llm_learning_rate": 4.368600682593857e-06,
+      "loss": 0.9815,
+      "output_text_mask_loss": 0.1602,
+      "predictor_learning_rate": 0.0004368600682593857,
+      "step": 256
+    },
+    {
+      "epoch": 0.03,
+      "gumbel_tau": 0.9414,
+      "image_masks_loss": 0.0252,
+      "llm_learning_rate": 4.385665529010239e-06,
+      "loss": 0.9434,
+      "output_text_mask_loss": 0.2991,
+      "predictor_learning_rate": 0.0004385665529010239,
+      "step": 257
+    },
+    {
+      "epoch": 0.03,
+      "gumbel_tau": 0.9412,
+      "image_masks_loss": 0.0232,
+      "llm_learning_rate": 4.402730375426622e-06,
+      "loss": 0.9817,
+      "output_text_mask_loss": 0.1495,
+      "predictor_learning_rate": 0.0004402730375426621,
+      "step": 258
+    },
+    {
+      "epoch": 0.03,
+      "gumbel_tau": 0.9409,
+      "image_masks_loss": 0.0169,
+      "llm_learning_rate": 4.419795221843004e-06,
+      "loss": 0.8758,
+      "output_text_mask_loss": 0.0694,
+      "predictor_learning_rate": 0.00044197952218430036,
+      "step": 259
+    },
+    {
+      "epoch": 0.03,
+      "gumbel_tau": 0.9407,
+      "image_masks_loss": 0.0273,
+      "llm_learning_rate": 4.436860068259386e-06,
+      "loss": 1.0058,
+      "output_text_mask_loss": 0.1244,
+      "predictor_learning_rate": 0.00044368600682593857,
+      "step": 260
+    },
+    {
+      "epoch": 0.03,
+      "gumbel_tau": 0.9405,
+      "image_masks_loss": 0.0292,
+      "llm_learning_rate": 4.453924914675768e-06,
+      "loss": 0.8361,
+      "output_text_mask_loss": 0.0927,
+      "predictor_learning_rate": 0.00044539249146757683,
+      "step": 261
+    },
+    {
+      "epoch": 0.03,
+      "gumbel_tau": 0.9403,
+      "image_masks_loss": 0.0629,
+      "llm_learning_rate": 4.47098976109215e-06,
+      "loss": 1.0199,
+      "output_text_mask_loss": 0.0153,
+      "predictor_learning_rate": 0.00044709897610921504,
+      "step": 262
+    },
+    {
+      "epoch": 0.03,
+      "gumbel_tau": 0.9401,
+      "image_masks_loss": 0.1472,
+      "llm_learning_rate": 4.488054607508533e-06,
+      "loss": 0.9094,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00044880546075085325,
+      "step": 263
+    },
+    {
+      "epoch": 0.03,
+      "gumbel_tau": 0.9398,
+      "image_masks_loss": 0.0553,
+      "llm_learning_rate": 4.505119453924915e-06,
+      "loss": 0.8607,
+      "output_text_mask_loss": 0.1083,
+      "predictor_learning_rate": 0.00045051194539249146,
+      "step": 264
+    },
+    {
+      "epoch": 0.03,
+      "gumbel_tau": 0.9396,
+      "image_masks_loss": 0.0908,
+      "llm_learning_rate": 4.522184300341297e-06,
+      "loss": 1.0556,
+      "output_text_mask_loss": 0.1717,
+      "predictor_learning_rate": 0.0004522184300341297,
+      "step": 265
+    },
+    {
+      "epoch": 0.03,
+      "gumbel_tau": 0.9394,
+      "image_masks_loss": 0.0877,
+      "llm_learning_rate": 4.539249146757679e-06,
+      "loss": 1.2271,
+      "output_text_mask_loss": 0.4486,
+      "predictor_learning_rate": 0.00045392491467576793,
+      "step": 266
+    },
+    {
+      "epoch": 0.03,
+      "gumbel_tau": 0.9392,
+      "image_masks_loss": 0.0549,
+      "llm_learning_rate": 4.556313993174062e-06,
+      "loss": 1.0062,
+      "output_text_mask_loss": 0.1564,
+      "predictor_learning_rate": 0.00045563139931740614,
+      "step": 267
+    },
+    {
+      "epoch": 0.03,
+      "gumbel_tau": 0.939,
+      "image_masks_loss": 0.0385,
+      "llm_learning_rate": 4.573378839590444e-06,
+      "loss": 0.95,
+      "output_text_mask_loss": 0.2838,
+      "predictor_learning_rate": 0.00045733788395904435,
+      "step": 268
+    },
+    {
+      "epoch": 0.03,
+      "gumbel_tau": 0.9387,
+      "image_masks_loss": 0.046,
+      "llm_learning_rate": 4.590443686006827e-06,
+      "loss": 1.0818,
+      "output_text_mask_loss": 0.47,
+      "predictor_learning_rate": 0.0004590443686006826,
+      "step": 269
+    },
+    {
+      "epoch": 0.03,
+      "gumbel_tau": 0.9385,
+      "image_masks_loss": 0.0835,
+      "llm_learning_rate": 4.6075085324232085e-06,
+      "loss": 0.8902,
+      "output_text_mask_loss": 0.1892,
+      "predictor_learning_rate": 0.0004607508532423209,
+      "step": 270
+    },
+    {
+      "epoch": 0.03,
+      "gumbel_tau": 0.9383,
+      "image_masks_loss": 0.0231,
+      "llm_learning_rate": 4.6245733788395904e-06,
+      "loss": 0.8734,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 0.00046245733788395903,
+      "step": 271
+    },
+    {
+      "epoch": 0.03,
+      "gumbel_tau": 0.9381,
+      "image_masks_loss": 0.0188,
+      "llm_learning_rate": 4.641638225255973e-06,
+      "loss": 0.8601,
+      "output_text_mask_loss": 0.0056,
+      "predictor_learning_rate": 0.0004641638225255973,
+      "step": 272
+    },
+    {
+      "epoch": 0.03,
+      "gumbel_tau": 0.9378,
+      "image_masks_loss": 0.0054,
+      "llm_learning_rate": 4.658703071672355e-06,
+      "loss": 0.9113,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.0004658703071672355,
+      "step": 273
+    },
+    {
+      "epoch": 0.03,
+      "gumbel_tau": 0.9376,
+      "image_masks_loss": 0.0549,
+      "llm_learning_rate": 4.675767918088738e-06,
+      "loss": 0.9157,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.00046757679180887377,
+      "step": 274
+    },
+    {
+      "epoch": 0.03,
+      "gumbel_tau": 0.9374,
+      "image_masks_loss": 0.0793,
+      "llm_learning_rate": 4.69283276450512e-06,
+      "loss": 0.8396,
+      "output_text_mask_loss": 0.0437,
+      "predictor_learning_rate": 0.0004692832764505119,
+      "step": 275
+    },
+    {
+      "epoch": 0.03,
+      "gumbel_tau": 0.9372,
+      "image_masks_loss": 0.0341,
+      "llm_learning_rate": 4.709897610921502e-06,
+      "loss": 0.84,
+      "output_text_mask_loss": 0.0004,
+      "predictor_learning_rate": 0.0004709897610921502,
+      "step": 276
+    },
+    {
+      "epoch": 0.03,
+      "gumbel_tau": 0.937,
+      "image_masks_loss": 0.0385,
+      "llm_learning_rate": 4.726962457337884e-06,
+      "loss": 0.7764,
+      "output_text_mask_loss": 0.1282,
+      "predictor_learning_rate": 0.0004726962457337884,
+      "step": 277
+    },
+    {
+      "epoch": 0.03,
+      "gumbel_tau": 0.9367,
+      "image_masks_loss": 0.079,
+      "llm_learning_rate": 4.744027303754267e-06,
+      "loss": 0.8857,
+      "output_text_mask_loss": 0.0641,
+      "predictor_learning_rate": 0.00047440273037542666,
+      "step": 278
+    },
+    {
+      "epoch": 0.03,
+      "gumbel_tau": 0.9365,
+      "image_masks_loss": 0.058,
+      "llm_learning_rate": 4.761092150170649e-06,
+      "loss": 0.8904,
+      "output_text_mask_loss": 0.2792,
+      "predictor_learning_rate": 0.0004761092150170648,
+      "step": 279
+    },
+    {
+      "epoch": 0.03,
+      "gumbel_tau": 0.9363,
+      "image_masks_loss": 0.0218,
+      "llm_learning_rate": 4.778156996587031e-06,
+      "loss": 1.009,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004778156996587031,
+      "step": 280
+    },
+    {
+      "epoch": 0.03,
+      "gumbel_tau": 0.9361,
+      "image_masks_loss": 0.0204,
+      "llm_learning_rate": 4.795221843003413e-06,
+      "loss": 1.1617,
+      "output_text_mask_loss": 0.2411,
+      "predictor_learning_rate": 0.0004795221843003413,
+      "step": 281
+    },
+    {
+      "epoch": 0.03,
+      "gumbel_tau": 0.9359,
+      "image_masks_loss": 0.0452,
+      "llm_learning_rate": 4.812286689419795e-06,
+      "loss": 1.0645,
+      "output_text_mask_loss": 0.4303,
+      "predictor_learning_rate": 0.00048122866894197955,
+      "step": 282
+    },
+    {
+      "epoch": 0.03,
+      "gumbel_tau": 0.9356,
+      "image_masks_loss": 0.0132,
+      "llm_learning_rate": 4.829351535836178e-06,
+      "loss": 0.9589,
+      "output_text_mask_loss": 0.1587,
+      "predictor_learning_rate": 0.00048293515358361776,
+      "step": 283
+    },
+    {
+      "epoch": 0.03,
+      "gumbel_tau": 0.9354,
+      "image_masks_loss": 0.0259,
+      "llm_learning_rate": 4.84641638225256e-06,
+      "loss": 0.7861,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00048464163822525597,
+      "step": 284
+    },
+    {
+      "epoch": 0.03,
+      "gumbel_tau": 0.9352,
+      "image_masks_loss": 0.0214,
+      "llm_learning_rate": 4.863481228668943e-06,
+      "loss": 0.7513,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00048634812286689423,
+      "step": 285
+    },
+    {
+      "epoch": 0.03,
+      "gumbel_tau": 0.935,
+      "image_masks_loss": 0.0483,
+      "llm_learning_rate": 4.880546075085325e-06,
+      "loss": 0.8173,
+      "output_text_mask_loss": 0.0248,
+      "predictor_learning_rate": 0.00048805460750853244,
+      "step": 286
+    },
+    {
+      "epoch": 0.03,
+      "gumbel_tau": 0.9348,
+      "image_masks_loss": 0.0244,
+      "llm_learning_rate": 4.897610921501707e-06,
+      "loss": 0.8429,
+      "output_text_mask_loss": 0.238,
+      "predictor_learning_rate": 0.0004897610921501706,
+      "step": 287
+    },
+    {
+      "epoch": 0.03,
+      "gumbel_tau": 0.9345,
+      "image_masks_loss": 0.0626,
+      "llm_learning_rate": 4.914675767918089e-06,
+      "loss": 0.8542,
+      "output_text_mask_loss": 0.0526,
+      "predictor_learning_rate": 0.0004914675767918089,
+      "step": 288
+    },
+    {
+      "epoch": 0.03,
+      "gumbel_tau": 0.9343,
+      "image_masks_loss": 0.0166,
+      "llm_learning_rate": 4.931740614334471e-06,
+      "loss": 0.898,
+      "output_text_mask_loss": 0.1625,
+      "predictor_learning_rate": 0.0004931740614334471,
+      "step": 289
+    },
+    {
+      "epoch": 0.03,
+      "gumbel_tau": 0.9341,
+      "image_masks_loss": 0.0131,
+      "llm_learning_rate": 4.948805460750853e-06,
+      "loss": 0.9843,
+      "output_text_mask_loss": 0.0326,
+      "predictor_learning_rate": 0.0004948805460750853,
+      "step": 290
+    },
+    {
+      "epoch": 0.03,
+      "gumbel_tau": 0.9339,
+      "image_masks_loss": 0.0345,
+      "llm_learning_rate": 4.965870307167236e-06,
+      "loss": 0.8433,
+      "output_text_mask_loss": 0.0058,
+      "predictor_learning_rate": 0.0004965870307167236,
+      "step": 291
+    },
+    {
+      "epoch": 0.03,
+      "gumbel_tau": 0.9337,
+      "image_masks_loss": 0.0124,
+      "llm_learning_rate": 4.982935153583618e-06,
+      "loss": 0.9095,
+      "output_text_mask_loss": 0.0004,
+      "predictor_learning_rate": 0.0004982935153583618,
+      "step": 292
+    },
+    {
+      "epoch": 0.03,
+      "gumbel_tau": 0.9334,
+      "image_masks_loss": 0.0284,
+      "llm_learning_rate": 5e-06,
+      "loss": 0.8123,
+      "output_text_mask_loss": 0.0275,
+      "predictor_learning_rate": 0.0005,
+      "step": 293
+    },
+    {
+      "epoch": 0.03,
+      "gumbel_tau": 0.9332,
+      "image_masks_loss": 0.0277,
+      "llm_learning_rate": 4.999999862347216e-06,
+      "loss": 0.7864,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004999999862347216,
+      "step": 294
+    },
+    {
+      "epoch": 0.03,
+      "gumbel_tau": 0.933,
+      "image_masks_loss": 0.0196,
+      "llm_learning_rate": 4.999999449388878e-06,
+      "loss": 0.8175,
+      "output_text_mask_loss": 0.1717,
+      "predictor_learning_rate": 0.0004999999449388878,
+      "step": 295
+    },
+    {
+      "epoch": 0.03,
+      "gumbel_tau": 0.9328,
+      "image_masks_loss": 0.0305,
+      "llm_learning_rate": 4.9999987611250335e-06,
+      "loss": 0.8707,
+      "output_text_mask_loss": 0.0687,
+      "predictor_learning_rate": 0.0004999998761125033,
+      "step": 296
+    },
+    {
+      "epoch": 0.03,
+      "gumbel_tau": 0.9326,
+      "image_masks_loss": 0.0157,
+      "llm_learning_rate": 4.999997797555756e-06,
+      "loss": 0.7325,
+      "output_text_mask_loss": 0.0763,
+      "predictor_learning_rate": 0.0004999997797555755,
+      "step": 297
+    },
+    {
+      "epoch": 0.03,
+      "gumbel_tau": 0.9323,
+      "image_masks_loss": 0.0324,
+      "llm_learning_rate": 4.999996558681152e-06,
+      "loss": 0.9039,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004999996558681151,
+      "step": 298
+    },
+    {
+      "epoch": 0.03,
+      "gumbel_tau": 0.9321,
+      "image_masks_loss": 0.045,
+      "llm_learning_rate": 4.999995044501359e-06,
+      "loss": 0.7929,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.0004999995044501358,
+      "step": 299
+    },
+    {
+      "epoch": 0.03,
+      "gumbel_tau": 0.9319,
+      "image_masks_loss": 0.0145,
+      "llm_learning_rate": 4.999993255016543e-06,
+      "loss": 0.8932,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 0.0004999993255016543,
+      "step": 300
+    },
+    {
+      "epoch": 0.03,
+      "gumbel_tau": 0.9317,
+      "image_masks_loss": 0.0197,
+      "llm_learning_rate": 4.999991190226901e-06,
+      "loss": 0.8906,
+      "output_text_mask_loss": 0.1434,
+      "predictor_learning_rate": 0.00049999911902269,
+      "step": 301
+    },
+    {
+      "epoch": 0.03,
+      "gumbel_tau": 0.9315,
+      "image_masks_loss": 0.0113,
+      "llm_learning_rate": 4.99998885013266e-06,
+      "loss": 0.962,
+      "output_text_mask_loss": 0.0259,
+      "predictor_learning_rate": 0.000499998885013266,
+      "step": 302
+    },
+    {
+      "epoch": 0.03,
+      "gumbel_tau": 0.9312,
+      "image_masks_loss": 0.0374,
+      "llm_learning_rate": 4.99998623473408e-06,
+      "loss": 0.9058,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.000499998623473408,
+      "step": 303
+    },
+    {
+      "epoch": 0.03,
+      "gumbel_tau": 0.931,
+      "image_masks_loss": 0.0275,
+      "llm_learning_rate": 4.999983344031446e-06,
+      "loss": 0.7237,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004999983344031446,
+      "step": 304
+    },
+    {
+      "epoch": 0.03,
+      "gumbel_tau": 0.9308,
+      "image_masks_loss": 0.0105,
+      "llm_learning_rate": 4.999980178025078e-06,
+      "loss": 1.0101,
+      "output_text_mask_loss": 0.0549,
+      "predictor_learning_rate": 0.0004999980178025078,
+      "step": 305
+    },
+    {
+      "epoch": 0.03,
+      "gumbel_tau": 0.9306,
+      "image_masks_loss": 0.02,
+      "llm_learning_rate": 4.999976736715325e-06,
+      "loss": 0.9071,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004999976736715325,
+      "step": 306
+    },
+    {
+      "epoch": 0.03,
+      "gumbel_tau": 0.9304,
+      "image_masks_loss": 0.0246,
+      "llm_learning_rate": 4.999973020102564e-06,
+      "loss": 0.8728,
+      "output_text_mask_loss": 0.0065,
+      "predictor_learning_rate": 0.0004999973020102564,
+      "step": 307
+    },
+    {
+      "epoch": 0.03,
+      "gumbel_tau": 0.9301,
+      "image_masks_loss": 0.021,
+      "llm_learning_rate": 4.999969028187207e-06,
+      "loss": 0.864,
+      "output_text_mask_loss": 0.0193,
+      "predictor_learning_rate": 0.0004999969028187207,
+      "step": 308
+    },
+    {
+      "epoch": 0.03,
+      "gumbel_tau": 0.9299,
+      "image_masks_loss": 0.0196,
+      "llm_learning_rate": 4.999964760969691e-06,
+      "loss": 0.8931,
+      "output_text_mask_loss": 0.0763,
+      "predictor_learning_rate": 0.0004999964760969691,
+      "step": 309
+    },
+    {
+      "epoch": 0.03,
+      "gumbel_tau": 0.9297,
+      "image_masks_loss": 0.0277,
+      "llm_learning_rate": 4.9999602184504884e-06,
+      "loss": 0.8917,
+      "output_text_mask_loss": 0.0087,
+      "predictor_learning_rate": 0.0004999960218450488,
+      "step": 310
+    },
+    {
+      "epoch": 0.03,
+      "gumbel_tau": 0.9295,
+      "image_masks_loss": 0.0267,
+      "llm_learning_rate": 4.999955400630096e-06,
+      "loss": 0.8354,
+      "output_text_mask_loss": 0.0433,
+      "predictor_learning_rate": 0.0004999955400630097,
+      "step": 311
+    },
+    {
+      "epoch": 0.03,
+      "gumbel_tau": 0.9293,
+      "image_masks_loss": 0.0507,
+      "llm_learning_rate": 4.999950307509048e-06,
+      "loss": 0.7159,
+      "output_text_mask_loss": 0.0118,
+      "predictor_learning_rate": 0.0004999950307509048,
+      "step": 312
+    },
+    {
+      "epoch": 0.03,
+      "gumbel_tau": 0.929,
+      "image_masks_loss": 0.032,
+      "llm_learning_rate": 4.999944939087903e-06,
+      "loss": 0.8964,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004999944939087903,
+      "step": 313
+    },
+    {
+      "epoch": 0.03,
+      "gumbel_tau": 0.9288,
+      "image_masks_loss": 0.0402,
+      "llm_learning_rate": 4.999939295367254e-06,
+      "loss": 0.8996,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.0004999939295367254,
+      "step": 314
+    },
+    {
+      "epoch": 0.03,
+      "gumbel_tau": 0.9286,
+      "image_masks_loss": 0.0418,
+      "llm_learning_rate": 4.999933376347721e-06,
+      "loss": 0.9399,
+      "output_text_mask_loss": 0.0565,
+      "predictor_learning_rate": 0.000499993337634772,
+      "step": 315
+    },
+    {
+      "epoch": 0.03,
+      "gumbel_tau": 0.9284,
+      "image_masks_loss": 0.0383,
+      "llm_learning_rate": 4.999927182029955e-06,
+      "loss": 0.8665,
+      "output_text_mask_loss": 0.0629,
+      "predictor_learning_rate": 0.0004999927182029955,
+      "step": 316
+    },
+    {
+      "epoch": 0.03,
+      "gumbel_tau": 0.9282,
+      "image_masks_loss": 0.0092,
+      "llm_learning_rate": 4.999920712414641e-06,
+      "loss": 0.7372,
+      "output_text_mask_loss": 0.0553,
+      "predictor_learning_rate": 0.0004999920712414641,
+      "step": 317
+    },
+    {
+      "epoch": 0.03,
+      "gumbel_tau": 0.9279,
+      "image_masks_loss": 0.0198,
+      "llm_learning_rate": 4.9999139675024885e-06,
+      "loss": 0.8607,
+      "output_text_mask_loss": 0.1534,
+      "predictor_learning_rate": 0.0004999913967502488,
+      "step": 318
+    },
+    {
+      "epoch": 0.03,
+      "gumbel_tau": 0.9277,
+      "image_masks_loss": 0.0085,
+      "llm_learning_rate": 4.9999069472942415e-06,
+      "loss": 0.8246,
+      "output_text_mask_loss": 0.1953,
+      "predictor_learning_rate": 0.0004999906947294241,
+      "step": 319
+    },
+    {
+      "epoch": 0.03,
+      "gumbel_tau": 0.9275,
+      "image_masks_loss": 0.0173,
+      "llm_learning_rate": 4.999899651790675e-06,
+      "loss": 0.8511,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.0004999899651790674,
+      "step": 320
+    },
+    {
+      "epoch": 0.03,
+      "gumbel_tau": 0.9273,
+      "image_masks_loss": 0.0368,
+      "llm_learning_rate": 4.999892080992589e-06,
+      "loss": 0.9007,
+      "output_text_mask_loss": 0.025,
+      "predictor_learning_rate": 0.0004999892080992588,
+      "step": 321
+    },
+    {
+      "epoch": 0.03,
+      "gumbel_tau": 0.9271,
+      "image_masks_loss": 0.0271,
+      "llm_learning_rate": 4.999884234900818e-06,
+      "loss": 0.8497,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.0004999884234900818,
+      "step": 322
+    },
+    {
+      "epoch": 0.03,
+      "gumbel_tau": 0.9268,
+      "image_masks_loss": 0.021,
+      "llm_learning_rate": 4.999876113516229e-06,
+      "loss": 0.9316,
+      "output_text_mask_loss": 0.1518,
+      "predictor_learning_rate": 0.0004999876113516229,
+      "step": 323
+    },
+    {
+      "epoch": 0.03,
+      "gumbel_tau": 0.9266,
+      "image_masks_loss": 0.0155,
+      "llm_learning_rate": 4.9998677168397124e-06,
+      "loss": 0.7027,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004999867716839712,
+      "step": 324
+    },
+    {
+      "epoch": 0.03,
+      "gumbel_tau": 0.9264,
+      "image_masks_loss": 0.0263,
+      "llm_learning_rate": 4.999859044872195e-06,
+      "loss": 0.79,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004999859044872195,
+      "step": 325
+    },
+    {
+      "epoch": 0.03,
+      "gumbel_tau": 0.9262,
+      "image_masks_loss": 0.0161,
+      "llm_learning_rate": 4.9998500976146316e-06,
+      "loss": 0.9394,
+      "output_text_mask_loss": 0.0717,
+      "predictor_learning_rate": 0.0004999850097614631,
+      "step": 326
+    },
+    {
+      "epoch": 0.03,
+      "gumbel_tau": 0.926,
+      "image_masks_loss": 0.0187,
+      "llm_learning_rate": 4.9998408750680075e-06,
+      "loss": 0.8599,
+      "output_text_mask_loss": 0.0549,
+      "predictor_learning_rate": 0.0004999840875068007,
+      "step": 327
+    },
+    {
+      "epoch": 0.03,
+      "gumbel_tau": 0.9258,
+      "image_masks_loss": 0.0572,
+      "llm_learning_rate": 4.999831377233337e-06,
+      "loss": 0.8828,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004999831377233337,
+      "step": 328
+    },
+    {
+      "epoch": 0.03,
+      "gumbel_tau": 0.9255,
+      "image_masks_loss": 0.0397,
+      "llm_learning_rate": 4.999821604111668e-06,
+      "loss": 0.8645,
+      "output_text_mask_loss": 0.0111,
+      "predictor_learning_rate": 0.0004999821604111667,
+      "step": 329
+    },
+    {
+      "epoch": 0.03,
+      "gumbel_tau": 0.9253,
+      "image_masks_loss": 0.0402,
+      "llm_learning_rate": 4.9998115557040754e-06,
+      "loss": 0.7591,
+      "output_text_mask_loss": 0.0393,
+      "predictor_learning_rate": 0.0004999811555704075,
+      "step": 330
+    },
+    {
+      "epoch": 0.03,
+      "gumbel_tau": 0.9251,
+      "image_masks_loss": 0.0364,
+      "llm_learning_rate": 4.999801232011666e-06,
+      "loss": 0.8883,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 0.0004999801232011666,
+      "step": 331
+    },
+    {
+      "epoch": 0.03,
+      "gumbel_tau": 0.9249,
+      "image_masks_loss": 0.01,
+      "llm_learning_rate": 4.999790633035577e-06,
+      "loss": 0.9099,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.0004999790633035577,
+      "step": 332
+    },
+    {
+      "epoch": 0.03,
+      "gumbel_tau": 0.9247,
+      "image_masks_loss": 0.0047,
+      "llm_learning_rate": 4.999779758776976e-06,
+      "loss": 0.8214,
+      "output_text_mask_loss": 0.0299,
+      "predictor_learning_rate": 0.0004999779758776975,
+      "step": 333
+    },
+    {
+      "epoch": 0.03,
+      "gumbel_tau": 0.9244,
+      "image_masks_loss": 0.0218,
+      "llm_learning_rate": 4.99976860923706e-06,
+      "loss": 0.814,
+      "output_text_mask_loss": 0.1495,
+      "predictor_learning_rate": 0.000499976860923706,
+      "step": 334
+    },
+    {
+      "epoch": 0.03,
+      "gumbel_tau": 0.9242,
+      "image_masks_loss": 0.0376,
+      "llm_learning_rate": 4.999757184417056e-06,
+      "loss": 0.8577,
+      "output_text_mask_loss": 0.0488,
+      "predictor_learning_rate": 0.0004999757184417055,
+      "step": 335
+    },
+    {
+      "epoch": 0.03,
+      "gumbel_tau": 0.924,
+      "image_masks_loss": 0.0353,
+      "llm_learning_rate": 4.999745484318223e-06,
+      "loss": 0.801,
+      "output_text_mask_loss": 0.1091,
+      "predictor_learning_rate": 0.0004999745484318222,
+      "step": 336
+    },
+    {
+      "epoch": 0.03,
+      "gumbel_tau": 0.9238,
+      "image_masks_loss": 0.0162,
+      "llm_learning_rate": 4.9997335089418496e-06,
+      "loss": 0.7833,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.0004999733508941849,
+      "step": 337
+    },
+    {
+      "epoch": 0.03,
+      "gumbel_tau": 0.9236,
+      "image_masks_loss": 0.0307,
+      "llm_learning_rate": 4.999721258289254e-06,
+      "loss": 0.7735,
+      "output_text_mask_loss": 0.0648,
+      "predictor_learning_rate": 0.0004999721258289254,
+      "step": 338
+    },
+    {
+      "epoch": 0.03,
+      "gumbel_tau": 0.9234,
+      "image_masks_loss": 0.0192,
+      "llm_learning_rate": 4.999708732361785e-06,
+      "loss": 0.8779,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004999708732361785,
+      "step": 339
+    },
+    {
+      "epoch": 0.03,
+      "gumbel_tau": 0.9231,
+      "image_masks_loss": 0.0133,
+      "llm_learning_rate": 4.999695931160824e-06,
+      "loss": 0.9111,
+      "output_text_mask_loss": 0.0012,
+      "predictor_learning_rate": 0.0004999695931160824,
+      "step": 340
+    },
+    {
+      "epoch": 0.03,
+      "gumbel_tau": 0.9229,
+      "image_masks_loss": 0.0303,
+      "llm_learning_rate": 4.999682854687778e-06,
+      "loss": 0.8551,
+      "output_text_mask_loss": 0.0603,
+      "predictor_learning_rate": 0.0004999682854687778,
+      "step": 341
+    },
+    {
+      "epoch": 0.04,
+      "gumbel_tau": 0.9227,
+      "image_masks_loss": 0.0353,
+      "llm_learning_rate": 4.999669502944089e-06,
+      "loss": 0.908,
+      "output_text_mask_loss": 0.058,
+      "predictor_learning_rate": 0.0004999669502944088,
+      "step": 342
+    },
+    {
+      "epoch": 0.04,
+      "gumbel_tau": 0.9225,
+      "image_masks_loss": 0.0244,
+      "llm_learning_rate": 4.999655875931225e-06,
+      "loss": 0.6894,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004999655875931225,
+      "step": 343
+    },
+    {
+      "epoch": 0.04,
+      "gumbel_tau": 0.9223,
+      "image_masks_loss": 0.0374,
+      "llm_learning_rate": 4.99964197365069e-06,
+      "loss": 0.8282,
+      "output_text_mask_loss": 0.2167,
+      "predictor_learning_rate": 0.0004999641973650689,
+      "step": 344
+    },
+    {
+      "epoch": 0.04,
+      "gumbel_tau": 0.922,
+      "image_masks_loss": 0.0202,
+      "llm_learning_rate": 4.999627796104012e-06,
+      "loss": 0.7943,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.0004999627796104013,
+      "step": 345
+    },
+    {
+      "epoch": 0.04,
+      "gumbel_tau": 0.9218,
+      "image_masks_loss": 0.0698,
+      "llm_learning_rate": 4.999613343292754e-06,
+      "loss": 0.9849,
+      "output_text_mask_loss": 0.0065,
+      "predictor_learning_rate": 0.0004999613343292754,
+      "step": 346
+    },
+    {
+      "epoch": 0.04,
+      "gumbel_tau": 0.9216,
+      "image_masks_loss": 0.0343,
+      "llm_learning_rate": 4.9995986152185075e-06,
+      "loss": 0.8854,
+      "output_text_mask_loss": 0.0174,
+      "predictor_learning_rate": 0.0004999598615218507,
+      "step": 347
+    },
+    {
+      "epoch": 0.04,
+      "gumbel_tau": 0.9214,
+      "image_masks_loss": 0.0066,
+      "llm_learning_rate": 4.999583611882892e-06,
+      "loss": 0.8579,
+      "output_text_mask_loss": 0.0391,
+      "predictor_learning_rate": 0.0004999583611882892,
+      "step": 348
+    },
+    {
+      "epoch": 0.04,
+      "gumbel_tau": 0.9212,
+      "image_masks_loss": 0.0124,
+      "llm_learning_rate": 4.999568333287563e-06,
+      "loss": 0.7768,
+      "output_text_mask_loss": 0.0095,
+      "predictor_learning_rate": 0.0004999568333287563,
+      "step": 349
+    },
+    {
+      "epoch": 0.04,
+      "gumbel_tau": 0.921,
+      "image_masks_loss": 0.0226,
+      "llm_learning_rate": 4.999552779434201e-06,
+      "loss": 0.7951,
+      "output_text_mask_loss": 0.0139,
+      "predictor_learning_rate": 0.0004999552779434201,
+      "step": 350
+    },
+    {
+      "epoch": 0.04,
+      "gumbel_tau": 0.9207,
+      "image_masks_loss": 0.0132,
+      "llm_learning_rate": 4.99953695032452e-06,
+      "loss": 0.8058,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.000499953695032452,
+      "step": 351
+    },
+    {
+      "epoch": 0.04,
+      "gumbel_tau": 0.9205,
+      "image_masks_loss": 0.0194,
+      "llm_learning_rate": 4.999520845960262e-06,
+      "loss": 0.8554,
+      "output_text_mask_loss": 0.0081,
+      "predictor_learning_rate": 0.0004999520845960262,
+      "step": 352
+    },
+    {
+      "epoch": 0.04,
+      "gumbel_tau": 0.9203,
+      "image_masks_loss": 0.0203,
+      "llm_learning_rate": 4.999504466343202e-06,
+      "loss": 0.7921,
+      "output_text_mask_loss": 0.0885,
+      "predictor_learning_rate": 0.0004999504466343201,
+      "step": 353
+    },
+    {
+      "epoch": 0.04,
+      "gumbel_tau": 0.9201,
+      "image_masks_loss": 0.0294,
+      "llm_learning_rate": 4.999487811475141e-06,
+      "loss": 0.908,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 0.000499948781147514,
+      "step": 354
+    },
+    {
+      "epoch": 0.04,
+      "gumbel_tau": 0.9199,
+      "image_masks_loss": 0.0252,
+      "llm_learning_rate": 4.999470881357915e-06,
+      "loss": 0.8269,
+      "output_text_mask_loss": 0.0475,
+      "predictor_learning_rate": 0.0004999470881357916,
+      "step": 355
+    },
+    {
+      "epoch": 0.04,
+      "gumbel_tau": 0.9197,
+      "image_masks_loss": 0.0134,
+      "llm_learning_rate": 4.99945367599339e-06,
+      "loss": 0.9074,
+      "output_text_mask_loss": 0.0238,
+      "predictor_learning_rate": 0.000499945367599339,
+      "step": 356
+    },
+    {
+      "epoch": 0.04,
+      "gumbel_tau": 0.9194,
+      "image_masks_loss": 0.0286,
+      "llm_learning_rate": 4.999436195383457e-06,
+      "loss": 0.936,
+      "output_text_mask_loss": 0.1312,
+      "predictor_learning_rate": 0.0004999436195383456,
+      "step": 357
+    },
+    {
+      "epoch": 0.04,
+      "gumbel_tau": 0.9192,
+      "image_masks_loss": 0.0238,
+      "llm_learning_rate": 4.999418439530044e-06,
+      "loss": 0.7935,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.0004999418439530043,
+      "step": 358
+    },
+    {
+      "epoch": 0.04,
+      "gumbel_tau": 0.919,
+      "image_masks_loss": 0.0389,
+      "llm_learning_rate": 4.9994004084351045e-06,
+      "loss": 0.9073,
+      "output_text_mask_loss": 0.0549,
+      "predictor_learning_rate": 0.0004999400408435104,
+      "step": 359
+    },
+    {
+      "epoch": 0.04,
+      "gumbel_tau": 0.9188,
+      "image_masks_loss": 0.037,
+      "llm_learning_rate": 4.999382102100625e-06,
+      "loss": 0.7435,
+      "output_text_mask_loss": 0.0374,
+      "predictor_learning_rate": 0.0004999382102100625,
+      "step": 360
+    },
+    {
+      "epoch": 0.04,
+      "gumbel_tau": 0.9186,
+      "image_masks_loss": 0.012,
+      "llm_learning_rate": 4.999363520528622e-06,
+      "loss": 0.8273,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.0004999363520528621,
+      "step": 361
+    },
+    {
+      "epoch": 0.04,
+      "gumbel_tau": 0.9184,
+      "image_masks_loss": 0.0096,
+      "llm_learning_rate": 4.999344663721141e-06,
+      "loss": 0.906,
+      "output_text_mask_loss": 0.046,
+      "predictor_learning_rate": 0.000499934466372114,
+      "step": 362
+    },
+    {
+      "epoch": 0.04,
+      "gumbel_tau": 0.9181,
+      "image_masks_loss": 0.0111,
+      "llm_learning_rate": 4.9993255316802565e-06,
+      "loss": 0.8343,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.0004999325531680256,
+      "step": 363
+    },
+    {
+      "epoch": 0.04,
+      "gumbel_tau": 0.9179,
+      "image_masks_loss": 0.0117,
+      "llm_learning_rate": 4.99930612440808e-06,
+      "loss": 0.7641,
+      "output_text_mask_loss": 0.0012,
+      "predictor_learning_rate": 0.0004999306124408079,
+      "step": 364
+    },
+    {
+      "epoch": 0.04,
+      "gumbel_tau": 0.9177,
+      "image_masks_loss": 0.0538,
+      "llm_learning_rate": 4.9992864419067445e-06,
+      "loss": 0.8156,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004999286441906744,
+      "step": 365
+    },
+    {
+      "epoch": 0.04,
+      "gumbel_tau": 0.9175,
+      "image_masks_loss": 0.0299,
+      "llm_learning_rate": 4.9992664841784196e-06,
+      "loss": 0.7456,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004999266484178419,
+      "step": 366
+    },
+    {
+      "epoch": 0.04,
+      "gumbel_tau": 0.9173,
+      "image_masks_loss": 0.0805,
+      "llm_learning_rate": 4.999246251225303e-06,
+      "loss": 0.856,
+      "output_text_mask_loss": 0.1831,
+      "predictor_learning_rate": 0.0004999246251225302,
+      "step": 367
+    },
+    {
+      "epoch": 0.04,
+      "gumbel_tau": 0.9171,
+      "image_masks_loss": 0.0309,
+      "llm_learning_rate": 4.999225743049622e-06,
+      "loss": 0.7465,
+      "output_text_mask_loss": 0.071,
+      "predictor_learning_rate": 0.0004999225743049622,
+      "step": 368
+    },
+    {
+      "epoch": 0.04,
+      "gumbel_tau": 0.9168,
+      "image_masks_loss": 0.0137,
+      "llm_learning_rate": 4.999204959653635e-06,
+      "loss": 0.8088,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004999204959653634,
+      "step": 369
+    },
+    {
+      "epoch": 0.04,
+      "gumbel_tau": 0.9166,
+      "image_masks_loss": 0.0082,
+      "llm_learning_rate": 4.999183901039631e-06,
+      "loss": 0.7836,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004999183901039632,
+      "step": 370
+    },
+    {
+      "epoch": 0.04,
+      "gumbel_tau": 0.9164,
+      "image_masks_loss": 0.0265,
+      "llm_learning_rate": 4.99916256720993e-06,
+      "loss": 0.8242,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.000499916256720993,
+      "step": 371
+    },
+    {
+      "epoch": 0.04,
+      "gumbel_tau": 0.9162,
+      "image_masks_loss": 0.0189,
+      "llm_learning_rate": 4.999140958166879e-06,
+      "loss": 0.9184,
+      "output_text_mask_loss": 0.0549,
+      "predictor_learning_rate": 0.0004999140958166879,
+      "step": 372
+    },
+    {
+      "epoch": 0.04,
+      "gumbel_tau": 0.916,
+      "image_masks_loss": 0.0216,
+      "llm_learning_rate": 4.99911907391286e-06,
+      "loss": 0.8167,
+      "output_text_mask_loss": 0.0256,
+      "predictor_learning_rate": 0.000499911907391286,
+      "step": 373
+    },
+    {
+      "epoch": 0.04,
+      "gumbel_tau": 0.9158,
+      "image_masks_loss": 0.0259,
+      "llm_learning_rate": 4.9990969144502824e-06,
+      "loss": 0.8355,
+      "output_text_mask_loss": 0.0332,
+      "predictor_learning_rate": 0.0004999096914450282,
+      "step": 374
+    },
+    {
+      "epoch": 0.04,
+      "gumbel_tau": 0.9155,
+      "image_masks_loss": 0.014,
+      "llm_learning_rate": 4.999074479781586e-06,
+      "loss": 0.7632,
+      "output_text_mask_loss": 0.0406,
+      "predictor_learning_rate": 0.0004999074479781585,
+      "step": 375
+    },
+    {
+      "epoch": 0.04,
+      "gumbel_tau": 0.9153,
+      "image_masks_loss": 0.0075,
+      "llm_learning_rate": 4.999051769909241e-06,
+      "loss": 0.9124,
+      "output_text_mask_loss": 0.0313,
+      "predictor_learning_rate": 0.0004999051769909241,
+      "step": 376
+    },
+    {
+      "epoch": 0.04,
+      "gumbel_tau": 0.9151,
+      "image_masks_loss": 0.032,
+      "llm_learning_rate": 4.999028784835749e-06,
+      "loss": 0.9003,
+      "output_text_mask_loss": 0.0603,
+      "predictor_learning_rate": 0.0004999028784835748,
+      "step": 377
+    },
+    {
+      "epoch": 0.04,
+      "gumbel_tau": 0.9149,
+      "image_masks_loss": 0.0307,
+      "llm_learning_rate": 4.99900552456364e-06,
+      "loss": 0.8423,
+      "output_text_mask_loss": 0.1297,
+      "predictor_learning_rate": 0.000499900552456364,
+      "step": 378
+    },
+    {
+      "epoch": 0.04,
+      "gumbel_tau": 0.9147,
+      "image_masks_loss": 0.0126,
+      "llm_learning_rate": 4.998981989095477e-06,
+      "loss": 0.7539,
+      "output_text_mask_loss": 0.0877,
+      "predictor_learning_rate": 0.0004998981989095477,
+      "step": 379
+    },
+    {
+      "epoch": 0.04,
+      "gumbel_tau": 0.9145,
+      "image_masks_loss": 0.0225,
+      "llm_learning_rate": 4.998958178433851e-06,
+      "loss": 0.8778,
+      "output_text_mask_loss": 0.0242,
+      "predictor_learning_rate": 0.000499895817843385,
+      "step": 380
+    },
+    {
+      "epoch": 0.04,
+      "gumbel_tau": 0.9143,
+      "image_masks_loss": 0.0706,
+      "llm_learning_rate": 4.998934092581384e-06,
+      "loss": 0.8819,
+      "output_text_mask_loss": 0.0185,
+      "predictor_learning_rate": 0.0004998934092581384,
+      "step": 381
+    },
+    {
+      "epoch": 0.04,
+      "gumbel_tau": 0.914,
+      "image_masks_loss": 0.0511,
+      "llm_learning_rate": 4.998909731540728e-06,
+      "loss": 0.7686,
+      "output_text_mask_loss": 0.0496,
+      "predictor_learning_rate": 0.0004998909731540729,
+      "step": 382
+    },
+    {
+      "epoch": 0.04,
+      "gumbel_tau": 0.9138,
+      "image_masks_loss": 0.0416,
+      "llm_learning_rate": 4.998885095314568e-06,
+      "loss": 0.8191,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.0004998885095314567,
+      "step": 383
+    },
+    {
+      "epoch": 0.04,
+      "gumbel_tau": 0.9136,
+      "image_masks_loss": 0.0138,
+      "llm_learning_rate": 4.9988601839056135e-06,
+      "loss": 0.9105,
+      "output_text_mask_loss": 0.0954,
+      "predictor_learning_rate": 0.0004998860183905614,
+      "step": 384
+    },
+    {
+      "epoch": 0.04,
+      "gumbel_tau": 0.9134,
+      "image_masks_loss": 0.0481,
+      "llm_learning_rate": 4.998834997316611e-06,
+      "loss": 0.8114,
+      "output_text_mask_loss": 0.0496,
+      "predictor_learning_rate": 0.000499883499731661,
+      "step": 385
+    },
+    {
+      "epoch": 0.04,
+      "gumbel_tau": 0.9132,
+      "image_masks_loss": 0.0418,
+      "llm_learning_rate": 4.9988095355503316e-06,
+      "loss": 0.7924,
+      "output_text_mask_loss": 0.0202,
+      "predictor_learning_rate": 0.0004998809535550332,
+      "step": 386
+    },
+    {
+      "epoch": 0.04,
+      "gumbel_tau": 0.913,
+      "image_masks_loss": 0.0164,
+      "llm_learning_rate": 4.998783798609581e-06,
+      "loss": 0.8128,
+      "output_text_mask_loss": 0.0023,
+      "predictor_learning_rate": 0.000499878379860958,
+      "step": 387
+    },
+    {
+      "epoch": 0.04,
+      "gumbel_tau": 0.9127,
+      "image_masks_loss": 0.0277,
+      "llm_learning_rate": 4.998757786497192e-06,
+      "loss": 0.9912,
+      "output_text_mask_loss": 0.0467,
+      "predictor_learning_rate": 0.0004998757786497191,
+      "step": 388
+    },
+    {
+      "epoch": 0.04,
+      "gumbel_tau": 0.9125,
+      "image_masks_loss": 0.042,
+      "llm_learning_rate": 4.99873149921603e-06,
+      "loss": 1.1995,
+      "output_text_mask_loss": 0.3876,
+      "predictor_learning_rate": 0.000499873149921603,
+      "step": 389
+    },
+    {
+      "epoch": 0.04,
+      "gumbel_tau": 0.9123,
+      "image_masks_loss": 0.0299,
+      "llm_learning_rate": 4.99870493676899e-06,
+      "loss": 1.215,
+      "output_text_mask_loss": 0.1854,
+      "predictor_learning_rate": 0.000499870493676899,
+      "step": 390
+    },
+    {
+      "epoch": 0.04,
+      "gumbel_tau": 0.9121,
+      "image_masks_loss": 0.0178,
+      "llm_learning_rate": 4.998678099158997e-06,
+      "loss": 0.9877,
+      "output_text_mask_loss": 0.3296,
+      "predictor_learning_rate": 0.0004998678099158996,
+      "step": 391
+    },
+    {
+      "epoch": 0.04,
+      "gumbel_tau": 0.9119,
+      "image_masks_loss": 0.061,
+      "llm_learning_rate": 4.998650986389005e-06,
+      "loss": 0.8004,
+      "output_text_mask_loss": 0.1282,
+      "predictor_learning_rate": 0.0004998650986389005,
+      "step": 392
+    },
+    {
+      "epoch": 0.04,
+      "gumbel_tau": 0.9117,
+      "image_masks_loss": 0.0595,
+      "llm_learning_rate": 4.9986235984620016e-06,
+      "loss": 0.8807,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.0004998623598462002,
+      "step": 393
+    },
+    {
+      "epoch": 0.04,
+      "gumbel_tau": 0.9115,
+      "image_masks_loss": 0.0153,
+      "llm_learning_rate": 4.998595935381001e-06,
+      "loss": 0.8469,
+      "output_text_mask_loss": 0.0942,
+      "predictor_learning_rate": 0.0004998595935381001,
+      "step": 394
+    },
+    {
+      "epoch": 0.04,
+      "gumbel_tau": 0.9112,
+      "image_masks_loss": 0.0168,
+      "llm_learning_rate": 4.998567997149052e-06,
+      "loss": 0.8371,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004998567997149052,
+      "step": 395
+    },
+    {
+      "epoch": 0.04,
+      "gumbel_tau": 0.911,
+      "image_masks_loss": 0.0324,
+      "llm_learning_rate": 4.998539783769229e-06,
+      "loss": 1.0029,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.000499853978376923,
+      "step": 396
+    },
+    {
+      "epoch": 0.04,
+      "gumbel_tau": 0.9108,
+      "image_masks_loss": 0.0591,
+      "llm_learning_rate": 4.9985112952446406e-06,
+      "loss": 1.0179,
+      "output_text_mask_loss": 0.2075,
+      "predictor_learning_rate": 0.000499851129524464,
+      "step": 397
+    },
+    {
+      "epoch": 0.04,
+      "gumbel_tau": 0.9106,
+      "image_masks_loss": 0.0175,
+      "llm_learning_rate": 4.998482531578423e-06,
+      "loss": 1.1628,
+      "output_text_mask_loss": 0.1793,
+      "predictor_learning_rate": 0.0004998482531578422,
+      "step": 398
+    },
+    {
+      "epoch": 0.04,
+      "gumbel_tau": 0.9104,
+      "image_masks_loss": 0.0429,
+      "llm_learning_rate": 4.9984534927737425e-06,
+      "loss": 1.2669,
+      "output_text_mask_loss": 0.8179,
+      "predictor_learning_rate": 0.0004998453492773742,
+      "step": 399
+    },
+    {
+      "epoch": 0.04,
+      "gumbel_tau": 0.9102,
+      "image_masks_loss": 0.0319,
+      "llm_learning_rate": 4.9984241788337994e-06,
+      "loss": 1.1858,
+      "output_text_mask_loss": 0.3052,
+      "predictor_learning_rate": 0.0004998424178833799,
+      "step": 400
+    },
+    {
+      "epoch": 0.04,
+      "gumbel_tau": 0.9099,
+      "image_masks_loss": 0.0228,
+      "llm_learning_rate": 4.998394589761821e-06,
+      "loss": 1.1818,
+      "output_text_mask_loss": 0.3937,
+      "predictor_learning_rate": 0.000499839458976182,
+      "step": 401
+    },
+    {
+      "epoch": 0.04,
+      "gumbel_tau": 0.9097,
+      "image_masks_loss": 0.0234,
+      "llm_learning_rate": 4.998364725561065e-06,
+      "loss": 0.878,
+      "output_text_mask_loss": 0.0618,
+      "predictor_learning_rate": 0.0004998364725561065,
+      "step": 402
+    },
+    {
+      "epoch": 0.04,
+      "gumbel_tau": 0.9095,
+      "image_masks_loss": 0.0191,
+      "llm_learning_rate": 4.9983345862348195e-06,
+      "loss": 0.792,
+      "output_text_mask_loss": 0.1411,
+      "predictor_learning_rate": 0.000499833458623482,
+      "step": 403
+    },
+    {
+      "epoch": 0.04,
+      "gumbel_tau": 0.9093,
+      "image_masks_loss": 0.0431,
+      "llm_learning_rate": 4.998304171786406e-06,
+      "loss": 0.8618,
+      "output_text_mask_loss": 0.0039,
+      "predictor_learning_rate": 0.0004998304171786405,
+      "step": 404
+    },
+    {
+      "epoch": 0.04,
+      "gumbel_tau": 0.9091,
+      "image_masks_loss": 0.0139,
+      "llm_learning_rate": 4.998273482219172e-06,
+      "loss": 0.854,
+      "output_text_mask_loss": 0.0119,
+      "predictor_learning_rate": 0.0004998273482219171,
+      "step": 405
+    },
+    {
+      "epoch": 0.04,
+      "gumbel_tau": 0.9089,
+      "image_masks_loss": 0.0355,
+      "llm_learning_rate": 4.998242517536497e-06,
+      "loss": 0.9303,
+      "output_text_mask_loss": 0.1022,
+      "predictor_learning_rate": 0.0004998242517536497,
+      "step": 406
+    },
+    {
+      "epoch": 0.04,
+      "gumbel_tau": 0.9087,
+      "image_masks_loss": 0.0197,
+      "llm_learning_rate": 4.998211277741792e-06,
+      "loss": 0.8426,
+      "output_text_mask_loss": 0.0942,
+      "predictor_learning_rate": 0.000499821127774179,
+      "step": 407
+    },
+    {
+      "epoch": 0.04,
+      "gumbel_tau": 0.9084,
+      "image_masks_loss": 0.0343,
+      "llm_learning_rate": 4.998179762838496e-06,
+      "loss": 0.7287,
+      "output_text_mask_loss": 0.0549,
+      "predictor_learning_rate": 0.0004998179762838495,
+      "step": 408
+    },
+    {
+      "epoch": 0.04,
+      "gumbel_tau": 0.9082,
+      "image_masks_loss": 0.0294,
+      "llm_learning_rate": 4.99814797283008e-06,
+      "loss": 0.7631,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.000499814797283008,
+      "step": 409
+    },
+    {
+      "epoch": 0.04,
+      "gumbel_tau": 0.908,
+      "image_masks_loss": 0.0219,
+      "llm_learning_rate": 4.998115907720045e-06,
+      "loss": 0.845,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.0004998115907720044,
+      "step": 410
+    },
+    {
+      "epoch": 0.04,
+      "gumbel_tau": 0.9078,
+      "image_masks_loss": 0.0185,
+      "llm_learning_rate": 4.998083567511922e-06,
+      "loss": 0.9346,
+      "output_text_mask_loss": 0.0277,
+      "predictor_learning_rate": 0.0004998083567511922,
+      "step": 411
+    },
+    {
+      "epoch": 0.04,
+      "gumbel_tau": 0.9076,
+      "image_masks_loss": 0.0106,
+      "llm_learning_rate": 4.998050952209273e-06,
+      "loss": 0.7939,
+      "output_text_mask_loss": 0.0429,
+      "predictor_learning_rate": 0.0004998050952209273,
+      "step": 412
+    },
+    {
+      "epoch": 0.04,
+      "gumbel_tau": 0.9074,
+      "image_masks_loss": 0.0286,
+      "llm_learning_rate": 4.998018061815687e-06,
+      "loss": 0.7401,
+      "output_text_mask_loss": 0.0397,
+      "predictor_learning_rate": 0.0004998018061815687,
+      "step": 413
+    },
+    {
+      "epoch": 0.04,
+      "gumbel_tau": 0.9072,
+      "image_masks_loss": 0.0119,
+      "llm_learning_rate": 4.99798489633479e-06,
+      "loss": 0.8074,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.0004997984896334789,
+      "step": 414
+    },
+    {
+      "epoch": 0.04,
+      "gumbel_tau": 0.9069,
+      "image_masks_loss": 0.0138,
+      "llm_learning_rate": 4.997951455770231e-06,
+      "loss": 0.8602,
+      "output_text_mask_loss": 0.1419,
+      "predictor_learning_rate": 0.000499795145577023,
+      "step": 415
+    },
+    {
+      "epoch": 0.04,
+      "gumbel_tau": 0.9067,
+      "image_masks_loss": 0.0273,
+      "llm_learning_rate": 4.997917740125694e-06,
+      "loss": 0.8683,
+      "output_text_mask_loss": 0.0618,
+      "predictor_learning_rate": 0.0004997917740125694,
+      "step": 416
+    },
+    {
+      "epoch": 0.04,
+      "gumbel_tau": 0.9065,
+      "image_masks_loss": 0.0246,
+      "llm_learning_rate": 4.997883749404891e-06,
+      "loss": 0.863,
+      "output_text_mask_loss": 0.0778,
+      "predictor_learning_rate": 0.0004997883749404892,
+      "step": 417
+    },
+    {
+      "epoch": 0.04,
+      "gumbel_tau": 0.9063,
+      "image_masks_loss": 0.0225,
+      "llm_learning_rate": 4.997849483611566e-06,
+      "loss": 0.9017,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004997849483611566,
+      "step": 418
+    },
+    {
+      "epoch": 0.04,
+      "gumbel_tau": 0.9061,
+      "image_masks_loss": 0.0148,
+      "llm_learning_rate": 4.997814942749493e-06,
+      "loss": 0.8212,
+      "output_text_mask_loss": 0.3387,
+      "predictor_learning_rate": 0.0004997814942749492,
+      "step": 419
+    },
+    {
+      "epoch": 0.04,
+      "gumbel_tau": 0.9059,
+      "image_masks_loss": 0.0128,
+      "llm_learning_rate": 4.997780126822475e-06,
+      "loss": 0.946,
+      "output_text_mask_loss": 0.2975,
+      "predictor_learning_rate": 0.0004997780126822474,
+      "step": 420
+    },
+    {
+      "epoch": 0.04,
+      "gumbel_tau": 0.9057,
+      "image_masks_loss": 0.0355,
+      "llm_learning_rate": 4.997745035834344e-06,
+      "loss": 0.8377,
+      "output_text_mask_loss": 0.021,
+      "predictor_learning_rate": 0.0004997745035834344,
+      "step": 421
+    },
+    {
+      "epoch": 0.04,
+      "gumbel_tau": 0.9055,
+      "image_masks_loss": 0.0446,
+      "llm_learning_rate": 4.997709669788968e-06,
+      "loss": 0.7924,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004997709669788968,
+      "step": 422
+    },
+    {
+      "epoch": 0.04,
+      "gumbel_tau": 0.9052,
+      "image_masks_loss": 0.0395,
+      "llm_learning_rate": 4.9976740286902385e-06,
+      "loss": 0.8698,
+      "output_text_mask_loss": 0.058,
+      "predictor_learning_rate": 0.0004997674028690239,
+      "step": 423
+    },
+    {
+      "epoch": 0.04,
+      "gumbel_tau": 0.905,
+      "image_masks_loss": 0.0059,
+      "llm_learning_rate": 4.997638112542083e-06,
+      "loss": 0.8597,
+      "output_text_mask_loss": 0.0231,
+      "predictor_learning_rate": 0.0004997638112542082,
+      "step": 424
+    },
+    {
+      "epoch": 0.04,
+      "gumbel_tau": 0.9048,
+      "image_masks_loss": 0.0151,
+      "llm_learning_rate": 4.997601921348454e-06,
+      "loss": 0.8212,
+      "output_text_mask_loss": 0.116,
+      "predictor_learning_rate": 0.0004997601921348454,
+      "step": 425
+    },
+    {
+      "epoch": 0.04,
+      "gumbel_tau": 0.9046,
+      "image_masks_loss": 0.0175,
+      "llm_learning_rate": 4.997565455113338e-06,
+      "loss": 0.8727,
+      "output_text_mask_loss": 0.0481,
+      "predictor_learning_rate": 0.0004997565455113338,
+      "step": 426
+    },
+    {
+      "epoch": 0.04,
+      "gumbel_tau": 0.9044,
+      "image_masks_loss": 0.0292,
+      "llm_learning_rate": 4.997528713840752e-06,
+      "loss": 0.8893,
+      "output_text_mask_loss": 0.1053,
+      "predictor_learning_rate": 0.0004997528713840751,
+      "step": 427
+    },
+    {
+      "epoch": 0.04,
+      "gumbel_tau": 0.9042,
+      "image_masks_loss": 0.0058,
+      "llm_learning_rate": 4.99749169753474e-06,
+      "loss": 0.8938,
+      "output_text_mask_loss": 0.0204,
+      "predictor_learning_rate": 0.000499749169753474,
+      "step": 428
+    },
+    {
+      "epoch": 0.04,
+      "gumbel_tau": 0.904,
+      "image_masks_loss": 0.0239,
+      "llm_learning_rate": 4.9974544061993805e-06,
+      "loss": 0.9037,
+      "output_text_mask_loss": 0.046,
+      "predictor_learning_rate": 0.000499745440619938,
+      "step": 429
+    },
+    {
+      "epoch": 0.04,
+      "gumbel_tau": 0.9037,
+      "image_masks_loss": 0.0717,
+      "llm_learning_rate": 4.997416839838778e-06,
+      "loss": 0.7443,
+      "output_text_mask_loss": 0.0305,
+      "predictor_learning_rate": 0.0004997416839838777,
+      "step": 430
+    },
+    {
+      "epoch": 0.04,
+      "gumbel_tau": 0.9035,
+      "image_masks_loss": 0.0261,
+      "llm_learning_rate": 4.997378998457071e-06,
+      "loss": 0.9455,
+      "output_text_mask_loss": 0.0149,
+      "predictor_learning_rate": 0.000499737899845707,
+      "step": 431
+    },
+    {
+      "epoch": 0.04,
+      "gumbel_tau": 0.9033,
+      "image_masks_loss": 0.025,
+      "llm_learning_rate": 4.997340882058425e-06,
+      "loss": 0.8482,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 0.0004997340882058425,
+      "step": 432
+    },
+    {
+      "epoch": 0.04,
+      "gumbel_tau": 0.9031,
+      "image_masks_loss": 0.0227,
+      "llm_learning_rate": 4.997302490647039e-06,
+      "loss": 0.8956,
+      "output_text_mask_loss": 0.0076,
+      "predictor_learning_rate": 0.0004997302490647039,
+      "step": 433
+    },
+    {
+      "epoch": 0.04,
+      "gumbel_tau": 0.9029,
+      "image_masks_loss": 0.0244,
+      "llm_learning_rate": 4.997263824227141e-06,
+      "loss": 0.863,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004997263824227141,
+      "step": 434
+    },
+    {
+      "epoch": 0.04,
+      "gumbel_tau": 0.9027,
+      "image_masks_loss": 0.0683,
+      "llm_learning_rate": 4.997224882802987e-06,
+      "loss": 1.0046,
+      "output_text_mask_loss": 0.3357,
+      "predictor_learning_rate": 0.0004997224882802987,
+      "step": 435
+    },
+    {
+      "epoch": 0.04,
+      "gumbel_tau": 0.9025,
+      "image_masks_loss": 0.0108,
+      "llm_learning_rate": 4.997185666378867e-06,
+      "loss": 0.8675,
+      "output_text_mask_loss": 0.0626,
+      "predictor_learning_rate": 0.0004997185666378867,
+      "step": 436
+    },
+    {
+      "epoch": 0.04,
+      "gumbel_tau": 0.9023,
+      "image_masks_loss": 0.042,
+      "llm_learning_rate": 4.9971461749591e-06,
+      "loss": 0.859,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 0.0004997146174959099,
+      "step": 437
+    },
+    {
+      "epoch": 0.04,
+      "gumbel_tau": 0.902,
+      "image_masks_loss": 0.0591,
+      "llm_learning_rate": 4.997106408548032e-06,
+      "loss": 0.9703,
+      "output_text_mask_loss": 0.0496,
+      "predictor_learning_rate": 0.0004997106408548033,
+      "step": 438
+    },
+    {
+      "epoch": 0.04,
+      "gumbel_tau": 0.9018,
+      "image_masks_loss": 0.028,
+      "llm_learning_rate": 4.9970663671500464e-06,
+      "loss": 0.8304,
+      "output_text_mask_loss": 0.1427,
+      "predictor_learning_rate": 0.0004997066367150047,
+      "step": 439
+    },
+    {
+      "epoch": 0.05,
+      "gumbel_tau": 0.9016,
+      "image_masks_loss": 0.0341,
+      "llm_learning_rate": 4.997026050769549e-06,
+      "loss": 0.9341,
+      "output_text_mask_loss": 0.0202,
+      "predictor_learning_rate": 0.0004997026050769549,
+      "step": 440
+    },
+    {
+      "epoch": 0.05,
+      "gumbel_tau": 0.9014,
+      "image_masks_loss": 0.0349,
+      "llm_learning_rate": 4.996985459410982e-06,
+      "loss": 0.8122,
+      "output_text_mask_loss": 0.0322,
+      "predictor_learning_rate": 0.0004996985459410981,
+      "step": 441
+    },
+    {
+      "epoch": 0.05,
+      "gumbel_tau": 0.9012,
+      "image_masks_loss": 0.0614,
+      "llm_learning_rate": 4.996944593078814e-06,
+      "loss": 0.8893,
+      "output_text_mask_loss": 0.0488,
+      "predictor_learning_rate": 0.0004996944593078814,
+      "step": 442
+    },
+    {
+      "epoch": 0.05,
+      "gumbel_tau": 0.901,
+      "image_masks_loss": 0.0156,
+      "llm_learning_rate": 4.996903451777545e-06,
+      "loss": 0.7428,
+      "output_text_mask_loss": 0.0809,
+      "predictor_learning_rate": 0.0004996903451777544,
+      "step": 443
+    },
+    {
+      "epoch": 0.05,
+      "gumbel_tau": 0.9008,
+      "image_masks_loss": 0.028,
+      "llm_learning_rate": 4.996862035511707e-06,
+      "loss": 0.9257,
+      "output_text_mask_loss": 0.1083,
+      "predictor_learning_rate": 0.0004996862035511707,
+      "step": 444
+    },
+    {
+      "epoch": 0.05,
+      "gumbel_tau": 0.9006,
+      "image_masks_loss": 0.05,
+      "llm_learning_rate": 4.9968203442858595e-06,
+      "loss": 0.8302,
+      "output_text_mask_loss": 0.0076,
+      "predictor_learning_rate": 0.000499682034428586,
+      "step": 445
+    },
+    {
+      "epoch": 0.05,
+      "gumbel_tau": 0.9003,
+      "image_masks_loss": 0.0123,
+      "llm_learning_rate": 4.996778378104595e-06,
+      "loss": 0.743,
+      "output_text_mask_loss": 0.0309,
+      "predictor_learning_rate": 0.0004996778378104595,
+      "step": 446
+    },
+    {
+      "epoch": 0.05,
+      "gumbel_tau": 0.9001,
+      "image_masks_loss": 0.0496,
+      "llm_learning_rate": 4.996736136972534e-06,
+      "loss": 0.7406,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.0004996736136972533,
+      "step": 447
+    },
+    {
+      "epoch": 0.05,
+      "gumbel_tau": 0.8999,
+      "image_masks_loss": 0.0359,
+      "llm_learning_rate": 4.996693620894329e-06,
+      "loss": 0.8146,
+      "output_text_mask_loss": 0.0935,
+      "predictor_learning_rate": 0.0004996693620894328,
+      "step": 448
+    },
+    {
+      "epoch": 0.05,
+      "gumbel_tau": 0.8997,
+      "image_masks_loss": 0.0299,
+      "llm_learning_rate": 4.99665082987466e-06,
+      "loss": 0.7452,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.0004996650829874659,
+      "step": 449
+    },
+    {
+      "epoch": 0.05,
+      "gumbel_tau": 0.8995,
+      "image_masks_loss": 0.0267,
+      "llm_learning_rate": 4.9966077639182405e-06,
+      "loss": 0.8395,
+      "output_text_mask_loss": 0.0282,
+      "predictor_learning_rate": 0.0004996607763918241,
+      "step": 450
+    },
+    {
+      "epoch": 0.05,
+      "gumbel_tau": 0.8993,
+      "image_masks_loss": 0.0254,
+      "llm_learning_rate": 4.996564423029814e-06,
+      "loss": 0.7543,
+      "output_text_mask_loss": 0.0313,
+      "predictor_learning_rate": 0.0004996564423029814,
+      "step": 451
+    },
+    {
+      "epoch": 0.05,
+      "gumbel_tau": 0.8991,
+      "image_masks_loss": 0.0237,
+      "llm_learning_rate": 4.996520807214152e-06,
+      "loss": 0.8217,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004996520807214152,
+      "step": 452
+    },
+    {
+      "epoch": 0.05,
+      "gumbel_tau": 0.8989,
+      "image_masks_loss": 0.0523,
+      "llm_learning_rate": 4.996476916476057e-06,
+      "loss": 0.7949,
+      "output_text_mask_loss": 0.0172,
+      "predictor_learning_rate": 0.0004996476916476057,
+      "step": 453
+    },
+    {
+      "epoch": 0.05,
+      "gumbel_tau": 0.8986,
+      "image_masks_loss": 0.0721,
+      "llm_learning_rate": 4.996432750820365e-06,
+      "loss": 0.8756,
+      "output_text_mask_loss": 0.0052,
+      "predictor_learning_rate": 0.0004996432750820365,
+      "step": 454
+    },
+    {
+      "epoch": 0.05,
+      "gumbel_tau": 0.8984,
+      "image_masks_loss": 0.0252,
+      "llm_learning_rate": 4.996388310251937e-06,
+      "loss": 0.76,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.0004996388310251936,
+      "step": 455
+    },
+    {
+      "epoch": 0.05,
+      "gumbel_tau": 0.8982,
+      "image_masks_loss": 0.0199,
+      "llm_learning_rate": 4.996343594775668e-06,
+      "loss": 0.7143,
+      "output_text_mask_loss": 0.1114,
+      "predictor_learning_rate": 0.0004996343594775667,
+      "step": 456
+    },
+    {
+      "epoch": 0.05,
+      "gumbel_tau": 0.898,
+      "image_masks_loss": 0.0267,
+      "llm_learning_rate": 4.996298604396481e-06,
+      "loss": 0.8919,
+      "output_text_mask_loss": 0.0378,
+      "predictor_learning_rate": 0.000499629860439648,
+      "step": 457
+    },
+    {
+      "epoch": 0.05,
+      "gumbel_tau": 0.8978,
+      "image_masks_loss": 0.0239,
+      "llm_learning_rate": 4.996253339119333e-06,
+      "loss": 0.8022,
+      "output_text_mask_loss": 0.0309,
+      "predictor_learning_rate": 0.0004996253339119332,
+      "step": 458
+    },
+    {
+      "epoch": 0.05,
+      "gumbel_tau": 0.8976,
+      "image_masks_loss": 0.0614,
+      "llm_learning_rate": 4.996207798949206e-06,
+      "loss": 0.748,
+      "output_text_mask_loss": 0.0277,
+      "predictor_learning_rate": 0.0004996207798949206,
+      "step": 459
+    },
+    {
+      "epoch": 0.05,
+      "gumbel_tau": 0.8974,
+      "image_masks_loss": 0.0254,
+      "llm_learning_rate": 4.996161983891116e-06,
+      "loss": 0.8239,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004996161983891116,
+      "step": 460
+    },
+    {
+      "epoch": 0.05,
+      "gumbel_tau": 0.8972,
+      "image_masks_loss": 0.02,
+      "llm_learning_rate": 4.996115893950109e-06,
+      "loss": 0.8528,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.0004996115893950108,
+      "step": 461
+    },
+    {
+      "epoch": 0.05,
+      "gumbel_tau": 0.8969,
+      "image_masks_loss": 0.0244,
+      "llm_learning_rate": 4.99606952913126e-06,
+      "loss": 0.9182,
+      "output_text_mask_loss": 0.0969,
+      "predictor_learning_rate": 0.0004996069529131259,
+      "step": 462
+    },
+    {
+      "epoch": 0.05,
+      "gumbel_tau": 0.8967,
+      "image_masks_loss": 0.0275,
+      "llm_learning_rate": 4.996022889439674e-06,
+      "loss": 0.8261,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004996022889439674,
+      "step": 463
+    },
+    {
+      "epoch": 0.05,
+      "gumbel_tau": 0.8965,
+      "image_masks_loss": 0.0319,
+      "llm_learning_rate": 4.995975974880488e-06,
+      "loss": 0.9738,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004995975974880488,
+      "step": 464
+    },
+    {
+      "epoch": 0.05,
+      "gumbel_tau": 0.8963,
+      "image_masks_loss": 0.0546,
+      "llm_learning_rate": 4.995928785458868e-06,
+      "loss": 0.8601,
+      "output_text_mask_loss": 0.1831,
+      "predictor_learning_rate": 0.0004995928785458867,
+      "step": 465
+    },
+    {
+      "epoch": 0.05,
+      "gumbel_tau": 0.8961,
+      "image_masks_loss": 0.0109,
+      "llm_learning_rate": 4.995881321180011e-06,
+      "loss": 0.7991,
+      "output_text_mask_loss": 0.0618,
+      "predictor_learning_rate": 0.000499588132118001,
+      "step": 466
+    },
+    {
+      "epoch": 0.05,
+      "gumbel_tau": 0.8959,
+      "image_masks_loss": 0.0148,
+      "llm_learning_rate": 4.995833582049143e-06,
+      "loss": 0.896,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004995833582049142,
+      "step": 467
+    },
+    {
+      "epoch": 0.05,
+      "gumbel_tau": 0.8957,
+      "image_masks_loss": 0.0265,
+      "llm_learning_rate": 4.995785568071522e-06,
+      "loss": 0.9268,
+      "output_text_mask_loss": 0.0847,
+      "predictor_learning_rate": 0.0004995785568071521,
+      "step": 468
+    },
+    {
+      "epoch": 0.05,
+      "gumbel_tau": 0.8955,
+      "image_masks_loss": 0.0267,
+      "llm_learning_rate": 4.995737279252435e-06,
+      "loss": 0.7723,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004995737279252434,
+      "step": 469
+    },
+    {
+      "epoch": 0.05,
+      "gumbel_tau": 0.8953,
+      "image_masks_loss": 0.025,
+      "llm_learning_rate": 4.9956887155972e-06,
+      "loss": 0.9,
+      "output_text_mask_loss": 0.1045,
+      "predictor_learning_rate": 0.0004995688715597199,
+      "step": 470
+    },
+    {
+      "epoch": 0.05,
+      "gumbel_tau": 0.895,
+      "image_masks_loss": 0.0479,
+      "llm_learning_rate": 4.995639877111164e-06,
+      "loss": 0.8341,
+      "output_text_mask_loss": 0.0275,
+      "predictor_learning_rate": 0.0004995639877111164,
+      "step": 471
+    },
+    {
+      "epoch": 0.05,
+      "gumbel_tau": 0.8948,
+      "image_masks_loss": 0.0175,
+      "llm_learning_rate": 4.995590763799707e-06,
+      "loss": 0.8564,
+      "output_text_mask_loss": 0.0198,
+      "predictor_learning_rate": 0.0004995590763799707,
+      "step": 472
+    },
+    {
+      "epoch": 0.05,
+      "gumbel_tau": 0.8946,
+      "image_masks_loss": 0.05,
+      "llm_learning_rate": 4.995541375668236e-06,
+      "loss": 0.8782,
+      "output_text_mask_loss": 0.0233,
+      "predictor_learning_rate": 0.0004995541375668235,
+      "step": 473
+    },
+    {
+      "epoch": 0.05,
+      "gumbel_tau": 0.8944,
+      "image_masks_loss": 0.0326,
+      "llm_learning_rate": 4.99549171272219e-06,
+      "loss": 0.8456,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.000499549171272219,
+      "step": 474
+    },
+    {
+      "epoch": 0.05,
+      "gumbel_tau": 0.8942,
+      "image_masks_loss": 0.0338,
+      "llm_learning_rate": 4.995441774967038e-06,
+      "loss": 0.8121,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.0004995441774967037,
+      "step": 475
+    },
+    {
+      "epoch": 0.05,
+      "gumbel_tau": 0.894,
+      "image_masks_loss": 0.05,
+      "llm_learning_rate": 4.995391562408279e-06,
+      "loss": 0.732,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.0004995391562408279,
+      "step": 476
+    },
+    {
+      "epoch": 0.05,
+      "gumbel_tau": 0.8938,
+      "image_masks_loss": 0.0309,
+      "llm_learning_rate": 4.995341075051444e-06,
+      "loss": 0.8189,
+      "output_text_mask_loss": 0.0042,
+      "predictor_learning_rate": 0.0004995341075051444,
+      "step": 477
+    },
+    {
+      "epoch": 0.05,
+      "gumbel_tau": 0.8936,
+      "image_masks_loss": 0.0387,
+      "llm_learning_rate": 4.995290312902092e-06,
+      "loss": 0.8401,
+      "output_text_mask_loss": 0.0368,
+      "predictor_learning_rate": 0.0004995290312902091,
+      "step": 478
+    },
+    {
+      "epoch": 0.05,
+      "gumbel_tau": 0.8934,
+      "image_masks_loss": 0.0284,
+      "llm_learning_rate": 4.995239275965811e-06,
+      "loss": 0.8491,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.000499523927596581,
+      "step": 479
+    },
+    {
+      "epoch": 0.05,
+      "gumbel_tau": 0.8931,
+      "image_masks_loss": 0.1175,
+      "llm_learning_rate": 4.995187964248224e-06,
+      "loss": 0.9773,
+      "output_text_mask_loss": 0.0271,
+      "predictor_learning_rate": 0.0004995187964248224,
+      "step": 480
+    },
+    {
+      "epoch": 0.05,
+      "gumbel_tau": 0.8929,
+      "image_masks_loss": 0.0347,
+      "llm_learning_rate": 4.99513637775498e-06,
+      "loss": 0.906,
+      "output_text_mask_loss": 0.0744,
+      "predictor_learning_rate": 0.0004995136377754979,
+      "step": 481
+    },
+    {
+      "epoch": 0.05,
+      "gumbel_tau": 0.8927,
+      "image_masks_loss": 0.0263,
+      "llm_learning_rate": 4.99508451649176e-06,
+      "loss": 0.8441,
+      "output_text_mask_loss": 0.0198,
+      "predictor_learning_rate": 0.000499508451649176,
+      "step": 482
+    },
+    {
+      "epoch": 0.05,
+      "gumbel_tau": 0.8925,
+      "image_masks_loss": 0.0109,
+      "llm_learning_rate": 4.995032380464276e-06,
+      "loss": 0.8526,
+      "output_text_mask_loss": 0.0219,
+      "predictor_learning_rate": 0.0004995032380464276,
+      "step": 483
+    },
+    {
+      "epoch": 0.05,
+      "gumbel_tau": 0.8923,
+      "image_masks_loss": 0.0166,
+      "llm_learning_rate": 4.994979969678269e-06,
+      "loss": 0.8167,
+      "output_text_mask_loss": 0.0156,
+      "predictor_learning_rate": 0.000499497996967827,
+      "step": 484
+    },
+    {
+      "epoch": 0.05,
+      "gumbel_tau": 0.8921,
+      "image_masks_loss": 0.0062,
+      "llm_learning_rate": 4.994927284139511e-06,
+      "loss": 0.9307,
+      "output_text_mask_loss": 0.1434,
+      "predictor_learning_rate": 0.000499492728413951,
+      "step": 485
+    },
+    {
+      "epoch": 0.05,
+      "gumbel_tau": 0.8919,
+      "image_masks_loss": 0.015,
+      "llm_learning_rate": 4.994874323853802e-06,
+      "loss": 0.8253,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 0.0004994874323853802,
+      "step": 486
+    },
+    {
+      "epoch": 0.05,
+      "gumbel_tau": 0.8917,
+      "image_masks_loss": 0.0391,
+      "llm_learning_rate": 4.994821088826976e-06,
+      "loss": 0.7826,
+      "output_text_mask_loss": 0.1083,
+      "predictor_learning_rate": 0.0004994821088826977,
+      "step": 487
+    },
+    {
+      "epoch": 0.05,
+      "gumbel_tau": 0.8915,
+      "image_masks_loss": 0.0164,
+      "llm_learning_rate": 4.994767579064895e-06,
+      "loss": 0.9418,
+      "output_text_mask_loss": 0.0374,
+      "predictor_learning_rate": 0.0004994767579064894,
+      "step": 488
+    },
+    {
+      "epoch": 0.05,
+      "gumbel_tau": 0.8913,
+      "image_masks_loss": 0.0364,
+      "llm_learning_rate": 4.994713794573451e-06,
+      "loss": 0.8053,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.000499471379457345,
+      "step": 489
+    },
+    {
+      "epoch": 0.05,
+      "gumbel_tau": 0.891,
+      "image_masks_loss": 0.0067,
+      "llm_learning_rate": 4.994659735358567e-06,
+      "loss": 0.8953,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004994659735358567,
+      "step": 490
+    },
+    {
+      "epoch": 0.05,
+      "gumbel_tau": 0.8908,
+      "image_masks_loss": 0.0157,
+      "llm_learning_rate": 4.9946054014261965e-06,
+      "loss": 0.8135,
+      "output_text_mask_loss": 0.0324,
+      "predictor_learning_rate": 0.0004994605401426196,
+      "step": 491
+    },
+    {
+      "epoch": 0.05,
+      "gumbel_tau": 0.8906,
+      "image_masks_loss": 0.0515,
+      "llm_learning_rate": 4.994550792782322e-06,
+      "loss": 0.7509,
+      "output_text_mask_loss": 0.0557,
+      "predictor_learning_rate": 0.0004994550792782322,
+      "step": 492
+    },
+    {
+      "epoch": 0.05,
+      "gumbel_tau": 0.8904,
+      "image_masks_loss": 0.0301,
+      "llm_learning_rate": 4.994495909432959e-06,
+      "loss": 0.7996,
+      "output_text_mask_loss": 0.0938,
+      "predictor_learning_rate": 0.0004994495909432958,
+      "step": 493
+    },
+    {
+      "epoch": 0.05,
+      "gumbel_tau": 0.8902,
+      "image_masks_loss": 0.0277,
+      "llm_learning_rate": 4.994440751384149e-06,
+      "loss": 0.7926,
+      "output_text_mask_loss": 0.0504,
+      "predictor_learning_rate": 0.0004994440751384148,
+      "step": 494
+    },
+    {
+      "epoch": 0.05,
+      "gumbel_tau": 0.89,
+      "image_masks_loss": 0.0103,
+      "llm_learning_rate": 4.994385318641968e-06,
+      "loss": 0.9132,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004994385318641968,
+      "step": 495
+    },
+    {
+      "epoch": 0.05,
+      "gumbel_tau": 0.8898,
+      "image_masks_loss": 0.0223,
+      "llm_learning_rate": 4.9943296112125194e-06,
+      "loss": 0.7197,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004994329611212519,
+      "step": 496
+    },
+    {
+      "epoch": 0.05,
+      "gumbel_tau": 0.8896,
+      "image_masks_loss": 0.0196,
+      "llm_learning_rate": 4.994273629101939e-06,
+      "loss": 0.8335,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004994273629101939,
+      "step": 497
+    },
+    {
+      "epoch": 0.05,
+      "gumbel_tau": 0.8894,
+      "image_masks_loss": 0.0166,
+      "llm_learning_rate": 4.99421737231639e-06,
+      "loss": 0.8688,
+      "output_text_mask_loss": 0.1755,
+      "predictor_learning_rate": 0.000499421737231639,
+      "step": 498
+    },
+    {
+      "epoch": 0.05,
+      "gumbel_tau": 0.8892,
+      "image_masks_loss": 0.0115,
+      "llm_learning_rate": 4.994160840862068e-06,
+      "loss": 0.7343,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004994160840862068,
+      "step": 499
+    },
+    {
+      "epoch": 0.05,
+      "gumbel_tau": 0.8889,
+      "image_masks_loss": 0.0162,
+      "llm_learning_rate": 4.9941040347451995e-06,
+      "loss": 0.8898,
+      "output_text_mask_loss": 0.0858,
+      "predictor_learning_rate": 0.0004994104034745199,
+      "step": 500
+    },
+    {
+      "epoch": 0.05,
+      "gumbel_tau": 0.8887,
+      "image_masks_loss": 0.0402,
+      "llm_learning_rate": 4.994046953972038e-06,
+      "loss": 0.8379,
+      "output_text_mask_loss": 0.0172,
+      "predictor_learning_rate": 0.0004994046953972037,
+      "step": 501
+    },
+    {
+      "epoch": 0.05,
+      "gumbel_tau": 0.8885,
+      "image_masks_loss": 0.0374,
+      "llm_learning_rate": 4.993989598548872e-06,
+      "loss": 0.7704,
+      "output_text_mask_loss": 0.0383,
+      "predictor_learning_rate": 0.0004993989598548872,
+      "step": 502
+    },
+    {
+      "epoch": 0.05,
+      "gumbel_tau": 0.8883,
+      "image_masks_loss": 0.0083,
+      "llm_learning_rate": 4.993931968482016e-06,
+      "loss": 0.785,
+      "output_text_mask_loss": 0.0086,
+      "predictor_learning_rate": 0.0004993931968482015,
+      "step": 503
+    },
+    {
+      "epoch": 0.05,
+      "gumbel_tau": 0.8881,
+      "image_masks_loss": 0.0751,
+      "llm_learning_rate": 4.993874063777816e-06,
+      "loss": 0.8759,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004993874063777816,
+      "step": 504
+    },
+    {
+      "epoch": 0.05,
+      "gumbel_tau": 0.8879,
+      "image_masks_loss": 0.0372,
+      "llm_learning_rate": 4.99381588444265e-06,
+      "loss": 0.8498,
+      "output_text_mask_loss": 0.0324,
+      "predictor_learning_rate": 0.0004993815884442649,
+      "step": 505
+    },
+    {
+      "epoch": 0.05,
+      "gumbel_tau": 0.8877,
+      "image_masks_loss": 0.0393,
+      "llm_learning_rate": 4.9937574304829234e-06,
+      "loss": 0.8549,
+      "output_text_mask_loss": 0.0221,
+      "predictor_learning_rate": 0.0004993757430482923,
+      "step": 506
+    },
+    {
+      "epoch": 0.05,
+      "gumbel_tau": 0.8875,
+      "image_masks_loss": 0.0877,
+      "llm_learning_rate": 4.993698701905075e-06,
+      "loss": 0.8481,
+      "output_text_mask_loss": 0.0374,
+      "predictor_learning_rate": 0.0004993698701905074,
+      "step": 507
+    },
+    {
+      "epoch": 0.05,
+      "gumbel_tau": 0.8873,
+      "image_masks_loss": 0.0131,
+      "llm_learning_rate": 4.993639698715569e-06,
+      "loss": 0.8164,
+      "output_text_mask_loss": 0.0299,
+      "predictor_learning_rate": 0.0004993639698715568,
+      "step": 508
+    },
+    {
+      "epoch": 0.05,
+      "gumbel_tau": 0.8871,
+      "image_masks_loss": 0.0801,
+      "llm_learning_rate": 4.9935804209209066e-06,
+      "loss": 0.8121,
+      "output_text_mask_loss": 0.0488,
+      "predictor_learning_rate": 0.0004993580420920906,
+      "step": 509
+    },
+    {
+      "epoch": 0.05,
+      "gumbel_tau": 0.8868,
+      "image_masks_loss": 0.0523,
+      "llm_learning_rate": 4.993520868527614e-06,
+      "loss": 0.815,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.0004993520868527613,
+      "step": 510
+    },
+    {
+      "epoch": 0.05,
+      "gumbel_tau": 0.8866,
+      "image_masks_loss": 0.0511,
+      "llm_learning_rate": 4.993461041542248e-06,
+      "loss": 0.7546,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 0.0004993461041542248,
+      "step": 511
+    },
+    {
+      "epoch": 0.05,
+      "gumbel_tau": 0.8864,
+      "image_masks_loss": 0.0284,
+      "llm_learning_rate": 4.9934009399714e-06,
+      "loss": 0.8688,
+      "output_text_mask_loss": 0.0809,
+      "predictor_learning_rate": 0.00049934009399714,
+      "step": 512
+    },
+    {
+      "epoch": 0.05,
+      "gumbel_tau": 0.8862,
+      "image_masks_loss": 0.0957,
+      "llm_learning_rate": 4.993340563821685e-06,
+      "loss": 0.8074,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004993340563821684,
+      "step": 513
+    },
+    {
+      "epoch": 0.05,
+      "gumbel_tau": 0.886,
+      "image_masks_loss": 0.087,
+      "llm_learning_rate": 4.993279913099755e-06,
+      "loss": 0.9126,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004993279913099754,
+      "step": 514
+    },
+    {
+      "epoch": 0.05,
+      "gumbel_tau": 0.8858,
+      "image_masks_loss": 0.0215,
+      "llm_learning_rate": 4.993218987812286e-06,
+      "loss": 0.9533,
+      "output_text_mask_loss": 0.037,
+      "predictor_learning_rate": 0.0004993218987812285,
+      "step": 515
+    },
+    {
+      "epoch": 0.05,
+      "gumbel_tau": 0.8856,
+      "image_masks_loss": 0.0393,
+      "llm_learning_rate": 4.993157787965989e-06,
+      "loss": 0.8763,
+      "output_text_mask_loss": 0.0368,
+      "predictor_learning_rate": 0.0004993157787965989,
+      "step": 516
+    },
+    {
+      "epoch": 0.05,
+      "gumbel_tau": 0.8854,
+      "image_masks_loss": 0.1045,
+      "llm_learning_rate": 4.993096313567603e-06,
+      "loss": 0.867,
+      "output_text_mask_loss": 0.0992,
+      "predictor_learning_rate": 0.0004993096313567602,
+      "step": 517
+    },
+    {
+      "epoch": 0.05,
+      "gumbel_tau": 0.8852,
+      "image_masks_loss": 0.0402,
+      "llm_learning_rate": 4.993034564623898e-06,
+      "loss": 0.8381,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004993034564623898,
+      "step": 518
+    },
+    {
+      "epoch": 0.05,
+      "gumbel_tau": 0.885,
+      "image_masks_loss": 0.0197,
+      "llm_learning_rate": 4.9929725411416744e-06,
+      "loss": 0.7604,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004992972541141674,
+      "step": 519
+    },
+    {
+      "epoch": 0.05,
+      "gumbel_tau": 0.8848,
+      "image_masks_loss": 0.0231,
+      "llm_learning_rate": 4.992910243127761e-06,
+      "loss": 0.8121,
+      "output_text_mask_loss": 0.1083,
+      "predictor_learning_rate": 0.0004992910243127761,
+      "step": 520
+    },
+    {
+      "epoch": 0.05,
+      "gumbel_tau": 0.8845,
+      "image_masks_loss": 0.1381,
+      "llm_learning_rate": 4.9928476705890184e-06,
+      "loss": 1.1425,
+      "output_text_mask_loss": 0.119,
+      "predictor_learning_rate": 0.0004992847670589018,
+      "step": 521
+    },
+    {
+      "epoch": 0.05,
+      "gumbel_tau": 0.8843,
+      "image_masks_loss": 0.0081,
+      "llm_learning_rate": 4.992784823532339e-06,
+      "loss": 1.3797,
+      "output_text_mask_loss": 0.4395,
+      "predictor_learning_rate": 0.0004992784823532338,
+      "step": 522
+    },
+    {
+      "epoch": 0.05,
+      "gumbel_tau": 0.8841,
+      "image_masks_loss": 0.0197,
+      "llm_learning_rate": 4.992721701964641e-06,
+      "loss": 1.1055,
+      "output_text_mask_loss": 0.4578,
+      "predictor_learning_rate": 0.0004992721701964641,
+      "step": 523
+    },
+    {
+      "epoch": 0.05,
+      "gumbel_tau": 0.8839,
+      "image_masks_loss": 0.0301,
+      "llm_learning_rate": 4.992658305892878e-06,
+      "loss": 1.8041,
+      "output_text_mask_loss": 2.1973,
+      "predictor_learning_rate": 0.0004992658305892877,
+      "step": 524
+    },
+    {
+      "epoch": 0.05,
+      "gumbel_tau": 0.8837,
+      "image_masks_loss": 0.0486,
+      "llm_learning_rate": 4.9925946353240295e-06,
+      "loss": 2.3226,
+      "output_text_mask_loss": 2.2583,
+      "predictor_learning_rate": 0.0004992594635324029,
+      "step": 525
+    },
+    {
+      "epoch": 0.05,
+      "gumbel_tau": 0.8835,
+      "image_masks_loss": 0.0383,
+      "llm_learning_rate": 4.992530690265107e-06,
+      "loss": 1.7845,
+      "output_text_mask_loss": 2.3804,
+      "predictor_learning_rate": 0.0004992530690265107,
+      "step": 526
+    },
+    {
+      "epoch": 0.05,
+      "gumbel_tau": 0.8833,
+      "image_masks_loss": 0.0299,
+      "llm_learning_rate": 4.992466470723154e-06,
+      "loss": 1.7782,
+      "output_text_mask_loss": 0.7812,
+      "predictor_learning_rate": 0.0004992466470723154,
+      "step": 527
+    },
+    {
+      "epoch": 0.05,
+      "gumbel_tau": 0.8831,
+      "image_masks_loss": 0.0492,
+      "llm_learning_rate": 4.9924019767052416e-06,
+      "loss": 2.1135,
+      "output_text_mask_loss": 0.6195,
+      "predictor_learning_rate": 0.0004992401976705241,
+      "step": 528
+    },
+    {
+      "epoch": 0.05,
+      "gumbel_tau": 0.8829,
+      "image_masks_loss": 0.0671,
+      "llm_learning_rate": 4.992337208218472e-06,
+      "loss": 2.2303,
+      "output_text_mask_loss": 2.5879,
+      "predictor_learning_rate": 0.0004992337208218472,
+      "step": 529
+    },
+    {
+      "epoch": 0.05,
+      "gumbel_tau": 0.8827,
+      "image_masks_loss": 0.0286,
+      "llm_learning_rate": 4.992272165269977e-06,
+      "loss": 2.2591,
+      "output_text_mask_loss": 0.4028,
+      "predictor_learning_rate": 0.0004992272165269977,
+      "step": 530
+    },
+    {
+      "epoch": 0.05,
+      "gumbel_tau": 0.8825,
+      "image_masks_loss": 0.0121,
+      "llm_learning_rate": 4.992206847866921e-06,
+      "loss": 1.4675,
+      "output_text_mask_loss": 0.2747,
+      "predictor_learning_rate": 0.0004992206847866921,
+      "step": 531
+    },
+    {
+      "epoch": 0.05,
+      "gumbel_tau": 0.8823,
+      "image_masks_loss": 0.0215,
+      "llm_learning_rate": 4.992141256016494e-06,
+      "loss": 1.2956,
+      "output_text_mask_loss": 0.6165,
+      "predictor_learning_rate": 0.0004992141256016494,
+      "step": 532
+    },
+    {
+      "epoch": 0.05,
+      "gumbel_tau": 0.882,
+      "image_masks_loss": 0.0301,
+      "llm_learning_rate": 4.992075389725922e-06,
+      "loss": 1.0866,
+      "output_text_mask_loss": 0.264,
+      "predictor_learning_rate": 0.0004992075389725922,
+      "step": 533
+    },
+    {
+      "epoch": 0.05,
+      "gumbel_tau": 0.8818,
+      "image_masks_loss": 0.0294,
+      "llm_learning_rate": 4.992009249002456e-06,
+      "loss": 1.0492,
+      "output_text_mask_loss": 0.3113,
+      "predictor_learning_rate": 0.0004992009249002456,
+      "step": 534
+    },
+    {
+      "epoch": 0.05,
+      "gumbel_tau": 0.8816,
+      "image_masks_loss": 0.032,
+      "llm_learning_rate": 4.991942833853383e-06,
+      "loss": 0.6835,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.0004991942833853382,
+      "step": 535
+    },
+    {
+      "epoch": 0.05,
+      "gumbel_tau": 0.8814,
+      "image_masks_loss": 0.0341,
+      "llm_learning_rate": 4.991876144286013e-06,
+      "loss": 0.7909,
+      "output_text_mask_loss": 0.0141,
+      "predictor_learning_rate": 0.0004991876144286012,
+      "step": 536
+    },
+    {
+      "epoch": 0.06,
+      "gumbel_tau": 0.8812,
+      "image_masks_loss": 0.045,
+      "llm_learning_rate": 4.9918091803076916e-06,
+      "loss": 0.8748,
+      "output_text_mask_loss": 0.069,
+      "predictor_learning_rate": 0.0004991809180307691,
+      "step": 537
+    },
+    {
+      "epoch": 0.06,
+      "gumbel_tau": 0.881,
+      "image_masks_loss": 0.0178,
+      "llm_learning_rate": 4.991741941925793e-06,
+      "loss": 0.9941,
+      "output_text_mask_loss": 0.0322,
+      "predictor_learning_rate": 0.0004991741941925792,
+      "step": 538
+    },
+    {
+      "epoch": 0.06,
+      "gumbel_tau": 0.8808,
+      "image_masks_loss": 0.0235,
+      "llm_learning_rate": 4.991674429147721e-06,
+      "loss": 1.0253,
+      "output_text_mask_loss": 0.19,
+      "predictor_learning_rate": 0.0004991674429147721,
+      "step": 539
+    },
+    {
+      "epoch": 0.06,
+      "gumbel_tau": 0.8806,
+      "image_masks_loss": 0.0154,
+      "llm_learning_rate": 4.991606641980912e-06,
+      "loss": 1.0113,
+      "output_text_mask_loss": 0.1694,
+      "predictor_learning_rate": 0.0004991606641980912,
+      "step": 540
+    },
+    {
+      "epoch": 0.06,
+      "gumbel_tau": 0.8804,
+      "image_masks_loss": 0.0286,
+      "llm_learning_rate": 4.99153858043283e-06,
+      "loss": 0.679,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004991538580432829,
+      "step": 541
+    },
+    {
+      "epoch": 0.06,
+      "gumbel_tau": 0.8802,
+      "image_masks_loss": 0.0246,
+      "llm_learning_rate": 4.991470244510969e-06,
+      "loss": 0.9802,
+      "output_text_mask_loss": 0.0763,
+      "predictor_learning_rate": 0.0004991470244510969,
+      "step": 542
+    },
+    {
+      "epoch": 0.06,
+      "gumbel_tau": 0.88,
+      "image_masks_loss": 0.0549,
+      "llm_learning_rate": 4.991401634222855e-06,
+      "loss": 0.8949,
+      "output_text_mask_loss": 0.0101,
+      "predictor_learning_rate": 0.0004991401634222855,
+      "step": 543
+    },
+    {
+      "epoch": 0.06,
+      "gumbel_tau": 0.8798,
+      "image_masks_loss": 0.0534,
+      "llm_learning_rate": 4.991332749576044e-06,
+      "loss": 0.8833,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004991332749576044,
+      "step": 544
+    },
+    {
+      "epoch": 0.06,
+      "gumbel_tau": 0.8796,
+      "image_masks_loss": 0.0439,
+      "llm_learning_rate": 4.991263590578121e-06,
+      "loss": 0.8497,
+      "output_text_mask_loss": 0.0313,
+      "predictor_learning_rate": 0.0004991263590578121,
+      "step": 545
+    },
+    {
+      "epoch": 0.06,
+      "gumbel_tau": 0.8793,
+      "image_masks_loss": 0.0702,
+      "llm_learning_rate": 4.991194157236704e-06,
+      "loss": 0.9433,
+      "output_text_mask_loss": 0.0152,
+      "predictor_learning_rate": 0.0004991194157236703,
+      "step": 546
+    },
+    {
+      "epoch": 0.06,
+      "gumbel_tau": 0.8791,
+      "image_masks_loss": 0.0334,
+      "llm_learning_rate": 4.991124449559436e-06,
+      "loss": 0.8893,
+      "output_text_mask_loss": 0.066,
+      "predictor_learning_rate": 0.0004991124449559435,
+      "step": 547
+    },
+    {
+      "epoch": 0.06,
+      "gumbel_tau": 0.8789,
+      "image_masks_loss": 0.0149,
+      "llm_learning_rate": 4.991054467553995e-06,
+      "loss": 0.7129,
+      "output_text_mask_loss": 0.1114,
+      "predictor_learning_rate": 0.0004991054467553995,
+      "step": 548
+    },
+    {
+      "epoch": 0.06,
+      "gumbel_tau": 0.8787,
+      "image_masks_loss": 0.0243,
+      "llm_learning_rate": 4.990984211228087e-06,
+      "loss": 0.7975,
+      "output_text_mask_loss": 0.106,
+      "predictor_learning_rate": 0.0004990984211228087,
+      "step": 549
+    },
+    {
+      "epoch": 0.06,
+      "gumbel_tau": 0.8785,
+      "image_masks_loss": 0.0184,
+      "llm_learning_rate": 4.99091368058945e-06,
+      "loss": 0.8987,
+      "output_text_mask_loss": 0.0923,
+      "predictor_learning_rate": 0.000499091368058945,
+      "step": 550
+    },
+    {
+      "epoch": 0.06,
+      "gumbel_tau": 0.8783,
+      "image_masks_loss": 0.0483,
+      "llm_learning_rate": 4.99084287564585e-06,
+      "loss": 0.834,
+      "output_text_mask_loss": 0.0322,
+      "predictor_learning_rate": 0.000499084287564585,
+      "step": 551
+    },
+    {
+      "epoch": 0.06,
+      "gumbel_tau": 0.8781,
+      "image_masks_loss": 0.0029,
+      "llm_learning_rate": 4.9907717964050845e-06,
+      "loss": 0.8924,
+      "output_text_mask_loss": 0.1831,
+      "predictor_learning_rate": 0.0004990771796405084,
+      "step": 552
+    },
+    {
+      "epoch": 0.06,
+      "gumbel_tau": 0.8779,
+      "image_masks_loss": 0.0182,
+      "llm_learning_rate": 4.990700442874981e-06,
+      "loss": 1.0774,
+      "output_text_mask_loss": 0.2625,
+      "predictor_learning_rate": 0.000499070044287498,
+      "step": 553
+    },
+    {
+      "epoch": 0.06,
+      "gumbel_tau": 0.8777,
+      "image_masks_loss": 0.0136,
+      "llm_learning_rate": 4.990628815063396e-06,
+      "loss": 0.9317,
+      "output_text_mask_loss": 0.0309,
+      "predictor_learning_rate": 0.0004990628815063396,
+      "step": 554
+    },
+    {
+      "epoch": 0.06,
+      "gumbel_tau": 0.8775,
+      "image_masks_loss": 0.0347,
+      "llm_learning_rate": 4.9905569129782185e-06,
+      "loss": 1.0035,
+      "output_text_mask_loss": 0.2487,
+      "predictor_learning_rate": 0.0004990556912978218,
+      "step": 555
+    },
+    {
+      "epoch": 0.06,
+      "gumbel_tau": 0.8773,
+      "image_masks_loss": 0.0176,
+      "llm_learning_rate": 4.990484736627366e-06,
+      "loss": 0.9649,
+      "output_text_mask_loss": 0.0687,
+      "predictor_learning_rate": 0.0004990484736627366,
+      "step": 556
+    },
+    {
+      "epoch": 0.06,
+      "gumbel_tau": 0.8771,
+      "image_masks_loss": 0.0154,
+      "llm_learning_rate": 4.9904122860187885e-06,
+      "loss": 0.7811,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004990412286018788,
+      "step": 557
+    },
+    {
+      "epoch": 0.06,
+      "gumbel_tau": 0.8769,
+      "image_masks_loss": 0.0242,
+      "llm_learning_rate": 4.990339561160461e-06,
+      "loss": 0.9114,
+      "output_text_mask_loss": 0.1198,
+      "predictor_learning_rate": 0.0004990339561160461,
+      "step": 558
+    },
+    {
+      "epoch": 0.06,
+      "gumbel_tau": 0.8767,
+      "image_masks_loss": 0.0189,
+      "llm_learning_rate": 4.990266562060395e-06,
+      "loss": 0.8579,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004990266562060394,
+      "step": 559
+    },
+    {
+      "epoch": 0.06,
+      "gumbel_tau": 0.8764,
+      "image_masks_loss": 0.0256,
+      "llm_learning_rate": 4.990193288726629e-06,
+      "loss": 0.8904,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 0.0004990193288726629,
+      "step": 560
+    },
+    {
+      "epoch": 0.06,
+      "gumbel_tau": 0.8762,
+      "image_masks_loss": 0.0378,
+      "llm_learning_rate": 4.99011974116723e-06,
+      "loss": 0.9644,
+      "output_text_mask_loss": 0.119,
+      "predictor_learning_rate": 0.000499011974116723,
+      "step": 561
+    },
+    {
+      "epoch": 0.06,
+      "gumbel_tau": 0.876,
+      "image_masks_loss": 0.0134,
+      "llm_learning_rate": 4.990045919390299e-06,
+      "loss": 0.8718,
+      "output_text_mask_loss": 0.1266,
+      "predictor_learning_rate": 0.00049900459193903,
+      "step": 562
+    },
+    {
+      "epoch": 0.06,
+      "gumbel_tau": 0.8758,
+      "image_masks_loss": 0.0144,
+      "llm_learning_rate": 4.989971823403966e-06,
+      "loss": 0.9169,
+      "output_text_mask_loss": 0.0519,
+      "predictor_learning_rate": 0.0004989971823403965,
+      "step": 563
+    },
+    {
+      "epoch": 0.06,
+      "gumbel_tau": 0.8756,
+      "image_masks_loss": 0.0336,
+      "llm_learning_rate": 4.989897453216388e-06,
+      "loss": 0.8,
+      "output_text_mask_loss": 0.0012,
+      "predictor_learning_rate": 0.0004989897453216389,
+      "step": 564
+    },
+    {
+      "epoch": 0.06,
+      "gumbel_tau": 0.8754,
+      "image_masks_loss": 0.0232,
+      "llm_learning_rate": 4.989822808835758e-06,
+      "loss": 0.7949,
+      "output_text_mask_loss": 0.0263,
+      "predictor_learning_rate": 0.0004989822808835758,
+      "step": 565
+    },
+    {
+      "epoch": 0.06,
+      "gumbel_tau": 0.8752,
+      "image_masks_loss": 0.0206,
+      "llm_learning_rate": 4.989747890270293e-06,
+      "loss": 0.8378,
+      "output_text_mask_loss": 0.0572,
+      "predictor_learning_rate": 0.0004989747890270293,
+      "step": 566
+    },
+    {
+      "epoch": 0.06,
+      "gumbel_tau": 0.875,
+      "image_masks_loss": 0.045,
+      "llm_learning_rate": 4.9896726975282455e-06,
+      "loss": 0.7257,
+      "output_text_mask_loss": 0.0039,
+      "predictor_learning_rate": 0.0004989672697528246,
+      "step": 567
+    },
+    {
+      "epoch": 0.06,
+      "gumbel_tau": 0.8748,
+      "image_masks_loss": 0.0309,
+      "llm_learning_rate": 4.989597230617894e-06,
+      "loss": 0.9649,
+      "output_text_mask_loss": 0.0196,
+      "predictor_learning_rate": 0.0004989597230617895,
+      "step": 568
+    },
+    {
+      "epoch": 0.06,
+      "gumbel_tau": 0.8746,
+      "image_masks_loss": 0.0626,
+      "llm_learning_rate": 4.9895214895475515e-06,
+      "loss": 0.8688,
+      "output_text_mask_loss": 0.0261,
+      "predictor_learning_rate": 0.0004989521489547552,
+      "step": 569
+    },
+    {
+      "epoch": 0.06,
+      "gumbel_tau": 0.8744,
+      "image_masks_loss": 0.0154,
+      "llm_learning_rate": 4.9894454743255576e-06,
+      "loss": 0.8803,
+      "output_text_mask_loss": 0.0862,
+      "predictor_learning_rate": 0.0004989445474325556,
+      "step": 570
+    },
+    {
+      "epoch": 0.06,
+      "gumbel_tau": 0.8742,
+      "image_masks_loss": 0.0629,
+      "llm_learning_rate": 4.989369184960281e-06,
+      "loss": 0.9166,
+      "output_text_mask_loss": 0.0216,
+      "predictor_learning_rate": 0.0004989369184960281,
+      "step": 571
+    },
+    {
+      "epoch": 0.06,
+      "gumbel_tau": 0.874,
+      "image_masks_loss": 0.0338,
+      "llm_learning_rate": 4.989292621460127e-06,
+      "loss": 0.8411,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004989292621460126,
+      "step": 572
+    },
+    {
+      "epoch": 0.06,
+      "gumbel_tau": 0.8738,
+      "image_masks_loss": 0.0135,
+      "llm_learning_rate": 4.989215783833523e-06,
+      "loss": 0.9005,
+      "output_text_mask_loss": 0.0687,
+      "predictor_learning_rate": 0.0004989215783833523,
+      "step": 573
+    },
+    {
+      "epoch": 0.06,
+      "gumbel_tau": 0.8736,
+      "image_masks_loss": 0.037,
+      "llm_learning_rate": 4.989138672088933e-06,
+      "loss": 0.9159,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.0004989138672088933,
+      "step": 574
+    },
+    {
+      "epoch": 0.06,
+      "gumbel_tau": 0.8734,
+      "image_masks_loss": 0.0359,
+      "llm_learning_rate": 4.989061286234849e-06,
+      "loss": 0.8137,
+      "output_text_mask_loss": 0.0809,
+      "predictor_learning_rate": 0.0004989061286234849,
+      "step": 575
+    },
+    {
+      "epoch": 0.06,
+      "gumbel_tau": 0.8731,
+      "image_masks_loss": 0.0496,
+      "llm_learning_rate": 4.988983626279792e-06,
+      "loss": 0.8139,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 0.0004988983626279791,
+      "step": 576
+    },
+    {
+      "epoch": 0.06,
+      "gumbel_tau": 0.8729,
+      "image_masks_loss": 0.036,
+      "llm_learning_rate": 4.988905692232313e-06,
+      "loss": 0.916,
+      "output_text_mask_loss": 0.074,
+      "predictor_learning_rate": 0.0004988905692232313,
+      "step": 577
+    },
+    {
+      "epoch": 0.06,
+      "gumbel_tau": 0.8727,
+      "image_masks_loss": 0.0286,
+      "llm_learning_rate": 4.988827484100996e-06,
+      "loss": 0.9721,
+      "output_text_mask_loss": 0.1068,
+      "predictor_learning_rate": 0.0004988827484100995,
+      "step": 578
+    },
+    {
+      "epoch": 0.06,
+      "gumbel_tau": 0.8725,
+      "image_masks_loss": 0.0334,
+      "llm_learning_rate": 4.988749001894452e-06,
+      "loss": 0.8103,
+      "output_text_mask_loss": 0.0322,
+      "predictor_learning_rate": 0.0004988749001894452,
+      "step": 579
+    },
+    {
+      "epoch": 0.06,
+      "gumbel_tau": 0.8723,
+      "image_masks_loss": 0.0343,
+      "llm_learning_rate": 4.9886702456213255e-06,
+      "loss": 0.8505,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004988670245621326,
+      "step": 580
+    },
+    {
+      "epoch": 0.06,
+      "gumbel_tau": 0.8721,
+      "image_masks_loss": 0.0467,
+      "llm_learning_rate": 4.9885912152902875e-06,
+      "loss": 0.8699,
+      "output_text_mask_loss": 0.119,
+      "predictor_learning_rate": 0.0004988591215290287,
+      "step": 581
+    },
+    {
+      "epoch": 0.06,
+      "gumbel_tau": 0.8719,
+      "image_masks_loss": 0.0572,
+      "llm_learning_rate": 4.988511910910042e-06,
+      "loss": 0.9041,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004988511910910042,
+      "step": 582
+    },
+    {
+      "epoch": 0.06,
+      "gumbel_tau": 0.8717,
+      "image_masks_loss": 0.0408,
+      "llm_learning_rate": 4.988432332489323e-06,
+      "loss": 0.8706,
+      "output_text_mask_loss": 0.033,
+      "predictor_learning_rate": 0.0004988432332489322,
+      "step": 583
+    },
+    {
+      "epoch": 0.06,
+      "gumbel_tau": 0.8715,
+      "image_masks_loss": 0.0526,
+      "llm_learning_rate": 4.988352480036891e-06,
+      "loss": 0.7975,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.0004988352480036891,
+      "step": 584
+    },
+    {
+      "epoch": 0.06,
+      "gumbel_tau": 0.8713,
+      "image_masks_loss": 0.0641,
+      "llm_learning_rate": 4.9882723535615415e-06,
+      "loss": 0.738,
+      "output_text_mask_loss": 0.0969,
+      "predictor_learning_rate": 0.0004988272353561541,
+      "step": 585
+    },
+    {
+      "epoch": 0.06,
+      "gumbel_tau": 0.8711,
+      "image_masks_loss": 0.0254,
+      "llm_learning_rate": 4.988191953072099e-06,
+      "loss": 0.9265,
+      "output_text_mask_loss": 0.1678,
+      "predictor_learning_rate": 0.0004988191953072099,
+      "step": 586
+    },
+    {
+      "epoch": 0.06,
+      "gumbel_tau": 0.8709,
+      "image_masks_loss": 0.0127,
+      "llm_learning_rate": 4.988111278577417e-06,
+      "loss": 0.9041,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004988111278577416,
+      "step": 587
+    },
+    {
+      "epoch": 0.06,
+      "gumbel_tau": 0.8707,
+      "image_masks_loss": 0.0664,
+      "llm_learning_rate": 4.988030330086377e-06,
+      "loss": 0.8522,
+      "output_text_mask_loss": 0.0572,
+      "predictor_learning_rate": 0.0004988030330086377,
+      "step": 588
+    },
+    {
+      "epoch": 0.06,
+      "gumbel_tau": 0.8705,
+      "image_masks_loss": 0.0412,
+      "llm_learning_rate": 4.9879491076078966e-06,
+      "loss": 0.8031,
+      "output_text_mask_loss": 0.0923,
+      "predictor_learning_rate": 0.0004987949107607897,
+      "step": 589
+    },
+    {
+      "epoch": 0.06,
+      "gumbel_tau": 0.8703,
+      "image_masks_loss": 0.0254,
+      "llm_learning_rate": 4.987867611150918e-06,
+      "loss": 0.816,
+      "output_text_mask_loss": 0.0511,
+      "predictor_learning_rate": 0.0004987867611150918,
+      "step": 590
+    },
+    {
+      "epoch": 0.06,
+      "gumbel_tau": 0.8701,
+      "image_masks_loss": 0.0443,
+      "llm_learning_rate": 4.9877858407244175e-06,
+      "loss": 0.8146,
+      "output_text_mask_loss": 0.0006,
+      "predictor_learning_rate": 0.0004987785840724417,
+      "step": 591
+    },
+    {
+      "epoch": 0.06,
+      "gumbel_tau": 0.8699,
+      "image_masks_loss": 0.0414,
+      "llm_learning_rate": 4.987703796337397e-06,
+      "loss": 0.901,
+      "output_text_mask_loss": 0.0294,
+      "predictor_learning_rate": 0.0004987703796337398,
+      "step": 592
+    },
+    {
+      "epoch": 0.06,
+      "gumbel_tau": 0.8696,
+      "image_masks_loss": 0.0736,
+      "llm_learning_rate": 4.987621477998895e-06,
+      "loss": 0.8412,
+      "output_text_mask_loss": 0.0004,
+      "predictor_learning_rate": 0.0004987621477998894,
+      "step": 593
+    },
+    {
+      "epoch": 0.06,
+      "gumbel_tau": 0.8694,
+      "image_masks_loss": 0.0343,
+      "llm_learning_rate": 4.987538885717974e-06,
+      "loss": 0.8518,
+      "output_text_mask_loss": 0.0107,
+      "predictor_learning_rate": 0.0004987538885717974,
+      "step": 594
+    },
+    {
+      "epoch": 0.06,
+      "gumbel_tau": 0.8692,
+      "image_masks_loss": 0.174,
+      "llm_learning_rate": 4.987456019503731e-06,
+      "loss": 0.8985,
+      "output_text_mask_loss": 0.0908,
+      "predictor_learning_rate": 0.0004987456019503731,
+      "step": 595
+    },
+    {
+      "epoch": 0.06,
+      "gumbel_tau": 0.869,
+      "image_masks_loss": 0.1381,
+      "llm_learning_rate": 4.98737287936529e-06,
+      "loss": 0.978,
+      "output_text_mask_loss": 0.0416,
+      "predictor_learning_rate": 0.000498737287936529,
+      "step": 596
+    },
+    {
+      "epoch": 0.06,
+      "gumbel_tau": 0.8688,
+      "image_masks_loss": 0.0179,
+      "llm_learning_rate": 4.987289465311808e-06,
+      "loss": 0.8013,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.0004987289465311807,
+      "step": 597
+    },
+    {
+      "epoch": 0.06,
+      "gumbel_tau": 0.8686,
+      "image_masks_loss": 0.0142,
+      "llm_learning_rate": 4.987205777352469e-06,
+      "loss": 0.7583,
+      "output_text_mask_loss": 0.0095,
+      "predictor_learning_rate": 0.0004987205777352468,
+      "step": 598
+    },
+    {
+      "epoch": 0.06,
+      "gumbel_tau": 0.8684,
+      "image_masks_loss": 0.0427,
+      "llm_learning_rate": 4.987121815496489e-06,
+      "loss": 0.8099,
+      "output_text_mask_loss": 0.0618,
+      "predictor_learning_rate": 0.0004987121815496489,
+      "step": 599
+    },
+    {
+      "epoch": 0.06,
+      "gumbel_tau": 0.8682,
+      "image_masks_loss": 0.0633,
+      "llm_learning_rate": 4.987037579753116e-06,
+      "loss": 0.7071,
+      "output_text_mask_loss": 0.0004,
+      "predictor_learning_rate": 0.0004987037579753116,
+      "step": 600
+    },
+    {
+      "epoch": 0.06,
+      "gumbel_tau": 0.868,
+      "image_masks_loss": 0.0433,
+      "llm_learning_rate": 4.986953070131625e-06,
+      "loss": 0.8427,
+      "output_text_mask_loss": 0.0025,
+      "predictor_learning_rate": 0.0004986953070131625,
+      "step": 601
+    },
+    {
+      "epoch": 0.06,
+      "gumbel_tau": 0.8678,
+      "image_masks_loss": 0.0332,
+      "llm_learning_rate": 4.986868286641322e-06,
+      "loss": 0.9268,
+      "output_text_mask_loss": 0.0324,
+      "predictor_learning_rate": 0.0004986868286641321,
+      "step": 602
+    },
+    {
+      "epoch": 0.06,
+      "gumbel_tau": 0.8676,
+      "image_masks_loss": 0.0496,
+      "llm_learning_rate": 4.986783229291544e-06,
+      "loss": 0.7854,
+      "output_text_mask_loss": 0.1312,
+      "predictor_learning_rate": 0.0004986783229291543,
+      "step": 603
+    },
+    {
+      "epoch": 0.06,
+      "gumbel_tau": 0.8674,
+      "image_masks_loss": 0.0347,
+      "llm_learning_rate": 4.986697898091658e-06,
+      "loss": 0.8166,
+      "output_text_mask_loss": 0.0488,
+      "predictor_learning_rate": 0.0004986697898091657,
+      "step": 604
+    },
+    {
+      "epoch": 0.06,
+      "gumbel_tau": 0.8672,
+      "image_masks_loss": 0.0881,
+      "llm_learning_rate": 4.986612293051061e-06,
+      "loss": 0.9111,
+      "output_text_mask_loss": 0.0012,
+      "predictor_learning_rate": 0.000498661229305106,
+      "step": 605
+    },
+    {
+      "epoch": 0.06,
+      "gumbel_tau": 0.867,
+      "image_masks_loss": 0.0418,
+      "llm_learning_rate": 4.9865264141791796e-06,
+      "loss": 0.7419,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.0004986526414179178,
+      "step": 606
+    },
+    {
+      "epoch": 0.06,
+      "gumbel_tau": 0.8668,
+      "image_masks_loss": 0.0235,
+      "llm_learning_rate": 4.98644026148547e-06,
+      "loss": 0.8371,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004986440261485469,
+      "step": 607
+    },
+    {
+      "epoch": 0.06,
+      "gumbel_tau": 0.8666,
+      "image_masks_loss": 0.0218,
+      "llm_learning_rate": 4.98635383497942e-06,
+      "loss": 0.9429,
+      "output_text_mask_loss": 0.1099,
+      "predictor_learning_rate": 0.000498635383497942,
+      "step": 608
+    },
+    {
+      "epoch": 0.06,
+      "gumbel_tau": 0.8664,
+      "image_masks_loss": 0.0404,
+      "llm_learning_rate": 4.986267134670549e-06,
+      "loss": 0.8454,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.0004986267134670548,
+      "step": 609
+    },
+    {
+      "epoch": 0.06,
+      "gumbel_tau": 0.8662,
+      "image_masks_loss": 0.0381,
+      "llm_learning_rate": 4.986180160568402e-06,
+      "loss": 0.8906,
+      "output_text_mask_loss": 0.0916,
+      "predictor_learning_rate": 0.0004986180160568403,
+      "step": 610
+    },
+    {
+      "epoch": 0.06,
+      "gumbel_tau": 0.866,
+      "image_masks_loss": 0.0313,
+      "llm_learning_rate": 4.986092912682559e-06,
+      "loss": 0.8199,
+      "output_text_mask_loss": 0.1076,
+      "predictor_learning_rate": 0.0004986092912682558,
+      "step": 611
+    },
+    {
+      "epoch": 0.06,
+      "gumbel_tau": 0.8658,
+      "image_masks_loss": 0.074,
+      "llm_learning_rate": 4.986005391022627e-06,
+      "loss": 0.9006,
+      "output_text_mask_loss": 0.1701,
+      "predictor_learning_rate": 0.0004986005391022626,
+      "step": 612
+    },
+    {
+      "epoch": 0.06,
+      "gumbel_tau": 0.8656,
+      "image_masks_loss": 0.1038,
+      "llm_learning_rate": 4.985917595598243e-06,
+      "loss": 0.8277,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004985917595598243,
+      "step": 613
+    },
+    {
+      "epoch": 0.06,
+      "gumbel_tau": 0.8654,
+      "image_masks_loss": 0.053,
+      "llm_learning_rate": 4.985829526419077e-06,
+      "loss": 0.8724,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004985829526419077,
+      "step": 614
+    },
+    {
+      "epoch": 0.06,
+      "gumbel_tau": 0.8651,
+      "image_masks_loss": 0.011,
+      "llm_learning_rate": 4.985741183494825e-06,
+      "loss": 0.8088,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004985741183494826,
+      "step": 615
+    },
+    {
+      "epoch": 0.06,
+      "gumbel_tau": 0.8649,
+      "image_masks_loss": 0.0422,
+      "llm_learning_rate": 4.985652566835219e-06,
+      "loss": 0.9159,
+      "output_text_mask_loss": 0.001,
+      "predictor_learning_rate": 0.0004985652566835218,
+      "step": 616
+    },
+    {
+      "epoch": 0.06,
+      "gumbel_tau": 0.8647,
+      "image_masks_loss": 0.0098,
+      "llm_learning_rate": 4.985563676450015e-06,
+      "loss": 0.7422,
+      "output_text_mask_loss": 0.0374,
+      "predictor_learning_rate": 0.0004985563676450014,
+      "step": 617
+    },
+    {
+      "epoch": 0.06,
+      "gumbel_tau": 0.8645,
+      "image_masks_loss": 0.0248,
+      "llm_learning_rate": 4.985474512349003e-06,
+      "loss": 0.8995,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004985474512349003,
+      "step": 618
+    },
+    {
+      "epoch": 0.06,
+      "gumbel_tau": 0.8643,
+      "image_masks_loss": 0.0427,
+      "llm_learning_rate": 4.985385074542001e-06,
+      "loss": 0.8788,
+      "output_text_mask_loss": 0.0124,
+      "predictor_learning_rate": 0.0004985385074542001,
+      "step": 619
+    },
+    {
+      "epoch": 0.06,
+      "gumbel_tau": 0.8641,
+      "image_masks_loss": 0.0227,
+      "llm_learning_rate": 4.985295363038859e-06,
+      "loss": 0.8863,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.0004985295363038859,
+      "step": 620
+    },
+    {
+      "epoch": 0.06,
+      "gumbel_tau": 0.8639,
+      "image_masks_loss": 0.0231,
+      "llm_learning_rate": 4.985205377849456e-06,
+      "loss": 0.7796,
+      "output_text_mask_loss": 0.001,
+      "predictor_learning_rate": 0.0004985205377849455,
+      "step": 621
+    },
+    {
+      "epoch": 0.06,
+      "gumbel_tau": 0.8637,
+      "image_masks_loss": 0.0324,
+      "llm_learning_rate": 4.9851151189837e-06,
+      "loss": 0.853,
+      "output_text_mask_loss": 0.0081,
+      "predictor_learning_rate": 0.0004985115118983699,
+      "step": 622
+    },
+    {
+      "epoch": 0.06,
+      "gumbel_tau": 0.8635,
+      "image_masks_loss": 0.0446,
+      "llm_learning_rate": 4.985024586451533e-06,
+      "loss": 0.8473,
+      "output_text_mask_loss": 0.0078,
+      "predictor_learning_rate": 0.0004985024586451533,
+      "step": 623
+    },
+    {
+      "epoch": 0.06,
+      "gumbel_tau": 0.8633,
+      "image_masks_loss": 0.0441,
+      "llm_learning_rate": 4.984933780262923e-06,
+      "loss": 0.7903,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004984933780262923,
+      "step": 624
+    },
+    {
+      "epoch": 0.06,
+      "gumbel_tau": 0.8631,
+      "image_masks_loss": 0.0294,
+      "llm_learning_rate": 4.98484270042787e-06,
+      "loss": 0.8028,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.000498484270042787,
+      "step": 625
+    },
+    {
+      "epoch": 0.06,
+      "gumbel_tau": 0.8629,
+      "image_masks_loss": 0.0492,
+      "llm_learning_rate": 4.984751346956404e-06,
+      "loss": 0.9066,
+      "output_text_mask_loss": 0.1579,
+      "predictor_learning_rate": 0.0004984751346956403,
+      "step": 626
+    },
+    {
+      "epoch": 0.06,
+      "gumbel_tau": 0.8627,
+      "image_masks_loss": 0.0094,
+      "llm_learning_rate": 4.984659719858586e-06,
+      "loss": 0.6869,
+      "output_text_mask_loss": 0.0511,
+      "predictor_learning_rate": 0.0004984659719858586,
+      "step": 627
+    },
+    {
+      "epoch": 0.06,
+      "gumbel_tau": 0.8625,
+      "image_masks_loss": 0.0412,
+      "llm_learning_rate": 4.984567819144504e-06,
+      "loss": 0.9286,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 0.0004984567819144503,
+      "step": 628
+    },
+    {
+      "epoch": 0.06,
+      "gumbel_tau": 0.8623,
+      "image_masks_loss": 0.0084,
+      "llm_learning_rate": 4.98447564482428e-06,
+      "loss": 0.8614,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.000498447564482428,
+      "step": 629
+    },
+    {
+      "epoch": 0.06,
+      "gumbel_tau": 0.8621,
+      "image_masks_loss": 0.0167,
+      "llm_learning_rate": 4.984383196908064e-06,
+      "loss": 0.9707,
+      "output_text_mask_loss": 0.1099,
+      "predictor_learning_rate": 0.0004984383196908064,
+      "step": 630
+    },
+    {
+      "epoch": 0.06,
+      "gumbel_tau": 0.8619,
+      "image_masks_loss": 0.0204,
+      "llm_learning_rate": 4.984290475406037e-06,
+      "loss": 0.8409,
+      "output_text_mask_loss": 0.0759,
+      "predictor_learning_rate": 0.0004984290475406037,
+      "step": 631
+    },
+    {
+      "epoch": 0.06,
+      "gumbel_tau": 0.8617,
+      "image_masks_loss": 0.0152,
+      "llm_learning_rate": 4.984197480328409e-06,
+      "loss": 0.8373,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.0004984197480328408,
+      "step": 632
+    },
+    {
+      "epoch": 0.06,
+      "gumbel_tau": 0.8615,
+      "image_masks_loss": 0.032,
+      "llm_learning_rate": 4.984104211685421e-06,
+      "loss": 0.8505,
+      "output_text_mask_loss": 0.087,
+      "predictor_learning_rate": 0.0004984104211685421,
+      "step": 633
+    },
+    {
+      "epoch": 0.06,
+      "gumbel_tau": 0.8613,
+      "image_masks_loss": 0.0341,
+      "llm_learning_rate": 4.984010669487344e-06,
+      "loss": 0.8392,
+      "output_text_mask_loss": 0.0353,
+      "predictor_learning_rate": 0.0004984010669487344,
+      "step": 634
+    },
+    {
+      "epoch": 0.07,
+      "gumbel_tau": 0.8611,
+      "image_masks_loss": 0.0114,
+      "llm_learning_rate": 4.98391685374448e-06,
+      "loss": 0.7651,
+      "output_text_mask_loss": 0.0162,
+      "predictor_learning_rate": 0.000498391685374448,
+      "step": 635
+    },
+    {
+      "epoch": 0.07,
+      "gumbel_tau": 0.8609,
+      "image_masks_loss": 0.025,
+      "llm_learning_rate": 4.983822764467159e-06,
+      "loss": 0.8246,
+      "output_text_mask_loss": 0.0071,
+      "predictor_learning_rate": 0.0004983822764467159,
+      "step": 636
+    },
+    {
+      "epoch": 0.07,
+      "gumbel_tau": 0.8607,
+      "image_masks_loss": 0.0492,
+      "llm_learning_rate": 4.983728401665743e-06,
+      "loss": 0.8266,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.0004983728401665743,
+      "step": 637
+    },
+    {
+      "epoch": 0.07,
+      "gumbel_tau": 0.8605,
+      "image_masks_loss": 0.0113,
+      "llm_learning_rate": 4.983633765350622e-06,
+      "loss": 0.7148,
+      "output_text_mask_loss": 0.0446,
+      "predictor_learning_rate": 0.0004983633765350622,
+      "step": 638
+    },
+    {
+      "epoch": 0.07,
+      "gumbel_tau": 0.8603,
+      "image_masks_loss": 0.0188,
+      "llm_learning_rate": 4.9835388555322194e-06,
+      "loss": 0.7729,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004983538855532219,
+      "step": 639
+    },
+    {
+      "epoch": 0.07,
+      "gumbel_tau": 0.8601,
+      "image_masks_loss": 0.0087,
+      "llm_learning_rate": 4.983443672220986e-06,
+      "loss": 0.889,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.0004983443672220986,
+      "step": 640
+    },
+    {
+      "epoch": 0.07,
+      "gumbel_tau": 0.8599,
+      "image_masks_loss": 0.0368,
+      "llm_learning_rate": 4.983348215427404e-06,
+      "loss": 0.7893,
+      "output_text_mask_loss": 0.0284,
+      "predictor_learning_rate": 0.0004983348215427403,
+      "step": 641
+    },
+    {
+      "epoch": 0.07,
+      "gumbel_tau": 0.8597,
+      "image_masks_loss": 0.0309,
+      "llm_learning_rate": 4.983252485161985e-06,
+      "loss": 0.8578,
+      "output_text_mask_loss": 0.0557,
+      "predictor_learning_rate": 0.0004983252485161984,
+      "step": 642
+    },
+    {
+      "epoch": 0.07,
+      "gumbel_tau": 0.8595,
+      "image_masks_loss": 0.021,
+      "llm_learning_rate": 4.983156481435271e-06,
+      "loss": 0.7471,
+      "output_text_mask_loss": 0.0105,
+      "predictor_learning_rate": 0.0004983156481435271,
+      "step": 643
+    },
+    {
+      "epoch": 0.07,
+      "gumbel_tau": 0.8592,
+      "image_masks_loss": 0.0119,
+      "llm_learning_rate": 4.983060204257834e-06,
+      "loss": 0.8747,
+      "output_text_mask_loss": 0.079,
+      "predictor_learning_rate": 0.0004983060204257834,
+      "step": 644
+    },
+    {
+      "epoch": 0.07,
+      "gumbel_tau": 0.859,
+      "image_masks_loss": 0.0109,
+      "llm_learning_rate": 4.982963653640277e-06,
+      "loss": 0.7743,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 0.0004982963653640276,
+      "step": 645
+    },
+    {
+      "epoch": 0.07,
+      "gumbel_tau": 0.8588,
+      "image_masks_loss": 0.0237,
+      "llm_learning_rate": 4.982866829593231e-06,
+      "loss": 0.8906,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.0004982866829593231,
+      "step": 646
+    },
+    {
+      "epoch": 0.07,
+      "gumbel_tau": 0.8586,
+      "image_masks_loss": 0.0282,
+      "llm_learning_rate": 4.982769732127361e-06,
+      "loss": 0.786,
+      "output_text_mask_loss": 0.0347,
+      "predictor_learning_rate": 0.0004982769732127361,
+      "step": 647
+    },
+    {
+      "epoch": 0.07,
+      "gumbel_tau": 0.8584,
+      "image_masks_loss": 0.0078,
+      "llm_learning_rate": 4.982672361253357e-06,
+      "loss": 0.8656,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004982672361253357,
+      "step": 648
+    },
+    {
+      "epoch": 0.07,
+      "gumbel_tau": 0.8582,
+      "image_masks_loss": 0.0113,
+      "llm_learning_rate": 4.982574716981942e-06,
+      "loss": 0.7749,
+      "output_text_mask_loss": 0.0055,
+      "predictor_learning_rate": 0.0004982574716981942,
+      "step": 649
+    },
+    {
+      "epoch": 0.07,
+      "gumbel_tau": 0.858,
+      "image_masks_loss": 0.0211,
+      "llm_learning_rate": 4.98247679932387e-06,
+      "loss": 0.8475,
+      "output_text_mask_loss": 0.0046,
+      "predictor_learning_rate": 0.0004982476799323869,
+      "step": 650
+    },
+    {
+      "epoch": 0.07,
+      "gumbel_tau": 0.8578,
+      "image_masks_loss": 0.0194,
+      "llm_learning_rate": 4.9823786082899235e-06,
+      "loss": 0.8767,
+      "output_text_mask_loss": 0.0471,
+      "predictor_learning_rate": 0.0004982378608289924,
+      "step": 651
+    },
+    {
+      "epoch": 0.07,
+      "gumbel_tau": 0.8576,
+      "image_masks_loss": 0.1022,
+      "llm_learning_rate": 4.982280143890915e-06,
+      "loss": 0.8252,
+      "output_text_mask_loss": 0.0207,
+      "predictor_learning_rate": 0.0004982280143890915,
+      "step": 652
+    },
+    {
+      "epoch": 0.07,
+      "gumbel_tau": 0.8574,
+      "image_masks_loss": 0.0546,
+      "llm_learning_rate": 4.982181406137689e-06,
+      "loss": 0.8959,
+      "output_text_mask_loss": 0.0946,
+      "predictor_learning_rate": 0.0004982181406137688,
+      "step": 653
+    },
+    {
+      "epoch": 0.07,
+      "gumbel_tau": 0.8572,
+      "image_masks_loss": 0.0431,
+      "llm_learning_rate": 4.9820823950411166e-06,
+      "loss": 0.8044,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 0.0004982082395041116,
+      "step": 654
+    },
+    {
+      "epoch": 0.07,
+      "gumbel_tau": 0.857,
+      "image_masks_loss": 0.0265,
+      "llm_learning_rate": 4.981983110612103e-06,
+      "loss": 0.8164,
+      "output_text_mask_loss": 0.0286,
+      "predictor_learning_rate": 0.0004981983110612103,
+      "step": 655
+    },
+    {
+      "epoch": 0.07,
+      "gumbel_tau": 0.8568,
+      "image_masks_loss": 0.0226,
+      "llm_learning_rate": 4.98188355286158e-06,
+      "loss": 0.8998,
+      "output_text_mask_loss": 0.1465,
+      "predictor_learning_rate": 0.000498188355286158,
+      "step": 656
+    },
+    {
+      "epoch": 0.07,
+      "gumbel_tau": 0.8566,
+      "image_masks_loss": 0.0091,
+      "llm_learning_rate": 4.981783721800513e-06,
+      "loss": 0.7299,
+      "output_text_mask_loss": 0.0055,
+      "predictor_learning_rate": 0.0004981783721800513,
+      "step": 657
+    },
+    {
+      "epoch": 0.07,
+      "gumbel_tau": 0.8564,
+      "image_masks_loss": 0.0307,
+      "llm_learning_rate": 4.981683617439893e-06,
+      "loss": 0.8503,
+      "output_text_mask_loss": 0.046,
+      "predictor_learning_rate": 0.0004981683617439894,
+      "step": 658
+    },
+    {
+      "epoch": 0.07,
+      "gumbel_tau": 0.8562,
+      "image_masks_loss": 0.037,
+      "llm_learning_rate": 4.981583239790747e-06,
+      "loss": 0.7887,
+      "output_text_mask_loss": 0.0221,
+      "predictor_learning_rate": 0.0004981583239790746,
+      "step": 659
+    },
+    {
+      "epoch": 0.07,
+      "gumbel_tau": 0.856,
+      "image_masks_loss": 0.0063,
+      "llm_learning_rate": 4.9814825888641265e-06,
+      "loss": 0.7919,
+      "output_text_mask_loss": 0.0416,
+      "predictor_learning_rate": 0.0004981482588864126,
+      "step": 660
+    },
+    {
+      "epoch": 0.07,
+      "gumbel_tau": 0.8558,
+      "image_masks_loss": 0.0235,
+      "llm_learning_rate": 4.981381664671116e-06,
+      "loss": 0.8192,
+      "output_text_mask_loss": 0.0039,
+      "predictor_learning_rate": 0.0004981381664671115,
+      "step": 661
+    },
+    {
+      "epoch": 0.07,
+      "gumbel_tau": 0.8556,
+      "image_masks_loss": 0.0083,
+      "llm_learning_rate": 4.98128046722283e-06,
+      "loss": 0.7794,
+      "output_text_mask_loss": 0.1015,
+      "predictor_learning_rate": 0.000498128046722283,
+      "step": 662
+    },
+    {
+      "epoch": 0.07,
+      "gumbel_tau": 0.8554,
+      "image_masks_loss": 0.0073,
+      "llm_learning_rate": 4.9811789965304115e-06,
+      "loss": 0.8087,
+      "output_text_mask_loss": 0.0374,
+      "predictor_learning_rate": 0.0004981178996530412,
+      "step": 663
+    },
+    {
+      "epoch": 0.07,
+      "gumbel_tau": 0.8552,
+      "image_masks_loss": 0.0203,
+      "llm_learning_rate": 4.981077252605035e-06,
+      "loss": 0.7688,
+      "output_text_mask_loss": 0.0269,
+      "predictor_learning_rate": 0.0004981077252605035,
+      "step": 664
+    },
+    {
+      "epoch": 0.07,
+      "gumbel_tau": 0.855,
+      "image_masks_loss": 0.0071,
+      "llm_learning_rate": 4.9809752354579055e-06,
+      "loss": 0.7547,
+      "output_text_mask_loss": 0.0103,
+      "predictor_learning_rate": 0.0004980975235457906,
+      "step": 665
+    },
+    {
+      "epoch": 0.07,
+      "gumbel_tau": 0.8548,
+      "image_masks_loss": 0.0057,
+      "llm_learning_rate": 4.980872945100257e-06,
+      "loss": 0.747,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004980872945100257,
+      "step": 666
+    },
+    {
+      "epoch": 0.07,
+      "gumbel_tau": 0.8546,
+      "image_masks_loss": 0.0433,
+      "llm_learning_rate": 4.980770381543354e-06,
+      "loss": 0.916,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.0004980770381543354,
+      "step": 667
+    },
+    {
+      "epoch": 0.07,
+      "gumbel_tau": 0.8544,
+      "image_masks_loss": 0.0193,
+      "llm_learning_rate": 4.980667544798491e-06,
+      "loss": 0.7805,
+      "output_text_mask_loss": 0.0736,
+      "predictor_learning_rate": 0.0004980667544798491,
+      "step": 668
+    },
+    {
+      "epoch": 0.07,
+      "gumbel_tau": 0.8542,
+      "image_masks_loss": 0.0181,
+      "llm_learning_rate": 4.980564434876991e-06,
+      "loss": 0.8095,
+      "output_text_mask_loss": 0.0343,
+      "predictor_learning_rate": 0.0004980564434876991,
+      "step": 669
+    },
+    {
+      "epoch": 0.07,
+      "gumbel_tau": 0.854,
+      "image_masks_loss": 0.046,
+      "llm_learning_rate": 4.980461051790212e-06,
+      "loss": 0.8448,
+      "output_text_mask_loss": 0.0334,
+      "predictor_learning_rate": 0.0004980461051790211,
+      "step": 670
+    },
+    {
+      "epoch": 0.07,
+      "gumbel_tau": 0.8538,
+      "image_masks_loss": 0.0263,
+      "llm_learning_rate": 4.980357395549537e-06,
+      "loss": 0.859,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004980357395549536,
+      "step": 671
+    },
+    {
+      "epoch": 0.07,
+      "gumbel_tau": 0.8536,
+      "image_masks_loss": 0.0225,
+      "llm_learning_rate": 4.980253466166379e-06,
+      "loss": 0.7933,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.0004980253466166379,
+      "step": 672
+    },
+    {
+      "epoch": 0.07,
+      "gumbel_tau": 0.8534,
+      "image_masks_loss": 0.036,
+      "llm_learning_rate": 4.980149263652186e-06,
+      "loss": 0.7862,
+      "output_text_mask_loss": 0.0128,
+      "predictor_learning_rate": 0.0004980149263652187,
+      "step": 673
+    },
+    {
+      "epoch": 0.07,
+      "gumbel_tau": 0.8532,
+      "image_masks_loss": 0.0401,
+      "llm_learning_rate": 4.980044788018432e-06,
+      "loss": 0.8582,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004980044788018431,
+      "step": 674
+    },
+    {
+      "epoch": 0.07,
+      "gumbel_tau": 0.853,
+      "image_masks_loss": 0.0462,
+      "llm_learning_rate": 4.979940039276622e-06,
+      "loss": 0.6898,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.0004979940039276621,
+      "step": 675
+    },
+    {
+      "epoch": 0.07,
+      "gumbel_tau": 0.8528,
+      "image_masks_loss": 0.0168,
+      "llm_learning_rate": 4.97983501743829e-06,
+      "loss": 0.7758,
+      "output_text_mask_loss": 0.0139,
+      "predictor_learning_rate": 0.000497983501743829,
+      "step": 676
+    },
+    {
+      "epoch": 0.07,
+      "gumbel_tau": 0.8526,
+      "image_masks_loss": 0.0117,
+      "llm_learning_rate": 4.979729722515003e-06,
+      "loss": 0.6818,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004979729722515003,
+      "step": 677
+    },
+    {
+      "epoch": 0.07,
+      "gumbel_tau": 0.8524,
+      "image_masks_loss": 0.0549,
+      "llm_learning_rate": 4.979624154518355e-06,
+      "loss": 0.915,
+      "output_text_mask_loss": 0.0771,
+      "predictor_learning_rate": 0.0004979624154518355,
+      "step": 678
+    },
+    {
+      "epoch": 0.07,
+      "gumbel_tau": 0.8522,
+      "image_masks_loss": 0.0397,
+      "llm_learning_rate": 4.979518313459973e-06,
+      "loss": 0.7331,
+      "output_text_mask_loss": 0.0004,
+      "predictor_learning_rate": 0.0004979518313459972,
+      "step": 679
+    },
+    {
+      "epoch": 0.07,
+      "gumbel_tau": 0.852,
+      "image_masks_loss": 0.0195,
+      "llm_learning_rate": 4.97941219935151e-06,
+      "loss": 0.705,
+      "output_text_mask_loss": 0.0025,
+      "predictor_learning_rate": 0.000497941219935151,
+      "step": 680
+    },
+    {
+      "epoch": 0.07,
+      "gumbel_tau": 0.8518,
+      "image_masks_loss": 0.0148,
+      "llm_learning_rate": 4.979305812204654e-06,
+      "loss": 0.8212,
+      "output_text_mask_loss": 0.074,
+      "predictor_learning_rate": 0.0004979305812204654,
+      "step": 681
+    },
+    {
+      "epoch": 0.07,
+      "gumbel_tau": 0.8516,
+      "image_masks_loss": 0.0221,
+      "llm_learning_rate": 4.979199152031119e-06,
+      "loss": 0.8705,
+      "output_text_mask_loss": 0.0694,
+      "predictor_learning_rate": 0.000497919915203112,
+      "step": 682
+    },
+    {
+      "epoch": 0.07,
+      "gumbel_tau": 0.8514,
+      "image_masks_loss": 0.0185,
+      "llm_learning_rate": 4.979092218842652e-06,
+      "loss": 0.8584,
+      "output_text_mask_loss": 0.0618,
+      "predictor_learning_rate": 0.0004979092218842651,
+      "step": 683
+    },
+    {
+      "epoch": 0.07,
+      "gumbel_tau": 0.8512,
+      "image_masks_loss": 0.0362,
+      "llm_learning_rate": 4.978985012651027e-06,
+      "loss": 0.8464,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.0004978985012651028,
+      "step": 684
+    },
+    {
+      "epoch": 0.07,
+      "gumbel_tau": 0.851,
+      "image_masks_loss": 0.0226,
+      "llm_learning_rate": 4.978877533468051e-06,
+      "loss": 0.87,
+      "output_text_mask_loss": 0.0212,
+      "predictor_learning_rate": 0.0004978877533468051,
+      "step": 685
+    },
+    {
+      "epoch": 0.07,
+      "gumbel_tau": 0.8508,
+      "image_masks_loss": 0.0292,
+      "llm_learning_rate": 4.9787697813055596e-06,
+      "loss": 0.8831,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004978769781305559,
+      "step": 686
+    },
+    {
+      "epoch": 0.07,
+      "gumbel_tau": 0.8506,
+      "image_masks_loss": 0.0164,
+      "llm_learning_rate": 4.978661756175419e-06,
+      "loss": 0.8447,
+      "output_text_mask_loss": 0.0763,
+      "predictor_learning_rate": 0.0004978661756175418,
+      "step": 687
+    },
+    {
+      "epoch": 0.07,
+      "gumbel_tau": 0.8504,
+      "image_masks_loss": 0.0265,
+      "llm_learning_rate": 4.978553458089524e-06,
+      "loss": 0.7843,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.0004978553458089524,
+      "step": 688
+    },
+    {
+      "epoch": 0.07,
+      "gumbel_tau": 0.8502,
+      "image_masks_loss": 0.0169,
+      "llm_learning_rate": 4.978444887059802e-06,
+      "loss": 0.7417,
+      "output_text_mask_loss": 0.0116,
+      "predictor_learning_rate": 0.0004978444887059802,
+      "step": 689
+    },
+    {
+      "epoch": 0.07,
+      "gumbel_tau": 0.85,
+      "image_masks_loss": 0.0197,
+      "llm_learning_rate": 4.978336043098208e-06,
+      "loss": 0.8235,
+      "output_text_mask_loss": 0.0284,
+      "predictor_learning_rate": 0.0004978336043098208,
+      "step": 690
+    },
+    {
+      "epoch": 0.07,
+      "gumbel_tau": 0.8498,
+      "image_masks_loss": 0.0328,
+      "llm_learning_rate": 4.978226926216729e-06,
+      "loss": 0.7612,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.0004978226926216729,
+      "step": 691
+    },
+    {
+      "epoch": 0.07,
+      "gumbel_tau": 0.8496,
+      "image_masks_loss": 0.018,
+      "llm_learning_rate": 4.9781175364273805e-06,
+      "loss": 0.7779,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004978117536427381,
+      "step": 692
+    },
+    {
+      "epoch": 0.07,
+      "gumbel_tau": 0.8494,
+      "image_masks_loss": 0.0171,
+      "llm_learning_rate": 4.97800787374221e-06,
+      "loss": 0.7398,
+      "output_text_mask_loss": 0.0233,
+      "predictor_learning_rate": 0.000497800787374221,
+      "step": 693
+    },
+    {
+      "epoch": 0.07,
+      "gumbel_tau": 0.8492,
+      "image_masks_loss": 0.0322,
+      "llm_learning_rate": 4.977897938173292e-06,
+      "loss": 0.7993,
+      "output_text_mask_loss": 0.0618,
+      "predictor_learning_rate": 0.0004977897938173291,
+      "step": 694
+    },
+    {
+      "epoch": 0.07,
+      "gumbel_tau": 0.849,
+      "image_masks_loss": 0.0121,
+      "llm_learning_rate": 4.977787729732734e-06,
+      "loss": 0.8604,
+      "output_text_mask_loss": 0.0374,
+      "predictor_learning_rate": 0.0004977787729732734,
+      "step": 695
+    },
+    {
+      "epoch": 0.07,
+      "gumbel_tau": 0.8488,
+      "image_masks_loss": 0.0305,
+      "llm_learning_rate": 4.977677248432672e-06,
+      "loss": 0.7878,
+      "output_text_mask_loss": 0.0275,
+      "predictor_learning_rate": 0.0004977677248432672,
+      "step": 696
+    },
+    {
+      "epoch": 0.07,
+      "gumbel_tau": 0.8486,
+      "image_masks_loss": 0.0549,
+      "llm_learning_rate": 4.977566494285273e-06,
+      "loss": 0.8279,
+      "output_text_mask_loss": 0.0763,
+      "predictor_learning_rate": 0.0004977566494285273,
+      "step": 697
+    },
+    {
+      "epoch": 0.07,
+      "gumbel_tau": 0.8484,
+      "image_masks_loss": 0.0267,
+      "llm_learning_rate": 4.9774554673027335e-06,
+      "loss": 0.8414,
+      "output_text_mask_loss": 0.1099,
+      "predictor_learning_rate": 0.0004977455467302733,
+      "step": 698
+    },
+    {
+      "epoch": 0.07,
+      "gumbel_tau": 0.8482,
+      "image_masks_loss": 0.0393,
+      "llm_learning_rate": 4.977344167497279e-06,
+      "loss": 0.9254,
+      "output_text_mask_loss": 0.0282,
+      "predictor_learning_rate": 0.0004977344167497278,
+      "step": 699
+    },
+    {
+      "epoch": 0.07,
+      "gumbel_tau": 0.848,
+      "image_masks_loss": 0.0418,
+      "llm_learning_rate": 4.9772325948811665e-06,
+      "loss": 0.802,
+      "output_text_mask_loss": 0.0477,
+      "predictor_learning_rate": 0.0004977232594881166,
+      "step": 700
+    },
+    {
+      "epoch": 0.07,
+      "gumbel_tau": 0.8478,
+      "image_masks_loss": 0.0167,
+      "llm_learning_rate": 4.977120749466683e-06,
+      "loss": 0.9364,
+      "output_text_mask_loss": 0.0353,
+      "predictor_learning_rate": 0.0004977120749466682,
+      "step": 701
+    },
+    {
+      "epoch": 0.07,
+      "gumbel_tau": 0.8476,
+      "image_masks_loss": 0.0216,
+      "llm_learning_rate": 4.977008631266145e-06,
+      "loss": 0.8238,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004977008631266145,
+      "step": 702
+    },
+    {
+      "epoch": 0.07,
+      "gumbel_tau": 0.8474,
+      "image_masks_loss": 0.0119,
+      "llm_learning_rate": 4.976896240291899e-06,
+      "loss": 0.8254,
+      "output_text_mask_loss": 0.0039,
+      "predictor_learning_rate": 0.0004976896240291899,
+      "step": 703
+    },
+    {
+      "epoch": 0.07,
+      "gumbel_tau": 0.8472,
+      "image_masks_loss": 0.0197,
+      "llm_learning_rate": 4.976783576556322e-06,
+      "loss": 0.8344,
+      "output_text_mask_loss": 0.0877,
+      "predictor_learning_rate": 0.0004976783576556322,
+      "step": 704
+    },
+    {
+      "epoch": 0.07,
+      "gumbel_tau": 0.847,
+      "image_masks_loss": 0.0257,
+      "llm_learning_rate": 4.976670640071821e-06,
+      "loss": 0.7857,
+      "output_text_mask_loss": 0.0374,
+      "predictor_learning_rate": 0.0004976670640071821,
+      "step": 705
+    },
+    {
+      "epoch": 0.07,
+      "gumbel_tau": 0.8468,
+      "image_masks_loss": 0.0171,
+      "llm_learning_rate": 4.976557430850833e-06,
+      "loss": 0.7079,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.0004976557430850832,
+      "step": 706
+    },
+    {
+      "epoch": 0.07,
+      "gumbel_tau": 0.8466,
+      "image_masks_loss": 0.0122,
+      "llm_learning_rate": 4.976443948905824e-06,
+      "loss": 0.7234,
+      "output_text_mask_loss": 0.0345,
+      "predictor_learning_rate": 0.0004976443948905824,
+      "step": 707
+    },
+    {
+      "epoch": 0.07,
+      "gumbel_tau": 0.8464,
+      "image_masks_loss": 0.0178,
+      "llm_learning_rate": 4.976330194249291e-06,
+      "loss": 0.6684,
+      "output_text_mask_loss": 0.0101,
+      "predictor_learning_rate": 0.000497633019424929,
+      "step": 708
+    },
+    {
+      "epoch": 0.07,
+      "gumbel_tau": 0.8462,
+      "image_masks_loss": 0.0167,
+      "llm_learning_rate": 4.9762161668937615e-06,
+      "loss": 0.7176,
+      "output_text_mask_loss": 0.0043,
+      "predictor_learning_rate": 0.0004976216166893761,
+      "step": 709
+    },
+    {
+      "epoch": 0.07,
+      "gumbel_tau": 0.846,
+      "image_masks_loss": 0.0126,
+      "llm_learning_rate": 4.976101866851792e-06,
+      "loss": 0.7795,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004976101866851791,
+      "step": 710
+    },
+    {
+      "epoch": 0.07,
+      "gumbel_tau": 0.8458,
+      "image_masks_loss": 0.0207,
+      "llm_learning_rate": 4.975987294135969e-06,
+      "loss": 0.8465,
+      "output_text_mask_loss": 0.0081,
+      "predictor_learning_rate": 0.0004975987294135969,
+      "step": 711
+    },
+    {
+      "epoch": 0.07,
+      "gumbel_tau": 0.8456,
+      "image_masks_loss": 0.0212,
+      "llm_learning_rate": 4.975872448758911e-06,
+      "loss": 0.7697,
+      "output_text_mask_loss": 0.0282,
+      "predictor_learning_rate": 0.0004975872448758911,
+      "step": 712
+    },
+    {
+      "epoch": 0.07,
+      "gumbel_tau": 0.8454,
+      "image_masks_loss": 0.02,
+      "llm_learning_rate": 4.9757573307332644e-06,
+      "loss": 0.7049,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.0004975757330733264,
+      "step": 713
+    },
+    {
+      "epoch": 0.07,
+      "gumbel_tau": 0.8452,
+      "image_masks_loss": 0.0126,
+      "llm_learning_rate": 4.975641940071705e-06,
+      "loss": 0.868,
+      "output_text_mask_loss": 0.0465,
+      "predictor_learning_rate": 0.0004975641940071705,
+      "step": 714
+    },
+    {
+      "epoch": 0.07,
+      "gumbel_tau": 0.845,
+      "image_masks_loss": 0.0158,
+      "llm_learning_rate": 4.975526276786942e-06,
+      "loss": 0.7993,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.0004975526276786941,
+      "step": 715
+    },
+    {
+      "epoch": 0.07,
+      "gumbel_tau": 0.8448,
+      "image_masks_loss": 0.0252,
+      "llm_learning_rate": 4.97541034089171e-06,
+      "loss": 0.733,
+      "output_text_mask_loss": 0.0301,
+      "predictor_learning_rate": 0.000497541034089171,
+      "step": 716
+    },
+    {
+      "epoch": 0.07,
+      "gumbel_tau": 0.8446,
+      "image_masks_loss": 0.0282,
+      "llm_learning_rate": 4.975294132398779e-06,
+      "loss": 0.7819,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.0004975294132398778,
+      "step": 717
+    },
+    {
+      "epoch": 0.07,
+      "gumbel_tau": 0.8444,
+      "image_masks_loss": 0.0027,
+      "llm_learning_rate": 4.975177651320944e-06,
+      "loss": 0.7774,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.0004975177651320943,
+      "step": 718
+    },
+    {
+      "epoch": 0.07,
+      "gumbel_tau": 0.8442,
+      "image_masks_loss": 0.0157,
+      "llm_learning_rate": 4.975060897671032e-06,
+      "loss": 0.7434,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 0.0004975060897671032,
+      "step": 719
+    },
+    {
+      "epoch": 0.07,
+      "gumbel_tau": 0.844,
+      "image_masks_loss": 0.0116,
+      "llm_learning_rate": 4.974943871461902e-06,
+      "loss": 0.7889,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004974943871461901,
+      "step": 720
+    },
+    {
+      "epoch": 0.07,
+      "gumbel_tau": 0.8438,
+      "image_masks_loss": 0.0215,
+      "llm_learning_rate": 4.9748265727064395e-06,
+      "loss": 0.8041,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.000497482657270644,
+      "step": 721
+    },
+    {
+      "epoch": 0.07,
+      "gumbel_tau": 0.8436,
+      "image_masks_loss": 0.0324,
+      "llm_learning_rate": 4.974709001417563e-06,
+      "loss": 0.8003,
+      "output_text_mask_loss": 0.0429,
+      "predictor_learning_rate": 0.0004974709001417562,
+      "step": 722
+    },
+    {
+      "epoch": 0.07,
+      "gumbel_tau": 0.8434,
+      "image_masks_loss": 0.0236,
+      "llm_learning_rate": 4.9745911576082186e-06,
+      "loss": 0.8224,
+      "output_text_mask_loss": 0.0322,
+      "predictor_learning_rate": 0.0004974591157608218,
+      "step": 723
+    },
+    {
+      "epoch": 0.07,
+      "gumbel_tau": 0.8432,
+      "image_masks_loss": 0.0164,
+      "llm_learning_rate": 4.974473041291385e-06,
+      "loss": 0.8719,
+      "output_text_mask_loss": 0.0717,
+      "predictor_learning_rate": 0.0004974473041291384,
+      "step": 724
+    },
+    {
+      "epoch": 0.07,
+      "gumbel_tau": 0.843,
+      "image_masks_loss": 0.0246,
+      "llm_learning_rate": 4.9743546524800665e-06,
+      "loss": 0.7955,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004974354652480066,
+      "step": 725
+    },
+    {
+      "epoch": 0.07,
+      "gumbel_tau": 0.8428,
+      "image_masks_loss": 0.0052,
+      "llm_learning_rate": 4.974235991187304e-06,
+      "loss": 0.7758,
+      "output_text_mask_loss": 0.0202,
+      "predictor_learning_rate": 0.0004974235991187304,
+      "step": 726
+    },
+    {
+      "epoch": 0.07,
+      "gumbel_tau": 0.8426,
+      "image_masks_loss": 0.0244,
+      "llm_learning_rate": 4.9741170574261625e-06,
+      "loss": 0.829,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.0004974117057426162,
+      "step": 727
+    },
+    {
+      "epoch": 0.07,
+      "gumbel_tau": 0.8424,
+      "image_masks_loss": 0.0215,
+      "llm_learning_rate": 4.973997851209739e-06,
+      "loss": 0.7391,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004973997851209738,
+      "step": 728
+    },
+    {
+      "epoch": 0.07,
+      "gumbel_tau": 0.8422,
+      "image_masks_loss": 0.0507,
+      "llm_learning_rate": 4.973878372551162e-06,
+      "loss": 0.898,
+      "output_text_mask_loss": 0.0381,
+      "predictor_learning_rate": 0.0004973878372551161,
+      "step": 729
+    },
+    {
+      "epoch": 0.07,
+      "gumbel_tau": 0.842,
+      "image_masks_loss": 0.0221,
+      "llm_learning_rate": 4.973758621463588e-06,
+      "loss": 0.7738,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.0004973758621463588,
+      "step": 730
+    },
+    {
+      "epoch": 0.07,
+      "gumbel_tau": 0.8418,
+      "image_masks_loss": 0.0163,
+      "llm_learning_rate": 4.973638597960205e-06,
+      "loss": 0.7549,
+      "output_text_mask_loss": 0.0107,
+      "predictor_learning_rate": 0.0004973638597960205,
+      "step": 731
+    },
+    {
+      "epoch": 0.07,
+      "gumbel_tau": 0.8416,
+      "image_masks_loss": 0.0175,
+      "llm_learning_rate": 4.97351830205423e-06,
+      "loss": 0.8025,
+      "output_text_mask_loss": 0.0824,
+      "predictor_learning_rate": 0.000497351830205423,
+      "step": 732
+    },
+    {
+      "epoch": 0.08,
+      "gumbel_tau": 0.8414,
+      "image_masks_loss": 0.038,
+      "llm_learning_rate": 4.973397733758909e-06,
+      "loss": 0.7491,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.000497339773375891,
+      "step": 733
+    },
+    {
+      "epoch": 0.08,
+      "gumbel_tau": 0.8412,
+      "image_masks_loss": 0.0206,
+      "llm_learning_rate": 4.973276893087522e-06,
+      "loss": 0.7841,
+      "output_text_mask_loss": 0.0058,
+      "predictor_learning_rate": 0.0004973276893087522,
+      "step": 734
+    },
+    {
+      "epoch": 0.08,
+      "gumbel_tau": 0.841,
+      "image_masks_loss": 0.0202,
+      "llm_learning_rate": 4.9731557800533734e-06,
+      "loss": 0.8116,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 0.0004973155780053373,
+      "step": 735
+    },
+    {
+      "epoch": 0.08,
+      "gumbel_tau": 0.8408,
+      "image_masks_loss": 0.0257,
+      "llm_learning_rate": 4.973034394669801e-06,
+      "loss": 0.8397,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.0004973034394669801,
+      "step": 736
+    },
+    {
+      "epoch": 0.08,
+      "gumbel_tau": 0.8406,
+      "image_masks_loss": 0.0355,
+      "llm_learning_rate": 4.972912736950174e-06,
+      "loss": 0.8547,
+      "output_text_mask_loss": 0.0404,
+      "predictor_learning_rate": 0.0004972912736950174,
+      "step": 737
+    },
+    {
+      "epoch": 0.08,
+      "gumbel_tau": 0.8404,
+      "image_masks_loss": 0.0133,
+      "llm_learning_rate": 4.972790806907887e-06,
+      "loss": 0.8669,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004972790806907887,
+      "step": 738
+    },
+    {
+      "epoch": 0.08,
+      "gumbel_tau": 0.8402,
+      "image_masks_loss": 0.0303,
+      "llm_learning_rate": 4.972668604556368e-06,
+      "loss": 0.7801,
+      "output_text_mask_loss": 0.0132,
+      "predictor_learning_rate": 0.0004972668604556369,
+      "step": 739
+    },
+    {
+      "epoch": 0.08,
+      "gumbel_tau": 0.84,
+      "image_masks_loss": 0.0156,
+      "llm_learning_rate": 4.972546129909077e-06,
+      "loss": 0.8469,
+      "output_text_mask_loss": 0.1122,
+      "predictor_learning_rate": 0.0004972546129909076,
+      "step": 740
+    },
+    {
+      "epoch": 0.08,
+      "gumbel_tau": 0.8398,
+      "image_masks_loss": 0.0223,
+      "llm_learning_rate": 4.9724233829794966e-06,
+      "loss": 0.7922,
+      "output_text_mask_loss": 0.001,
+      "predictor_learning_rate": 0.0004972423382979496,
+      "step": 741
+    },
+    {
+      "epoch": 0.08,
+      "gumbel_tau": 0.8396,
+      "image_masks_loss": 0.038,
+      "llm_learning_rate": 4.972300363781148e-06,
+      "loss": 0.8158,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.0004972300363781147,
+      "step": 742
+    },
+    {
+      "epoch": 0.08,
+      "gumbel_tau": 0.8394,
+      "image_masks_loss": 0.0519,
+      "llm_learning_rate": 4.972177072327574e-06,
+      "loss": 0.8342,
+      "output_text_mask_loss": 0.0092,
+      "predictor_learning_rate": 0.0004972177072327574,
+      "step": 743
+    },
+    {
+      "epoch": 0.08,
+      "gumbel_tau": 0.8392,
+      "image_masks_loss": 0.0427,
+      "llm_learning_rate": 4.972053508632356e-06,
+      "loss": 0.8887,
+      "output_text_mask_loss": 0.0374,
+      "predictor_learning_rate": 0.0004972053508632356,
+      "step": 744
+    },
+    {
+      "epoch": 0.08,
+      "gumbel_tau": 0.839,
+      "image_masks_loss": 0.0305,
+      "llm_learning_rate": 4.9719296727090986e-06,
+      "loss": 0.8401,
+      "output_text_mask_loss": 0.0193,
+      "predictor_learning_rate": 0.0004971929672709099,
+      "step": 745
+    },
+    {
+      "epoch": 0.08,
+      "gumbel_tau": 0.8388,
+      "image_masks_loss": 0.0203,
+      "llm_learning_rate": 4.9718055645714405e-06,
+      "loss": 0.7448,
+      "output_text_mask_loss": 0.0778,
+      "predictor_learning_rate": 0.000497180556457144,
+      "step": 746
+    },
+    {
+      "epoch": 0.08,
+      "gumbel_tau": 0.8386,
+      "image_masks_loss": 0.0383,
+      "llm_learning_rate": 4.971681184233047e-06,
+      "loss": 0.788,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 0.0004971681184233047,
+      "step": 747
+    },
+    {
+      "epoch": 0.08,
+      "gumbel_tau": 0.8384,
+      "image_masks_loss": 0.0362,
+      "llm_learning_rate": 4.971556531707616e-06,
+      "loss": 0.7757,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.0004971556531707616,
+      "step": 748
+    },
+    {
+      "epoch": 0.08,
+      "gumbel_tau": 0.8382,
+      "image_masks_loss": 0.0081,
+      "llm_learning_rate": 4.9714316070088756e-06,
+      "loss": 0.7744,
+      "output_text_mask_loss": 0.0214,
+      "predictor_learning_rate": 0.0004971431607008875,
+      "step": 749
+    },
+    {
+      "epoch": 0.08,
+      "gumbel_tau": 0.838,
+      "image_masks_loss": 0.0324,
+      "llm_learning_rate": 4.971306410150581e-06,
+      "loss": 0.7939,
+      "output_text_mask_loss": 0.0284,
+      "predictor_learning_rate": 0.000497130641015058,
+      "step": 750
+    },
+    {
+      "epoch": 0.08,
+      "gumbel_tau": 0.8378,
+      "image_masks_loss": 0.0282,
+      "llm_learning_rate": 4.97118094114652e-06,
+      "loss": 0.7698,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.000497118094114652,
+      "step": 751
+    },
+    {
+      "epoch": 0.08,
+      "gumbel_tau": 0.8376,
+      "image_masks_loss": 0.0183,
+      "llm_learning_rate": 4.97105520001051e-06,
+      "loss": 0.7788,
+      "output_text_mask_loss": 0.0149,
+      "predictor_learning_rate": 0.0004971055200010509,
+      "step": 752
+    },
+    {
+      "epoch": 0.08,
+      "gumbel_tau": 0.8374,
+      "image_masks_loss": 0.0269,
+      "llm_learning_rate": 4.970929186756396e-06,
+      "loss": 0.8066,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004970929186756396,
+      "step": 753
+    },
+    {
+      "epoch": 0.08,
+      "gumbel_tau": 0.8372,
+      "image_masks_loss": 0.0175,
+      "llm_learning_rate": 4.9708029013980565e-06,
+      "loss": 0.8434,
+      "output_text_mask_loss": 0.0376,
+      "predictor_learning_rate": 0.0004970802901398056,
+      "step": 754
+    },
+    {
+      "epoch": 0.08,
+      "gumbel_tau": 0.837,
+      "image_masks_loss": 0.0188,
+      "llm_learning_rate": 4.9706763439493985e-06,
+      "loss": 0.8164,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.0004970676343949398,
+      "step": 755
+    },
+    {
+      "epoch": 0.08,
+      "gumbel_tau": 0.8368,
+      "image_masks_loss": 0.0078,
+      "llm_learning_rate": 4.970549514424358e-06,
+      "loss": 0.7068,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004970549514424358,
+      "step": 756
+    },
+    {
+      "epoch": 0.08,
+      "gumbel_tau": 0.8366,
+      "image_masks_loss": 0.0207,
+      "llm_learning_rate": 4.970422412836903e-06,
+      "loss": 0.8491,
+      "output_text_mask_loss": 0.0322,
+      "predictor_learning_rate": 0.0004970422412836902,
+      "step": 757
+    },
+    {
+      "epoch": 0.08,
+      "gumbel_tau": 0.8364,
+      "image_masks_loss": 0.0089,
+      "llm_learning_rate": 4.970295039201029e-06,
+      "loss": 0.7364,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.0004970295039201029,
+      "step": 758
+    },
+    {
+      "epoch": 0.08,
+      "gumbel_tau": 0.8362,
+      "image_masks_loss": 0.0127,
+      "llm_learning_rate": 4.970167393530762e-06,
+      "loss": 0.759,
+      "output_text_mask_loss": 0.0404,
+      "predictor_learning_rate": 0.0004970167393530761,
+      "step": 759
+    },
+    {
+      "epoch": 0.08,
+      "gumbel_tau": 0.8361,
+      "image_masks_loss": 0.0408,
+      "llm_learning_rate": 4.970039475840159e-06,
+      "loss": 0.8758,
+      "output_text_mask_loss": 0.0095,
+      "predictor_learning_rate": 0.0004970039475840159,
+      "step": 760
+    },
+    {
+      "epoch": 0.08,
+      "gumbel_tau": 0.8359,
+      "image_masks_loss": 0.0374,
+      "llm_learning_rate": 4.969911286143309e-06,
+      "loss": 0.6774,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.0004969911286143308,
+      "step": 761
+    },
+    {
+      "epoch": 0.08,
+      "gumbel_tau": 0.8357,
+      "image_masks_loss": 0.0058,
+      "llm_learning_rate": 4.969782824454325e-06,
+      "loss": 0.8369,
+      "output_text_mask_loss": 0.0162,
+      "predictor_learning_rate": 0.0004969782824454325,
+      "step": 762
+    },
+    {
+      "epoch": 0.08,
+      "gumbel_tau": 0.8355,
+      "image_masks_loss": 0.0271,
+      "llm_learning_rate": 4.969654090787356e-06,
+      "loss": 0.746,
+      "output_text_mask_loss": 0.0214,
+      "predictor_learning_rate": 0.0004969654090787355,
+      "step": 763
+    },
+    {
+      "epoch": 0.08,
+      "gumbel_tau": 0.8353,
+      "image_masks_loss": 0.0349,
+      "llm_learning_rate": 4.969525085156578e-06,
+      "loss": 0.7757,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004969525085156578,
+      "step": 764
+    },
+    {
+      "epoch": 0.08,
+      "gumbel_tau": 0.8351,
+      "image_masks_loss": 0.0126,
+      "llm_learning_rate": 4.969395807576196e-06,
+      "loss": 0.8237,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.0004969395807576196,
+      "step": 765
+    },
+    {
+      "epoch": 0.08,
+      "gumbel_tau": 0.8349,
+      "image_masks_loss": 0.0203,
+      "llm_learning_rate": 4.969266258060448e-06,
+      "loss": 0.7281,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004969266258060448,
+      "step": 766
+    },
+    {
+      "epoch": 0.08,
+      "gumbel_tau": 0.8347,
+      "image_masks_loss": 0.0252,
+      "llm_learning_rate": 4.969136436623599e-06,
+      "loss": 0.8226,
+      "output_text_mask_loss": 0.0076,
+      "predictor_learning_rate": 0.0004969136436623599,
+      "step": 767
+    },
+    {
+      "epoch": 0.08,
+      "gumbel_tau": 0.8345,
+      "image_masks_loss": 0.025,
+      "llm_learning_rate": 4.969006343279946e-06,
+      "loss": 0.7179,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 0.0004969006343279946,
+      "step": 768
+    },
+    {
+      "epoch": 0.08,
+      "gumbel_tau": 0.8343,
+      "image_masks_loss": 0.019,
+      "llm_learning_rate": 4.968875978043815e-06,
+      "loss": 0.6959,
+      "output_text_mask_loss": 0.021,
+      "predictor_learning_rate": 0.0004968875978043815,
+      "step": 769
+    },
+    {
+      "epoch": 0.08,
+      "gumbel_tau": 0.8341,
+      "image_masks_loss": 0.0157,
+      "llm_learning_rate": 4.968745340929563e-06,
+      "loss": 0.8561,
+      "output_text_mask_loss": 0.0181,
+      "predictor_learning_rate": 0.0004968745340929563,
+      "step": 770
+    },
+    {
+      "epoch": 0.08,
+      "gumbel_tau": 0.8339,
+      "image_masks_loss": 0.0057,
+      "llm_learning_rate": 4.968614431951574e-06,
+      "loss": 0.6856,
+      "output_text_mask_loss": 0.0919,
+      "predictor_learning_rate": 0.0004968614431951574,
+      "step": 771
+    },
+    {
+      "epoch": 0.08,
+      "gumbel_tau": 0.8337,
+      "image_masks_loss": 0.0078,
+      "llm_learning_rate": 4.968483251124265e-06,
+      "loss": 0.7013,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.0004968483251124265,
+      "step": 772
+    },
+    {
+      "epoch": 0.08,
+      "gumbel_tau": 0.8335,
+      "image_masks_loss": 0.0053,
+      "llm_learning_rate": 4.968351798462082e-06,
+      "loss": 0.7933,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 0.0004968351798462083,
+      "step": 773
+    },
+    {
+      "epoch": 0.08,
+      "gumbel_tau": 0.8333,
+      "image_masks_loss": 0.0408,
+      "llm_learning_rate": 4.9682200739795025e-06,
+      "loss": 0.8579,
+      "output_text_mask_loss": 0.0162,
+      "predictor_learning_rate": 0.0004968220073979502,
+      "step": 774
+    },
+    {
+      "epoch": 0.08,
+      "gumbel_tau": 0.8331,
+      "image_masks_loss": 0.037,
+      "llm_learning_rate": 4.968088077691029e-06,
+      "loss": 0.8151,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.0004968088077691029,
+      "step": 775
+    },
+    {
+      "epoch": 0.08,
+      "gumbel_tau": 0.8329,
+      "image_masks_loss": 0.0246,
+      "llm_learning_rate": 4.9679558096112e-06,
+      "loss": 0.8366,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00049679558096112,
+      "step": 776
+    },
+    {
+      "epoch": 0.08,
+      "gumbel_tau": 0.8327,
+      "image_masks_loss": 0.0175,
+      "llm_learning_rate": 4.96782326975458e-06,
+      "loss": 0.8274,
+      "output_text_mask_loss": 0.0023,
+      "predictor_learning_rate": 0.000496782326975458,
+      "step": 777
+    },
+    {
+      "epoch": 0.08,
+      "gumbel_tau": 0.8325,
+      "image_masks_loss": 0.0167,
+      "llm_learning_rate": 4.967690458135764e-06,
+      "loss": 0.7309,
+      "output_text_mask_loss": 0.0039,
+      "predictor_learning_rate": 0.0004967690458135764,
+      "step": 778
+    },
+    {
+      "epoch": 0.08,
+      "gumbel_tau": 0.8323,
+      "image_masks_loss": 0.0687,
+      "llm_learning_rate": 4.967557374769378e-06,
+      "loss": 0.8401,
+      "output_text_mask_loss": 0.0015,
+      "predictor_learning_rate": 0.0004967557374769378,
+      "step": 779
+    },
+    {
+      "epoch": 0.08,
+      "gumbel_tau": 0.8321,
+      "image_masks_loss": 0.0259,
+      "llm_learning_rate": 4.9674240196700785e-06,
+      "loss": 0.8336,
+      "output_text_mask_loss": 0.0341,
+      "predictor_learning_rate": 0.0004967424019670079,
+      "step": 780
+    },
+    {
+      "epoch": 0.08,
+      "gumbel_tau": 0.8319,
+      "image_masks_loss": 0.0083,
+      "llm_learning_rate": 4.967290392852549e-06,
+      "loss": 0.8623,
+      "output_text_mask_loss": 0.0526,
+      "predictor_learning_rate": 0.000496729039285255,
+      "step": 781
+    },
+    {
+      "epoch": 0.08,
+      "gumbel_tau": 0.8317,
+      "image_masks_loss": 0.0328,
+      "llm_learning_rate": 4.967156494331507e-06,
+      "loss": 0.8437,
+      "output_text_mask_loss": 0.0223,
+      "predictor_learning_rate": 0.0004967156494331507,
+      "step": 782
+    },
+    {
+      "epoch": 0.08,
+      "gumbel_tau": 0.8315,
+      "image_masks_loss": 0.0248,
+      "llm_learning_rate": 4.967022324121697e-06,
+      "loss": 0.8236,
+      "output_text_mask_loss": 0.0248,
+      "predictor_learning_rate": 0.0004967022324121696,
+      "step": 783
+    },
+    {
+      "epoch": 0.08,
+      "gumbel_tau": 0.8313,
+      "image_masks_loss": 0.0213,
+      "llm_learning_rate": 4.9668878822378914e-06,
+      "loss": 0.7861,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.0004966887882237891,
+      "step": 784
+    },
+    {
+      "epoch": 0.08,
+      "gumbel_tau": 0.8311,
+      "image_masks_loss": 0.0213,
+      "llm_learning_rate": 4.966753168694899e-06,
+      "loss": 0.7425,
+      "output_text_mask_loss": 0.0101,
+      "predictor_learning_rate": 0.0004966753168694899,
+      "step": 785
+    },
+    {
+      "epoch": 0.08,
+      "gumbel_tau": 0.8309,
+      "image_masks_loss": 0.0216,
+      "llm_learning_rate": 4.966618183507552e-06,
+      "loss": 0.8368,
+      "output_text_mask_loss": 0.0151,
+      "predictor_learning_rate": 0.0004966618183507552,
+      "step": 786
+    },
+    {
+      "epoch": 0.08,
+      "gumbel_tau": 0.8307,
+      "image_masks_loss": 0.0216,
+      "llm_learning_rate": 4.966482926690718e-06,
+      "loss": 0.8306,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.0004966482926690717,
+      "step": 787
+    },
+    {
+      "epoch": 0.08,
+      "gumbel_tau": 0.8305,
+      "image_masks_loss": 0.0142,
+      "llm_learning_rate": 4.966347398259288e-06,
+      "loss": 0.8152,
+      "output_text_mask_loss": 0.0839,
+      "predictor_learning_rate": 0.0004966347398259288,
+      "step": 788
+    },
+    {
+      "epoch": 0.08,
+      "gumbel_tau": 0.8304,
+      "image_masks_loss": 0.0471,
+      "llm_learning_rate": 4.9662115982281914e-06,
+      "loss": 0.8807,
+      "output_text_mask_loss": 0.0149,
+      "predictor_learning_rate": 0.0004966211598228191,
+      "step": 789
+    },
+    {
+      "epoch": 0.08,
+      "gumbel_tau": 0.8302,
+      "image_masks_loss": 0.0215,
+      "llm_learning_rate": 4.966075526612379e-06,
+      "loss": 0.7774,
+      "output_text_mask_loss": 0.074,
+      "predictor_learning_rate": 0.0004966075526612379,
+      "step": 790
+    },
+    {
+      "epoch": 0.08,
+      "gumbel_tau": 0.83,
+      "image_masks_loss": 0.0229,
+      "llm_learning_rate": 4.9659391834268365e-06,
+      "loss": 0.8541,
+      "output_text_mask_loss": 0.0916,
+      "predictor_learning_rate": 0.0004965939183426836,
+      "step": 791
+    },
+    {
+      "epoch": 0.08,
+      "gumbel_tau": 0.8298,
+      "image_masks_loss": 0.0298,
+      "llm_learning_rate": 4.9658025686865795e-06,
+      "loss": 0.832,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004965802568686579,
+      "step": 792
+    },
+    {
+      "epoch": 0.08,
+      "gumbel_tau": 0.8296,
+      "image_masks_loss": 0.0364,
+      "llm_learning_rate": 4.96566568240665e-06,
+      "loss": 0.8092,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.000496566568240665,
+      "step": 793
+    },
+    {
+      "epoch": 0.08,
+      "gumbel_tau": 0.8294,
+      "image_masks_loss": 0.0144,
+      "llm_learning_rate": 4.965528524602124e-06,
+      "loss": 0.7413,
+      "output_text_mask_loss": 0.0954,
+      "predictor_learning_rate": 0.0004965528524602124,
+      "step": 794
+    },
+    {
+      "epoch": 0.08,
+      "gumbel_tau": 0.8292,
+      "image_masks_loss": 0.0357,
+      "llm_learning_rate": 4.965391095288106e-06,
+      "loss": 0.7795,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.0004965391095288106,
+      "step": 795
+    },
+    {
+      "epoch": 0.08,
+      "gumbel_tau": 0.829,
+      "image_masks_loss": 0.0431,
+      "llm_learning_rate": 4.965253394479728e-06,
+      "loss": 0.8008,
+      "output_text_mask_loss": 0.0041,
+      "predictor_learning_rate": 0.0004965253394479727,
+      "step": 796
+    },
+    {
+      "epoch": 0.08,
+      "gumbel_tau": 0.8288,
+      "image_masks_loss": 0.0153,
+      "llm_learning_rate": 4.965115422192155e-06,
+      "loss": 0.8041,
+      "output_text_mask_loss": 0.0843,
+      "predictor_learning_rate": 0.0004965115422192154,
+      "step": 797
+    },
+    {
+      "epoch": 0.08,
+      "gumbel_tau": 0.8286,
+      "image_masks_loss": 0.0205,
+      "llm_learning_rate": 4.9649771784405816e-06,
+      "loss": 0.7155,
+      "output_text_mask_loss": 0.0298,
+      "predictor_learning_rate": 0.0004964977178440581,
+      "step": 798
+    },
+    {
+      "epoch": 0.08,
+      "gumbel_tau": 0.8284,
+      "image_masks_loss": 0.0193,
+      "llm_learning_rate": 4.96483866324023e-06,
+      "loss": 0.8369,
+      "output_text_mask_loss": 0.0288,
+      "predictor_learning_rate": 0.000496483866324023,
+      "step": 799
+    },
+    {
+      "epoch": 0.08,
+      "gumbel_tau": 0.8282,
+      "image_masks_loss": 0.0227,
+      "llm_learning_rate": 4.964699876606356e-06,
+      "loss": 0.7948,
+      "output_text_mask_loss": 0.1083,
+      "predictor_learning_rate": 0.0004964699876606356,
+      "step": 800
+    },
+    {
+      "epoch": 0.08,
+      "gumbel_tau": 0.828,
+      "image_masks_loss": 0.0237,
+      "llm_learning_rate": 4.964560818554241e-06,
+      "loss": 0.8668,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.000496456081855424,
+      "step": 801
+    },
+    {
+      "epoch": 0.08,
+      "gumbel_tau": 0.8278,
+      "image_masks_loss": 0.0127,
+      "llm_learning_rate": 4.964421489099198e-06,
+      "loss": 0.8655,
+      "output_text_mask_loss": 0.0851,
+      "predictor_learning_rate": 0.0004964421489099198,
+      "step": 802
+    },
+    {
+      "epoch": 0.08,
+      "gumbel_tau": 0.8276,
+      "image_masks_loss": 0.0273,
+      "llm_learning_rate": 4.964281888256573e-06,
+      "loss": 0.8092,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 0.0004964281888256572,
+      "step": 803
+    },
+    {
+      "epoch": 0.08,
+      "gumbel_tau": 0.8274,
+      "image_masks_loss": 0.019,
+      "llm_learning_rate": 4.964142016041736e-06,
+      "loss": 0.9016,
+      "output_text_mask_loss": 0.0465,
+      "predictor_learning_rate": 0.0004964142016041736,
+      "step": 804
+    },
+    {
+      "epoch": 0.08,
+      "gumbel_tau": 0.8272,
+      "image_masks_loss": 0.024,
+      "llm_learning_rate": 4.9640018724700925e-06,
+      "loss": 0.6737,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004964001872470092,
+      "step": 805
+    },
+    {
+      "epoch": 0.08,
+      "gumbel_tau": 0.827,
+      "image_masks_loss": 0.0096,
+      "llm_learning_rate": 4.9638614575570735e-06,
+      "loss": 0.7796,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.0004963861457557073,
+      "step": 806
+    },
+    {
+      "epoch": 0.08,
+      "gumbel_tau": 0.8268,
+      "image_masks_loss": 0.034,
+      "llm_learning_rate": 4.9637207713181435e-06,
+      "loss": 0.9148,
+      "output_text_mask_loss": 0.0015,
+      "predictor_learning_rate": 0.0004963720771318143,
+      "step": 807
+    },
+    {
+      "epoch": 0.08,
+      "gumbel_tau": 0.8266,
+      "image_masks_loss": 0.011,
+      "llm_learning_rate": 4.963579813768795e-06,
+      "loss": 0.8138,
+      "output_text_mask_loss": 0.1122,
+      "predictor_learning_rate": 0.0004963579813768794,
+      "step": 808
+    },
+    {
+      "epoch": 0.08,
+      "gumbel_tau": 0.8264,
+      "image_masks_loss": 0.0187,
+      "llm_learning_rate": 4.963438584924549e-06,
+      "loss": 0.7513,
+      "output_text_mask_loss": 0.0519,
+      "predictor_learning_rate": 0.0004963438584924549,
+      "step": 809
+    },
+    {
+      "epoch": 0.08,
+      "gumbel_tau": 0.8262,
+      "image_masks_loss": 0.0214,
+      "llm_learning_rate": 4.963297084800959e-06,
+      "loss": 0.7724,
+      "output_text_mask_loss": 0.0429,
+      "predictor_learning_rate": 0.0004963297084800959,
+      "step": 810
+    },
+    {
+      "epoch": 0.08,
+      "gumbel_tau": 0.8261,
+      "image_masks_loss": 0.0234,
+      "llm_learning_rate": 4.963155313413608e-06,
+      "loss": 0.8317,
+      "output_text_mask_loss": 0.0626,
+      "predictor_learning_rate": 0.0004963155313413608,
+      "step": 811
+    },
+    {
+      "epoch": 0.08,
+      "gumbel_tau": 0.8259,
+      "image_masks_loss": 0.0326,
+      "llm_learning_rate": 4.963013270778106e-06,
+      "loss": 0.914,
+      "output_text_mask_loss": 0.0519,
+      "predictor_learning_rate": 0.0004963013270778106,
+      "step": 812
+    },
+    {
+      "epoch": 0.08,
+      "gumbel_tau": 0.8257,
+      "image_masks_loss": 0.0542,
+      "llm_learning_rate": 4.962870956910098e-06,
+      "loss": 0.8923,
+      "output_text_mask_loss": 0.0549,
+      "predictor_learning_rate": 0.0004962870956910098,
+      "step": 813
+    },
+    {
+      "epoch": 0.08,
+      "gumbel_tau": 0.8255,
+      "image_masks_loss": 0.023,
+      "llm_learning_rate": 4.962728371825254e-06,
+      "loss": 0.7638,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004962728371825253,
+      "step": 814
+    },
+    {
+      "epoch": 0.08,
+      "gumbel_tau": 0.8253,
+      "image_masks_loss": 0.0408,
+      "llm_learning_rate": 4.962585515539275e-06,
+      "loss": 0.8061,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004962585515539275,
+      "step": 815
+    },
+    {
+      "epoch": 0.08,
+      "gumbel_tau": 0.8251,
+      "image_masks_loss": 0.0192,
+      "llm_learning_rate": 4.962442388067896e-06,
+      "loss": 0.7434,
+      "output_text_mask_loss": 0.0275,
+      "predictor_learning_rate": 0.0004962442388067895,
+      "step": 816
+    },
+    {
+      "epoch": 0.08,
+      "gumbel_tau": 0.8249,
+      "image_masks_loss": 0.0097,
+      "llm_learning_rate": 4.9622989894268745e-06,
+      "loss": 0.818,
+      "output_text_mask_loss": 0.0999,
+      "predictor_learning_rate": 0.0004962298989426875,
+      "step": 817
+    },
+    {
+      "epoch": 0.08,
+      "gumbel_tau": 0.8247,
+      "image_masks_loss": 0.0157,
+      "llm_learning_rate": 4.962155319632005e-06,
+      "loss": 0.8026,
+      "output_text_mask_loss": 0.0137,
+      "predictor_learning_rate": 0.0004962155319632004,
+      "step": 818
+    },
+    {
+      "epoch": 0.08,
+      "gumbel_tau": 0.8245,
+      "image_masks_loss": 0.042,
+      "llm_learning_rate": 4.962011378699107e-06,
+      "loss": 0.7401,
+      "output_text_mask_loss": 0.0124,
+      "predictor_learning_rate": 0.0004962011378699106,
+      "step": 819
+    },
+    {
+      "epoch": 0.08,
+      "gumbel_tau": 0.8243,
+      "image_masks_loss": 0.0252,
+      "llm_learning_rate": 4.961867166644032e-06,
+      "loss": 0.7269,
+      "output_text_mask_loss": 0.0202,
+      "predictor_learning_rate": 0.0004961867166644031,
+      "step": 820
+    },
+    {
+      "epoch": 0.08,
+      "gumbel_tau": 0.8241,
+      "image_masks_loss": 0.0515,
+      "llm_learning_rate": 4.96172268348266e-06,
+      "loss": 0.7039,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.000496172268348266,
+      "step": 821
+    },
+    {
+      "epoch": 0.08,
+      "gumbel_tau": 0.8239,
+      "image_masks_loss": 0.0165,
+      "llm_learning_rate": 4.9615779292309045e-06,
+      "loss": 0.7814,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004961577929230904,
+      "step": 822
+    },
+    {
+      "epoch": 0.08,
+      "gumbel_tau": 0.8237,
+      "image_masks_loss": 0.0177,
+      "llm_learning_rate": 4.961432903904704e-06,
+      "loss": 0.8027,
+      "output_text_mask_loss": 0.0145,
+      "predictor_learning_rate": 0.0004961432903904703,
+      "step": 823
+    },
+    {
+      "epoch": 0.08,
+      "gumbel_tau": 0.8235,
+      "image_masks_loss": 0.0079,
+      "llm_learning_rate": 4.961287607520029e-06,
+      "loss": 0.9118,
+      "output_text_mask_loss": 0.0305,
+      "predictor_learning_rate": 0.0004961287607520028,
+      "step": 824
+    },
+    {
+      "epoch": 0.08,
+      "gumbel_tau": 0.8233,
+      "image_masks_loss": 0.006,
+      "llm_learning_rate": 4.96114204009288e-06,
+      "loss": 0.7812,
+      "output_text_mask_loss": 0.0227,
+      "predictor_learning_rate": 0.000496114204009288,
+      "step": 825
+    },
+    {
+      "epoch": 0.08,
+      "gumbel_tau": 0.8231,
+      "image_masks_loss": 0.0322,
+      "llm_learning_rate": 4.9609962016392885e-06,
+      "loss": 0.8184,
+      "output_text_mask_loss": 0.0645,
+      "predictor_learning_rate": 0.0004960996201639288,
+      "step": 826
+    },
+    {
+      "epoch": 0.08,
+      "gumbel_tau": 0.8229,
+      "image_masks_loss": 0.0175,
+      "llm_learning_rate": 4.960850092175313e-06,
+      "loss": 0.7856,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004960850092175313,
+      "step": 827
+    },
+    {
+      "epoch": 0.08,
+      "gumbel_tau": 0.8227,
+      "image_masks_loss": 0.037,
+      "llm_learning_rate": 4.960703711717045e-06,
+      "loss": 0.8443,
+      "output_text_mask_loss": 0.0071,
+      "predictor_learning_rate": 0.0004960703711717044,
+      "step": 828
+    },
+    {
+      "epoch": 0.08,
+      "gumbel_tau": 0.8226,
+      "image_masks_loss": 0.0246,
+      "llm_learning_rate": 4.960557060280602e-06,
+      "loss": 0.7818,
+      "output_text_mask_loss": 0.033,
+      "predictor_learning_rate": 0.0004960557060280601,
+      "step": 829
+    },
+    {
+      "epoch": 0.09,
+      "gumbel_tau": 0.8224,
+      "image_masks_loss": 0.0172,
+      "llm_learning_rate": 4.9604101378821355e-06,
+      "loss": 0.8261,
+      "output_text_mask_loss": 0.033,
+      "predictor_learning_rate": 0.0004960410137882135,
+      "step": 830
+    },
+    {
+      "epoch": 0.09,
+      "gumbel_tau": 0.8222,
+      "image_masks_loss": 0.0319,
+      "llm_learning_rate": 4.9602629445378244e-06,
+      "loss": 0.7237,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004960262944537824,
+      "step": 831
+    },
+    {
+      "epoch": 0.09,
+      "gumbel_tau": 0.822,
+      "image_masks_loss": 0.0319,
+      "llm_learning_rate": 4.9601154802638774e-06,
+      "loss": 0.7705,
+      "output_text_mask_loss": 0.0938,
+      "predictor_learning_rate": 0.0004960115480263877,
+      "step": 832
+    },
+    {
+      "epoch": 0.09,
+      "gumbel_tau": 0.8218,
+      "image_masks_loss": 0.0198,
+      "llm_learning_rate": 4.959967745076535e-06,
+      "loss": 0.8051,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004959967745076534,
+      "step": 833
+    },
+    {
+      "epoch": 0.09,
+      "gumbel_tau": 0.8216,
+      "image_masks_loss": 0.0275,
+      "llm_learning_rate": 4.959819738992064e-06,
+      "loss": 0.8327,
+      "output_text_mask_loss": 0.0216,
+      "predictor_learning_rate": 0.0004959819738992063,
+      "step": 834
+    },
+    {
+      "epoch": 0.09,
+      "gumbel_tau": 0.8214,
+      "image_masks_loss": 0.0065,
+      "llm_learning_rate": 4.959671462026765e-06,
+      "loss": 0.7601,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004959671462026764,
+      "step": 835
+    },
+    {
+      "epoch": 0.09,
+      "gumbel_tau": 0.8212,
+      "image_masks_loss": 0.0315,
+      "llm_learning_rate": 4.959522914196964e-06,
+      "loss": 0.7589,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 0.0004959522914196965,
+      "step": 836
+    },
+    {
+      "epoch": 0.09,
+      "gumbel_tau": 0.821,
+      "image_masks_loss": 0.0376,
+      "llm_learning_rate": 4.959374095519024e-06,
+      "loss": 0.8977,
+      "output_text_mask_loss": 0.0923,
+      "predictor_learning_rate": 0.0004959374095519024,
+      "step": 837
+    },
+    {
+      "epoch": 0.09,
+      "gumbel_tau": 0.8208,
+      "image_masks_loss": 0.0157,
+      "llm_learning_rate": 4.959225006009329e-06,
+      "loss": 0.7417,
+      "output_text_mask_loss": 0.0294,
+      "predictor_learning_rate": 0.0004959225006009329,
+      "step": 838
+    },
+    {
+      "epoch": 0.09,
+      "gumbel_tau": 0.8206,
+      "image_masks_loss": 0.0047,
+      "llm_learning_rate": 4.9590756456842994e-06,
+      "loss": 0.8413,
+      "output_text_mask_loss": 0.174,
+      "predictor_learning_rate": 0.0004959075645684298,
+      "step": 839
+    },
+    {
+      "epoch": 0.09,
+      "gumbel_tau": 0.8204,
+      "image_masks_loss": 0.0217,
+      "llm_learning_rate": 4.958926014560382e-06,
+      "loss": 0.8432,
+      "output_text_mask_loss": 0.0376,
+      "predictor_learning_rate": 0.0004958926014560382,
+      "step": 840
+    },
+    {
+      "epoch": 0.09,
+      "gumbel_tau": 0.8202,
+      "image_masks_loss": 0.0277,
+      "llm_learning_rate": 4.958776112654055e-06,
+      "loss": 0.8376,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004958776112654055,
+      "step": 841
+    },
+    {
+      "epoch": 0.09,
+      "gumbel_tau": 0.82,
+      "image_masks_loss": 0.0129,
+      "llm_learning_rate": 4.958625939981826e-06,
+      "loss": 0.813,
+      "output_text_mask_loss": 0.0065,
+      "predictor_learning_rate": 0.0004958625939981826,
+      "step": 842
+    },
+    {
+      "epoch": 0.09,
+      "gumbel_tau": 0.8198,
+      "image_masks_loss": 0.0284,
+      "llm_learning_rate": 4.958475496560231e-06,
+      "loss": 0.8305,
+      "output_text_mask_loss": 0.034,
+      "predictor_learning_rate": 0.0004958475496560231,
+      "step": 843
+    },
+    {
+      "epoch": 0.09,
+      "gumbel_tau": 0.8196,
+      "image_masks_loss": 0.0214,
+      "llm_learning_rate": 4.95832478240584e-06,
+      "loss": 0.7789,
+      "output_text_mask_loss": 0.0763,
+      "predictor_learning_rate": 0.0004958324782405839,
+      "step": 844
+    },
+    {
+      "epoch": 0.09,
+      "gumbel_tau": 0.8195,
+      "image_masks_loss": 0.0182,
+      "llm_learning_rate": 4.958173797535247e-06,
+      "loss": 0.7053,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.0004958173797535246,
+      "step": 845
+    },
+    {
+      "epoch": 0.09,
+      "gumbel_tau": 0.8193,
+      "image_masks_loss": 0.0282,
+      "llm_learning_rate": 4.95802254196508e-06,
+      "loss": 0.8643,
+      "output_text_mask_loss": 0.0248,
+      "predictor_learning_rate": 0.000495802254196508,
+      "step": 846
+    },
+    {
+      "epoch": 0.09,
+      "gumbel_tau": 0.8191,
+      "image_masks_loss": 0.0202,
+      "llm_learning_rate": 4.9578710157119956e-06,
+      "loss": 0.7618,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.0004957871015711995,
+      "step": 847
+    },
+    {
+      "epoch": 0.09,
+      "gumbel_tau": 0.8189,
+      "image_masks_loss": 0.0139,
+      "llm_learning_rate": 4.9577192187926806e-06,
+      "loss": 0.8313,
+      "output_text_mask_loss": 0.1053,
+      "predictor_learning_rate": 0.000495771921879268,
+      "step": 848
+    },
+    {
+      "epoch": 0.09,
+      "gumbel_tau": 0.8187,
+      "image_masks_loss": 0.0119,
+      "llm_learning_rate": 4.957567151223851e-06,
+      "loss": 0.7428,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 0.0004957567151223851,
+      "step": 849
+    },
+    {
+      "epoch": 0.09,
+      "gumbel_tau": 0.8185,
+      "image_masks_loss": 0.0226,
+      "llm_learning_rate": 4.957414813022253e-06,
+      "loss": 0.7132,
+      "output_text_mask_loss": 0.0004,
+      "predictor_learning_rate": 0.0004957414813022252,
+      "step": 850
+    },
+    {
+      "epoch": 0.09,
+      "gumbel_tau": 0.8183,
+      "image_masks_loss": 0.0322,
+      "llm_learning_rate": 4.9572622042046614e-06,
+      "loss": 0.6715,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.0004957262204204661,
+      "step": 851
+    },
+    {
+      "epoch": 0.09,
+      "gumbel_tau": 0.8181,
+      "image_masks_loss": 0.0112,
+      "llm_learning_rate": 4.957109324787883e-06,
+      "loss": 0.7183,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.0004957109324787883,
+      "step": 852
+    },
+    {
+      "epoch": 0.09,
+      "gumbel_tau": 0.8179,
+      "image_masks_loss": 0.0275,
+      "llm_learning_rate": 4.956956174788753e-06,
+      "loss": 0.8682,
+      "output_text_mask_loss": 0.0299,
+      "predictor_learning_rate": 0.0004956956174788753,
+      "step": 853
+    },
+    {
+      "epoch": 0.09,
+      "gumbel_tau": 0.8177,
+      "image_masks_loss": 0.0181,
+      "llm_learning_rate": 4.956802754224136e-06,
+      "loss": 0.8596,
+      "output_text_mask_loss": 0.0248,
+      "predictor_learning_rate": 0.0004956802754224136,
+      "step": 854
+    },
+    {
+      "epoch": 0.09,
+      "gumbel_tau": 0.8175,
+      "image_masks_loss": 0.0092,
+      "llm_learning_rate": 4.956649063110927e-06,
+      "loss": 0.7331,
+      "output_text_mask_loss": 0.0511,
+      "predictor_learning_rate": 0.0004956649063110926,
+      "step": 855
+    },
+    {
+      "epoch": 0.09,
+      "gumbel_tau": 0.8173,
+      "image_masks_loss": 0.0129,
+      "llm_learning_rate": 4.956495101466051e-06,
+      "loss": 0.7818,
+      "output_text_mask_loss": 0.0641,
+      "predictor_learning_rate": 0.0004956495101466051,
+      "step": 856
+    },
+    {
+      "epoch": 0.09,
+      "gumbel_tau": 0.8171,
+      "image_masks_loss": 0.0381,
+      "llm_learning_rate": 4.956340869306464e-06,
+      "loss": 0.8031,
+      "output_text_mask_loss": 0.1434,
+      "predictor_learning_rate": 0.0004956340869306464,
+      "step": 857
+    },
+    {
+      "epoch": 0.09,
+      "gumbel_tau": 0.8169,
+      "image_masks_loss": 0.0176,
+      "llm_learning_rate": 4.956186366649148e-06,
+      "loss": 0.8427,
+      "output_text_mask_loss": 0.017,
+      "predictor_learning_rate": 0.0004956186366649148,
+      "step": 858
+    },
+    {
+      "epoch": 0.09,
+      "gumbel_tau": 0.8168,
+      "image_masks_loss": 0.0385,
+      "llm_learning_rate": 4.956031593511119e-06,
+      "loss": 0.7833,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.0004956031593511119,
+      "step": 859
+    },
+    {
+      "epoch": 0.09,
+      "gumbel_tau": 0.8166,
+      "image_masks_loss": 0.0242,
+      "llm_learning_rate": 4.955876549909421e-06,
+      "loss": 0.8398,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.000495587654990942,
+      "step": 860
+    },
+    {
+      "epoch": 0.09,
+      "gumbel_tau": 0.8164,
+      "image_masks_loss": 0.0256,
+      "llm_learning_rate": 4.955721235861126e-06,
+      "loss": 0.8122,
+      "output_text_mask_loss": 0.0725,
+      "predictor_learning_rate": 0.0004955721235861126,
+      "step": 861
+    },
+    {
+      "epoch": 0.09,
+      "gumbel_tau": 0.8162,
+      "image_masks_loss": 0.0313,
+      "llm_learning_rate": 4.955565651383339e-06,
+      "loss": 0.9379,
+      "output_text_mask_loss": 0.0399,
+      "predictor_learning_rate": 0.0004955565651383339,
+      "step": 862
+    },
+    {
+      "epoch": 0.09,
+      "gumbel_tau": 0.816,
+      "image_masks_loss": 0.0492,
+      "llm_learning_rate": 4.955409796493192e-06,
+      "loss": 0.6948,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004955409796493192,
+      "step": 863
+    },
+    {
+      "epoch": 0.09,
+      "gumbel_tau": 0.8158,
+      "image_masks_loss": 0.0324,
+      "llm_learning_rate": 4.95525367120785e-06,
+      "loss": 0.8131,
+      "output_text_mask_loss": 0.0271,
+      "predictor_learning_rate": 0.000495525367120785,
+      "step": 864
+    },
+    {
+      "epoch": 0.09,
+      "gumbel_tau": 0.8156,
+      "image_masks_loss": 0.0227,
+      "llm_learning_rate": 4.955097275544505e-06,
+      "loss": 0.8416,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004955097275544505,
+      "step": 865
+    },
+    {
+      "epoch": 0.09,
+      "gumbel_tau": 0.8154,
+      "image_masks_loss": 0.0154,
+      "llm_learning_rate": 4.954940609520379e-06,
+      "loss": 0.7159,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 0.0004954940609520379,
+      "step": 866
+    },
+    {
+      "epoch": 0.09,
+      "gumbel_tau": 0.8152,
+      "image_masks_loss": 0.0067,
+      "llm_learning_rate": 4.954783673152725e-06,
+      "loss": 0.7359,
+      "output_text_mask_loss": 0.0198,
+      "predictor_learning_rate": 0.0004954783673152725,
+      "step": 867
+    },
+    {
+      "epoch": 0.09,
+      "gumbel_tau": 0.815,
+      "image_masks_loss": 0.0137,
+      "llm_learning_rate": 4.954626466458826e-06,
+      "loss": 0.7945,
+      "output_text_mask_loss": 0.0364,
+      "predictor_learning_rate": 0.0004954626466458826,
+      "step": 868
+    },
+    {
+      "epoch": 0.09,
+      "gumbel_tau": 0.8148,
+      "image_masks_loss": 0.0144,
+      "llm_learning_rate": 4.9544689894559915e-06,
+      "loss": 0.6948,
+      "output_text_mask_loss": 0.0488,
+      "predictor_learning_rate": 0.0004954468989455991,
+      "step": 869
+    },
+    {
+      "epoch": 0.09,
+      "gumbel_tau": 0.8146,
+      "image_masks_loss": 0.025,
+      "llm_learning_rate": 4.954311242161566e-06,
+      "loss": 0.8484,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.0004954311242161565,
+      "step": 870
+    },
+    {
+      "epoch": 0.09,
+      "gumbel_tau": 0.8144,
+      "image_masks_loss": 0.0399,
+      "llm_learning_rate": 4.954153224592919e-06,
+      "loss": 0.712,
+      "output_text_mask_loss": 0.0618,
+      "predictor_learning_rate": 0.0004954153224592919,
+      "step": 871
+    },
+    {
+      "epoch": 0.09,
+      "gumbel_tau": 0.8143,
+      "image_masks_loss": 0.0093,
+      "llm_learning_rate": 4.953994936767453e-06,
+      "loss": 0.7254,
+      "output_text_mask_loss": 0.0278,
+      "predictor_learning_rate": 0.0004953994936767452,
+      "step": 872
+    },
+    {
+      "epoch": 0.09,
+      "gumbel_tau": 0.8141,
+      "image_masks_loss": 0.0242,
+      "llm_learning_rate": 4.953836378702598e-06,
+      "loss": 0.7859,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 0.0004953836378702597,
+      "step": 873
+    },
+    {
+      "epoch": 0.09,
+      "gumbel_tau": 0.8139,
+      "image_masks_loss": 0.0414,
+      "llm_learning_rate": 4.953677550415815e-06,
+      "loss": 0.8117,
+      "output_text_mask_loss": 0.0061,
+      "predictor_learning_rate": 0.0004953677550415815,
+      "step": 874
+    },
+    {
+      "epoch": 0.09,
+      "gumbel_tau": 0.8137,
+      "image_masks_loss": 0.0229,
+      "llm_learning_rate": 4.953518451924596e-06,
+      "loss": 0.7998,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.0004953518451924595,
+      "step": 875
+    },
+    {
+      "epoch": 0.09,
+      "gumbel_tau": 0.8135,
+      "image_masks_loss": 0.0359,
+      "llm_learning_rate": 4.953359083246459e-06,
+      "loss": 0.787,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.0004953359083246459,
+      "step": 876
+    },
+    {
+      "epoch": 0.09,
+      "gumbel_tau": 0.8133,
+      "image_masks_loss": 0.0154,
+      "llm_learning_rate": 4.953199444398956e-06,
+      "loss": 0.7107,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004953199444398956,
+      "step": 877
+    },
+    {
+      "epoch": 0.09,
+      "gumbel_tau": 0.8131,
+      "image_masks_loss": 0.0101,
+      "llm_learning_rate": 4.9530395353996655e-06,
+      "loss": 0.8534,
+      "output_text_mask_loss": 0.0477,
+      "predictor_learning_rate": 0.0004953039535399665,
+      "step": 878
+    },
+    {
+      "epoch": 0.09,
+      "gumbel_tau": 0.8129,
+      "image_masks_loss": 0.0233,
+      "llm_learning_rate": 4.952879356266197e-06,
+      "loss": 0.8171,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.0004952879356266197,
+      "step": 879
+    },
+    {
+      "epoch": 0.09,
+      "gumbel_tau": 0.8127,
+      "image_masks_loss": 0.0614,
+      "llm_learning_rate": 4.952718907016191e-06,
+      "loss": 0.7641,
+      "output_text_mask_loss": 0.0406,
+      "predictor_learning_rate": 0.000495271890701619,
+      "step": 880
+    },
+    {
+      "epoch": 0.09,
+      "gumbel_tau": 0.8125,
+      "image_masks_loss": 0.0243,
+      "llm_learning_rate": 4.9525581876673144e-06,
+      "loss": 0.8032,
+      "output_text_mask_loss": 0.0067,
+      "predictor_learning_rate": 0.0004952558187667314,
+      "step": 881
+    },
+    {
+      "epoch": 0.09,
+      "gumbel_tau": 0.8123,
+      "image_masks_loss": 0.0197,
+      "llm_learning_rate": 4.952397198237269e-06,
+      "loss": 0.6368,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004952397198237268,
+      "step": 882
+    },
+    {
+      "epoch": 0.09,
+      "gumbel_tau": 0.8121,
+      "image_masks_loss": 0.0439,
+      "llm_learning_rate": 4.95223593874378e-06,
+      "loss": 0.852,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.000495223593874378,
+      "step": 883
+    },
+    {
+      "epoch": 0.09,
+      "gumbel_tau": 0.8119,
+      "image_masks_loss": 0.0366,
+      "llm_learning_rate": 4.952074409204609e-06,
+      "loss": 0.6592,
+      "output_text_mask_loss": 0.0065,
+      "predictor_learning_rate": 0.0004952074409204608,
+      "step": 884
+    },
+    {
+      "epoch": 0.09,
+      "gumbel_tau": 0.8118,
+      "image_masks_loss": 0.0199,
+      "llm_learning_rate": 4.951912609637542e-06,
+      "loss": 0.782,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.0004951912609637542,
+      "step": 885
+    },
+    {
+      "epoch": 0.09,
+      "gumbel_tau": 0.8116,
+      "image_masks_loss": 0.0192,
+      "llm_learning_rate": 4.9517505400603955e-06,
+      "loss": 0.7918,
+      "output_text_mask_loss": 0.0193,
+      "predictor_learning_rate": 0.0004951750540060395,
+      "step": 886
+    },
+    {
+      "epoch": 0.09,
+      "gumbel_tau": 0.8114,
+      "image_masks_loss": 0.0322,
+      "llm_learning_rate": 4.95158820049102e-06,
+      "loss": 0.7751,
+      "output_text_mask_loss": 0.0763,
+      "predictor_learning_rate": 0.0004951588200491019,
+      "step": 887
+    },
+    {
+      "epoch": 0.09,
+      "gumbel_tau": 0.8112,
+      "image_masks_loss": 0.0473,
+      "llm_learning_rate": 4.95142559094729e-06,
+      "loss": 0.7983,
+      "output_text_mask_loss": 0.016,
+      "predictor_learning_rate": 0.0004951425590947289,
+      "step": 888
+    },
+    {
+      "epoch": 0.09,
+      "gumbel_tau": 0.811,
+      "image_masks_loss": 0.0257,
+      "llm_learning_rate": 4.951262711447114e-06,
+      "loss": 0.7446,
+      "output_text_mask_loss": 0.037,
+      "predictor_learning_rate": 0.0004951262711447114,
+      "step": 889
+    },
+    {
+      "epoch": 0.09,
+      "gumbel_tau": 0.8108,
+      "image_masks_loss": 0.0263,
+      "llm_learning_rate": 4.951099562008429e-06,
+      "loss": 0.8252,
+      "output_text_mask_loss": 0.0004,
+      "predictor_learning_rate": 0.0004951099562008428,
+      "step": 890
+    },
+    {
+      "epoch": 0.09,
+      "gumbel_tau": 0.8106,
+      "image_masks_loss": 0.0159,
+      "llm_learning_rate": 4.9509361426491994e-06,
+      "loss": 0.8666,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.0004950936142649199,
+      "step": 891
+    },
+    {
+      "epoch": 0.09,
+      "gumbel_tau": 0.8104,
+      "image_masks_loss": 0.0381,
+      "llm_learning_rate": 4.950772453387423e-06,
+      "loss": 0.8193,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.0004950772453387423,
+      "step": 892
+    },
+    {
+      "epoch": 0.09,
+      "gumbel_tau": 0.8102,
+      "image_masks_loss": 0.025,
+      "llm_learning_rate": 4.950608494241125e-06,
+      "loss": 0.6976,
+      "output_text_mask_loss": 0.0076,
+      "predictor_learning_rate": 0.0004950608494241125,
+      "step": 893
+    },
+    {
+      "epoch": 0.09,
+      "gumbel_tau": 0.81,
+      "image_masks_loss": 0.0284,
+      "llm_learning_rate": 4.950444265228361e-06,
+      "loss": 0.9611,
+      "output_text_mask_loss": 0.0077,
+      "predictor_learning_rate": 0.0004950444265228361,
+      "step": 894
+    },
+    {
+      "epoch": 0.09,
+      "gumbel_tau": 0.8098,
+      "image_masks_loss": 0.0444,
+      "llm_learning_rate": 4.950279766367218e-06,
+      "loss": 0.7933,
+      "output_text_mask_loss": 0.0248,
+      "predictor_learning_rate": 0.0004950279766367217,
+      "step": 895
+    },
+    {
+      "epoch": 0.09,
+      "gumbel_tau": 0.8097,
+      "image_masks_loss": 0.0115,
+      "llm_learning_rate": 4.950114997675808e-06,
+      "loss": 0.7718,
+      "output_text_mask_loss": 0.0263,
+      "predictor_learning_rate": 0.0004950114997675808,
+      "step": 896
+    },
+    {
+      "epoch": 0.09,
+      "gumbel_tau": 0.8095,
+      "image_masks_loss": 0.0338,
+      "llm_learning_rate": 4.949949959172277e-06,
+      "loss": 0.7962,
+      "output_text_mask_loss": 0.0004,
+      "predictor_learning_rate": 0.0004949949959172276,
+      "step": 897
+    },
+    {
+      "epoch": 0.09,
+      "gumbel_tau": 0.8093,
+      "image_masks_loss": 0.0131,
+      "llm_learning_rate": 4.949784650874799e-06,
+      "loss": 0.9204,
+      "output_text_mask_loss": 0.0448,
+      "predictor_learning_rate": 0.0004949784650874799,
+      "step": 898
+    },
+    {
+      "epoch": 0.09,
+      "gumbel_tau": 0.8091,
+      "image_masks_loss": 0.0282,
+      "llm_learning_rate": 4.949619072801579e-06,
+      "loss": 0.8368,
+      "output_text_mask_loss": 0.0156,
+      "predictor_learning_rate": 0.0004949619072801579,
+      "step": 899
+    },
+    {
+      "epoch": 0.09,
+      "gumbel_tau": 0.8089,
+      "image_masks_loss": 0.0226,
+      "llm_learning_rate": 4.94945322497085e-06,
+      "loss": 0.7803,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.000494945322497085,
+      "step": 900
+    },
+    {
+      "epoch": 0.09,
+      "gumbel_tau": 0.8087,
+      "image_masks_loss": 0.0196,
+      "llm_learning_rate": 4.949287107400877e-06,
+      "loss": 0.7935,
+      "output_text_mask_loss": 0.0376,
+      "predictor_learning_rate": 0.0004949287107400877,
+      "step": 901
+    },
+    {
+      "epoch": 0.09,
+      "gumbel_tau": 0.8085,
+      "image_masks_loss": 0.0164,
+      "llm_learning_rate": 4.949120720109952e-06,
+      "loss": 0.8477,
+      "output_text_mask_loss": 0.0429,
+      "predictor_learning_rate": 0.0004949120720109952,
+      "step": 902
+    },
+    {
+      "epoch": 0.09,
+      "gumbel_tau": 0.8083,
+      "image_masks_loss": 0.0114,
+      "llm_learning_rate": 4.948954063116397e-06,
+      "loss": 0.8009,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004948954063116397,
+      "step": 903
+    },
+    {
+      "epoch": 0.09,
+      "gumbel_tau": 0.8081,
+      "image_masks_loss": 0.0414,
+      "llm_learning_rate": 4.9487871364385664e-06,
+      "loss": 0.6975,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.0004948787136438566,
+      "step": 904
+    },
+    {
+      "epoch": 0.09,
+      "gumbel_tau": 0.8079,
+      "image_masks_loss": 0.0161,
+      "llm_learning_rate": 4.948619940094842e-06,
+      "loss": 0.732,
+      "output_text_mask_loss": 0.0219,
+      "predictor_learning_rate": 0.0004948619940094842,
+      "step": 905
+    },
+    {
+      "epoch": 0.09,
+      "gumbel_tau": 0.8077,
+      "image_masks_loss": 0.0225,
+      "llm_learning_rate": 4.948452474103636e-06,
+      "loss": 0.8501,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004948452474103636,
+      "step": 906
+    },
+    {
+      "epoch": 0.09,
+      "gumbel_tau": 0.8076,
+      "image_masks_loss": 0.0114,
+      "llm_learning_rate": 4.948284738483389e-06,
+      "loss": 0.8062,
+      "output_text_mask_loss": 0.0565,
+      "predictor_learning_rate": 0.0004948284738483389,
+      "step": 907
+    },
+    {
+      "epoch": 0.09,
+      "gumbel_tau": 0.8074,
+      "image_masks_loss": 0.0303,
+      "llm_learning_rate": 4.948116733252574e-06,
+      "loss": 0.7707,
+      "output_text_mask_loss": 0.0526,
+      "predictor_learning_rate": 0.0004948116733252574,
+      "step": 908
+    },
+    {
+      "epoch": 0.09,
+      "gumbel_tau": 0.8072,
+      "image_masks_loss": 0.0117,
+      "llm_learning_rate": 4.947948458429691e-06,
+      "loss": 0.7233,
+      "output_text_mask_loss": 0.0496,
+      "predictor_learning_rate": 0.000494794845842969,
+      "step": 909
+    },
+    {
+      "epoch": 0.09,
+      "gumbel_tau": 0.807,
+      "image_masks_loss": 0.0355,
+      "llm_learning_rate": 4.9477799140332715e-06,
+      "loss": 0.7792,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.0004947779914033271,
+      "step": 910
+    },
+    {
+      "epoch": 0.09,
+      "gumbel_tau": 0.8068,
+      "image_masks_loss": 0.0191,
+      "llm_learning_rate": 4.947611100081875e-06,
+      "loss": 0.6996,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004947611100081875,
+      "step": 911
+    },
+    {
+      "epoch": 0.09,
+      "gumbel_tau": 0.8066,
+      "image_masks_loss": 0.0125,
+      "llm_learning_rate": 4.947442016594093e-06,
+      "loss": 0.7741,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004947442016594092,
+      "step": 912
+    },
+    {
+      "epoch": 0.09,
+      "gumbel_tau": 0.8064,
+      "image_masks_loss": 0.0269,
+      "llm_learning_rate": 4.947272663588544e-06,
+      "loss": 0.6667,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004947272663588544,
+      "step": 913
+    },
+    {
+      "epoch": 0.09,
+      "gumbel_tau": 0.8062,
+      "image_masks_loss": 0.0261,
+      "llm_learning_rate": 4.947103041083879e-06,
+      "loss": 0.826,
+      "output_text_mask_loss": 0.0256,
+      "predictor_learning_rate": 0.0004947103041083879,
+      "step": 914
+    },
+    {
+      "epoch": 0.09,
+      "gumbel_tau": 0.806,
+      "image_masks_loss": 0.0105,
+      "llm_learning_rate": 4.946933149098776e-06,
+      "loss": 0.7518,
+      "output_text_mask_loss": 0.0282,
+      "predictor_learning_rate": 0.0004946933149098775,
+      "step": 915
+    },
+    {
+      "epoch": 0.09,
+      "gumbel_tau": 0.8058,
+      "image_masks_loss": 0.0252,
+      "llm_learning_rate": 4.946762987651944e-06,
+      "loss": 0.6325,
+      "output_text_mask_loss": 0.0202,
+      "predictor_learning_rate": 0.0004946762987651943,
+      "step": 916
+    },
+    {
+      "epoch": 0.09,
+      "gumbel_tau": 0.8057,
+      "image_masks_loss": 0.0152,
+      "llm_learning_rate": 4.946592556762122e-06,
+      "loss": 0.804,
+      "output_text_mask_loss": 0.0009,
+      "predictor_learning_rate": 0.0004946592556762121,
+      "step": 917
+    },
+    {
+      "epoch": 0.09,
+      "gumbel_tau": 0.8055,
+      "image_masks_loss": 0.0149,
+      "llm_learning_rate": 4.946421856448078e-06,
+      "loss": 0.8085,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004946421856448078,
+      "step": 918
+    },
+    {
+      "epoch": 0.09,
+      "gumbel_tau": 0.8053,
+      "image_masks_loss": 0.0067,
+      "llm_learning_rate": 4.9462508867286104e-06,
+      "loss": 0.7472,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.000494625088672861,
+      "step": 919
+    },
+    {
+      "epoch": 0.09,
+      "gumbel_tau": 0.8051,
+      "image_masks_loss": 0.0169,
+      "llm_learning_rate": 4.946079647622547e-06,
+      "loss": 0.8075,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004946079647622547,
+      "step": 920
+    },
+    {
+      "epoch": 0.09,
+      "gumbel_tau": 0.8049,
+      "image_masks_loss": 0.0164,
+      "llm_learning_rate": 4.945908139148743e-06,
+      "loss": 0.8423,
+      "output_text_mask_loss": 0.0043,
+      "predictor_learning_rate": 0.0004945908139148742,
+      "step": 921
+    },
+    {
+      "epoch": 0.09,
+      "gumbel_tau": 0.8047,
+      "image_masks_loss": 0.0263,
+      "llm_learning_rate": 4.9457363613260865e-06,
+      "loss": 0.85,
+      "output_text_mask_loss": 0.0147,
+      "predictor_learning_rate": 0.0004945736361326086,
+      "step": 922
+    },
+    {
+      "epoch": 0.09,
+      "gumbel_tau": 0.8045,
+      "image_masks_loss": 0.0368,
+      "llm_learning_rate": 4.945564314173496e-06,
+      "loss": 0.6841,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004945564314173495,
+      "step": 923
+    },
+    {
+      "epoch": 0.09,
+      "gumbel_tau": 0.8043,
+      "image_masks_loss": 0.0232,
+      "llm_learning_rate": 4.945391997709914e-06,
+      "loss": 0.7416,
+      "output_text_mask_loss": 0.0269,
+      "predictor_learning_rate": 0.0004945391997709914,
+      "step": 924
+    },
+    {
+      "epoch": 0.09,
+      "gumbel_tau": 0.8041,
+      "image_masks_loss": 0.0141,
+      "llm_learning_rate": 4.94521941195432e-06,
+      "loss": 0.7407,
+      "output_text_mask_loss": 0.1694,
+      "predictor_learning_rate": 0.000494521941195432,
+      "step": 925
+    },
+    {
+      "epoch": 0.09,
+      "gumbel_tau": 0.8039,
+      "image_masks_loss": 0.0147,
+      "llm_learning_rate": 4.945046556925716e-06,
+      "loss": 0.8643,
+      "output_text_mask_loss": 0.0414,
+      "predictor_learning_rate": 0.0004945046556925716,
+      "step": 926
+    },
+    {
+      "epoch": 0.09,
+      "gumbel_tau": 0.8038,
+      "image_masks_loss": 0.0099,
+      "llm_learning_rate": 4.94487343264314e-06,
+      "loss": 0.8119,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.0004944873432643139,
+      "step": 927
+    },
+    {
+      "epoch": 0.1,
+      "gumbel_tau": 0.8036,
+      "image_masks_loss": 0.011,
+      "llm_learning_rate": 4.944700039125656e-06,
+      "loss": 0.7879,
+      "output_text_mask_loss": 0.0116,
+      "predictor_learning_rate": 0.0004944700039125656,
+      "step": 928
+    },
+    {
+      "epoch": 0.1,
+      "gumbel_tau": 0.8034,
+      "image_masks_loss": 0.0168,
+      "llm_learning_rate": 4.944526376392358e-06,
+      "loss": 0.7845,
+      "output_text_mask_loss": 0.0671,
+      "predictor_learning_rate": 0.0004944526376392358,
+      "step": 929
+    },
+    {
+      "epoch": 0.1,
+      "gumbel_tau": 0.8032,
+      "image_masks_loss": 0.0284,
+      "llm_learning_rate": 4.944352444462371e-06,
+      "loss": 0.7191,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004944352444462371,
+      "step": 930
+    },
+    {
+      "epoch": 0.1,
+      "gumbel_tau": 0.803,
+      "image_masks_loss": 0.0166,
+      "llm_learning_rate": 4.944178243354848e-06,
+      "loss": 0.8,
+      "output_text_mask_loss": 0.0284,
+      "predictor_learning_rate": 0.0004944178243354848,
+      "step": 931
+    },
+    {
+      "epoch": 0.1,
+      "gumbel_tau": 0.8028,
+      "image_masks_loss": 0.0152,
+      "llm_learning_rate": 4.944003773088972e-06,
+      "loss": 0.8423,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 0.0004944003773088972,
+      "step": 932
+    },
+    {
+      "epoch": 0.1,
+      "gumbel_tau": 0.8026,
+      "image_masks_loss": 0.0178,
+      "llm_learning_rate": 4.943829033683958e-06,
+      "loss": 0.7804,
+      "output_text_mask_loss": 0.1434,
+      "predictor_learning_rate": 0.0004943829033683957,
+      "step": 933
+    },
+    {
+      "epoch": 0.1,
+      "gumbel_tau": 0.8024,
+      "image_masks_loss": 0.0326,
+      "llm_learning_rate": 4.943654025159047e-06,
+      "loss": 0.8126,
+      "output_text_mask_loss": 0.0454,
+      "predictor_learning_rate": 0.0004943654025159046,
+      "step": 934
+    },
+    {
+      "epoch": 0.1,
+      "gumbel_tau": 0.8022,
+      "image_masks_loss": 0.0208,
+      "llm_learning_rate": 4.943478747533511e-06,
+      "loss": 0.7669,
+      "output_text_mask_loss": 0.0322,
+      "predictor_learning_rate": 0.000494347874753351,
+      "step": 935
+    },
+    {
+      "epoch": 0.1,
+      "gumbel_tau": 0.802,
+      "image_masks_loss": 0.005,
+      "llm_learning_rate": 4.943303200826653e-06,
+      "loss": 0.728,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.0004943303200826653,
+      "step": 936
+    },
+    {
+      "epoch": 0.1,
+      "gumbel_tau": 0.8019,
+      "image_masks_loss": 0.0154,
+      "llm_learning_rate": 4.943127385057804e-06,
+      "loss": 0.8788,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 0.0004943127385057804,
+      "step": 937
+    },
+    {
+      "epoch": 0.1,
+      "gumbel_tau": 0.8017,
+      "image_masks_loss": 0.016,
+      "llm_learning_rate": 4.9429513002463255e-06,
+      "loss": 0.9239,
+      "output_text_mask_loss": 0.0422,
+      "predictor_learning_rate": 0.0004942951300246325,
+      "step": 938
+    },
+    {
+      "epoch": 0.1,
+      "gumbel_tau": 0.8015,
+      "image_masks_loss": 0.0372,
+      "llm_learning_rate": 4.94277494641161e-06,
+      "loss": 0.7668,
+      "output_text_mask_loss": 0.0035,
+      "predictor_learning_rate": 0.0004942774946411609,
+      "step": 939
+    },
+    {
+      "epoch": 0.1,
+      "gumbel_tau": 0.8013,
+      "image_masks_loss": 0.0311,
+      "llm_learning_rate": 4.942598323573074e-06,
+      "loss": 0.8292,
+      "output_text_mask_loss": 0.021,
+      "predictor_learning_rate": 0.0004942598323573075,
+      "step": 940
+    },
+    {
+      "epoch": 0.1,
+      "gumbel_tau": 0.8011,
+      "image_masks_loss": 0.0414,
+      "llm_learning_rate": 4.942421431750171e-06,
+      "loss": 0.8129,
+      "output_text_mask_loss": 0.0992,
+      "predictor_learning_rate": 0.0004942421431750171,
+      "step": 941
+    },
+    {
+      "epoch": 0.1,
+      "gumbel_tau": 0.8009,
+      "image_masks_loss": 0.0059,
+      "llm_learning_rate": 4.94224427096238e-06,
+      "loss": 0.7252,
+      "output_text_mask_loss": 0.0771,
+      "predictor_learning_rate": 0.000494224427096238,
+      "step": 942
+    },
+    {
+      "epoch": 0.1,
+      "gumbel_tau": 0.8007,
+      "image_masks_loss": 0.0191,
+      "llm_learning_rate": 4.94206684122921e-06,
+      "loss": 0.78,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004942066841229209,
+      "step": 943
+    },
+    {
+      "epoch": 0.1,
+      "gumbel_tau": 0.8005,
+      "image_masks_loss": 0.0142,
+      "llm_learning_rate": 4.941889142570199e-06,
+      "loss": 0.7856,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004941889142570198,
+      "step": 944
+    },
+    {
+      "epoch": 0.1,
+      "gumbel_tau": 0.8003,
+      "image_masks_loss": 0.0241,
+      "llm_learning_rate": 4.941711175004917e-06,
+      "loss": 0.7489,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 0.0004941711175004917,
+      "step": 945
+    },
+    {
+      "epoch": 0.1,
+      "gumbel_tau": 0.8002,
+      "image_masks_loss": 0.0376,
+      "llm_learning_rate": 4.941532938552962e-06,
+      "loss": 0.8594,
+      "output_text_mask_loss": 0.0549,
+      "predictor_learning_rate": 0.0004941532938552962,
+      "step": 946
+    },
+    {
+      "epoch": 0.1,
+      "gumbel_tau": 0.8,
+      "image_masks_loss": 0.0099,
+      "llm_learning_rate": 4.9413544332339615e-06,
+      "loss": 0.7576,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.0004941354433233962,
+      "step": 947
+    },
+    {
+      "epoch": 0.1,
+      "gumbel_tau": 0.7998,
+      "image_masks_loss": 0.0149,
+      "llm_learning_rate": 4.941175659067572e-06,
+      "loss": 0.8596,
+      "output_text_mask_loss": 0.0275,
+      "predictor_learning_rate": 0.0004941175659067572,
+      "step": 948
+    },
+    {
+      "epoch": 0.1,
+      "gumbel_tau": 0.7996,
+      "image_masks_loss": 0.0243,
+      "llm_learning_rate": 4.940996616073482e-06,
+      "loss": 0.8395,
+      "output_text_mask_loss": 0.0238,
+      "predictor_learning_rate": 0.0004940996616073482,
+      "step": 949
+    },
+    {
+      "epoch": 0.1,
+      "gumbel_tau": 0.7994,
+      "image_masks_loss": 0.0121,
+      "llm_learning_rate": 4.940817304271408e-06,
+      "loss": 0.7642,
+      "output_text_mask_loss": 0.0923,
+      "predictor_learning_rate": 0.0004940817304271408,
+      "step": 950
+    },
+    {
+      "epoch": 0.1,
+      "gumbel_tau": 0.7992,
+      "image_masks_loss": 0.0273,
+      "llm_learning_rate": 4.940637723681095e-06,
+      "loss": 0.8059,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 0.0004940637723681096,
+      "step": 951
+    },
+    {
+      "epoch": 0.1,
+      "gumbel_tau": 0.799,
+      "image_masks_loss": 0.0104,
+      "llm_learning_rate": 4.9404578743223204e-06,
+      "loss": 0.8614,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004940457874322321,
+      "step": 952
+    },
+    {
+      "epoch": 0.1,
+      "gumbel_tau": 0.7988,
+      "image_masks_loss": 0.0189,
+      "llm_learning_rate": 4.940277756214888e-06,
+      "loss": 0.7634,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004940277756214888,
+      "step": 953
+    },
+    {
+      "epoch": 0.1,
+      "gumbel_tau": 0.7987,
+      "image_masks_loss": 0.0121,
+      "llm_learning_rate": 4.940097369378634e-06,
+      "loss": 0.8277,
+      "output_text_mask_loss": 0.0038,
+      "predictor_learning_rate": 0.0004940097369378635,
+      "step": 954
+    },
+    {
+      "epoch": 0.1,
+      "gumbel_tau": 0.7985,
+      "image_masks_loss": 0.0292,
+      "llm_learning_rate": 4.939916713833423e-06,
+      "loss": 0.7622,
+      "output_text_mask_loss": 0.0763,
+      "predictor_learning_rate": 0.0004939916713833423,
+      "step": 955
+    },
+    {
+      "epoch": 0.1,
+      "gumbel_tau": 0.7983,
+      "image_masks_loss": 0.0206,
+      "llm_learning_rate": 4.939735789599148e-06,
+      "loss": 0.8145,
+      "output_text_mask_loss": 0.0271,
+      "predictor_learning_rate": 0.0004939735789599148,
+      "step": 956
+    },
+    {
+      "epoch": 0.1,
+      "gumbel_tau": 0.7981,
+      "image_masks_loss": 0.0106,
+      "llm_learning_rate": 4.9395545966957335e-06,
+      "loss": 0.7696,
+      "output_text_mask_loss": 0.0141,
+      "predictor_learning_rate": 0.0004939554596695734,
+      "step": 957
+    },
+    {
+      "epoch": 0.1,
+      "gumbel_tau": 0.7979,
+      "image_masks_loss": 0.0427,
+      "llm_learning_rate": 4.939373135143134e-06,
+      "loss": 0.7844,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004939373135143133,
+      "step": 958
+    },
+    {
+      "epoch": 0.1,
+      "gumbel_tau": 0.7977,
+      "image_masks_loss": 0.0209,
+      "llm_learning_rate": 4.9391914049613305e-06,
+      "loss": 0.8301,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.0004939191404961331,
+      "step": 959
+    },
+    {
+      "epoch": 0.1,
+      "gumbel_tau": 0.7975,
+      "image_masks_loss": 0.0301,
+      "llm_learning_rate": 4.939009406170337e-06,
+      "loss": 0.8675,
+      "output_text_mask_loss": 0.0916,
+      "predictor_learning_rate": 0.0004939009406170336,
+      "step": 960
+    },
+    {
+      "epoch": 0.1,
+      "gumbel_tau": 0.7973,
+      "image_masks_loss": 0.0082,
+      "llm_learning_rate": 4.9388271387901945e-06,
+      "loss": 0.762,
+      "output_text_mask_loss": 0.0019,
+      "predictor_learning_rate": 0.0004938827138790194,
+      "step": 961
+    },
+    {
+      "epoch": 0.1,
+      "gumbel_tau": 0.7971,
+      "image_masks_loss": 0.0177,
+      "llm_learning_rate": 4.938644602840974e-06,
+      "loss": 0.8564,
+      "output_text_mask_loss": 0.0137,
+      "predictor_learning_rate": 0.0004938644602840974,
+      "step": 962
+    },
+    {
+      "epoch": 0.1,
+      "gumbel_tau": 0.797,
+      "image_masks_loss": 0.0234,
+      "llm_learning_rate": 4.93846179834278e-06,
+      "loss": 0.82,
+      "output_text_mask_loss": 0.0004,
+      "predictor_learning_rate": 0.0004938461798342779,
+      "step": 963
+    },
+    {
+      "epoch": 0.1,
+      "gumbel_tau": 0.7968,
+      "image_masks_loss": 0.0137,
+      "llm_learning_rate": 4.9382787253157404e-06,
+      "loss": 0.8675,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.000493827872531574,
+      "step": 964
+    },
+    {
+      "epoch": 0.1,
+      "gumbel_tau": 0.7966,
+      "image_masks_loss": 0.0131,
+      "llm_learning_rate": 4.938095383780017e-06,
+      "loss": 0.7337,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004938095383780017,
+      "step": 965
+    },
+    {
+      "epoch": 0.1,
+      "gumbel_tau": 0.7964,
+      "image_masks_loss": 0.0228,
+      "llm_learning_rate": 4.937911773755798e-06,
+      "loss": 0.6821,
+      "output_text_mask_loss": 0.0595,
+      "predictor_learning_rate": 0.0004937911773755799,
+      "step": 966
+    },
+    {
+      "epoch": 0.1,
+      "gumbel_tau": 0.7962,
+      "image_masks_loss": 0.01,
+      "llm_learning_rate": 4.9377278952633055e-06,
+      "loss": 0.7919,
+      "output_text_mask_loss": 0.0109,
+      "predictor_learning_rate": 0.0004937727895263306,
+      "step": 967
+    },
+    {
+      "epoch": 0.1,
+      "gumbel_tau": 0.796,
+      "image_masks_loss": 0.0212,
+      "llm_learning_rate": 4.937543748322787e-06,
+      "loss": 0.8752,
+      "output_text_mask_loss": 0.0126,
+      "predictor_learning_rate": 0.0004937543748322787,
+      "step": 968
+    },
+    {
+      "epoch": 0.1,
+      "gumbel_tau": 0.7958,
+      "image_masks_loss": 0.0143,
+      "llm_learning_rate": 4.937359332954521e-06,
+      "loss": 0.806,
+      "output_text_mask_loss": 0.0086,
+      "predictor_learning_rate": 0.0004937359332954521,
+      "step": 969
+    },
+    {
+      "epoch": 0.1,
+      "gumbel_tau": 0.7956,
+      "image_masks_loss": 0.0068,
+      "llm_learning_rate": 4.937174649178818e-06,
+      "loss": 0.7511,
+      "output_text_mask_loss": 0.0322,
+      "predictor_learning_rate": 0.0004937174649178817,
+      "step": 970
+    },
+    {
+      "epoch": 0.1,
+      "gumbel_tau": 0.7955,
+      "image_masks_loss": 0.0269,
+      "llm_learning_rate": 4.936989697016013e-06,
+      "loss": 0.8125,
+      "output_text_mask_loss": 0.02,
+      "predictor_learning_rate": 0.0004936989697016012,
+      "step": 971
+    },
+    {
+      "epoch": 0.1,
+      "gumbel_tau": 0.7953,
+      "image_masks_loss": 0.0186,
+      "llm_learning_rate": 4.9368044764864735e-06,
+      "loss": 0.8519,
+      "output_text_mask_loss": 0.0305,
+      "predictor_learning_rate": 0.0004936804476486472,
+      "step": 972
+    },
+    {
+      "epoch": 0.1,
+      "gumbel_tau": 0.7951,
+      "image_masks_loss": 0.0169,
+      "llm_learning_rate": 4.936618987610597e-06,
+      "loss": 0.667,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004936618987610597,
+      "step": 973
+    },
+    {
+      "epoch": 0.1,
+      "gumbel_tau": 0.7949,
+      "image_masks_loss": 0.0133,
+      "llm_learning_rate": 4.936433230408812e-06,
+      "loss": 0.8384,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.0004936433230408811,
+      "step": 974
+    },
+    {
+      "epoch": 0.1,
+      "gumbel_tau": 0.7947,
+      "image_masks_loss": 0.0246,
+      "llm_learning_rate": 4.936247204901572e-06,
+      "loss": 0.7922,
+      "output_text_mask_loss": 0.0061,
+      "predictor_learning_rate": 0.0004936247204901571,
+      "step": 975
+    },
+    {
+      "epoch": 0.1,
+      "gumbel_tau": 0.7945,
+      "image_masks_loss": 0.0275,
+      "llm_learning_rate": 4.936060911109362e-06,
+      "loss": 0.8327,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.0004936060911109362,
+      "step": 976
+    },
+    {
+      "epoch": 0.1,
+      "gumbel_tau": 0.7943,
+      "image_masks_loss": 0.0236,
+      "llm_learning_rate": 4.9358743490527e-06,
+      "loss": 0.8189,
+      "output_text_mask_loss": 0.0412,
+      "predictor_learning_rate": 0.00049358743490527,
+      "step": 977
+    },
+    {
+      "epoch": 0.1,
+      "gumbel_tau": 0.7941,
+      "image_masks_loss": 0.0125,
+      "llm_learning_rate": 4.9356875187521295e-06,
+      "loss": 0.803,
+      "output_text_mask_loss": 0.0231,
+      "predictor_learning_rate": 0.0004935687518752129,
+      "step": 978
+    },
+    {
+      "epoch": 0.1,
+      "gumbel_tau": 0.794,
+      "image_masks_loss": 0.0211,
+      "llm_learning_rate": 4.935500420228223e-06,
+      "loss": 0.8374,
+      "output_text_mask_loss": 0.0393,
+      "predictor_learning_rate": 0.0004935500420228222,
+      "step": 979
+    },
+    {
+      "epoch": 0.1,
+      "gumbel_tau": 0.7938,
+      "image_masks_loss": 0.0187,
+      "llm_learning_rate": 4.935313053501586e-06,
+      "loss": 0.6618,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.0004935313053501585,
+      "step": 980
+    },
+    {
+      "epoch": 0.1,
+      "gumbel_tau": 0.7936,
+      "image_masks_loss": 0.0387,
+      "llm_learning_rate": 4.93512541859285e-06,
+      "loss": 0.8467,
+      "output_text_mask_loss": 0.0015,
+      "predictor_learning_rate": 0.000493512541859285,
+      "step": 981
+    },
+    {
+      "epoch": 0.1,
+      "gumbel_tau": 0.7934,
+      "image_masks_loss": 0.0288,
+      "llm_learning_rate": 4.93493751552268e-06,
+      "loss": 0.6798,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004934937515522679,
+      "step": 982
+    },
+    {
+      "epoch": 0.1,
+      "gumbel_tau": 0.7932,
+      "image_masks_loss": 0.016,
+      "llm_learning_rate": 4.934749344311766e-06,
+      "loss": 0.8034,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004934749344311766,
+      "step": 983
+    },
+    {
+      "epoch": 0.1,
+      "gumbel_tau": 0.793,
+      "image_masks_loss": 0.0481,
+      "llm_learning_rate": 4.934560904980832e-06,
+      "loss": 0.7421,
+      "output_text_mask_loss": 0.0189,
+      "predictor_learning_rate": 0.0004934560904980832,
+      "step": 984
+    },
+    {
+      "epoch": 0.1,
+      "gumbel_tau": 0.7928,
+      "image_masks_loss": 0.0484,
+      "llm_learning_rate": 4.934372197550628e-06,
+      "loss": 0.8084,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.0004934372197550628,
+      "step": 985
+    },
+    {
+      "epoch": 0.1,
+      "gumbel_tau": 0.7926,
+      "image_masks_loss": 0.0201,
+      "llm_learning_rate": 4.934183222041935e-06,
+      "loss": 0.8003,
+      "output_text_mask_loss": 0.0032,
+      "predictor_learning_rate": 0.0004934183222041935,
+      "step": 986
+    },
+    {
+      "epoch": 0.1,
+      "gumbel_tau": 0.7925,
+      "image_masks_loss": 0.0082,
+      "llm_learning_rate": 4.933993978475564e-06,
+      "loss": 0.7152,
+      "output_text_mask_loss": 0.0254,
+      "predictor_learning_rate": 0.0004933993978475564,
+      "step": 987
+    },
+    {
+      "epoch": 0.1,
+      "gumbel_tau": 0.7923,
+      "image_masks_loss": 0.0307,
+      "llm_learning_rate": 4.933804466872355e-06,
+      "loss": 0.8596,
+      "output_text_mask_loss": 0.0092,
+      "predictor_learning_rate": 0.0004933804466872355,
+      "step": 988
+    },
+    {
+      "epoch": 0.1,
+      "gumbel_tau": 0.7921,
+      "image_masks_loss": 0.017,
+      "llm_learning_rate": 4.9336146872531766e-06,
+      "loss": 0.7828,
+      "output_text_mask_loss": 0.0353,
+      "predictor_learning_rate": 0.0004933614687253177,
+      "step": 989
+    },
+    {
+      "epoch": 0.1,
+      "gumbel_tau": 0.7919,
+      "image_masks_loss": 0.0359,
+      "llm_learning_rate": 4.933424639638928e-06,
+      "loss": 0.8216,
+      "output_text_mask_loss": 0.0429,
+      "predictor_learning_rate": 0.0004933424639638927,
+      "step": 990
+    },
+    {
+      "epoch": 0.1,
+      "gumbel_tau": 0.7917,
+      "image_masks_loss": 0.0246,
+      "llm_learning_rate": 4.933234324050537e-06,
+      "loss": 0.8065,
+      "output_text_mask_loss": 0.0305,
+      "predictor_learning_rate": 0.0004933234324050538,
+      "step": 991
+    },
+    {
+      "epoch": 0.1,
+      "gumbel_tau": 0.7915,
+      "image_masks_loss": 0.0163,
+      "llm_learning_rate": 4.933043740508964e-06,
+      "loss": 0.6957,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.0004933043740508964,
+      "step": 992
+    },
+    {
+      "epoch": 0.1,
+      "gumbel_tau": 0.7913,
+      "image_masks_loss": 0.0091,
+      "llm_learning_rate": 4.9328528890351946e-06,
+      "loss": 0.7722,
+      "output_text_mask_loss": 0.0771,
+      "predictor_learning_rate": 0.0004932852889035194,
+      "step": 993
+    },
+    {
+      "epoch": 0.1,
+      "gumbel_tau": 0.7911,
+      "image_masks_loss": 0.0143,
+      "llm_learning_rate": 4.932661769650245e-06,
+      "loss": 0.8414,
+      "output_text_mask_loss": 0.0141,
+      "predictor_learning_rate": 0.0004932661769650245,
+      "step": 994
+    },
+    {
+      "epoch": 0.1,
+      "gumbel_tau": 0.791,
+      "image_masks_loss": 0.0177,
+      "llm_learning_rate": 4.932470382375164e-06,
+      "loss": 0.8745,
+      "output_text_mask_loss": 0.0492,
+      "predictor_learning_rate": 0.0004932470382375164,
+      "step": 995
+    },
+    {
+      "epoch": 0.1,
+      "gumbel_tau": 0.7908,
+      "image_masks_loss": 0.0149,
+      "llm_learning_rate": 4.932278727231027e-06,
+      "loss": 0.7978,
+      "output_text_mask_loss": 0.0557,
+      "predictor_learning_rate": 0.0004932278727231026,
+      "step": 996
+    },
+    {
+      "epoch": 0.1,
+      "gumbel_tau": 0.7906,
+      "image_masks_loss": 0.0275,
+      "llm_learning_rate": 4.932086804238937e-06,
+      "loss": 0.6856,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.0004932086804238936,
+      "step": 997
+    },
+    {
+      "epoch": 0.1,
+      "gumbel_tau": 0.7904,
+      "image_masks_loss": 0.018,
+      "llm_learning_rate": 4.931894613420032e-06,
+      "loss": 0.7739,
+      "output_text_mask_loss": 0.0149,
+      "predictor_learning_rate": 0.0004931894613420031,
+      "step": 998
+    },
+    {
+      "epoch": 0.1,
+      "gumbel_tau": 0.7902,
+      "image_masks_loss": 0.0143,
+      "llm_learning_rate": 4.931702154795475e-06,
+      "loss": 0.7705,
+      "output_text_mask_loss": 0.119,
+      "predictor_learning_rate": 0.0004931702154795474,
+      "step": 999
+    },
+    {
+      "epoch": 0.1,
+      "gumbel_tau": 0.79,
+      "image_masks_loss": 0.0156,
+      "llm_learning_rate": 4.9315094283864605e-06,
+      "loss": 0.7831,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004931509428386461,
+      "step": 1000
+    },
+    {
+      "epoch": 0.1,
+      "gumbel_tau": 0.7898,
+      "image_masks_loss": 0.0278,
+      "llm_learning_rate": 4.931316434214212e-06,
+      "loss": 0.8472,
+      "output_text_mask_loss": 0.0004,
+      "predictor_learning_rate": 0.0004931316434214211,
+      "step": 1001
+    },
+    {
+      "epoch": 0.1,
+      "gumbel_tau": 0.7897,
+      "image_masks_loss": 0.0171,
+      "llm_learning_rate": 4.931123172299982e-06,
+      "loss": 0.7867,
+      "output_text_mask_loss": 0.0179,
+      "predictor_learning_rate": 0.0004931123172299981,
+      "step": 1002
+    },
+    {
+      "epoch": 0.1,
+      "gumbel_tau": 0.7895,
+      "image_masks_loss": 0.0115,
+      "llm_learning_rate": 4.930929642665052e-06,
+      "loss": 0.7162,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004930929642665052,
+      "step": 1003
+    },
+    {
+      "epoch": 0.1,
+      "gumbel_tau": 0.7893,
+      "image_masks_loss": 0.0139,
+      "llm_learning_rate": 4.930735845330736e-06,
+      "loss": 0.8248,
+      "output_text_mask_loss": 0.0248,
+      "predictor_learning_rate": 0.0004930735845330737,
+      "step": 1004
+    },
+    {
+      "epoch": 0.1,
+      "gumbel_tau": 0.7891,
+      "image_masks_loss": 0.0301,
+      "llm_learning_rate": 4.930541780318375e-06,
+      "loss": 0.8278,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 0.0004930541780318375,
+      "step": 1005
+    },
+    {
+      "epoch": 0.1,
+      "gumbel_tau": 0.7889,
+      "image_masks_loss": 0.0191,
+      "llm_learning_rate": 4.930347447649339e-06,
+      "loss": 0.8019,
+      "output_text_mask_loss": 0.0591,
+      "predictor_learning_rate": 0.0004930347447649338,
+      "step": 1006
+    },
+    {
+      "epoch": 0.1,
+      "gumbel_tau": 0.7887,
+      "image_masks_loss": 0.0148,
+      "llm_learning_rate": 4.930152847345028e-06,
+      "loss": 0.7247,
+      "output_text_mask_loss": 0.0565,
+      "predictor_learning_rate": 0.0004930152847345028,
+      "step": 1007
+    },
+    {
+      "epoch": 0.1,
+      "gumbel_tau": 0.7885,
+      "image_masks_loss": 0.0275,
+      "llm_learning_rate": 4.929957979426874e-06,
+      "loss": 0.8264,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004929957979426874,
+      "step": 1008
+    },
+    {
+      "epoch": 0.1,
+      "gumbel_tau": 0.7884,
+      "image_masks_loss": 0.0119,
+      "llm_learning_rate": 4.929762843916333e-06,
+      "loss": 0.7535,
+      "output_text_mask_loss": 0.0019,
+      "predictor_learning_rate": 0.0004929762843916333,
+      "step": 1009
+    },
+    {
+      "epoch": 0.1,
+      "gumbel_tau": 0.7882,
+      "image_masks_loss": 0.0238,
+      "llm_learning_rate": 4.9295674408348965e-06,
+      "loss": 0.8067,
+      "output_text_mask_loss": 0.0273,
+      "predictor_learning_rate": 0.0004929567440834896,
+      "step": 1010
+    },
+    {
+      "epoch": 0.1,
+      "gumbel_tau": 0.788,
+      "image_masks_loss": 0.0129,
+      "llm_learning_rate": 4.929371770204082e-06,
+      "loss": 0.8705,
+      "output_text_mask_loss": 0.0416,
+      "predictor_learning_rate": 0.0004929371770204081,
+      "step": 1011
+    },
+    {
+      "epoch": 0.1,
+      "gumbel_tau": 0.7878,
+      "image_masks_loss": 0.0296,
+      "llm_learning_rate": 4.929175832045436e-06,
+      "loss": 0.7503,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004929175832045436,
+      "step": 1012
+    },
+    {
+      "epoch": 0.1,
+      "gumbel_tau": 0.7876,
+      "image_masks_loss": 0.0142,
+      "llm_learning_rate": 4.928979626380537e-06,
+      "loss": 0.803,
+      "output_text_mask_loss": 0.0256,
+      "predictor_learning_rate": 0.0004928979626380537,
+      "step": 1013
+    },
+    {
+      "epoch": 0.1,
+      "gumbel_tau": 0.7874,
+      "image_masks_loss": 0.0203,
+      "llm_learning_rate": 4.928783153230991e-06,
+      "loss": 0.8228,
+      "output_text_mask_loss": 0.0187,
+      "predictor_learning_rate": 0.0004928783153230991,
+      "step": 1014
+    },
+    {
+      "epoch": 0.1,
+      "gumbel_tau": 0.7872,
+      "image_masks_loss": 0.0284,
+      "llm_learning_rate": 4.928586412618436e-06,
+      "loss": 0.838,
+      "output_text_mask_loss": 0.034,
+      "predictor_learning_rate": 0.0004928586412618435,
+      "step": 1015
+    },
+    {
+      "epoch": 0.1,
+      "gumbel_tau": 0.7871,
+      "image_masks_loss": 0.0122,
+      "llm_learning_rate": 4.928389404564533e-06,
+      "loss": 0.8214,
+      "output_text_mask_loss": 0.013,
+      "predictor_learning_rate": 0.0004928389404564534,
+      "step": 1016
+    },
+    {
+      "epoch": 0.1,
+      "gumbel_tau": 0.7869,
+      "image_masks_loss": 0.0221,
+      "llm_learning_rate": 4.928192129090982e-06,
+      "loss": 0.8585,
+      "output_text_mask_loss": 0.0946,
+      "predictor_learning_rate": 0.0004928192129090981,
+      "step": 1017
+    },
+    {
+      "epoch": 0.1,
+      "gumbel_tau": 0.7867,
+      "image_masks_loss": 0.0254,
+      "llm_learning_rate": 4.927994586219504e-06,
+      "loss": 0.8411,
+      "output_text_mask_loss": 0.033,
+      "predictor_learning_rate": 0.0004927994586219505,
+      "step": 1018
+    },
+    {
+      "epoch": 0.1,
+      "gumbel_tau": 0.7865,
+      "image_masks_loss": 0.0077,
+      "llm_learning_rate": 4.927796775971855e-06,
+      "loss": 0.7268,
+      "output_text_mask_loss": 0.0012,
+      "predictor_learning_rate": 0.0004927796775971854,
+      "step": 1019
+    },
+    {
+      "epoch": 0.1,
+      "gumbel_tau": 0.7863,
+      "image_masks_loss": 0.0085,
+      "llm_learning_rate": 4.9275986983698165e-06,
+      "loss": 0.8738,
+      "output_text_mask_loss": 0.0587,
+      "predictor_learning_rate": 0.0004927598698369816,
+      "step": 1020
+    },
+    {
+      "epoch": 0.1,
+      "gumbel_tau": 0.7861,
+      "image_masks_loss": 0.0123,
+      "llm_learning_rate": 4.9274003534352025e-06,
+      "loss": 0.6479,
+      "output_text_mask_loss": 0.0961,
+      "predictor_learning_rate": 0.0004927400353435202,
+      "step": 1021
+    },
+    {
+      "epoch": 0.1,
+      "gumbel_tau": 0.7859,
+      "image_masks_loss": 0.0217,
+      "llm_learning_rate": 4.927201741189854e-06,
+      "loss": 0.756,
+      "output_text_mask_loss": 0.0324,
+      "predictor_learning_rate": 0.0004927201741189854,
+      "step": 1022
+    },
+    {
+      "epoch": 0.1,
+      "gumbel_tau": 0.7858,
+      "image_masks_loss": 0.0278,
+      "llm_learning_rate": 4.927002861655644e-06,
+      "loss": 0.7986,
+      "output_text_mask_loss": 0.0652,
+      "predictor_learning_rate": 0.0004927002861655643,
+      "step": 1023
+    },
+    {
+      "epoch": 0.1,
+      "gumbel_tau": 0.7856,
+      "image_masks_loss": 0.0366,
+      "llm_learning_rate": 4.926803714854472e-06,
+      "loss": 0.7941,
+      "output_text_mask_loss": 0.0111,
+      "predictor_learning_rate": 0.0004926803714854472,
+      "step": 1024
+    },
+    {
+      "epoch": 0.11,
+      "gumbel_tau": 0.7854,
+      "image_masks_loss": 0.0385,
+      "llm_learning_rate": 4.9266043008082705e-06,
+      "loss": 0.805,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004926604300808271,
+      "step": 1025
+    },
+    {
+      "epoch": 0.11,
+      "gumbel_tau": 0.7852,
+      "image_masks_loss": 0.0112,
+      "llm_learning_rate": 4.926404619538998e-06,
+      "loss": 0.8466,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.0004926404619538997,
+      "step": 1026
+    },
+    {
+      "epoch": 0.11,
+      "gumbel_tau": 0.785,
+      "image_masks_loss": 0.0108,
+      "llm_learning_rate": 4.926204671068645e-06,
+      "loss": 0.6433,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.0004926204671068645,
+      "step": 1027
+    },
+    {
+      "epoch": 0.11,
+      "gumbel_tau": 0.7848,
+      "image_masks_loss": 0.0154,
+      "llm_learning_rate": 4.9260044554192285e-06,
+      "loss": 0.7157,
+      "output_text_mask_loss": 0.0416,
+      "predictor_learning_rate": 0.0004926004455419228,
+      "step": 1028
+    },
+    {
+      "epoch": 0.11,
+      "gumbel_tau": 0.7846,
+      "image_masks_loss": 0.0096,
+      "llm_learning_rate": 4.925803972612798e-06,
+      "loss": 0.7883,
+      "output_text_mask_loss": 0.1442,
+      "predictor_learning_rate": 0.0004925803972612798,
+      "step": 1029
+    },
+    {
+      "epoch": 0.11,
+      "gumbel_tau": 0.7845,
+      "image_masks_loss": 0.0112,
+      "llm_learning_rate": 4.92560322267143e-06,
+      "loss": 0.7162,
+      "output_text_mask_loss": 0.0039,
+      "predictor_learning_rate": 0.0004925603222671429,
+      "step": 1030
+    },
+    {
+      "epoch": 0.11,
+      "gumbel_tau": 0.7843,
+      "image_masks_loss": 0.0147,
+      "llm_learning_rate": 4.9254022056172335e-06,
+      "loss": 0.847,
+      "output_text_mask_loss": 0.0854,
+      "predictor_learning_rate": 0.0004925402205617234,
+      "step": 1031
+    },
+    {
+      "epoch": 0.11,
+      "gumbel_tau": 0.7841,
+      "image_masks_loss": 0.0303,
+      "llm_learning_rate": 4.925200921472343e-06,
+      "loss": 0.8808,
+      "output_text_mask_loss": 0.0458,
+      "predictor_learning_rate": 0.0004925200921472343,
+      "step": 1032
+    },
+    {
+      "epoch": 0.11,
+      "gumbel_tau": 0.7839,
+      "image_masks_loss": 0.0372,
+      "llm_learning_rate": 4.924999370258926e-06,
+      "loss": 0.7277,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.0004924999370258925,
+      "step": 1033
+    },
+    {
+      "epoch": 0.11,
+      "gumbel_tau": 0.7837,
+      "image_masks_loss": 0.0263,
+      "llm_learning_rate": 4.924797551999175e-06,
+      "loss": 0.8272,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.0004924797551999175,
+      "step": 1034
+    },
+    {
+      "epoch": 0.11,
+      "gumbel_tau": 0.7835,
+      "image_masks_loss": 0.004,
+      "llm_learning_rate": 4.924595466715318e-06,
+      "loss": 0.7803,
+      "output_text_mask_loss": 0.0349,
+      "predictor_learning_rate": 0.0004924595466715318,
+      "step": 1035
+    },
+    {
+      "epoch": 0.11,
+      "gumbel_tau": 0.7833,
+      "image_masks_loss": 0.02,
+      "llm_learning_rate": 4.9243931144296075e-06,
+      "loss": 0.7591,
+      "output_text_mask_loss": 0.0248,
+      "predictor_learning_rate": 0.0004924393114429607,
+      "step": 1036
+    },
+    {
+      "epoch": 0.11,
+      "gumbel_tau": 0.7832,
+      "image_masks_loss": 0.0246,
+      "llm_learning_rate": 4.924190495164327e-06,
+      "loss": 0.8173,
+      "output_text_mask_loss": 0.0324,
+      "predictor_learning_rate": 0.0004924190495164326,
+      "step": 1037
+    },
+    {
+      "epoch": 0.11,
+      "gumbel_tau": 0.783,
+      "image_masks_loss": 0.0116,
+      "llm_learning_rate": 4.923987608941789e-06,
+      "loss": 0.8556,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.000492398760894179,
+      "step": 1038
+    },
+    {
+      "epoch": 0.11,
+      "gumbel_tau": 0.7828,
+      "image_masks_loss": 0.0349,
+      "llm_learning_rate": 4.9237844557843375e-06,
+      "loss": 0.7563,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004923784455784337,
+      "step": 1039
+    },
+    {
+      "epoch": 0.11,
+      "gumbel_tau": 0.7826,
+      "image_masks_loss": 0.0237,
+      "llm_learning_rate": 4.923581035714342e-06,
+      "loss": 0.7734,
+      "output_text_mask_loss": 0.0278,
+      "predictor_learning_rate": 0.0004923581035714342,
+      "step": 1040
+    },
+    {
+      "epoch": 0.11,
+      "gumbel_tau": 0.7824,
+      "image_masks_loss": 0.0336,
+      "llm_learning_rate": 4.923377348754205e-06,
+      "loss": 0.8212,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.0004923377348754204,
+      "step": 1041
+    },
+    {
+      "epoch": 0.11,
+      "gumbel_tau": 0.7822,
+      "image_masks_loss": 0.033,
+      "llm_learning_rate": 4.923173394926356e-06,
+      "loss": 0.7035,
+      "output_text_mask_loss": 0.0736,
+      "predictor_learning_rate": 0.0004923173394926355,
+      "step": 1042
+    },
+    {
+      "epoch": 0.11,
+      "gumbel_tau": 0.7821,
+      "image_masks_loss": 0.0136,
+      "llm_learning_rate": 4.9229691742532556e-06,
+      "loss": 0.6574,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004922969174253255,
+      "step": 1043
+    },
+    {
+      "epoch": 0.11,
+      "gumbel_tau": 0.7819,
+      "image_masks_loss": 0.0366,
+      "llm_learning_rate": 4.922764686757393e-06,
+      "loss": 0.7539,
+      "output_text_mask_loss": 0.0488,
+      "predictor_learning_rate": 0.0004922764686757393,
+      "step": 1044
+    },
+    {
+      "epoch": 0.11,
+      "gumbel_tau": 0.7817,
+      "image_masks_loss": 0.0165,
+      "llm_learning_rate": 4.922559932461286e-06,
+      "loss": 0.7111,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 0.0004922559932461286,
+      "step": 1045
+    },
+    {
+      "epoch": 0.11,
+      "gumbel_tau": 0.7815,
+      "image_masks_loss": 0.0119,
+      "llm_learning_rate": 4.922354911387484e-06,
+      "loss": 0.7032,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.0004922354911387483,
+      "step": 1046
+    },
+    {
+      "epoch": 0.11,
+      "gumbel_tau": 0.7813,
+      "image_masks_loss": 0.0059,
+      "llm_learning_rate": 4.922149623558563e-06,
+      "loss": 0.8235,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004922149623558563,
+      "step": 1047
+    },
+    {
+      "epoch": 0.11,
+      "gumbel_tau": 0.7811,
+      "image_masks_loss": 0.0177,
+      "llm_learning_rate": 4.92194406899713e-06,
+      "loss": 0.8386,
+      "output_text_mask_loss": 0.0223,
+      "predictor_learning_rate": 0.000492194406899713,
+      "step": 1048
+    },
+    {
+      "epoch": 0.11,
+      "gumbel_tau": 0.7809,
+      "image_masks_loss": 0.0099,
+      "llm_learning_rate": 4.921738247725823e-06,
+      "loss": 0.8669,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.0004921738247725822,
+      "step": 1049
+    },
+    {
+      "epoch": 0.11,
+      "gumbel_tau": 0.7808,
+      "image_masks_loss": 0.0122,
+      "llm_learning_rate": 4.921532159767304e-06,
+      "loss": 0.8342,
+      "output_text_mask_loss": 0.0269,
+      "predictor_learning_rate": 0.0004921532159767304,
+      "step": 1050
+    },
+    {
+      "epoch": 0.11,
+      "gumbel_tau": 0.7806,
+      "image_masks_loss": 0.018,
+      "llm_learning_rate": 4.921325805144271e-06,
+      "loss": 0.7311,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.000492132580514427,
+      "step": 1051
+    },
+    {
+      "epoch": 0.11,
+      "gumbel_tau": 0.7804,
+      "image_masks_loss": 0.0158,
+      "llm_learning_rate": 4.921119183879447e-06,
+      "loss": 0.8645,
+      "output_text_mask_loss": 0.0313,
+      "predictor_learning_rate": 0.0004921119183879446,
+      "step": 1052
+    },
+    {
+      "epoch": 0.11,
+      "gumbel_tau": 0.7802,
+      "image_masks_loss": 0.0092,
+      "llm_learning_rate": 4.920912295995585e-06,
+      "loss": 0.7201,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004920912295995584,
+      "step": 1053
+    },
+    {
+      "epoch": 0.11,
+      "gumbel_tau": 0.78,
+      "image_masks_loss": 0.0149,
+      "llm_learning_rate": 4.920705141515469e-06,
+      "loss": 0.7468,
+      "output_text_mask_loss": 0.0088,
+      "predictor_learning_rate": 0.0004920705141515469,
+      "step": 1054
+    },
+    {
+      "epoch": 0.11,
+      "gumbel_tau": 0.7798,
+      "image_masks_loss": 0.0119,
+      "llm_learning_rate": 4.920497720461911e-06,
+      "loss": 0.7724,
+      "output_text_mask_loss": 0.0603,
+      "predictor_learning_rate": 0.000492049772046191,
+      "step": 1055
+    },
+    {
+      "epoch": 0.11,
+      "gumbel_tau": 0.7797,
+      "image_masks_loss": 0.0087,
+      "llm_learning_rate": 4.920290032857752e-06,
+      "loss": 0.792,
+      "output_text_mask_loss": 0.0231,
+      "predictor_learning_rate": 0.0004920290032857751,
+      "step": 1056
+    },
+    {
+      "epoch": 0.11,
+      "gumbel_tau": 0.7795,
+      "image_masks_loss": 0.0111,
+      "llm_learning_rate": 4.920082078725864e-06,
+      "loss": 0.8197,
+      "output_text_mask_loss": 0.0488,
+      "predictor_learning_rate": 0.0004920082078725864,
+      "step": 1057
+    },
+    {
+      "epoch": 0.11,
+      "gumbel_tau": 0.7793,
+      "image_masks_loss": 0.0149,
+      "llm_learning_rate": 4.919873858089147e-06,
+      "loss": 0.7845,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.0004919873858089147,
+      "step": 1058
+    },
+    {
+      "epoch": 0.11,
+      "gumbel_tau": 0.7791,
+      "image_masks_loss": 0.0205,
+      "llm_learning_rate": 4.919665370970531e-06,
+      "loss": 0.8222,
+      "output_text_mask_loss": 0.0087,
+      "predictor_learning_rate": 0.0004919665370970531,
+      "step": 1059
+    },
+    {
+      "epoch": 0.11,
+      "gumbel_tau": 0.7789,
+      "image_masks_loss": 0.0186,
+      "llm_learning_rate": 4.919456617392975e-06,
+      "loss": 0.7463,
+      "output_text_mask_loss": 0.0071,
+      "predictor_learning_rate": 0.0004919456617392974,
+      "step": 1060
+    },
+    {
+      "epoch": 0.11,
+      "gumbel_tau": 0.7787,
+      "image_masks_loss": 0.0176,
+      "llm_learning_rate": 4.9192475973794666e-06,
+      "loss": 0.7315,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004919247597379466,
+      "step": 1061
+    },
+    {
+      "epoch": 0.11,
+      "gumbel_tau": 0.7786,
+      "image_masks_loss": 0.0145,
+      "llm_learning_rate": 4.919038310953024e-06,
+      "loss": 0.806,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.0004919038310953023,
+      "step": 1062
+    },
+    {
+      "epoch": 0.11,
+      "gumbel_tau": 0.7784,
+      "image_masks_loss": 0.0383,
+      "llm_learning_rate": 4.918828758136694e-06,
+      "loss": 0.7443,
+      "output_text_mask_loss": 0.0012,
+      "predictor_learning_rate": 0.0004918828758136693,
+      "step": 1063
+    },
+    {
+      "epoch": 0.11,
+      "gumbel_tau": 0.7782,
+      "image_masks_loss": 0.0116,
+      "llm_learning_rate": 4.918618938953555e-06,
+      "loss": 0.85,
+      "output_text_mask_loss": 0.036,
+      "predictor_learning_rate": 0.0004918618938953554,
+      "step": 1064
+    },
+    {
+      "epoch": 0.11,
+      "gumbel_tau": 0.778,
+      "image_masks_loss": 0.0301,
+      "llm_learning_rate": 4.918408853426709e-06,
+      "loss": 0.8179,
+      "output_text_mask_loss": 0.0193,
+      "predictor_learning_rate": 0.0004918408853426709,
+      "step": 1065
+    },
+    {
+      "epoch": 0.11,
+      "gumbel_tau": 0.7778,
+      "image_masks_loss": 0.0246,
+      "llm_learning_rate": 4.9181985015792945e-06,
+      "loss": 0.8299,
+      "output_text_mask_loss": 0.0242,
+      "predictor_learning_rate": 0.0004918198501579294,
+      "step": 1066
+    },
+    {
+      "epoch": 0.11,
+      "gumbel_tau": 0.7776,
+      "image_masks_loss": 0.0179,
+      "llm_learning_rate": 4.917987883434475e-06,
+      "loss": 0.7334,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.0004917987883434475,
+      "step": 1067
+    },
+    {
+      "epoch": 0.11,
+      "gumbel_tau": 0.7775,
+      "image_masks_loss": 0.029,
+      "llm_learning_rate": 4.917776999015443e-06,
+      "loss": 0.9328,
+      "output_text_mask_loss": 0.1442,
+      "predictor_learning_rate": 0.0004917776999015443,
+      "step": 1068
+    },
+    {
+      "epoch": 0.11,
+      "gumbel_tau": 0.7773,
+      "image_masks_loss": 0.0174,
+      "llm_learning_rate": 4.917565848345423e-06,
+      "loss": 0.7868,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 0.0004917565848345423,
+      "step": 1069
+    },
+    {
+      "epoch": 0.11,
+      "gumbel_tau": 0.7771,
+      "image_masks_loss": 0.0224,
+      "llm_learning_rate": 4.917354431447667e-06,
+      "loss": 0.8958,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004917354431447667,
+      "step": 1070
+    },
+    {
+      "epoch": 0.11,
+      "gumbel_tau": 0.7769,
+      "image_masks_loss": 0.023,
+      "llm_learning_rate": 4.917142748345457e-06,
+      "loss": 0.8439,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 0.0004917142748345456,
+      "step": 1071
+    },
+    {
+      "epoch": 0.11,
+      "gumbel_tau": 0.7767,
+      "image_masks_loss": 0.0191,
+      "llm_learning_rate": 4.916930799062103e-06,
+      "loss": 0.7227,
+      "output_text_mask_loss": 0.0355,
+      "predictor_learning_rate": 0.0004916930799062102,
+      "step": 1072
+    },
+    {
+      "epoch": 0.11,
+      "gumbel_tau": 0.7765,
+      "image_masks_loss": 0.0104,
+      "llm_learning_rate": 4.9167185836209455e-06,
+      "loss": 0.7548,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.0004916718583620945,
+      "step": 1073
+    },
+    {
+      "epoch": 0.11,
+      "gumbel_tau": 0.7764,
+      "image_masks_loss": 0.0131,
+      "llm_learning_rate": 4.916506102045355e-06,
+      "loss": 0.8493,
+      "output_text_mask_loss": 0.0025,
+      "predictor_learning_rate": 0.0004916506102045354,
+      "step": 1074
+    },
+    {
+      "epoch": 0.11,
+      "gumbel_tau": 0.7762,
+      "image_masks_loss": 0.0269,
+      "llm_learning_rate": 4.9162933543587295e-06,
+      "loss": 0.7211,
+      "output_text_mask_loss": 0.0353,
+      "predictor_learning_rate": 0.0004916293354358729,
+      "step": 1075
+    },
+    {
+      "epoch": 0.11,
+      "gumbel_tau": 0.776,
+      "image_masks_loss": 0.0072,
+      "llm_learning_rate": 4.916080340584498e-06,
+      "loss": 0.8409,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.0004916080340584498,
+      "step": 1076
+    },
+    {
+      "epoch": 0.11,
+      "gumbel_tau": 0.7758,
+      "image_masks_loss": 0.0134,
+      "llm_learning_rate": 4.915867060746118e-06,
+      "loss": 0.7858,
+      "output_text_mask_loss": 0.0261,
+      "predictor_learning_rate": 0.0004915867060746118,
+      "step": 1077
+    },
+    {
+      "epoch": 0.11,
+      "gumbel_tau": 0.7756,
+      "image_masks_loss": 0.0343,
+      "llm_learning_rate": 4.915653514867076e-06,
+      "loss": 0.7399,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.0004915653514867076,
+      "step": 1078
+    },
+    {
+      "epoch": 0.11,
+      "gumbel_tau": 0.7754,
+      "image_masks_loss": 0.029,
+      "llm_learning_rate": 4.915439702970889e-06,
+      "loss": 0.8416,
+      "output_text_mask_loss": 0.0046,
+      "predictor_learning_rate": 0.0004915439702970889,
+      "step": 1079
+    },
+    {
+      "epoch": 0.11,
+      "gumbel_tau": 0.7753,
+      "image_masks_loss": 0.0265,
+      "llm_learning_rate": 4.9152256250811e-06,
+      "loss": 0.8535,
+      "output_text_mask_loss": 0.0307,
+      "predictor_learning_rate": 0.00049152256250811,
+      "step": 1080
+    },
+    {
+      "epoch": 0.11,
+      "gumbel_tau": 0.7751,
+      "image_masks_loss": 0.0089,
+      "llm_learning_rate": 4.915011281221287e-06,
+      "loss": 0.7945,
+      "output_text_mask_loss": 0.0117,
+      "predictor_learning_rate": 0.0004915011281221287,
+      "step": 1081
+    },
+    {
+      "epoch": 0.11,
+      "gumbel_tau": 0.7749,
+      "image_masks_loss": 0.008,
+      "llm_learning_rate": 4.914796671415052e-06,
+      "loss": 0.8142,
+      "output_text_mask_loss": 0.0126,
+      "predictor_learning_rate": 0.0004914796671415052,
+      "step": 1082
+    },
+    {
+      "epoch": 0.11,
+      "gumbel_tau": 0.7747,
+      "image_masks_loss": 0.0273,
+      "llm_learning_rate": 4.9145817956860286e-06,
+      "loss": 0.6463,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 0.0004914581795686029,
+      "step": 1083
+    },
+    {
+      "epoch": 0.11,
+      "gumbel_tau": 0.7745,
+      "image_masks_loss": 0.0208,
+      "llm_learning_rate": 4.914366654057881e-06,
+      "loss": 0.7822,
+      "output_text_mask_loss": 0.0151,
+      "predictor_learning_rate": 0.000491436665405788,
+      "step": 1084
+    },
+    {
+      "epoch": 0.11,
+      "gumbel_tau": 0.7743,
+      "image_masks_loss": 0.0584,
+      "llm_learning_rate": 4.914151246554298e-06,
+      "loss": 0.8066,
+      "output_text_mask_loss": 0.0231,
+      "predictor_learning_rate": 0.0004914151246554298,
+      "step": 1085
+    },
+    {
+      "epoch": 0.11,
+      "gumbel_tau": 0.7742,
+      "image_masks_loss": 0.0479,
+      "llm_learning_rate": 4.913935573199003e-06,
+      "loss": 0.8394,
+      "output_text_mask_loss": 0.0385,
+      "predictor_learning_rate": 0.0004913935573199002,
+      "step": 1086
+    },
+    {
+      "epoch": 0.11,
+      "gumbel_tau": 0.774,
+      "image_masks_loss": 0.0267,
+      "llm_learning_rate": 4.913719634015745e-06,
+      "loss": 0.8014,
+      "output_text_mask_loss": 0.0877,
+      "predictor_learning_rate": 0.0004913719634015745,
+      "step": 1087
+    },
+    {
+      "epoch": 0.11,
+      "gumbel_tau": 0.7738,
+      "image_masks_loss": 0.0248,
+      "llm_learning_rate": 4.913503429028305e-06,
+      "loss": 0.8676,
+      "output_text_mask_loss": 0.0877,
+      "predictor_learning_rate": 0.0004913503429028306,
+      "step": 1088
+    },
+    {
+      "epoch": 0.11,
+      "gumbel_tau": 0.7736,
+      "image_masks_loss": 0.02,
+      "llm_learning_rate": 4.913286958260493e-06,
+      "loss": 0.7339,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004913286958260492,
+      "step": 1089
+    },
+    {
+      "epoch": 0.11,
+      "gumbel_tau": 0.7734,
+      "image_masks_loss": 0.0359,
+      "llm_learning_rate": 4.913070221736144e-06,
+      "loss": 0.8598,
+      "output_text_mask_loss": 0.0208,
+      "predictor_learning_rate": 0.0004913070221736144,
+      "step": 1090
+    },
+    {
+      "epoch": 0.11,
+      "gumbel_tau": 0.7732,
+      "image_masks_loss": 0.0154,
+      "llm_learning_rate": 4.912853219479128e-06,
+      "loss": 0.7695,
+      "output_text_mask_loss": 0.0112,
+      "predictor_learning_rate": 0.0004912853219479127,
+      "step": 1091
+    },
+    {
+      "epoch": 0.11,
+      "gumbel_tau": 0.7731,
+      "image_masks_loss": 0.0228,
+      "llm_learning_rate": 4.912635951513341e-06,
+      "loss": 0.6625,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004912635951513341,
+      "step": 1092
+    },
+    {
+      "epoch": 0.11,
+      "gumbel_tau": 0.7729,
+      "image_masks_loss": 0.0066,
+      "llm_learning_rate": 4.912418417862709e-06,
+      "loss": 0.8088,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004912418417862708,
+      "step": 1093
+    },
+    {
+      "epoch": 0.11,
+      "gumbel_tau": 0.7727,
+      "image_masks_loss": 0.0277,
+      "llm_learning_rate": 4.912200618551187e-06,
+      "loss": 0.7647,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004912200618551187,
+      "step": 1094
+    },
+    {
+      "epoch": 0.11,
+      "gumbel_tau": 0.7725,
+      "image_masks_loss": 0.0278,
+      "llm_learning_rate": 4.911982553602761e-06,
+      "loss": 0.8084,
+      "output_text_mask_loss": 0.0145,
+      "predictor_learning_rate": 0.0004911982553602761,
+      "step": 1095
+    },
+    {
+      "epoch": 0.11,
+      "gumbel_tau": 0.7723,
+      "image_masks_loss": 0.0317,
+      "llm_learning_rate": 4.911764223041443e-06,
+      "loss": 0.7168,
+      "output_text_mask_loss": 0.0248,
+      "predictor_learning_rate": 0.0004911764223041443,
+      "step": 1096
+    },
+    {
+      "epoch": 0.11,
+      "gumbel_tau": 0.7722,
+      "image_masks_loss": 0.0133,
+      "llm_learning_rate": 4.9115456268912765e-06,
+      "loss": 0.7502,
+      "output_text_mask_loss": 0.0099,
+      "predictor_learning_rate": 0.0004911545626891277,
+      "step": 1097
+    },
+    {
+      "epoch": 0.11,
+      "gumbel_tau": 0.772,
+      "image_masks_loss": 0.0177,
+      "llm_learning_rate": 4.911326765176334e-06,
+      "loss": 0.8359,
+      "output_text_mask_loss": 0.2121,
+      "predictor_learning_rate": 0.0004911326765176334,
+      "step": 1098
+    },
+    {
+      "epoch": 0.11,
+      "gumbel_tau": 0.7718,
+      "image_masks_loss": 0.0244,
+      "llm_learning_rate": 4.911107637920718e-06,
+      "loss": 0.8962,
+      "output_text_mask_loss": 0.0443,
+      "predictor_learning_rate": 0.0004911107637920718,
+      "step": 1099
+    },
+    {
+      "epoch": 0.11,
+      "gumbel_tau": 0.7716,
+      "image_masks_loss": 0.0397,
+      "llm_learning_rate": 4.910888245148557e-06,
+      "loss": 0.8477,
+      "output_text_mask_loss": 0.0385,
+      "predictor_learning_rate": 0.0004910888245148557,
+      "step": 1100
+    },
+    {
+      "epoch": 0.11,
+      "gumbel_tau": 0.7714,
+      "image_masks_loss": 0.0039,
+      "llm_learning_rate": 4.910668586884014e-06,
+      "loss": 0.7788,
+      "output_text_mask_loss": 0.1404,
+      "predictor_learning_rate": 0.0004910668586884014,
+      "step": 1101
+    },
+    {
+      "epoch": 0.11,
+      "gumbel_tau": 0.7712,
+      "image_masks_loss": 0.014,
+      "llm_learning_rate": 4.910448663151276e-06,
+      "loss": 0.8079,
+      "output_text_mask_loss": 0.0851,
+      "predictor_learning_rate": 0.0004910448663151275,
+      "step": 1102
+    },
+    {
+      "epoch": 0.11,
+      "gumbel_tau": 0.7711,
+      "image_masks_loss": 0.0305,
+      "llm_learning_rate": 4.910228473974562e-06,
+      "loss": 0.8448,
+      "output_text_mask_loss": 0.038,
+      "predictor_learning_rate": 0.0004910228473974562,
+      "step": 1103
+    },
+    {
+      "epoch": 0.11,
+      "gumbel_tau": 0.7709,
+      "image_masks_loss": 0.0263,
+      "llm_learning_rate": 4.91000801937812e-06,
+      "loss": 0.6759,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.000491000801937812,
+      "step": 1104
+    },
+    {
+      "epoch": 0.11,
+      "gumbel_tau": 0.7707,
+      "image_masks_loss": 0.0425,
+      "llm_learning_rate": 4.9097872993862265e-06,
+      "loss": 0.8338,
+      "output_text_mask_loss": 0.1434,
+      "predictor_learning_rate": 0.0004909787299386226,
+      "step": 1105
+    },
+    {
+      "epoch": 0.11,
+      "gumbel_tau": 0.7705,
+      "image_masks_loss": 0.0238,
+      "llm_learning_rate": 4.909566314023188e-06,
+      "loss": 0.6622,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.0004909566314023187,
+      "step": 1106
+    },
+    {
+      "epoch": 0.11,
+      "gumbel_tau": 0.7703,
+      "image_masks_loss": 0.0151,
+      "llm_learning_rate": 4.90934506331334e-06,
+      "loss": 0.8627,
+      "output_text_mask_loss": 0.0412,
+      "predictor_learning_rate": 0.000490934506331334,
+      "step": 1107
+    },
+    {
+      "epoch": 0.11,
+      "gumbel_tau": 0.7702,
+      "image_masks_loss": 0.0198,
+      "llm_learning_rate": 4.909123547281047e-06,
+      "loss": 0.8082,
+      "output_text_mask_loss": 0.0353,
+      "predictor_learning_rate": 0.0004909123547281047,
+      "step": 1108
+    },
+    {
+      "epoch": 0.11,
+      "gumbel_tau": 0.77,
+      "image_masks_loss": 0.0184,
+      "llm_learning_rate": 4.908901765950703e-06,
+      "loss": 0.826,
+      "output_text_mask_loss": 0.0235,
+      "predictor_learning_rate": 0.0004908901765950703,
+      "step": 1109
+    },
+    {
+      "epoch": 0.11,
+      "gumbel_tau": 0.7698,
+      "image_masks_loss": 0.0193,
+      "llm_learning_rate": 4.908679719346732e-06,
+      "loss": 0.8036,
+      "output_text_mask_loss": 0.0549,
+      "predictor_learning_rate": 0.0004908679719346731,
+      "step": 1110
+    },
+    {
+      "epoch": 0.11,
+      "gumbel_tau": 0.7696,
+      "image_masks_loss": 0.0248,
+      "llm_learning_rate": 4.908457407493583e-06,
+      "loss": 0.7548,
+      "output_text_mask_loss": 0.0626,
+      "predictor_learning_rate": 0.0004908457407493583,
+      "step": 1111
+    },
+    {
+      "epoch": 0.11,
+      "gumbel_tau": 0.7694,
+      "image_masks_loss": 0.029,
+      "llm_learning_rate": 4.908234830415741e-06,
+      "loss": 0.6821,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004908234830415741,
+      "step": 1112
+    },
+    {
+      "epoch": 0.11,
+      "gumbel_tau": 0.7692,
+      "image_masks_loss": 0.0429,
+      "llm_learning_rate": 4.908011988137715e-06,
+      "loss": 0.8817,
+      "output_text_mask_loss": 0.0847,
+      "predictor_learning_rate": 0.0004908011988137715,
+      "step": 1113
+    },
+    {
+      "epoch": 0.11,
+      "gumbel_tau": 0.7691,
+      "image_masks_loss": 0.0481,
+      "llm_learning_rate": 4.907788880684045e-06,
+      "loss": 0.8048,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.0004907788880684045,
+      "step": 1114
+    },
+    {
+      "epoch": 0.11,
+      "gumbel_tau": 0.7689,
+      "image_masks_loss": 0.0271,
+      "llm_learning_rate": 4.907565508079302e-06,
+      "loss": 0.8042,
+      "output_text_mask_loss": 0.0322,
+      "predictor_learning_rate": 0.0004907565508079302,
+      "step": 1115
+    },
+    {
+      "epoch": 0.11,
+      "gumbel_tau": 0.7687,
+      "image_masks_loss": 0.0319,
+      "llm_learning_rate": 4.90734187034808e-06,
+      "loss": 0.8044,
+      "output_text_mask_loss": 0.033,
+      "predictor_learning_rate": 0.0004907341870348079,
+      "step": 1116
+    },
+    {
+      "epoch": 0.11,
+      "gumbel_tau": 0.7685,
+      "image_masks_loss": 0.0232,
+      "llm_learning_rate": 4.907117967515011e-06,
+      "loss": 0.804,
+      "output_text_mask_loss": 0.0391,
+      "predictor_learning_rate": 0.000490711796751501,
+      "step": 1117
+    },
+    {
+      "epoch": 0.11,
+      "gumbel_tau": 0.7683,
+      "image_masks_loss": 0.0448,
+      "llm_learning_rate": 4.90689379960475e-06,
+      "loss": 0.746,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.000490689379960475,
+      "step": 1118
+    },
+    {
+      "epoch": 0.11,
+      "gumbel_tau": 0.7682,
+      "image_masks_loss": 0.0123,
+      "llm_learning_rate": 4.906669366641981e-06,
+      "loss": 0.7618,
+      "output_text_mask_loss": 0.0324,
+      "predictor_learning_rate": 0.0004906669366641981,
+      "step": 1119
+    },
+    {
+      "epoch": 0.11,
+      "gumbel_tau": 0.768,
+      "image_masks_loss": 0.0092,
+      "llm_learning_rate": 4.906444668651421e-06,
+      "loss": 0.7622,
+      "output_text_mask_loss": 0.0429,
+      "predictor_learning_rate": 0.0004906444668651421,
+      "step": 1120
+    },
+    {
+      "epoch": 0.11,
+      "gumbel_tau": 0.7678,
+      "image_masks_loss": 0.032,
+      "llm_learning_rate": 4.906219705657815e-06,
+      "loss": 0.8947,
+      "output_text_mask_loss": 0.0275,
+      "predictor_learning_rate": 0.0004906219705657815,
+      "step": 1121
+    },
+    {
+      "epoch": 0.11,
+      "gumbel_tau": 0.7676,
+      "image_masks_loss": 0.0241,
+      "llm_learning_rate": 4.905994477685934e-06,
+      "loss": 0.7512,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.0004905994477685935,
+      "step": 1122
+    },
+    {
+      "epoch": 0.12,
+      "gumbel_tau": 0.7674,
+      "image_masks_loss": 0.0196,
+      "llm_learning_rate": 4.905768984760583e-06,
+      "loss": 0.749,
+      "output_text_mask_loss": 0.0347,
+      "predictor_learning_rate": 0.0004905768984760583,
+      "step": 1123
+    },
+    {
+      "epoch": 0.12,
+      "gumbel_tau": 0.7673,
+      "image_masks_loss": 0.0128,
+      "llm_learning_rate": 4.905543226906593e-06,
+      "loss": 0.7638,
+      "output_text_mask_loss": 0.0275,
+      "predictor_learning_rate": 0.0004905543226906593,
+      "step": 1124
+    },
+    {
+      "epoch": 0.12,
+      "gumbel_tau": 0.7671,
+      "image_masks_loss": 0.0278,
+      "llm_learning_rate": 4.905317204148824e-06,
+      "loss": 0.7505,
+      "output_text_mask_loss": 0.0139,
+      "predictor_learning_rate": 0.0004905317204148824,
+      "step": 1125
+    },
+    {
+      "epoch": 0.12,
+      "gumbel_tau": 0.7669,
+      "image_masks_loss": 0.0128,
+      "llm_learning_rate": 4.905090916512167e-06,
+      "loss": 0.7139,
+      "output_text_mask_loss": 0.0231,
+      "predictor_learning_rate": 0.0004905090916512168,
+      "step": 1126
+    },
+    {
+      "epoch": 0.12,
+      "gumbel_tau": 0.7667,
+      "image_masks_loss": 0.0301,
+      "llm_learning_rate": 4.904864364021541e-06,
+      "loss": 0.7739,
+      "output_text_mask_loss": 0.0618,
+      "predictor_learning_rate": 0.000490486436402154,
+      "step": 1127
+    },
+    {
+      "epoch": 0.12,
+      "gumbel_tau": 0.7665,
+      "image_masks_loss": 0.0235,
+      "llm_learning_rate": 4.904637546701895e-06,
+      "loss": 0.8849,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.0004904637546701894,
+      "step": 1128
+    },
+    {
+      "epoch": 0.12,
+      "gumbel_tau": 0.7663,
+      "image_masks_loss": 0.0132,
+      "llm_learning_rate": 4.9044104645782066e-06,
+      "loss": 0.7011,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 0.0004904410464578207,
+      "step": 1129
+    },
+    {
+      "epoch": 0.12,
+      "gumbel_tau": 0.7662,
+      "image_masks_loss": 0.0315,
+      "llm_learning_rate": 4.904183117675481e-06,
+      "loss": 0.8161,
+      "output_text_mask_loss": 0.0039,
+      "predictor_learning_rate": 0.000490418311767548,
+      "step": 1130
+    },
+    {
+      "epoch": 0.12,
+      "gumbel_tau": 0.766,
+      "image_masks_loss": 0.0391,
+      "llm_learning_rate": 4.903955506018756e-06,
+      "loss": 0.8023,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.0004903955506018755,
+      "step": 1131
+    },
+    {
+      "epoch": 0.12,
+      "gumbel_tau": 0.7658,
+      "image_masks_loss": 0.0353,
+      "llm_learning_rate": 4.9037276296330965e-06,
+      "loss": 0.8086,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 0.0004903727629633096,
+      "step": 1132
+    },
+    {
+      "epoch": 0.12,
+      "gumbel_tau": 0.7656,
+      "image_masks_loss": 0.0299,
+      "llm_learning_rate": 4.903499488543596e-06,
+      "loss": 0.7569,
+      "output_text_mask_loss": 0.0385,
+      "predictor_learning_rate": 0.0004903499488543596,
+      "step": 1133
+    },
+    {
+      "epoch": 0.12,
+      "gumbel_tau": 0.7654,
+      "image_masks_loss": 0.0431,
+      "llm_learning_rate": 4.903271082775378e-06,
+      "loss": 0.7741,
+      "output_text_mask_loss": 0.006,
+      "predictor_learning_rate": 0.0004903271082775378,
+      "step": 1134
+    },
+    {
+      "epoch": 0.12,
+      "gumbel_tau": 0.7653,
+      "image_masks_loss": 0.0292,
+      "llm_learning_rate": 4.903042412353596e-06,
+      "loss": 0.7479,
+      "output_text_mask_loss": 0.0277,
+      "predictor_learning_rate": 0.0004903042412353596,
+      "step": 1135
+    },
+    {
+      "epoch": 0.12,
+      "gumbel_tau": 0.7651,
+      "image_masks_loss": 0.0216,
+      "llm_learning_rate": 4.9028134773034305e-06,
+      "loss": 0.8502,
+      "output_text_mask_loss": 0.0429,
+      "predictor_learning_rate": 0.000490281347730343,
+      "step": 1136
+    },
+    {
+      "epoch": 0.12,
+      "gumbel_tau": 0.7649,
+      "image_masks_loss": 0.0218,
+      "llm_learning_rate": 4.902584277650092e-06,
+      "loss": 0.7727,
+      "output_text_mask_loss": 0.0271,
+      "predictor_learning_rate": 0.0004902584277650092,
+      "step": 1137
+    },
+    {
+      "epoch": 0.12,
+      "gumbel_tau": 0.7647,
+      "image_masks_loss": 0.0193,
+      "llm_learning_rate": 4.902354813418822e-06,
+      "loss": 0.8642,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.0004902354813418823,
+      "step": 1138
+    },
+    {
+      "epoch": 0.12,
+      "gumbel_tau": 0.7645,
+      "image_masks_loss": 0.0299,
+      "llm_learning_rate": 4.902125084634889e-06,
+      "loss": 0.8105,
+      "output_text_mask_loss": 0.0729,
+      "predictor_learning_rate": 0.0004902125084634889,
+      "step": 1139
+    },
+    {
+      "epoch": 0.12,
+      "gumbel_tau": 0.7644,
+      "image_masks_loss": 0.0256,
+      "llm_learning_rate": 4.901895091323592e-06,
+      "loss": 0.8409,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.0004901895091323591,
+      "step": 1140
+    },
+    {
+      "epoch": 0.12,
+      "gumbel_tau": 0.7642,
+      "image_masks_loss": 0.034,
+      "llm_learning_rate": 4.901664833510257e-06,
+      "loss": 0.7963,
+      "output_text_mask_loss": 0.0012,
+      "predictor_learning_rate": 0.0004901664833510256,
+      "step": 1141
+    },
+    {
+      "epoch": 0.12,
+      "gumbel_tau": 0.764,
+      "image_masks_loss": 0.0154,
+      "llm_learning_rate": 4.90143431122024e-06,
+      "loss": 0.8109,
+      "output_text_mask_loss": 0.0039,
+      "predictor_learning_rate": 0.000490143431122024,
+      "step": 1142
+    },
+    {
+      "epoch": 0.12,
+      "gumbel_tau": 0.7638,
+      "image_masks_loss": 0.0338,
+      "llm_learning_rate": 4.901203524478929e-06,
+      "loss": 0.7597,
+      "output_text_mask_loss": 0.0496,
+      "predictor_learning_rate": 0.0004901203524478928,
+      "step": 1143
+    },
+    {
+      "epoch": 0.12,
+      "gumbel_tau": 0.7636,
+      "image_masks_loss": 0.0173,
+      "llm_learning_rate": 4.900972473311737e-06,
+      "loss": 0.7284,
+      "output_text_mask_loss": 0.0322,
+      "predictor_learning_rate": 0.0004900972473311737,
+      "step": 1144
+    },
+    {
+      "epoch": 0.12,
+      "gumbel_tau": 0.7635,
+      "image_masks_loss": 0.0383,
+      "llm_learning_rate": 4.900741157744108e-06,
+      "loss": 0.8172,
+      "output_text_mask_loss": 0.0429,
+      "predictor_learning_rate": 0.0004900741157744107,
+      "step": 1145
+    },
+    {
+      "epoch": 0.12,
+      "gumbel_tau": 0.7633,
+      "image_masks_loss": 0.016,
+      "llm_learning_rate": 4.900509577801515e-06,
+      "loss": 0.786,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.0004900509577801514,
+      "step": 1146
+    },
+    {
+      "epoch": 0.12,
+      "gumbel_tau": 0.7631,
+      "image_masks_loss": 0.028,
+      "llm_learning_rate": 4.900277733509461e-06,
+      "loss": 0.738,
+      "output_text_mask_loss": 0.0429,
+      "predictor_learning_rate": 0.0004900277733509461,
+      "step": 1147
+    },
+    {
+      "epoch": 0.12,
+      "gumbel_tau": 0.7629,
+      "image_masks_loss": 0.038,
+      "llm_learning_rate": 4.900045624893476e-06,
+      "loss": 0.8011,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.0004900045624893476,
+      "step": 1148
+    },
+    {
+      "epoch": 0.12,
+      "gumbel_tau": 0.7627,
+      "image_masks_loss": 0.0237,
+      "llm_learning_rate": 4.899813251979121e-06,
+      "loss": 0.881,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 0.000489981325197912,
+      "step": 1149
+    },
+    {
+      "epoch": 0.12,
+      "gumbel_tau": 0.7626,
+      "image_masks_loss": 0.0147,
+      "llm_learning_rate": 4.899580614791986e-06,
+      "loss": 0.7684,
+      "output_text_mask_loss": 0.0095,
+      "predictor_learning_rate": 0.0004899580614791986,
+      "step": 1150
+    },
+    {
+      "epoch": 0.12,
+      "gumbel_tau": 0.7624,
+      "image_masks_loss": 0.022,
+      "llm_learning_rate": 4.899347713357688e-06,
+      "loss": 0.732,
+      "output_text_mask_loss": 0.0523,
+      "predictor_learning_rate": 0.0004899347713357688,
+      "step": 1151
+    },
+    {
+      "epoch": 0.12,
+      "gumbel_tau": 0.7622,
+      "image_masks_loss": 0.022,
+      "llm_learning_rate": 4.899114547701876e-06,
+      "loss": 0.7863,
+      "output_text_mask_loss": 0.0202,
+      "predictor_learning_rate": 0.0004899114547701875,
+      "step": 1152
+    },
+    {
+      "epoch": 0.12,
+      "gumbel_tau": 0.762,
+      "image_masks_loss": 0.0149,
+      "llm_learning_rate": 4.898881117850226e-06,
+      "loss": 0.8253,
+      "output_text_mask_loss": 0.0698,
+      "predictor_learning_rate": 0.0004898881117850226,
+      "step": 1153
+    },
+    {
+      "epoch": 0.12,
+      "gumbel_tau": 0.7618,
+      "image_masks_loss": 0.0226,
+      "llm_learning_rate": 4.898647423828444e-06,
+      "loss": 0.6214,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004898647423828443,
+      "step": 1154
+    },
+    {
+      "epoch": 0.12,
+      "gumbel_tau": 0.7617,
+      "image_masks_loss": 0.0161,
+      "llm_learning_rate": 4.898413465662265e-06,
+      "loss": 0.7243,
+      "output_text_mask_loss": 0.0202,
+      "predictor_learning_rate": 0.0004898413465662264,
+      "step": 1155
+    },
+    {
+      "epoch": 0.12,
+      "gumbel_tau": 0.7615,
+      "image_masks_loss": 0.0399,
+      "llm_learning_rate": 4.8981792433774524e-06,
+      "loss": 0.7005,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 0.0004898179243377452,
+      "step": 1156
+    },
+    {
+      "epoch": 0.12,
+      "gumbel_tau": 0.7613,
+      "image_masks_loss": 0.0364,
+      "llm_learning_rate": 4.897944756999801e-06,
+      "loss": 0.8428,
+      "output_text_mask_loss": 0.0189,
+      "predictor_learning_rate": 0.00048979447569998,
+      "step": 1157
+    },
+    {
+      "epoch": 0.12,
+      "gumbel_tau": 0.7611,
+      "image_masks_loss": 0.0244,
+      "llm_learning_rate": 4.89771000655513e-06,
+      "loss": 0.8416,
+      "output_text_mask_loss": 0.0227,
+      "predictor_learning_rate": 0.000489771000655513,
+      "step": 1158
+    },
+    {
+      "epoch": 0.12,
+      "gumbel_tau": 0.7609,
+      "image_masks_loss": 0.0028,
+      "llm_learning_rate": 4.897474992069294e-06,
+      "loss": 0.6967,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.0004897474992069294,
+      "step": 1159
+    },
+    {
+      "epoch": 0.12,
+      "gumbel_tau": 0.7608,
+      "image_masks_loss": 0.0205,
+      "llm_learning_rate": 4.89723971356817e-06,
+      "loss": 0.8471,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.000489723971356817,
+      "step": 1160
+    },
+    {
+      "epoch": 0.12,
+      "gumbel_tau": 0.7606,
+      "image_masks_loss": 0.034,
+      "llm_learning_rate": 4.89700417107767e-06,
+      "loss": 0.6902,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.000489700417107767,
+      "step": 1161
+    },
+    {
+      "epoch": 0.12,
+      "gumbel_tau": 0.7604,
+      "image_masks_loss": 0.0168,
+      "llm_learning_rate": 4.896768364623733e-06,
+      "loss": 0.8039,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.0004896768364623733,
+      "step": 1162
+    },
+    {
+      "epoch": 0.12,
+      "gumbel_tau": 0.7602,
+      "image_masks_loss": 0.0294,
+      "llm_learning_rate": 4.8965322942323234e-06,
+      "loss": 0.8218,
+      "output_text_mask_loss": 0.025,
+      "predictor_learning_rate": 0.0004896532294232323,
+      "step": 1163
+    },
+    {
+      "epoch": 0.12,
+      "gumbel_tau": 0.76,
+      "image_masks_loss": 0.0224,
+      "llm_learning_rate": 4.89629595992944e-06,
+      "loss": 0.7968,
+      "output_text_mask_loss": 0.0763,
+      "predictor_learning_rate": 0.000489629595992944,
+      "step": 1164
+    },
+    {
+      "epoch": 0.12,
+      "gumbel_tau": 0.7599,
+      "image_masks_loss": 0.0166,
+      "llm_learning_rate": 4.896059361741108e-06,
+      "loss": 0.7817,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004896059361741107,
+      "step": 1165
+    },
+    {
+      "epoch": 0.12,
+      "gumbel_tau": 0.7597,
+      "image_masks_loss": 0.024,
+      "llm_learning_rate": 4.895822499693381e-06,
+      "loss": 0.7963,
+      "output_text_mask_loss": 0.0248,
+      "predictor_learning_rate": 0.0004895822499693381,
+      "step": 1166
+    },
+    {
+      "epoch": 0.12,
+      "gumbel_tau": 0.7595,
+      "image_masks_loss": 0.0299,
+      "llm_learning_rate": 4.895585373812346e-06,
+      "loss": 0.7986,
+      "output_text_mask_loss": 0.013,
+      "predictor_learning_rate": 0.0004895585373812345,
+      "step": 1167
+    },
+    {
+      "epoch": 0.12,
+      "gumbel_tau": 0.7593,
+      "image_masks_loss": 0.0114,
+      "llm_learning_rate": 4.895347984124112e-06,
+      "loss": 0.7434,
+      "output_text_mask_loss": 0.0086,
+      "predictor_learning_rate": 0.0004895347984124112,
+      "step": 1168
+    },
+    {
+      "epoch": 0.12,
+      "gumbel_tau": 0.7592,
+      "image_masks_loss": 0.0319,
+      "llm_learning_rate": 4.895110330654823e-06,
+      "loss": 0.863,
+      "output_text_mask_loss": 0.0012,
+      "predictor_learning_rate": 0.0004895110330654823,
+      "step": 1169
+    },
+    {
+      "epoch": 0.12,
+      "gumbel_tau": 0.759,
+      "image_masks_loss": 0.0299,
+      "llm_learning_rate": 4.8948724134306495e-06,
+      "loss": 0.8022,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.0004894872413430649,
+      "step": 1170
+    },
+    {
+      "epoch": 0.12,
+      "gumbel_tau": 0.7588,
+      "image_masks_loss": 0.0126,
+      "llm_learning_rate": 4.894634232477792e-06,
+      "loss": 0.7855,
+      "output_text_mask_loss": 0.0277,
+      "predictor_learning_rate": 0.0004894634232477791,
+      "step": 1171
+    },
+    {
+      "epoch": 0.12,
+      "gumbel_tau": 0.7586,
+      "image_masks_loss": 0.02,
+      "llm_learning_rate": 4.894395787822478e-06,
+      "loss": 0.8221,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 0.0004894395787822478,
+      "step": 1172
+    },
+    {
+      "epoch": 0.12,
+      "gumbel_tau": 0.7584,
+      "image_masks_loss": 0.0172,
+      "llm_learning_rate": 4.894157079490967e-06,
+      "loss": 0.7409,
+      "output_text_mask_loss": 0.0132,
+      "predictor_learning_rate": 0.0004894157079490967,
+      "step": 1173
+    },
+    {
+      "epoch": 0.12,
+      "gumbel_tau": 0.7583,
+      "image_masks_loss": 0.0086,
+      "llm_learning_rate": 4.893918107509546e-06,
+      "loss": 0.8053,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.0004893918107509545,
+      "step": 1174
+    },
+    {
+      "epoch": 0.12,
+      "gumbel_tau": 0.7581,
+      "image_masks_loss": 0.0374,
+      "llm_learning_rate": 4.89367887190453e-06,
+      "loss": 0.7906,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 0.000489367887190453,
+      "step": 1175
+    },
+    {
+      "epoch": 0.12,
+      "gumbel_tau": 0.7579,
+      "image_masks_loss": 0.0372,
+      "llm_learning_rate": 4.893439372702265e-06,
+      "loss": 0.8097,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.0004893439372702265,
+      "step": 1176
+    },
+    {
+      "epoch": 0.12,
+      "gumbel_tau": 0.7577,
+      "image_masks_loss": 0.028,
+      "llm_learning_rate": 4.8931996099291254e-06,
+      "loss": 0.8149,
+      "output_text_mask_loss": 0.0353,
+      "predictor_learning_rate": 0.0004893199609929125,
+      "step": 1177
+    },
+    {
+      "epoch": 0.12,
+      "gumbel_tau": 0.7575,
+      "image_masks_loss": 0.0322,
+      "llm_learning_rate": 4.892959583611514e-06,
+      "loss": 0.7333,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 0.0004892959583611513,
+      "step": 1178
+    },
+    {
+      "epoch": 0.12,
+      "gumbel_tau": 0.7574,
+      "image_masks_loss": 0.0326,
+      "llm_learning_rate": 4.892719293775864e-06,
+      "loss": 0.8674,
+      "output_text_mask_loss": 0.0725,
+      "predictor_learning_rate": 0.0004892719293775864,
+      "step": 1179
+    },
+    {
+      "epoch": 0.12,
+      "gumbel_tau": 0.7572,
+      "image_masks_loss": 0.0224,
+      "llm_learning_rate": 4.892478740448635e-06,
+      "loss": 0.8304,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.0004892478740448634,
+      "step": 1180
+    },
+    {
+      "epoch": 0.12,
+      "gumbel_tau": 0.757,
+      "image_masks_loss": 0.0132,
+      "llm_learning_rate": 4.8922379236563175e-06,
+      "loss": 0.7462,
+      "output_text_mask_loss": 0.0755,
+      "predictor_learning_rate": 0.0004892237923656318,
+      "step": 1181
+    },
+    {
+      "epoch": 0.12,
+      "gumbel_tau": 0.7568,
+      "image_masks_loss": 0.0315,
+      "llm_learning_rate": 4.891996843425432e-06,
+      "loss": 0.8361,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.0004891996843425432,
+      "step": 1182
+    },
+    {
+      "epoch": 0.12,
+      "gumbel_tau": 0.7566,
+      "image_masks_loss": 0.0191,
+      "llm_learning_rate": 4.891755499782527e-06,
+      "loss": 0.8205,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.0004891755499782526,
+      "step": 1183
+    },
+    {
+      "epoch": 0.12,
+      "gumbel_tau": 0.7565,
+      "image_masks_loss": 0.0156,
+      "llm_learning_rate": 4.8915138927541775e-06,
+      "loss": 0.7453,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004891513892754177,
+      "step": 1184
+    },
+    {
+      "epoch": 0.12,
+      "gumbel_tau": 0.7563,
+      "image_masks_loss": 0.0128,
+      "llm_learning_rate": 4.891272022366992e-06,
+      "loss": 0.6413,
+      "output_text_mask_loss": 0.0507,
+      "predictor_learning_rate": 0.0004891272022366992,
+      "step": 1185
+    },
+    {
+      "epoch": 0.12,
+      "gumbel_tau": 0.7561,
+      "image_masks_loss": 0.0114,
+      "llm_learning_rate": 4.891029888647605e-06,
+      "loss": 0.7682,
+      "output_text_mask_loss": 0.0259,
+      "predictor_learning_rate": 0.0004891029888647604,
+      "step": 1186
+    },
+    {
+      "epoch": 0.12,
+      "gumbel_tau": 0.7559,
+      "image_masks_loss": 0.0173,
+      "llm_learning_rate": 4.8907874916226805e-06,
+      "loss": 0.7514,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.000489078749162268,
+      "step": 1187
+    },
+    {
+      "epoch": 0.12,
+      "gumbel_tau": 0.7558,
+      "image_masks_loss": 0.0393,
+      "llm_learning_rate": 4.890544831318912e-06,
+      "loss": 0.778,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.0004890544831318912,
+      "step": 1188
+    },
+    {
+      "epoch": 0.12,
+      "gumbel_tau": 0.7556,
+      "image_masks_loss": 0.0305,
+      "llm_learning_rate": 4.890301907763022e-06,
+      "loss": 0.7903,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004890301907763022,
+      "step": 1189
+    },
+    {
+      "epoch": 0.12,
+      "gumbel_tau": 0.7554,
+      "image_masks_loss": 0.0218,
+      "llm_learning_rate": 4.890058720981762e-06,
+      "loss": 0.8148,
+      "output_text_mask_loss": 0.0793,
+      "predictor_learning_rate": 0.0004890058720981762,
+      "step": 1190
+    },
+    {
+      "epoch": 0.12,
+      "gumbel_tau": 0.7552,
+      "image_masks_loss": 0.0241,
+      "llm_learning_rate": 4.889815271001912e-06,
+      "loss": 0.8048,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.0004889815271001912,
+      "step": 1191
+    },
+    {
+      "epoch": 0.12,
+      "gumbel_tau": 0.755,
+      "image_masks_loss": 0.0273,
+      "llm_learning_rate": 4.889571557850281e-06,
+      "loss": 0.8346,
+      "output_text_mask_loss": 0.0004,
+      "predictor_learning_rate": 0.000488957155785028,
+      "step": 1192
+    },
+    {
+      "epoch": 0.12,
+      "gumbel_tau": 0.7549,
+      "image_masks_loss": 0.0049,
+      "llm_learning_rate": 4.889327581553709e-06,
+      "loss": 0.7939,
+      "output_text_mask_loss": 0.0397,
+      "predictor_learning_rate": 0.0004889327581553708,
+      "step": 1193
+    },
+    {
+      "epoch": 0.12,
+      "gumbel_tau": 0.7547,
+      "image_masks_loss": 0.0193,
+      "llm_learning_rate": 4.889083342139059e-06,
+      "loss": 0.8616,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004889083342139059,
+      "step": 1194
+    },
+    {
+      "epoch": 0.12,
+      "gumbel_tau": 0.7545,
+      "image_masks_loss": 0.0441,
+      "llm_learning_rate": 4.888838839633232e-06,
+      "loss": 0.7362,
+      "output_text_mask_loss": 0.0012,
+      "predictor_learning_rate": 0.0004888838839633232,
+      "step": 1195
+    },
+    {
+      "epoch": 0.12,
+      "gumbel_tau": 0.7543,
+      "image_masks_loss": 0.0328,
+      "llm_learning_rate": 4.88859407406315e-06,
+      "loss": 0.8226,
+      "output_text_mask_loss": 0.0141,
+      "predictor_learning_rate": 0.000488859407406315,
+      "step": 1196
+    },
+    {
+      "epoch": 0.12,
+      "gumbel_tau": 0.7542,
+      "image_masks_loss": 0.0122,
+      "llm_learning_rate": 4.888349045455769e-06,
+      "loss": 0.7102,
+      "output_text_mask_loss": 0.0072,
+      "predictor_learning_rate": 0.0004888349045455768,
+      "step": 1197
+    },
+    {
+      "epoch": 0.12,
+      "gumbel_tau": 0.754,
+      "image_masks_loss": 0.0183,
+      "llm_learning_rate": 4.888103753838071e-06,
+      "loss": 0.8289,
+      "output_text_mask_loss": 0.0113,
+      "predictor_learning_rate": 0.000488810375383807,
+      "step": 1198
+    },
+    {
+      "epoch": 0.12,
+      "gumbel_tau": 0.7538,
+      "image_masks_loss": 0.0137,
+      "llm_learning_rate": 4.8878581992370675e-06,
+      "loss": 0.8528,
+      "output_text_mask_loss": 0.0111,
+      "predictor_learning_rate": 0.0004887858199237068,
+      "step": 1199
+    },
+    {
+      "epoch": 0.12,
+      "gumbel_tau": 0.7536,
+      "image_masks_loss": 0.061,
+      "llm_learning_rate": 4.887612381679802e-06,
+      "loss": 0.6867,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004887612381679801,
+      "step": 1200
+    },
+    {
+      "epoch": 0.12,
+      "gumbel_tau": 0.7534,
+      "image_masks_loss": 0.0129,
+      "llm_learning_rate": 4.887366301193342e-06,
+      "loss": 0.8225,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.0004887366301193341,
+      "step": 1201
+    },
+    {
+      "epoch": 0.12,
+      "gumbel_tau": 0.7533,
+      "image_masks_loss": 0.0041,
+      "llm_learning_rate": 4.887119957804787e-06,
+      "loss": 0.7508,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.0004887119957804787,
+      "step": 1202
+    },
+    {
+      "epoch": 0.12,
+      "gumbel_tau": 0.7531,
+      "image_masks_loss": 0.0165,
+      "llm_learning_rate": 4.886873351541264e-06,
+      "loss": 0.7107,
+      "output_text_mask_loss": 0.0076,
+      "predictor_learning_rate": 0.0004886873351541265,
+      "step": 1203
+    },
+    {
+      "epoch": 0.12,
+      "gumbel_tau": 0.7529,
+      "image_masks_loss": 0.0069,
+      "llm_learning_rate": 4.886626482429933e-06,
+      "loss": 0.7994,
+      "output_text_mask_loss": 0.0025,
+      "predictor_learning_rate": 0.0004886626482429932,
+      "step": 1204
+    },
+    {
+      "epoch": 0.12,
+      "gumbel_tau": 0.7527,
+      "image_masks_loss": 0.0134,
+      "llm_learning_rate": 4.886379350497978e-06,
+      "loss": 0.7385,
+      "output_text_mask_loss": 0.0397,
+      "predictor_learning_rate": 0.0004886379350497977,
+      "step": 1205
+    },
+    {
+      "epoch": 0.12,
+      "gumbel_tau": 0.7526,
+      "image_masks_loss": 0.0123,
+      "llm_learning_rate": 4.8861319557726115e-06,
+      "loss": 0.8102,
+      "output_text_mask_loss": 0.0246,
+      "predictor_learning_rate": 0.0004886131955772611,
+      "step": 1206
+    },
+    {
+      "epoch": 0.12,
+      "gumbel_tau": 0.7524,
+      "image_masks_loss": 0.0158,
+      "llm_learning_rate": 4.88588429828108e-06,
+      "loss": 0.8492,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.000488588429828108,
+      "step": 1207
+    },
+    {
+      "epoch": 0.12,
+      "gumbel_tau": 0.7522,
+      "image_masks_loss": 0.0305,
+      "llm_learning_rate": 4.885636378050656e-06,
+      "loss": 0.7765,
+      "output_text_mask_loss": 0.0633,
+      "predictor_learning_rate": 0.0004885636378050655,
+      "step": 1208
+    },
+    {
+      "epoch": 0.12,
+      "gumbel_tau": 0.752,
+      "image_masks_loss": 0.029,
+      "llm_learning_rate": 4.885388195108639e-06,
+      "loss": 0.84,
+      "output_text_mask_loss": 0.0473,
+      "predictor_learning_rate": 0.0004885388195108639,
+      "step": 1209
+    },
+    {
+      "epoch": 0.12,
+      "gumbel_tau": 0.7518,
+      "image_masks_loss": 0.0122,
+      "llm_learning_rate": 4.885139749482361e-06,
+      "loss": 0.7903,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.0004885139749482361,
+      "step": 1210
+    },
+    {
+      "epoch": 0.12,
+      "gumbel_tau": 0.7517,
+      "image_masks_loss": 0.0359,
+      "llm_learning_rate": 4.884891041199181e-06,
+      "loss": 0.8944,
+      "output_text_mask_loss": 0.0095,
+      "predictor_learning_rate": 0.0004884891041199181,
+      "step": 1211
+    },
+    {
+      "epoch": 0.12,
+      "gumbel_tau": 0.7515,
+      "image_masks_loss": 0.0175,
+      "llm_learning_rate": 4.884642070286489e-06,
+      "loss": 0.8159,
+      "output_text_mask_loss": 0.0549,
+      "predictor_learning_rate": 0.0004884642070286488,
+      "step": 1212
+    },
+    {
+      "epoch": 0.12,
+      "gumbel_tau": 0.7513,
+      "image_masks_loss": 0.011,
+      "llm_learning_rate": 4.884392836771699e-06,
+      "loss": 0.8286,
+      "output_text_mask_loss": 0.0072,
+      "predictor_learning_rate": 0.0004884392836771699,
+      "step": 1213
+    },
+    {
+      "epoch": 0.12,
+      "gumbel_tau": 0.7511,
+      "image_masks_loss": 0.0429,
+      "llm_learning_rate": 4.884143340682259e-06,
+      "loss": 0.7962,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.0004884143340682259,
+      "step": 1214
+    },
+    {
+      "epoch": 0.12,
+      "gumbel_tau": 0.751,
+      "image_masks_loss": 0.0166,
+      "llm_learning_rate": 4.883893582045644e-06,
+      "loss": 0.7771,
+      "output_text_mask_loss": 0.0422,
+      "predictor_learning_rate": 0.0004883893582045644,
+      "step": 1215
+    },
+    {
+      "epoch": 0.12,
+      "gumbel_tau": 0.7508,
+      "image_masks_loss": 0.0408,
+      "llm_learning_rate": 4.883643560889358e-06,
+      "loss": 0.8217,
+      "output_text_mask_loss": 0.0221,
+      "predictor_learning_rate": 0.0004883643560889357,
+      "step": 1216
+    },
+    {
+      "epoch": 0.12,
+      "gumbel_tau": 0.7506,
+      "image_masks_loss": 0.0203,
+      "llm_learning_rate": 4.883393277240933e-06,
+      "loss": 0.6504,
+      "output_text_mask_loss": 0.1076,
+      "predictor_learning_rate": 0.0004883393277240933,
+      "step": 1217
+    },
+    {
+      "epoch": 0.12,
+      "gumbel_tau": 0.7504,
+      "image_masks_loss": 0.0284,
+      "llm_learning_rate": 4.8831427311279325e-06,
+      "loss": 0.8591,
+      "output_text_mask_loss": 0.0319,
+      "predictor_learning_rate": 0.0004883142731127932,
+      "step": 1218
+    },
+    {
+      "epoch": 0.12,
+      "gumbel_tau": 0.7502,
+      "image_masks_loss": 0.0313,
+      "llm_learning_rate": 4.882891922577946e-06,
+      "loss": 0.8293,
+      "output_text_mask_loss": 0.0257,
+      "predictor_learning_rate": 0.0004882891922577946,
+      "step": 1219
+    },
+    {
+      "epoch": 0.12,
+      "gumbel_tau": 0.7501,
+      "image_masks_loss": 0.0202,
+      "llm_learning_rate": 4.882640851618593e-06,
+      "loss": 0.6726,
+      "output_text_mask_loss": 0.0185,
+      "predictor_learning_rate": 0.00048826408516185925,
+      "step": 1220
+    },
+    {
+      "epoch": 0.13,
+      "gumbel_tau": 0.7499,
+      "image_masks_loss": 0.0101,
+      "llm_learning_rate": 4.882389518277523e-06,
+      "loss": 0.7991,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.00048823895182775225,
+      "step": 1221
+    },
+    {
+      "epoch": 0.13,
+      "gumbel_tau": 0.7497,
+      "image_masks_loss": 0.0109,
+      "llm_learning_rate": 4.882137922582412e-06,
+      "loss": 0.805,
+      "output_text_mask_loss": 0.0542,
+      "predictor_learning_rate": 0.00048821379225824114,
+      "step": 1222
+    },
+    {
+      "epoch": 0.13,
+      "gumbel_tau": 0.7495,
+      "image_masks_loss": 0.0206,
+      "llm_learning_rate": 4.881886064560967e-06,
+      "loss": 0.7342,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.0004881886064560966,
+      "step": 1223
+    },
+    {
+      "epoch": 0.13,
+      "gumbel_tau": 0.7494,
+      "image_masks_loss": 0.0345,
+      "llm_learning_rate": 4.881633944240923e-06,
+      "loss": 0.6771,
+      "output_text_mask_loss": 0.0523,
+      "predictor_learning_rate": 0.00048816339442409225,
+      "step": 1224
+    },
+    {
+      "epoch": 0.13,
+      "gumbel_tau": 0.7492,
+      "image_masks_loss": 0.058,
+      "llm_learning_rate": 4.8813815616500435e-06,
+      "loss": 0.7802,
+      "output_text_mask_loss": 0.0054,
+      "predictor_learning_rate": 0.0004881381561650043,
+      "step": 1225
+    },
+    {
+      "epoch": 0.13,
+      "gumbel_tau": 0.749,
+      "image_masks_loss": 0.0132,
+      "llm_learning_rate": 4.881128916816124e-06,
+      "loss": 0.7403,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.0004881128916816123,
+      "step": 1226
+    },
+    {
+      "epoch": 0.13,
+      "gumbel_tau": 0.7488,
+      "image_masks_loss": 0.0265,
+      "llm_learning_rate": 4.880876009766983e-06,
+      "loss": 0.7197,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00048808760097669826,
+      "step": 1227
+    },
+    {
+      "epoch": 0.13,
+      "gumbel_tau": 0.7487,
+      "image_masks_loss": 0.0233,
+      "llm_learning_rate": 4.880622840530473e-06,
+      "loss": 0.7853,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00048806228405304724,
+      "step": 1228
+    },
+    {
+      "epoch": 0.13,
+      "gumbel_tau": 0.7485,
+      "image_masks_loss": 0.0122,
+      "llm_learning_rate": 4.8803694091344735e-06,
+      "loss": 0.7867,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 0.0004880369409134473,
+      "step": 1229
+    },
+    {
+      "epoch": 0.13,
+      "gumbel_tau": 0.7483,
+      "image_masks_loss": 0.0079,
+      "llm_learning_rate": 4.880115715606892e-06,
+      "loss": 0.8314,
+      "output_text_mask_loss": 0.004,
+      "predictor_learning_rate": 0.0004880115715606892,
+      "step": 1230
+    },
+    {
+      "epoch": 0.13,
+      "gumbel_tau": 0.7481,
+      "image_masks_loss": 0.0263,
+      "llm_learning_rate": 4.879861759975666e-06,
+      "loss": 0.8009,
+      "output_text_mask_loss": 0.0124,
+      "predictor_learning_rate": 0.0004879861759975666,
+      "step": 1231
+    },
+    {
+      "epoch": 0.13,
+      "gumbel_tau": 0.748,
+      "image_masks_loss": 0.0338,
+      "llm_learning_rate": 4.879607542268763e-06,
+      "loss": 0.8667,
+      "output_text_mask_loss": 0.1099,
+      "predictor_learning_rate": 0.0004879607542268763,
+      "step": 1232
+    },
+    {
+      "epoch": 0.13,
+      "gumbel_tau": 0.7478,
+      "image_masks_loss": 0.0311,
+      "llm_learning_rate": 4.879353062514177e-06,
+      "loss": 0.7249,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00048793530625141767,
+      "step": 1233
+    },
+    {
+      "epoch": 0.13,
+      "gumbel_tau": 0.7476,
+      "image_masks_loss": 0.0154,
+      "llm_learning_rate": 4.879098320739933e-06,
+      "loss": 0.7802,
+      "output_text_mask_loss": 0.0095,
+      "predictor_learning_rate": 0.0004879098320739932,
+      "step": 1234
+    },
+    {
+      "epoch": 0.13,
+      "gumbel_tau": 0.7474,
+      "image_masks_loss": 0.009,
+      "llm_learning_rate": 4.87884331697408e-06,
+      "loss": 0.8276,
+      "output_text_mask_loss": 0.0116,
+      "predictor_learning_rate": 0.000487884331697408,
+      "step": 1235
+    },
+    {
+      "epoch": 0.13,
+      "gumbel_tau": 0.7472,
+      "image_masks_loss": 0.0103,
+      "llm_learning_rate": 4.878588051244704e-06,
+      "loss": 0.7336,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004878588051244703,
+      "step": 1236
+    },
+    {
+      "epoch": 0.13,
+      "gumbel_tau": 0.7471,
+      "image_masks_loss": 0.0463,
+      "llm_learning_rate": 4.878332523579913e-06,
+      "loss": 0.8903,
+      "output_text_mask_loss": 0.0401,
+      "predictor_learning_rate": 0.0004878332523579913,
+      "step": 1237
+    },
+    {
+      "epoch": 0.13,
+      "gumbel_tau": 0.7469,
+      "image_masks_loss": 0.024,
+      "llm_learning_rate": 4.878076734007848e-06,
+      "loss": 0.7405,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00048780767340078483,
+      "step": 1238
+    },
+    {
+      "epoch": 0.13,
+      "gumbel_tau": 0.7467,
+      "image_masks_loss": 0.0093,
+      "llm_learning_rate": 4.8778206825566755e-06,
+      "loss": 0.7342,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00048778206825566753,
+      "step": 1239
+    },
+    {
+      "epoch": 0.13,
+      "gumbel_tau": 0.7465,
+      "image_masks_loss": 0.0412,
+      "llm_learning_rate": 4.877564369254593e-06,
+      "loss": 0.7767,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00048775643692545924,
+      "step": 1240
+    },
+    {
+      "epoch": 0.13,
+      "gumbel_tau": 0.7464,
+      "image_masks_loss": 0.0248,
+      "llm_learning_rate": 4.877307794129826e-06,
+      "loss": 0.8274,
+      "output_text_mask_loss": 0.0126,
+      "predictor_learning_rate": 0.0004877307794129826,
+      "step": 1241
+    },
+    {
+      "epoch": 0.13,
+      "gumbel_tau": 0.7462,
+      "image_masks_loss": 0.0208,
+      "llm_learning_rate": 4.87705095721063e-06,
+      "loss": 0.8502,
+      "output_text_mask_loss": 0.0633,
+      "predictor_learning_rate": 0.000487705095721063,
+      "step": 1242
+    },
+    {
+      "epoch": 0.13,
+      "gumbel_tau": 0.746,
+      "image_masks_loss": 0.0227,
+      "llm_learning_rate": 4.876793858525288e-06,
+      "loss": 0.9114,
+      "output_text_mask_loss": 0.177,
+      "predictor_learning_rate": 0.0004876793858525288,
+      "step": 1243
+    },
+    {
+      "epoch": 0.13,
+      "gumbel_tau": 0.7458,
+      "image_masks_loss": 0.0137,
+      "llm_learning_rate": 4.876536498102112e-06,
+      "loss": 0.6677,
+      "output_text_mask_loss": 0.0385,
+      "predictor_learning_rate": 0.0004876536498102112,
+      "step": 1244
+    },
+    {
+      "epoch": 0.13,
+      "gumbel_tau": 0.7457,
+      "image_masks_loss": 0.0389,
+      "llm_learning_rate": 4.8762788759694435e-06,
+      "loss": 0.7238,
+      "output_text_mask_loss": NaN,
+      "predictor_learning_rate": 0.0004876278875969443,
+      "step": 1245
+    },
+    {
+      "epoch": 0.13,
+      "gumbel_tau": 0.7455,
+      "image_masks_loss": 0.0209,
+      "llm_learning_rate": 4.876020992155652e-06,
+      "loss": 0.7848,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.00048760209921556514,
+      "step": 1246
+    },
+    {
+      "epoch": 0.13,
+      "gumbel_tau": 0.7453,
+      "image_masks_loss": 0.0259,
+      "llm_learning_rate": 4.875762846689136e-06,
+      "loss": 0.785,
+      "output_text_mask_loss": 0.0507,
+      "predictor_learning_rate": 0.0004875762846689136,
+      "step": 1247
+    },
+    {
+      "epoch": 0.13,
+      "gumbel_tau": 0.7451,
+      "image_masks_loss": 0.0136,
+      "llm_learning_rate": 4.8755044395983244e-06,
+      "loss": 0.6388,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004875504439598324,
+      "step": 1248
+    },
+    {
+      "epoch": 0.13,
+      "gumbel_tau": 0.745,
+      "image_masks_loss": 0.0401,
+      "llm_learning_rate": 4.875245770911672e-06,
+      "loss": 0.8005,
+      "output_text_mask_loss": 0.006,
+      "predictor_learning_rate": 0.00048752457709116715,
+      "step": 1249
+    },
+    {
+      "epoch": 0.13,
+      "gumbel_tau": 0.7448,
+      "image_masks_loss": 0.0133,
+      "llm_learning_rate": 4.874986840657665e-06,
+      "loss": 0.7112,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004874986840657665,
+      "step": 1250
+    },
+    {
+      "epoch": 0.13,
+      "gumbel_tau": 0.7446,
+      "image_masks_loss": 0.0387,
+      "llm_learning_rate": 4.874727648864817e-06,
+      "loss": 0.7878,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004874727648864816,
+      "step": 1251
+    },
+    {
+      "epoch": 0.13,
+      "gumbel_tau": 0.7444,
+      "image_masks_loss": 0.0082,
+      "llm_learning_rate": 4.8744681955616705e-06,
+      "loss": 0.7706,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.000487446819556167,
+      "step": 1252
+    },
+    {
+      "epoch": 0.13,
+      "gumbel_tau": 0.7443,
+      "image_masks_loss": 0.0334,
+      "llm_learning_rate": 4.874208480776798e-06,
+      "loss": 0.7893,
+      "output_text_mask_loss": 0.0277,
+      "predictor_learning_rate": 0.0004874208480776797,
+      "step": 1253
+    },
+    {
+      "epoch": 0.13,
+      "gumbel_tau": 0.7441,
+      "image_masks_loss": 0.0349,
+      "llm_learning_rate": 4.873948504538798e-06,
+      "loss": 0.8407,
+      "output_text_mask_loss": 0.0113,
+      "predictor_learning_rate": 0.00048739485045387983,
+      "step": 1254
+    },
+    {
+      "epoch": 0.13,
+      "gumbel_tau": 0.7439,
+      "image_masks_loss": 0.0252,
+      "llm_learning_rate": 4.873688266876302e-06,
+      "loss": 0.7886,
+      "output_text_mask_loss": 0.0725,
+      "predictor_learning_rate": 0.00048736882668763015,
+      "step": 1255
+    },
+    {
+      "epoch": 0.13,
+      "gumbel_tau": 0.7437,
+      "image_masks_loss": 0.0399,
+      "llm_learning_rate": 4.873427767817967e-06,
+      "loss": 0.788,
+      "output_text_mask_loss": 0.0124,
+      "predictor_learning_rate": 0.0004873427767817967,
+      "step": 1256
+    },
+    {
+      "epoch": 0.13,
+      "gumbel_tau": 0.7436,
+      "image_masks_loss": 0.0061,
+      "llm_learning_rate": 4.873167007392479e-06,
+      "loss": 0.7867,
+      "output_text_mask_loss": 0.0408,
+      "predictor_learning_rate": 0.0004873167007392479,
+      "step": 1257
+    },
+    {
+      "epoch": 0.13,
+      "gumbel_tau": 0.7434,
+      "image_masks_loss": 0.0137,
+      "llm_learning_rate": 4.872905985628555e-06,
+      "loss": 0.8016,
+      "output_text_mask_loss": 0.0429,
+      "predictor_learning_rate": 0.0004872905985628554,
+      "step": 1258
+    },
+    {
+      "epoch": 0.13,
+      "gumbel_tau": 0.7432,
+      "image_masks_loss": 0.0186,
+      "llm_learning_rate": 4.872644702554937e-06,
+      "loss": 0.8208,
+      "output_text_mask_loss": 0.0771,
+      "predictor_learning_rate": 0.0004872644702554937,
+      "step": 1259
+    },
+    {
+      "epoch": 0.13,
+      "gumbel_tau": 0.743,
+      "image_masks_loss": 0.0483,
+      "llm_learning_rate": 4.8723831582004015e-06,
+      "loss": 0.814,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00048723831582004007,
+      "step": 1260
+    },
+    {
+      "epoch": 0.13,
+      "gumbel_tau": 0.7429,
+      "image_masks_loss": 0.0174,
+      "llm_learning_rate": 4.872121352593746e-06,
+      "loss": 0.8598,
+      "output_text_mask_loss": 0.0117,
+      "predictor_learning_rate": 0.00048721213525937455,
+      "step": 1261
+    },
+    {
+      "epoch": 0.13,
+      "gumbel_tau": 0.7427,
+      "image_masks_loss": 0.0324,
+      "llm_learning_rate": 4.871859285763805e-06,
+      "loss": 0.7568,
+      "output_text_mask_loss": 0.0446,
+      "predictor_learning_rate": 0.00048718592857638045,
+      "step": 1262
+    },
+    {
+      "epoch": 0.13,
+      "gumbel_tau": 0.7425,
+      "image_masks_loss": 0.025,
+      "llm_learning_rate": 4.871596957739436e-06,
+      "loss": 0.8337,
+      "output_text_mask_loss": 0.09,
+      "predictor_learning_rate": 0.0004871596957739436,
+      "step": 1263
+    },
+    {
+      "epoch": 0.13,
+      "gumbel_tau": 0.7423,
+      "image_masks_loss": 0.0259,
+      "llm_learning_rate": 4.8713343685495274e-06,
+      "loss": 0.8093,
+      "output_text_mask_loss": 0.0629,
+      "predictor_learning_rate": 0.0004871334368549527,
+      "step": 1264
+    },
+    {
+      "epoch": 0.13,
+      "gumbel_tau": 0.7422,
+      "image_masks_loss": 0.0196,
+      "llm_learning_rate": 4.871071518222996e-06,
+      "loss": 0.7271,
+      "output_text_mask_loss": 0.0056,
+      "predictor_learning_rate": 0.0004871071518222996,
+      "step": 1265
+    },
+    {
+      "epoch": 0.13,
+      "gumbel_tau": 0.742,
+      "image_masks_loss": 0.0221,
+      "llm_learning_rate": 4.870808406788789e-06,
+      "loss": 0.7658,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.0004870808406788788,
+      "step": 1266
+    },
+    {
+      "epoch": 0.13,
+      "gumbel_tau": 0.7418,
+      "image_masks_loss": 0.0234,
+      "llm_learning_rate": 4.870545034275878e-06,
+      "loss": 0.8567,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 0.0004870545034275877,
+      "step": 1267
+    },
+    {
+      "epoch": 0.13,
+      "gumbel_tau": 0.7416,
+      "image_masks_loss": 0.0322,
+      "llm_learning_rate": 4.870281400713268e-06,
+      "loss": 0.8057,
+      "output_text_mask_loss": 0.0376,
+      "predictor_learning_rate": 0.00048702814007132674,
+      "step": 1268
+    },
+    {
+      "epoch": 0.13,
+      "gumbel_tau": 0.7415,
+      "image_masks_loss": 0.0203,
+      "llm_learning_rate": 4.8700175061299905e-06,
+      "loss": 0.7415,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.000487001750612999,
+      "step": 1269
+    },
+    {
+      "epoch": 0.13,
+      "gumbel_tau": 0.7413,
+      "image_masks_loss": 0.0299,
+      "llm_learning_rate": 4.8697533505551055e-06,
+      "loss": 0.8222,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.0004869753350555106,
+      "step": 1270
+    },
+    {
+      "epoch": 0.13,
+      "gumbel_tau": 0.7411,
+      "image_masks_loss": 0.0217,
+      "llm_learning_rate": 4.869488934017704e-06,
+      "loss": 0.6937,
+      "output_text_mask_loss": 0.0362,
+      "predictor_learning_rate": 0.0004869488934017704,
+      "step": 1271
+    },
+    {
+      "epoch": 0.13,
+      "gumbel_tau": 0.7409,
+      "image_masks_loss": 0.0084,
+      "llm_learning_rate": 4.8692242565469036e-06,
+      "loss": 0.7244,
+      "output_text_mask_loss": 0.0422,
+      "predictor_learning_rate": 0.00048692242565469035,
+      "step": 1272
+    },
+    {
+      "epoch": 0.13,
+      "gumbel_tau": 0.7408,
+      "image_masks_loss": 0.019,
+      "llm_learning_rate": 4.868959318171851e-06,
+      "loss": 0.8755,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.000486895931817185,
+      "step": 1273
+    },
+    {
+      "epoch": 0.13,
+      "gumbel_tau": 0.7406,
+      "image_masks_loss": 0.033,
+      "llm_learning_rate": 4.868694118921721e-06,
+      "loss": 0.8272,
+      "output_text_mask_loss": 0.0098,
+      "predictor_learning_rate": 0.000486869411892172,
+      "step": 1274
+    },
+    {
+      "epoch": 0.13,
+      "gumbel_tau": 0.7404,
+      "image_masks_loss": 0.0431,
+      "llm_learning_rate": 4.868428658825718e-06,
+      "loss": 0.7726,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 0.0004868428658825718,
+      "step": 1275
+    },
+    {
+      "epoch": 0.13,
+      "gumbel_tau": 0.7402,
+      "image_masks_loss": 0.0546,
+      "llm_learning_rate": 4.868162937913076e-06,
+      "loss": 0.8049,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.0004868162937913076,
+      "step": 1276
+    },
+    {
+      "epoch": 0.13,
+      "gumbel_tau": 0.7401,
+      "image_masks_loss": 0.0179,
+      "llm_learning_rate": 4.867896956213057e-06,
+      "loss": 0.7677,
+      "output_text_mask_loss": 0.0288,
+      "predictor_learning_rate": 0.0004867896956213056,
+      "step": 1277
+    },
+    {
+      "epoch": 0.13,
+      "gumbel_tau": 0.7399,
+      "image_masks_loss": 0.0538,
+      "llm_learning_rate": 4.867630713754951e-06,
+      "loss": 0.7834,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.000486763071375495,
+      "step": 1278
+    },
+    {
+      "epoch": 0.13,
+      "gumbel_tau": 0.7397,
+      "image_masks_loss": 0.0357,
+      "llm_learning_rate": 4.867364210568076e-06,
+      "loss": 0.6868,
+      "output_text_mask_loss": 0.0261,
+      "predictor_learning_rate": 0.0004867364210568076,
+      "step": 1279
+    },
+    {
+      "epoch": 0.13,
+      "gumbel_tau": 0.7395,
+      "image_masks_loss": 0.0104,
+      "llm_learning_rate": 4.867097446681782e-06,
+      "loss": 0.8013,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.00048670974466817813,
+      "step": 1280
+    },
+    {
+      "epoch": 0.13,
+      "gumbel_tau": 0.7394,
+      "image_masks_loss": 0.0334,
+      "llm_learning_rate": 4.866830422125444e-06,
+      "loss": 0.796,
+      "output_text_mask_loss": 0.0137,
+      "predictor_learning_rate": 0.00048668304221254435,
+      "step": 1281
+    },
+    {
+      "epoch": 0.13,
+      "gumbel_tau": 0.7392,
+      "image_masks_loss": 0.069,
+      "llm_learning_rate": 4.866563136928469e-06,
+      "loss": 0.7997,
+      "output_text_mask_loss": 0.0437,
+      "predictor_learning_rate": 0.0004866563136928468,
+      "step": 1282
+    },
+    {
+      "epoch": 0.13,
+      "gumbel_tau": 0.739,
+      "image_masks_loss": 0.0083,
+      "llm_learning_rate": 4.866295591120289e-06,
+      "loss": 0.7949,
+      "output_text_mask_loss": 0.0019,
+      "predictor_learning_rate": 0.00048662955911202884,
+      "step": 1283
+    },
+    {
+      "epoch": 0.13,
+      "gumbel_tau": 0.7388,
+      "image_masks_loss": 0.0278,
+      "llm_learning_rate": 4.866027784730367e-06,
+      "loss": 0.7366,
+      "output_text_mask_loss": 0.0275,
+      "predictor_learning_rate": 0.0004866027784730367,
+      "step": 1284
+    },
+    {
+      "epoch": 0.13,
+      "gumbel_tau": 0.7387,
+      "image_masks_loss": 0.0271,
+      "llm_learning_rate": 4.8657597177881976e-06,
+      "loss": 0.7096,
+      "output_text_mask_loss": 0.0565,
+      "predictor_learning_rate": 0.00048657597177881974,
+      "step": 1285
+    },
+    {
+      "epoch": 0.13,
+      "gumbel_tau": 0.7385,
+      "image_masks_loss": 0.0263,
+      "llm_learning_rate": 4.865491390323297e-06,
+      "loss": 0.8536,
+      "output_text_mask_loss": 0.0216,
+      "predictor_learning_rate": 0.00048654913903232967,
+      "step": 1286
+    },
+    {
+      "epoch": 0.13,
+      "gumbel_tau": 0.7383,
+      "image_masks_loss": 0.0173,
+      "llm_learning_rate": 4.865222802365216e-06,
+      "loss": 0.734,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00048652228023652157,
+      "step": 1287
+    },
+    {
+      "epoch": 0.13,
+      "gumbel_tau": 0.7381,
+      "image_masks_loss": 0.033,
+      "llm_learning_rate": 4.864953953943532e-06,
+      "loss": 0.736,
+      "output_text_mask_loss": 0.0261,
+      "predictor_learning_rate": 0.00048649539539435315,
+      "step": 1288
+    },
+    {
+      "epoch": 0.13,
+      "gumbel_tau": 0.738,
+      "image_masks_loss": 0.0278,
+      "llm_learning_rate": 4.86468484508785e-06,
+      "loss": 0.7455,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.000486468484508785,
+      "step": 1289
+    },
+    {
+      "epoch": 0.13,
+      "gumbel_tau": 0.7378,
+      "image_masks_loss": 0.0104,
+      "llm_learning_rate": 4.864415475827806e-06,
+      "loss": 0.754,
+      "output_text_mask_loss": 0.0078,
+      "predictor_learning_rate": 0.0004864415475827806,
+      "step": 1290
+    },
+    {
+      "epoch": 0.13,
+      "gumbel_tau": 0.7376,
+      "image_masks_loss": 0.0744,
+      "llm_learning_rate": 4.864145846193064e-06,
+      "loss": 0.8673,
+      "output_text_mask_loss": 0.0923,
+      "predictor_learning_rate": 0.0004864145846193063,
+      "step": 1291
+    },
+    {
+      "epoch": 0.13,
+      "gumbel_tau": 0.7374,
+      "image_masks_loss": 0.0132,
+      "llm_learning_rate": 4.863875956213315e-06,
+      "loss": 0.7577,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00048638759562133146,
+      "step": 1292
+    },
+    {
+      "epoch": 0.13,
+      "gumbel_tau": 0.7373,
+      "image_masks_loss": 0.0235,
+      "llm_learning_rate": 4.86360580591828e-06,
+      "loss": 0.8342,
+      "output_text_mask_loss": 0.0231,
+      "predictor_learning_rate": 0.00048636058059182795,
+      "step": 1293
+    },
+    {
+      "epoch": 0.13,
+      "gumbel_tau": 0.7371,
+      "image_masks_loss": 0.0208,
+      "llm_learning_rate": 4.86333539533771e-06,
+      "loss": 0.8206,
+      "output_text_mask_loss": 0.0138,
+      "predictor_learning_rate": 0.0004863335395337709,
+      "step": 1294
+    },
+    {
+      "epoch": 0.13,
+      "gumbel_tau": 0.7369,
+      "image_masks_loss": 0.0135,
+      "llm_learning_rate": 4.86306472450138e-06,
+      "loss": 0.7937,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.00048630647245013803,
+      "step": 1295
+    },
+    {
+      "epoch": 0.13,
+      "gumbel_tau": 0.7367,
+      "image_masks_loss": 0.0119,
+      "llm_learning_rate": 4.862793793439101e-06,
+      "loss": 0.7982,
+      "output_text_mask_loss": 0.0317,
+      "predictor_learning_rate": 0.00048627937934391007,
+      "step": 1296
+    },
+    {
+      "epoch": 0.13,
+      "gumbel_tau": 0.7366,
+      "image_masks_loss": 0.0275,
+      "llm_learning_rate": 4.862522602180706e-06,
+      "loss": 0.8744,
+      "output_text_mask_loss": 0.0193,
+      "predictor_learning_rate": 0.00048625226021807057,
+      "step": 1297
+    },
+    {
+      "epoch": 0.13,
+      "gumbel_tau": 0.7364,
+      "image_masks_loss": 0.021,
+      "llm_learning_rate": 4.86225115075606e-06,
+      "loss": 0.8694,
+      "output_text_mask_loss": 0.0124,
+      "predictor_learning_rate": 0.00048622511507560596,
+      "step": 1298
+    },
+    {
+      "epoch": 0.13,
+      "gumbel_tau": 0.7362,
+      "image_masks_loss": 0.0053,
+      "llm_learning_rate": 4.861979439195055e-06,
+      "loss": 0.7358,
+      "output_text_mask_loss": 0.0056,
+      "predictor_learning_rate": 0.00048619794391950545,
+      "step": 1299
+    },
+    {
+      "epoch": 0.13,
+      "gumbel_tau": 0.736,
+      "image_masks_loss": 0.0149,
+      "llm_learning_rate": 4.861707467527613e-06,
+      "loss": 0.7334,
+      "output_text_mask_loss": 0.0381,
+      "predictor_learning_rate": 0.0004861707467527613,
+      "step": 1300
+    },
+    {
+      "epoch": 0.13,
+      "gumbel_tau": 0.7359,
+      "image_masks_loss": 0.018,
+      "llm_learning_rate": 4.861435235783685e-06,
+      "loss": 0.7184,
+      "output_text_mask_loss": 0.005,
+      "predictor_learning_rate": 0.00048614352357836846,
+      "step": 1301
+    },
+    {
+      "epoch": 0.13,
+      "gumbel_tau": 0.7357,
+      "image_masks_loss": 0.009,
+      "llm_learning_rate": 4.861162743993249e-06,
+      "loss": 0.7492,
+      "output_text_mask_loss": 0.0193,
+      "predictor_learning_rate": 0.0004861162743993248,
+      "step": 1302
+    },
+    {
+      "epoch": 0.13,
+      "gumbel_tau": 0.7355,
+      "image_masks_loss": 0.0195,
+      "llm_learning_rate": 4.860889992186311e-06,
+      "loss": 0.8128,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 0.0004860889992186311,
+      "step": 1303
+    },
+    {
+      "epoch": 0.13,
+      "gumbel_tau": 0.7354,
+      "image_masks_loss": 0.0385,
+      "llm_learning_rate": 4.86061698039291e-06,
+      "loss": 0.8039,
+      "output_text_mask_loss": 0.0073,
+      "predictor_learning_rate": 0.0004860616980392909,
+      "step": 1304
+    },
+    {
+      "epoch": 0.13,
+      "gumbel_tau": 0.7352,
+      "image_masks_loss": 0.0137,
+      "llm_learning_rate": 4.860343708643107e-06,
+      "loss": 0.7897,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.00048603437086431066,
+      "step": 1305
+    },
+    {
+      "epoch": 0.13,
+      "gumbel_tau": 0.735,
+      "image_masks_loss": 0.0292,
+      "llm_learning_rate": 4.860070176966999e-06,
+      "loss": 0.6783,
+      "output_text_mask_loss": 0.0248,
+      "predictor_learning_rate": 0.00048600701769669984,
+      "step": 1306
+    },
+    {
+      "epoch": 0.13,
+      "gumbel_tau": 0.7348,
+      "image_masks_loss": 0.0178,
+      "llm_learning_rate": 4.859796385394706e-06,
+      "loss": 0.7822,
+      "output_text_mask_loss": 0.0061,
+      "predictor_learning_rate": 0.0004859796385394706,
+      "step": 1307
+    },
+    {
+      "epoch": 0.13,
+      "gumbel_tau": 0.7347,
+      "image_masks_loss": 0.0355,
+      "llm_learning_rate": 4.859522333956378e-06,
+      "loss": 0.8326,
+      "output_text_mask_loss": 0.0039,
+      "predictor_learning_rate": 0.00048595223339563783,
+      "step": 1308
+    },
+    {
+      "epoch": 0.13,
+      "gumbel_tau": 0.7345,
+      "image_masks_loss": 0.0144,
+      "llm_learning_rate": 4.859248022682196e-06,
+      "loss": 0.7612,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.00048592480226821957,
+      "step": 1309
+    },
+    {
+      "epoch": 0.13,
+      "gumbel_tau": 0.7343,
+      "image_masks_loss": 0.0178,
+      "llm_learning_rate": 4.8589734516023666e-06,
+      "loss": 0.7815,
+      "output_text_mask_loss": 0.0671,
+      "predictor_learning_rate": 0.00048589734516023666,
+      "step": 1310
+    },
+    {
+      "epoch": 0.13,
+      "gumbel_tau": 0.7341,
+      "image_masks_loss": 0.033,
+      "llm_learning_rate": 4.858698620747126e-06,
+      "loss": 0.8153,
+      "output_text_mask_loss": 0.1572,
+      "predictor_learning_rate": 0.00048586986207471255,
+      "step": 1311
+    },
+    {
+      "epoch": 0.13,
+      "gumbel_tau": 0.734,
+      "image_masks_loss": 0.0232,
+      "llm_learning_rate": 4.8584235301467395e-06,
+      "loss": 0.7906,
+      "output_text_mask_loss": 0.0124,
+      "predictor_learning_rate": 0.0004858423530146739,
+      "step": 1312
+    },
+    {
+      "epoch": 0.13,
+      "gumbel_tau": 0.7338,
+      "image_masks_loss": 0.012,
+      "llm_learning_rate": 4.8581481798315006e-06,
+      "loss": 0.7272,
+      "output_text_mask_loss": 0.0032,
+      "predictor_learning_rate": 0.00048581481798315,
+      "step": 1313
+    },
+    {
+      "epoch": 0.13,
+      "gumbel_tau": 0.7336,
+      "image_masks_loss": 0.017,
+      "llm_learning_rate": 4.857872569831732e-06,
+      "loss": 0.8174,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004857872569831732,
+      "step": 1314
+    },
+    {
+      "epoch": 0.13,
+      "gumbel_tau": 0.7334,
+      "image_masks_loss": 0.0408,
+      "llm_learning_rate": 4.857596700177784e-06,
+      "loss": 0.7755,
+      "output_text_mask_loss": 0.0223,
+      "predictor_learning_rate": 0.00048575967001777834,
+      "step": 1315
+    },
+    {
+      "epoch": 0.13,
+      "gumbel_tau": 0.7333,
+      "image_masks_loss": 0.0144,
+      "llm_learning_rate": 4.857320570900035e-06,
+      "loss": 0.7042,
+      "output_text_mask_loss": 0.0193,
+      "predictor_learning_rate": 0.00048573205709000346,
+      "step": 1316
+    },
+    {
+      "epoch": 0.13,
+      "gumbel_tau": 0.7331,
+      "image_masks_loss": 0.01,
+      "llm_learning_rate": 4.857044182028895e-06,
+      "loss": 0.7014,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00048570441820288947,
+      "step": 1317
+    },
+    {
+      "epoch": 0.14,
+      "gumbel_tau": 0.7329,
+      "image_masks_loss": 0.0269,
+      "llm_learning_rate": 4.8567675335948e-06,
+      "loss": 0.7733,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 0.00048567675335947993,
+      "step": 1318
+    },
+    {
+      "epoch": 0.14,
+      "gumbel_tau": 0.7328,
+      "image_masks_loss": 0.0233,
+      "llm_learning_rate": 4.856490625628213e-06,
+      "loss": 0.7625,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00048564906256282124,
+      "step": 1319
+    },
+    {
+      "epoch": 0.14,
+      "gumbel_tau": 0.7326,
+      "image_masks_loss": 0.0175,
+      "llm_learning_rate": 4.856213458159631e-06,
+      "loss": 0.7762,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00048562134581596304,
+      "step": 1320
+    },
+    {
+      "epoch": 0.14,
+      "gumbel_tau": 0.7324,
+      "image_masks_loss": 0.0584,
+      "llm_learning_rate": 4.855936031219573e-06,
+      "loss": 0.8456,
+      "output_text_mask_loss": 0.0307,
+      "predictor_learning_rate": 0.0004855936031219573,
+      "step": 1321
+    },
+    {
+      "epoch": 0.14,
+      "gumbel_tau": 0.7322,
+      "image_masks_loss": 0.0102,
+      "llm_learning_rate": 4.8556583448385936e-06,
+      "loss": 0.7863,
+      "output_text_mask_loss": 0.0397,
+      "predictor_learning_rate": 0.0004855658344838593,
+      "step": 1322
+    },
+    {
+      "epoch": 0.14,
+      "gumbel_tau": 0.7321,
+      "image_masks_loss": 0.0046,
+      "llm_learning_rate": 4.855380399047269e-06,
+      "loss": 0.688,
+      "output_text_mask_loss": 0.0347,
+      "predictor_learning_rate": 0.0004855380399047269,
+      "step": 1323
+    },
+    {
+      "epoch": 0.14,
+      "gumbel_tau": 0.7319,
+      "image_masks_loss": 0.0332,
+      "llm_learning_rate": 4.855102193876209e-06,
+      "loss": 0.8324,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.00048551021938762084,
+      "step": 1324
+    },
+    {
+      "epoch": 0.14,
+      "gumbel_tau": 0.7317,
+      "image_masks_loss": 0.0383,
+      "llm_learning_rate": 4.85482372935605e-06,
+      "loss": 0.8114,
+      "output_text_mask_loss": 0.0193,
+      "predictor_learning_rate": 0.00048548237293560495,
+      "step": 1325
+    },
+    {
+      "epoch": 0.14,
+      "gumbel_tau": 0.7315,
+      "image_masks_loss": 0.042,
+      "llm_learning_rate": 4.854545005517456e-06,
+      "loss": 0.7757,
+      "output_text_mask_loss": 0.0626,
+      "predictor_learning_rate": 0.0004854545005517456,
+      "step": 1326
+    },
+    {
+      "epoch": 0.14,
+      "gumbel_tau": 0.7314,
+      "image_masks_loss": 0.0225,
+      "llm_learning_rate": 4.854266022391122e-06,
+      "loss": 0.7757,
+      "output_text_mask_loss": 0.0572,
+      "predictor_learning_rate": 0.00048542660223911216,
+      "step": 1327
+    },
+    {
+      "epoch": 0.14,
+      "gumbel_tau": 0.7312,
+      "image_masks_loss": 0.0134,
+      "llm_learning_rate": 4.8539867800077685e-06,
+      "loss": 0.6208,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.00048539867800077684,
+      "step": 1328
+    },
+    {
+      "epoch": 0.14,
+      "gumbel_tau": 0.731,
+      "image_masks_loss": 0.0105,
+      "llm_learning_rate": 4.853707278398149e-06,
+      "loss": 0.7614,
+      "output_text_mask_loss": 0.0086,
+      "predictor_learning_rate": 0.00048537072783981484,
+      "step": 1329
+    },
+    {
+      "epoch": 0.14,
+      "gumbel_tau": 0.7309,
+      "image_masks_loss": 0.0152,
+      "llm_learning_rate": 4.8534275175930405e-06,
+      "loss": 0.7879,
+      "output_text_mask_loss": 0.0012,
+      "predictor_learning_rate": 0.00048534275175930405,
+      "step": 1330
+    },
+    {
+      "epoch": 0.14,
+      "gumbel_tau": 0.7307,
+      "image_masks_loss": 0.0254,
+      "llm_learning_rate": 4.853147497623252e-06,
+      "loss": 0.775,
+      "output_text_mask_loss": 0.0219,
+      "predictor_learning_rate": 0.0004853147497623251,
+      "step": 1331
+    },
+    {
+      "epoch": 0.14,
+      "gumbel_tau": 0.7305,
+      "image_masks_loss": 0.015,
+      "llm_learning_rate": 4.85286721851962e-06,
+      "loss": 0.723,
+      "output_text_mask_loss": 0.1335,
+      "predictor_learning_rate": 0.0004852867218519619,
+      "step": 1332
+    },
+    {
+      "epoch": 0.14,
+      "gumbel_tau": 0.7303,
+      "image_masks_loss": 0.0201,
+      "llm_learning_rate": 4.852586680313008e-06,
+      "loss": 0.8317,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 0.0004852586680313008,
+      "step": 1333
+    },
+    {
+      "epoch": 0.14,
+      "gumbel_tau": 0.7302,
+      "image_masks_loss": 0.0096,
+      "llm_learning_rate": 4.852305883034312e-06,
+      "loss": 0.7991,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00048523058830343116,
+      "step": 1334
+    },
+    {
+      "epoch": 0.14,
+      "gumbel_tau": 0.73,
+      "image_masks_loss": 0.0123,
+      "llm_learning_rate": 4.852024826714451e-06,
+      "loss": 0.7464,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.0004852024826714451,
+      "step": 1335
+    },
+    {
+      "epoch": 0.14,
+      "gumbel_tau": 0.7298,
+      "image_masks_loss": 0.0187,
+      "llm_learning_rate": 4.851743511384378e-06,
+      "loss": 0.7922,
+      "output_text_mask_loss": 0.0208,
+      "predictor_learning_rate": 0.00048517435113843777,
+      "step": 1336
+    },
+    {
+      "epoch": 0.14,
+      "gumbel_tau": 0.7297,
+      "image_masks_loss": 0.0108,
+      "llm_learning_rate": 4.851461937075071e-06,
+      "loss": 0.7483,
+      "output_text_mask_loss": 0.0038,
+      "predictor_learning_rate": 0.0004851461937075071,
+      "step": 1337
+    },
+    {
+      "epoch": 0.14,
+      "gumbel_tau": 0.7295,
+      "image_masks_loss": 0.0235,
+      "llm_learning_rate": 4.851180103817539e-06,
+      "loss": 0.7645,
+      "output_text_mask_loss": 0.0687,
+      "predictor_learning_rate": 0.0004851180103817539,
+      "step": 1338
+    },
+    {
+      "epoch": 0.14,
+      "gumbel_tau": 0.7293,
+      "image_masks_loss": 0.0116,
+      "llm_learning_rate": 4.850898011642815e-06,
+      "loss": 0.6461,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 0.00048508980116428147,
+      "step": 1339
+    },
+    {
+      "epoch": 0.14,
+      "gumbel_tau": 0.7291,
+      "image_masks_loss": 0.0144,
+      "llm_learning_rate": 4.850615660581967e-06,
+      "loss": 0.7091,
+      "output_text_mask_loss": 0.0557,
+      "predictor_learning_rate": 0.00048506156605819664,
+      "step": 1340
+    },
+    {
+      "epoch": 0.14,
+      "gumbel_tau": 0.729,
+      "image_masks_loss": 0.0687,
+      "llm_learning_rate": 4.850333050666087e-06,
+      "loss": 0.7775,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.0004850333050666086,
+      "step": 1341
+    },
+    {
+      "epoch": 0.14,
+      "gumbel_tau": 0.7288,
+      "image_masks_loss": 0.0026,
+      "llm_learning_rate": 4.850050181926295e-06,
+      "loss": 0.733,
+      "output_text_mask_loss": 0.0099,
+      "predictor_learning_rate": 0.0004850050181926294,
+      "step": 1342
+    },
+    {
+      "epoch": 0.14,
+      "gumbel_tau": 0.7286,
+      "image_masks_loss": 0.0393,
+      "llm_learning_rate": 4.849767054393742e-06,
+      "loss": 0.7809,
+      "output_text_mask_loss": 0.0275,
+      "predictor_learning_rate": 0.0004849767054393742,
+      "step": 1343
+    },
+    {
+      "epoch": 0.14,
+      "gumbel_tau": 0.7284,
+      "image_masks_loss": 0.0265,
+      "llm_learning_rate": 4.849483668099608e-06,
+      "loss": 0.7905,
+      "output_text_mask_loss": 0.0391,
+      "predictor_learning_rate": 0.0004849483668099608,
+      "step": 1344
+    },
+    {
+      "epoch": 0.14,
+      "gumbel_tau": 0.7283,
+      "image_masks_loss": 0.0362,
+      "llm_learning_rate": 4.8492000230750995e-06,
+      "loss": 0.7423,
+      "output_text_mask_loss": 0.0275,
+      "predictor_learning_rate": 0.0004849200023075099,
+      "step": 1345
+    },
+    {
+      "epoch": 0.14,
+      "gumbel_tau": 0.7281,
+      "image_masks_loss": 0.0028,
+      "llm_learning_rate": 4.848916119351451e-06,
+      "loss": 0.6543,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.00048489161193514507,
+      "step": 1346
+    },
+    {
+      "epoch": 0.14,
+      "gumbel_tau": 0.7279,
+      "image_masks_loss": 0.0201,
+      "llm_learning_rate": 4.848631956959928e-06,
+      "loss": 0.7274,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.0004848631956959928,
+      "step": 1347
+    },
+    {
+      "epoch": 0.14,
+      "gumbel_tau": 0.7278,
+      "image_masks_loss": 0.0486,
+      "llm_learning_rate": 4.848347535931823e-06,
+      "loss": 0.8815,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.00048483475359318223,
+      "step": 1348
+    },
+    {
+      "epoch": 0.14,
+      "gumbel_tau": 0.7276,
+      "image_masks_loss": 0.0288,
+      "llm_learning_rate": 4.848062856298456e-06,
+      "loss": 0.8079,
+      "output_text_mask_loss": 0.0086,
+      "predictor_learning_rate": 0.0004848062856298455,
+      "step": 1349
+    },
+    {
+      "epoch": 0.14,
+      "gumbel_tau": 0.7274,
+      "image_masks_loss": 0.02,
+      "llm_learning_rate": 4.8477779180911765e-06,
+      "loss": 0.8048,
+      "output_text_mask_loss": 0.0056,
+      "predictor_learning_rate": 0.0004847777918091176,
+      "step": 1350
+    },
+    {
+      "epoch": 0.14,
+      "gumbel_tau": 0.7272,
+      "image_masks_loss": 0.0157,
+      "llm_learning_rate": 4.847492721341364e-06,
+      "loss": 0.7997,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.0004847492721341364,
+      "step": 1351
+    },
+    {
+      "epoch": 0.14,
+      "gumbel_tau": 0.7271,
+      "image_masks_loss": 0.05,
+      "llm_learning_rate": 4.847207266080424e-06,
+      "loss": 0.7709,
+      "output_text_mask_loss": 0.0164,
+      "predictor_learning_rate": 0.0004847207266080424,
+      "step": 1352
+    },
+    {
+      "epoch": 0.14,
+      "gumbel_tau": 0.7269,
+      "image_masks_loss": 0.0404,
+      "llm_learning_rate": 4.846921552339792e-06,
+      "loss": 0.8517,
+      "output_text_mask_loss": 0.0504,
+      "predictor_learning_rate": 0.00048469215523397915,
+      "step": 1353
+    },
+    {
+      "epoch": 0.14,
+      "gumbel_tau": 0.7267,
+      "image_masks_loss": 0.0045,
+      "llm_learning_rate": 4.8466355801509305e-06,
+      "loss": 0.8602,
+      "output_text_mask_loss": 0.0885,
+      "predictor_learning_rate": 0.00048466355801509304,
+      "step": 1354
+    },
+    {
+      "epoch": 0.14,
+      "gumbel_tau": 0.7266,
+      "image_masks_loss": 0.0052,
+      "llm_learning_rate": 4.846349349545333e-06,
+      "loss": 0.815,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00048463493495453324,
+      "step": 1355
+    },
+    {
+      "epoch": 0.14,
+      "gumbel_tau": 0.7264,
+      "image_masks_loss": 0.0145,
+      "llm_learning_rate": 4.846062860554518e-06,
+      "loss": 0.7392,
+      "output_text_mask_loss": 0.0221,
+      "predictor_learning_rate": 0.0004846062860554518,
+      "step": 1356
+    },
+    {
+      "epoch": 0.14,
+      "gumbel_tau": 0.7262,
+      "image_masks_loss": 0.0269,
+      "llm_learning_rate": 4.845776113210035e-06,
+      "loss": 0.8038,
+      "output_text_mask_loss": 0.0418,
+      "predictor_learning_rate": 0.00048457761132100345,
+      "step": 1357
+    },
+    {
+      "epoch": 0.14,
+      "gumbel_tau": 0.726,
+      "image_masks_loss": 0.0317,
+      "llm_learning_rate": 4.8454891075434614e-06,
+      "loss": 0.7575,
+      "output_text_mask_loss": 0.0115,
+      "predictor_learning_rate": 0.00048454891075434614,
+      "step": 1358
+    },
+    {
+      "epoch": 0.14,
+      "gumbel_tau": 0.7259,
+      "image_masks_loss": 0.0296,
+      "llm_learning_rate": 4.845201843586403e-06,
+      "loss": 0.6966,
+      "output_text_mask_loss": 0.0248,
+      "predictor_learning_rate": 0.00048452018435864034,
+      "step": 1359
+    },
+    {
+      "epoch": 0.14,
+      "gumbel_tau": 0.7257,
+      "image_masks_loss": 0.0214,
+      "llm_learning_rate": 4.8449143213704945e-06,
+      "loss": 0.656,
+      "output_text_mask_loss": 0.0145,
+      "predictor_learning_rate": 0.0004844914321370494,
+      "step": 1360
+    },
+    {
+      "epoch": 0.14,
+      "gumbel_tau": 0.7255,
+      "image_masks_loss": 0.0199,
+      "llm_learning_rate": 4.844626540927397e-06,
+      "loss": 0.7142,
+      "output_text_mask_loss": 0.0423,
+      "predictor_learning_rate": 0.0004844626540927397,
+      "step": 1361
+    },
+    {
+      "epoch": 0.14,
+      "gumbel_tau": 0.7254,
+      "image_masks_loss": 0.029,
+      "llm_learning_rate": 4.844338502288803e-06,
+      "loss": 0.6796,
+      "output_text_mask_loss": 0.002,
+      "predictor_learning_rate": 0.0004844338502288803,
+      "step": 1362
+    },
+    {
+      "epoch": 0.14,
+      "gumbel_tau": 0.7252,
+      "image_masks_loss": 0.0039,
+      "llm_learning_rate": 4.844050205486431e-06,
+      "loss": 0.7315,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004844050205486431,
+      "step": 1363
+    },
+    {
+      "epoch": 0.14,
+      "gumbel_tau": 0.725,
+      "image_masks_loss": 0.0146,
+      "llm_learning_rate": 4.843761650552029e-06,
+      "loss": 0.7955,
+      "output_text_mask_loss": 0.0305,
+      "predictor_learning_rate": 0.00048437616505520286,
+      "step": 1364
+    },
+    {
+      "epoch": 0.14,
+      "gumbel_tau": 0.7248,
+      "image_masks_loss": 0.0519,
+      "llm_learning_rate": 4.843472837517374e-06,
+      "loss": 0.7925,
+      "output_text_mask_loss": 0.0078,
+      "predictor_learning_rate": 0.00048434728375173743,
+      "step": 1365
+    },
+    {
+      "epoch": 0.14,
+      "gumbel_tau": 0.7247,
+      "image_masks_loss": 0.0248,
+      "llm_learning_rate": 4.84318376641427e-06,
+      "loss": 0.6844,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.000484318376641427,
+      "step": 1366
+    },
+    {
+      "epoch": 0.14,
+      "gumbel_tau": 0.7245,
+      "image_masks_loss": 0.0364,
+      "llm_learning_rate": 4.84289443727455e-06,
+      "loss": 0.8038,
+      "output_text_mask_loss": 0.0023,
+      "predictor_learning_rate": 0.00048428944372745506,
+      "step": 1367
+    },
+    {
+      "epoch": 0.14,
+      "gumbel_tau": 0.7243,
+      "image_masks_loss": 0.0305,
+      "llm_learning_rate": 4.842604850130077e-06,
+      "loss": 0.8059,
+      "output_text_mask_loss": 0.0132,
+      "predictor_learning_rate": 0.0004842604850130077,
+      "step": 1368
+    },
+    {
+      "epoch": 0.14,
+      "gumbel_tau": 0.7242,
+      "image_masks_loss": 0.0177,
+      "llm_learning_rate": 4.84231500501274e-06,
+      "loss": 0.8314,
+      "output_text_mask_loss": 0.071,
+      "predictor_learning_rate": 0.00048423150050127395,
+      "step": 1369
+    },
+    {
+      "epoch": 0.14,
+      "gumbel_tau": 0.724,
+      "image_masks_loss": 0.0184,
+      "llm_learning_rate": 4.842024901954456e-06,
+      "loss": 0.758,
+      "output_text_mask_loss": 0.0938,
+      "predictor_learning_rate": 0.0004842024901954456,
+      "step": 1370
+    },
+    {
+      "epoch": 0.14,
+      "gumbel_tau": 0.7238,
+      "image_masks_loss": 0.0179,
+      "llm_learning_rate": 4.841734540987174e-06,
+      "loss": 0.8088,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.00048417345409871736,
+      "step": 1371
+    },
+    {
+      "epoch": 0.14,
+      "gumbel_tau": 0.7237,
+      "image_masks_loss": 0.018,
+      "llm_learning_rate": 4.841443922142868e-06,
+      "loss": 0.731,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004841443922142868,
+      "step": 1372
+    },
+    {
+      "epoch": 0.14,
+      "gumbel_tau": 0.7235,
+      "image_masks_loss": 0.0214,
+      "llm_learning_rate": 4.841153045453543e-06,
+      "loss": 0.7368,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 0.00048411530454535424,
+      "step": 1373
+    },
+    {
+      "epoch": 0.14,
+      "gumbel_tau": 0.7233,
+      "image_masks_loss": 0.0169,
+      "llm_learning_rate": 4.840861910951228e-06,
+      "loss": 0.6668,
+      "output_text_mask_loss": NaN,
+      "predictor_learning_rate": 0.0004840861910951228,
+      "step": 1374
+    },
+    {
+      "epoch": 0.14,
+      "gumbel_tau": 0.7231,
+      "image_masks_loss": 0.0309,
+      "llm_learning_rate": 4.840570518667987e-06,
+      "loss": 0.787,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.00048405705186679865,
+      "step": 1375
+    },
+    {
+      "epoch": 0.14,
+      "gumbel_tau": 0.723,
+      "image_masks_loss": 0.0056,
+      "llm_learning_rate": 4.840278868635906e-06,
+      "loss": 0.6518,
+      "output_text_mask_loss": 0.0231,
+      "predictor_learning_rate": 0.0004840278868635906,
+      "step": 1376
+    },
+    {
+      "epoch": 0.14,
+      "gumbel_tau": 0.7228,
+      "image_masks_loss": 0.0162,
+      "llm_learning_rate": 4.839986960887104e-06,
+      "loss": 0.7663,
+      "output_text_mask_loss": 0.0099,
+      "predictor_learning_rate": 0.0004839986960887104,
+      "step": 1377
+    },
+    {
+      "epoch": 0.14,
+      "gumbel_tau": 0.7226,
+      "image_masks_loss": 0.0058,
+      "llm_learning_rate": 4.839694795453725e-06,
+      "loss": 0.8295,
+      "output_text_mask_loss": 0.0035,
+      "predictor_learning_rate": 0.00048396947954537254,
+      "step": 1378
+    },
+    {
+      "epoch": 0.14,
+      "gumbel_tau": 0.7225,
+      "image_masks_loss": 0.0135,
+      "llm_learning_rate": 4.839402372367944e-06,
+      "loss": 0.7221,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.00048394023723679437,
+      "step": 1379
+    },
+    {
+      "epoch": 0.14,
+      "gumbel_tau": 0.7223,
+      "image_masks_loss": 0.021,
+      "llm_learning_rate": 4.839109691661963e-06,
+      "loss": 0.8408,
+      "output_text_mask_loss": 0.0153,
+      "predictor_learning_rate": 0.00048391096916619624,
+      "step": 1380
+    },
+    {
+      "epoch": 0.14,
+      "gumbel_tau": 0.7221,
+      "image_masks_loss": 0.019,
+      "llm_learning_rate": 4.838816753368013e-06,
+      "loss": 0.8592,
+      "output_text_mask_loss": 0.0433,
+      "predictor_learning_rate": 0.00048388167533680126,
+      "step": 1381
+    },
+    {
+      "epoch": 0.14,
+      "gumbel_tau": 0.7219,
+      "image_masks_loss": 0.0191,
+      "llm_learning_rate": 4.838523557518352e-06,
+      "loss": 0.7416,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00048385235575183516,
+      "step": 1382
+    },
+    {
+      "epoch": 0.14,
+      "gumbel_tau": 0.7218,
+      "image_masks_loss": 0.0174,
+      "llm_learning_rate": 4.838230104145268e-06,
+      "loss": 0.6973,
+      "output_text_mask_loss": 0.002,
+      "predictor_learning_rate": 0.0004838230104145268,
+      "step": 1383
+    },
+    {
+      "epoch": 0.14,
+      "gumbel_tau": 0.7216,
+      "image_masks_loss": 0.0299,
+      "llm_learning_rate": 4.837936393281078e-06,
+      "loss": 0.654,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004837936393281077,
+      "step": 1384
+    },
+    {
+      "epoch": 0.14,
+      "gumbel_tau": 0.7214,
+      "image_masks_loss": 0.0277,
+      "llm_learning_rate": 4.837642424958123e-06,
+      "loss": 0.8026,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.00048376424249581225,
+      "step": 1385
+    },
+    {
+      "epoch": 0.14,
+      "gumbel_tau": 0.7213,
+      "image_masks_loss": 0.02,
+      "llm_learning_rate": 4.837348199208778e-06,
+      "loss": 0.9071,
+      "output_text_mask_loss": 0.0511,
+      "predictor_learning_rate": 0.00048373481992087776,
+      "step": 1386
+    },
+    {
+      "epoch": 0.14,
+      "gumbel_tau": 0.7211,
+      "image_masks_loss": 0.0263,
+      "llm_learning_rate": 4.837053716065443e-06,
+      "loss": 0.7167,
+      "output_text_mask_loss": 0.0284,
+      "predictor_learning_rate": 0.0004837053716065443,
+      "step": 1387
+    },
+    {
+      "epoch": 0.14,
+      "gumbel_tau": 0.7209,
+      "image_masks_loss": 0.0232,
+      "llm_learning_rate": 4.836758975560547e-06,
+      "loss": 0.7193,
+      "output_text_mask_loss": 0.0124,
+      "predictor_learning_rate": 0.0004836758975560547,
+      "step": 1388
+    },
+    {
+      "epoch": 0.14,
+      "gumbel_tau": 0.7208,
+      "image_masks_loss": 0.0236,
+      "llm_learning_rate": 4.836463977726548e-06,
+      "loss": 0.7144,
+      "output_text_mask_loss": 0.0462,
+      "predictor_learning_rate": 0.0004836463977726547,
+      "step": 1389
+    },
+    {
+      "epoch": 0.14,
+      "gumbel_tau": 0.7206,
+      "image_masks_loss": 0.0081,
+      "llm_learning_rate": 4.836168722595931e-06,
+      "loss": 0.8881,
+      "output_text_mask_loss": 0.0025,
+      "predictor_learning_rate": 0.0004836168722595931,
+      "step": 1390
+    },
+    {
+      "epoch": 0.14,
+      "gumbel_tau": 0.7204,
+      "image_masks_loss": 0.0213,
+      "llm_learning_rate": 4.8358732102012116e-06,
+      "loss": 0.8104,
+      "output_text_mask_loss": 0.0187,
+      "predictor_learning_rate": 0.00048358732102012113,
+      "step": 1391
+    },
+    {
+      "epoch": 0.14,
+      "gumbel_tau": 0.7202,
+      "image_masks_loss": 0.0278,
+      "llm_learning_rate": 4.8355774405749315e-06,
+      "loss": 0.8582,
+      "output_text_mask_loss": 0.0353,
+      "predictor_learning_rate": 0.0004835577440574931,
+      "step": 1392
+    },
+    {
+      "epoch": 0.14,
+      "gumbel_tau": 0.7201,
+      "image_masks_loss": 0.0178,
+      "llm_learning_rate": 4.83528141374966e-06,
+      "loss": 0.7752,
+      "output_text_mask_loss": 0.053,
+      "predictor_learning_rate": 0.000483528141374966,
+      "step": 1393
+    },
+    {
+      "epoch": 0.14,
+      "gumbel_tau": 0.7199,
+      "image_masks_loss": 0.0317,
+      "llm_learning_rate": 4.834985129757999e-06,
+      "loss": 0.8622,
+      "output_text_mask_loss": 0.0296,
+      "predictor_learning_rate": 0.0004834985129757998,
+      "step": 1394
+    },
+    {
+      "epoch": 0.14,
+      "gumbel_tau": 0.7197,
+      "image_masks_loss": 0.0241,
+      "llm_learning_rate": 4.834688588632574e-06,
+      "loss": 0.7656,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.0004834688588632573,
+      "step": 1395
+    },
+    {
+      "epoch": 0.14,
+      "gumbel_tau": 0.7196,
+      "image_masks_loss": 0.0095,
+      "llm_learning_rate": 4.83439179040604e-06,
+      "loss": 0.7408,
+      "output_text_mask_loss": 0.0196,
+      "predictor_learning_rate": 0.00048343917904060396,
+      "step": 1396
+    },
+    {
+      "epoch": 0.14,
+      "gumbel_tau": 0.7194,
+      "image_masks_loss": 0.0246,
+      "llm_learning_rate": 4.834094735111083e-06,
+      "loss": 0.719,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004834094735111083,
+      "step": 1397
+    },
+    {
+      "epoch": 0.14,
+      "gumbel_tau": 0.7192,
+      "image_masks_loss": 0.0389,
+      "llm_learning_rate": 4.833797422780416e-06,
+      "loss": 0.8549,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.0004833797422780416,
+      "step": 1398
+    },
+    {
+      "epoch": 0.14,
+      "gumbel_tau": 0.7191,
+      "image_masks_loss": 0.0208,
+      "llm_learning_rate": 4.833499853446778e-06,
+      "loss": 0.8251,
+      "output_text_mask_loss": 0.0216,
+      "predictor_learning_rate": 0.00048334998534467775,
+      "step": 1399
+    },
+    {
+      "epoch": 0.14,
+      "gumbel_tau": 0.7189,
+      "image_masks_loss": 0.0218,
+      "llm_learning_rate": 4.833202027142938e-06,
+      "loss": 0.7597,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004833202027142938,
+      "step": 1400
+    },
+    {
+      "epoch": 0.14,
+      "gumbel_tau": 0.7187,
+      "image_masks_loss": 0.0362,
+      "llm_learning_rate": 4.832903943901694e-06,
+      "loss": 0.8685,
+      "output_text_mask_loss": 0.0252,
+      "predictor_learning_rate": 0.0004832903943901694,
+      "step": 1401
+    },
+    {
+      "epoch": 0.14,
+      "gumbel_tau": 0.7185,
+      "image_masks_loss": 0.0248,
+      "llm_learning_rate": 4.832605603755871e-06,
+      "loss": 0.7914,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.0004832605603755871,
+      "step": 1402
+    },
+    {
+      "epoch": 0.14,
+      "gumbel_tau": 0.7184,
+      "image_masks_loss": 0.0277,
+      "llm_learning_rate": 4.8323070067383236e-06,
+      "loss": 0.8091,
+      "output_text_mask_loss": 0.0347,
+      "predictor_learning_rate": 0.00048323070067383234,
+      "step": 1403
+    },
+    {
+      "epoch": 0.14,
+      "gumbel_tau": 0.7182,
+      "image_masks_loss": 0.0265,
+      "llm_learning_rate": 4.832008152881934e-06,
+      "loss": 0.8421,
+      "output_text_mask_loss": 0.0351,
+      "predictor_learning_rate": 0.0004832008152881934,
+      "step": 1404
+    },
+    {
+      "epoch": 0.14,
+      "gumbel_tau": 0.718,
+      "image_masks_loss": 0.0126,
+      "llm_learning_rate": 4.831709042219611e-06,
+      "loss": 0.8256,
+      "output_text_mask_loss": 0.0116,
+      "predictor_learning_rate": 0.00048317090422196114,
+      "step": 1405
+    },
+    {
+      "epoch": 0.14,
+      "gumbel_tau": 0.7179,
+      "image_masks_loss": 0.0089,
+      "llm_learning_rate": 4.831409674784297e-06,
+      "loss": 0.6956,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00048314096747842963,
+      "step": 1406
+    },
+    {
+      "epoch": 0.14,
+      "gumbel_tau": 0.7177,
+      "image_masks_loss": 0.0162,
+      "llm_learning_rate": 4.831110050608954e-06,
+      "loss": 0.6776,
+      "output_text_mask_loss": 0.0065,
+      "predictor_learning_rate": 0.0004831110050608954,
+      "step": 1407
+    },
+    {
+      "epoch": 0.14,
+      "gumbel_tau": 0.7175,
+      "image_masks_loss": 0.0098,
+      "llm_learning_rate": 4.830810169726581e-06,
+      "loss": 0.7335,
+      "output_text_mask_loss": 0.0275,
+      "predictor_learning_rate": 0.00048308101697265804,
+      "step": 1408
+    },
+    {
+      "epoch": 0.14,
+      "gumbel_tau": 0.7174,
+      "image_masks_loss": 0.0278,
+      "llm_learning_rate": 4.8305100321702e-06,
+      "loss": 0.738,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.00048305100321702,
+      "step": 1409
+    },
+    {
+      "epoch": 0.14,
+      "gumbel_tau": 0.7172,
+      "image_masks_loss": 0.0186,
+      "llm_learning_rate": 4.830209637972864e-06,
+      "loss": 0.7604,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 0.00048302096379728636,
+      "step": 1410
+    },
+    {
+      "epoch": 0.14,
+      "gumbel_tau": 0.717,
+      "image_masks_loss": 0.004,
+      "llm_learning_rate": 4.829908987167652e-06,
+      "loss": 0.6886,
+      "output_text_mask_loss": 0.0275,
+      "predictor_learning_rate": 0.00048299089871676516,
+      "step": 1411
+    },
+    {
+      "epoch": 0.14,
+      "gumbel_tau": 0.7169,
+      "image_masks_loss": 0.0092,
+      "llm_learning_rate": 4.8296080797876725e-06,
+      "loss": 0.7961,
+      "output_text_mask_loss": 0.0246,
+      "predictor_learning_rate": 0.0004829608079787672,
+      "step": 1412
+    },
+    {
+      "epoch": 0.14,
+      "gumbel_tau": 0.7167,
+      "image_masks_loss": 0.0213,
+      "llm_learning_rate": 4.829306915866062e-06,
+      "loss": 0.8349,
+      "output_text_mask_loss": 0.0946,
+      "predictor_learning_rate": 0.0004829306915866062,
+      "step": 1413
+    },
+    {
+      "epoch": 0.14,
+      "gumbel_tau": 0.7165,
+      "image_masks_loss": 0.0121,
+      "llm_learning_rate": 4.8290054954359865e-06,
+      "loss": 0.8069,
+      "output_text_mask_loss": 0.0362,
+      "predictor_learning_rate": 0.00048290054954359864,
+      "step": 1414
+    },
+    {
+      "epoch": 0.14,
+      "gumbel_tau": 0.7163,
+      "image_masks_loss": 0.0092,
+      "llm_learning_rate": 4.8287038185306375e-06,
+      "loss": 0.7854,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.00048287038185306375,
+      "step": 1415
+    },
+    {
+      "epoch": 0.15,
+      "gumbel_tau": 0.7162,
+      "image_masks_loss": 0.0082,
+      "llm_learning_rate": 4.828401885183237e-06,
+      "loss": 0.742,
+      "output_text_mask_loss": 0.0429,
+      "predictor_learning_rate": 0.0004828401885183237,
+      "step": 1416
+    },
+    {
+      "epoch": 0.15,
+      "gumbel_tau": 0.716,
+      "image_masks_loss": 0.0118,
+      "llm_learning_rate": 4.828099695427035e-06,
+      "loss": 0.8646,
+      "output_text_mask_loss": 0.0081,
+      "predictor_learning_rate": 0.00048280996954270345,
+      "step": 1417
+    },
+    {
+      "epoch": 0.15,
+      "gumbel_tau": 0.7158,
+      "image_masks_loss": 0.0211,
+      "llm_learning_rate": 4.827797249295308e-06,
+      "loss": 0.7442,
+      "output_text_mask_loss": 0.0042,
+      "predictor_learning_rate": 0.0004827797249295308,
+      "step": 1418
+    },
+    {
+      "epoch": 0.15,
+      "gumbel_tau": 0.7157,
+      "image_masks_loss": 0.0133,
+      "llm_learning_rate": 4.827494546821364e-06,
+      "loss": 0.7708,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.0004827494546821364,
+      "step": 1419
+    },
+    {
+      "epoch": 0.15,
+      "gumbel_tau": 0.7155,
+      "image_masks_loss": 0.0239,
+      "llm_learning_rate": 4.8271915880385355e-06,
+      "loss": 0.8462,
+      "output_text_mask_loss": 0.0231,
+      "predictor_learning_rate": 0.00048271915880385353,
+      "step": 1420
+    },
+    {
+      "epoch": 0.15,
+      "gumbel_tau": 0.7153,
+      "image_masks_loss": 0.0429,
+      "llm_learning_rate": 4.826888372980187e-06,
+      "loss": 0.8531,
+      "output_text_mask_loss": 0.0078,
+      "predictor_learning_rate": 0.0004826888372980186,
+      "step": 1421
+    },
+    {
+      "epoch": 0.15,
+      "gumbel_tau": 0.7152,
+      "image_masks_loss": 0.019,
+      "llm_learning_rate": 4.8265849016797065e-06,
+      "loss": 0.8045,
+      "output_text_mask_loss": 0.0095,
+      "predictor_learning_rate": 0.00048265849016797057,
+      "step": 1422
+    },
+    {
+      "epoch": 0.15,
+      "gumbel_tau": 0.715,
+      "image_masks_loss": 0.0162,
+      "llm_learning_rate": 4.826281174170515e-06,
+      "loss": 0.769,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004826281174170515,
+      "step": 1423
+    },
+    {
+      "epoch": 0.15,
+      "gumbel_tau": 0.7148,
+      "image_masks_loss": 0.0166,
+      "llm_learning_rate": 4.8259771904860585e-06,
+      "loss": 0.8039,
+      "output_text_mask_loss": 0.0259,
+      "predictor_learning_rate": 0.00048259771904860584,
+      "step": 1424
+    },
+    {
+      "epoch": 0.15,
+      "gumbel_tau": 0.7147,
+      "image_masks_loss": 0.0198,
+      "llm_learning_rate": 4.825672950659813e-06,
+      "loss": 0.6872,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.0004825672950659813,
+      "step": 1425
+    },
+    {
+      "epoch": 0.15,
+      "gumbel_tau": 0.7145,
+      "image_masks_loss": 0.0206,
+      "llm_learning_rate": 4.825368454725282e-06,
+      "loss": 0.7941,
+      "output_text_mask_loss": 0.001,
+      "predictor_learning_rate": 0.0004825368454725282,
+      "step": 1426
+    },
+    {
+      "epoch": 0.15,
+      "gumbel_tau": 0.7143,
+      "image_masks_loss": 0.0067,
+      "llm_learning_rate": 4.825063702715997e-06,
+      "loss": 0.8064,
+      "output_text_mask_loss": 0.0275,
+      "predictor_learning_rate": 0.00048250637027159974,
+      "step": 1427
+    },
+    {
+      "epoch": 0.15,
+      "gumbel_tau": 0.7142,
+      "image_masks_loss": 0.0073,
+      "llm_learning_rate": 4.824758694665519e-06,
+      "loss": 0.7369,
+      "output_text_mask_loss": 0.0244,
+      "predictor_learning_rate": 0.0004824758694665519,
+      "step": 1428
+    },
+    {
+      "epoch": 0.15,
+      "gumbel_tau": 0.714,
+      "image_masks_loss": 0.0121,
+      "llm_learning_rate": 4.824453430607435e-06,
+      "loss": 0.7404,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.0004824453430607435,
+      "step": 1429
+    },
+    {
+      "epoch": 0.15,
+      "gumbel_tau": 0.7138,
+      "image_masks_loss": 0.0066,
+      "llm_learning_rate": 4.824147910575362e-06,
+      "loss": 0.7599,
+      "output_text_mask_loss": 0.0292,
+      "predictor_learning_rate": 0.0004824147910575362,
+      "step": 1430
+    },
+    {
+      "epoch": 0.15,
+      "gumbel_tau": 0.7136,
+      "image_masks_loss": 0.0296,
+      "llm_learning_rate": 4.8238421346029435e-06,
+      "loss": 0.7834,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00048238421346029434,
+      "step": 1431
+    },
+    {
+      "epoch": 0.15,
+      "gumbel_tau": 0.7135,
+      "image_masks_loss": 0.0232,
+      "llm_learning_rate": 4.8235361027238536e-06,
+      "loss": 0.7844,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 0.00048235361027238534,
+      "step": 1432
+    },
+    {
+      "epoch": 0.15,
+      "gumbel_tau": 0.7133,
+      "image_masks_loss": 0.0136,
+      "llm_learning_rate": 4.823229814971792e-06,
+      "loss": 0.7552,
+      "output_text_mask_loss": 0.0269,
+      "predictor_learning_rate": 0.0004823229814971792,
+      "step": 1433
+    },
+    {
+      "epoch": 0.15,
+      "gumbel_tau": 0.7131,
+      "image_masks_loss": 0.0252,
+      "llm_learning_rate": 4.822923271380488e-06,
+      "loss": 0.7412,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 0.00048229232713804885,
+      "step": 1434
+    },
+    {
+      "epoch": 0.15,
+      "gumbel_tau": 0.713,
+      "image_masks_loss": 0.0254,
+      "llm_learning_rate": 4.822616471983701e-06,
+      "loss": 0.8637,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004822616471983701,
+      "step": 1435
+    },
+    {
+      "epoch": 0.15,
+      "gumbel_tau": 0.7128,
+      "image_masks_loss": 0.02,
+      "llm_learning_rate": 4.822309416815213e-06,
+      "loss": 0.7884,
+      "output_text_mask_loss": 0.0004,
+      "predictor_learning_rate": 0.0004822309416815213,
+      "step": 1436
+    },
+    {
+      "epoch": 0.15,
+      "gumbel_tau": 0.7126,
+      "image_masks_loss": 0.0137,
+      "llm_learning_rate": 4.82200210590884e-06,
+      "loss": 0.8211,
+      "output_text_mask_loss": 0.0058,
+      "predictor_learning_rate": 0.000482200210590884,
+      "step": 1437
+    },
+    {
+      "epoch": 0.15,
+      "gumbel_tau": 0.7125,
+      "image_masks_loss": 0.0086,
+      "llm_learning_rate": 4.821694539298423e-06,
+      "loss": 0.6899,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 0.00048216945392984224,
+      "step": 1438
+    },
+    {
+      "epoch": 0.15,
+      "gumbel_tau": 0.7123,
+      "image_masks_loss": 0.0265,
+      "llm_learning_rate": 4.821386717017832e-06,
+      "loss": 0.7892,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.0004821386717017832,
+      "step": 1439
+    },
+    {
+      "epoch": 0.15,
+      "gumbel_tau": 0.7121,
+      "image_masks_loss": 0.0188,
+      "llm_learning_rate": 4.821078639100964e-06,
+      "loss": 0.7646,
+      "output_text_mask_loss": 0.0098,
+      "predictor_learning_rate": 0.0004821078639100964,
+      "step": 1440
+    },
+    {
+      "epoch": 0.15,
+      "gumbel_tau": 0.712,
+      "image_masks_loss": 0.0311,
+      "llm_learning_rate": 4.820770305581748e-06,
+      "loss": 0.7055,
+      "output_text_mask_loss": 0.0263,
+      "predictor_learning_rate": 0.0004820770305581748,
+      "step": 1441
+    },
+    {
+      "epoch": 0.15,
+      "gumbel_tau": 0.7118,
+      "image_masks_loss": 0.0334,
+      "llm_learning_rate": 4.820461716494136e-06,
+      "loss": 0.7918,
+      "output_text_mask_loss": 0.1633,
+      "predictor_learning_rate": 0.00048204617164941357,
+      "step": 1442
+    },
+    {
+      "epoch": 0.15,
+      "gumbel_tau": 0.7116,
+      "image_masks_loss": 0.0224,
+      "llm_learning_rate": 4.820152871872111e-06,
+      "loss": 0.7256,
+      "output_text_mask_loss": 0.041,
+      "predictor_learning_rate": 0.00048201528718721105,
+      "step": 1443
+    },
+    {
+      "epoch": 0.15,
+      "gumbel_tau": 0.7115,
+      "image_masks_loss": 0.0301,
+      "llm_learning_rate": 4.819843771749684e-06,
+      "loss": 0.8244,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.00048198437717496837,
+      "step": 1444
+    },
+    {
+      "epoch": 0.15,
+      "gumbel_tau": 0.7113,
+      "image_masks_loss": 0.025,
+      "llm_learning_rate": 4.819534416160893e-06,
+      "loss": 0.726,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004819534416160892,
+      "step": 1445
+    },
+    {
+      "epoch": 0.15,
+      "gumbel_tau": 0.7111,
+      "image_masks_loss": 0.0198,
+      "llm_learning_rate": 4.819224805139806e-06,
+      "loss": 0.743,
+      "output_text_mask_loss": 0.0254,
+      "predictor_learning_rate": 0.00048192248051398057,
+      "step": 1446
+    },
+    {
+      "epoch": 0.15,
+      "gumbel_tau": 0.711,
+      "image_masks_loss": 0.0244,
+      "llm_learning_rate": 4.818914938720518e-06,
+      "loss": 0.7206,
+      "output_text_mask_loss": 0.0208,
+      "predictor_learning_rate": 0.00048189149387205173,
+      "step": 1447
+    },
+    {
+      "epoch": 0.15,
+      "gumbel_tau": 0.7108,
+      "image_masks_loss": 0.0188,
+      "llm_learning_rate": 4.818604816937151e-06,
+      "loss": 0.8297,
+      "output_text_mask_loss": 0.0275,
+      "predictor_learning_rate": 0.00048186048169371505,
+      "step": 1448
+    },
+    {
+      "epoch": 0.15,
+      "gumbel_tau": 0.7106,
+      "image_masks_loss": 0.0052,
+      "llm_learning_rate": 4.818294439823857e-06,
+      "loss": 0.8269,
+      "output_text_mask_loss": 0.0416,
+      "predictor_learning_rate": 0.0004818294439823857,
+      "step": 1449
+    },
+    {
+      "epoch": 0.15,
+      "gumbel_tau": 0.7105,
+      "image_masks_loss": 0.0091,
+      "llm_learning_rate": 4.8179838074148154e-06,
+      "loss": 0.6648,
+      "output_text_mask_loss": 0.0041,
+      "predictor_learning_rate": 0.0004817983807414815,
+      "step": 1450
+    },
+    {
+      "epoch": 0.15,
+      "gumbel_tau": 0.7103,
+      "image_masks_loss": 0.0254,
+      "llm_learning_rate": 4.817672919744235e-06,
+      "loss": 0.7675,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.00048176729197442345,
+      "step": 1451
+    },
+    {
+      "epoch": 0.15,
+      "gumbel_tau": 0.7101,
+      "image_masks_loss": 0.0402,
+      "llm_learning_rate": 4.8173617768463495e-06,
+      "loss": 0.8019,
+      "output_text_mask_loss": 0.0072,
+      "predictor_learning_rate": 0.0004817361776846349,
+      "step": 1452
+    },
+    {
+      "epoch": 0.15,
+      "gumbel_tau": 0.71,
+      "image_masks_loss": 0.0136,
+      "llm_learning_rate": 4.817050378755423e-06,
+      "loss": 0.8454,
+      "output_text_mask_loss": 0.034,
+      "predictor_learning_rate": 0.0004817050378755423,
+      "step": 1453
+    },
+    {
+      "epoch": 0.15,
+      "gumbel_tau": 0.7098,
+      "image_masks_loss": 0.0309,
+      "llm_learning_rate": 4.816738725505749e-06,
+      "loss": 0.784,
+      "output_text_mask_loss": 0.0025,
+      "predictor_learning_rate": 0.0004816738725505748,
+      "step": 1454
+    },
+    {
+      "epoch": 0.15,
+      "gumbel_tau": 0.7096,
+      "image_masks_loss": 0.0211,
+      "llm_learning_rate": 4.8164268171316445e-06,
+      "loss": 0.709,
+      "output_text_mask_loss": 0.016,
+      "predictor_learning_rate": 0.00048164268171316444,
+      "step": 1455
+    },
+    {
+      "epoch": 0.15,
+      "gumbel_tau": 0.7095,
+      "image_masks_loss": 0.0078,
+      "llm_learning_rate": 4.816114653667461e-06,
+      "loss": 0.8165,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.00048161146536674607,
+      "step": 1456
+    },
+    {
+      "epoch": 0.15,
+      "gumbel_tau": 0.7093,
+      "image_masks_loss": 0.032,
+      "llm_learning_rate": 4.815802235147571e-06,
+      "loss": 0.7444,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.0004815802235147571,
+      "step": 1457
+    },
+    {
+      "epoch": 0.15,
+      "gumbel_tau": 0.7091,
+      "image_masks_loss": 0.0187,
+      "llm_learning_rate": 4.815489561606382e-06,
+      "loss": 0.7051,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00048154895616063823,
+      "step": 1458
+    },
+    {
+      "epoch": 0.15,
+      "gumbel_tau": 0.7089,
+      "image_masks_loss": 0.0228,
+      "llm_learning_rate": 4.815176633078325e-06,
+      "loss": 0.7381,
+      "output_text_mask_loss": 0.0153,
+      "predictor_learning_rate": 0.0004815176633078324,
+      "step": 1459
+    },
+    {
+      "epoch": 0.15,
+      "gumbel_tau": 0.7088,
+      "image_masks_loss": 0.028,
+      "llm_learning_rate": 4.81486344959786e-06,
+      "loss": 0.7752,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.00048148634495978593,
+      "step": 1460
+    },
+    {
+      "epoch": 0.15,
+      "gumbel_tau": 0.7086,
+      "image_masks_loss": 0.0919,
+      "llm_learning_rate": 4.814550011199475e-06,
+      "loss": 0.8592,
+      "output_text_mask_loss": 0.0019,
+      "predictor_learning_rate": 0.0004814550011199475,
+      "step": 1461
+    },
+    {
+      "epoch": 0.15,
+      "gumbel_tau": 0.7084,
+      "image_masks_loss": 0.0378,
+      "llm_learning_rate": 4.814236317917689e-06,
+      "loss": 0.7745,
+      "output_text_mask_loss": 0.0128,
+      "predictor_learning_rate": 0.00048142363179176883,
+      "step": 1462
+    },
+    {
+      "epoch": 0.15,
+      "gumbel_tau": 0.7083,
+      "image_masks_loss": 0.0215,
+      "llm_learning_rate": 4.813922369787043e-06,
+      "loss": 0.8053,
+      "output_text_mask_loss": 0.0275,
+      "predictor_learning_rate": 0.00048139223697870427,
+      "step": 1463
+    },
+    {
+      "epoch": 0.15,
+      "gumbel_tau": 0.7081,
+      "image_masks_loss": 0.061,
+      "llm_learning_rate": 4.813608166842113e-06,
+      "loss": 0.8305,
+      "output_text_mask_loss": 0.0322,
+      "predictor_learning_rate": 0.00048136081668421126,
+      "step": 1464
+    },
+    {
+      "epoch": 0.15,
+      "gumbel_tau": 0.7079,
+      "image_masks_loss": 0.0429,
+      "llm_learning_rate": 4.813293709117497e-06,
+      "loss": 0.78,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00048132937091174973,
+      "step": 1465
+    },
+    {
+      "epoch": 0.15,
+      "gumbel_tau": 0.7078,
+      "image_masks_loss": 0.0107,
+      "llm_learning_rate": 4.812978996647827e-06,
+      "loss": 0.7937,
+      "output_text_mask_loss": 0.0248,
+      "predictor_learning_rate": 0.00048129789966478266,
+      "step": 1466
+    },
+    {
+      "epoch": 0.15,
+      "gumbel_tau": 0.7076,
+      "image_masks_loss": 0.1236,
+      "llm_learning_rate": 4.8126640294677565e-06,
+      "loss": 0.736,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.0004812664029467756,
+      "step": 1467
+    },
+    {
+      "epoch": 0.15,
+      "gumbel_tau": 0.7074,
+      "image_masks_loss": 0.0182,
+      "llm_learning_rate": 4.812348807611972e-06,
+      "loss": 0.8341,
+      "output_text_mask_loss": 0.0113,
+      "predictor_learning_rate": 0.00048123488076119714,
+      "step": 1468
+    },
+    {
+      "epoch": 0.15,
+      "gumbel_tau": 0.7073,
+      "image_masks_loss": 0.0222,
+      "llm_learning_rate": 4.812033331115186e-06,
+      "loss": 0.8051,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.0004812033331115186,
+      "step": 1469
+    },
+    {
+      "epoch": 0.15,
+      "gumbel_tau": 0.7071,
+      "image_masks_loss": 0.0622,
+      "llm_learning_rate": 4.81171760001214e-06,
+      "loss": 0.7532,
+      "output_text_mask_loss": 0.0023,
+      "predictor_learning_rate": 0.000481171760001214,
+      "step": 1470
+    },
+    {
+      "epoch": 0.15,
+      "gumbel_tau": 0.7069,
+      "image_masks_loss": 0.0721,
+      "llm_learning_rate": 4.811401614337603e-06,
+      "loss": 0.7975,
+      "output_text_mask_loss": 0.0282,
+      "predictor_learning_rate": 0.00048114016143376026,
+      "step": 1471
+    },
+    {
+      "epoch": 0.15,
+      "gumbel_tau": 0.7068,
+      "image_masks_loss": 0.0374,
+      "llm_learning_rate": 4.811085374126372e-06,
+      "loss": 0.8068,
+      "output_text_mask_loss": 0.001,
+      "predictor_learning_rate": 0.0004811085374126372,
+      "step": 1472
+    },
+    {
+      "epoch": 0.15,
+      "gumbel_tau": 0.7066,
+      "image_masks_loss": 0.0137,
+      "llm_learning_rate": 4.810768879413271e-06,
+      "loss": 0.7357,
+      "output_text_mask_loss": 0.0257,
+      "predictor_learning_rate": 0.0004810768879413271,
+      "step": 1473
+    },
+    {
+      "epoch": 0.15,
+      "gumbel_tau": 0.7064,
+      "image_masks_loss": 0.036,
+      "llm_learning_rate": 4.810452130233155e-06,
+      "loss": 0.8921,
+      "output_text_mask_loss": 0.0298,
+      "predictor_learning_rate": 0.0004810452130233155,
+      "step": 1474
+    },
+    {
+      "epoch": 0.15,
+      "gumbel_tau": 0.7063,
+      "image_masks_loss": 0.0271,
+      "llm_learning_rate": 4.810135126620904e-06,
+      "loss": 0.761,
+      "output_text_mask_loss": 0.0187,
+      "predictor_learning_rate": 0.0004810135126620903,
+      "step": 1475
+    },
+    {
+      "epoch": 0.15,
+      "gumbel_tau": 0.7061,
+      "image_masks_loss": 0.0397,
+      "llm_learning_rate": 4.809817868611426e-06,
+      "loss": 0.7888,
+      "output_text_mask_loss": 0.0019,
+      "predictor_learning_rate": 0.0004809817868611426,
+      "step": 1476
+    },
+    {
+      "epoch": 0.15,
+      "gumbel_tau": 0.7059,
+      "image_masks_loss": 0.0185,
+      "llm_learning_rate": 4.809500356239661e-06,
+      "loss": 0.6945,
+      "output_text_mask_loss": 0.0324,
+      "predictor_learning_rate": 0.000480950035623966,
+      "step": 1477
+    },
+    {
+      "epoch": 0.15,
+      "gumbel_tau": 0.7058,
+      "image_masks_loss": 0.0199,
+      "llm_learning_rate": 4.809182589540571e-06,
+      "loss": 0.6953,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00048091825895405703,
+      "step": 1478
+    },
+    {
+      "epoch": 0.15,
+      "gumbel_tau": 0.7056,
+      "image_masks_loss": 0.0154,
+      "llm_learning_rate": 4.808864568549152e-06,
+      "loss": 0.8263,
+      "output_text_mask_loss": 0.0832,
+      "predictor_learning_rate": 0.00048088645685491515,
+      "step": 1479
+    },
+    {
+      "epoch": 0.15,
+      "gumbel_tau": 0.7054,
+      "image_masks_loss": 0.0179,
+      "llm_learning_rate": 4.808546293300423e-06,
+      "loss": 0.8438,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 0.00048085462933004225,
+      "step": 1480
+    },
+    {
+      "epoch": 0.15,
+      "gumbel_tau": 0.7053,
+      "image_masks_loss": 0.0336,
+      "llm_learning_rate": 4.808227763829434e-06,
+      "loss": 0.7813,
+      "output_text_mask_loss": 0.0062,
+      "predictor_learning_rate": 0.0004808227763829434,
+      "step": 1481
+    },
+    {
+      "epoch": 0.15,
+      "gumbel_tau": 0.7051,
+      "image_masks_loss": 0.0292,
+      "llm_learning_rate": 4.807908980171263e-06,
+      "loss": 0.7338,
+      "output_text_mask_loss": 0.0137,
+      "predictor_learning_rate": 0.0004807908980171263,
+      "step": 1482
+    },
+    {
+      "epoch": 0.15,
+      "gumbel_tau": 0.7049,
+      "image_masks_loss": 0.0315,
+      "llm_learning_rate": 4.8075899423610136e-06,
+      "loss": 0.7732,
+      "output_text_mask_loss": 0.0078,
+      "predictor_learning_rate": 0.00048075899423610135,
+      "step": 1483
+    },
+    {
+      "epoch": 0.15,
+      "gumbel_tau": 0.7048,
+      "image_masks_loss": 0.0229,
+      "llm_learning_rate": 4.807270650433821e-06,
+      "loss": 0.7425,
+      "output_text_mask_loss": 0.0322,
+      "predictor_learning_rate": 0.0004807270650433821,
+      "step": 1484
+    },
+    {
+      "epoch": 0.15,
+      "gumbel_tau": 0.7046,
+      "image_masks_loss": 0.0151,
+      "llm_learning_rate": 4.806951104424844e-06,
+      "loss": 0.7567,
+      "output_text_mask_loss": 0.0145,
+      "predictor_learning_rate": 0.0004806951104424844,
+      "step": 1485
+    },
+    {
+      "epoch": 0.15,
+      "gumbel_tau": 0.7044,
+      "image_masks_loss": 0.0401,
+      "llm_learning_rate": 4.806631304369274e-06,
+      "loss": 0.8015,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.00048066313043692736,
+      "step": 1486
+    },
+    {
+      "epoch": 0.15,
+      "gumbel_tau": 0.7043,
+      "image_masks_loss": 0.0557,
+      "llm_learning_rate": 4.806311250302326e-06,
+      "loss": 0.7426,
+      "output_text_mask_loss": 0.0488,
+      "predictor_learning_rate": 0.0004806311250302326,
+      "step": 1487
+    },
+    {
+      "epoch": 0.15,
+      "gumbel_tau": 0.7041,
+      "image_masks_loss": 0.0228,
+      "llm_learning_rate": 4.805990942259247e-06,
+      "loss": 0.8078,
+      "output_text_mask_loss": 0.037,
+      "predictor_learning_rate": 0.0004805990942259246,
+      "step": 1488
+    },
+    {
+      "epoch": 0.15,
+      "gumbel_tau": 0.7039,
+      "image_masks_loss": 0.0259,
+      "llm_learning_rate": 4.805670380275308e-06,
+      "loss": 0.7633,
+      "output_text_mask_loss": 0.0137,
+      "predictor_learning_rate": 0.00048056703802753077,
+      "step": 1489
+    },
+    {
+      "epoch": 0.15,
+      "gumbel_tau": 0.7038,
+      "image_masks_loss": 0.0149,
+      "llm_learning_rate": 4.805349564385812e-06,
+      "loss": 0.7637,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.00048053495643858113,
+      "step": 1490
+    },
+    {
+      "epoch": 0.15,
+      "gumbel_tau": 0.7036,
+      "image_masks_loss": 0.0215,
+      "llm_learning_rate": 4.805028494626086e-06,
+      "loss": 0.7348,
+      "output_text_mask_loss": 0.0751,
+      "predictor_learning_rate": 0.0004805028494626086,
+      "step": 1491
+    },
+    {
+      "epoch": 0.15,
+      "gumbel_tau": 0.7035,
+      "image_masks_loss": 0.0048,
+      "llm_learning_rate": 4.804707171031488e-06,
+      "loss": 0.6735,
+      "output_text_mask_loss": NaN,
+      "predictor_learning_rate": 0.0004804707171031488,
+      "step": 1492
+    },
+    {
+      "epoch": 0.15,
+      "gumbel_tau": 0.7033,
+      "image_masks_loss": 0.0683,
+      "llm_learning_rate": 4.804385593637403e-06,
+      "loss": 0.8722,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.00048043855936374025,
+      "step": 1493
+    },
+    {
+      "epoch": 0.15,
+      "gumbel_tau": 0.7031,
+      "image_masks_loss": 0.0156,
+      "llm_learning_rate": 4.804063762479244e-06,
+      "loss": 0.9129,
+      "output_text_mask_loss": 0.0324,
+      "predictor_learning_rate": 0.00048040637624792437,
+      "step": 1494
+    },
+    {
+      "epoch": 0.15,
+      "gumbel_tau": 0.703,
+      "image_masks_loss": 0.0127,
+      "llm_learning_rate": 4.803741677592451e-06,
+      "loss": 0.7704,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00048037416775924507,
+      "step": 1495
+    },
+    {
+      "epoch": 0.15,
+      "gumbel_tau": 0.7028,
+      "image_masks_loss": 0.0123,
+      "llm_learning_rate": 4.803419339012493e-06,
+      "loss": 0.8467,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 0.0004803419339012493,
+      "step": 1496
+    },
+    {
+      "epoch": 0.15,
+      "gumbel_tau": 0.7026,
+      "image_masks_loss": 0.0149,
+      "llm_learning_rate": 4.803096746774867e-06,
+      "loss": 0.6501,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004803096746774867,
+      "step": 1497
+    },
+    {
+      "epoch": 0.15,
+      "gumbel_tau": 0.7025,
+      "image_masks_loss": 0.0159,
+      "llm_learning_rate": 4.802773900915098e-06,
+      "loss": 0.7526,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00048027739009150975,
+      "step": 1498
+    },
+    {
+      "epoch": 0.15,
+      "gumbel_tau": 0.7023,
+      "image_masks_loss": 0.0576,
+      "llm_learning_rate": 4.802450801468737e-06,
+      "loss": 0.838,
+      "output_text_mask_loss": 0.0231,
+      "predictor_learning_rate": 0.00048024508014687364,
+      "step": 1499
+    },
+    {
+      "epoch": 0.15,
+      "gumbel_tau": 0.7021,
+      "image_masks_loss": 0.0127,
+      "llm_learning_rate": 4.802127448471364e-06,
+      "loss": 0.8477,
+      "output_text_mask_loss": 0.0174,
+      "predictor_learning_rate": 0.0004802127448471364,
+      "step": 1500
+    },
+    {
+      "epoch": 0.15,
+      "gumbel_tau": 0.702,
+      "image_masks_loss": 0.0269,
+      "llm_learning_rate": 4.801803841958591e-06,
+      "loss": 0.7997,
+      "output_text_mask_loss": 0.013,
+      "predictor_learning_rate": 0.000480180384195859,
+      "step": 1501
+    },
+    {
+      "epoch": 0.15,
+      "gumbel_tau": 0.7018,
+      "image_masks_loss": 0.0324,
+      "llm_learning_rate": 4.80147998196605e-06,
+      "loss": 0.7727,
+      "output_text_mask_loss": 0.0065,
+      "predictor_learning_rate": 0.00048014799819660496,
+      "step": 1502
+    },
+    {
+      "epoch": 0.15,
+      "gumbel_tau": 0.7016,
+      "image_masks_loss": 0.0406,
+      "llm_learning_rate": 4.801155868529408e-06,
+      "loss": 0.7052,
+      "output_text_mask_loss": 0.0012,
+      "predictor_learning_rate": 0.0004801155868529408,
+      "step": 1503
+    },
+    {
+      "epoch": 0.15,
+      "gumbel_tau": 0.7015,
+      "image_masks_loss": 0.045,
+      "llm_learning_rate": 4.800831501684356e-06,
+      "loss": 0.8,
+      "output_text_mask_loss": 0.1068,
+      "predictor_learning_rate": 0.00048008315016843554,
+      "step": 1504
+    },
+    {
+      "epoch": 0.15,
+      "gumbel_tau": 0.7013,
+      "image_masks_loss": 0.0614,
+      "llm_learning_rate": 4.800506881466613e-06,
+      "loss": 0.7299,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00048005068814666133,
+      "step": 1505
+    },
+    {
+      "epoch": 0.15,
+      "gumbel_tau": 0.7011,
+      "image_masks_loss": 0.0221,
+      "llm_learning_rate": 4.80018200791193e-06,
+      "loss": 0.7868,
+      "output_text_mask_loss": 0.0076,
+      "predictor_learning_rate": 0.00048001820079119294,
+      "step": 1506
+    },
+    {
+      "epoch": 0.15,
+      "gumbel_tau": 0.701,
+      "image_masks_loss": 0.0246,
+      "llm_learning_rate": 4.79985688105608e-06,
+      "loss": 0.7667,
+      "output_text_mask_loss": 0.0641,
+      "predictor_learning_rate": 0.00047998568810560797,
+      "step": 1507
+    },
+    {
+      "epoch": 0.15,
+      "gumbel_tau": 0.7008,
+      "image_masks_loss": 0.0317,
+      "llm_learning_rate": 4.799531500934867e-06,
+      "loss": 0.7253,
+      "output_text_mask_loss": 0.103,
+      "predictor_learning_rate": 0.00047995315009348666,
+      "step": 1508
+    },
+    {
+      "epoch": 0.15,
+      "gumbel_tau": 0.7006,
+      "image_masks_loss": 0.0147,
+      "llm_learning_rate": 4.799205867584123e-06,
+      "loss": 0.703,
+      "output_text_mask_loss": 0.0099,
+      "predictor_learning_rate": 0.0004799205867584123,
+      "step": 1509
+    },
+    {
+      "epoch": 0.15,
+      "gumbel_tau": 0.7005,
+      "image_masks_loss": 0.008,
+      "llm_learning_rate": 4.798879981039709e-06,
+      "loss": 0.7316,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.0004798879981039708,
+      "step": 1510
+    },
+    {
+      "epoch": 0.15,
+      "gumbel_tau": 0.7003,
+      "image_masks_loss": 0.0265,
+      "llm_learning_rate": 4.79855384133751e-06,
+      "loss": 0.7883,
+      "output_text_mask_loss": 0.0488,
+      "predictor_learning_rate": 0.00047985538413375096,
+      "step": 1511
+    },
+    {
+      "epoch": 0.15,
+      "gumbel_tau": 0.7001,
+      "image_masks_loss": 0.0338,
+      "llm_learning_rate": 4.798227448513441e-06,
+      "loss": 0.7343,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.00047982274485134415,
+      "step": 1512
+    },
+    {
+      "epoch": 0.16,
+      "gumbel_tau": 0.7,
+      "image_masks_loss": 0.0148,
+      "llm_learning_rate": 4.797900802603448e-06,
+      "loss": 0.8544,
+      "output_text_mask_loss": 0.0374,
+      "predictor_learning_rate": 0.00047979008026034477,
+      "step": 1513
+    },
+    {
+      "epoch": 0.16,
+      "gumbel_tau": 0.6998,
+      "image_masks_loss": 0.0062,
+      "llm_learning_rate": 4.7975739036435e-06,
+      "loss": 0.72,
+      "output_text_mask_loss": 0.0206,
+      "predictor_learning_rate": 0.00047975739036434993,
+      "step": 1514
+    },
+    {
+      "epoch": 0.16,
+      "gumbel_tau": 0.6996,
+      "image_masks_loss": 0.0121,
+      "llm_learning_rate": 4.797246751669595e-06,
+      "loss": 0.85,
+      "output_text_mask_loss": 0.0099,
+      "predictor_learning_rate": 0.0004797246751669595,
+      "step": 1515
+    },
+    {
+      "epoch": 0.16,
+      "gumbel_tau": 0.6995,
+      "image_masks_loss": 0.0177,
+      "llm_learning_rate": 4.7969193467177624e-06,
+      "loss": 0.8039,
+      "output_text_mask_loss": 0.0019,
+      "predictor_learning_rate": 0.0004796919346717762,
+      "step": 1516
+    },
+    {
+      "epoch": 0.16,
+      "gumbel_tau": 0.6993,
+      "image_masks_loss": 0.0572,
+      "llm_learning_rate": 4.796591688824054e-06,
+      "loss": 0.7549,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00047965916888240536,
+      "step": 1517
+    },
+    {
+      "epoch": 0.16,
+      "gumbel_tau": 0.6991,
+      "image_masks_loss": 0.0175,
+      "llm_learning_rate": 4.796263778024554e-06,
+      "loss": 0.8026,
+      "output_text_mask_loss": 0.0032,
+      "predictor_learning_rate": 0.0004796263778024553,
+      "step": 1518
+    },
+    {
+      "epoch": 0.16,
+      "gumbel_tau": 0.699,
+      "image_masks_loss": 0.0082,
+      "llm_learning_rate": 4.795935614355372e-06,
+      "loss": 0.7667,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.0004795935614355371,
+      "step": 1519
+    },
+    {
+      "epoch": 0.16,
+      "gumbel_tau": 0.6988,
+      "image_masks_loss": 0.0317,
+      "llm_learning_rate": 4.7956071978526454e-06,
+      "loss": 0.8822,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 0.0004795607197852645,
+      "step": 1520
+    },
+    {
+      "epoch": 0.16,
+      "gumbel_tau": 0.6987,
+      "image_masks_loss": 0.0166,
+      "llm_learning_rate": 4.795278528552541e-06,
+      "loss": 0.7703,
+      "output_text_mask_loss": 0.0058,
+      "predictor_learning_rate": 0.00047952785285525413,
+      "step": 1521
+    },
+    {
+      "epoch": 0.16,
+      "gumbel_tau": 0.6985,
+      "image_masks_loss": 0.0324,
+      "llm_learning_rate": 4.794949606491253e-06,
+      "loss": 0.7466,
+      "output_text_mask_loss": 0.0177,
+      "predictor_learning_rate": 0.0004794949606491253,
+      "step": 1522
+    },
+    {
+      "epoch": 0.16,
+      "gumbel_tau": 0.6983,
+      "image_masks_loss": 0.0099,
+      "llm_learning_rate": 4.794620431705002e-06,
+      "loss": 0.6743,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.0004794620431705002,
+      "step": 1523
+    },
+    {
+      "epoch": 0.16,
+      "gumbel_tau": 0.6982,
+      "image_masks_loss": 0.0066,
+      "llm_learning_rate": 4.79429100423004e-06,
+      "loss": 0.7663,
+      "output_text_mask_loss": 0.0223,
+      "predictor_learning_rate": 0.0004794291004230039,
+      "step": 1524
+    },
+    {
+      "epoch": 0.16,
+      "gumbel_tau": 0.698,
+      "image_masks_loss": 0.0162,
+      "llm_learning_rate": 4.793961324102639e-06,
+      "loss": 0.751,
+      "output_text_mask_loss": 0.0065,
+      "predictor_learning_rate": 0.0004793961324102639,
+      "step": 1525
+    },
+    {
+      "epoch": 0.16,
+      "gumbel_tau": 0.6978,
+      "image_masks_loss": 0.0199,
+      "llm_learning_rate": 4.79363139135911e-06,
+      "loss": 0.7486,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.000479363139135911,
+      "step": 1526
+    },
+    {
+      "epoch": 0.16,
+      "gumbel_tau": 0.6977,
+      "image_masks_loss": 0.0182,
+      "llm_learning_rate": 4.793301206035782e-06,
+      "loss": 0.7713,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 0.00047933012060357816,
+      "step": 1527
+    },
+    {
+      "epoch": 0.16,
+      "gumbel_tau": 0.6975,
+      "image_masks_loss": 0.0248,
+      "llm_learning_rate": 4.792970768169017e-06,
+      "loss": 0.7448,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004792970768169017,
+      "step": 1528
+    },
+    {
+      "epoch": 0.16,
+      "gumbel_tau": 0.6973,
+      "image_masks_loss": 0.0136,
+      "llm_learning_rate": 4.792640077795205e-06,
+      "loss": 0.7009,
+      "output_text_mask_loss": 0.0282,
+      "predictor_learning_rate": 0.00047926400777952043,
+      "step": 1529
+    },
+    {
+      "epoch": 0.16,
+      "gumbel_tau": 0.6972,
+      "image_masks_loss": 0.0317,
+      "llm_learning_rate": 4.79230913495076e-06,
+      "loss": 0.7558,
+      "output_text_mask_loss": 0.0725,
+      "predictor_learning_rate": 0.00047923091349507595,
+      "step": 1530
+    },
+    {
+      "epoch": 0.16,
+      "gumbel_tau": 0.697,
+      "image_masks_loss": 0.0107,
+      "llm_learning_rate": 4.791977939672127e-06,
+      "loss": 0.6596,
+      "output_text_mask_loss": 0.0012,
+      "predictor_learning_rate": 0.00047919779396721267,
+      "step": 1531
+    },
+    {
+      "epoch": 0.16,
+      "gumbel_tau": 0.6968,
+      "image_masks_loss": 0.0136,
+      "llm_learning_rate": 4.791646491995778e-06,
+      "loss": 0.8003,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004791646491995778,
+      "step": 1532
+    },
+    {
+      "epoch": 0.16,
+      "gumbel_tau": 0.6967,
+      "image_masks_loss": 0.0176,
+      "llm_learning_rate": 4.791314791958214e-06,
+      "loss": 0.7579,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00047913147919582136,
+      "step": 1533
+    },
+    {
+      "epoch": 0.16,
+      "gumbel_tau": 0.6965,
+      "image_masks_loss": 0.021,
+      "llm_learning_rate": 4.79098283959596e-06,
+      "loss": 0.7605,
+      "output_text_mask_loss": 0.0004,
+      "predictor_learning_rate": 0.000479098283959596,
+      "step": 1534
+    },
+    {
+      "epoch": 0.16,
+      "gumbel_tau": 0.6964,
+      "image_masks_loss": 0.0188,
+      "llm_learning_rate": 4.790650634945573e-06,
+      "loss": 0.7379,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004790650634945573,
+      "step": 1535
+    },
+    {
+      "epoch": 0.16,
+      "gumbel_tau": 0.6962,
+      "image_masks_loss": 0.0303,
+      "llm_learning_rate": 4.790318178043637e-06,
+      "loss": 0.8678,
+      "output_text_mask_loss": 0.0139,
+      "predictor_learning_rate": 0.00047903181780436366,
+      "step": 1536
+    },
+    {
+      "epoch": 0.16,
+      "gumbel_tau": 0.696,
+      "image_masks_loss": 0.0254,
+      "llm_learning_rate": 4.789985468926761e-06,
+      "loss": 0.7688,
+      "output_text_mask_loss": 0.0393,
+      "predictor_learning_rate": 0.000478998546892676,
+      "step": 1537
+    },
+    {
+      "epoch": 0.16,
+      "gumbel_tau": 0.6959,
+      "image_masks_loss": 0.0145,
+      "llm_learning_rate": 4.789652507631584e-06,
+      "loss": 0.8082,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.0004789652507631583,
+      "step": 1538
+    },
+    {
+      "epoch": 0.16,
+      "gumbel_tau": 0.6957,
+      "image_masks_loss": 0.0179,
+      "llm_learning_rate": 4.789319294194773e-06,
+      "loss": 0.7718,
+      "output_text_mask_loss": 0.0231,
+      "predictor_learning_rate": 0.0004789319294194773,
+      "step": 1539
+    },
+    {
+      "epoch": 0.16,
+      "gumbel_tau": 0.6955,
+      "image_masks_loss": 0.0202,
+      "llm_learning_rate": 4.788985828653022e-06,
+      "loss": 0.7899,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004788985828653022,
+      "step": 1540
+    },
+    {
+      "epoch": 0.16,
+      "gumbel_tau": 0.6954,
+      "image_masks_loss": 0.0169,
+      "llm_learning_rate": 4.788652111043054e-06,
+      "loss": 0.8559,
+      "output_text_mask_loss": 0.0269,
+      "predictor_learning_rate": 0.00047886521110430535,
+      "step": 1541
+    },
+    {
+      "epoch": 0.16,
+      "gumbel_tau": 0.6952,
+      "image_masks_loss": 0.0093,
+      "llm_learning_rate": 4.7883181414016165e-06,
+      "loss": 0.9246,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.00047883181414016163,
+      "step": 1542
+    },
+    {
+      "epoch": 0.16,
+      "gumbel_tau": 0.695,
+      "image_masks_loss": 0.0246,
+      "llm_learning_rate": 4.787983919765489e-06,
+      "loss": 0.7357,
+      "output_text_mask_loss": 0.046,
+      "predictor_learning_rate": 0.0004787983919765488,
+      "step": 1543
+    },
+    {
+      "epoch": 0.16,
+      "gumbel_tau": 0.6949,
+      "image_masks_loss": 0.0399,
+      "llm_learning_rate": 4.787649446171476e-06,
+      "loss": 0.8206,
+      "output_text_mask_loss": 0.036,
+      "predictor_learning_rate": 0.00047876494461714754,
+      "step": 1544
+    },
+    {
+      "epoch": 0.16,
+      "gumbel_tau": 0.6947,
+      "image_masks_loss": 0.0154,
+      "llm_learning_rate": 4.78731472065641e-06,
+      "loss": 0.8151,
+      "output_text_mask_loss": 0.0565,
+      "predictor_learning_rate": 0.00047873147206564095,
+      "step": 1545
+    },
+    {
+      "epoch": 0.16,
+      "gumbel_tau": 0.6945,
+      "image_masks_loss": 0.0194,
+      "llm_learning_rate": 4.786979743257152e-06,
+      "loss": 0.7578,
+      "output_text_mask_loss": 0.0111,
+      "predictor_learning_rate": 0.0004786979743257152,
+      "step": 1546
+    },
+    {
+      "epoch": 0.16,
+      "gumbel_tau": 0.6944,
+      "image_masks_loss": 0.0622,
+      "llm_learning_rate": 4.786644514010591e-06,
+      "loss": 0.7905,
+      "output_text_mask_loss": 0.0075,
+      "predictor_learning_rate": 0.00047866445140105907,
+      "step": 1547
+    },
+    {
+      "epoch": 0.16,
+      "gumbel_tau": 0.6942,
+      "image_masks_loss": 0.0341,
+      "llm_learning_rate": 4.786309032953643e-06,
+      "loss": 0.7892,
+      "output_text_mask_loss": 0.0376,
+      "predictor_learning_rate": 0.0004786309032953643,
+      "step": 1548
+    },
+    {
+      "epoch": 0.16,
+      "gumbel_tau": 0.6941,
+      "image_masks_loss": 0.01,
+      "llm_learning_rate": 4.7859733001232515e-06,
+      "loss": 0.7422,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 0.0004785973300123251,
+      "step": 1549
+    },
+    {
+      "epoch": 0.16,
+      "gumbel_tau": 0.6939,
+      "image_masks_loss": 0.0246,
+      "llm_learning_rate": 4.785637315556388e-06,
+      "loss": 0.8009,
+      "output_text_mask_loss": 0.061,
+      "predictor_learning_rate": 0.0004785637315556388,
+      "step": 1550
+    },
+    {
+      "epoch": 0.16,
+      "gumbel_tau": 0.6937,
+      "image_masks_loss": 0.028,
+      "llm_learning_rate": 4.785301079290054e-06,
+      "loss": 0.7121,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.0004785301079290053,
+      "step": 1551
+    },
+    {
+      "epoch": 0.16,
+      "gumbel_tau": 0.6936,
+      "image_masks_loss": 0.0162,
+      "llm_learning_rate": 4.784964591361273e-06,
+      "loss": 0.749,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00047849645913612725,
+      "step": 1552
+    },
+    {
+      "epoch": 0.16,
+      "gumbel_tau": 0.6934,
+      "image_masks_loss": 0.0204,
+      "llm_learning_rate": 4.784627851807102e-06,
+      "loss": 0.7886,
+      "output_text_mask_loss": 0.0109,
+      "predictor_learning_rate": 0.0004784627851807102,
+      "step": 1553
+    },
+    {
+      "epoch": 0.16,
+      "gumbel_tau": 0.6932,
+      "image_masks_loss": 0.0141,
+      "llm_learning_rate": 4.784290860664623e-06,
+      "loss": 0.7412,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.0004784290860664623,
+      "step": 1554
+    },
+    {
+      "epoch": 0.16,
+      "gumbel_tau": 0.6931,
+      "image_masks_loss": 0.0076,
+      "llm_learning_rate": 4.783953617970946e-06,
+      "loss": 0.6318,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004783953617970946,
+      "step": 1555
+    },
+    {
+      "epoch": 0.16,
+      "gumbel_tau": 0.6929,
+      "image_masks_loss": 0.0252,
+      "llm_learning_rate": 4.783616123763211e-06,
+      "loss": 0.7918,
+      "output_text_mask_loss": 0.0158,
+      "predictor_learning_rate": 0.00047836161237632105,
+      "step": 1556
+    },
+    {
+      "epoch": 0.16,
+      "gumbel_tau": 0.6927,
+      "image_masks_loss": 0.008,
+      "llm_learning_rate": 4.783278378078581e-06,
+      "loss": 0.8009,
+      "output_text_mask_loss": 0.0345,
+      "predictor_learning_rate": 0.00047832783780785805,
+      "step": 1557
+    },
+    {
+      "epoch": 0.16,
+      "gumbel_tau": 0.6926,
+      "image_masks_loss": 0.0119,
+      "llm_learning_rate": 4.7829403809542495e-06,
+      "loss": 0.7456,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.0004782940380954249,
+      "step": 1558
+    },
+    {
+      "epoch": 0.16,
+      "gumbel_tau": 0.6924,
+      "image_masks_loss": 0.0507,
+      "llm_learning_rate": 4.782602132427439e-06,
+      "loss": 0.8659,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 0.0004782602132427439,
+      "step": 1559
+    },
+    {
+      "epoch": 0.16,
+      "gumbel_tau": 0.6923,
+      "image_masks_loss": 0.0317,
+      "llm_learning_rate": 4.782263632535397e-06,
+      "loss": 0.8388,
+      "output_text_mask_loss": 0.0448,
+      "predictor_learning_rate": 0.0004782263632535397,
+      "step": 1560
+    },
+    {
+      "epoch": 0.16,
+      "gumbel_tau": 0.6921,
+      "image_masks_loss": 0.0036,
+      "llm_learning_rate": 4.781924881315401e-06,
+      "loss": 0.751,
+      "output_text_mask_loss": 0.002,
+      "predictor_learning_rate": 0.00047819248813154,
+      "step": 1561
+    },
+    {
+      "epoch": 0.16,
+      "gumbel_tau": 0.6919,
+      "image_masks_loss": 0.0408,
+      "llm_learning_rate": 4.781585878804753e-06,
+      "loss": 0.7989,
+      "output_text_mask_loss": 0.0277,
+      "predictor_learning_rate": 0.00047815858788047527,
+      "step": 1562
+    },
+    {
+      "epoch": 0.16,
+      "gumbel_tau": 0.6918,
+      "image_masks_loss": 0.0148,
+      "llm_learning_rate": 4.7812466250407864e-06,
+      "loss": 0.7286,
+      "output_text_mask_loss": 0.0687,
+      "predictor_learning_rate": 0.00047812466250407863,
+      "step": 1563
+    },
+    {
+      "epoch": 0.16,
+      "gumbel_tau": 0.6916,
+      "image_masks_loss": 0.0184,
+      "llm_learning_rate": 4.780907120060861e-06,
+      "loss": 0.7166,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.00047809071200608606,
+      "step": 1564
+    },
+    {
+      "epoch": 0.16,
+      "gumbel_tau": 0.6914,
+      "image_masks_loss": 0.0395,
+      "llm_learning_rate": 4.780567363902363e-06,
+      "loss": 0.7465,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.0004780567363902362,
+      "step": 1565
+    },
+    {
+      "epoch": 0.16,
+      "gumbel_tau": 0.6913,
+      "image_masks_loss": 0.0153,
+      "llm_learning_rate": 4.7802273566027056e-06,
+      "loss": 0.8155,
+      "output_text_mask_loss": 0.0271,
+      "predictor_learning_rate": 0.00047802273566027053,
+      "step": 1566
+    },
+    {
+      "epoch": 0.16,
+      "gumbel_tau": 0.6911,
+      "image_masks_loss": 0.0229,
+      "llm_learning_rate": 4.779887098199334e-06,
+      "loss": 0.7923,
+      "output_text_mask_loss": 0.0303,
+      "predictor_learning_rate": 0.00047798870981993335,
+      "step": 1567
+    },
+    {
+      "epoch": 0.16,
+      "gumbel_tau": 0.691,
+      "image_masks_loss": 0.0101,
+      "llm_learning_rate": 4.779546588729717e-06,
+      "loss": 0.7417,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 0.00047795465887297164,
+      "step": 1568
+    },
+    {
+      "epoch": 0.16,
+      "gumbel_tau": 0.6908,
+      "image_masks_loss": 0.0016,
+      "llm_learning_rate": 4.779205828231351e-06,
+      "loss": 0.67,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 0.0004779205828231351,
+      "step": 1569
+    },
+    {
+      "epoch": 0.16,
+      "gumbel_tau": 0.6906,
+      "image_masks_loss": 0.017,
+      "llm_learning_rate": 4.778864816741763e-06,
+      "loss": 0.7771,
+      "output_text_mask_loss": 0.032,
+      "predictor_learning_rate": 0.0004778864816741763,
+      "step": 1570
+    },
+    {
+      "epoch": 0.16,
+      "gumbel_tau": 0.6905,
+      "image_masks_loss": 0.0267,
+      "llm_learning_rate": 4.778523554298506e-06,
+      "loss": 0.7799,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.00047785235542985055,
+      "step": 1571
+    },
+    {
+      "epoch": 0.16,
+      "gumbel_tau": 0.6903,
+      "image_masks_loss": 0.0572,
+      "llm_learning_rate": 4.77818204093916e-06,
+      "loss": 0.9395,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.00047781820409391595,
+      "step": 1572
+    },
+    {
+      "epoch": 0.16,
+      "gumbel_tau": 0.6901,
+      "image_masks_loss": 0.0074,
+      "llm_learning_rate": 4.7778402767013325e-06,
+      "loss": 0.7606,
+      "output_text_mask_loss": 0.1175,
+      "predictor_learning_rate": 0.0004777840276701332,
+      "step": 1573
+    },
+    {
+      "epoch": 0.16,
+      "gumbel_tau": 0.69,
+      "image_masks_loss": 0.0061,
+      "llm_learning_rate": 4.77749826162266e-06,
+      "loss": 0.6807,
+      "output_text_mask_loss": 0.0614,
+      "predictor_learning_rate": 0.00047774982616226595,
+      "step": 1574
+    },
+    {
+      "epoch": 0.16,
+      "gumbel_tau": 0.6898,
+      "image_masks_loss": 0.0132,
+      "llm_learning_rate": 4.777155995740807e-06,
+      "loss": 0.6939,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 0.0004777155995740806,
+      "step": 1575
+    },
+    {
+      "epoch": 0.16,
+      "gumbel_tau": 0.6896,
+      "image_masks_loss": 0.0072,
+      "llm_learning_rate": 4.7768134790934625e-06,
+      "loss": 0.7412,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.00047768134790934617,
+      "step": 1576
+    },
+    {
+      "epoch": 0.16,
+      "gumbel_tau": 0.6895,
+      "image_masks_loss": 0.0526,
+      "llm_learning_rate": 4.776470711718346e-06,
+      "loss": 0.7639,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004776470711718346,
+      "step": 1577
+    },
+    {
+      "epoch": 0.16,
+      "gumbel_tau": 0.6893,
+      "image_masks_loss": 0.0519,
+      "llm_learning_rate": 4.776127693653205e-06,
+      "loss": 0.7844,
+      "output_text_mask_loss": 0.0618,
+      "predictor_learning_rate": 0.0004776127693653205,
+      "step": 1578
+    },
+    {
+      "epoch": 0.16,
+      "gumbel_tau": 0.6892,
+      "image_masks_loss": 0.0102,
+      "llm_learning_rate": 4.7757844249358125e-06,
+      "loss": 0.7188,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00047757844249358126,
+      "step": 1579
+    },
+    {
+      "epoch": 0.16,
+      "gumbel_tau": 0.689,
+      "image_masks_loss": 0.0144,
+      "llm_learning_rate": 4.77544090560397e-06,
+      "loss": 0.7439,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00047754409056039696,
+      "step": 1580
+    },
+    {
+      "epoch": 0.16,
+      "gumbel_tau": 0.6888,
+      "image_masks_loss": 0.0217,
+      "llm_learning_rate": 4.775097135695506e-06,
+      "loss": 0.8125,
+      "output_text_mask_loss": 0.0172,
+      "predictor_learning_rate": 0.0004775097135695506,
+      "step": 1581
+    },
+    {
+      "epoch": 0.16,
+      "gumbel_tau": 0.6887,
+      "image_masks_loss": 0.0292,
+      "llm_learning_rate": 4.774753115248279e-06,
+      "loss": 0.6976,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.0004774753115248278,
+      "step": 1582
+    },
+    {
+      "epoch": 0.16,
+      "gumbel_tau": 0.6885,
+      "image_masks_loss": 0.0121,
+      "llm_learning_rate": 4.774408844300171e-06,
+      "loss": 0.7329,
+      "output_text_mask_loss": 0.0504,
+      "predictor_learning_rate": 0.0004774408844300171,
+      "step": 1583
+    },
+    {
+      "epoch": 0.16,
+      "gumbel_tau": 0.6883,
+      "image_masks_loss": 0.0286,
+      "llm_learning_rate": 4.774064322889095e-06,
+      "loss": 0.7443,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.00047740643228890946,
+      "step": 1584
+    },
+    {
+      "epoch": 0.16,
+      "gumbel_tau": 0.6882,
+      "image_masks_loss": 0.0168,
+      "llm_learning_rate": 4.773719551052991e-06,
+      "loss": 0.7641,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004773719551052991,
+      "step": 1585
+    },
+    {
+      "epoch": 0.16,
+      "gumbel_tau": 0.688,
+      "image_masks_loss": 0.0381,
+      "llm_learning_rate": 4.773374528829825e-06,
+      "loss": 0.8942,
+      "output_text_mask_loss": 0.0187,
+      "predictor_learning_rate": 0.0004773374528829825,
+      "step": 1586
+    },
+    {
+      "epoch": 0.16,
+      "gumbel_tau": 0.6879,
+      "image_masks_loss": 0.066,
+      "llm_learning_rate": 4.773029256257593e-06,
+      "loss": 0.8016,
+      "output_text_mask_loss": 0.0058,
+      "predictor_learning_rate": 0.0004773029256257593,
+      "step": 1587
+    },
+    {
+      "epoch": 0.16,
+      "gumbel_tau": 0.6877,
+      "image_masks_loss": 0.02,
+      "llm_learning_rate": 4.772683733374316e-06,
+      "loss": 0.8736,
+      "output_text_mask_loss": 0.0019,
+      "predictor_learning_rate": 0.00047726837333743154,
+      "step": 1588
+    },
+    {
+      "epoch": 0.16,
+      "gumbel_tau": 0.6875,
+      "image_masks_loss": 0.0093,
+      "llm_learning_rate": 4.772337960218045e-06,
+      "loss": 0.6937,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 0.0004772337960218044,
+      "step": 1589
+    },
+    {
+      "epoch": 0.16,
+      "gumbel_tau": 0.6874,
+      "image_masks_loss": 0.021,
+      "llm_learning_rate": 4.771991936826854e-06,
+      "loss": 0.7277,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004771991936826854,
+      "step": 1590
+    },
+    {
+      "epoch": 0.16,
+      "gumbel_tau": 0.6872,
+      "image_masks_loss": 0.0169,
+      "llm_learning_rate": 4.771645663238852e-06,
+      "loss": 0.7049,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00047716456632388517,
+      "step": 1591
+    },
+    {
+      "epoch": 0.16,
+      "gumbel_tau": 0.687,
+      "image_masks_loss": 0.015,
+      "llm_learning_rate": 4.771299139492169e-06,
+      "loss": 0.7908,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.0004771299139492169,
+      "step": 1592
+    },
+    {
+      "epoch": 0.16,
+      "gumbel_tau": 0.6869,
+      "image_masks_loss": 0.0418,
+      "llm_learning_rate": 4.770952365624966e-06,
+      "loss": 0.7863,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.00047709523656249656,
+      "step": 1593
+    },
+    {
+      "epoch": 0.16,
+      "gumbel_tau": 0.6867,
+      "image_masks_loss": 0.0214,
+      "llm_learning_rate": 4.77060534167543e-06,
+      "loss": 0.7864,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.00047706053416754295,
+      "step": 1594
+    },
+    {
+      "epoch": 0.16,
+      "gumbel_tau": 0.6866,
+      "image_masks_loss": 0.0161,
+      "llm_learning_rate": 4.770258067681776e-06,
+      "loss": 0.8341,
+      "output_text_mask_loss": 0.0004,
+      "predictor_learning_rate": 0.0004770258067681776,
+      "step": 1595
+    },
+    {
+      "epoch": 0.16,
+      "gumbel_tau": 0.6864,
+      "image_masks_loss": 0.0105,
+      "llm_learning_rate": 4.769910543682247e-06,
+      "loss": 0.8387,
+      "output_text_mask_loss": 0.0338,
+      "predictor_learning_rate": 0.00047699105436822467,
+      "step": 1596
+    },
+    {
+      "epoch": 0.16,
+      "gumbel_tau": 0.6862,
+      "image_masks_loss": 0.0142,
+      "llm_learning_rate": 4.7695627697151135e-06,
+      "loss": 0.8048,
+      "output_text_mask_loss": 0.001,
+      "predictor_learning_rate": 0.00047695627697151127,
+      "step": 1597
+    },
+    {
+      "epoch": 0.16,
+      "gumbel_tau": 0.6861,
+      "image_masks_loss": 0.0129,
+      "llm_learning_rate": 4.769214745818671e-06,
+      "loss": 0.8272,
+      "output_text_mask_loss": 0.0286,
+      "predictor_learning_rate": 0.00047692147458186713,
+      "step": 1598
+    },
+    {
+      "epoch": 0.16,
+      "gumbel_tau": 0.6859,
+      "image_masks_loss": 0.0273,
+      "llm_learning_rate": 4.768866472031248e-06,
+      "loss": 0.7542,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 0.00047688664720312474,
+      "step": 1599
+    },
+    {
+      "epoch": 0.16,
+      "gumbel_tau": 0.6858,
+      "image_masks_loss": 0.0216,
+      "llm_learning_rate": 4.768517948391195e-06,
+      "loss": 0.7079,
+      "output_text_mask_loss": NaN,
+      "predictor_learning_rate": 0.00047685179483911947,
+      "step": 1600
+    },
+    {
+      "epoch": 0.16,
+      "gumbel_tau": 0.6856,
+      "image_masks_loss": 0.0219,
+      "llm_learning_rate": 4.768169174936892e-06,
+      "loss": 0.6301,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004768169174936892,
+      "step": 1601
+    },
+    {
+      "epoch": 0.16,
+      "gumbel_tau": 0.6854,
+      "image_masks_loss": 0.038,
+      "llm_learning_rate": 4.767820151706747e-06,
+      "loss": 0.8043,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 0.00047678201517067464,
+      "step": 1602
+    },
+    {
+      "epoch": 0.16,
+      "gumbel_tau": 0.6853,
+      "image_masks_loss": 0.0214,
+      "llm_learning_rate": 4.767470878739196e-06,
+      "loss": 0.7792,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 0.0004767470878739196,
+      "step": 1603
+    },
+    {
+      "epoch": 0.16,
+      "gumbel_tau": 0.6851,
+      "image_masks_loss": 0.019,
+      "llm_learning_rate": 4.767121356072701e-06,
+      "loss": 0.8583,
+      "output_text_mask_loss": 0.0492,
+      "predictor_learning_rate": 0.0004767121356072701,
+      "step": 1604
+    },
+    {
+      "epoch": 0.16,
+      "gumbel_tau": 0.6849,
+      "image_masks_loss": 0.0144,
+      "llm_learning_rate": 4.766771583745753e-06,
+      "loss": 0.7056,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004766771583745753,
+      "step": 1605
+    },
+    {
+      "epoch": 0.16,
+      "gumbel_tau": 0.6848,
+      "image_masks_loss": 0.0072,
+      "llm_learning_rate": 4.766421561796869e-06,
+      "loss": 0.832,
+      "output_text_mask_loss": 0.0353,
+      "predictor_learning_rate": 0.00047664215617968685,
+      "step": 1606
+    },
+    {
+      "epoch": 0.16,
+      "gumbel_tau": 0.6846,
+      "image_masks_loss": 0.013,
+      "llm_learning_rate": 4.766071290264594e-06,
+      "loss": 0.7845,
+      "output_text_mask_loss": 0.0429,
+      "predictor_learning_rate": 0.0004766071290264594,
+      "step": 1607
+    },
+    {
+      "epoch": 0.16,
+      "gumbel_tau": 0.6845,
+      "image_masks_loss": 0.0103,
+      "llm_learning_rate": 4.765720769187502e-06,
+      "loss": 0.8174,
+      "output_text_mask_loss": 0.0292,
+      "predictor_learning_rate": 0.00047657207691875015,
+      "step": 1608
+    },
+    {
+      "epoch": 0.16,
+      "gumbel_tau": 0.6843,
+      "image_masks_loss": 0.0214,
+      "llm_learning_rate": 4.765369998604191e-06,
+      "loss": 0.7769,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.0004765369998604191,
+      "step": 1609
+    },
+    {
+      "epoch": 0.16,
+      "gumbel_tau": 0.6841,
+      "image_masks_loss": 0.0324,
+      "llm_learning_rate": 4.7650189785532905e-06,
+      "loss": 0.7057,
+      "output_text_mask_loss": 0.0212,
+      "predictor_learning_rate": 0.000476501897855329,
+      "step": 1610
+    },
+    {
+      "epoch": 0.17,
+      "gumbel_tau": 0.684,
+      "image_masks_loss": 0.0345,
+      "llm_learning_rate": 4.764667709073455e-06,
+      "loss": 0.8576,
+      "output_text_mask_loss": 0.0362,
+      "predictor_learning_rate": 0.00047646677090734543,
+      "step": 1611
+    },
+    {
+      "epoch": 0.17,
+      "gumbel_tau": 0.6838,
+      "image_masks_loss": 0.0298,
+      "llm_learning_rate": 4.764316190203367e-06,
+      "loss": 0.682,
+      "output_text_mask_loss": NaN,
+      "predictor_learning_rate": 0.0004764316190203366,
+      "step": 1612
+    },
+    {
+      "epoch": 0.17,
+      "gumbel_tau": 0.6837,
+      "image_masks_loss": 0.0143,
+      "llm_learning_rate": 4.763964421981736e-06,
+      "loss": 0.8295,
+      "output_text_mask_loss": 0.0282,
+      "predictor_learning_rate": 0.0004763964421981735,
+      "step": 1613
+    },
+    {
+      "epoch": 0.17,
+      "gumbel_tau": 0.6835,
+      "image_masks_loss": 0.0471,
+      "llm_learning_rate": 4.7636124044473e-06,
+      "loss": 0.7043,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.00047636124044472995,
+      "step": 1614
+    },
+    {
+      "epoch": 0.17,
+      "gumbel_tau": 0.6833,
+      "image_masks_loss": 0.0181,
+      "llm_learning_rate": 4.763260137638824e-06,
+      "loss": 0.709,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004763260137638824,
+      "step": 1615
+    },
+    {
+      "epoch": 0.17,
+      "gumbel_tau": 0.6832,
+      "image_masks_loss": 0.009,
+      "llm_learning_rate": 4.762907621595102e-06,
+      "loss": 0.8487,
+      "output_text_mask_loss": 0.0039,
+      "predictor_learning_rate": 0.00047629076215951013,
+      "step": 1616
+    },
+    {
+      "epoch": 0.17,
+      "gumbel_tau": 0.683,
+      "image_masks_loss": 0.0473,
+      "llm_learning_rate": 4.7625548563549505e-06,
+      "loss": 0.8069,
+      "output_text_mask_loss": 0.0065,
+      "predictor_learning_rate": 0.00047625548563549506,
+      "step": 1617
+    },
+    {
+      "epoch": 0.17,
+      "gumbel_tau": 0.6828,
+      "image_masks_loss": 0.034,
+      "llm_learning_rate": 4.762201841957219e-06,
+      "loss": 0.7447,
+      "output_text_mask_loss": 0.0229,
+      "predictor_learning_rate": 0.0004762201841957219,
+      "step": 1618
+    },
+    {
+      "epoch": 0.17,
+      "gumbel_tau": 0.6827,
+      "image_masks_loss": 0.0402,
+      "llm_learning_rate": 4.761848578440782e-06,
+      "loss": 0.8158,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.00047618485784407816,
+      "step": 1619
+    },
+    {
+      "epoch": 0.17,
+      "gumbel_tau": 0.6825,
+      "image_masks_loss": 0.0614,
+      "llm_learning_rate": 4.761495065844542e-06,
+      "loss": 0.8409,
+      "output_text_mask_loss": 0.0099,
+      "predictor_learning_rate": 0.00047614950658445417,
+      "step": 1620
+    },
+    {
+      "epoch": 0.17,
+      "gumbel_tau": 0.6824,
+      "image_masks_loss": 0.015,
+      "llm_learning_rate": 4.761141304207427e-06,
+      "loss": 0.7527,
+      "output_text_mask_loss": 0.0117,
+      "predictor_learning_rate": 0.0004761141304207427,
+      "step": 1621
+    },
+    {
+      "epoch": 0.17,
+      "gumbel_tau": 0.6822,
+      "image_masks_loss": 0.0093,
+      "llm_learning_rate": 4.760787293568396e-06,
+      "loss": 0.7659,
+      "output_text_mask_loss": 0.0322,
+      "predictor_learning_rate": 0.0004760787293568396,
+      "step": 1622
+    },
+    {
+      "epoch": 0.17,
+      "gumbel_tau": 0.682,
+      "image_masks_loss": 0.0347,
+      "llm_learning_rate": 4.760433033966432e-06,
+      "loss": 0.727,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00047604330339664315,
+      "step": 1623
+    },
+    {
+      "epoch": 0.17,
+      "gumbel_tau": 0.6819,
+      "image_masks_loss": 0.0103,
+      "llm_learning_rate": 4.760078525440548e-06,
+      "loss": 0.7338,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 0.00047600785254405475,
+      "step": 1624
+    },
+    {
+      "epoch": 0.17,
+      "gumbel_tau": 0.6817,
+      "image_masks_loss": 0.032,
+      "llm_learning_rate": 4.759723768029782e-06,
+      "loss": 0.7161,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.00047597237680297816,
+      "step": 1625
+    },
+    {
+      "epoch": 0.17,
+      "gumbel_tau": 0.6816,
+      "image_masks_loss": 0.0107,
+      "llm_learning_rate": 4.7593687617732015e-06,
+      "loss": 0.8109,
+      "output_text_mask_loss": 0.0412,
+      "predictor_learning_rate": 0.00047593687617732014,
+      "step": 1626
+    },
+    {
+      "epoch": 0.17,
+      "gumbel_tau": 0.6814,
+      "image_masks_loss": 0.0151,
+      "llm_learning_rate": 4.7590135067099e-06,
+      "loss": 0.772,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 0.00047590135067099,
+      "step": 1627
+    },
+    {
+      "epoch": 0.17,
+      "gumbel_tau": 0.6812,
+      "image_masks_loss": 0.0345,
+      "llm_learning_rate": 4.758658002879001e-06,
+      "loss": 0.7602,
+      "output_text_mask_loss": 0.0004,
+      "predictor_learning_rate": 0.0004758658002879,
+      "step": 1628
+    },
+    {
+      "epoch": 0.17,
+      "gumbel_tau": 0.6811,
+      "image_masks_loss": 0.0456,
+      "llm_learning_rate": 4.75830225031965e-06,
+      "loss": 0.7826,
+      "output_text_mask_loss": 0.0284,
+      "predictor_learning_rate": 0.00047583022503196496,
+      "step": 1629
+    },
+    {
+      "epoch": 0.17,
+      "gumbel_tau": 0.6809,
+      "image_masks_loss": 0.09,
+      "llm_learning_rate": 4.7579462490710256e-06,
+      "loss": 0.8814,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.0004757946249071025,
+      "step": 1630
+    },
+    {
+      "epoch": 0.17,
+      "gumbel_tau": 0.6808,
+      "image_masks_loss": 0.0046,
+      "llm_learning_rate": 4.75758999917233e-06,
+      "loss": 0.7837,
+      "output_text_mask_loss": 0.0234,
+      "predictor_learning_rate": 0.00047575899991723306,
+      "step": 1631
+    },
+    {
+      "epoch": 0.17,
+      "gumbel_tau": 0.6806,
+      "image_masks_loss": 0.0223,
+      "llm_learning_rate": 4.7572335006627964e-06,
+      "loss": 0.8075,
+      "output_text_mask_loss": 0.0094,
+      "predictor_learning_rate": 0.0004757233500662797,
+      "step": 1632
+    },
+    {
+      "epoch": 0.17,
+      "gumbel_tau": 0.6804,
+      "image_masks_loss": 0.0471,
+      "llm_learning_rate": 4.756876753581682e-06,
+      "loss": 0.7848,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 0.0004756876753581681,
+      "step": 1633
+    },
+    {
+      "epoch": 0.17,
+      "gumbel_tau": 0.6803,
+      "image_masks_loss": 0.0387,
+      "llm_learning_rate": 4.756519757968272e-06,
+      "loss": 0.7819,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.0004756519757968272,
+      "step": 1634
+    },
+    {
+      "epoch": 0.17,
+      "gumbel_tau": 0.6801,
+      "image_masks_loss": 0.0259,
+      "llm_learning_rate": 4.75616251386188e-06,
+      "loss": 0.8697,
+      "output_text_mask_loss": 0.0416,
+      "predictor_learning_rate": 0.00047561625138618796,
+      "step": 1635
+    },
+    {
+      "epoch": 0.17,
+      "gumbel_tau": 0.68,
+      "image_masks_loss": 0.0172,
+      "llm_learning_rate": 4.755805021301847e-06,
+      "loss": 0.8084,
+      "output_text_mask_loss": 0.0278,
+      "predictor_learning_rate": 0.0004755805021301847,
+      "step": 1636
+    },
+    {
+      "epoch": 0.17,
+      "gumbel_tau": 0.6798,
+      "image_masks_loss": 0.0153,
+      "llm_learning_rate": 4.75544728032754e-06,
+      "loss": 0.7247,
+      "output_text_mask_loss": 0.0484,
+      "predictor_learning_rate": 0.00047554472803275393,
+      "step": 1637
+    },
+    {
+      "epoch": 0.17,
+      "gumbel_tau": 0.6796,
+      "image_masks_loss": 0.0717,
+      "llm_learning_rate": 4.755089290978354e-06,
+      "loss": 0.8555,
+      "output_text_mask_loss": 0.0034,
+      "predictor_learning_rate": 0.00047550892909783537,
+      "step": 1638
+    },
+    {
+      "epoch": 0.17,
+      "gumbel_tau": 0.6795,
+      "image_masks_loss": 0.0534,
+      "llm_learning_rate": 4.754731053293713e-06,
+      "loss": 0.7969,
+      "output_text_mask_loss": 0.0248,
+      "predictor_learning_rate": 0.00047547310532937133,
+      "step": 1639
+    },
+    {
+      "epoch": 0.17,
+      "gumbel_tau": 0.6793,
+      "image_masks_loss": 0.0246,
+      "llm_learning_rate": 4.754372567313067e-06,
+      "loss": 0.7432,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004754372567313067,
+      "step": 1640
+    },
+    {
+      "epoch": 0.17,
+      "gumbel_tau": 0.6792,
+      "image_masks_loss": 0.0182,
+      "llm_learning_rate": 4.754013833075891e-06,
+      "loss": 0.7508,
+      "output_text_mask_loss": 0.0095,
+      "predictor_learning_rate": 0.0004754013833075891,
+      "step": 1641
+    },
+    {
+      "epoch": 0.17,
+      "gumbel_tau": 0.679,
+      "image_masks_loss": 0.01,
+      "llm_learning_rate": 4.753654850621692e-06,
+      "loss": 0.7574,
+      "output_text_mask_loss": 0.017,
+      "predictor_learning_rate": 0.00047536548506216914,
+      "step": 1642
+    },
+    {
+      "epoch": 0.17,
+      "gumbel_tau": 0.6788,
+      "image_masks_loss": 0.0591,
+      "llm_learning_rate": 4.753295619990001e-06,
+      "loss": 0.8414,
+      "output_text_mask_loss": 0.0098,
+      "predictor_learning_rate": 0.00047532956199900005,
+      "step": 1643
+    },
+    {
+      "epoch": 0.17,
+      "gumbel_tau": 0.6787,
+      "image_masks_loss": 0.0231,
+      "llm_learning_rate": 4.752936141220376e-06,
+      "loss": 0.8407,
+      "output_text_mask_loss": 0.0618,
+      "predictor_learning_rate": 0.0004752936141220376,
+      "step": 1644
+    },
+    {
+      "epoch": 0.17,
+      "gumbel_tau": 0.6785,
+      "image_masks_loss": 0.0254,
+      "llm_learning_rate": 4.752576414352407e-06,
+      "loss": 0.7687,
+      "output_text_mask_loss": 0.0021,
+      "predictor_learning_rate": 0.0004752576414352406,
+      "step": 1645
+    },
+    {
+      "epoch": 0.17,
+      "gumbel_tau": 0.6784,
+      "image_masks_loss": 0.0284,
+      "llm_learning_rate": 4.752216439425704e-06,
+      "loss": 0.8268,
+      "output_text_mask_loss": 0.0116,
+      "predictor_learning_rate": 0.00047522164394257045,
+      "step": 1646
+    },
+    {
+      "epoch": 0.17,
+      "gumbel_tau": 0.6782,
+      "image_masks_loss": 0.0273,
+      "llm_learning_rate": 4.751856216479912e-06,
+      "loss": 0.7245,
+      "output_text_mask_loss": 0.0116,
+      "predictor_learning_rate": 0.00047518562164799115,
+      "step": 1647
+    },
+    {
+      "epoch": 0.17,
+      "gumbel_tau": 0.678,
+      "image_masks_loss": 0.0317,
+      "llm_learning_rate": 4.751495745554696e-06,
+      "loss": 0.7576,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00047514957455546957,
+      "step": 1648
+    },
+    {
+      "epoch": 0.17,
+      "gumbel_tau": 0.6779,
+      "image_masks_loss": 0.0395,
+      "llm_learning_rate": 4.751135026689754e-06,
+      "loss": 0.7347,
+      "output_text_mask_loss": 0.0429,
+      "predictor_learning_rate": 0.0004751135026689754,
+      "step": 1649
+    },
+    {
+      "epoch": 0.17,
+      "gumbel_tau": 0.6777,
+      "image_masks_loss": 0.0587,
+      "llm_learning_rate": 4.7507740599248096e-06,
+      "loss": 0.6762,
+      "output_text_mask_loss": 0.0116,
+      "predictor_learning_rate": 0.0004750774059924809,
+      "step": 1650
+    },
+    {
+      "epoch": 0.17,
+      "gumbel_tau": 0.6776,
+      "image_masks_loss": 0.0103,
+      "llm_learning_rate": 4.750412845299611e-06,
+      "loss": 0.7563,
+      "output_text_mask_loss": 0.0248,
+      "predictor_learning_rate": 0.0004750412845299611,
+      "step": 1651
+    },
+    {
+      "epoch": 0.17,
+      "gumbel_tau": 0.6774,
+      "image_masks_loss": 0.0113,
+      "llm_learning_rate": 4.750051382853938e-06,
+      "loss": 0.7189,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00047500513828539374,
+      "step": 1652
+    },
+    {
+      "epoch": 0.17,
+      "gumbel_tau": 0.6772,
+      "image_masks_loss": 0.0347,
+      "llm_learning_rate": 4.749689672627594e-06,
+      "loss": 0.8173,
+      "output_text_mask_loss": 0.0015,
+      "predictor_learning_rate": 0.0004749689672627594,
+      "step": 1653
+    },
+    {
+      "epoch": 0.17,
+      "gumbel_tau": 0.6771,
+      "image_masks_loss": 0.0153,
+      "llm_learning_rate": 4.749327714660413e-06,
+      "loss": 0.8108,
+      "output_text_mask_loss": 0.0004,
+      "predictor_learning_rate": 0.0004749327714660413,
+      "step": 1654
+    },
+    {
+      "epoch": 0.17,
+      "gumbel_tau": 0.6769,
+      "image_masks_loss": 0.0435,
+      "llm_learning_rate": 4.748965508992254e-06,
+      "loss": 0.8158,
+      "output_text_mask_loss": 0.0374,
+      "predictor_learning_rate": 0.00047489655089922537,
+      "step": 1655
+    },
+    {
+      "epoch": 0.17,
+      "gumbel_tau": 0.6768,
+      "image_masks_loss": 0.0687,
+      "llm_learning_rate": 4.7486030556630035e-06,
+      "loss": 0.769,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004748603055663003,
+      "step": 1656
+    },
+    {
+      "epoch": 0.17,
+      "gumbel_tau": 0.6766,
+      "image_masks_loss": 0.0191,
+      "llm_learning_rate": 4.748240354712575e-06,
+      "loss": 0.8144,
+      "output_text_mask_loss": 0.0076,
+      "predictor_learning_rate": 0.00047482403547125754,
+      "step": 1657
+    },
+    {
+      "epoch": 0.17,
+      "gumbel_tau": 0.6764,
+      "image_masks_loss": 0.0162,
+      "llm_learning_rate": 4.747877406180913e-06,
+      "loss": 0.8032,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.0004747877406180913,
+      "step": 1658
+    },
+    {
+      "epoch": 0.17,
+      "gumbel_tau": 0.6763,
+      "image_masks_loss": 0.0144,
+      "llm_learning_rate": 4.747514210107983e-06,
+      "loss": 0.768,
+      "output_text_mask_loss": 0.0427,
+      "predictor_learning_rate": 0.00047475142101079825,
+      "step": 1659
+    },
+    {
+      "epoch": 0.17,
+      "gumbel_tau": 0.6761,
+      "image_masks_loss": 0.0164,
+      "llm_learning_rate": 4.747150766533782e-06,
+      "loss": 0.7237,
+      "output_text_mask_loss": 0.0618,
+      "predictor_learning_rate": 0.00047471507665337815,
+      "step": 1660
+    },
+    {
+      "epoch": 0.17,
+      "gumbel_tau": 0.676,
+      "image_masks_loss": 0.0202,
+      "llm_learning_rate": 4.746787075498332e-06,
+      "loss": 0.8612,
+      "output_text_mask_loss": 0.0488,
+      "predictor_learning_rate": 0.0004746787075498332,
+      "step": 1661
+    },
+    {
+      "epoch": 0.17,
+      "gumbel_tau": 0.6758,
+      "image_masks_loss": 0.0222,
+      "llm_learning_rate": 4.746423137041687e-06,
+      "loss": 0.8403,
+      "output_text_mask_loss": 0.0015,
+      "predictor_learning_rate": 0.00047464231370416865,
+      "step": 1662
+    },
+    {
+      "epoch": 0.17,
+      "gumbel_tau": 0.6756,
+      "image_masks_loss": 0.0345,
+      "llm_learning_rate": 4.7460589512039215e-06,
+      "loss": 0.775,
+      "output_text_mask_loss": 0.0134,
+      "predictor_learning_rate": 0.0004746058951203921,
+      "step": 1663
+    },
+    {
+      "epoch": 0.17,
+      "gumbel_tau": 0.6755,
+      "image_masks_loss": 0.0155,
+      "llm_learning_rate": 4.745694518025142e-06,
+      "loss": 0.7753,
+      "output_text_mask_loss": 0.0418,
+      "predictor_learning_rate": 0.0004745694518025141,
+      "step": 1664
+    },
+    {
+      "epoch": 0.17,
+      "gumbel_tau": 0.6753,
+      "image_masks_loss": 0.0261,
+      "llm_learning_rate": 4.745329837545478e-06,
+      "loss": 0.6884,
+      "output_text_mask_loss": 0.0572,
+      "predictor_learning_rate": 0.00047453298375454777,
+      "step": 1665
+    },
+    {
+      "epoch": 0.17,
+      "gumbel_tau": 0.6752,
+      "image_masks_loss": 0.0112,
+      "llm_learning_rate": 4.744964909805093e-06,
+      "loss": 0.8724,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.00047449649098050927,
+      "step": 1666
+    },
+    {
+      "epoch": 0.17,
+      "gumbel_tau": 0.675,
+      "image_masks_loss": 0.0129,
+      "llm_learning_rate": 4.744599734844171e-06,
+      "loss": 0.784,
+      "output_text_mask_loss": 0.0664,
+      "predictor_learning_rate": 0.00047445997348441704,
+      "step": 1667
+    },
+    {
+      "epoch": 0.17,
+      "gumbel_tau": 0.6748,
+      "image_masks_loss": 0.0277,
+      "llm_learning_rate": 4.744234312702927e-06,
+      "loss": 0.6635,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004744234312702926,
+      "step": 1668
+    },
+    {
+      "epoch": 0.17,
+      "gumbel_tau": 0.6747,
+      "image_masks_loss": 0.0177,
+      "llm_learning_rate": 4.743868643421601e-06,
+      "loss": 0.7796,
+      "output_text_mask_loss": 0.0124,
+      "predictor_learning_rate": 0.00047438686434216005,
+      "step": 1669
+    },
+    {
+      "epoch": 0.17,
+      "gumbel_tau": 0.6745,
+      "image_masks_loss": 0.0248,
+      "llm_learning_rate": 4.743502727040462e-06,
+      "loss": 0.7643,
+      "output_text_mask_loss": 0.0141,
+      "predictor_learning_rate": 0.0004743502727040462,
+      "step": 1670
+    },
+    {
+      "epoch": 0.17,
+      "gumbel_tau": 0.6744,
+      "image_masks_loss": 0.0188,
+      "llm_learning_rate": 4.743136563599806e-06,
+      "loss": 0.7902,
+      "output_text_mask_loss": 0.0202,
+      "predictor_learning_rate": 0.0004743136563599806,
+      "step": 1671
+    },
+    {
+      "epoch": 0.17,
+      "gumbel_tau": 0.6742,
+      "image_masks_loss": 0.0233,
+      "llm_learning_rate": 4.7427701531399554e-06,
+      "loss": 0.7845,
+      "output_text_mask_loss": 0.0282,
+      "predictor_learning_rate": 0.00047427701531399546,
+      "step": 1672
+    },
+    {
+      "epoch": 0.17,
+      "gumbel_tau": 0.674,
+      "image_masks_loss": 0.0395,
+      "llm_learning_rate": 4.742403495701259e-06,
+      "loss": 0.785,
+      "output_text_mask_loss": 0.0368,
+      "predictor_learning_rate": 0.00047424034957012584,
+      "step": 1673
+    },
+    {
+      "epoch": 0.17,
+      "gumbel_tau": 0.6739,
+      "image_masks_loss": 0.0313,
+      "llm_learning_rate": 4.742036591324095e-06,
+      "loss": 0.9102,
+      "output_text_mask_loss": 0.0019,
+      "predictor_learning_rate": 0.0004742036591324095,
+      "step": 1674
+    },
+    {
+      "epoch": 0.17,
+      "gumbel_tau": 0.6737,
+      "image_masks_loss": 0.0244,
+      "llm_learning_rate": 4.741669440048869e-06,
+      "loss": 0.7932,
+      "output_text_mask_loss": 0.0193,
+      "predictor_learning_rate": 0.0004741669440048868,
+      "step": 1675
+    },
+    {
+      "epoch": 0.17,
+      "gumbel_tau": 0.6736,
+      "image_masks_loss": 0.0127,
+      "llm_learning_rate": 4.741302041916009e-06,
+      "loss": 0.7933,
+      "output_text_mask_loss": 0.0072,
+      "predictor_learning_rate": 0.0004741302041916009,
+      "step": 1676
+    },
+    {
+      "epoch": 0.17,
+      "gumbel_tau": 0.6734,
+      "image_masks_loss": 0.036,
+      "llm_learning_rate": 4.740934396965977e-06,
+      "loss": 0.8752,
+      "output_text_mask_loss": 0.0372,
+      "predictor_learning_rate": 0.0004740934396965977,
+      "step": 1677
+    },
+    {
+      "epoch": 0.17,
+      "gumbel_tau": 0.6732,
+      "image_masks_loss": 0.02,
+      "llm_learning_rate": 4.740566505239258e-06,
+      "loss": 0.8447,
+      "output_text_mask_loss": 0.0025,
+      "predictor_learning_rate": 0.0004740566505239258,
+      "step": 1678
+    },
+    {
+      "epoch": 0.17,
+      "gumbel_tau": 0.6731,
+      "image_masks_loss": 0.017,
+      "llm_learning_rate": 4.740198366776364e-06,
+      "loss": 0.7809,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.0004740198366776364,
+      "step": 1679
+    },
+    {
+      "epoch": 0.17,
+      "gumbel_tau": 0.6729,
+      "image_masks_loss": 0.009,
+      "llm_learning_rate": 4.739829981617837e-06,
+      "loss": 0.8003,
+      "output_text_mask_loss": 0.0433,
+      "predictor_learning_rate": 0.0004739829981617837,
+      "step": 1680
+    },
+    {
+      "epoch": 0.17,
+      "gumbel_tau": 0.6728,
+      "image_masks_loss": 0.0288,
+      "llm_learning_rate": 4.739461349804243e-06,
+      "loss": 0.7949,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 0.0004739461349804243,
+      "step": 1681
+    },
+    {
+      "epoch": 0.17,
+      "gumbel_tau": 0.6726,
+      "image_masks_loss": 0.0261,
+      "llm_learning_rate": 4.739092471376176e-06,
+      "loss": 0.8274,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.0004739092471376176,
+      "step": 1682
+    },
+    {
+      "epoch": 0.17,
+      "gumbel_tau": 0.6725,
+      "image_masks_loss": 0.0315,
+      "llm_learning_rate": 4.7387233463742605e-06,
+      "loss": 0.736,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.000473872334637426,
+      "step": 1683
+    },
+    {
+      "epoch": 0.17,
+      "gumbel_tau": 0.6723,
+      "image_masks_loss": 0.0265,
+      "llm_learning_rate": 4.738353974839142e-06,
+      "loss": 0.7963,
+      "output_text_mask_loss": 0.0809,
+      "predictor_learning_rate": 0.00047383539748391425,
+      "step": 1684
+    },
+    {
+      "epoch": 0.17,
+      "gumbel_tau": 0.6721,
+      "image_masks_loss": 0.0053,
+      "llm_learning_rate": 4.737984356811499e-06,
+      "loss": 0.7524,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.0004737984356811499,
+      "step": 1685
+    },
+    {
+      "epoch": 0.17,
+      "gumbel_tau": 0.672,
+      "image_masks_loss": 0.0119,
+      "llm_learning_rate": 4.737614492332033e-06,
+      "loss": 0.728,
+      "output_text_mask_loss": 0.017,
+      "predictor_learning_rate": 0.0004737614492332033,
+      "step": 1686
+    },
+    {
+      "epoch": 0.17,
+      "gumbel_tau": 0.6718,
+      "image_masks_loss": 0.0341,
+      "llm_learning_rate": 4.737244381441476e-06,
+      "loss": 0.8162,
+      "output_text_mask_loss": 0.0067,
+      "predictor_learning_rate": 0.00047372443814414756,
+      "step": 1687
+    },
+    {
+      "epoch": 0.17,
+      "gumbel_tau": 0.6717,
+      "image_masks_loss": 0.0315,
+      "llm_learning_rate": 4.736874024180583e-06,
+      "loss": 0.8002,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.0004736874024180583,
+      "step": 1688
+    },
+    {
+      "epoch": 0.17,
+      "gumbel_tau": 0.6715,
+      "image_masks_loss": 0.0165,
+      "llm_learning_rate": 4.736503420590142e-06,
+      "loss": 0.8002,
+      "output_text_mask_loss": 0.0153,
+      "predictor_learning_rate": 0.00047365034205901414,
+      "step": 1689
+    },
+    {
+      "epoch": 0.17,
+      "gumbel_tau": 0.6713,
+      "image_masks_loss": 0.0246,
+      "llm_learning_rate": 4.736132570710961e-06,
+      "loss": 0.7661,
+      "output_text_mask_loss": 0.0231,
+      "predictor_learning_rate": 0.000473613257071096,
+      "step": 1690
+    },
+    {
+      "epoch": 0.17,
+      "gumbel_tau": 0.6712,
+      "image_masks_loss": 0.0313,
+      "llm_learning_rate": 4.735761474583881e-06,
+      "loss": 0.796,
+      "output_text_mask_loss": 0.0105,
+      "predictor_learning_rate": 0.0004735761474583881,
+      "step": 1691
+    },
+    {
+      "epoch": 0.17,
+      "gumbel_tau": 0.671,
+      "image_masks_loss": 0.0263,
+      "llm_learning_rate": 4.735390132249768e-06,
+      "loss": 0.8126,
+      "output_text_mask_loss": 0.0099,
+      "predictor_learning_rate": 0.0004735390132249767,
+      "step": 1692
+    },
+    {
+      "epoch": 0.17,
+      "gumbel_tau": 0.6709,
+      "image_masks_loss": 0.0345,
+      "llm_learning_rate": 4.735018543749514e-06,
+      "loss": 0.7586,
+      "output_text_mask_loss": 0.01,
+      "predictor_learning_rate": 0.0004735018543749513,
+      "step": 1693
+    },
+    {
+      "epoch": 0.17,
+      "gumbel_tau": 0.6707,
+      "image_masks_loss": 0.0134,
+      "llm_learning_rate": 4.734646709124039e-06,
+      "loss": 0.6798,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.00047346467091240386,
+      "step": 1694
+    },
+    {
+      "epoch": 0.17,
+      "gumbel_tau": 0.6706,
+      "image_masks_loss": 0.0213,
+      "llm_learning_rate": 4.734274628414292e-06,
+      "loss": 0.7871,
+      "output_text_mask_loss": 0.0137,
+      "predictor_learning_rate": 0.00047342746284142914,
+      "step": 1695
+    },
+    {
+      "epoch": 0.17,
+      "gumbel_tau": 0.6704,
+      "image_masks_loss": 0.0115,
+      "llm_learning_rate": 4.733902301661245e-06,
+      "loss": 0.8551,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00047339023016612446,
+      "step": 1696
+    },
+    {
+      "epoch": 0.17,
+      "gumbel_tau": 0.6702,
+      "image_masks_loss": 0.016,
+      "llm_learning_rate": 4.733529728905902e-06,
+      "loss": 0.768,
+      "output_text_mask_loss": 0.0217,
+      "predictor_learning_rate": 0.00047335297289059016,
+      "step": 1697
+    },
+    {
+      "epoch": 0.17,
+      "gumbel_tau": 0.6701,
+      "image_masks_loss": 0.0222,
+      "llm_learning_rate": 4.7331569101892894e-06,
+      "loss": 0.7071,
+      "output_text_mask_loss": 0.0322,
+      "predictor_learning_rate": 0.0004733156910189289,
+      "step": 1698
+    },
+    {
+      "epoch": 0.17,
+      "gumbel_tau": 0.6699,
+      "image_masks_loss": 0.0231,
+      "llm_learning_rate": 4.732783845552464e-06,
+      "loss": 0.8236,
+      "output_text_mask_loss": 0.0763,
+      "predictor_learning_rate": 0.0004732783845552464,
+      "step": 1699
+    },
+    {
+      "epoch": 0.17,
+      "gumbel_tau": 0.6698,
+      "image_masks_loss": 0.0292,
+      "llm_learning_rate": 4.732410535036509e-06,
+      "loss": 0.7478,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.0004732410535036509,
+      "step": 1700
+    },
+    {
+      "epoch": 0.17,
+      "gumbel_tau": 0.6696,
+      "image_masks_loss": 0.0103,
+      "llm_learning_rate": 4.732036978682532e-06,
+      "loss": 0.7769,
+      "output_text_mask_loss": 0.0504,
+      "predictor_learning_rate": 0.00047320369786825326,
+      "step": 1701
+    },
+    {
+      "epoch": 0.17,
+      "gumbel_tau": 0.6694,
+      "image_masks_loss": 0.0128,
+      "llm_learning_rate": 4.731663176531673e-06,
+      "loss": 0.6934,
+      "output_text_mask_loss": 0.0265,
+      "predictor_learning_rate": 0.0004731663176531672,
+      "step": 1702
+    },
+    {
+      "epoch": 0.17,
+      "gumbel_tau": 0.6693,
+      "image_masks_loss": 0.0221,
+      "llm_learning_rate": 4.731289128625092e-06,
+      "loss": 0.8368,
+      "output_text_mask_loss": 0.0029,
+      "predictor_learning_rate": 0.0004731289128625092,
+      "step": 1703
+    },
+    {
+      "epoch": 0.17,
+      "gumbel_tau": 0.6691,
+      "image_masks_loss": 0.0068,
+      "llm_learning_rate": 4.730914835003983e-06,
+      "loss": 0.7459,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.00047309148350039833,
+      "step": 1704
+    },
+    {
+      "epoch": 0.17,
+      "gumbel_tau": 0.669,
+      "image_masks_loss": 0.0221,
+      "llm_learning_rate": 4.730540295709564e-06,
+      "loss": 0.7077,
+      "output_text_mask_loss": 0.0058,
+      "predictor_learning_rate": 0.00047305402957095635,
+      "step": 1705
+    },
+    {
+      "epoch": 0.17,
+      "gumbel_tau": 0.6688,
+      "image_masks_loss": 0.0134,
+      "llm_learning_rate": 4.730165510783079e-06,
+      "loss": 0.7744,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.00047301655107830786,
+      "step": 1706
+    },
+    {
+      "epoch": 0.17,
+      "gumbel_tau": 0.6687,
+      "image_masks_loss": 0.0115,
+      "llm_learning_rate": 4.7297904802658e-06,
+      "loss": 0.6699,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.00047297904802658,
+      "step": 1707
+    },
+    {
+      "epoch": 0.17,
+      "gumbel_tau": 0.6685,
+      "image_masks_loss": 0.0237,
+      "llm_learning_rate": 4.729415204199027e-06,
+      "loss": 0.735,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004729415204199027,
+      "step": 1708
+    },
+    {
+      "epoch": 0.18,
+      "gumbel_tau": 0.6683,
+      "image_masks_loss": 0.0168,
+      "llm_learning_rate": 4.729039682624086e-06,
+      "loss": 0.7823,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.00047290396826240856,
+      "step": 1709
+    },
+    {
+      "epoch": 0.18,
+      "gumbel_tau": 0.6682,
+      "image_masks_loss": 0.0126,
+      "llm_learning_rate": 4.728663915582329e-06,
+      "loss": 0.8658,
+      "output_text_mask_loss": 0.0216,
+      "predictor_learning_rate": 0.0004728663915582329,
+      "step": 1710
+    },
+    {
+      "epoch": 0.18,
+      "gumbel_tau": 0.668,
+      "image_masks_loss": 0.0175,
+      "llm_learning_rate": 4.728287903115139e-06,
+      "loss": 0.7539,
+      "output_text_mask_loss": 0.0763,
+      "predictor_learning_rate": 0.00047282879031151383,
+      "step": 1711
+    },
+    {
+      "epoch": 0.18,
+      "gumbel_tau": 0.6679,
+      "image_masks_loss": 0.0097,
+      "llm_learning_rate": 4.727911645263921e-06,
+      "loss": 0.7818,
+      "output_text_mask_loss": 0.0347,
+      "predictor_learning_rate": 0.00047279116452639204,
+      "step": 1712
+    },
+    {
+      "epoch": 0.18,
+      "gumbel_tau": 0.6677,
+      "image_masks_loss": 0.023,
+      "llm_learning_rate": 4.72753514207011e-06,
+      "loss": 0.7779,
+      "output_text_mask_loss": 0.009,
+      "predictor_learning_rate": 0.000472753514207011,
+      "step": 1713
+    },
+    {
+      "epoch": 0.18,
+      "gumbel_tau": 0.6676,
+      "image_masks_loss": 0.0307,
+      "llm_learning_rate": 4.727158393575167e-06,
+      "loss": 0.753,
+      "output_text_mask_loss": 0.0507,
+      "predictor_learning_rate": 0.0004727158393575167,
+      "step": 1714
+    },
+    {
+      "epoch": 0.18,
+      "gumbel_tau": 0.6674,
+      "image_masks_loss": 0.0104,
+      "llm_learning_rate": 4.726781399820583e-06,
+      "loss": 0.7507,
+      "output_text_mask_loss": 0.042,
+      "predictor_learning_rate": 0.0004726781399820582,
+      "step": 1715
+    },
+    {
+      "epoch": 0.18,
+      "gumbel_tau": 0.6672,
+      "image_masks_loss": 0.0117,
+      "llm_learning_rate": 4.72640416084787e-06,
+      "loss": 0.7676,
+      "output_text_mask_loss": 0.0095,
+      "predictor_learning_rate": 0.00047264041608478693,
+      "step": 1716
+    },
+    {
+      "epoch": 0.18,
+      "gumbel_tau": 0.6671,
+      "image_masks_loss": 0.0097,
+      "llm_learning_rate": 4.726026676698571e-06,
+      "loss": 0.8167,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004726026676698571,
+      "step": 1717
+    },
+    {
+      "epoch": 0.18,
+      "gumbel_tau": 0.6669,
+      "image_masks_loss": 0.0244,
+      "llm_learning_rate": 4.725648947414257e-06,
+      "loss": 0.662,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.00047256489474142567,
+      "step": 1718
+    },
+    {
+      "epoch": 0.18,
+      "gumbel_tau": 0.6668,
+      "image_masks_loss": 0.0196,
+      "llm_learning_rate": 4.725270973036523e-06,
+      "loss": 0.7508,
+      "output_text_mask_loss": 0.0012,
+      "predictor_learning_rate": 0.00047252709730365226,
+      "step": 1719
+    },
+    {
+      "epoch": 0.18,
+      "gumbel_tau": 0.6666,
+      "image_masks_loss": 0.0111,
+      "llm_learning_rate": 4.724892753606994e-06,
+      "loss": 0.7579,
+      "output_text_mask_loss": 0.0039,
+      "predictor_learning_rate": 0.0004724892753606993,
+      "step": 1720
+    },
+    {
+      "epoch": 0.18,
+      "gumbel_tau": 0.6665,
+      "image_masks_loss": 0.0127,
+      "llm_learning_rate": 4.724514289167318e-06,
+      "loss": 0.7496,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00047245142891673176,
+      "step": 1721
+    },
+    {
+      "epoch": 0.18,
+      "gumbel_tau": 0.6663,
+      "image_masks_loss": 0.0166,
+      "llm_learning_rate": 4.724135579759173e-06,
+      "loss": 0.7558,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00047241355797591735,
+      "step": 1722
+    },
+    {
+      "epoch": 0.18,
+      "gumbel_tau": 0.6661,
+      "image_masks_loss": 0.0052,
+      "llm_learning_rate": 4.723756625424265e-06,
+      "loss": 0.6677,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.0004723756625424265,
+      "step": 1723
+    },
+    {
+      "epoch": 0.18,
+      "gumbel_tau": 0.666,
+      "image_masks_loss": 0.0309,
+      "llm_learning_rate": 4.723377426204324e-06,
+      "loss": 0.8269,
+      "output_text_mask_loss": 0.1091,
+      "predictor_learning_rate": 0.0004723377426204324,
+      "step": 1724
+    },
+    {
+      "epoch": 0.18,
+      "gumbel_tau": 0.6658,
+      "image_masks_loss": 0.028,
+      "llm_learning_rate": 4.722997982141109e-06,
+      "loss": 0.8519,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.00047229979821411083,
+      "step": 1725
+    },
+    {
+      "epoch": 0.18,
+      "gumbel_tau": 0.6657,
+      "image_masks_loss": 0.0231,
+      "llm_learning_rate": 4.7226182932764034e-06,
+      "loss": 0.7544,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.0004722618293276403,
+      "step": 1726
+    },
+    {
+      "epoch": 0.18,
+      "gumbel_tau": 0.6655,
+      "image_masks_loss": 0.037,
+      "llm_learning_rate": 4.722238359652022e-06,
+      "loss": 0.8209,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.00047222383596520214,
+      "step": 1727
+    },
+    {
+      "epoch": 0.18,
+      "gumbel_tau": 0.6654,
+      "image_masks_loss": 0.0086,
+      "llm_learning_rate": 4.721858181309802e-06,
+      "loss": 0.7701,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.0004721858181309801,
+      "step": 1728
+    },
+    {
+      "epoch": 0.18,
+      "gumbel_tau": 0.6652,
+      "image_masks_loss": 0.0273,
+      "llm_learning_rate": 4.721477758291609e-06,
+      "loss": 0.8,
+      "output_text_mask_loss": 0.0896,
+      "predictor_learning_rate": 0.0004721477758291609,
+      "step": 1729
+    },
+    {
+      "epoch": 0.18,
+      "gumbel_tau": 0.665,
+      "image_masks_loss": 0.0372,
+      "llm_learning_rate": 4.721097090639339e-06,
+      "loss": 0.7975,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004721097090639338,
+      "step": 1730
+    },
+    {
+      "epoch": 0.18,
+      "gumbel_tau": 0.6649,
+      "image_masks_loss": 0.0073,
+      "llm_learning_rate": 4.720716178394908e-06,
+      "loss": 0.7271,
+      "output_text_mask_loss": 0.0473,
+      "predictor_learning_rate": 0.00047207161783949085,
+      "step": 1731
+    },
+    {
+      "epoch": 0.18,
+      "gumbel_tau": 0.6647,
+      "image_masks_loss": 0.028,
+      "llm_learning_rate": 4.720335021600266e-06,
+      "loss": 0.8115,
+      "output_text_mask_loss": 0.0158,
+      "predictor_learning_rate": 0.0004720335021600266,
+      "step": 1732
+    },
+    {
+      "epoch": 0.18,
+      "gumbel_tau": 0.6646,
+      "image_masks_loss": 0.0246,
+      "llm_learning_rate": 4.719953620297387e-06,
+      "loss": 0.7542,
+      "output_text_mask_loss": 0.0668,
+      "predictor_learning_rate": 0.00047199536202973866,
+      "step": 1733
+    },
+    {
+      "epoch": 0.18,
+      "gumbel_tau": 0.6644,
+      "image_masks_loss": 0.0196,
+      "llm_learning_rate": 4.719571974528269e-06,
+      "loss": 0.8071,
+      "output_text_mask_loss": 0.0771,
+      "predictor_learning_rate": 0.0004719571974528269,
+      "step": 1734
+    },
+    {
+      "epoch": 0.18,
+      "gumbel_tau": 0.6643,
+      "image_masks_loss": 0.0214,
+      "llm_learning_rate": 4.719190084334942e-06,
+      "loss": 0.672,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.00047191900843349414,
+      "step": 1735
+    },
+    {
+      "epoch": 0.18,
+      "gumbel_tau": 0.6641,
+      "image_masks_loss": 0.0167,
+      "llm_learning_rate": 4.7188079497594594e-06,
+      "loss": 0.7908,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.00047188079497594595,
+      "step": 1736
+    },
+    {
+      "epoch": 0.18,
+      "gumbel_tau": 0.6639,
+      "image_masks_loss": 0.0385,
+      "llm_learning_rate": 4.718425570843904e-06,
+      "loss": 0.8429,
+      "output_text_mask_loss": 0.0599,
+      "predictor_learning_rate": 0.0004718425570843903,
+      "step": 1737
+    },
+    {
+      "epoch": 0.18,
+      "gumbel_tau": 0.6638,
+      "image_masks_loss": 0.0138,
+      "llm_learning_rate": 4.718042947630383e-06,
+      "loss": 0.8368,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.00047180429476303823,
+      "step": 1738
+    },
+    {
+      "epoch": 0.18,
+      "gumbel_tau": 0.6636,
+      "image_masks_loss": 0.0267,
+      "llm_learning_rate": 4.717660080161032e-06,
+      "loss": 0.7925,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.0004717660080161031,
+      "step": 1739
+    },
+    {
+      "epoch": 0.18,
+      "gumbel_tau": 0.6635,
+      "image_masks_loss": 0.0307,
+      "llm_learning_rate": 4.717276968478013e-06,
+      "loss": 0.732,
+      "output_text_mask_loss": 0.024,
+      "predictor_learning_rate": 0.00047172769684780127,
+      "step": 1740
+    },
+    {
+      "epoch": 0.18,
+      "gumbel_tau": 0.6633,
+      "image_masks_loss": 0.0177,
+      "llm_learning_rate": 4.716893612623515e-06,
+      "loss": 0.7942,
+      "output_text_mask_loss": 0.0465,
+      "predictor_learning_rate": 0.00047168936126235154,
+      "step": 1741
+    },
+    {
+      "epoch": 0.18,
+      "gumbel_tau": 0.6632,
+      "image_masks_loss": 0.0257,
+      "llm_learning_rate": 4.7165100126397555e-06,
+      "loss": 0.6749,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.00047165100126397554,
+      "step": 1742
+    },
+    {
+      "epoch": 0.18,
+      "gumbel_tau": 0.663,
+      "image_masks_loss": 0.0332,
+      "llm_learning_rate": 4.716126168568976e-06,
+      "loss": 0.8188,
+      "output_text_mask_loss": 0.0231,
+      "predictor_learning_rate": 0.0004716126168568976,
+      "step": 1743
+    },
+    {
+      "epoch": 0.18,
+      "gumbel_tau": 0.6628,
+      "image_masks_loss": 0.0129,
+      "llm_learning_rate": 4.715742080453446e-06,
+      "loss": 0.7963,
+      "output_text_mask_loss": 0.0141,
+      "predictor_learning_rate": 0.0004715742080453446,
+      "step": 1744
+    },
+    {
+      "epoch": 0.18,
+      "gumbel_tau": 0.6627,
+      "image_masks_loss": 0.0159,
+      "llm_learning_rate": 4.715357748335464e-06,
+      "loss": 0.7843,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 0.00047153577483354637,
+      "step": 1745
+    },
+    {
+      "epoch": 0.18,
+      "gumbel_tau": 0.6625,
+      "image_masks_loss": 0.0671,
+      "llm_learning_rate": 4.714973172257351e-06,
+      "loss": 0.7568,
+      "output_text_mask_loss": 0.0231,
+      "predictor_learning_rate": 0.0004714973172257351,
+      "step": 1746
+    },
+    {
+      "epoch": 0.18,
+      "gumbel_tau": 0.6624,
+      "image_masks_loss": 0.0161,
+      "llm_learning_rate": 4.71458835226146e-06,
+      "loss": 0.8082,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.00047145883522614596,
+      "step": 1747
+    },
+    {
+      "epoch": 0.18,
+      "gumbel_tau": 0.6622,
+      "image_masks_loss": 0.0102,
+      "llm_learning_rate": 4.714203288390166e-06,
+      "loss": 0.7103,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.0004714203288390165,
+      "step": 1748
+    },
+    {
+      "epoch": 0.18,
+      "gumbel_tau": 0.6621,
+      "image_masks_loss": 0.0296,
+      "llm_learning_rate": 4.713817980685874e-06,
+      "loss": 0.7481,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.00047138179806858734,
+      "step": 1749
+    },
+    {
+      "epoch": 0.18,
+      "gumbel_tau": 0.6619,
+      "image_masks_loss": 0.0101,
+      "llm_learning_rate": 4.713432429191015e-06,
+      "loss": 0.7242,
+      "output_text_mask_loss": 0.0729,
+      "predictor_learning_rate": 0.00047134324291910144,
+      "step": 1750
+    },
+    {
+      "epoch": 0.18,
+      "gumbel_tau": 0.6618,
+      "image_masks_loss": 0.0207,
+      "llm_learning_rate": 4.7130466339480455e-06,
+      "loss": 0.8226,
+      "output_text_mask_loss": 0.0523,
+      "predictor_learning_rate": 0.00047130466339480456,
+      "step": 1751
+    },
+    {
+      "epoch": 0.18,
+      "gumbel_tau": 0.6616,
+      "image_masks_loss": 0.0633,
+      "llm_learning_rate": 4.712660594999453e-06,
+      "loss": 0.8056,
+      "output_text_mask_loss": 0.0081,
+      "predictor_learning_rate": 0.0004712660594999452,
+      "step": 1752
+    },
+    {
+      "epoch": 0.18,
+      "gumbel_tau": 0.6614,
+      "image_masks_loss": 0.0351,
+      "llm_learning_rate": 4.712274312387746e-06,
+      "loss": 0.6499,
+      "output_text_mask_loss": 0.0014,
+      "predictor_learning_rate": 0.0004712274312387746,
+      "step": 1753
+    },
+    {
+      "epoch": 0.18,
+      "gumbel_tau": 0.6613,
+      "image_masks_loss": 0.0212,
+      "llm_learning_rate": 4.7118877861554645e-06,
+      "loss": 0.8366,
+      "output_text_mask_loss": 0.0231,
+      "predictor_learning_rate": 0.00047118877861554643,
+      "step": 1754
+    },
+    {
+      "epoch": 0.18,
+      "gumbel_tau": 0.6611,
+      "image_masks_loss": 0.045,
+      "llm_learning_rate": 4.7115010163451735e-06,
+      "loss": 0.7507,
+      "output_text_mask_loss": 0.0425,
+      "predictor_learning_rate": 0.0004711501016345173,
+      "step": 1755
+    },
+    {
+      "epoch": 0.18,
+      "gumbel_tau": 0.661,
+      "image_masks_loss": 0.032,
+      "llm_learning_rate": 4.711114002999464e-06,
+      "loss": 0.7827,
+      "output_text_mask_loss": 0.0238,
+      "predictor_learning_rate": 0.0004711114002999464,
+      "step": 1756
+    },
+    {
+      "epoch": 0.18,
+      "gumbel_tau": 0.6608,
+      "image_masks_loss": 0.0275,
+      "llm_learning_rate": 4.7107267461609565e-06,
+      "loss": 0.9,
+      "output_text_mask_loss": 0.0511,
+      "predictor_learning_rate": 0.0004710726746160956,
+      "step": 1757
+    },
+    {
+      "epoch": 0.18,
+      "gumbel_tau": 0.6607,
+      "image_masks_loss": 0.0081,
+      "llm_learning_rate": 4.710339245872295e-06,
+      "loss": 0.78,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00047103392458722947,
+      "step": 1758
+    },
+    {
+      "epoch": 0.18,
+      "gumbel_tau": 0.6605,
+      "image_masks_loss": 0.0155,
+      "llm_learning_rate": 4.709951502176152e-06,
+      "loss": 0.7938,
+      "output_text_mask_loss": 0.033,
+      "predictor_learning_rate": 0.0004709951502176152,
+      "step": 1759
+    },
+    {
+      "epoch": 0.18,
+      "gumbel_tau": 0.6604,
+      "image_masks_loss": 0.037,
+      "llm_learning_rate": 4.7095635151152275e-06,
+      "loss": 0.7718,
+      "output_text_mask_loss": 0.0671,
+      "predictor_learning_rate": 0.0004709563515115227,
+      "step": 1760
+    },
+    {
+      "epoch": 0.18,
+      "gumbel_tau": 0.6602,
+      "image_masks_loss": 0.0492,
+      "llm_learning_rate": 4.709175284732247e-06,
+      "loss": 0.8543,
+      "output_text_mask_loss": 0.0071,
+      "predictor_learning_rate": 0.00047091752847322466,
+      "step": 1761
+    },
+    {
+      "epoch": 0.18,
+      "gumbel_tau": 0.66,
+      "image_masks_loss": 0.0061,
+      "llm_learning_rate": 4.7087868110699635e-06,
+      "loss": 0.7549,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004708786811069963,
+      "step": 1762
+    },
+    {
+      "epoch": 0.18,
+      "gumbel_tau": 0.6599,
+      "image_masks_loss": 0.0284,
+      "llm_learning_rate": 4.708398094171156e-06,
+      "loss": 0.7695,
+      "output_text_mask_loss": 0.033,
+      "predictor_learning_rate": 0.0004708398094171156,
+      "step": 1763
+    },
+    {
+      "epoch": 0.18,
+      "gumbel_tau": 0.6597,
+      "image_masks_loss": 0.0129,
+      "llm_learning_rate": 4.7080091340786326e-06,
+      "loss": 0.6857,
+      "output_text_mask_loss": 0.0061,
+      "predictor_learning_rate": 0.00047080091340786324,
+      "step": 1764
+    },
+    {
+      "epoch": 0.18,
+      "gumbel_tau": 0.6596,
+      "image_masks_loss": 0.0101,
+      "llm_learning_rate": 4.707619930835224e-06,
+      "loss": 0.854,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004707619930835224,
+      "step": 1765
+    },
+    {
+      "epoch": 0.18,
+      "gumbel_tau": 0.6594,
+      "image_masks_loss": 0.0257,
+      "llm_learning_rate": 4.707230484483792e-06,
+      "loss": 0.8165,
+      "output_text_mask_loss": 0.0277,
+      "predictor_learning_rate": 0.00047072304844837914,
+      "step": 1766
+    },
+    {
+      "epoch": 0.18,
+      "gumbel_tau": 0.6593,
+      "image_masks_loss": 0.015,
+      "llm_learning_rate": 4.7068407950672225e-06,
+      "loss": 0.7602,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.0004706840795067222,
+      "step": 1767
+    },
+    {
+      "epoch": 0.18,
+      "gumbel_tau": 0.6591,
+      "image_masks_loss": 0.0121,
+      "llm_learning_rate": 4.706450862628429e-06,
+      "loss": 0.796,
+      "output_text_mask_loss": 0.0078,
+      "predictor_learning_rate": 0.00047064508626284285,
+      "step": 1768
+    },
+    {
+      "epoch": 0.18,
+      "gumbel_tau": 0.6589,
+      "image_masks_loss": 0.0135,
+      "llm_learning_rate": 4.706060687210351e-06,
+      "loss": 0.8097,
+      "output_text_mask_loss": 0.0231,
+      "predictor_learning_rate": 0.00047060606872103514,
+      "step": 1769
+    },
+    {
+      "epoch": 0.18,
+      "gumbel_tau": 0.6588,
+      "image_masks_loss": 0.0136,
+      "llm_learning_rate": 4.705670268855957e-06,
+      "loss": 0.8964,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004705670268855957,
+      "step": 1770
+    },
+    {
+      "epoch": 0.18,
+      "gumbel_tau": 0.6586,
+      "image_masks_loss": 0.0156,
+      "llm_learning_rate": 4.705279607608241e-06,
+      "loss": 0.7826,
+      "output_text_mask_loss": 0.0086,
+      "predictor_learning_rate": 0.000470527960760824,
+      "step": 1771
+    },
+    {
+      "epoch": 0.18,
+      "gumbel_tau": 0.6585,
+      "image_masks_loss": 0.0349,
+      "llm_learning_rate": 4.704888703510221e-06,
+      "loss": 0.7764,
+      "output_text_mask_loss": 0.0208,
+      "predictor_learning_rate": 0.00047048887035102213,
+      "step": 1772
+    },
+    {
+      "epoch": 0.18,
+      "gumbel_tau": 0.6583,
+      "image_masks_loss": 0.0082,
+      "llm_learning_rate": 4.704497556604946e-06,
+      "loss": 0.7612,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.00047044975566049464,
+      "step": 1773
+    },
+    {
+      "epoch": 0.18,
+      "gumbel_tau": 0.6582,
+      "image_masks_loss": 0.0132,
+      "llm_learning_rate": 4.704106166935491e-06,
+      "loss": 0.7704,
+      "output_text_mask_loss": 0.0278,
+      "predictor_learning_rate": 0.00047041061669354914,
+      "step": 1774
+    },
+    {
+      "epoch": 0.18,
+      "gumbel_tau": 0.658,
+      "image_masks_loss": 0.0077,
+      "llm_learning_rate": 4.703714534544955e-06,
+      "loss": 0.8145,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.00047037145345449544,
+      "step": 1775
+    },
+    {
+      "epoch": 0.18,
+      "gumbel_tau": 0.6579,
+      "image_masks_loss": 0.0177,
+      "llm_learning_rate": 4.703322659476465e-06,
+      "loss": 0.7253,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.0004703322659476465,
+      "step": 1776
+    },
+    {
+      "epoch": 0.18,
+      "gumbel_tau": 0.6577,
+      "image_masks_loss": 0.045,
+      "llm_learning_rate": 4.702930541773177e-06,
+      "loss": 0.7933,
+      "output_text_mask_loss": 0.0004,
+      "predictor_learning_rate": 0.0004702930541773177,
+      "step": 1777
+    },
+    {
+      "epoch": 0.18,
+      "gumbel_tau": 0.6576,
+      "image_masks_loss": 0.0124,
+      "llm_learning_rate": 4.70253818147827e-06,
+      "loss": 0.775,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 0.00047025381814782695,
+      "step": 1778
+    },
+    {
+      "epoch": 0.18,
+      "gumbel_tau": 0.6574,
+      "image_masks_loss": 0.0584,
+      "llm_learning_rate": 4.702145578634953e-06,
+      "loss": 0.7382,
+      "output_text_mask_loss": 0.0427,
+      "predictor_learning_rate": 0.0004702145578634953,
+      "step": 1779
+    },
+    {
+      "epoch": 0.18,
+      "gumbel_tau": 0.6572,
+      "image_masks_loss": 0.0549,
+      "llm_learning_rate": 4.70175273328646e-06,
+      "loss": 0.7076,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.00047017527332864593,
+      "step": 1780
+    },
+    {
+      "epoch": 0.18,
+      "gumbel_tau": 0.6571,
+      "image_masks_loss": 0.0303,
+      "llm_learning_rate": 4.701359645476052e-06,
+      "loss": 0.7632,
+      "output_text_mask_loss": 0.0198,
+      "predictor_learning_rate": 0.0004701359645476051,
+      "step": 1781
+    },
+    {
+      "epoch": 0.18,
+      "gumbel_tau": 0.6569,
+      "image_masks_loss": 0.0134,
+      "llm_learning_rate": 4.700966315247015e-06,
+      "loss": 0.6275,
+      "output_text_mask_loss": 0.0029,
+      "predictor_learning_rate": 0.0004700966315247015,
+      "step": 1782
+    },
+    {
+      "epoch": 0.18,
+      "gumbel_tau": 0.6568,
+      "image_masks_loss": 0.0618,
+      "llm_learning_rate": 4.700572742642666e-06,
+      "loss": 0.7825,
+      "output_text_mask_loss": 0.0081,
+      "predictor_learning_rate": 0.0004700572742642666,
+      "step": 1783
+    },
+    {
+      "epoch": 0.18,
+      "gumbel_tau": 0.6566,
+      "image_masks_loss": 0.0237,
+      "llm_learning_rate": 4.700178927706346e-06,
+      "loss": 0.7344,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.0004700178927706346,
+      "step": 1784
+    },
+    {
+      "epoch": 0.18,
+      "gumbel_tau": 0.6565,
+      "image_masks_loss": 0.0122,
+      "llm_learning_rate": 4.69978487048142e-06,
+      "loss": 0.7559,
+      "output_text_mask_loss": 0.0488,
+      "predictor_learning_rate": 0.000469978487048142,
+      "step": 1785
+    },
+    {
+      "epoch": 0.18,
+      "gumbel_tau": 0.6563,
+      "image_masks_loss": 0.0149,
+      "llm_learning_rate": 4.699390571011285e-06,
+      "loss": 0.7687,
+      "output_text_mask_loss": 0.0114,
+      "predictor_learning_rate": 0.0004699390571011285,
+      "step": 1786
+    },
+    {
+      "epoch": 0.18,
+      "gumbel_tau": 0.6562,
+      "image_masks_loss": 0.0391,
+      "llm_learning_rate": 4.698996029339361e-06,
+      "loss": 0.7645,
+      "output_text_mask_loss": 0.0381,
+      "predictor_learning_rate": 0.00046989960293393614,
+      "step": 1787
+    },
+    {
+      "epoch": 0.18,
+      "gumbel_tau": 0.656,
+      "image_masks_loss": 0.0143,
+      "llm_learning_rate": 4.698601245509097e-06,
+      "loss": 0.7843,
+      "output_text_mask_loss": 0.0132,
+      "predictor_learning_rate": 0.00046986012455090963,
+      "step": 1788
+    },
+    {
+      "epoch": 0.18,
+      "gumbel_tau": 0.6558,
+      "image_masks_loss": 0.0391,
+      "llm_learning_rate": 4.698206219563966e-06,
+      "loss": 0.7708,
+      "output_text_mask_loss": 0.0599,
+      "predictor_learning_rate": 0.0004698206219563966,
+      "step": 1789
+    },
+    {
+      "epoch": 0.18,
+      "gumbel_tau": 0.6557,
+      "image_masks_loss": 0.0241,
+      "llm_learning_rate": 4.69781095154747e-06,
+      "loss": 0.846,
+      "output_text_mask_loss": 0.1411,
+      "predictor_learning_rate": 0.000469781095154747,
+      "step": 1790
+    },
+    {
+      "epoch": 0.18,
+      "gumbel_tau": 0.6555,
+      "image_masks_loss": 0.0134,
+      "llm_learning_rate": 4.697415441503137e-06,
+      "loss": 0.7812,
+      "output_text_mask_loss": 0.017,
+      "predictor_learning_rate": 0.00046974154415031365,
+      "step": 1791
+    },
+    {
+      "epoch": 0.18,
+      "gumbel_tau": 0.6554,
+      "image_masks_loss": 0.0134,
+      "llm_learning_rate": 4.69701968947452e-06,
+      "loss": 0.826,
+      "output_text_mask_loss": 0.0374,
+      "predictor_learning_rate": 0.000469701968947452,
+      "step": 1792
+    },
+    {
+      "epoch": 0.18,
+      "gumbel_tau": 0.6552,
+      "image_masks_loss": 0.0106,
+      "llm_learning_rate": 4.696623695505202e-06,
+      "loss": 0.8039,
+      "output_text_mask_loss": 0.0275,
+      "predictor_learning_rate": 0.00046966236955052015,
+      "step": 1793
+    },
+    {
+      "epoch": 0.18,
+      "gumbel_tau": 0.6551,
+      "image_masks_loss": 0.0256,
+      "llm_learning_rate": 4.696227459638789e-06,
+      "loss": 0.7821,
+      "output_text_mask_loss": 0.0299,
+      "predictor_learning_rate": 0.0004696227459638789,
+      "step": 1794
+    },
+    {
+      "epoch": 0.18,
+      "gumbel_tau": 0.6549,
+      "image_masks_loss": 0.0115,
+      "llm_learning_rate": 4.695830981918918e-06,
+      "loss": 0.845,
+      "output_text_mask_loss": 0.0185,
+      "predictor_learning_rate": 0.0004695830981918917,
+      "step": 1795
+    },
+    {
+      "epoch": 0.18,
+      "gumbel_tau": 0.6548,
+      "image_masks_loss": 0.0572,
+      "llm_learning_rate": 4.695434262389246e-06,
+      "loss": 0.753,
+      "output_text_mask_loss": 0.0215,
+      "predictor_learning_rate": 0.0004695434262389246,
+      "step": 1796
+    },
+    {
+      "epoch": 0.18,
+      "gumbel_tau": 0.6546,
+      "image_masks_loss": 0.016,
+      "llm_learning_rate": 4.695037301093465e-06,
+      "loss": 0.7298,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00046950373010934643,
+      "step": 1797
+    },
+    {
+      "epoch": 0.18,
+      "gumbel_tau": 0.6545,
+      "image_masks_loss": 0.0292,
+      "llm_learning_rate": 4.6946400980752854e-06,
+      "loss": 0.8474,
+      "output_text_mask_loss": 0.0034,
+      "predictor_learning_rate": 0.0004694640098075285,
+      "step": 1798
+    },
+    {
+      "epoch": 0.18,
+      "gumbel_tau": 0.6543,
+      "image_masks_loss": 0.0263,
+      "llm_learning_rate": 4.6942426533784504e-06,
+      "loss": 0.8384,
+      "output_text_mask_loss": 0.0052,
+      "predictor_learning_rate": 0.000469424265337845,
+      "step": 1799
+    },
+    {
+      "epoch": 0.18,
+      "gumbel_tau": 0.6541,
+      "image_masks_loss": 0.0271,
+      "llm_learning_rate": 4.693844967046727e-06,
+      "loss": 0.7243,
+      "output_text_mask_loss": 0.0216,
+      "predictor_learning_rate": 0.00046938449670467266,
+      "step": 1800
+    },
+    {
+      "epoch": 0.18,
+      "gumbel_tau": 0.654,
+      "image_masks_loss": 0.0071,
+      "llm_learning_rate": 4.69344703912391e-06,
+      "loss": 0.76,
+      "output_text_mask_loss": 0.0035,
+      "predictor_learning_rate": 0.0004693447039123909,
+      "step": 1801
+    },
+    {
+      "epoch": 0.18,
+      "gumbel_tau": 0.6538,
+      "image_masks_loss": 0.0081,
+      "llm_learning_rate": 4.693048869653818e-06,
+      "loss": 0.7561,
+      "output_text_mask_loss": 0.0858,
+      "predictor_learning_rate": 0.00046930488696538176,
+      "step": 1802
+    },
+    {
+      "epoch": 0.18,
+      "gumbel_tau": 0.6537,
+      "image_masks_loss": 0.0256,
+      "llm_learning_rate": 4.6926504586803e-06,
+      "loss": 0.7115,
+      "output_text_mask_loss": 0.0572,
+      "predictor_learning_rate": 0.00046926504586803,
+      "step": 1803
+    },
+    {
+      "epoch": 0.18,
+      "gumbel_tau": 0.6535,
+      "image_masks_loss": 0.0263,
+      "llm_learning_rate": 4.69225180624723e-06,
+      "loss": 0.8286,
+      "output_text_mask_loss": 0.0885,
+      "predictor_learning_rate": 0.00046922518062472295,
+      "step": 1804
+    },
+    {
+      "epoch": 0.18,
+      "gumbel_tau": 0.6534,
+      "image_masks_loss": 0.0092,
+      "llm_learning_rate": 4.691852912398508e-06,
+      "loss": 0.892,
+      "output_text_mask_loss": 0.0124,
+      "predictor_learning_rate": 0.00046918529123985076,
+      "step": 1805
+    },
+    {
+      "epoch": 0.19,
+      "gumbel_tau": 0.6532,
+      "image_masks_loss": 0.0404,
+      "llm_learning_rate": 4.691453777178061e-06,
+      "loss": 0.8164,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 0.00046914537771780606,
+      "step": 1806
+    },
+    {
+      "epoch": 0.19,
+      "gumbel_tau": 0.6531,
+      "image_masks_loss": 0.0239,
+      "llm_learning_rate": 4.691054400629843e-06,
+      "loss": 0.7569,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.0004691054400629843,
+      "step": 1807
+    },
+    {
+      "epoch": 0.19,
+      "gumbel_tau": 0.6529,
+      "image_masks_loss": 0.0334,
+      "llm_learning_rate": 4.690654782797833e-06,
+      "loss": 0.8365,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.0004690654782797833,
+      "step": 1808
+    },
+    {
+      "epoch": 0.19,
+      "gumbel_tau": 0.6528,
+      "image_masks_loss": 0.0125,
+      "llm_learning_rate": 4.69025492372604e-06,
+      "loss": 0.7599,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.00046902549237260393,
+      "step": 1809
+    },
+    {
+      "epoch": 0.19,
+      "gumbel_tau": 0.6526,
+      "image_masks_loss": 0.0263,
+      "llm_learning_rate": 4.689854823458496e-06,
+      "loss": 0.7659,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00046898548234584956,
+      "step": 1810
+    },
+    {
+      "epoch": 0.19,
+      "gumbel_tau": 0.6525,
+      "image_masks_loss": 0.0117,
+      "llm_learning_rate": 4.68945448203926e-06,
+      "loss": 0.7373,
+      "output_text_mask_loss": 0.0012,
+      "predictor_learning_rate": 0.00046894544820392603,
+      "step": 1811
+    },
+    {
+      "epoch": 0.19,
+      "gumbel_tau": 0.6523,
+      "image_masks_loss": 0.0136,
+      "llm_learning_rate": 4.689053899512421e-06,
+      "loss": 0.7524,
+      "output_text_mask_loss": 0.0174,
+      "predictor_learning_rate": 0.0004689053899512421,
+      "step": 1812
+    },
+    {
+      "epoch": 0.19,
+      "gumbel_tau": 0.6521,
+      "image_masks_loss": 0.0103,
+      "llm_learning_rate": 4.688653075922091e-06,
+      "loss": 0.8365,
+      "output_text_mask_loss": 0.0992,
+      "predictor_learning_rate": 0.00046886530759220903,
+      "step": 1813
+    },
+    {
+      "epoch": 0.19,
+      "gumbel_tau": 0.652,
+      "image_masks_loss": 0.0216,
+      "llm_learning_rate": 4.688252011312408e-06,
+      "loss": 0.8778,
+      "output_text_mask_loss": 0.0679,
+      "predictor_learning_rate": 0.00046882520113124076,
+      "step": 1814
+    },
+    {
+      "epoch": 0.19,
+      "gumbel_tau": 0.6518,
+      "image_masks_loss": 0.0178,
+      "llm_learning_rate": 4.68785070572754e-06,
+      "loss": 0.7806,
+      "output_text_mask_loss": 0.0166,
+      "predictor_learning_rate": 0.000468785070572754,
+      "step": 1815
+    },
+    {
+      "epoch": 0.19,
+      "gumbel_tau": 0.6517,
+      "image_masks_loss": 0.0072,
+      "llm_learning_rate": 4.687449159211679e-06,
+      "loss": 0.6819,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.00046874491592116785,
+      "step": 1816
+    },
+    {
+      "epoch": 0.19,
+      "gumbel_tau": 0.6515,
+      "image_masks_loss": 0.0216,
+      "llm_learning_rate": 4.687047371809044e-06,
+      "loss": 0.8717,
+      "output_text_mask_loss": 0.0391,
+      "predictor_learning_rate": 0.00046870473718090434,
+      "step": 1817
+    },
+    {
+      "epoch": 0.19,
+      "gumbel_tau": 0.6514,
+      "image_masks_loss": 0.0116,
+      "llm_learning_rate": 4.686645343563881e-06,
+      "loss": 0.763,
+      "output_text_mask_loss": 0.0172,
+      "predictor_learning_rate": 0.00046866453435638806,
+      "step": 1818
+    },
+    {
+      "epoch": 0.19,
+      "gumbel_tau": 0.6512,
+      "image_masks_loss": 0.0156,
+      "llm_learning_rate": 4.686243074520462e-06,
+      "loss": 0.6971,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004686243074520462,
+      "step": 1819
+    },
+    {
+      "epoch": 0.19,
+      "gumbel_tau": 0.6511,
+      "image_masks_loss": 0.017,
+      "llm_learning_rate": 4.685840564723086e-06,
+      "loss": 0.806,
+      "output_text_mask_loss": 0.0298,
+      "predictor_learning_rate": 0.0004685840564723086,
+      "step": 1820
+    },
+    {
+      "epoch": 0.19,
+      "gumbel_tau": 0.6509,
+      "image_masks_loss": 0.0152,
+      "llm_learning_rate": 4.685437814216079e-06,
+      "loss": 0.8329,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00046854378142160784,
+      "step": 1821
+    },
+    {
+      "epoch": 0.19,
+      "gumbel_tau": 0.6508,
+      "image_masks_loss": 0.02,
+      "llm_learning_rate": 4.6850348230437914e-06,
+      "loss": 0.796,
+      "output_text_mask_loss": 0.0124,
+      "predictor_learning_rate": 0.0004685034823043791,
+      "step": 1822
+    },
+    {
+      "epoch": 0.19,
+      "gumbel_tau": 0.6506,
+      "image_masks_loss": 0.0278,
+      "llm_learning_rate": 4.684631591250602e-06,
+      "loss": 0.7953,
+      "output_text_mask_loss": 0.0484,
+      "predictor_learning_rate": 0.0004684631591250602,
+      "step": 1823
+    },
+    {
+      "epoch": 0.19,
+      "gumbel_tau": 0.6505,
+      "image_masks_loss": 0.0053,
+      "llm_learning_rate": 4.6842281188809165e-06,
+      "loss": 0.6881,
+      "output_text_mask_loss": 0.0116,
+      "predictor_learning_rate": 0.0004684228118880916,
+      "step": 1824
+    },
+    {
+      "epoch": 0.19,
+      "gumbel_tau": 0.6503,
+      "image_masks_loss": 0.0062,
+      "llm_learning_rate": 4.683824405979165e-06,
+      "loss": 0.7671,
+      "output_text_mask_loss": 0.0353,
+      "predictor_learning_rate": 0.0004683824405979165,
+      "step": 1825
+    },
+    {
+      "epoch": 0.19,
+      "gumbel_tau": 0.6501,
+      "image_masks_loss": 0.0112,
+      "llm_learning_rate": 4.683420452589805e-06,
+      "loss": 0.8378,
+      "output_text_mask_loss": 0.0054,
+      "predictor_learning_rate": 0.0004683420452589805,
+      "step": 1826
+    },
+    {
+      "epoch": 0.19,
+      "gumbel_tau": 0.65,
+      "image_masks_loss": 0.025,
+      "llm_learning_rate": 4.683016258757321e-06,
+      "loss": 0.768,
+      "output_text_mask_loss": 0.0416,
+      "predictor_learning_rate": 0.0004683016258757321,
+      "step": 1827
+    },
+    {
+      "epoch": 0.19,
+      "gumbel_tau": 0.6498,
+      "image_masks_loss": 0.0374,
+      "llm_learning_rate": 4.6826118245262255e-06,
+      "loss": 0.7447,
+      "output_text_mask_loss": 0.0896,
+      "predictor_learning_rate": 0.0004682611824526225,
+      "step": 1828
+    },
+    {
+      "epoch": 0.19,
+      "gumbel_tau": 0.6497,
+      "image_masks_loss": 0.0265,
+      "llm_learning_rate": 4.682207149941054e-06,
+      "loss": 0.6563,
+      "output_text_mask_loss": 0.021,
+      "predictor_learning_rate": 0.00046822071499410536,
+      "step": 1829
+    },
+    {
+      "epoch": 0.19,
+      "gumbel_tau": 0.6495,
+      "image_masks_loss": 0.0237,
+      "llm_learning_rate": 4.6818022350463695e-06,
+      "loss": 0.6911,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00046818022350463695,
+      "step": 1830
+    },
+    {
+      "epoch": 0.19,
+      "gumbel_tau": 0.6494,
+      "image_masks_loss": 0.0113,
+      "llm_learning_rate": 4.681397079886765e-06,
+      "loss": 0.7066,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.00046813970798867643,
+      "step": 1831
+    },
+    {
+      "epoch": 0.19,
+      "gumbel_tau": 0.6492,
+      "image_masks_loss": 0.0233,
+      "llm_learning_rate": 4.680991684506853e-06,
+      "loss": 0.7661,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.0004680991684506853,
+      "step": 1832
+    },
+    {
+      "epoch": 0.19,
+      "gumbel_tau": 0.6491,
+      "image_masks_loss": 0.014,
+      "llm_learning_rate": 4.68058604895128e-06,
+      "loss": 0.8407,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.00046805860489512795,
+      "step": 1833
+    },
+    {
+      "epoch": 0.19,
+      "gumbel_tau": 0.6489,
+      "image_masks_loss": 0.0208,
+      "llm_learning_rate": 4.680180173264714e-06,
+      "loss": 0.811,
+      "output_text_mask_loss": 0.0193,
+      "predictor_learning_rate": 0.00046801801732647137,
+      "step": 1834
+    },
+    {
+      "epoch": 0.19,
+      "gumbel_tau": 0.6488,
+      "image_masks_loss": 0.0175,
+      "llm_learning_rate": 4.679774057491851e-06,
+      "loss": 0.8169,
+      "output_text_mask_loss": 0.0225,
+      "predictor_learning_rate": 0.0004679774057491851,
+      "step": 1835
+    },
+    {
+      "epoch": 0.19,
+      "gumbel_tau": 0.6486,
+      "image_masks_loss": 0.0359,
+      "llm_learning_rate": 4.679367701677414e-06,
+      "loss": 0.7595,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 0.0004679367701677414,
+      "step": 1836
+    },
+    {
+      "epoch": 0.19,
+      "gumbel_tau": 0.6485,
+      "image_masks_loss": 0.0226,
+      "llm_learning_rate": 4.678961105866151e-06,
+      "loss": 0.8098,
+      "output_text_mask_loss": 0.0132,
+      "predictor_learning_rate": 0.0004678961105866151,
+      "step": 1837
+    },
+    {
+      "epoch": 0.19,
+      "gumbel_tau": 0.6483,
+      "image_masks_loss": 0.0235,
+      "llm_learning_rate": 4.678554270102837e-06,
+      "loss": 0.8731,
+      "output_text_mask_loss": 0.0235,
+      "predictor_learning_rate": 0.0004678554270102837,
+      "step": 1838
+    },
+    {
+      "epoch": 0.19,
+      "gumbel_tau": 0.6482,
+      "image_masks_loss": 0.0328,
+      "llm_learning_rate": 4.678147194432276e-06,
+      "loss": 0.763,
+      "output_text_mask_loss": 0.0324,
+      "predictor_learning_rate": 0.00046781471944322753,
+      "step": 1839
+    },
+    {
+      "epoch": 0.19,
+      "gumbel_tau": 0.648,
+      "image_masks_loss": 0.0332,
+      "llm_learning_rate": 4.6777398788992925e-06,
+      "loss": 0.7633,
+      "output_text_mask_loss": 0.0032,
+      "predictor_learning_rate": 0.0004677739878899292,
+      "step": 1840
+    },
+    {
+      "epoch": 0.19,
+      "gumbel_tau": 0.6479,
+      "image_masks_loss": 0.0313,
+      "llm_learning_rate": 4.677332323548743e-06,
+      "loss": 0.6917,
+      "output_text_mask_loss": 0.005,
+      "predictor_learning_rate": 0.0004677332323548743,
+      "step": 1841
+    },
+    {
+      "epoch": 0.19,
+      "gumbel_tau": 0.6477,
+      "image_masks_loss": 0.0469,
+      "llm_learning_rate": 4.676924528425508e-06,
+      "loss": 0.8171,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 0.0004676924528425508,
+      "step": 1842
+    },
+    {
+      "epoch": 0.19,
+      "gumbel_tau": 0.6475,
+      "image_masks_loss": 0.0127,
+      "llm_learning_rate": 4.676516493574496e-06,
+      "loss": 0.7293,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.0004676516493574496,
+      "step": 1843
+    },
+    {
+      "epoch": 0.19,
+      "gumbel_tau": 0.6474,
+      "image_masks_loss": 0.0206,
+      "llm_learning_rate": 4.67610821904064e-06,
+      "loss": 0.8505,
+      "output_text_mask_loss": 0.0032,
+      "predictor_learning_rate": 0.00046761082190406394,
+      "step": 1844
+    },
+    {
+      "epoch": 0.19,
+      "gumbel_tau": 0.6472,
+      "image_masks_loss": 0.0084,
+      "llm_learning_rate": 4.675699704868899e-06,
+      "loss": 0.7708,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004675699704868899,
+      "step": 1845
+    },
+    {
+      "epoch": 0.19,
+      "gumbel_tau": 0.6471,
+      "image_masks_loss": 0.011,
+      "llm_learning_rate": 4.67529095110426e-06,
+      "loss": 0.8027,
+      "output_text_mask_loss": 0.0156,
+      "predictor_learning_rate": 0.000467529095110426,
+      "step": 1846
+    },
+    {
+      "epoch": 0.19,
+      "gumbel_tau": 0.6469,
+      "image_masks_loss": 0.0123,
+      "llm_learning_rate": 4.674881957791737e-06,
+      "loss": 0.8292,
+      "output_text_mask_loss": 0.0565,
+      "predictor_learning_rate": 0.0004674881957791737,
+      "step": 1847
+    },
+    {
+      "epoch": 0.19,
+      "gumbel_tau": 0.6468,
+      "image_masks_loss": 0.0332,
+      "llm_learning_rate": 4.6744727249763686e-06,
+      "loss": 0.7795,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.0004674472724976368,
+      "step": 1848
+    },
+    {
+      "epoch": 0.19,
+      "gumbel_tau": 0.6466,
+      "image_masks_loss": 0.0265,
+      "llm_learning_rate": 4.674063252703219e-06,
+      "loss": 0.8672,
+      "output_text_mask_loss": 0.0416,
+      "predictor_learning_rate": 0.00046740632527032187,
+      "step": 1849
+    },
+    {
+      "epoch": 0.19,
+      "gumbel_tau": 0.6465,
+      "image_masks_loss": 0.0221,
+      "llm_learning_rate": 4.6736535410173826e-06,
+      "loss": 0.8104,
+      "output_text_mask_loss": 0.0034,
+      "predictor_learning_rate": 0.00046736535410173824,
+      "step": 1850
+    },
+    {
+      "epoch": 0.19,
+      "gumbel_tau": 0.6463,
+      "image_masks_loss": 0.0116,
+      "llm_learning_rate": 4.673243589963977e-06,
+      "loss": 0.7809,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00046732435899639763,
+      "step": 1851
+    },
+    {
+      "epoch": 0.19,
+      "gumbel_tau": 0.6462,
+      "image_masks_loss": 0.0178,
+      "llm_learning_rate": 4.672833399588146e-06,
+      "loss": 0.742,
+      "output_text_mask_loss": 0.0427,
+      "predictor_learning_rate": 0.00046728333995881455,
+      "step": 1852
+    },
+    {
+      "epoch": 0.19,
+      "gumbel_tau": 0.646,
+      "image_masks_loss": 0.0047,
+      "llm_learning_rate": 4.672422969935061e-06,
+      "loss": 0.7046,
+      "output_text_mask_loss": 0.0019,
+      "predictor_learning_rate": 0.00046724229699350605,
+      "step": 1853
+    },
+    {
+      "epoch": 0.19,
+      "gumbel_tau": 0.6459,
+      "image_masks_loss": 0.0149,
+      "llm_learning_rate": 4.67201230104992e-06,
+      "loss": 0.6683,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 0.00046720123010499196,
+      "step": 1854
+    },
+    {
+      "epoch": 0.19,
+      "gumbel_tau": 0.6457,
+      "image_masks_loss": 0.0443,
+      "llm_learning_rate": 4.6716013929779455e-06,
+      "loss": 0.7876,
+      "output_text_mask_loss": 0.0446,
+      "predictor_learning_rate": 0.00046716013929779454,
+      "step": 1855
+    },
+    {
+      "epoch": 0.19,
+      "gumbel_tau": 0.6456,
+      "image_masks_loss": 0.0102,
+      "llm_learning_rate": 4.67119024576439e-06,
+      "loss": 0.8093,
+      "output_text_mask_loss": 0.0748,
+      "predictor_learning_rate": 0.00046711902457643897,
+      "step": 1856
+    },
+    {
+      "epoch": 0.19,
+      "gumbel_tau": 0.6454,
+      "image_masks_loss": 0.0133,
+      "llm_learning_rate": 4.670778859454528e-06,
+      "loss": 0.7594,
+      "output_text_mask_loss": 0.005,
+      "predictor_learning_rate": 0.00046707788594545273,
+      "step": 1857
+    },
+    {
+      "epoch": 0.19,
+      "gumbel_tau": 0.6453,
+      "image_masks_loss": 0.0243,
+      "llm_learning_rate": 4.670367234093662e-06,
+      "loss": 0.6773,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004670367234093662,
+      "step": 1858
+    },
+    {
+      "epoch": 0.19,
+      "gumbel_tau": 0.6451,
+      "image_masks_loss": 0.0257,
+      "llm_learning_rate": 4.669955369727123e-06,
+      "loss": 0.883,
+      "output_text_mask_loss": 0.0322,
+      "predictor_learning_rate": 0.0004669955369727122,
+      "step": 1859
+    },
+    {
+      "epoch": 0.19,
+      "gumbel_tau": 0.645,
+      "image_masks_loss": 0.014,
+      "llm_learning_rate": 4.6695432664002645e-06,
+      "loss": 0.7938,
+      "output_text_mask_loss": 0.0019,
+      "predictor_learning_rate": 0.0004669543266400264,
+      "step": 1860
+    },
+    {
+      "epoch": 0.19,
+      "gumbel_tau": 0.6448,
+      "image_masks_loss": 0.0122,
+      "llm_learning_rate": 4.6691309241584695e-06,
+      "loss": 0.7776,
+      "output_text_mask_loss": 0.024,
+      "predictor_learning_rate": 0.0004669130924158469,
+      "step": 1861
+    },
+    {
+      "epoch": 0.19,
+      "gumbel_tau": 0.6446,
+      "image_masks_loss": 0.0254,
+      "llm_learning_rate": 4.668718343047145e-06,
+      "loss": 0.6627,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00046687183430471445,
+      "step": 1862
+    },
+    {
+      "epoch": 0.19,
+      "gumbel_tau": 0.6445,
+      "image_masks_loss": 0.0113,
+      "llm_learning_rate": 4.668305523111727e-06,
+      "loss": 0.7985,
+      "output_text_mask_loss": 0.0305,
+      "predictor_learning_rate": 0.0004668305523111726,
+      "step": 1863
+    },
+    {
+      "epoch": 0.19,
+      "gumbel_tau": 0.6443,
+      "image_masks_loss": 0.0254,
+      "llm_learning_rate": 4.667892464397674e-06,
+      "loss": 0.8252,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004667892464397674,
+      "step": 1864
+    },
+    {
+      "epoch": 0.19,
+      "gumbel_tau": 0.6442,
+      "image_masks_loss": 0.0151,
+      "llm_learning_rate": 4.667479166950474e-06,
+      "loss": 0.7135,
+      "output_text_mask_loss": 0.0132,
+      "predictor_learning_rate": 0.0004667479166950474,
+      "step": 1865
+    },
+    {
+      "epoch": 0.19,
+      "gumbel_tau": 0.644,
+      "image_masks_loss": 0.0163,
+      "llm_learning_rate": 4.667065630815641e-06,
+      "loss": 0.8472,
+      "output_text_mask_loss": 0.0019,
+      "predictor_learning_rate": 0.00046670656308156407,
+      "step": 1866
+    },
+    {
+      "epoch": 0.19,
+      "gumbel_tau": 0.6439,
+      "image_masks_loss": 0.0087,
+      "llm_learning_rate": 4.666651856038713e-06,
+      "loss": 0.7587,
+      "output_text_mask_loss": 0.0572,
+      "predictor_learning_rate": 0.0004666651856038713,
+      "step": 1867
+    },
+    {
+      "epoch": 0.19,
+      "gumbel_tau": 0.6437,
+      "image_masks_loss": 0.014,
+      "llm_learning_rate": 4.666237842665257e-06,
+      "loss": 0.6855,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00046662378426652574,
+      "step": 1868
+    },
+    {
+      "epoch": 0.19,
+      "gumbel_tau": 0.6436,
+      "image_masks_loss": 0.0145,
+      "llm_learning_rate": 4.6658235907408655e-06,
+      "loss": 0.7147,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00046658235907408654,
+      "step": 1869
+    },
+    {
+      "epoch": 0.19,
+      "gumbel_tau": 0.6434,
+      "image_masks_loss": 0.0229,
+      "llm_learning_rate": 4.665409100311155e-06,
+      "loss": 0.7827,
+      "output_text_mask_loss": 0.0687,
+      "predictor_learning_rate": 0.0004665409100311155,
+      "step": 1870
+    },
+    {
+      "epoch": 0.19,
+      "gumbel_tau": 0.6433,
+      "image_masks_loss": 0.0153,
+      "llm_learning_rate": 4.664994371421772e-06,
+      "loss": 0.7903,
+      "output_text_mask_loss": 0.0172,
+      "predictor_learning_rate": 0.0004664994371421771,
+      "step": 1871
+    },
+    {
+      "epoch": 0.19,
+      "gumbel_tau": 0.6431,
+      "image_masks_loss": 0.0128,
+      "llm_learning_rate": 4.664579404118386e-06,
+      "loss": 0.8363,
+      "output_text_mask_loss": 0.0399,
+      "predictor_learning_rate": 0.00046645794041183855,
+      "step": 1872
+    },
+    {
+      "epoch": 0.19,
+      "gumbel_tau": 0.643,
+      "image_masks_loss": 0.0351,
+      "llm_learning_rate": 4.664164198446693e-06,
+      "loss": 0.7976,
+      "output_text_mask_loss": 0.0202,
+      "predictor_learning_rate": 0.0004664164198446693,
+      "step": 1873
+    },
+    {
+      "epoch": 0.19,
+      "gumbel_tau": 0.6428,
+      "image_masks_loss": 0.017,
+      "llm_learning_rate": 4.66374875445242e-06,
+      "loss": 0.758,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.000466374875445242,
+      "step": 1874
+    },
+    {
+      "epoch": 0.19,
+      "gumbel_tau": 0.6427,
+      "image_masks_loss": 0.0208,
+      "llm_learning_rate": 4.663333072181314e-06,
+      "loss": 0.7542,
+      "output_text_mask_loss": 0.0465,
+      "predictor_learning_rate": 0.0004663333072181314,
+      "step": 1875
+    },
+    {
+      "epoch": 0.19,
+      "gumbel_tau": 0.6425,
+      "image_masks_loss": 0.0284,
+      "llm_learning_rate": 4.662917151679152e-06,
+      "loss": 0.8219,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00046629171516791515,
+      "step": 1876
+    },
+    {
+      "epoch": 0.19,
+      "gumbel_tau": 0.6424,
+      "image_masks_loss": 0.0105,
+      "llm_learning_rate": 4.662500992991735e-06,
+      "loss": 0.7208,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.00046625009929917346,
+      "step": 1877
+    },
+    {
+      "epoch": 0.19,
+      "gumbel_tau": 0.6422,
+      "image_masks_loss": 0.0131,
+      "llm_learning_rate": 4.662084596164891e-06,
+      "loss": 0.7158,
+      "output_text_mask_loss": 0.0193,
+      "predictor_learning_rate": 0.0004662084596164891,
+      "step": 1878
+    },
+    {
+      "epoch": 0.19,
+      "gumbel_tau": 0.6421,
+      "image_masks_loss": 0.0364,
+      "llm_learning_rate": 4.661667961244477e-06,
+      "loss": 0.7543,
+      "output_text_mask_loss": 0.0229,
+      "predictor_learning_rate": 0.0004661667961244477,
+      "step": 1879
+    },
+    {
+      "epoch": 0.19,
+      "gumbel_tau": 0.6419,
+      "image_masks_loss": 0.0261,
+      "llm_learning_rate": 4.661251088276372e-06,
+      "loss": 0.6799,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.00046612510882763713,
+      "step": 1880
+    },
+    {
+      "epoch": 0.19,
+      "gumbel_tau": 0.6418,
+      "image_masks_loss": 0.0267,
+      "llm_learning_rate": 4.660833977306482e-06,
+      "loss": 0.7134,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00046608339773064817,
+      "step": 1881
+    },
+    {
+      "epoch": 0.19,
+      "gumbel_tau": 0.6416,
+      "image_masks_loss": 0.0171,
+      "llm_learning_rate": 4.660416628380741e-06,
+      "loss": 0.7656,
+      "output_text_mask_loss": 0.0055,
+      "predictor_learning_rate": 0.00046604166283807413,
+      "step": 1882
+    },
+    {
+      "epoch": 0.19,
+      "gumbel_tau": 0.6415,
+      "image_masks_loss": 0.0244,
+      "llm_learning_rate": 4.659999041545111e-06,
+      "loss": 0.7671,
+      "output_text_mask_loss": 0.0012,
+      "predictor_learning_rate": 0.00046599990415451106,
+      "step": 1883
+    },
+    {
+      "epoch": 0.19,
+      "gumbel_tau": 0.6413,
+      "image_masks_loss": 0.0282,
+      "llm_learning_rate": 4.659581216845574e-06,
+      "loss": 0.847,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.0004659581216845574,
+      "step": 1884
+    },
+    {
+      "epoch": 0.19,
+      "gumbel_tau": 0.6412,
+      "image_masks_loss": 0.033,
+      "llm_learning_rate": 4.659163154328143e-06,
+      "loss": 0.7383,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.00046591631543281434,
+      "step": 1885
+    },
+    {
+      "epoch": 0.19,
+      "gumbel_tau": 0.641,
+      "image_masks_loss": 0.0322,
+      "llm_learning_rate": 4.658744854038856e-06,
+      "loss": 0.7427,
+      "output_text_mask_loss": 0.033,
+      "predictor_learning_rate": 0.0004658744854038856,
+      "step": 1886
+    },
+    {
+      "epoch": 0.19,
+      "gumbel_tau": 0.6409,
+      "image_masks_loss": 0.0137,
+      "llm_learning_rate": 4.658326316023778e-06,
+      "loss": 0.8204,
+      "output_text_mask_loss": 0.0158,
+      "predictor_learning_rate": 0.0004658326316023778,
+      "step": 1887
+    },
+    {
+      "epoch": 0.19,
+      "gumbel_tau": 0.6407,
+      "image_masks_loss": 0.0125,
+      "llm_learning_rate": 4.657907540328998e-06,
+      "loss": 0.733,
+      "output_text_mask_loss": 0.0107,
+      "predictor_learning_rate": 0.00046579075403289976,
+      "step": 1888
+    },
+    {
+      "epoch": 0.19,
+      "gumbel_tau": 0.6406,
+      "image_masks_loss": 0.0203,
+      "llm_learning_rate": 4.657488527000633e-06,
+      "loss": 0.6316,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00046574885270006327,
+      "step": 1889
+    },
+    {
+      "epoch": 0.19,
+      "gumbel_tau": 0.6404,
+      "image_masks_loss": 0.0153,
+      "llm_learning_rate": 4.657069276084826e-06,
+      "loss": 0.6904,
+      "output_text_mask_loss": 0.0124,
+      "predictor_learning_rate": 0.0004657069276084826,
+      "step": 1890
+    },
+    {
+      "epoch": 0.19,
+      "gumbel_tau": 0.6403,
+      "image_masks_loss": 0.033,
+      "llm_learning_rate": 4.656649787627746e-06,
+      "loss": 0.8403,
+      "output_text_mask_loss": 0.0032,
+      "predictor_learning_rate": 0.00046566497876277457,
+      "step": 1891
+    },
+    {
+      "epoch": 0.19,
+      "gumbel_tau": 0.6401,
+      "image_masks_loss": 0.0197,
+      "llm_learning_rate": 4.656230061675588e-06,
+      "loss": 0.6778,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.00046562300616755873,
+      "step": 1892
+    },
+    {
+      "epoch": 0.19,
+      "gumbel_tau": 0.64,
+      "image_masks_loss": 0.015,
+      "llm_learning_rate": 4.655810098274571e-06,
+      "loss": 0.7802,
+      "output_text_mask_loss": 0.0039,
+      "predictor_learning_rate": 0.0004655810098274571,
+      "step": 1893
+    },
+    {
+      "epoch": 0.19,
+      "gumbel_tau": 0.6398,
+      "image_masks_loss": 0.0123,
+      "llm_learning_rate": 4.655389897470946e-06,
+      "loss": 0.7568,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.00046553898974709454,
+      "step": 1894
+    },
+    {
+      "epoch": 0.19,
+      "gumbel_tau": 0.6397,
+      "image_masks_loss": 0.0393,
+      "llm_learning_rate": 4.654969459310984e-06,
+      "loss": 0.8064,
+      "output_text_mask_loss": 0.0021,
+      "predictor_learning_rate": 0.0004654969459310983,
+      "step": 1895
+    },
+    {
+      "epoch": 0.19,
+      "gumbel_tau": 0.6395,
+      "image_masks_loss": 0.0168,
+      "llm_learning_rate": 4.654548783840984e-06,
+      "loss": 0.7087,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.0004654548783840984,
+      "step": 1896
+    },
+    {
+      "epoch": 0.19,
+      "gumbel_tau": 0.6393,
+      "image_masks_loss": 0.0123,
+      "llm_learning_rate": 4.654127871107274e-06,
+      "loss": 0.6071,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00046541278711072735,
+      "step": 1897
+    },
+    {
+      "epoch": 0.19,
+      "gumbel_tau": 0.6392,
+      "image_masks_loss": 0.0254,
+      "llm_learning_rate": 4.653706721156205e-06,
+      "loss": 0.6754,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.00046537067211562043,
+      "step": 1898
+    },
+    {
+      "epoch": 0.19,
+      "gumbel_tau": 0.639,
+      "image_masks_loss": 0.0216,
+      "llm_learning_rate": 4.6532853340341536e-06,
+      "loss": 0.7799,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.00046532853340341534,
+      "step": 1899
+    },
+    {
+      "epoch": 0.19,
+      "gumbel_tau": 0.6389,
+      "image_masks_loss": 0.0185,
+      "llm_learning_rate": 4.652863709787526e-06,
+      "loss": 0.7235,
+      "output_text_mask_loss": 0.0231,
+      "predictor_learning_rate": 0.0004652863709787525,
+      "step": 1900
+    },
+    {
+      "epoch": 0.19,
+      "gumbel_tau": 0.6387,
+      "image_masks_loss": 0.0153,
+      "llm_learning_rate": 4.65244184846275e-06,
+      "loss": 0.7824,
+      "output_text_mask_loss": 0.0374,
+      "predictor_learning_rate": 0.000465244184846275,
+      "step": 1901
+    },
+    {
+      "epoch": 0.19,
+      "gumbel_tau": 0.6386,
+      "image_masks_loss": 0.0123,
+      "llm_learning_rate": 4.652019750106284e-06,
+      "loss": 0.7428,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00046520197501062843,
+      "step": 1902
+    },
+    {
+      "epoch": 0.19,
+      "gumbel_tau": 0.6384,
+      "image_masks_loss": 0.0288,
+      "llm_learning_rate": 4.6515974147646105e-06,
+      "loss": 0.8374,
+      "output_text_mask_loss": 0.0065,
+      "predictor_learning_rate": 0.000465159741476461,
+      "step": 1903
+    },
+    {
+      "epoch": 0.2,
+      "gumbel_tau": 0.6383,
+      "image_masks_loss": 0.0116,
+      "llm_learning_rate": 4.6511748424842365e-06,
+      "loss": 0.8426,
+      "output_text_mask_loss": 0.0041,
+      "predictor_learning_rate": 0.0004651174842484236,
+      "step": 1904
+    },
+    {
+      "epoch": 0.2,
+      "gumbel_tau": 0.6381,
+      "image_masks_loss": 0.029,
+      "llm_learning_rate": 4.6507520333116975e-06,
+      "loss": 0.7878,
+      "output_text_mask_loss": 0.014,
+      "predictor_learning_rate": 0.0004650752033311697,
+      "step": 1905
+    },
+    {
+      "epoch": 0.2,
+      "gumbel_tau": 0.638,
+      "image_masks_loss": 0.0207,
+      "llm_learning_rate": 4.650328987293553e-06,
+      "loss": 0.829,
+      "output_text_mask_loss": 0.0549,
+      "predictor_learning_rate": 0.0004650328987293553,
+      "step": 1906
+    },
+    {
+      "epoch": 0.2,
+      "gumbel_tau": 0.6378,
+      "image_masks_loss": 0.0259,
+      "llm_learning_rate": 4.649905704476392e-06,
+      "loss": 0.8267,
+      "output_text_mask_loss": 0.0374,
+      "predictor_learning_rate": 0.00046499057044763914,
+      "step": 1907
+    },
+    {
+      "epoch": 0.2,
+      "gumbel_tau": 0.6377,
+      "image_masks_loss": 0.0054,
+      "llm_learning_rate": 4.649482184906826e-06,
+      "loss": 0.7734,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00046494821849068254,
+      "step": 1908
+    },
+    {
+      "epoch": 0.2,
+      "gumbel_tau": 0.6375,
+      "image_masks_loss": 0.0133,
+      "llm_learning_rate": 4.649058428631493e-06,
+      "loss": 0.7707,
+      "output_text_mask_loss": 0.013,
+      "predictor_learning_rate": 0.0004649058428631493,
+      "step": 1909
+    },
+    {
+      "epoch": 0.2,
+      "gumbel_tau": 0.6374,
+      "image_masks_loss": 0.0261,
+      "llm_learning_rate": 4.648634435697059e-06,
+      "loss": 0.7676,
+      "output_text_mask_loss": 0.0185,
+      "predictor_learning_rate": 0.00046486344356970596,
+      "step": 1910
+    },
+    {
+      "epoch": 0.2,
+      "gumbel_tau": 0.6372,
+      "image_masks_loss": 0.0145,
+      "llm_learning_rate": 4.648210206150217e-06,
+      "loss": 0.7218,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.00046482102061502165,
+      "step": 1911
+    },
+    {
+      "epoch": 0.2,
+      "gumbel_tau": 0.6371,
+      "image_masks_loss": 0.0288,
+      "llm_learning_rate": 4.64778574003768e-06,
+      "loss": 0.7617,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.000464778574003768,
+      "step": 1912
+    },
+    {
+      "epoch": 0.2,
+      "gumbel_tau": 0.6369,
+      "image_masks_loss": 0.013,
+      "llm_learning_rate": 4.647361037406194e-06,
+      "loss": 0.7626,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 0.00046473610374061944,
+      "step": 1913
+    },
+    {
+      "epoch": 0.2,
+      "gumbel_tau": 0.6368,
+      "image_masks_loss": 0.0093,
+      "llm_learning_rate": 4.646936098302528e-06,
+      "loss": 0.7335,
+      "output_text_mask_loss": 0.0012,
+      "predictor_learning_rate": 0.0004646936098302528,
+      "step": 1914
+    },
+    {
+      "epoch": 0.2,
+      "gumbel_tau": 0.6366,
+      "image_masks_loss": 0.0093,
+      "llm_learning_rate": 4.6465109227734765e-06,
+      "loss": 0.7819,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.0004646510922773476,
+      "step": 1915
+    },
+    {
+      "epoch": 0.2,
+      "gumbel_tau": 0.6365,
+      "image_masks_loss": 0.0599,
+      "llm_learning_rate": 4.646085510865861e-06,
+      "loss": 0.8006,
+      "output_text_mask_loss": 0.0185,
+      "predictor_learning_rate": 0.00046460855108658607,
+      "step": 1916
+    },
+    {
+      "epoch": 0.2,
+      "gumbel_tau": 0.6363,
+      "image_masks_loss": 0.0119,
+      "llm_learning_rate": 4.645659862626528e-06,
+      "loss": 0.8174,
+      "output_text_mask_loss": 0.0141,
+      "predictor_learning_rate": 0.00046456598626265277,
+      "step": 1917
+    },
+    {
+      "epoch": 0.2,
+      "gumbel_tau": 0.6362,
+      "image_masks_loss": 0.0203,
+      "llm_learning_rate": 4.645233978102352e-06,
+      "loss": 0.6839,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004645233978102352,
+      "step": 1918
+    },
+    {
+      "epoch": 0.2,
+      "gumbel_tau": 0.636,
+      "image_masks_loss": 0.0187,
+      "llm_learning_rate": 4.644807857340234e-06,
+      "loss": 0.6668,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004644807857340233,
+      "step": 1919
+    },
+    {
+      "epoch": 0.2,
+      "gumbel_tau": 0.6359,
+      "image_masks_loss": 0.009,
+      "llm_learning_rate": 4.644381500387095e-06,
+      "loss": 0.8336,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004644381500387095,
+      "step": 1920
+    },
+    {
+      "epoch": 0.2,
+      "gumbel_tau": 0.6357,
+      "image_masks_loss": 0.0103,
+      "llm_learning_rate": 4.643954907289889e-06,
+      "loss": 0.6488,
+      "output_text_mask_loss": 0.0549,
+      "predictor_learning_rate": 0.0004643954907289889,
+      "step": 1921
+    },
+    {
+      "epoch": 0.2,
+      "gumbel_tau": 0.6356,
+      "image_masks_loss": 0.036,
+      "llm_learning_rate": 4.643528078095594e-06,
+      "loss": 0.8235,
+      "output_text_mask_loss": 0.0055,
+      "predictor_learning_rate": 0.00046435280780955933,
+      "step": 1922
+    },
+    {
+      "epoch": 0.2,
+      "gumbel_tau": 0.6354,
+      "image_masks_loss": 0.0454,
+      "llm_learning_rate": 4.6431010128512126e-06,
+      "loss": 0.8309,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.0004643101012851212,
+      "step": 1923
+    },
+    {
+      "epoch": 0.2,
+      "gumbel_tau": 0.6353,
+      "image_masks_loss": 0.0401,
+      "llm_learning_rate": 4.642673711603773e-06,
+      "loss": 0.7337,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.0004642673711603773,
+      "step": 1924
+    },
+    {
+      "epoch": 0.2,
+      "gumbel_tau": 0.6351,
+      "image_masks_loss": 0.0064,
+      "llm_learning_rate": 4.6422461744003325e-06,
+      "loss": 0.7474,
+      "output_text_mask_loss": 0.0614,
+      "predictor_learning_rate": 0.0004642246174400332,
+      "step": 1925
+    },
+    {
+      "epoch": 0.2,
+      "gumbel_tau": 0.635,
+      "image_masks_loss": 0.0418,
+      "llm_learning_rate": 4.641818401287972e-06,
+      "loss": 0.7113,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.0004641818401287972,
+      "step": 1926
+    },
+    {
+      "epoch": 0.2,
+      "gumbel_tau": 0.6348,
+      "image_masks_loss": 0.0185,
+      "llm_learning_rate": 4.641390392313797e-06,
+      "loss": 0.7701,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 0.00046413903923137967,
+      "step": 1927
+    },
+    {
+      "epoch": 0.2,
+      "gumbel_tau": 0.6347,
+      "image_masks_loss": 0.0261,
+      "llm_learning_rate": 4.640962147524944e-06,
+      "loss": 0.7182,
+      "output_text_mask_loss": 0.0071,
+      "predictor_learning_rate": 0.0004640962147524943,
+      "step": 1928
+    },
+    {
+      "epoch": 0.2,
+      "gumbel_tau": 0.6345,
+      "image_masks_loss": 0.0637,
+      "llm_learning_rate": 4.640533666968569e-06,
+      "loss": 0.7912,
+      "output_text_mask_loss": 0.0351,
+      "predictor_learning_rate": 0.0004640533666968569,
+      "step": 1929
+    },
+    {
+      "epoch": 0.2,
+      "gumbel_tau": 0.6344,
+      "image_masks_loss": 0.0209,
+      "llm_learning_rate": 4.640104950691859e-06,
+      "loss": 0.6675,
+      "output_text_mask_loss": 0.0081,
+      "predictor_learning_rate": 0.0004640104950691859,
+      "step": 1930
+    },
+    {
+      "epoch": 0.2,
+      "gumbel_tau": 0.6342,
+      "image_masks_loss": 0.0035,
+      "llm_learning_rate": 4.639675998742026e-06,
+      "loss": 0.7403,
+      "output_text_mask_loss": 0.0305,
+      "predictor_learning_rate": 0.0004639675998742026,
+      "step": 1931
+    },
+    {
+      "epoch": 0.2,
+      "gumbel_tau": 0.6341,
+      "image_masks_loss": 0.0175,
+      "llm_learning_rate": 4.639246811166306e-06,
+      "loss": 0.8189,
+      "output_text_mask_loss": 0.0496,
+      "predictor_learning_rate": 0.0004639246811166305,
+      "step": 1932
+    },
+    {
+      "epoch": 0.2,
+      "gumbel_tau": 0.6339,
+      "image_masks_loss": 0.0378,
+      "llm_learning_rate": 4.6388173880119605e-06,
+      "loss": 0.6893,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.000463881738801196,
+      "step": 1933
+    },
+    {
+      "epoch": 0.2,
+      "gumbel_tau": 0.6338,
+      "image_masks_loss": 0.0057,
+      "llm_learning_rate": 4.6383877293262815e-06,
+      "loss": 0.7915,
+      "output_text_mask_loss": 0.0477,
+      "predictor_learning_rate": 0.0004638387729326281,
+      "step": 1934
+    },
+    {
+      "epoch": 0.2,
+      "gumbel_tau": 0.6336,
+      "image_masks_loss": 0.019,
+      "llm_learning_rate": 4.6379578351565824e-06,
+      "loss": 0.7679,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 0.0004637957835156582,
+      "step": 1935
+    },
+    {
+      "epoch": 0.2,
+      "gumbel_tau": 0.6335,
+      "image_masks_loss": 0.0269,
+      "llm_learning_rate": 4.637527705550205e-06,
+      "loss": 0.6877,
+      "output_text_mask_loss": 0.0618,
+      "predictor_learning_rate": 0.0004637527705550204,
+      "step": 1936
+    },
+    {
+      "epoch": 0.2,
+      "gumbel_tau": 0.6333,
+      "image_masks_loss": 0.0078,
+      "llm_learning_rate": 4.637097340554514e-06,
+      "loss": 0.7699,
+      "output_text_mask_loss": 0.0078,
+      "predictor_learning_rate": 0.0004637097340554514,
+      "step": 1937
+    },
+    {
+      "epoch": 0.2,
+      "gumbel_tau": 0.6332,
+      "image_masks_loss": 0.021,
+      "llm_learning_rate": 4.636666740216905e-06,
+      "loss": 0.8087,
+      "output_text_mask_loss": 0.0208,
+      "predictor_learning_rate": 0.00046366667402169046,
+      "step": 1938
+    },
+    {
+      "epoch": 0.2,
+      "gumbel_tau": 0.633,
+      "image_masks_loss": 0.0227,
+      "llm_learning_rate": 4.6362359045847945e-06,
+      "loss": 0.8202,
+      "output_text_mask_loss": 0.0248,
+      "predictor_learning_rate": 0.0004636235904584794,
+      "step": 1939
+    },
+    {
+      "epoch": 0.2,
+      "gumbel_tau": 0.6329,
+      "image_masks_loss": 0.02,
+      "llm_learning_rate": 4.635804833705628e-06,
+      "loss": 0.8123,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004635804833705628,
+      "step": 1940
+    },
+    {
+      "epoch": 0.2,
+      "gumbel_tau": 0.6327,
+      "image_masks_loss": 0.0246,
+      "llm_learning_rate": 4.635373527626876e-06,
+      "loss": 0.7841,
+      "output_text_mask_loss": 0.0029,
+      "predictor_learning_rate": 0.0004635373527626876,
+      "step": 1941
+    },
+    {
+      "epoch": 0.2,
+      "gumbel_tau": 0.6326,
+      "image_masks_loss": 0.0166,
+      "llm_learning_rate": 4.634941986396035e-06,
+      "loss": 0.7543,
+      "output_text_mask_loss": 0.046,
+      "predictor_learning_rate": 0.0004634941986396035,
+      "step": 1942
+    },
+    {
+      "epoch": 0.2,
+      "gumbel_tau": 0.6324,
+      "image_masks_loss": 0.0163,
+      "llm_learning_rate": 4.634510210060628e-06,
+      "loss": 0.7543,
+      "output_text_mask_loss": 0.0416,
+      "predictor_learning_rate": 0.00046345102100606275,
+      "step": 1943
+    },
+    {
+      "epoch": 0.2,
+      "gumbel_tau": 0.6323,
+      "image_masks_loss": 0.015,
+      "llm_learning_rate": 4.634078198668201e-06,
+      "loss": 0.7295,
+      "output_text_mask_loss": 0.0104,
+      "predictor_learning_rate": 0.00046340781986682,
+      "step": 1944
+    },
+    {
+      "epoch": 0.2,
+      "gumbel_tau": 0.6321,
+      "image_masks_loss": 0.0041,
+      "llm_learning_rate": 4.633645952266329e-06,
+      "loss": 0.7156,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00046336459522663287,
+      "step": 1945
+    },
+    {
+      "epoch": 0.2,
+      "gumbel_tau": 0.632,
+      "image_masks_loss": 0.0179,
+      "llm_learning_rate": 4.6332134709026135e-06,
+      "loss": 0.7388,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004633213470902613,
+      "step": 1946
+    },
+    {
+      "epoch": 0.2,
+      "gumbel_tau": 0.6319,
+      "image_masks_loss": 0.0256,
+      "llm_learning_rate": 4.632780754624678e-06,
+      "loss": 0.8127,
+      "output_text_mask_loss": 0.0233,
+      "predictor_learning_rate": 0.00046327807546246773,
+      "step": 1947
+    },
+    {
+      "epoch": 0.2,
+      "gumbel_tau": 0.6317,
+      "image_masks_loss": 0.013,
+      "llm_learning_rate": 4.632347803480175e-06,
+      "loss": 0.7353,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00046323478034801746,
+      "step": 1948
+    },
+    {
+      "epoch": 0.2,
+      "gumbel_tau": 0.6316,
+      "image_masks_loss": 0.01,
+      "llm_learning_rate": 4.631914617516783e-06,
+      "loss": 0.7644,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.00046319146175167824,
+      "step": 1949
+    },
+    {
+      "epoch": 0.2,
+      "gumbel_tau": 0.6314,
+      "image_masks_loss": 0.0085,
+      "llm_learning_rate": 4.631481196782205e-06,
+      "loss": 0.6595,
+      "output_text_mask_loss": 0.0056,
+      "predictor_learning_rate": 0.00046314811967822044,
+      "step": 1950
+    },
+    {
+      "epoch": 0.2,
+      "gumbel_tau": 0.6313,
+      "image_masks_loss": 0.0237,
+      "llm_learning_rate": 4.631047541324169e-06,
+      "loss": 0.779,
+      "output_text_mask_loss": 0.0374,
+      "predictor_learning_rate": 0.0004631047541324169,
+      "step": 1951
+    },
+    {
+      "epoch": 0.2,
+      "gumbel_tau": 0.6311,
+      "image_masks_loss": 0.0324,
+      "llm_learning_rate": 4.630613651190431e-06,
+      "loss": 0.7195,
+      "output_text_mask_loss": 0.1823,
+      "predictor_learning_rate": 0.0004630613651190431,
+      "step": 1952
+    },
+    {
+      "epoch": 0.2,
+      "gumbel_tau": 0.631,
+      "image_masks_loss": 0.0462,
+      "llm_learning_rate": 4.630179526428773e-06,
+      "loss": 0.7604,
+      "output_text_mask_loss": 0.0309,
+      "predictor_learning_rate": 0.00046301795264287726,
+      "step": 1953
+    },
+    {
+      "epoch": 0.2,
+      "gumbel_tau": 0.6308,
+      "image_masks_loss": 0.0362,
+      "llm_learning_rate": 4.6297451670869996e-06,
+      "loss": 0.7347,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.0004629745167086999,
+      "step": 1954
+    },
+    {
+      "epoch": 0.2,
+      "gumbel_tau": 0.6307,
+      "image_masks_loss": 0.0065,
+      "llm_learning_rate": 4.629310573212945e-06,
+      "loss": 0.7855,
+      "output_text_mask_loss": 0.0238,
+      "predictor_learning_rate": 0.0004629310573212945,
+      "step": 1955
+    },
+    {
+      "epoch": 0.2,
+      "gumbel_tau": 0.6305,
+      "image_masks_loss": 0.0044,
+      "llm_learning_rate": 4.628875744854467e-06,
+      "loss": 0.8291,
+      "output_text_mask_loss": 0.0511,
+      "predictor_learning_rate": 0.0004628875744854467,
+      "step": 1956
+    },
+    {
+      "epoch": 0.2,
+      "gumbel_tau": 0.6304,
+      "image_masks_loss": 0.0196,
+      "llm_learning_rate": 4.628440682059451e-06,
+      "loss": 0.6701,
+      "output_text_mask_loss": 0.0242,
+      "predictor_learning_rate": 0.00046284406820594503,
+      "step": 1957
+    },
+    {
+      "epoch": 0.2,
+      "gumbel_tau": 0.6302,
+      "image_masks_loss": 0.0146,
+      "llm_learning_rate": 4.6280053848758054e-06,
+      "loss": 0.818,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004628005384875805,
+      "step": 1958
+    },
+    {
+      "epoch": 0.2,
+      "gumbel_tau": 0.6301,
+      "image_masks_loss": 0.024,
+      "llm_learning_rate": 4.627569853351468e-06,
+      "loss": 0.7294,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 0.00046275698533514674,
+      "step": 1959
+    },
+    {
+      "epoch": 0.2,
+      "gumbel_tau": 0.6299,
+      "image_masks_loss": 0.0423,
+      "llm_learning_rate": 4.627134087534398e-06,
+      "loss": 0.7292,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 0.00046271340875343975,
+      "step": 1960
+    },
+    {
+      "epoch": 0.2,
+      "gumbel_tau": 0.6298,
+      "image_masks_loss": 0.0099,
+      "llm_learning_rate": 4.626698087472585e-06,
+      "loss": 0.7366,
+      "output_text_mask_loss": 0.013,
+      "predictor_learning_rate": 0.0004626698087472584,
+      "step": 1961
+    },
+    {
+      "epoch": 0.2,
+      "gumbel_tau": 0.6296,
+      "image_masks_loss": 0.0209,
+      "llm_learning_rate": 4.626261853214041e-06,
+      "loss": 0.5943,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00046262618532140405,
+      "step": 1962
+    },
+    {
+      "epoch": 0.2,
+      "gumbel_tau": 0.6295,
+      "image_masks_loss": 0.0381,
+      "llm_learning_rate": 4.625825384806806e-06,
+      "loss": 0.7493,
+      "output_text_mask_loss": 0.0004,
+      "predictor_learning_rate": 0.0004625825384806806,
+      "step": 1963
+    },
+    {
+      "epoch": 0.2,
+      "gumbel_tau": 0.6293,
+      "image_masks_loss": 0.0393,
+      "llm_learning_rate": 4.625388682298945e-06,
+      "loss": 0.8167,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.0004625388682298944,
+      "step": 1964
+    },
+    {
+      "epoch": 0.2,
+      "gumbel_tau": 0.6292,
+      "image_masks_loss": 0.0259,
+      "llm_learning_rate": 4.624951745738547e-06,
+      "loss": 0.7637,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.0004624951745738547,
+      "step": 1965
+    },
+    {
+      "epoch": 0.2,
+      "gumbel_tau": 0.629,
+      "image_masks_loss": 0.0167,
+      "llm_learning_rate": 4.62451457517373e-06,
+      "loss": 0.8527,
+      "output_text_mask_loss": 0.0139,
+      "predictor_learning_rate": 0.000462451457517373,
+      "step": 1966
+    },
+    {
+      "epoch": 0.2,
+      "gumbel_tau": 0.6289,
+      "image_masks_loss": 0.0187,
+      "llm_learning_rate": 4.624077170652637e-06,
+      "loss": 0.7203,
+      "output_text_mask_loss": 0.0012,
+      "predictor_learning_rate": 0.00046240771706526363,
+      "step": 1967
+    },
+    {
+      "epoch": 0.2,
+      "gumbel_tau": 0.6287,
+      "image_masks_loss": 0.0065,
+      "llm_learning_rate": 4.623639532223434e-06,
+      "loss": 0.8228,
+      "output_text_mask_loss": 0.0158,
+      "predictor_learning_rate": 0.0004623639532223433,
+      "step": 1968
+    },
+    {
+      "epoch": 0.2,
+      "gumbel_tau": 0.6286,
+      "image_masks_loss": 0.0031,
+      "llm_learning_rate": 4.623201659934314e-06,
+      "loss": 0.7492,
+      "output_text_mask_loss": 0.0322,
+      "predictor_learning_rate": 0.0004623201659934314,
+      "step": 1969
+    },
+    {
+      "epoch": 0.2,
+      "gumbel_tau": 0.6284,
+      "image_masks_loss": 0.018,
+      "llm_learning_rate": 4.6227635538335e-06,
+      "loss": 0.6951,
+      "output_text_mask_loss": 0.0162,
+      "predictor_learning_rate": 0.00046227635538335,
+      "step": 1970
+    },
+    {
+      "epoch": 0.2,
+      "gumbel_tau": 0.6283,
+      "image_masks_loss": 0.0454,
+      "llm_learning_rate": 4.622325213969234e-06,
+      "loss": 0.7705,
+      "output_text_mask_loss": 0.0387,
+      "predictor_learning_rate": 0.0004622325213969234,
+      "step": 1971
+    },
+    {
+      "epoch": 0.2,
+      "gumbel_tau": 0.6281,
+      "image_masks_loss": 0.0116,
+      "llm_learning_rate": 4.62188664038979e-06,
+      "loss": 0.6579,
+      "output_text_mask_loss": 0.0113,
+      "predictor_learning_rate": 0.0004621886640389789,
+      "step": 1972
+    },
+    {
+      "epoch": 0.2,
+      "gumbel_tau": 0.628,
+      "image_masks_loss": 0.0125,
+      "llm_learning_rate": 4.621447833143461e-06,
+      "loss": 0.8296,
+      "output_text_mask_loss": 0.0275,
+      "predictor_learning_rate": 0.00046214478331434606,
+      "step": 1973
+    },
+    {
+      "epoch": 0.2,
+      "gumbel_tau": 0.6278,
+      "image_masks_loss": 0.0131,
+      "llm_learning_rate": 4.621008792278572e-06,
+      "loss": 0.9745,
+      "output_text_mask_loss": 0.161,
+      "predictor_learning_rate": 0.00046210087922785714,
+      "step": 1974
+    },
+    {
+      "epoch": 0.2,
+      "gumbel_tau": 0.6277,
+      "image_masks_loss": 0.0154,
+      "llm_learning_rate": 4.620569517843469e-06,
+      "loss": 0.9967,
+      "output_text_mask_loss": 0.1007,
+      "predictor_learning_rate": 0.00046205695178434694,
+      "step": 1975
+    },
+    {
+      "epoch": 0.2,
+      "gumbel_tau": 0.6275,
+      "image_masks_loss": 0.0167,
+      "llm_learning_rate": 4.620130009886529e-06,
+      "loss": 1.2464,
+      "output_text_mask_loss": 0.3525,
+      "predictor_learning_rate": 0.0004620130009886529,
+      "step": 1976
+    },
+    {
+      "epoch": 0.2,
+      "gumbel_tau": 0.6274,
+      "image_masks_loss": 0.0246,
+      "llm_learning_rate": 4.61969026845615e-06,
+      "loss": 0.8836,
+      "output_text_mask_loss": 0.0429,
+      "predictor_learning_rate": 0.00046196902684561493,
+      "step": 1977
+    },
+    {
+      "epoch": 0.2,
+      "gumbel_tau": 0.6272,
+      "image_masks_loss": 0.0126,
+      "llm_learning_rate": 4.619250293600756e-06,
+      "loss": 0.9582,
+      "output_text_mask_loss": 0.0687,
+      "predictor_learning_rate": 0.00046192502936007564,
+      "step": 1978
+    },
+    {
+      "epoch": 0.2,
+      "gumbel_tau": 0.6271,
+      "image_masks_loss": 0.0125,
+      "llm_learning_rate": 4.6188100853688005e-06,
+      "loss": 0.9284,
+      "output_text_mask_loss": 0.2472,
+      "predictor_learning_rate": 0.00046188100853688,
+      "step": 1979
+    },
+    {
+      "epoch": 0.2,
+      "gumbel_tau": 0.627,
+      "image_masks_loss": 0.0091,
+      "llm_learning_rate": 4.618369643808758e-06,
+      "loss": 0.8684,
+      "output_text_mask_loss": 0.1038,
+      "predictor_learning_rate": 0.00046183696438087575,
+      "step": 1980
+    },
+    {
+      "epoch": 0.2,
+      "gumbel_tau": 0.6268,
+      "image_masks_loss": 0.0138,
+      "llm_learning_rate": 4.617928968969132e-06,
+      "loss": 0.8475,
+      "output_text_mask_loss": 0.2182,
+      "predictor_learning_rate": 0.00046179289689691317,
+      "step": 1981
+    },
+    {
+      "epoch": 0.2,
+      "gumbel_tau": 0.6267,
+      "image_masks_loss": 0.0167,
+      "llm_learning_rate": 4.617488060898451e-06,
+      "loss": 0.8047,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 0.00046174880608984504,
+      "step": 1982
+    },
+    {
+      "epoch": 0.2,
+      "gumbel_tau": 0.6265,
+      "image_masks_loss": 0.0353,
+      "llm_learning_rate": 4.617046919645267e-06,
+      "loss": 0.7812,
+      "output_text_mask_loss": 0.0292,
+      "predictor_learning_rate": 0.0004617046919645267,
+      "step": 1983
+    },
+    {
+      "epoch": 0.2,
+      "gumbel_tau": 0.6264,
+      "image_masks_loss": 0.0179,
+      "llm_learning_rate": 4.616605545258161e-06,
+      "loss": 0.7637,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.00046166055452581613,
+      "step": 1984
+    },
+    {
+      "epoch": 0.2,
+      "gumbel_tau": 0.6262,
+      "image_masks_loss": 0.0219,
+      "llm_learning_rate": 4.616163937785739e-06,
+      "loss": 0.7238,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004616163937785739,
+      "step": 1985
+    },
+    {
+      "epoch": 0.2,
+      "gumbel_tau": 0.6261,
+      "image_masks_loss": 0.0185,
+      "llm_learning_rate": 4.6157220972766305e-06,
+      "loss": 0.8106,
+      "output_text_mask_loss": 0.0056,
+      "predictor_learning_rate": 0.000461572209727663,
+      "step": 1986
+    },
+    {
+      "epoch": 0.2,
+      "gumbel_tau": 0.6259,
+      "image_masks_loss": 0.0273,
+      "llm_learning_rate": 4.615280023779491e-06,
+      "loss": 0.8264,
+      "output_text_mask_loss": 0.0401,
+      "predictor_learning_rate": 0.0004615280023779491,
+      "step": 1987
+    },
+    {
+      "epoch": 0.2,
+      "gumbel_tau": 0.6258,
+      "image_masks_loss": 0.0182,
+      "llm_learning_rate": 4.614837717343004e-06,
+      "loss": 0.7854,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004614837717343003,
+      "step": 1988
+    },
+    {
+      "epoch": 0.2,
+      "gumbel_tau": 0.6256,
+      "image_masks_loss": 0.0311,
+      "llm_learning_rate": 4.614395178015877e-06,
+      "loss": 0.823,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00046143951780158764,
+      "step": 1989
+    },
+    {
+      "epoch": 0.2,
+      "gumbel_tau": 0.6255,
+      "image_masks_loss": 0.0175,
+      "llm_learning_rate": 4.6139524058468435e-06,
+      "loss": 0.7671,
+      "output_text_mask_loss": 0.0526,
+      "predictor_learning_rate": 0.0004613952405846843,
+      "step": 1990
+    },
+    {
+      "epoch": 0.2,
+      "gumbel_tau": 0.6253,
+      "image_masks_loss": 0.0065,
+      "llm_learning_rate": 4.6135094008846626e-06,
+      "loss": 0.785,
+      "output_text_mask_loss": 0.1228,
+      "predictor_learning_rate": 0.00046135094008846624,
+      "step": 1991
+    },
+    {
+      "epoch": 0.2,
+      "gumbel_tau": 0.6252,
+      "image_masks_loss": 0.0132,
+      "llm_learning_rate": 4.6130661631781184e-06,
+      "loss": 0.7555,
+      "output_text_mask_loss": 0.0271,
+      "predictor_learning_rate": 0.0004613066163178118,
+      "step": 1992
+    },
+    {
+      "epoch": 0.2,
+      "gumbel_tau": 0.625,
+      "image_masks_loss": 0.0702,
+      "llm_learning_rate": 4.612622692776021e-06,
+      "loss": 0.7289,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.0004612622692776021,
+      "step": 1993
+    },
+    {
+      "epoch": 0.2,
+      "gumbel_tau": 0.6249,
+      "image_masks_loss": 0.0141,
+      "llm_learning_rate": 4.612178989727209e-06,
+      "loss": 0.8095,
+      "output_text_mask_loss": 0.0231,
+      "predictor_learning_rate": 0.0004612178989727208,
+      "step": 1994
+    },
+    {
+      "epoch": 0.2,
+      "gumbel_tau": 0.6247,
+      "image_masks_loss": 0.0265,
+      "llm_learning_rate": 4.61173505408054e-06,
+      "loss": 0.8424,
+      "output_text_mask_loss": 0.0202,
+      "predictor_learning_rate": 0.00046117350540805395,
+      "step": 1995
+    },
+    {
+      "epoch": 0.2,
+      "gumbel_tau": 0.6246,
+      "image_masks_loss": 0.0154,
+      "llm_learning_rate": 4.611290885884904e-06,
+      "loss": 0.743,
+      "output_text_mask_loss": 0.0278,
+      "predictor_learning_rate": 0.00046112908858849034,
+      "step": 1996
+    },
+    {
+      "epoch": 0.2,
+      "gumbel_tau": 0.6244,
+      "image_masks_loss": 0.0068,
+      "llm_learning_rate": 4.6108464851892135e-06,
+      "loss": 0.7474,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004610846485189213,
+      "step": 1997
+    },
+    {
+      "epoch": 0.2,
+      "gumbel_tau": 0.6243,
+      "image_masks_loss": 0.0196,
+      "llm_learning_rate": 4.6104018520424045e-06,
+      "loss": 0.8592,
+      "output_text_mask_loss": 0.0847,
+      "predictor_learning_rate": 0.00046104018520424047,
+      "step": 1998
+    },
+    {
+      "epoch": 0.2,
+      "gumbel_tau": 0.6241,
+      "image_masks_loss": 0.0111,
+      "llm_learning_rate": 4.6099569864934445e-06,
+      "loss": 0.7868,
+      "output_text_mask_loss": 0.0549,
+      "predictor_learning_rate": 0.00046099569864934445,
+      "step": 1999
+    },
+    {
+      "epoch": 0.2,
+      "gumbel_tau": 0.624,
+      "image_masks_loss": 0.0156,
+      "llm_learning_rate": 4.609511888591321e-06,
+      "loss": 0.7529,
+      "output_text_mask_loss": 0.1572,
+      "predictor_learning_rate": 0.0004609511888591321,
+      "step": 2000
+    },
+    {
+      "epoch": 0.21,
+      "gumbel_tau": 0.6239,
+      "image_masks_loss": 0.033,
+      "llm_learning_rate": 4.609066558385049e-06,
+      "loss": 0.8651,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 0.0004609066558385049,
+      "step": 2001
+    },
+    {
+      "epoch": 0.21,
+      "gumbel_tau": 0.6237,
+      "image_masks_loss": 0.0089,
+      "llm_learning_rate": 4.60862099592367e-06,
+      "loss": 0.8154,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.00046086209959236705,
+      "step": 2002
+    },
+    {
+      "epoch": 0.21,
+      "gumbel_tau": 0.6236,
+      "image_masks_loss": 0.0204,
+      "llm_learning_rate": 4.608175201256251e-06,
+      "loss": 0.8029,
+      "output_text_mask_loss": 0.1244,
+      "predictor_learning_rate": 0.00046081752012562505,
+      "step": 2003
+    },
+    {
+      "epoch": 0.21,
+      "gumbel_tau": 0.6234,
+      "image_masks_loss": 0.0178,
+      "llm_learning_rate": 4.607729174431882e-06,
+      "loss": 0.8088,
+      "output_text_mask_loss": 0.0202,
+      "predictor_learning_rate": 0.0004607729174431882,
+      "step": 2004
+    },
+    {
+      "epoch": 0.21,
+      "gumbel_tau": 0.6233,
+      "image_masks_loss": 0.0166,
+      "llm_learning_rate": 4.607282915499682e-06,
+      "loss": 0.794,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 0.0004607282915499682,
+      "step": 2005
+    },
+    {
+      "epoch": 0.21,
+      "gumbel_tau": 0.6231,
+      "image_masks_loss": 0.0275,
+      "llm_learning_rate": 4.606836424508794e-06,
+      "loss": 0.8047,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.00046068364245087937,
+      "step": 2006
+    },
+    {
+      "epoch": 0.21,
+      "gumbel_tau": 0.623,
+      "image_masks_loss": 0.0192,
+      "llm_learning_rate": 4.606389701508386e-06,
+      "loss": 0.8041,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.00046063897015083855,
+      "step": 2007
+    },
+    {
+      "epoch": 0.21,
+      "gumbel_tau": 0.6228,
+      "image_masks_loss": 0.0037,
+      "llm_learning_rate": 4.605942746547652e-06,
+      "loss": 0.8262,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.0004605942746547651,
+      "step": 2008
+    },
+    {
+      "epoch": 0.21,
+      "gumbel_tau": 0.6227,
+      "image_masks_loss": 0.0086,
+      "llm_learning_rate": 4.605495559675811e-06,
+      "loss": 0.8649,
+      "output_text_mask_loss": 0.0151,
+      "predictor_learning_rate": 0.0004605495559675811,
+      "step": 2009
+    },
+    {
+      "epoch": 0.21,
+      "gumbel_tau": 0.6225,
+      "image_masks_loss": 0.014,
+      "llm_learning_rate": 4.60504814094211e-06,
+      "loss": 0.7712,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.000460504814094211,
+      "step": 2010
+    },
+    {
+      "epoch": 0.21,
+      "gumbel_tau": 0.6224,
+      "image_masks_loss": 0.0216,
+      "llm_learning_rate": 4.60460049039582e-06,
+      "loss": 0.796,
+      "output_text_mask_loss": 0.0005,
+      "predictor_learning_rate": 0.00046046004903958196,
+      "step": 2011
+    },
+    {
+      "epoch": 0.21,
+      "gumbel_tau": 0.6222,
+      "image_masks_loss": 0.0235,
+      "llm_learning_rate": 4.604152608086235e-06,
+      "loss": 0.7327,
+      "output_text_mask_loss": 0.0039,
+      "predictor_learning_rate": 0.00046041526080862344,
+      "step": 2012
+    },
+    {
+      "epoch": 0.21,
+      "gumbel_tau": 0.6221,
+      "image_masks_loss": 0.0091,
+      "llm_learning_rate": 4.603704494062678e-06,
+      "loss": 0.7332,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.00046037044940626777,
+      "step": 2013
+    },
+    {
+      "epoch": 0.21,
+      "gumbel_tau": 0.6219,
+      "image_masks_loss": 0.0225,
+      "llm_learning_rate": 4.603256148374496e-06,
+      "loss": 0.7475,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.0004603256148374496,
+      "step": 2014
+    },
+    {
+      "epoch": 0.21,
+      "gumbel_tau": 0.6218,
+      "image_masks_loss": 0.0069,
+      "llm_learning_rate": 4.6028075710710625e-06,
+      "loss": 0.7597,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.00046028075710710626,
+      "step": 2015
+    },
+    {
+      "epoch": 0.21,
+      "gumbel_tau": 0.6216,
+      "image_masks_loss": 0.0172,
+      "llm_learning_rate": 4.602358762201776e-06,
+      "loss": 0.7957,
+      "output_text_mask_loss": 0.0039,
+      "predictor_learning_rate": 0.0004602358762201776,
+      "step": 2016
+    },
+    {
+      "epoch": 0.21,
+      "gumbel_tau": 0.6215,
+      "image_masks_loss": 0.0221,
+      "llm_learning_rate": 4.6019097218160594e-06,
+      "loss": 0.7887,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 0.00046019097218160586,
+      "step": 2017
+    },
+    {
+      "epoch": 0.21,
+      "gumbel_tau": 0.6214,
+      "image_masks_loss": 0.0296,
+      "llm_learning_rate": 4.601460449963362e-06,
+      "loss": 0.6686,
+      "output_text_mask_loss": 0.0095,
+      "predictor_learning_rate": 0.0004601460449963362,
+      "step": 2018
+    },
+    {
+      "epoch": 0.21,
+      "gumbel_tau": 0.6212,
+      "image_masks_loss": 0.0236,
+      "llm_learning_rate": 4.6010109466931594e-06,
+      "loss": 0.7805,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 0.00046010109466931593,
+      "step": 2019
+    },
+    {
+      "epoch": 0.21,
+      "gumbel_tau": 0.6211,
+      "image_masks_loss": 0.0364,
+      "llm_learning_rate": 4.600561212054951e-06,
+      "loss": 0.8403,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004600561212054951,
+      "step": 2020
+    },
+    {
+      "epoch": 0.21,
+      "gumbel_tau": 0.6209,
+      "image_masks_loss": 0.05,
+      "llm_learning_rate": 4.600111246098264e-06,
+      "loss": 0.7672,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.0004600111246098264,
+      "step": 2021
+    },
+    {
+      "epoch": 0.21,
+      "gumbel_tau": 0.6208,
+      "image_masks_loss": 0.003,
+      "llm_learning_rate": 4.599661048872648e-06,
+      "loss": 0.8675,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.0004599661048872648,
+      "step": 2022
+    },
+    {
+      "epoch": 0.21,
+      "gumbel_tau": 0.6206,
+      "image_masks_loss": 0.0183,
+      "llm_learning_rate": 4.599210620427681e-06,
+      "loss": 0.7824,
+      "output_text_mask_loss": 0.0071,
+      "predictor_learning_rate": 0.00045992106204276813,
+      "step": 2023
+    },
+    {
+      "epoch": 0.21,
+      "gumbel_tau": 0.6205,
+      "image_masks_loss": 0.0152,
+      "llm_learning_rate": 4.598759960812965e-06,
+      "loss": 0.8125,
+      "output_text_mask_loss": 0.0197,
+      "predictor_learning_rate": 0.0004598759960812965,
+      "step": 2024
+    },
+    {
+      "epoch": 0.21,
+      "gumbel_tau": 0.6203,
+      "image_masks_loss": 0.0106,
+      "llm_learning_rate": 4.598309070078127e-06,
+      "loss": 0.7318,
+      "output_text_mask_loss": 0.0025,
+      "predictor_learning_rate": 0.0004598309070078126,
+      "step": 2025
+    },
+    {
+      "epoch": 0.21,
+      "gumbel_tau": 0.6202,
+      "image_masks_loss": 0.0153,
+      "llm_learning_rate": 4.59785794827282e-06,
+      "loss": 0.7552,
+      "output_text_mask_loss": 0.0141,
+      "predictor_learning_rate": 0.000459785794827282,
+      "step": 2026
+    },
+    {
+      "epoch": 0.21,
+      "gumbel_tau": 0.62,
+      "image_masks_loss": 0.0183,
+      "llm_learning_rate": 4.597406595446723e-06,
+      "loss": 0.8186,
+      "output_text_mask_loss": 0.0162,
+      "predictor_learning_rate": 0.0004597406595446723,
+      "step": 2027
+    },
+    {
+      "epoch": 0.21,
+      "gumbel_tau": 0.6199,
+      "image_masks_loss": 0.0109,
+      "llm_learning_rate": 4.59695501164954e-06,
+      "loss": 0.7973,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00045969550116495403,
+      "step": 2028
+    },
+    {
+      "epoch": 0.21,
+      "gumbel_tau": 0.6197,
+      "image_masks_loss": 0.0208,
+      "llm_learning_rate": 4.596503196931e-06,
+      "loss": 0.8485,
+      "output_text_mask_loss": 0.0538,
+      "predictor_learning_rate": 0.0004596503196931,
+      "step": 2029
+    },
+    {
+      "epoch": 0.21,
+      "gumbel_tau": 0.6196,
+      "image_masks_loss": 0.0259,
+      "llm_learning_rate": 4.596051151340859e-06,
+      "loss": 0.7496,
+      "output_text_mask_loss": 0.1701,
+      "predictor_learning_rate": 0.0004596051151340859,
+      "step": 2030
+    },
+    {
+      "epoch": 0.21,
+      "gumbel_tau": 0.6195,
+      "image_masks_loss": 0.0218,
+      "llm_learning_rate": 4.595598874928896e-06,
+      "loss": 0.749,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00045955988749288957,
+      "step": 2031
+    },
+    {
+      "epoch": 0.21,
+      "gumbel_tau": 0.6193,
+      "image_masks_loss": 0.0098,
+      "llm_learning_rate": 4.595146367744917e-06,
+      "loss": 0.7508,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004595146367744917,
+      "step": 2032
+    },
+    {
+      "epoch": 0.21,
+      "gumbel_tau": 0.6192,
+      "image_masks_loss": 0.0081,
+      "llm_learning_rate": 4.594693629838754e-06,
+      "loss": 0.759,
+      "output_text_mask_loss": 0.002,
+      "predictor_learning_rate": 0.0004594693629838754,
+      "step": 2033
+    },
+    {
+      "epoch": 0.21,
+      "gumbel_tau": 0.619,
+      "image_masks_loss": 0.0122,
+      "llm_learning_rate": 4.594240661260262e-06,
+      "loss": 0.8007,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00045942406612602615,
+      "step": 2034
+    },
+    {
+      "epoch": 0.21,
+      "gumbel_tau": 0.6189,
+      "image_masks_loss": 0.0096,
+      "llm_learning_rate": 4.593787462059324e-06,
+      "loss": 0.7745,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.00045937874620593236,
+      "step": 2035
+    },
+    {
+      "epoch": 0.21,
+      "gumbel_tau": 0.6187,
+      "image_masks_loss": 0.0112,
+      "llm_learning_rate": 4.593334032285847e-06,
+      "loss": 0.8144,
+      "output_text_mask_loss": 0.0172,
+      "predictor_learning_rate": 0.0004593334032285846,
+      "step": 2036
+    },
+    {
+      "epoch": 0.21,
+      "gumbel_tau": 0.6186,
+      "image_masks_loss": 0.0154,
+      "llm_learning_rate": 4.592880371989763e-06,
+      "loss": 0.7645,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.0004592880371989763,
+      "step": 2037
+    },
+    {
+      "epoch": 0.21,
+      "gumbel_tau": 0.6184,
+      "image_masks_loss": 0.0142,
+      "llm_learning_rate": 4.592426481221031e-06,
+      "loss": 0.7709,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.00045924264812210303,
+      "step": 2038
+    },
+    {
+      "epoch": 0.21,
+      "gumbel_tau": 0.6183,
+      "image_masks_loss": 0.02,
+      "llm_learning_rate": 4.591972360029634e-06,
+      "loss": 0.7962,
+      "output_text_mask_loss": 0.0504,
+      "predictor_learning_rate": 0.0004591972360029634,
+      "step": 2039
+    },
+    {
+      "epoch": 0.21,
+      "gumbel_tau": 0.6181,
+      "image_masks_loss": 0.0101,
+      "llm_learning_rate": 4.591518008465582e-06,
+      "loss": 0.6216,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00045915180084655816,
+      "step": 2040
+    },
+    {
+      "epoch": 0.21,
+      "gumbel_tau": 0.618,
+      "image_masks_loss": 0.0138,
+      "llm_learning_rate": 4.591063426578907e-06,
+      "loss": 0.795,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.0004591063426578907,
+      "step": 2041
+    },
+    {
+      "epoch": 0.21,
+      "gumbel_tau": 0.6178,
+      "image_masks_loss": 0.0305,
+      "llm_learning_rate": 4.590608614419671e-06,
+      "loss": 0.6962,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.00045906086144196703,
+      "step": 2042
+    },
+    {
+      "epoch": 0.21,
+      "gumbel_tau": 0.6177,
+      "image_masks_loss": 0.0286,
+      "llm_learning_rate": 4.590153572037957e-06,
+      "loss": 0.7267,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00045901535720379563,
+      "step": 2043
+    },
+    {
+      "epoch": 0.21,
+      "gumbel_tau": 0.6176,
+      "image_masks_loss": 0.0051,
+      "llm_learning_rate": 4.589698299483875e-06,
+      "loss": 0.6184,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00045896982994838757,
+      "step": 2044
+    },
+    {
+      "epoch": 0.21,
+      "gumbel_tau": 0.6174,
+      "image_masks_loss": 0.0188,
+      "llm_learning_rate": 4.589242796807564e-06,
+      "loss": 0.7999,
+      "output_text_mask_loss": 0.0156,
+      "predictor_learning_rate": 0.00045892427968075635,
+      "step": 2045
+    },
+    {
+      "epoch": 0.21,
+      "gumbel_tau": 0.6173,
+      "image_masks_loss": 0.0199,
+      "llm_learning_rate": 4.588787064059181e-06,
+      "loss": 0.8566,
+      "output_text_mask_loss": 0.0549,
+      "predictor_learning_rate": 0.00045887870640591805,
+      "step": 2046
+    },
+    {
+      "epoch": 0.21,
+      "gumbel_tau": 0.6171,
+      "image_masks_loss": 0.0433,
+      "llm_learning_rate": 4.588331101288914e-06,
+      "loss": 0.7565,
+      "output_text_mask_loss": 0.0322,
+      "predictor_learning_rate": 0.0004588331101288914,
+      "step": 2047
+    },
+    {
+      "epoch": 0.21,
+      "gumbel_tau": 0.617,
+      "image_masks_loss": 0.0172,
+      "llm_learning_rate": 4.587874908546975e-06,
+      "loss": 0.8659,
+      "output_text_mask_loss": 0.0519,
+      "predictor_learning_rate": 0.00045878749085469743,
+      "step": 2048
+    },
+    {
+      "epoch": 0.21,
+      "gumbel_tau": 0.6168,
+      "image_masks_loss": 0.0292,
+      "llm_learning_rate": 4.5874184858836006e-06,
+      "loss": 0.7426,
+      "output_text_mask_loss": 0.0322,
+      "predictor_learning_rate": 0.00045874184858836,
+      "step": 2049
+    },
+    {
+      "epoch": 0.21,
+      "gumbel_tau": 0.6167,
+      "image_masks_loss": 0.0147,
+      "llm_learning_rate": 4.586961833349052e-06,
+      "loss": 0.7834,
+      "output_text_mask_loss": 0.0889,
+      "predictor_learning_rate": 0.00045869618333490523,
+      "step": 2050
+    },
+    {
+      "epoch": 0.21,
+      "gumbel_tau": 0.6165,
+      "image_masks_loss": 0.0169,
+      "llm_learning_rate": 4.586504950993619e-06,
+      "loss": 0.8424,
+      "output_text_mask_loss": 0.0755,
+      "predictor_learning_rate": 0.00045865049509936185,
+      "step": 2051
+    },
+    {
+      "epoch": 0.21,
+      "gumbel_tau": 0.6164,
+      "image_masks_loss": 0.0414,
+      "llm_learning_rate": 4.586047838867612e-06,
+      "loss": 0.7304,
+      "output_text_mask_loss": 0.0429,
+      "predictor_learning_rate": 0.00045860478388676126,
+      "step": 2052
+    },
+    {
+      "epoch": 0.21,
+      "gumbel_tau": 0.6162,
+      "image_masks_loss": 0.012,
+      "llm_learning_rate": 4.585590497021372e-06,
+      "loss": 0.8355,
+      "output_text_mask_loss": 0.005,
+      "predictor_learning_rate": 0.00045855904970213714,
+      "step": 2053
+    },
+    {
+      "epoch": 0.21,
+      "gumbel_tau": 0.6161,
+      "image_masks_loss": 0.0101,
+      "llm_learning_rate": 4.5851329255052605e-06,
+      "loss": 0.8677,
+      "output_text_mask_loss": 0.02,
+      "predictor_learning_rate": 0.000458513292550526,
+      "step": 2054
+    },
+    {
+      "epoch": 0.21,
+      "gumbel_tau": 0.616,
+      "image_masks_loss": 0.014,
+      "llm_learning_rate": 4.584675124369667e-06,
+      "loss": 0.7757,
+      "output_text_mask_loss": 0.0137,
+      "predictor_learning_rate": 0.0004584675124369666,
+      "step": 2055
+    },
+    {
+      "epoch": 0.21,
+      "gumbel_tau": 0.6158,
+      "image_masks_loss": 0.0425,
+      "llm_learning_rate": 4.584217093665004e-06,
+      "loss": 0.6602,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004584217093665004,
+      "step": 2056
+    },
+    {
+      "epoch": 0.21,
+      "gumbel_tau": 0.6157,
+      "image_masks_loss": 0.037,
+      "llm_learning_rate": 4.583758833441713e-06,
+      "loss": 0.8237,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00045837588334417137,
+      "step": 2057
+    },
+    {
+      "epoch": 0.21,
+      "gumbel_tau": 0.6155,
+      "image_masks_loss": 0.0114,
+      "llm_learning_rate": 4.5833003437502585e-06,
+      "loss": 0.8393,
+      "output_text_mask_loss": 0.0109,
+      "predictor_learning_rate": 0.00045833003437502584,
+      "step": 2058
+    },
+    {
+      "epoch": 0.21,
+      "gumbel_tau": 0.6154,
+      "image_masks_loss": 0.011,
+      "llm_learning_rate": 4.582841624641129e-06,
+      "loss": 0.6173,
+      "output_text_mask_loss": 0.0032,
+      "predictor_learning_rate": 0.00045828416246411285,
+      "step": 2059
+    },
+    {
+      "epoch": 0.21,
+      "gumbel_tau": 0.6152,
+      "image_masks_loss": 0.0271,
+      "llm_learning_rate": 4.582382676164842e-06,
+      "loss": 0.8038,
+      "output_text_mask_loss": 0.0404,
+      "predictor_learning_rate": 0.0004582382676164841,
+      "step": 2060
+    },
+    {
+      "epoch": 0.21,
+      "gumbel_tau": 0.6151,
+      "image_masks_loss": 0.0123,
+      "llm_learning_rate": 4.581923498371934e-06,
+      "loss": 0.8323,
+      "output_text_mask_loss": 0.0238,
+      "predictor_learning_rate": 0.0004581923498371934,
+      "step": 2061
+    },
+    {
+      "epoch": 0.21,
+      "gumbel_tau": 0.6149,
+      "image_masks_loss": 0.0271,
+      "llm_learning_rate": 4.581464091312974e-06,
+      "loss": 0.8432,
+      "output_text_mask_loss": 0.05,
+      "predictor_learning_rate": 0.0004581464091312974,
+      "step": 2062
+    },
+    {
+      "epoch": 0.21,
+      "gumbel_tau": 0.6148,
+      "image_masks_loss": 0.0083,
+      "llm_learning_rate": 4.581004455038552e-06,
+      "loss": 0.7486,
+      "output_text_mask_loss": 0.0322,
+      "predictor_learning_rate": 0.00045810044550385524,
+      "step": 2063
+    },
+    {
+      "epoch": 0.21,
+      "gumbel_tau": 0.6146,
+      "image_masks_loss": 0.0161,
+      "llm_learning_rate": 4.5805445895992854e-06,
+      "loss": 0.6484,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004580544589599285,
+      "step": 2064
+    },
+    {
+      "epoch": 0.21,
+      "gumbel_tau": 0.6145,
+      "image_masks_loss": 0.0163,
+      "llm_learning_rate": 4.580084495045813e-06,
+      "loss": 0.8305,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.0004580084495045813,
+      "step": 2065
+    },
+    {
+      "epoch": 0.21,
+      "gumbel_tau": 0.6144,
+      "image_masks_loss": 0.0195,
+      "llm_learning_rate": 4.5796241714288035e-06,
+      "loss": 0.709,
+      "output_text_mask_loss": 0.0004,
+      "predictor_learning_rate": 0.00045796241714288026,
+      "step": 2066
+    },
+    {
+      "epoch": 0.21,
+      "gumbel_tau": 0.6142,
+      "image_masks_loss": 0.0077,
+      "llm_learning_rate": 4.579163618798948e-06,
+      "loss": 0.8121,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.0004579163618798947,
+      "step": 2067
+    },
+    {
+      "epoch": 0.21,
+      "gumbel_tau": 0.6141,
+      "image_masks_loss": 0.0094,
+      "llm_learning_rate": 4.578702837206962e-06,
+      "loss": 0.7635,
+      "output_text_mask_loss": 0.0062,
+      "predictor_learning_rate": 0.0004578702837206962,
+      "step": 2068
+    },
+    {
+      "epoch": 0.21,
+      "gumbel_tau": 0.6139,
+      "image_masks_loss": 0.0185,
+      "llm_learning_rate": 4.578241826703591e-06,
+      "loss": 0.8095,
+      "output_text_mask_loss": 0.0378,
+      "predictor_learning_rate": 0.0004578241826703591,
+      "step": 2069
+    },
+    {
+      "epoch": 0.21,
+      "gumbel_tau": 0.6138,
+      "image_masks_loss": 0.0098,
+      "llm_learning_rate": 4.5777805873396005e-06,
+      "loss": 0.744,
+      "output_text_mask_loss": 0.0507,
+      "predictor_learning_rate": 0.00045777805873396007,
+      "step": 2070
+    },
+    {
+      "epoch": 0.21,
+      "gumbel_tau": 0.6136,
+      "image_masks_loss": 0.0093,
+      "llm_learning_rate": 4.577319119165784e-06,
+      "loss": 0.882,
+      "output_text_mask_loss": 0.0927,
+      "predictor_learning_rate": 0.0004577319119165784,
+      "step": 2071
+    },
+    {
+      "epoch": 0.21,
+      "gumbel_tau": 0.6135,
+      "image_masks_loss": 0.0196,
+      "llm_learning_rate": 4.5768574222329596e-06,
+      "loss": 0.7455,
+      "output_text_mask_loss": 0.0343,
+      "predictor_learning_rate": 0.0004576857422232959,
+      "step": 2072
+    },
+    {
+      "epoch": 0.21,
+      "gumbel_tau": 0.6133,
+      "image_masks_loss": 0.0134,
+      "llm_learning_rate": 4.576395496591969e-06,
+      "loss": 0.769,
+      "output_text_mask_loss": 0.0231,
+      "predictor_learning_rate": 0.0004576395496591969,
+      "step": 2073
+    },
+    {
+      "epoch": 0.21,
+      "gumbel_tau": 0.6132,
+      "image_masks_loss": 0.0202,
+      "llm_learning_rate": 4.575933342293681e-06,
+      "loss": 0.7276,
+      "output_text_mask_loss": 0.0124,
+      "predictor_learning_rate": 0.0004575933342293681,
+      "step": 2074
+    },
+    {
+      "epoch": 0.21,
+      "gumbel_tau": 0.6131,
+      "image_masks_loss": 0.0095,
+      "llm_learning_rate": 4.5754709593889904e-06,
+      "loss": 0.8636,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 0.00045754709593889896,
+      "step": 2075
+    },
+    {
+      "epoch": 0.21,
+      "gumbel_tau": 0.6129,
+      "image_masks_loss": 0.0063,
+      "llm_learning_rate": 4.575008347928814e-06,
+      "loss": 0.7418,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.0004575008347928814,
+      "step": 2076
+    },
+    {
+      "epoch": 0.21,
+      "gumbel_tau": 0.6128,
+      "image_masks_loss": 0.0084,
+      "llm_learning_rate": 4.574545507964096e-06,
+      "loss": 0.7395,
+      "output_text_mask_loss": 0.0284,
+      "predictor_learning_rate": 0.0004574545507964096,
+      "step": 2077
+    },
+    {
+      "epoch": 0.21,
+      "gumbel_tau": 0.6126,
+      "image_masks_loss": 0.0201,
+      "llm_learning_rate": 4.574082439545807e-06,
+      "loss": 0.7226,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.0004574082439545807,
+      "step": 2078
+    },
+    {
+      "epoch": 0.21,
+      "gumbel_tau": 0.6125,
+      "image_masks_loss": 0.0151,
+      "llm_learning_rate": 4.57361914272494e-06,
+      "loss": 0.8187,
+      "output_text_mask_loss": 0.0177,
+      "predictor_learning_rate": 0.00045736191427249396,
+      "step": 2079
+    },
+    {
+      "epoch": 0.21,
+      "gumbel_tau": 0.6123,
+      "image_masks_loss": 0.0278,
+      "llm_learning_rate": 4.573155617552514e-06,
+      "loss": 0.6945,
+      "output_text_mask_loss": 0.0039,
+      "predictor_learning_rate": 0.0004573155617552513,
+      "step": 2080
+    },
+    {
+      "epoch": 0.21,
+      "gumbel_tau": 0.6122,
+      "image_masks_loss": 0.0142,
+      "llm_learning_rate": 4.572691864079572e-06,
+      "loss": 0.7996,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004572691864079572,
+      "step": 2081
+    },
+    {
+      "epoch": 0.21,
+      "gumbel_tau": 0.612,
+      "image_masks_loss": 0.0309,
+      "llm_learning_rate": 4.572227882357186e-06,
+      "loss": 0.7665,
+      "output_text_mask_loss": 0.0175,
+      "predictor_learning_rate": 0.00045722278823571863,
+      "step": 2082
+    },
+    {
+      "epoch": 0.21,
+      "gumbel_tau": 0.6119,
+      "image_masks_loss": 0.0216,
+      "llm_learning_rate": 4.571763672436451e-06,
+      "loss": 0.7006,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004571763672436451,
+      "step": 2083
+    },
+    {
+      "epoch": 0.21,
+      "gumbel_tau": 0.6118,
+      "image_masks_loss": 0.0234,
+      "llm_learning_rate": 4.571299234368485e-06,
+      "loss": 0.7877,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.00045712992343684844,
+      "step": 2084
+    },
+    {
+      "epoch": 0.21,
+      "gumbel_tau": 0.6116,
+      "image_masks_loss": 0.0106,
+      "llm_learning_rate": 4.570834568204434e-06,
+      "loss": 0.7933,
+      "output_text_mask_loss": 0.0782,
+      "predictor_learning_rate": 0.0004570834568204434,
+      "step": 2085
+    },
+    {
+      "epoch": 0.21,
+      "gumbel_tau": 0.6115,
+      "image_masks_loss": 0.0254,
+      "llm_learning_rate": 4.570369673995467e-06,
+      "loss": 0.7627,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 0.0004570369673995467,
+      "step": 2086
+    },
+    {
+      "epoch": 0.21,
+      "gumbel_tau": 0.6113,
+      "image_masks_loss": 0.0127,
+      "llm_learning_rate": 4.569904551792781e-06,
+      "loss": 0.7666,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00045699045517927806,
+      "step": 2087
+    },
+    {
+      "epoch": 0.21,
+      "gumbel_tau": 0.6112,
+      "image_masks_loss": 0.0271,
+      "llm_learning_rate": 4.569439201647594e-06,
+      "loss": 0.7218,
+      "output_text_mask_loss": 0.0343,
+      "predictor_learning_rate": 0.0004569439201647594,
+      "step": 2088
+    },
+    {
+      "epoch": 0.21,
+      "gumbel_tau": 0.611,
+      "image_masks_loss": 0.006,
+      "llm_learning_rate": 4.568973623611154e-06,
+      "loss": 0.8525,
+      "output_text_mask_loss": 0.0221,
+      "predictor_learning_rate": 0.00045689736236111534,
+      "step": 2089
+    },
+    {
+      "epoch": 0.21,
+      "gumbel_tau": 0.6109,
+      "image_masks_loss": 0.0104,
+      "llm_learning_rate": 4.5685078177347285e-06,
+      "loss": 0.7225,
+      "output_text_mask_loss": 0.009,
+      "predictor_learning_rate": 0.00045685078177347285,
+      "step": 2090
+    },
+    {
+      "epoch": 0.21,
+      "gumbel_tau": 0.6107,
+      "image_masks_loss": 0.0069,
+      "llm_learning_rate": 4.568041784069616e-06,
+      "loss": 0.6838,
+      "output_text_mask_loss": 0.0006,
+      "predictor_learning_rate": 0.0004568041784069615,
+      "step": 2091
+    },
+    {
+      "epoch": 0.21,
+      "gumbel_tau": 0.6106,
+      "image_masks_loss": 0.0273,
+      "llm_learning_rate": 4.567575522667135e-06,
+      "loss": 0.783,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00045675755226671344,
+      "step": 2092
+    },
+    {
+      "epoch": 0.21,
+      "gumbel_tau": 0.6105,
+      "image_masks_loss": 0.0296,
+      "llm_learning_rate": 4.567109033578633e-06,
+      "loss": 0.8189,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 0.00045671090335786327,
+      "step": 2093
+    },
+    {
+      "epoch": 0.21,
+      "gumbel_tau": 0.6103,
+      "image_masks_loss": 0.0248,
+      "llm_learning_rate": 4.566642316855479e-06,
+      "loss": 0.7376,
+      "output_text_mask_loss": 0.0004,
+      "predictor_learning_rate": 0.0004566642316855479,
+      "step": 2094
+    },
+    {
+      "epoch": 0.21,
+      "gumbel_tau": 0.6102,
+      "image_masks_loss": 0.0263,
+      "llm_learning_rate": 4.56617537254907e-06,
+      "loss": 0.7945,
+      "output_text_mask_loss": 0.0839,
+      "predictor_learning_rate": 0.00045661753725490694,
+      "step": 2095
+    },
+    {
+      "epoch": 0.21,
+      "gumbel_tau": 0.61,
+      "image_masks_loss": 0.0157,
+      "llm_learning_rate": 4.565708200710827e-06,
+      "loss": 0.6961,
+      "output_text_mask_loss": 0.0055,
+      "predictor_learning_rate": 0.0004565708200710827,
+      "step": 2096
+    },
+    {
+      "epoch": 0.21,
+      "gumbel_tau": 0.6099,
+      "image_masks_loss": 0.0523,
+      "llm_learning_rate": 4.565240801392196e-06,
+      "loss": 0.8684,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.0004565240801392195,
+      "step": 2097
+    },
+    {
+      "epoch": 0.21,
+      "gumbel_tau": 0.6097,
+      "image_masks_loss": 0.019,
+      "llm_learning_rate": 4.564773174644647e-06,
+      "loss": 0.8021,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00045647731746446476,
+      "step": 2098
+    },
+    {
+      "epoch": 0.22,
+      "gumbel_tau": 0.6096,
+      "image_masks_loss": 0.0051,
+      "llm_learning_rate": 4.564305320519677e-06,
+      "loss": 0.8586,
+      "output_text_mask_loss": 0.0488,
+      "predictor_learning_rate": 0.0004564305320519677,
+      "step": 2099
+    },
+    {
+      "epoch": 0.22,
+      "gumbel_tau": 0.6095,
+      "image_masks_loss": 0.0172,
+      "llm_learning_rate": 4.563837239068808e-06,
+      "loss": 0.809,
+      "output_text_mask_loss": 0.0277,
+      "predictor_learning_rate": 0.00045638372390688077,
+      "step": 2100
+    },
+    {
+      "epoch": 0.22,
+      "gumbel_tau": 0.6093,
+      "image_masks_loss": 0.032,
+      "llm_learning_rate": 4.5633689303435845e-06,
+      "loss": 0.6338,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00045633689303435843,
+      "step": 2101
+    },
+    {
+      "epoch": 0.22,
+      "gumbel_tau": 0.6092,
+      "image_masks_loss": 0.0083,
+      "llm_learning_rate": 4.562900394395579e-06,
+      "loss": 0.7228,
+      "output_text_mask_loss": 0.0847,
+      "predictor_learning_rate": 0.0004562900394395579,
+      "step": 2102
+    },
+    {
+      "epoch": 0.22,
+      "gumbel_tau": 0.609,
+      "image_masks_loss": 0.0084,
+      "llm_learning_rate": 4.562431631276387e-06,
+      "loss": 0.8062,
+      "output_text_mask_loss": 0.0164,
+      "predictor_learning_rate": 0.0004562431631276386,
+      "step": 2103
+    },
+    {
+      "epoch": 0.22,
+      "gumbel_tau": 0.6089,
+      "image_masks_loss": 0.0157,
+      "llm_learning_rate": 4.5619626410376295e-06,
+      "loss": 0.8142,
+      "output_text_mask_loss": 0.0019,
+      "predictor_learning_rate": 0.0004561962641037629,
+      "step": 2104
+    },
+    {
+      "epoch": 0.22,
+      "gumbel_tau": 0.6087,
+      "image_masks_loss": 0.0069,
+      "llm_learning_rate": 4.561493423730953e-06,
+      "loss": 0.7948,
+      "output_text_mask_loss": 0.033,
+      "predictor_learning_rate": 0.00045614934237309525,
+      "step": 2105
+    },
+    {
+      "epoch": 0.22,
+      "gumbel_tau": 0.6086,
+      "image_masks_loss": 0.0347,
+      "llm_learning_rate": 4.561023979408029e-06,
+      "loss": 0.7723,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 0.0004561023979408029,
+      "step": 2106
+    },
+    {
+      "epoch": 0.22,
+      "gumbel_tau": 0.6084,
+      "image_masks_loss": 0.0189,
+      "llm_learning_rate": 4.560554308120554e-06,
+      "loss": 0.8695,
+      "output_text_mask_loss": 0.016,
+      "predictor_learning_rate": 0.0004560554308120554,
+      "step": 2107
+    },
+    {
+      "epoch": 0.22,
+      "gumbel_tau": 0.6083,
+      "image_masks_loss": 0.0163,
+      "llm_learning_rate": 4.560084409920248e-06,
+      "loss": 0.8155,
+      "output_text_mask_loss": 0.0108,
+      "predictor_learning_rate": 0.00045600844099202484,
+      "step": 2108
+    },
+    {
+      "epoch": 0.22,
+      "gumbel_tau": 0.6082,
+      "image_masks_loss": 0.0275,
+      "llm_learning_rate": 4.559614284858859e-06,
+      "loss": 0.7628,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.0004559614284858859,
+      "step": 2109
+    },
+    {
+      "epoch": 0.22,
+      "gumbel_tau": 0.608,
+      "image_masks_loss": 0.0182,
+      "llm_learning_rate": 4.559143932988157e-06,
+      "loss": 0.786,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.0004559143932988157,
+      "step": 2110
+    },
+    {
+      "epoch": 0.22,
+      "gumbel_tau": 0.6079,
+      "image_masks_loss": 0.0299,
+      "llm_learning_rate": 4.55867335435994e-06,
+      "loss": 0.8731,
+      "output_text_mask_loss": 0.0116,
+      "predictor_learning_rate": 0.0004558673354359939,
+      "step": 2111
+    },
+    {
+      "epoch": 0.22,
+      "gumbel_tau": 0.6077,
+      "image_masks_loss": 0.0098,
+      "llm_learning_rate": 4.558202549026025e-06,
+      "loss": 0.8054,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004558202549026025,
+      "step": 2112
+    },
+    {
+      "epoch": 0.22,
+      "gumbel_tau": 0.6076,
+      "image_masks_loss": 0.0071,
+      "llm_learning_rate": 4.5577315170382625e-06,
+      "loss": 0.7344,
+      "output_text_mask_loss": 0.0298,
+      "predictor_learning_rate": 0.00045577315170382623,
+      "step": 2113
+    },
+    {
+      "epoch": 0.22,
+      "gumbel_tau": 0.6074,
+      "image_masks_loss": 0.0176,
+      "llm_learning_rate": 4.557260258448522e-06,
+      "loss": 0.8388,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 0.00045572602584485213,
+      "step": 2114
+    },
+    {
+      "epoch": 0.22,
+      "gumbel_tau": 0.6073,
+      "image_masks_loss": 0.036,
+      "llm_learning_rate": 4.556788773308699e-06,
+      "loss": 0.8019,
+      "output_text_mask_loss": 0.0019,
+      "predictor_learning_rate": 0.00045567887733086984,
+      "step": 2115
+    },
+    {
+      "epoch": 0.22,
+      "gumbel_tau": 0.6072,
+      "image_masks_loss": 0.0144,
+      "llm_learning_rate": 4.556317061670714e-06,
+      "loss": 0.8226,
+      "output_text_mask_loss": 0.0012,
+      "predictor_learning_rate": 0.00045563170616707145,
+      "step": 2116
+    },
+    {
+      "epoch": 0.22,
+      "gumbel_tau": 0.607,
+      "image_masks_loss": 0.0242,
+      "llm_learning_rate": 4.555845123586515e-06,
+      "loss": 0.8165,
+      "output_text_mask_loss": 0.0023,
+      "predictor_learning_rate": 0.0004555845123586515,
+      "step": 2117
+    },
+    {
+      "epoch": 0.22,
+      "gumbel_tau": 0.6069,
+      "image_masks_loss": 0.0278,
+      "llm_learning_rate": 4.555372959108072e-06,
+      "loss": 0.7365,
+      "output_text_mask_loss": 0.0313,
+      "predictor_learning_rate": 0.0004555372959108072,
+      "step": 2118
+    },
+    {
+      "epoch": 0.22,
+      "gumbel_tau": 0.6067,
+      "image_masks_loss": 0.0078,
+      "llm_learning_rate": 4.55490056828738e-06,
+      "loss": 0.8054,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.000455490056828738,
+      "step": 2119
+    },
+    {
+      "epoch": 0.22,
+      "gumbel_tau": 0.6066,
+      "image_masks_loss": 0.0162,
+      "llm_learning_rate": 4.554427951176461e-06,
+      "loss": 0.7786,
+      "output_text_mask_loss": 0.0422,
+      "predictor_learning_rate": 0.00045544279511764606,
+      "step": 2120
+    },
+    {
+      "epoch": 0.22,
+      "gumbel_tau": 0.6064,
+      "image_masks_loss": 0.0098,
+      "llm_learning_rate": 4.55395510782736e-06,
+      "loss": 0.8288,
+      "output_text_mask_loss": 0.0021,
+      "predictor_learning_rate": 0.0004553955107827359,
+      "step": 2121
+    },
+    {
+      "epoch": 0.22,
+      "gumbel_tau": 0.6063,
+      "image_masks_loss": 0.0519,
+      "llm_learning_rate": 4.553482038292147e-06,
+      "loss": 0.7905,
+      "output_text_mask_loss": 0.0111,
+      "predictor_learning_rate": 0.00045534820382921463,
+      "step": 2122
+    },
+    {
+      "epoch": 0.22,
+      "gumbel_tau": 0.6062,
+      "image_masks_loss": 0.0149,
+      "llm_learning_rate": 4.553008742622918e-06,
+      "loss": 0.6404,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004553008742622918,
+      "step": 2123
+    },
+    {
+      "epoch": 0.22,
+      "gumbel_tau": 0.606,
+      "image_masks_loss": 0.0172,
+      "llm_learning_rate": 4.552535220871793e-06,
+      "loss": 0.7368,
+      "output_text_mask_loss": 0.006,
+      "predictor_learning_rate": 0.0004552535220871793,
+      "step": 2124
+    },
+    {
+      "epoch": 0.22,
+      "gumbel_tau": 0.6059,
+      "image_masks_loss": 0.0121,
+      "llm_learning_rate": 4.552061473090918e-06,
+      "loss": 0.7536,
+      "output_text_mask_loss": 0.0153,
+      "predictor_learning_rate": 0.00045520614730909184,
+      "step": 2125
+    },
+    {
+      "epoch": 0.22,
+      "gumbel_tau": 0.6057,
+      "image_masks_loss": 0.0222,
+      "llm_learning_rate": 4.551587499332463e-06,
+      "loss": 0.7855,
+      "output_text_mask_loss": 0.0139,
+      "predictor_learning_rate": 0.0004551587499332463,
+      "step": 2126
+    },
+    {
+      "epoch": 0.22,
+      "gumbel_tau": 0.6056,
+      "image_masks_loss": 0.0406,
+      "llm_learning_rate": 4.551113299648623e-06,
+      "loss": 0.7864,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00045511132996486225,
+      "step": 2127
+    },
+    {
+      "epoch": 0.22,
+      "gumbel_tau": 0.6054,
+      "image_masks_loss": 0.0183,
+      "llm_learning_rate": 4.550638874091617e-06,
+      "loss": 0.7627,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.0004550638874091617,
+      "step": 2128
+    },
+    {
+      "epoch": 0.22,
+      "gumbel_tau": 0.6053,
+      "image_masks_loss": 0.024,
+      "llm_learning_rate": 4.5501642227136914e-06,
+      "loss": 0.7731,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00045501642227136906,
+      "step": 2129
+    },
+    {
+      "epoch": 0.22,
+      "gumbel_tau": 0.6052,
+      "image_masks_loss": 0.02,
+      "llm_learning_rate": 4.5496893455671144e-06,
+      "loss": 0.7325,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004549689345567114,
+      "step": 2130
+    },
+    {
+      "epoch": 0.22,
+      "gumbel_tau": 0.605,
+      "image_masks_loss": 0.0128,
+      "llm_learning_rate": 4.549214242704182e-06,
+      "loss": 0.7924,
+      "output_text_mask_loss": 0.0412,
+      "predictor_learning_rate": 0.00045492142427041817,
+      "step": 2131
+    },
+    {
+      "epoch": 0.22,
+      "gumbel_tau": 0.6049,
+      "image_masks_loss": 0.0267,
+      "llm_learning_rate": 4.548738914177212e-06,
+      "loss": 0.7384,
+      "output_text_mask_loss": 0.021,
+      "predictor_learning_rate": 0.00045487389141772117,
+      "step": 2132
+    },
+    {
+      "epoch": 0.22,
+      "gumbel_tau": 0.6047,
+      "image_masks_loss": 0.0364,
+      "llm_learning_rate": 4.5482633600385495e-06,
+      "loss": 0.7766,
+      "output_text_mask_loss": 0.0303,
+      "predictor_learning_rate": 0.0004548263360038549,
+      "step": 2133
+    },
+    {
+      "epoch": 0.22,
+      "gumbel_tau": 0.6046,
+      "image_masks_loss": 0.0097,
+      "llm_learning_rate": 4.547787580340563e-06,
+      "loss": 0.7307,
+      "output_text_mask_loss": 0.0534,
+      "predictor_learning_rate": 0.00045477875803405634,
+      "step": 2134
+    },
+    {
+      "epoch": 0.22,
+      "gumbel_tau": 0.6044,
+      "image_masks_loss": 0.0241,
+      "llm_learning_rate": 4.547311575135648e-06,
+      "loss": 0.7841,
+      "output_text_mask_loss": 0.0111,
+      "predictor_learning_rate": 0.00045473115751356474,
+      "step": 2135
+    },
+    {
+      "epoch": 0.22,
+      "gumbel_tau": 0.6043,
+      "image_masks_loss": 0.0079,
+      "llm_learning_rate": 4.546835344476221e-06,
+      "loss": 0.7406,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.0004546835344476221,
+      "step": 2136
+    },
+    {
+      "epoch": 0.22,
+      "gumbel_tau": 0.6042,
+      "image_masks_loss": 0.0351,
+      "llm_learning_rate": 4.546358888414728e-06,
+      "loss": 0.7401,
+      "output_text_mask_loss": 0.0132,
+      "predictor_learning_rate": 0.00045463588884147275,
+      "step": 2137
+    },
+    {
+      "epoch": 0.22,
+      "gumbel_tau": 0.604,
+      "image_masks_loss": 0.0239,
+      "llm_learning_rate": 4.545882207003635e-06,
+      "loss": 0.7726,
+      "output_text_mask_loss": 0.0565,
+      "predictor_learning_rate": 0.0004545882207003635,
+      "step": 2138
+    },
+    {
+      "epoch": 0.22,
+      "gumbel_tau": 0.6039,
+      "image_masks_loss": 0.0138,
+      "llm_learning_rate": 4.5454053002954365e-06,
+      "loss": 0.6998,
+      "output_text_mask_loss": 0.0019,
+      "predictor_learning_rate": 0.00045454053002954366,
+      "step": 2139
+    },
+    {
+      "epoch": 0.22,
+      "gumbel_tau": 0.6037,
+      "image_masks_loss": 0.0095,
+      "llm_learning_rate": 4.544928168342651e-06,
+      "loss": 0.7648,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004544928168342651,
+      "step": 2140
+    },
+    {
+      "epoch": 0.22,
+      "gumbel_tau": 0.6036,
+      "image_masks_loss": 0.0542,
+      "llm_learning_rate": 4.54445081119782e-06,
+      "loss": 0.736,
+      "output_text_mask_loss": 0.0086,
+      "predictor_learning_rate": 0.000454445081119782,
+      "step": 2141
+    },
+    {
+      "epoch": 0.22,
+      "gumbel_tau": 0.6034,
+      "image_masks_loss": 0.0196,
+      "llm_learning_rate": 4.543973228913512e-06,
+      "loss": 0.7428,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.00045439732289135124,
+      "step": 2142
+    },
+    {
+      "epoch": 0.22,
+      "gumbel_tau": 0.6033,
+      "image_masks_loss": 0.0144,
+      "llm_learning_rate": 4.54349542154232e-06,
+      "loss": 0.6845,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.000454349542154232,
+      "step": 2143
+    },
+    {
+      "epoch": 0.22,
+      "gumbel_tau": 0.6032,
+      "image_masks_loss": 0.0094,
+      "llm_learning_rate": 4.54301738913686e-06,
+      "loss": 0.72,
+      "output_text_mask_loss": 0.0021,
+      "predictor_learning_rate": 0.000454301738913686,
+      "step": 2144
+    },
+    {
+      "epoch": 0.22,
+      "gumbel_tau": 0.603,
+      "image_masks_loss": 0.0309,
+      "llm_learning_rate": 4.542539131749774e-06,
+      "loss": 0.8133,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004542539131749773,
+      "step": 2145
+    },
+    {
+      "epoch": 0.22,
+      "gumbel_tau": 0.6029,
+      "image_masks_loss": 0.0129,
+      "llm_learning_rate": 4.542060649433729e-06,
+      "loss": 0.8266,
+      "output_text_mask_loss": 0.0187,
+      "predictor_learning_rate": 0.0004542060649433729,
+      "step": 2146
+    },
+    {
+      "epoch": 0.22,
+      "gumbel_tau": 0.6027,
+      "image_masks_loss": 0.0232,
+      "llm_learning_rate": 4.541581942241418e-06,
+      "loss": 0.8977,
+      "output_text_mask_loss": 0.0301,
+      "predictor_learning_rate": 0.00045415819422414173,
+      "step": 2147
+    },
+    {
+      "epoch": 0.22,
+      "gumbel_tau": 0.6026,
+      "image_masks_loss": 0.0256,
+      "llm_learning_rate": 4.541103010225555e-06,
+      "loss": 0.7606,
+      "output_text_mask_loss": 0.0086,
+      "predictor_learning_rate": 0.00045411030102255543,
+      "step": 2148
+    },
+    {
+      "epoch": 0.22,
+      "gumbel_tau": 0.6024,
+      "image_masks_loss": 0.0376,
+      "llm_learning_rate": 4.540623853438883e-06,
+      "loss": 0.7844,
+      "output_text_mask_loss": 0.0526,
+      "predictor_learning_rate": 0.0004540623853438883,
+      "step": 2149
+    },
+    {
+      "epoch": 0.22,
+      "gumbel_tau": 0.6023,
+      "image_masks_loss": 0.0107,
+      "llm_learning_rate": 4.540144471934166e-06,
+      "loss": 0.821,
+      "output_text_mask_loss": 0.0116,
+      "predictor_learning_rate": 0.0004540144471934166,
+      "step": 2150
+    },
+    {
+      "epoch": 0.22,
+      "gumbel_tau": 0.6022,
+      "image_masks_loss": 0.0252,
+      "llm_learning_rate": 4.539664865764196e-06,
+      "loss": 0.7819,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004539664865764196,
+      "step": 2151
+    },
+    {
+      "epoch": 0.22,
+      "gumbel_tau": 0.602,
+      "image_masks_loss": 0.0256,
+      "llm_learning_rate": 4.539185034981788e-06,
+      "loss": 0.8314,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.00045391850349817874,
+      "step": 2152
+    },
+    {
+      "epoch": 0.22,
+      "gumbel_tau": 0.6019,
+      "image_masks_loss": 0.0151,
+      "llm_learning_rate": 4.5387049796397816e-06,
+      "loss": 0.8469,
+      "output_text_mask_loss": 0.0139,
+      "predictor_learning_rate": 0.00045387049796397815,
+      "step": 2153
+    },
+    {
+      "epoch": 0.22,
+      "gumbel_tau": 0.6017,
+      "image_masks_loss": 0.0458,
+      "llm_learning_rate": 4.5382246997910415e-06,
+      "loss": 0.8337,
+      "output_text_mask_loss": 0.0141,
+      "predictor_learning_rate": 0.00045382246997910413,
+      "step": 2154
+    },
+    {
+      "epoch": 0.22,
+      "gumbel_tau": 0.6016,
+      "image_masks_loss": 0.0257,
+      "llm_learning_rate": 4.537744195488459e-06,
+      "loss": 0.8085,
+      "output_text_mask_loss": 0.0071,
+      "predictor_learning_rate": 0.0004537744195488458,
+      "step": 2155
+    },
+    {
+      "epoch": 0.22,
+      "gumbel_tau": 0.6015,
+      "image_masks_loss": 0.0131,
+      "llm_learning_rate": 4.537263466784945e-06,
+      "loss": 0.745,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004537263466784945,
+      "step": 2156
+    },
+    {
+      "epoch": 0.22,
+      "gumbel_tau": 0.6013,
+      "image_masks_loss": 0.0166,
+      "llm_learning_rate": 4.536782513733441e-06,
+      "loss": 0.7836,
+      "output_text_mask_loss": 0.0162,
+      "predictor_learning_rate": 0.00045367825137334407,
+      "step": 2157
+    },
+    {
+      "epoch": 0.22,
+      "gumbel_tau": 0.6012,
+      "image_masks_loss": 0.0097,
+      "llm_learning_rate": 4.536301336386909e-06,
+      "loss": 0.7801,
+      "output_text_mask_loss": 0.0381,
+      "predictor_learning_rate": 0.00045363013363869094,
+      "step": 2158
+    },
+    {
+      "epoch": 0.22,
+      "gumbel_tau": 0.601,
+      "image_masks_loss": 0.0114,
+      "llm_learning_rate": 4.535819934798339e-06,
+      "loss": 0.7388,
+      "output_text_mask_loss": 0.0134,
+      "predictor_learning_rate": 0.00045358199347983393,
+      "step": 2159
+    },
+    {
+      "epoch": 0.22,
+      "gumbel_tau": 0.6009,
+      "image_masks_loss": 0.025,
+      "llm_learning_rate": 4.535338309020743e-06,
+      "loss": 0.8531,
+      "output_text_mask_loss": 0.0523,
+      "predictor_learning_rate": 0.00045353383090207435,
+      "step": 2160
+    },
+    {
+      "epoch": 0.22,
+      "gumbel_tau": 0.6007,
+      "image_masks_loss": 0.0094,
+      "llm_learning_rate": 4.53485645910716e-06,
+      "loss": 0.7522,
+      "output_text_mask_loss": 0.0151,
+      "predictor_learning_rate": 0.00045348564591071597,
+      "step": 2161
+    },
+    {
+      "epoch": 0.22,
+      "gumbel_tau": 0.6006,
+      "image_masks_loss": 0.0177,
+      "llm_learning_rate": 4.5343743851106495e-06,
+      "loss": 0.7653,
+      "output_text_mask_loss": 0.0406,
+      "predictor_learning_rate": 0.0004534374385110649,
+      "step": 2162
+    },
+    {
+      "epoch": 0.22,
+      "gumbel_tau": 0.6005,
+      "image_masks_loss": 0.025,
+      "llm_learning_rate": 4.533892087084301e-06,
+      "loss": 0.8447,
+      "output_text_mask_loss": 0.009,
+      "predictor_learning_rate": 0.00045338920870843005,
+      "step": 2163
+    },
+    {
+      "epoch": 0.22,
+      "gumbel_tau": 0.6003,
+      "image_masks_loss": 0.0095,
+      "llm_learning_rate": 4.533409565081226e-06,
+      "loss": 0.7543,
+      "output_text_mask_loss": 0.0319,
+      "predictor_learning_rate": 0.00045334095650812254,
+      "step": 2164
+    },
+    {
+      "epoch": 0.22,
+      "gumbel_tau": 0.6002,
+      "image_masks_loss": 0.0119,
+      "llm_learning_rate": 4.532926819154559e-06,
+      "loss": 0.7244,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004532926819154559,
+      "step": 2165
+    },
+    {
+      "epoch": 0.22,
+      "gumbel_tau": 0.6,
+      "image_masks_loss": 0.0231,
+      "llm_learning_rate": 4.532443849357464e-06,
+      "loss": 0.7281,
+      "output_text_mask_loss": 0.0076,
+      "predictor_learning_rate": 0.00045324438493574636,
+      "step": 2166
+    },
+    {
+      "epoch": 0.22,
+      "gumbel_tau": 0.5999,
+      "image_masks_loss": 0.013,
+      "llm_learning_rate": 4.531960655743125e-06,
+      "loss": 0.8622,
+      "output_text_mask_loss": 0.0671,
+      "predictor_learning_rate": 0.0004531960655743125,
+      "step": 2167
+    },
+    {
+      "epoch": 0.22,
+      "gumbel_tau": 0.5998,
+      "image_masks_loss": 0.0075,
+      "llm_learning_rate": 4.531477238364752e-06,
+      "loss": 0.7779,
+      "output_text_mask_loss": 0.0149,
+      "predictor_learning_rate": 0.0004531477238364752,
+      "step": 2168
+    },
+    {
+      "epoch": 0.22,
+      "gumbel_tau": 0.5996,
+      "image_masks_loss": 0.0186,
+      "llm_learning_rate": 4.530993597275581e-06,
+      "loss": 0.7257,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.00045309935972755806,
+      "step": 2169
+    },
+    {
+      "epoch": 0.22,
+      "gumbel_tau": 0.5995,
+      "image_masks_loss": 0.0092,
+      "llm_learning_rate": 4.530509732528871e-06,
+      "loss": 0.681,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 0.000453050973252887,
+      "step": 2170
+    },
+    {
+      "epoch": 0.22,
+      "gumbel_tau": 0.5993,
+      "image_masks_loss": 0.012,
+      "llm_learning_rate": 4.530025644177906e-06,
+      "loss": 0.7992,
+      "output_text_mask_loss": 0.0282,
+      "predictor_learning_rate": 0.0004530025644177906,
+      "step": 2171
+    },
+    {
+      "epoch": 0.22,
+      "gumbel_tau": 0.5992,
+      "image_masks_loss": 0.0138,
+      "llm_learning_rate": 4.529541332275996e-06,
+      "loss": 0.8162,
+      "output_text_mask_loss": 0.045,
+      "predictor_learning_rate": 0.0004529541332275996,
+      "step": 2172
+    },
+    {
+      "epoch": 0.22,
+      "gumbel_tau": 0.599,
+      "image_masks_loss": 0.0184,
+      "llm_learning_rate": 4.529056796876474e-06,
+      "loss": 0.7369,
+      "output_text_mask_loss": 0.0294,
+      "predictor_learning_rate": 0.00045290567968764734,
+      "step": 2173
+    },
+    {
+      "epoch": 0.22,
+      "gumbel_tau": 0.5989,
+      "image_masks_loss": 0.0202,
+      "llm_learning_rate": 4.528572038032698e-06,
+      "loss": 0.7537,
+      "output_text_mask_loss": 0.0557,
+      "predictor_learning_rate": 0.0004528572038032697,
+      "step": 2174
+    },
+    {
+      "epoch": 0.22,
+      "gumbel_tau": 0.5988,
+      "image_masks_loss": 0.0137,
+      "llm_learning_rate": 4.52808705579805e-06,
+      "loss": 0.8122,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00045280870557980494,
+      "step": 2175
+    },
+    {
+      "epoch": 0.22,
+      "gumbel_tau": 0.5986,
+      "image_masks_loss": 0.0148,
+      "llm_learning_rate": 4.527601850225939e-06,
+      "loss": 0.8558,
+      "output_text_mask_loss": 0.0328,
+      "predictor_learning_rate": 0.0004527601850225938,
+      "step": 2176
+    },
+    {
+      "epoch": 0.22,
+      "gumbel_tau": 0.5985,
+      "image_masks_loss": 0.0156,
+      "llm_learning_rate": 4.527116421369795e-06,
+      "loss": 0.7572,
+      "output_text_mask_loss": 0.0488,
+      "predictor_learning_rate": 0.00045271164213697945,
+      "step": 2177
+    },
+    {
+      "epoch": 0.22,
+      "gumbel_tau": 0.5983,
+      "image_masks_loss": 0.0254,
+      "llm_learning_rate": 4.526630769283076e-06,
+      "loss": 0.6239,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00045266307692830754,
+      "step": 2178
+    },
+    {
+      "epoch": 0.22,
+      "gumbel_tau": 0.5982,
+      "image_masks_loss": 0.0292,
+      "llm_learning_rate": 4.526144894019262e-06,
+      "loss": 0.865,
+      "output_text_mask_loss": 0.0345,
+      "predictor_learning_rate": 0.00045261448940192616,
+      "step": 2179
+    },
+    {
+      "epoch": 0.22,
+      "gumbel_tau": 0.5981,
+      "image_masks_loss": 0.0153,
+      "llm_learning_rate": 4.52565879563186e-06,
+      "loss": 0.8553,
+      "output_text_mask_loss": 0.0496,
+      "predictor_learning_rate": 0.00045256587956318594,
+      "step": 2180
+    },
+    {
+      "epoch": 0.22,
+      "gumbel_tau": 0.5979,
+      "image_masks_loss": 0.0108,
+      "llm_learning_rate": 4.525172474174398e-06,
+      "loss": 0.7905,
+      "output_text_mask_loss": 0.0277,
+      "predictor_learning_rate": 0.0004525172474174398,
+      "step": 2181
+    },
+    {
+      "epoch": 0.22,
+      "gumbel_tau": 0.5978,
+      "image_masks_loss": 0.0228,
+      "llm_learning_rate": 4.5246859297004346e-06,
+      "loss": 0.8034,
+      "output_text_mask_loss": 0.0235,
+      "predictor_learning_rate": 0.0004524685929700434,
+      "step": 2182
+    },
+    {
+      "epoch": 0.22,
+      "gumbel_tau": 0.5976,
+      "image_masks_loss": 0.0034,
+      "llm_learning_rate": 4.524199162263545e-06,
+      "loss": 0.7668,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 0.00045241991622635446,
+      "step": 2183
+    },
+    {
+      "epoch": 0.22,
+      "gumbel_tau": 0.5975,
+      "image_masks_loss": 0.0125,
+      "llm_learning_rate": 4.523712171917335e-06,
+      "loss": 0.6922,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004523712171917335,
+      "step": 2184
+    },
+    {
+      "epoch": 0.22,
+      "gumbel_tau": 0.5974,
+      "image_masks_loss": 0.0078,
+      "llm_learning_rate": 4.523224958715434e-06,
+      "loss": 0.7725,
+      "output_text_mask_loss": 0.0058,
+      "predictor_learning_rate": 0.00045232249587154333,
+      "step": 2185
+    },
+    {
+      "epoch": 0.22,
+      "gumbel_tau": 0.5972,
+      "image_masks_loss": 0.0269,
+      "llm_learning_rate": 4.5227375227114936e-06,
+      "loss": 0.7196,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004522737522711493,
+      "step": 2186
+    },
+    {
+      "epoch": 0.22,
+      "gumbel_tau": 0.5971,
+      "image_masks_loss": 0.0271,
+      "llm_learning_rate": 4.522249863959191e-06,
+      "loss": 0.7723,
+      "output_text_mask_loss": 0.0039,
+      "predictor_learning_rate": 0.0004522249863959191,
+      "step": 2187
+    },
+    {
+      "epoch": 0.22,
+      "gumbel_tau": 0.5969,
+      "image_masks_loss": 0.0162,
+      "llm_learning_rate": 4.52176198251223e-06,
+      "loss": 0.7396,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00045217619825122294,
+      "step": 2188
+    },
+    {
+      "epoch": 0.22,
+      "gumbel_tau": 0.5968,
+      "image_masks_loss": 0.0057,
+      "llm_learning_rate": 4.521273878424335e-06,
+      "loss": 0.7545,
+      "output_text_mask_loss": 0.021,
+      "predictor_learning_rate": 0.0004521273878424335,
+      "step": 2189
+    },
+    {
+      "epoch": 0.22,
+      "gumbel_tau": 0.5966,
+      "image_masks_loss": 0.0334,
+      "llm_learning_rate": 4.52078555174926e-06,
+      "loss": 0.8278,
+      "output_text_mask_loss": 0.0109,
+      "predictor_learning_rate": 0.0004520785551749259,
+      "step": 2190
+    },
+    {
+      "epoch": 0.22,
+      "gumbel_tau": 0.5965,
+      "image_masks_loss": 0.0124,
+      "llm_learning_rate": 4.520297002540778e-06,
+      "loss": 0.7319,
+      "output_text_mask_loss": 0.0164,
+      "predictor_learning_rate": 0.00045202970025407776,
+      "step": 2191
+    },
+    {
+      "epoch": 0.22,
+      "gumbel_tau": 0.5964,
+      "image_masks_loss": 0.0103,
+      "llm_learning_rate": 4.519808230852689e-06,
+      "loss": 0.751,
+      "output_text_mask_loss": 0.0326,
+      "predictor_learning_rate": 0.0004519808230852689,
+      "step": 2192
+    },
+    {
+      "epoch": 0.22,
+      "gumbel_tau": 0.5962,
+      "image_masks_loss": 0.0218,
+      "llm_learning_rate": 4.51931923673882e-06,
+      "loss": 0.6923,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.000451931923673882,
+      "step": 2193
+    },
+    {
+      "epoch": 0.22,
+      "gumbel_tau": 0.5961,
+      "image_masks_loss": 0.0546,
+      "llm_learning_rate": 4.518830020253018e-06,
+      "loss": 0.735,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 0.0004518830020253018,
+      "step": 2194
+    },
+    {
+      "epoch": 0.22,
+      "gumbel_tau": 0.5959,
+      "image_masks_loss": 0.0248,
+      "llm_learning_rate": 4.518340581449158e-06,
+      "loss": 0.7785,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.0004518340581449158,
+      "step": 2195
+    },
+    {
+      "epoch": 0.23,
+      "gumbel_tau": 0.5958,
+      "image_masks_loss": 0.0096,
+      "llm_learning_rate": 4.517850920381138e-06,
+      "loss": 0.7218,
+      "output_text_mask_loss": 0.1129,
+      "predictor_learning_rate": 0.00045178509203811374,
+      "step": 2196
+    },
+    {
+      "epoch": 0.23,
+      "gumbel_tau": 0.5957,
+      "image_masks_loss": 0.0311,
+      "llm_learning_rate": 4.517361037102878e-06,
+      "loss": 0.8019,
+      "output_text_mask_loss": 0.0039,
+      "predictor_learning_rate": 0.00045173610371028786,
+      "step": 2197
+    },
+    {
+      "epoch": 0.23,
+      "gumbel_tau": 0.5955,
+      "image_masks_loss": 0.0072,
+      "llm_learning_rate": 4.51687093166833e-06,
+      "loss": 0.7321,
+      "output_text_mask_loss": 0.0441,
+      "predictor_learning_rate": 0.0004516870931668329,
+      "step": 2198
+    },
+    {
+      "epoch": 0.23,
+      "gumbel_tau": 0.5954,
+      "image_masks_loss": 0.0136,
+      "llm_learning_rate": 4.51638060413146e-06,
+      "loss": 0.7518,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.00045163806041314603,
+      "step": 2199
+    },
+    {
+      "epoch": 0.23,
+      "gumbel_tau": 0.5952,
+      "image_masks_loss": 0.0188,
+      "llm_learning_rate": 4.5158900545462684e-06,
+      "loss": 0.8303,
+      "output_text_mask_loss": 0.1968,
+      "predictor_learning_rate": 0.0004515890054546268,
+      "step": 2200
+    },
+    {
+      "epoch": 0.23,
+      "gumbel_tau": 0.5951,
+      "image_masks_loss": 0.0168,
+      "llm_learning_rate": 4.515399282966773e-06,
+      "loss": 0.7797,
+      "output_text_mask_loss": 0.0128,
+      "predictor_learning_rate": 0.00045153992829667725,
+      "step": 2201
+    },
+    {
+      "epoch": 0.23,
+      "gumbel_tau": 0.595,
+      "image_masks_loss": 0.042,
+      "llm_learning_rate": 4.51490828944702e-06,
+      "loss": 0.8304,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 0.0004514908289447019,
+      "step": 2202
+    },
+    {
+      "epoch": 0.23,
+      "gumbel_tau": 0.5948,
+      "image_masks_loss": 0.0694,
+      "llm_learning_rate": 4.514417074041077e-06,
+      "loss": 0.8009,
+      "output_text_mask_loss": 0.013,
+      "predictor_learning_rate": 0.0004514417074041077,
+      "step": 2203
+    },
+    {
+      "epoch": 0.23,
+      "gumbel_tau": 0.5947,
+      "image_masks_loss": 0.0127,
+      "llm_learning_rate": 4.513925636803039e-06,
+      "loss": 0.7247,
+      "output_text_mask_loss": 0.0446,
+      "predictor_learning_rate": 0.0004513925636803039,
+      "step": 2204
+    },
+    {
+      "epoch": 0.23,
+      "gumbel_tau": 0.5945,
+      "image_masks_loss": 0.0313,
+      "llm_learning_rate": 4.5134339777870255e-06,
+      "loss": 0.8306,
+      "output_text_mask_loss": 0.0107,
+      "predictor_learning_rate": 0.0004513433977787025,
+      "step": 2205
+    },
+    {
+      "epoch": 0.23,
+      "gumbel_tau": 0.5944,
+      "image_masks_loss": 0.0332,
+      "llm_learning_rate": 4.512942097047176e-06,
+      "loss": 0.7265,
+      "output_text_mask_loss": 0.0101,
+      "predictor_learning_rate": 0.0004512942097047176,
+      "step": 2206
+    },
+    {
+      "epoch": 0.23,
+      "gumbel_tau": 0.5943,
+      "image_masks_loss": 0.0401,
+      "llm_learning_rate": 4.51244999463766e-06,
+      "loss": 0.77,
+      "output_text_mask_loss": 0.0816,
+      "predictor_learning_rate": 0.000451244999463766,
+      "step": 2207
+    },
+    {
+      "epoch": 0.23,
+      "gumbel_tau": 0.5941,
+      "image_masks_loss": 0.0416,
+      "llm_learning_rate": 4.5119576706126675e-06,
+      "loss": 0.788,
+      "output_text_mask_loss": 0.0137,
+      "predictor_learning_rate": 0.00045119576706126676,
+      "step": 2208
+    },
+    {
+      "epoch": 0.23,
+      "gumbel_tau": 0.594,
+      "image_masks_loss": 0.0725,
+      "llm_learning_rate": 4.511465125026415e-06,
+      "loss": 0.8023,
+      "output_text_mask_loss": 0.0095,
+      "predictor_learning_rate": 0.0004511465125026415,
+      "step": 2209
+    },
+    {
+      "epoch": 0.23,
+      "gumbel_tau": 0.5938,
+      "image_masks_loss": 0.0219,
+      "llm_learning_rate": 4.510972357933143e-06,
+      "loss": 0.8142,
+      "output_text_mask_loss": 0.0086,
+      "predictor_learning_rate": 0.0004510972357933143,
+      "step": 2210
+    },
+    {
+      "epoch": 0.23,
+      "gumbel_tau": 0.5937,
+      "image_masks_loss": 0.0147,
+      "llm_learning_rate": 4.5104793693871155e-06,
+      "loss": 0.835,
+      "output_text_mask_loss": 0.0038,
+      "predictor_learning_rate": 0.0004510479369387115,
+      "step": 2211
+    },
+    {
+      "epoch": 0.23,
+      "gumbel_tau": 0.5936,
+      "image_masks_loss": 0.0229,
+      "llm_learning_rate": 4.509986159442621e-06,
+      "loss": 0.8014,
+      "output_text_mask_loss": 0.0039,
+      "predictor_learning_rate": 0.00045099861594426205,
+      "step": 2212
+    },
+    {
+      "epoch": 0.23,
+      "gumbel_tau": 0.5934,
+      "image_masks_loss": 0.0519,
+      "llm_learning_rate": 4.509492728153975e-06,
+      "loss": 0.8461,
+      "output_text_mask_loss": 0.033,
+      "predictor_learning_rate": 0.0004509492728153974,
+      "step": 2213
+    },
+    {
+      "epoch": 0.23,
+      "gumbel_tau": 0.5933,
+      "image_masks_loss": 0.0194,
+      "llm_learning_rate": 4.508999075575512e-06,
+      "loss": 0.7937,
+      "output_text_mask_loss": 0.0034,
+      "predictor_learning_rate": 0.0004508999075575512,
+      "step": 2214
+    },
+    {
+      "epoch": 0.23,
+      "gumbel_tau": 0.5931,
+      "image_masks_loss": 0.0224,
+      "llm_learning_rate": 4.508505201761597e-06,
+      "loss": 0.7817,
+      "output_text_mask_loss": 0.1251,
+      "predictor_learning_rate": 0.00045085052017615967,
+      "step": 2215
+    },
+    {
+      "epoch": 0.23,
+      "gumbel_tau": 0.593,
+      "image_masks_loss": 0.0477,
+      "llm_learning_rate": 4.508011106766616e-06,
+      "loss": 0.781,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 0.00045080111067666154,
+      "step": 2216
+    },
+    {
+      "epoch": 0.23,
+      "gumbel_tau": 0.5929,
+      "image_masks_loss": 0.0045,
+      "llm_learning_rate": 4.507516790644978e-06,
+      "loss": 0.7407,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00045075167906449777,
+      "step": 2217
+    },
+    {
+      "epoch": 0.23,
+      "gumbel_tau": 0.5927,
+      "image_masks_loss": 0.0197,
+      "llm_learning_rate": 4.50702225345112e-06,
+      "loss": 0.8316,
+      "output_text_mask_loss": 0.0011,
+      "predictor_learning_rate": 0.00045070222534511196,
+      "step": 2218
+    },
+    {
+      "epoch": 0.23,
+      "gumbel_tau": 0.5926,
+      "image_masks_loss": 0.0115,
+      "llm_learning_rate": 4.506527495239501e-06,
+      "loss": 0.7665,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004506527495239501,
+      "step": 2219
+    },
+    {
+      "epoch": 0.23,
+      "gumbel_tau": 0.5924,
+      "image_masks_loss": 0.0889,
+      "llm_learning_rate": 4.506032516064604e-06,
+      "loss": 0.9282,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 0.00045060325160646043,
+      "step": 2220
+    },
+    {
+      "epoch": 0.23,
+      "gumbel_tau": 0.5923,
+      "image_masks_loss": 0.0123,
+      "llm_learning_rate": 4.50553731598094e-06,
+      "loss": 0.7499,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.0004505537315980939,
+      "step": 2221
+    },
+    {
+      "epoch": 0.23,
+      "gumbel_tau": 0.5922,
+      "image_masks_loss": 0.0652,
+      "llm_learning_rate": 4.505041895043038e-06,
+      "loss": 0.7806,
+      "output_text_mask_loss": 0.0275,
+      "predictor_learning_rate": 0.0004505041895043037,
+      "step": 2222
+    },
+    {
+      "epoch": 0.23,
+      "gumbel_tau": 0.592,
+      "image_masks_loss": 0.0172,
+      "llm_learning_rate": 4.504546253305457e-06,
+      "loss": 0.7129,
+      "output_text_mask_loss": 0.0467,
+      "predictor_learning_rate": 0.0004504546253305457,
+      "step": 2223
+    },
+    {
+      "epoch": 0.23,
+      "gumbel_tau": 0.5919,
+      "image_masks_loss": 0.0546,
+      "llm_learning_rate": 4.504050390822777e-06,
+      "loss": 0.7107,
+      "output_text_mask_loss": 0.0153,
+      "predictor_learning_rate": 0.0004504050390822777,
+      "step": 2224
+    },
+    {
+      "epoch": 0.23,
+      "gumbel_tau": 0.5917,
+      "image_masks_loss": 0.0254,
+      "llm_learning_rate": 4.503554307649605e-06,
+      "loss": 0.8096,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004503554307649604,
+      "step": 2225
+    },
+    {
+      "epoch": 0.23,
+      "gumbel_tau": 0.5916,
+      "image_masks_loss": 0.0294,
+      "llm_learning_rate": 4.503058003840569e-06,
+      "loss": 0.789,
+      "output_text_mask_loss": 0.0103,
+      "predictor_learning_rate": 0.0004503058003840569,
+      "step": 2226
+    },
+    {
+      "epoch": 0.23,
+      "gumbel_tau": 0.5915,
+      "image_masks_loss": 0.0469,
+      "llm_learning_rate": 4.5025614794503255e-06,
+      "loss": 0.8311,
+      "output_text_mask_loss": 0.0322,
+      "predictor_learning_rate": 0.0004502561479450325,
+      "step": 2227
+    },
+    {
+      "epoch": 0.23,
+      "gumbel_tau": 0.5913,
+      "image_masks_loss": 0.0209,
+      "llm_learning_rate": 4.50206473453355e-06,
+      "loss": 0.6808,
+      "output_text_mask_loss": 0.016,
+      "predictor_learning_rate": 0.00045020647345335495,
+      "step": 2228
+    },
+    {
+      "epoch": 0.23,
+      "gumbel_tau": 0.5912,
+      "image_masks_loss": 0.0435,
+      "llm_learning_rate": 4.501567769144947e-06,
+      "loss": 0.7443,
+      "output_text_mask_loss": 0.0058,
+      "predictor_learning_rate": 0.0004501567769144947,
+      "step": 2229
+    },
+    {
+      "epoch": 0.23,
+      "gumbel_tau": 0.591,
+      "image_masks_loss": 0.0187,
+      "llm_learning_rate": 4.501070583339243e-06,
+      "loss": 0.8526,
+      "output_text_mask_loss": 0.1007,
+      "predictor_learning_rate": 0.00045010705833392427,
+      "step": 2230
+    },
+    {
+      "epoch": 0.23,
+      "gumbel_tau": 0.5909,
+      "image_masks_loss": 0.0269,
+      "llm_learning_rate": 4.500573177171189e-06,
+      "loss": 0.7896,
+      "output_text_mask_loss": 0.006,
+      "predictor_learning_rate": 0.0004500573177171189,
+      "step": 2231
+    },
+    {
+      "epoch": 0.23,
+      "gumbel_tau": 0.5908,
+      "image_masks_loss": 0.0271,
+      "llm_learning_rate": 4.5000755506955605e-06,
+      "loss": 0.7318,
+      "output_text_mask_loss": 0.0095,
+      "predictor_learning_rate": 0.00045000755506955606,
+      "step": 2232
+    },
+    {
+      "epoch": 0.23,
+      "gumbel_tau": 0.5906,
+      "image_masks_loss": 0.0256,
+      "llm_learning_rate": 4.499577703967158e-06,
+      "loss": 0.7964,
+      "output_text_mask_loss": 0.0058,
+      "predictor_learning_rate": 0.00044995777039671573,
+      "step": 2233
+    },
+    {
+      "epoch": 0.23,
+      "gumbel_tau": 0.5905,
+      "image_masks_loss": 0.0299,
+      "llm_learning_rate": 4.499079637040804e-06,
+      "loss": 0.8116,
+      "output_text_mask_loss": 0.0132,
+      "predictor_learning_rate": 0.0004499079637040804,
+      "step": 2234
+    },
+    {
+      "epoch": 0.23,
+      "gumbel_tau": 0.5903,
+      "image_masks_loss": 0.0345,
+      "llm_learning_rate": 4.498581349971348e-06,
+      "loss": 0.7007,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.0004498581349971348,
+      "step": 2235
+    },
+    {
+      "epoch": 0.23,
+      "gumbel_tau": 0.5902,
+      "image_masks_loss": 0.0492,
+      "llm_learning_rate": 4.498082842813662e-06,
+      "loss": 0.8179,
+      "output_text_mask_loss": 0.0385,
+      "predictor_learning_rate": 0.0004498082842813662,
+      "step": 2236
+    },
+    {
+      "epoch": 0.23,
+      "gumbel_tau": 0.5901,
+      "image_masks_loss": 0.0063,
+      "llm_learning_rate": 4.497584115622643e-06,
+      "loss": 0.7563,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.0004497584115622643,
+      "step": 2237
+    },
+    {
+      "epoch": 0.23,
+      "gumbel_tau": 0.5899,
+      "image_masks_loss": 0.0269,
+      "llm_learning_rate": 4.497085168453211e-06,
+      "loss": 0.8242,
+      "output_text_mask_loss": 0.0056,
+      "predictor_learning_rate": 0.0004497085168453211,
+      "step": 2238
+    },
+    {
+      "epoch": 0.23,
+      "gumbel_tau": 0.5898,
+      "image_masks_loss": 0.0433,
+      "llm_learning_rate": 4.496586001360312e-06,
+      "loss": 0.8497,
+      "output_text_mask_loss": 0.0231,
+      "predictor_learning_rate": 0.0004496586001360312,
+      "step": 2239
+    },
+    {
+      "epoch": 0.23,
+      "gumbel_tau": 0.5896,
+      "image_masks_loss": 0.0402,
+      "llm_learning_rate": 4.496086614398917e-06,
+      "loss": 0.6689,
+      "output_text_mask_loss": 0.0187,
+      "predictor_learning_rate": 0.0004496086614398917,
+      "step": 2240
+    },
+    {
+      "epoch": 0.23,
+      "gumbel_tau": 0.5895,
+      "image_masks_loss": 0.0252,
+      "llm_learning_rate": 4.495587007624016e-06,
+      "loss": 0.8065,
+      "output_text_mask_loss": 0.0282,
+      "predictor_learning_rate": 0.0004495587007624016,
+      "step": 2241
+    },
+    {
+      "epoch": 0.23,
+      "gumbel_tau": 0.5894,
+      "image_masks_loss": 0.0127,
+      "llm_learning_rate": 4.4950871810906295e-06,
+      "loss": 0.774,
+      "output_text_mask_loss": 0.0591,
+      "predictor_learning_rate": 0.00044950871810906293,
+      "step": 2242
+    },
+    {
+      "epoch": 0.23,
+      "gumbel_tau": 0.5892,
+      "image_masks_loss": 0.0576,
+      "llm_learning_rate": 4.494587134853799e-06,
+      "loss": 0.7675,
+      "output_text_mask_loss": 0.0231,
+      "predictor_learning_rate": 0.00044945871348537985,
+      "step": 2243
+    },
+    {
+      "epoch": 0.23,
+      "gumbel_tau": 0.5891,
+      "image_masks_loss": 0.0183,
+      "llm_learning_rate": 4.49408686896859e-06,
+      "loss": 0.697,
+      "output_text_mask_loss": 0.0158,
+      "predictor_learning_rate": 0.000449408686896859,
+      "step": 2244
+    },
+    {
+      "epoch": 0.23,
+      "gumbel_tau": 0.589,
+      "image_masks_loss": 0.0242,
+      "llm_learning_rate": 4.493586383490093e-06,
+      "loss": 0.7741,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004493586383490093,
+      "step": 2245
+    },
+    {
+      "epoch": 0.23,
+      "gumbel_tau": 0.5888,
+      "image_masks_loss": 0.0113,
+      "llm_learning_rate": 4.493085678473423e-06,
+      "loss": 0.6832,
+      "output_text_mask_loss": 0.0275,
+      "predictor_learning_rate": 0.0004493085678473423,
+      "step": 2246
+    },
+    {
+      "epoch": 0.23,
+      "gumbel_tau": 0.5887,
+      "image_masks_loss": 0.0214,
+      "llm_learning_rate": 4.49258475397372e-06,
+      "loss": 0.7126,
+      "output_text_mask_loss": 0.0256,
+      "predictor_learning_rate": 0.000449258475397372,
+      "step": 2247
+    },
+    {
+      "epoch": 0.23,
+      "gumbel_tau": 0.5885,
+      "image_masks_loss": 0.0221,
+      "llm_learning_rate": 4.492083610046144e-06,
+      "loss": 0.6875,
+      "output_text_mask_loss": 0.0246,
+      "predictor_learning_rate": 0.00044920836100461437,
+      "step": 2248
+    },
+    {
+      "epoch": 0.23,
+      "gumbel_tau": 0.5884,
+      "image_masks_loss": 0.0305,
+      "llm_learning_rate": 4.491582246745885e-06,
+      "loss": 0.75,
+      "output_text_mask_loss": 0.1305,
+      "predictor_learning_rate": 0.00044915822467458846,
+      "step": 2249
+    },
+    {
+      "epoch": 0.23,
+      "gumbel_tau": 0.5883,
+      "image_masks_loss": 0.0467,
+      "llm_learning_rate": 4.491080664128152e-06,
+      "loss": 0.738,
+      "output_text_mask_loss": 0.0162,
+      "predictor_learning_rate": 0.00044910806641281517,
+      "step": 2250
+    },
+    {
+      "epoch": 0.23,
+      "gumbel_tau": 0.5881,
+      "image_masks_loss": 0.0115,
+      "llm_learning_rate": 4.490578862248182e-06,
+      "loss": 0.8613,
+      "output_text_mask_loss": 0.0101,
+      "predictor_learning_rate": 0.00044905788622481813,
+      "step": 2251
+    },
+    {
+      "epoch": 0.23,
+      "gumbel_tau": 0.588,
+      "image_masks_loss": 0.0267,
+      "llm_learning_rate": 4.4900768411612325e-06,
+      "loss": 0.8933,
+      "output_text_mask_loss": 0.0353,
+      "predictor_learning_rate": 0.00044900768411612325,
+      "step": 2252
+    },
+    {
+      "epoch": 0.23,
+      "gumbel_tau": 0.5878,
+      "image_masks_loss": 0.0298,
+      "llm_learning_rate": 4.4895746009225895e-06,
+      "loss": 0.774,
+      "output_text_mask_loss": 0.0923,
+      "predictor_learning_rate": 0.00044895746009225893,
+      "step": 2253
+    },
+    {
+      "epoch": 0.23,
+      "gumbel_tau": 0.5877,
+      "image_masks_loss": 0.0399,
+      "llm_learning_rate": 4.489072141587559e-06,
+      "loss": 0.7512,
+      "output_text_mask_loss": 0.0231,
+      "predictor_learning_rate": 0.0004489072141587559,
+      "step": 2254
+    },
+    {
+      "epoch": 0.23,
+      "gumbel_tau": 0.5876,
+      "image_masks_loss": 0.0208,
+      "llm_learning_rate": 4.488569463211475e-06,
+      "loss": 0.773,
+      "output_text_mask_loss": 0.1106,
+      "predictor_learning_rate": 0.0004488569463211475,
+      "step": 2255
+    },
+    {
+      "epoch": 0.23,
+      "gumbel_tau": 0.5874,
+      "image_masks_loss": 0.0269,
+      "llm_learning_rate": 4.488066565849692e-06,
+      "loss": 0.7547,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.0004488066565849692,
+      "step": 2256
+    },
+    {
+      "epoch": 0.23,
+      "gumbel_tau": 0.5873,
+      "image_masks_loss": 0.0049,
+      "llm_learning_rate": 4.487563449557591e-06,
+      "loss": 0.8011,
+      "output_text_mask_loss": 0.0446,
+      "predictor_learning_rate": 0.000448756344955759,
+      "step": 2257
+    },
+    {
+      "epoch": 0.23,
+      "gumbel_tau": 0.5872,
+      "image_masks_loss": 0.0029,
+      "llm_learning_rate": 4.487060114390574e-06,
+      "loss": 0.7522,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.00044870601143905743,
+      "step": 2258
+    },
+    {
+      "epoch": 0.23,
+      "gumbel_tau": 0.587,
+      "image_masks_loss": 0.0232,
+      "llm_learning_rate": 4.486556560404073e-06,
+      "loss": 0.6695,
+      "output_text_mask_loss": 0.0277,
+      "predictor_learning_rate": 0.0004486556560404073,
+      "step": 2259
+    },
+    {
+      "epoch": 0.23,
+      "gumbel_tau": 0.5869,
+      "image_masks_loss": 0.0273,
+      "llm_learning_rate": 4.486052787653538e-06,
+      "loss": 0.7091,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.0004486052787653537,
+      "step": 2260
+    },
+    {
+      "epoch": 0.23,
+      "gumbel_tau": 0.5867,
+      "image_masks_loss": 0.0191,
+      "llm_learning_rate": 4.485548796194447e-06,
+      "loss": 0.7843,
+      "output_text_mask_loss": 0.0055,
+      "predictor_learning_rate": 0.00044855487961944465,
+      "step": 2261
+    },
+    {
+      "epoch": 0.23,
+      "gumbel_tau": 0.5866,
+      "image_masks_loss": 0.024,
+      "llm_learning_rate": 4.4850445860822986e-06,
+      "loss": 0.7783,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.0004485044586082298,
+      "step": 2262
+    },
+    {
+      "epoch": 0.23,
+      "gumbel_tau": 0.5865,
+      "image_masks_loss": 0.0259,
+      "llm_learning_rate": 4.484540157372619e-06,
+      "loss": 0.824,
+      "output_text_mask_loss": 0.0248,
+      "predictor_learning_rate": 0.0004484540157372619,
+      "step": 2263
+    },
+    {
+      "epoch": 0.23,
+      "gumbel_tau": 0.5863,
+      "image_masks_loss": 0.0089,
+      "llm_learning_rate": 4.484035510120958e-06,
+      "loss": 0.837,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 0.00044840355101209576,
+      "step": 2264
+    },
+    {
+      "epoch": 0.23,
+      "gumbel_tau": 0.5862,
+      "image_masks_loss": 0.0129,
+      "llm_learning_rate": 4.483530644382886e-06,
+      "loss": 0.772,
+      "output_text_mask_loss": 0.0072,
+      "predictor_learning_rate": 0.0004483530644382886,
+      "step": 2265
+    },
+    {
+      "epoch": 0.23,
+      "gumbel_tau": 0.586,
+      "image_masks_loss": 0.0166,
+      "llm_learning_rate": 4.4830255602140025e-06,
+      "loss": 0.8232,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 0.00044830255602140023,
+      "step": 2266
+    },
+    {
+      "epoch": 0.23,
+      "gumbel_tau": 0.5859,
+      "image_masks_loss": 0.0064,
+      "llm_learning_rate": 4.482520257669927e-06,
+      "loss": 0.755,
+      "output_text_mask_loss": 0.0058,
+      "predictor_learning_rate": 0.00044825202576699266,
+      "step": 2267
+    },
+    {
+      "epoch": 0.23,
+      "gumbel_tau": 0.5858,
+      "image_masks_loss": 0.029,
+      "llm_learning_rate": 4.482014736806305e-06,
+      "loss": 0.7264,
+      "output_text_mask_loss": 0.0039,
+      "predictor_learning_rate": 0.0004482014736806305,
+      "step": 2268
+    },
+    {
+      "epoch": 0.23,
+      "gumbel_tau": 0.5856,
+      "image_masks_loss": 0.0282,
+      "llm_learning_rate": 4.481508997678805e-06,
+      "loss": 0.7111,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.0004481508997678805,
+      "step": 2269
+    },
+    {
+      "epoch": 0.23,
+      "gumbel_tau": 0.5855,
+      "image_masks_loss": 0.0328,
+      "llm_learning_rate": 4.48100304034312e-06,
+      "loss": 0.7457,
+      "output_text_mask_loss": 0.0076,
+      "predictor_learning_rate": 0.000448100304034312,
+      "step": 2270
+    },
+    {
+      "epoch": 0.23,
+      "gumbel_tau": 0.5854,
+      "image_masks_loss": 0.0192,
+      "llm_learning_rate": 4.480496864854968e-06,
+      "loss": 0.8173,
+      "output_text_mask_loss": 0.0626,
+      "predictor_learning_rate": 0.0004480496864854968,
+      "step": 2271
+    },
+    {
+      "epoch": 0.23,
+      "gumbel_tau": 0.5852,
+      "image_masks_loss": 0.0143,
+      "llm_learning_rate": 4.479990471270091e-06,
+      "loss": 0.7196,
+      "output_text_mask_loss": 0.0143,
+      "predictor_learning_rate": 0.0004479990471270091,
+      "step": 2272
+    },
+    {
+      "epoch": 0.23,
+      "gumbel_tau": 0.5851,
+      "image_masks_loss": 0.0248,
+      "llm_learning_rate": 4.4794838596442525e-06,
+      "loss": 0.7869,
+      "output_text_mask_loss": 0.006,
+      "predictor_learning_rate": 0.00044794838596442523,
+      "step": 2273
+    },
+    {
+      "epoch": 0.23,
+      "gumbel_tau": 0.5849,
+      "image_masks_loss": 0.0372,
+      "llm_learning_rate": 4.4789770300332415e-06,
+      "loss": 0.6881,
+      "output_text_mask_loss": 0.0202,
+      "predictor_learning_rate": 0.0004478977030033241,
+      "step": 2274
+    },
+    {
+      "epoch": 0.23,
+      "gumbel_tau": 0.5848,
+      "image_masks_loss": 0.0378,
+      "llm_learning_rate": 4.4784699824928725e-06,
+      "loss": 0.7183,
+      "output_text_mask_loss": 0.0134,
+      "predictor_learning_rate": 0.0004478469982492872,
+      "step": 2275
+    },
+    {
+      "epoch": 0.23,
+      "gumbel_tau": 0.5847,
+      "image_masks_loss": 0.0225,
+      "llm_learning_rate": 4.477962717078982e-06,
+      "loss": 0.6747,
+      "output_text_mask_loss": 0.0202,
+      "predictor_learning_rate": 0.0004477962717078982,
+      "step": 2276
+    },
+    {
+      "epoch": 0.23,
+      "gumbel_tau": 0.5845,
+      "image_masks_loss": 0.0553,
+      "llm_learning_rate": 4.4774552338474315e-06,
+      "loss": 0.8612,
+      "output_text_mask_loss": 0.0322,
+      "predictor_learning_rate": 0.0004477455233847431,
+      "step": 2277
+    },
+    {
+      "epoch": 0.23,
+      "gumbel_tau": 0.5844,
+      "image_masks_loss": 0.0133,
+      "llm_learning_rate": 4.476947532854106e-06,
+      "loss": 0.7044,
+      "output_text_mask_loss": 0.0378,
+      "predictor_learning_rate": 0.00044769475328541056,
+      "step": 2278
+    },
+    {
+      "epoch": 0.23,
+      "gumbel_tau": 0.5842,
+      "image_masks_loss": 0.0192,
+      "llm_learning_rate": 4.476439614154915e-06,
+      "loss": 0.7277,
+      "output_text_mask_loss": 0.0124,
+      "predictor_learning_rate": 0.00044764396141549145,
+      "step": 2279
+    },
+    {
+      "epoch": 0.23,
+      "gumbel_tau": 0.5841,
+      "image_masks_loss": 0.0296,
+      "llm_learning_rate": 4.4759314778057904e-06,
+      "loss": 0.6223,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00044759314778057904,
+      "step": 2280
+    },
+    {
+      "epoch": 0.23,
+      "gumbel_tau": 0.584,
+      "image_masks_loss": 0.0275,
+      "llm_learning_rate": 4.475423123862691e-06,
+      "loss": 0.8403,
+      "output_text_mask_loss": 0.0488,
+      "predictor_learning_rate": 0.0004475423123862691,
+      "step": 2281
+    },
+    {
+      "epoch": 0.23,
+      "gumbel_tau": 0.5838,
+      "image_masks_loss": 0.0271,
+      "llm_learning_rate": 4.474914552381597e-06,
+      "loss": 0.7578,
+      "output_text_mask_loss": 0.0277,
+      "predictor_learning_rate": 0.0004474914552381597,
+      "step": 2282
+    },
+    {
+      "epoch": 0.23,
+      "gumbel_tau": 0.5837,
+      "image_masks_loss": 0.0271,
+      "llm_learning_rate": 4.474405763418514e-06,
+      "loss": 0.8089,
+      "output_text_mask_loss": 0.0111,
+      "predictor_learning_rate": 0.0004474405763418513,
+      "step": 2283
+    },
+    {
+      "epoch": 0.23,
+      "gumbel_tau": 0.5836,
+      "image_masks_loss": 0.0553,
+      "llm_learning_rate": 4.47389675702947e-06,
+      "loss": 0.8658,
+      "output_text_mask_loss": 0.0322,
+      "predictor_learning_rate": 0.00044738967570294696,
+      "step": 2284
+    },
+    {
+      "epoch": 0.23,
+      "gumbel_tau": 0.5834,
+      "image_masks_loss": 0.0172,
+      "llm_learning_rate": 4.4733875332705176e-06,
+      "loss": 0.7257,
+      "output_text_mask_loss": 0.0065,
+      "predictor_learning_rate": 0.00044733875332705176,
+      "step": 2285
+    },
+    {
+      "epoch": 0.23,
+      "gumbel_tau": 0.5833,
+      "image_masks_loss": 0.0246,
+      "llm_learning_rate": 4.472878092197736e-06,
+      "loss": 0.6864,
+      "output_text_mask_loss": 0.0038,
+      "predictor_learning_rate": 0.0004472878092197736,
+      "step": 2286
+    },
+    {
+      "epoch": 0.23,
+      "gumbel_tau": 0.5831,
+      "image_masks_loss": 0.0109,
+      "llm_learning_rate": 4.472368433867224e-06,
+      "loss": 0.7321,
+      "output_text_mask_loss": 0.0025,
+      "predictor_learning_rate": 0.0004472368433867224,
+      "step": 2287
+    },
+    {
+      "epoch": 0.23,
+      "gumbel_tau": 0.583,
+      "image_masks_loss": 0.0222,
+      "llm_learning_rate": 4.471858558335107e-06,
+      "loss": 0.8118,
+      "output_text_mask_loss": 0.0656,
+      "predictor_learning_rate": 0.0004471858558335106,
+      "step": 2288
+    },
+    {
+      "epoch": 0.23,
+      "gumbel_tau": 0.5829,
+      "image_masks_loss": 0.0146,
+      "llm_learning_rate": 4.471348465657532e-06,
+      "loss": 0.7591,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.00044713484656575323,
+      "step": 2289
+    },
+    {
+      "epoch": 0.23,
+      "gumbel_tau": 0.5827,
+      "image_masks_loss": 0.0126,
+      "llm_learning_rate": 4.470838155890675e-06,
+      "loss": 0.8273,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 0.0004470838155890674,
+      "step": 2290
+    },
+    {
+      "epoch": 0.23,
+      "gumbel_tau": 0.5826,
+      "image_masks_loss": 0.0097,
+      "llm_learning_rate": 4.470327629090729e-06,
+      "loss": 0.7684,
+      "output_text_mask_loss": 0.0336,
+      "predictor_learning_rate": 0.00044703276290907287,
+      "step": 2291
+    },
+    {
+      "epoch": 0.23,
+      "gumbel_tau": 0.5825,
+      "image_masks_loss": 0.0061,
+      "llm_learning_rate": 4.4698168853139165e-06,
+      "loss": 0.7795,
+      "output_text_mask_loss": 0.0015,
+      "predictor_learning_rate": 0.0004469816885313916,
+      "step": 2292
+    },
+    {
+      "epoch": 0.23,
+      "gumbel_tau": 0.5823,
+      "image_masks_loss": 0.0315,
+      "llm_learning_rate": 4.46930592461648e-06,
+      "loss": 0.9366,
+      "output_text_mask_loss": 0.0056,
+      "predictor_learning_rate": 0.000446930592461648,
+      "step": 2293
+    },
+    {
+      "epoch": 0.24,
+      "gumbel_tau": 0.5822,
+      "image_masks_loss": 0.0148,
+      "llm_learning_rate": 4.4687947470546895e-06,
+      "loss": 0.7279,
+      "output_text_mask_loss": 0.0786,
+      "predictor_learning_rate": 0.0004468794747054689,
+      "step": 2294
+    },
+    {
+      "epoch": 0.24,
+      "gumbel_tau": 0.582,
+      "image_masks_loss": 0.0133,
+      "llm_learning_rate": 4.468283352684836e-06,
+      "loss": 0.8196,
+      "output_text_mask_loss": 0.0648,
+      "predictor_learning_rate": 0.0004468283352684835,
+      "step": 2295
+    },
+    {
+      "epoch": 0.24,
+      "gumbel_tau": 0.5819,
+      "image_masks_loss": 0.0117,
+      "llm_learning_rate": 4.467771741563235e-06,
+      "loss": 0.8462,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 0.0004467771741563235,
+      "step": 2296
+    },
+    {
+      "epoch": 0.24,
+      "gumbel_tau": 0.5818,
+      "image_masks_loss": 0.0142,
+      "llm_learning_rate": 4.467259913746228e-06,
+      "loss": 0.8376,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.0004467259913746228,
+      "step": 2297
+    },
+    {
+      "epoch": 0.24,
+      "gumbel_tau": 0.5816,
+      "image_masks_loss": 0.0254,
+      "llm_learning_rate": 4.466747869290176e-06,
+      "loss": 0.7645,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.00044667478692901765,
+      "step": 2298
+    },
+    {
+      "epoch": 0.24,
+      "gumbel_tau": 0.5815,
+      "image_masks_loss": 0.0073,
+      "llm_learning_rate": 4.466235608251469e-06,
+      "loss": 0.8121,
+      "output_text_mask_loss": 0.0162,
+      "predictor_learning_rate": 0.0004466235608251469,
+      "step": 2299
+    },
+    {
+      "epoch": 0.24,
+      "gumbel_tau": 0.5814,
+      "image_masks_loss": 0.0244,
+      "llm_learning_rate": 4.465723130686517e-06,
+      "loss": 0.7293,
+      "output_text_mask_loss": 0.0235,
+      "predictor_learning_rate": 0.00044657231306865163,
+      "step": 2300
+    },
+    {
+      "epoch": 0.24,
+      "gumbel_tau": 0.5812,
+      "image_masks_loss": 0.0188,
+      "llm_learning_rate": 4.465210436651755e-06,
+      "loss": 0.7912,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.00044652104366517544,
+      "step": 2301
+    },
+    {
+      "epoch": 0.24,
+      "gumbel_tau": 0.5811,
+      "image_masks_loss": 0.0175,
+      "llm_learning_rate": 4.4646975262036415e-06,
+      "loss": 0.7544,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 0.00044646975262036415,
+      "step": 2302
+    },
+    {
+      "epoch": 0.24,
+      "gumbel_tau": 0.581,
+      "image_masks_loss": 0.0131,
+      "llm_learning_rate": 4.464184399398661e-06,
+      "loss": 0.8331,
+      "output_text_mask_loss": 0.0231,
+      "predictor_learning_rate": 0.00044641843993986613,
+      "step": 2303
+    },
+    {
+      "epoch": 0.24,
+      "gumbel_tau": 0.5808,
+      "image_masks_loss": 0.0224,
+      "llm_learning_rate": 4.46367105629332e-06,
+      "loss": 0.8101,
+      "output_text_mask_loss": 0.0092,
+      "predictor_learning_rate": 0.00044636710562933197,
+      "step": 2304
+    },
+    {
+      "epoch": 0.24,
+      "gumbel_tau": 0.5807,
+      "image_masks_loss": 0.0224,
+      "llm_learning_rate": 4.4631574969441485e-06,
+      "loss": 0.7274,
+      "output_text_mask_loss": 0.0238,
+      "predictor_learning_rate": 0.0004463157496944148,
+      "step": 2305
+    },
+    {
+      "epoch": 0.24,
+      "gumbel_tau": 0.5805,
+      "image_masks_loss": 0.025,
+      "llm_learning_rate": 4.4626437214077e-06,
+      "loss": 0.7619,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 0.00044626437214077,
+      "step": 2306
+    },
+    {
+      "epoch": 0.24,
+      "gumbel_tau": 0.5804,
+      "image_masks_loss": 0.05,
+      "llm_learning_rate": 4.462129729740553e-06,
+      "loss": 0.8839,
+      "output_text_mask_loss": 0.0165,
+      "predictor_learning_rate": 0.00044621297297405525,
+      "step": 2307
+    },
+    {
+      "epoch": 0.24,
+      "gumbel_tau": 0.5803,
+      "image_masks_loss": 0.0171,
+      "llm_learning_rate": 4.4616155219993105e-06,
+      "loss": 0.8056,
+      "output_text_mask_loss": 0.1183,
+      "predictor_learning_rate": 0.000446161552199931,
+      "step": 2308
+    },
+    {
+      "epoch": 0.24,
+      "gumbel_tau": 0.5801,
+      "image_masks_loss": 0.0168,
+      "llm_learning_rate": 4.461101098240597e-06,
+      "loss": 0.8876,
+      "output_text_mask_loss": 0.0039,
+      "predictor_learning_rate": 0.0004461101098240596,
+      "step": 2309
+    },
+    {
+      "epoch": 0.24,
+      "gumbel_tau": 0.58,
+      "image_masks_loss": 0.0317,
+      "llm_learning_rate": 4.460586458521062e-06,
+      "loss": 0.814,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 0.0004460586458521062,
+      "step": 2310
+    },
+    {
+      "epoch": 0.24,
+      "gumbel_tau": 0.5799,
+      "image_masks_loss": 0.0171,
+      "llm_learning_rate": 4.46007160289738e-06,
+      "loss": 0.7813,
+      "output_text_mask_loss": 0.0086,
+      "predictor_learning_rate": 0.0004460071602897379,
+      "step": 2311
+    },
+    {
+      "epoch": 0.24,
+      "gumbel_tau": 0.5797,
+      "image_masks_loss": 0.032,
+      "llm_learning_rate": 4.459556531426246e-06,
+      "loss": 0.7023,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.0004459556531426246,
+      "step": 2312
+    },
+    {
+      "epoch": 0.24,
+      "gumbel_tau": 0.5796,
+      "image_masks_loss": 0.0201,
+      "llm_learning_rate": 4.459041244164383e-06,
+      "loss": 0.7922,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.00044590412441643826,
+      "step": 2313
+    },
+    {
+      "epoch": 0.24,
+      "gumbel_tau": 0.5794,
+      "image_masks_loss": 0.0191,
+      "llm_learning_rate": 4.458525741168534e-06,
+      "loss": 0.8642,
+      "output_text_mask_loss": 0.0139,
+      "predictor_learning_rate": 0.0004458525741168534,
+      "step": 2314
+    },
+    {
+      "epoch": 0.24,
+      "gumbel_tau": 0.5793,
+      "image_masks_loss": 0.0122,
+      "llm_learning_rate": 4.458010022495468e-06,
+      "loss": 0.87,
+      "output_text_mask_loss": 0.0216,
+      "predictor_learning_rate": 0.0004458010022495468,
+      "step": 2315
+    },
+    {
+      "epoch": 0.24,
+      "gumbel_tau": 0.5792,
+      "image_masks_loss": 0.0128,
+      "llm_learning_rate": 4.457494088201978e-06,
+      "loss": 0.8236,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.00044574940882019775,
+      "step": 2316
+    },
+    {
+      "epoch": 0.24,
+      "gumbel_tau": 0.579,
+      "image_masks_loss": 0.0156,
+      "llm_learning_rate": 4.456977938344877e-06,
+      "loss": 0.7835,
+      "output_text_mask_loss": 0.0221,
+      "predictor_learning_rate": 0.0004456977938344877,
+      "step": 2317
+    },
+    {
+      "epoch": 0.24,
+      "gumbel_tau": 0.5789,
+      "image_masks_loss": 0.0271,
+      "llm_learning_rate": 4.4564615729810085e-06,
+      "loss": 0.7015,
+      "output_text_mask_loss": 0.033,
+      "predictor_learning_rate": 0.0004456461572981008,
+      "step": 2318
+    },
+    {
+      "epoch": 0.24,
+      "gumbel_tau": 0.5788,
+      "image_masks_loss": 0.009,
+      "llm_learning_rate": 4.455944992167233e-06,
+      "loss": 0.6944,
+      "output_text_mask_loss": 0.0124,
+      "predictor_learning_rate": 0.00044559449921672325,
+      "step": 2319
+    },
+    {
+      "epoch": 0.24,
+      "gumbel_tau": 0.5786,
+      "image_masks_loss": 0.0267,
+      "llm_learning_rate": 4.4554281959604376e-06,
+      "loss": 0.7872,
+      "output_text_mask_loss": 0.0115,
+      "predictor_learning_rate": 0.00044554281959604375,
+      "step": 2320
+    },
+    {
+      "epoch": 0.24,
+      "gumbel_tau": 0.5785,
+      "image_masks_loss": 0.0248,
+      "llm_learning_rate": 4.454911184417534e-06,
+      "loss": 0.6653,
+      "output_text_mask_loss": 0.0111,
+      "predictor_learning_rate": 0.0004454911184417534,
+      "step": 2321
+    },
+    {
+      "epoch": 0.24,
+      "gumbel_tau": 0.5784,
+      "image_masks_loss": 0.0334,
+      "llm_learning_rate": 4.454393957595457e-06,
+      "loss": 0.7204,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.0004454393957595457,
+      "step": 2322
+    },
+    {
+      "epoch": 0.24,
+      "gumbel_tau": 0.5782,
+      "image_masks_loss": 0.0387,
+      "llm_learning_rate": 4.453876515551164e-06,
+      "loss": 0.8207,
+      "output_text_mask_loss": 0.0349,
+      "predictor_learning_rate": 0.00044538765155511633,
+      "step": 2323
+    },
+    {
+      "epoch": 0.24,
+      "gumbel_tau": 0.5781,
+      "image_masks_loss": 0.0133,
+      "llm_learning_rate": 4.453358858341636e-06,
+      "loss": 0.7656,
+      "output_text_mask_loss": 0.0542,
+      "predictor_learning_rate": 0.0004453358858341636,
+      "step": 2324
+    },
+    {
+      "epoch": 0.24,
+      "gumbel_tau": 0.5779,
+      "image_masks_loss": 0.0137,
+      "llm_learning_rate": 4.452840986023881e-06,
+      "loss": 0.781,
+      "output_text_mask_loss": 0.074,
+      "predictor_learning_rate": 0.00044528409860238803,
+      "step": 2325
+    },
+    {
+      "epoch": 0.24,
+      "gumbel_tau": 0.5778,
+      "image_masks_loss": 0.0271,
+      "llm_learning_rate": 4.452322898654926e-06,
+      "loss": 0.7183,
+      "output_text_mask_loss": 0.0141,
+      "predictor_learning_rate": 0.00044523228986549255,
+      "step": 2326
+    },
+    {
+      "epoch": 0.24,
+      "gumbel_tau": 0.5777,
+      "image_masks_loss": 0.0108,
+      "llm_learning_rate": 4.4518045962918245e-06,
+      "loss": 0.7594,
+      "output_text_mask_loss": 0.0168,
+      "predictor_learning_rate": 0.0004451804596291824,
+      "step": 2327
+    },
+    {
+      "epoch": 0.24,
+      "gumbel_tau": 0.5775,
+      "image_masks_loss": 0.0221,
+      "llm_learning_rate": 4.451286078991654e-06,
+      "loss": 0.7776,
+      "output_text_mask_loss": 0.0345,
+      "predictor_learning_rate": 0.0004451286078991654,
+      "step": 2328
+    },
+    {
+      "epoch": 0.24,
+      "gumbel_tau": 0.5774,
+      "image_masks_loss": 0.0079,
+      "llm_learning_rate": 4.450767346811514e-06,
+      "loss": 0.7348,
+      "output_text_mask_loss": 0.0462,
+      "predictor_learning_rate": 0.0004450767346811514,
+      "step": 2329
+    },
+    {
+      "epoch": 0.24,
+      "gumbel_tau": 0.5773,
+      "image_masks_loss": 0.022,
+      "llm_learning_rate": 4.450248399808529e-06,
+      "loss": 0.7316,
+      "output_text_mask_loss": 0.0023,
+      "predictor_learning_rate": 0.00044502483998085286,
+      "step": 2330
+    },
+    {
+      "epoch": 0.24,
+      "gumbel_tau": 0.5771,
+      "image_masks_loss": 0.0091,
+      "llm_learning_rate": 4.449729238039846e-06,
+      "loss": 0.806,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 0.0004449729238039845,
+      "step": 2331
+    },
+    {
+      "epoch": 0.24,
+      "gumbel_tau": 0.577,
+      "image_masks_loss": 0.0194,
+      "llm_learning_rate": 4.4492098615626355e-06,
+      "loss": 0.7749,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 0.00044492098615626354,
+      "step": 2332
+    },
+    {
+      "epoch": 0.24,
+      "gumbel_tau": 0.5769,
+      "image_masks_loss": 0.028,
+      "llm_learning_rate": 4.448690270434094e-06,
+      "loss": 0.76,
+      "output_text_mask_loss": 0.0004,
+      "predictor_learning_rate": 0.00044486902704340937,
+      "step": 2333
+    },
+    {
+      "epoch": 0.24,
+      "gumbel_tau": 0.5767,
+      "image_masks_loss": 0.0038,
+      "llm_learning_rate": 4.4481704647114395e-06,
+      "loss": 0.7709,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.00044481704647114396,
+      "step": 2334
+    },
+    {
+      "epoch": 0.24,
+      "gumbel_tau": 0.5766,
+      "image_masks_loss": 0.0242,
+      "llm_learning_rate": 4.447650444451914e-06,
+      "loss": 0.8028,
+      "output_text_mask_loss": 0.005,
+      "predictor_learning_rate": 0.0004447650444451914,
+      "step": 2335
+    },
+    {
+      "epoch": 0.24,
+      "gumbel_tau": 0.5764,
+      "image_masks_loss": 0.0286,
+      "llm_learning_rate": 4.447130209712783e-06,
+      "loss": 0.7508,
+      "output_text_mask_loss": 0.0168,
+      "predictor_learning_rate": 0.0004447130209712783,
+      "step": 2336
+    },
+    {
+      "epoch": 0.24,
+      "gumbel_tau": 0.5763,
+      "image_masks_loss": 0.0383,
+      "llm_learning_rate": 4.446609760551337e-06,
+      "loss": 0.7627,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.0004446609760551337,
+      "step": 2337
+    },
+    {
+      "epoch": 0.24,
+      "gumbel_tau": 0.5762,
+      "image_masks_loss": 0.0298,
+      "llm_learning_rate": 4.446089097024889e-06,
+      "loss": 0.8252,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 0.0004446089097024888,
+      "step": 2338
+    },
+    {
+      "epoch": 0.24,
+      "gumbel_tau": 0.576,
+      "image_masks_loss": 0.0269,
+      "llm_learning_rate": 4.4455682191907724e-06,
+      "loss": 0.8138,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.0004445568219190772,
+      "step": 2339
+    },
+    {
+      "epoch": 0.24,
+      "gumbel_tau": 0.5759,
+      "image_masks_loss": 0.0045,
+      "llm_learning_rate": 4.445047127106351e-06,
+      "loss": 0.6872,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004445047127106351,
+      "step": 2340
+    },
+    {
+      "epoch": 0.24,
+      "gumbel_tau": 0.5758,
+      "image_masks_loss": 0.0146,
+      "llm_learning_rate": 4.4445258208290085e-06,
+      "loss": 0.6667,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.00044445258208290083,
+      "step": 2341
+    },
+    {
+      "epoch": 0.24,
+      "gumbel_tau": 0.5756,
+      "image_masks_loss": 0.0107,
+      "llm_learning_rate": 4.444004300416151e-06,
+      "loss": 0.7613,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.000444400430041615,
+      "step": 2342
+    },
+    {
+      "epoch": 0.24,
+      "gumbel_tau": 0.5755,
+      "image_masks_loss": 0.0105,
+      "llm_learning_rate": 4.443482565925209e-06,
+      "loss": 0.7842,
+      "output_text_mask_loss": 0.1236,
+      "predictor_learning_rate": 0.0004443482565925209,
+      "step": 2343
+    },
+    {
+      "epoch": 0.24,
+      "gumbel_tau": 0.5754,
+      "image_masks_loss": 0.0336,
+      "llm_learning_rate": 4.442960617413639e-06,
+      "loss": 0.8193,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004442960617413638,
+      "step": 2344
+    },
+    {
+      "epoch": 0.24,
+      "gumbel_tau": 0.5752,
+      "image_masks_loss": 0.0443,
+      "llm_learning_rate": 4.4424384549389165e-06,
+      "loss": 0.8335,
+      "output_text_mask_loss": 0.0214,
+      "predictor_learning_rate": 0.0004442438454938916,
+      "step": 2345
+    },
+    {
+      "epoch": 0.24,
+      "gumbel_tau": 0.5751,
+      "image_masks_loss": 0.0175,
+      "llm_learning_rate": 4.441916078558546e-06,
+      "loss": 0.7303,
+      "output_text_mask_loss": 0.0202,
+      "predictor_learning_rate": 0.00044419160785585455,
+      "step": 2346
+    },
+    {
+      "epoch": 0.24,
+      "gumbel_tau": 0.575,
+      "image_masks_loss": 0.025,
+      "llm_learning_rate": 4.441393488330051e-06,
+      "loss": 0.7803,
+      "output_text_mask_loss": 0.0042,
+      "predictor_learning_rate": 0.0004441393488330051,
+      "step": 2347
+    },
+    {
+      "epoch": 0.24,
+      "gumbel_tau": 0.5748,
+      "image_masks_loss": 0.0083,
+      "llm_learning_rate": 4.440870684310981e-06,
+      "loss": 0.7648,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.00044408706843109805,
+      "step": 2348
+    },
+    {
+      "epoch": 0.24,
+      "gumbel_tau": 0.5747,
+      "image_masks_loss": 0.0216,
+      "llm_learning_rate": 4.440347666558908e-06,
+      "loss": 0.781,
+      "output_text_mask_loss": 0.0041,
+      "predictor_learning_rate": 0.00044403476665589084,
+      "step": 2349
+    },
+    {
+      "epoch": 0.24,
+      "gumbel_tau": 0.5745,
+      "image_masks_loss": 0.0111,
+      "llm_learning_rate": 4.439824435131428e-06,
+      "loss": 0.7572,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.00044398244351314276,
+      "step": 2350
+    },
+    {
+      "epoch": 0.24,
+      "gumbel_tau": 0.5744,
+      "image_masks_loss": 0.0694,
+      "llm_learning_rate": 4.439300990086161e-06,
+      "loss": 0.8244,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.000443930099008616,
+      "step": 2351
+    },
+    {
+      "epoch": 0.24,
+      "gumbel_tau": 0.5743,
+      "image_masks_loss": 0.0117,
+      "llm_learning_rate": 4.438777331480748e-06,
+      "loss": 0.7352,
+      "output_text_mask_loss": 0.0086,
+      "predictor_learning_rate": 0.00044387773314807477,
+      "step": 2352
+    },
+    {
+      "epoch": 0.24,
+      "gumbel_tau": 0.5741,
+      "image_masks_loss": 0.05,
+      "llm_learning_rate": 4.438253459372857e-06,
+      "loss": 0.7704,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004438253459372857,
+      "step": 2353
+    },
+    {
+      "epoch": 0.24,
+      "gumbel_tau": 0.574,
+      "image_masks_loss": 0.0288,
+      "llm_learning_rate": 4.437729373820179e-06,
+      "loss": 0.7938,
+      "output_text_mask_loss": 0.0786,
+      "predictor_learning_rate": 0.00044377293738201787,
+      "step": 2354
+    },
+    {
+      "epoch": 0.24,
+      "gumbel_tau": 0.5739,
+      "image_masks_loss": 0.0366,
+      "llm_learning_rate": 4.4372050748804256e-06,
+      "loss": 0.7529,
+      "output_text_mask_loss": 0.0259,
+      "predictor_learning_rate": 0.00044372050748804257,
+      "step": 2355
+    },
+    {
+      "epoch": 0.24,
+      "gumbel_tau": 0.5737,
+      "image_masks_loss": 0.019,
+      "llm_learning_rate": 4.436680562611335e-06,
+      "loss": 0.7692,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.00044366805626113343,
+      "step": 2356
+    },
+    {
+      "epoch": 0.24,
+      "gumbel_tau": 0.5736,
+      "image_masks_loss": 0.0132,
+      "llm_learning_rate": 4.436155837070666e-06,
+      "loss": 0.7707,
+      "output_text_mask_loss": 0.0004,
+      "predictor_learning_rate": 0.0004436155837070666,
+      "step": 2357
+    },
+    {
+      "epoch": 0.24,
+      "gumbel_tau": 0.5735,
+      "image_masks_loss": 0.0185,
+      "llm_learning_rate": 4.435630898316205e-06,
+      "loss": 0.8032,
+      "output_text_mask_loss": 0.0458,
+      "predictor_learning_rate": 0.0004435630898316204,
+      "step": 2358
+    },
+    {
+      "epoch": 0.24,
+      "gumbel_tau": 0.5733,
+      "image_masks_loss": 0.016,
+      "llm_learning_rate": 4.435105746405756e-06,
+      "loss": 0.7057,
+      "output_text_mask_loss": 0.0235,
+      "predictor_learning_rate": 0.0004435105746405756,
+      "step": 2359
+    },
+    {
+      "epoch": 0.24,
+      "gumbel_tau": 0.5732,
+      "image_masks_loss": 0.0212,
+      "llm_learning_rate": 4.434580381397153e-06,
+      "loss": 0.7464,
+      "output_text_mask_loss": 0.0378,
+      "predictor_learning_rate": 0.0004434580381397153,
+      "step": 2360
+    },
+    {
+      "epoch": 0.24,
+      "gumbel_tau": 0.5731,
+      "image_masks_loss": 0.0299,
+      "llm_learning_rate": 4.434054803348249e-06,
+      "loss": 0.7383,
+      "output_text_mask_loss": 0.0107,
+      "predictor_learning_rate": 0.00044340548033482493,
+      "step": 2361
+    },
+    {
+      "epoch": 0.24,
+      "gumbel_tau": 0.5729,
+      "image_masks_loss": 0.0178,
+      "llm_learning_rate": 4.433529012316923e-06,
+      "loss": 0.6545,
+      "output_text_mask_loss": 0.0177,
+      "predictor_learning_rate": 0.00044335290123169224,
+      "step": 2362
+    },
+    {
+      "epoch": 0.24,
+      "gumbel_tau": 0.5728,
+      "image_masks_loss": 0.0355,
+      "llm_learning_rate": 4.4330030083610745e-06,
+      "loss": 0.794,
+      "output_text_mask_loss": 0.0664,
+      "predictor_learning_rate": 0.0004433003008361075,
+      "step": 2363
+    },
+    {
+      "epoch": 0.24,
+      "gumbel_tau": 0.5727,
+      "image_masks_loss": 0.0299,
+      "llm_learning_rate": 4.43247679153863e-06,
+      "loss": 0.8025,
+      "output_text_mask_loss": 0.0029,
+      "predictor_learning_rate": 0.00044324767915386294,
+      "step": 2364
+    },
+    {
+      "epoch": 0.24,
+      "gumbel_tau": 0.5725,
+      "image_masks_loss": 0.022,
+      "llm_learning_rate": 4.431950361907536e-06,
+      "loss": 0.8431,
+      "output_text_mask_loss": 0.071,
+      "predictor_learning_rate": 0.0004431950361907536,
+      "step": 2365
+    },
+    {
+      "epoch": 0.24,
+      "gumbel_tau": 0.5724,
+      "image_masks_loss": 0.0252,
+      "llm_learning_rate": 4.431423719525765e-06,
+      "loss": 0.7642,
+      "output_text_mask_loss": 0.0137,
+      "predictor_learning_rate": 0.0004431423719525765,
+      "step": 2366
+    },
+    {
+      "epoch": 0.24,
+      "gumbel_tau": 0.5722,
+      "image_masks_loss": 0.0095,
+      "llm_learning_rate": 4.4308968644513126e-06,
+      "loss": 0.7232,
+      "output_text_mask_loss": 0.0141,
+      "predictor_learning_rate": 0.0004430896864451312,
+      "step": 2367
+    },
+    {
+      "epoch": 0.24,
+      "gumbel_tau": 0.5721,
+      "image_masks_loss": 0.0181,
+      "llm_learning_rate": 4.430369796742196e-06,
+      "loss": 0.8372,
+      "output_text_mask_loss": 0.0439,
+      "predictor_learning_rate": 0.0004430369796742195,
+      "step": 2368
+    },
+    {
+      "epoch": 0.24,
+      "gumbel_tau": 0.572,
+      "image_masks_loss": 0.0173,
+      "llm_learning_rate": 4.4298425164564576e-06,
+      "loss": 0.8106,
+      "output_text_mask_loss": 0.0429,
+      "predictor_learning_rate": 0.0004429842516456457,
+      "step": 2369
+    },
+    {
+      "epoch": 0.24,
+      "gumbel_tau": 0.5718,
+      "image_masks_loss": 0.0228,
+      "llm_learning_rate": 4.429315023652163e-06,
+      "loss": 0.7376,
+      "output_text_mask_loss": 0.0355,
+      "predictor_learning_rate": 0.00044293150236521624,
+      "step": 2370
+    },
+    {
+      "epoch": 0.24,
+      "gumbel_tau": 0.5717,
+      "image_masks_loss": 0.0143,
+      "llm_learning_rate": 4.4287873183874e-06,
+      "loss": 0.8801,
+      "output_text_mask_loss": 0.0271,
+      "predictor_learning_rate": 0.00044287873183873996,
+      "step": 2371
+    },
+    {
+      "epoch": 0.24,
+      "gumbel_tau": 0.5716,
+      "image_masks_loss": 0.0299,
+      "llm_learning_rate": 4.428259400720282e-06,
+      "loss": 0.8127,
+      "output_text_mask_loss": 0.034,
+      "predictor_learning_rate": 0.00044282594007202813,
+      "step": 2372
+    },
+    {
+      "epoch": 0.24,
+      "gumbel_tau": 0.5714,
+      "image_masks_loss": 0.046,
+      "llm_learning_rate": 4.427731270708943e-06,
+      "loss": 0.7574,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 0.00044277312707089426,
+      "step": 2373
+    },
+    {
+      "epoch": 0.24,
+      "gumbel_tau": 0.5713,
+      "image_masks_loss": 0.0186,
+      "llm_learning_rate": 4.427202928411544e-06,
+      "loss": 0.7978,
+      "output_text_mask_loss": 0.0475,
+      "predictor_learning_rate": 0.00044272029284115434,
+      "step": 2374
+    },
+    {
+      "epoch": 0.24,
+      "gumbel_tau": 0.5712,
+      "image_masks_loss": 0.0664,
+      "llm_learning_rate": 4.4266743738862646e-06,
+      "loss": 0.9127,
+      "output_text_mask_loss": 0.0025,
+      "predictor_learning_rate": 0.0004426674373886264,
+      "step": 2375
+    },
+    {
+      "epoch": 0.24,
+      "gumbel_tau": 0.571,
+      "image_masks_loss": 0.0656,
+      "llm_learning_rate": 4.4261456071913115e-06,
+      "loss": 0.7533,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00044261456071913115,
+      "step": 2376
+    },
+    {
+      "epoch": 0.24,
+      "gumbel_tau": 0.5709,
+      "image_masks_loss": 0.0267,
+      "llm_learning_rate": 4.425616628384915e-06,
+      "loss": 0.7992,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.0004425616628384914,
+      "step": 2377
+    },
+    {
+      "epoch": 0.24,
+      "gumbel_tau": 0.5708,
+      "image_masks_loss": 0.0874,
+      "llm_learning_rate": 4.425087437525325e-06,
+      "loss": 0.7961,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.0004425087437525325,
+      "step": 2378
+    },
+    {
+      "epoch": 0.24,
+      "gumbel_tau": 0.5706,
+      "image_masks_loss": 0.0462,
+      "llm_learning_rate": 4.424558034670819e-06,
+      "loss": 0.8531,
+      "output_text_mask_loss": 0.0025,
+      "predictor_learning_rate": 0.00044245580346708184,
+      "step": 2379
+    },
+    {
+      "epoch": 0.24,
+      "gumbel_tau": 0.5705,
+      "image_masks_loss": 0.0162,
+      "llm_learning_rate": 4.4240284198796956e-06,
+      "loss": 0.7231,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.0004424028419879695,
+      "step": 2380
+    },
+    {
+      "epoch": 0.24,
+      "gumbel_tau": 0.5704,
+      "image_masks_loss": 0.0326,
+      "llm_learning_rate": 4.423498593210277e-06,
+      "loss": 0.8294,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.00044234985932102767,
+      "step": 2381
+    },
+    {
+      "epoch": 0.24,
+      "gumbel_tau": 0.5702,
+      "image_masks_loss": 0.0132,
+      "llm_learning_rate": 4.422968554720908e-06,
+      "loss": 0.7473,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004422968554720908,
+      "step": 2382
+    },
+    {
+      "epoch": 0.24,
+      "gumbel_tau": 0.5701,
+      "image_masks_loss": 0.0408,
+      "llm_learning_rate": 4.4224383044699595e-06,
+      "loss": 0.7903,
+      "output_text_mask_loss": 0.0732,
+      "predictor_learning_rate": 0.00044224383044699593,
+      "step": 2383
+    },
+    {
+      "epoch": 0.24,
+      "gumbel_tau": 0.57,
+      "image_masks_loss": 0.0319,
+      "llm_learning_rate": 4.421907842515822e-06,
+      "loss": 0.783,
+      "output_text_mask_loss": 0.017,
+      "predictor_learning_rate": 0.0004421907842515822,
+      "step": 2384
+    },
+    {
+      "epoch": 0.24,
+      "gumbel_tau": 0.5698,
+      "image_masks_loss": 0.0127,
+      "llm_learning_rate": 4.421377168916913e-06,
+      "loss": 0.7992,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.0004421377168916913,
+      "step": 2385
+    },
+    {
+      "epoch": 0.24,
+      "gumbel_tau": 0.5697,
+      "image_masks_loss": 0.021,
+      "llm_learning_rate": 4.42084628373167e-06,
+      "loss": 0.7865,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 0.000442084628373167,
+      "step": 2386
+    },
+    {
+      "epoch": 0.24,
+      "gumbel_tau": 0.5696,
+      "image_masks_loss": 0.0108,
+      "llm_learning_rate": 4.420315187018556e-06,
+      "loss": 0.8442,
+      "output_text_mask_loss": 0.0126,
+      "predictor_learning_rate": 0.00044203151870185557,
+      "step": 2387
+    },
+    {
+      "epoch": 0.24,
+      "gumbel_tau": 0.5694,
+      "image_masks_loss": 0.0286,
+      "llm_learning_rate": 4.419783878836056e-06,
+      "loss": 0.8156,
+      "output_text_mask_loss": 0.0071,
+      "predictor_learning_rate": 0.00044197838788360554,
+      "step": 2388
+    },
+    {
+      "epoch": 0.24,
+      "gumbel_tau": 0.5693,
+      "image_masks_loss": 0.0183,
+      "llm_learning_rate": 4.41925235924268e-06,
+      "loss": 0.7378,
+      "output_text_mask_loss": 0.0305,
+      "predictor_learning_rate": 0.00044192523592426794,
+      "step": 2389
+    },
+    {
+      "epoch": 0.24,
+      "gumbel_tau": 0.5691,
+      "image_masks_loss": 0.0277,
+      "llm_learning_rate": 4.4187206282969575e-06,
+      "loss": 0.8087,
+      "output_text_mask_loss": 0.0368,
+      "predictor_learning_rate": 0.0004418720628296957,
+      "step": 2390
+    },
+    {
+      "epoch": 0.24,
+      "gumbel_tau": 0.569,
+      "image_masks_loss": 0.0156,
+      "llm_learning_rate": 4.418188686057446e-06,
+      "loss": 0.8109,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.0004418188686057446,
+      "step": 2391
+    },
+    {
+      "epoch": 0.25,
+      "gumbel_tau": 0.5689,
+      "image_masks_loss": 0.0595,
+      "llm_learning_rate": 4.417656532582725e-06,
+      "loss": 0.7603,
+      "output_text_mask_loss": 0.0275,
+      "predictor_learning_rate": 0.00044176565325827245,
+      "step": 2392
+    },
+    {
+      "epoch": 0.25,
+      "gumbel_tau": 0.5687,
+      "image_masks_loss": 0.0267,
+      "llm_learning_rate": 4.417124167931395e-06,
+      "loss": 0.7327,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004417124167931395,
+      "step": 2393
+    },
+    {
+      "epoch": 0.25,
+      "gumbel_tau": 0.5686,
+      "image_masks_loss": 0.0157,
+      "llm_learning_rate": 4.416591592162081e-06,
+      "loss": 0.7156,
+      "output_text_mask_loss": 0.0058,
+      "predictor_learning_rate": 0.00044165915921620804,
+      "step": 2394
+    },
+    {
+      "epoch": 0.25,
+      "gumbel_tau": 0.5685,
+      "image_masks_loss": 0.0317,
+      "llm_learning_rate": 4.416058805333431e-06,
+      "loss": 0.7539,
+      "output_text_mask_loss": 0.0004,
+      "predictor_learning_rate": 0.00044160588053334304,
+      "step": 2395
+    },
+    {
+      "epoch": 0.25,
+      "gumbel_tau": 0.5683,
+      "image_masks_loss": 0.0209,
+      "llm_learning_rate": 4.415525807504119e-06,
+      "loss": 0.7621,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.00044155258075041185,
+      "step": 2396
+    },
+    {
+      "epoch": 0.25,
+      "gumbel_tau": 0.5682,
+      "image_masks_loss": 0.012,
+      "llm_learning_rate": 4.4149925987328365e-06,
+      "loss": 0.7755,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00044149925987328363,
+      "step": 2397
+    },
+    {
+      "epoch": 0.25,
+      "gumbel_tau": 0.5681,
+      "image_masks_loss": 0.0092,
+      "llm_learning_rate": 4.414459179078305e-06,
+      "loss": 0.7544,
+      "output_text_mask_loss": 0.0113,
+      "predictor_learning_rate": 0.00044144591790783047,
+      "step": 2398
+    },
+    {
+      "epoch": 0.25,
+      "gumbel_tau": 0.5679,
+      "image_masks_loss": 0.0418,
+      "llm_learning_rate": 4.413925548599264e-06,
+      "loss": 0.7824,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 0.0004413925548599264,
+      "step": 2399
+    },
+    {
+      "epoch": 0.25,
+      "gumbel_tau": 0.5678,
+      "image_masks_loss": 0.029,
+      "llm_learning_rate": 4.413391707354479e-06,
+      "loss": 0.8345,
+      "output_text_mask_loss": 0.01,
+      "predictor_learning_rate": 0.0004413391707354478,
+      "step": 2400
+    },
+    {
+      "epoch": 0.25,
+      "gumbel_tau": 0.5677,
+      "image_masks_loss": 0.0053,
+      "llm_learning_rate": 4.412857655402737e-06,
+      "loss": 0.6965,
+      "output_text_mask_loss": 0.0107,
+      "predictor_learning_rate": 0.00044128576554027367,
+      "step": 2401
+    },
+    {
+      "epoch": 0.25,
+      "gumbel_tau": 0.5675,
+      "image_masks_loss": 0.0549,
+      "llm_learning_rate": 4.4123233928028486e-06,
+      "loss": 0.771,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004412323392802848,
+      "step": 2402
+    },
+    {
+      "epoch": 0.25,
+      "gumbel_tau": 0.5674,
+      "image_masks_loss": 0.0183,
+      "llm_learning_rate": 4.411788919613649e-06,
+      "loss": 0.766,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.0004411788919613649,
+      "step": 2403
+    },
+    {
+      "epoch": 0.25,
+      "gumbel_tau": 0.5673,
+      "image_masks_loss": 0.0111,
+      "llm_learning_rate": 4.4112542358939955e-06,
+      "loss": 0.8695,
+      "output_text_mask_loss": 0.0511,
+      "predictor_learning_rate": 0.00044112542358939956,
+      "step": 2404
+    },
+    {
+      "epoch": 0.25,
+      "gumbel_tau": 0.5671,
+      "image_masks_loss": 0.0207,
+      "llm_learning_rate": 4.4107193417027686e-06,
+      "loss": 0.8571,
+      "output_text_mask_loss": 0.0565,
+      "predictor_learning_rate": 0.0004410719341702768,
+      "step": 2405
+    },
+    {
+      "epoch": 0.25,
+      "gumbel_tau": 0.567,
+      "image_masks_loss": 0.0128,
+      "llm_learning_rate": 4.410184237098872e-06,
+      "loss": 0.8914,
+      "output_text_mask_loss": 0.0832,
+      "predictor_learning_rate": 0.0004410184237098871,
+      "step": 2406
+    },
+    {
+      "epoch": 0.25,
+      "gumbel_tau": 0.5669,
+      "image_masks_loss": 0.0122,
+      "llm_learning_rate": 4.409648922141231e-06,
+      "loss": 0.849,
+      "output_text_mask_loss": 0.0656,
+      "predictor_learning_rate": 0.00044096489221412313,
+      "step": 2407
+    },
+    {
+      "epoch": 0.25,
+      "gumbel_tau": 0.5667,
+      "image_masks_loss": 0.0151,
+      "llm_learning_rate": 4.409113396888798e-06,
+      "loss": 0.7603,
+      "output_text_mask_loss": 0.0452,
+      "predictor_learning_rate": 0.00044091133968887983,
+      "step": 2408
+    },
+    {
+      "epoch": 0.25,
+      "gumbel_tau": 0.5666,
+      "image_masks_loss": 0.0095,
+      "llm_learning_rate": 4.408577661400546e-06,
+      "loss": 0.8168,
+      "output_text_mask_loss": 0.0462,
+      "predictor_learning_rate": 0.00044085776614005457,
+      "step": 2409
+    },
+    {
+      "epoch": 0.25,
+      "gumbel_tau": 0.5665,
+      "image_masks_loss": 0.0387,
+      "llm_learning_rate": 4.40804171573547e-06,
+      "loss": 0.779,
+      "output_text_mask_loss": 0.0759,
+      "predictor_learning_rate": 0.00044080417157354697,
+      "step": 2410
+    },
+    {
+      "epoch": 0.25,
+      "gumbel_tau": 0.5663,
+      "image_masks_loss": 0.0163,
+      "llm_learning_rate": 4.40750555995259e-06,
+      "loss": 0.7869,
+      "output_text_mask_loss": 0.005,
+      "predictor_learning_rate": 0.00044075055599525906,
+      "step": 2411
+    },
+    {
+      "epoch": 0.25,
+      "gumbel_tau": 0.5662,
+      "image_masks_loss": 0.037,
+      "llm_learning_rate": 4.40696919411095e-06,
+      "loss": 0.7429,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00044069691941109497,
+      "step": 2412
+    },
+    {
+      "epoch": 0.25,
+      "gumbel_tau": 0.5661,
+      "image_masks_loss": 0.0137,
+      "llm_learning_rate": 4.406432618269613e-06,
+      "loss": 0.7042,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.0004406432618269614,
+      "step": 2413
+    },
+    {
+      "epoch": 0.25,
+      "gumbel_tau": 0.5659,
+      "image_masks_loss": 0.0244,
+      "llm_learning_rate": 4.405895832487671e-06,
+      "loss": 0.777,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004405895832487671,
+      "step": 2414
+    },
+    {
+      "epoch": 0.25,
+      "gumbel_tau": 0.5658,
+      "image_masks_loss": 0.0349,
+      "llm_learning_rate": 4.4053588368242354e-06,
+      "loss": 0.7691,
+      "output_text_mask_loss": 0.0698,
+      "predictor_learning_rate": 0.0004405358836824235,
+      "step": 2415
+    },
+    {
+      "epoch": 0.25,
+      "gumbel_tau": 0.5657,
+      "image_masks_loss": 0.0263,
+      "llm_learning_rate": 4.40482163133844e-06,
+      "loss": 0.8465,
+      "output_text_mask_loss": 0.0854,
+      "predictor_learning_rate": 0.0004404821631338439,
+      "step": 2416
+    },
+    {
+      "epoch": 0.25,
+      "gumbel_tau": 0.5655,
+      "image_masks_loss": 0.0368,
+      "llm_learning_rate": 4.404284216089443e-06,
+      "loss": 0.7895,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.0004404284216089443,
+      "step": 2417
+    },
+    {
+      "epoch": 0.25,
+      "gumbel_tau": 0.5654,
+      "image_masks_loss": 0.0226,
+      "llm_learning_rate": 4.403746591136427e-06,
+      "loss": 0.7183,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.00044037465911364273,
+      "step": 2418
+    },
+    {
+      "epoch": 0.25,
+      "gumbel_tau": 0.5653,
+      "image_masks_loss": 0.0105,
+      "llm_learning_rate": 4.4032087565385965e-06,
+      "loss": 0.8042,
+      "output_text_mask_loss": 0.0185,
+      "predictor_learning_rate": 0.0004403208756538596,
+      "step": 2419
+    },
+    {
+      "epoch": 0.25,
+      "gumbel_tau": 0.5651,
+      "image_masks_loss": 0.0137,
+      "llm_learning_rate": 4.402670712355177e-06,
+      "loss": 0.8308,
+      "output_text_mask_loss": 0.016,
+      "predictor_learning_rate": 0.0004402670712355177,
+      "step": 2420
+    },
+    {
+      "epoch": 0.25,
+      "gumbel_tau": 0.565,
+      "image_masks_loss": 0.0271,
+      "llm_learning_rate": 4.402132458645422e-06,
+      "loss": 0.7667,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00044021324586454217,
+      "step": 2421
+    },
+    {
+      "epoch": 0.25,
+      "gumbel_tau": 0.5649,
+      "image_masks_loss": 0.0519,
+      "llm_learning_rate": 4.401593995468603e-06,
+      "loss": 0.756,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.0004401593995468603,
+      "step": 2422
+    },
+    {
+      "epoch": 0.25,
+      "gumbel_tau": 0.5647,
+      "image_masks_loss": 0.0114,
+      "llm_learning_rate": 4.401055322884018e-06,
+      "loss": 0.8359,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.00044010553228840177,
+      "step": 2423
+    },
+    {
+      "epoch": 0.25,
+      "gumbel_tau": 0.5646,
+      "image_masks_loss": 0.0027,
+      "llm_learning_rate": 4.400516440950987e-06,
+      "loss": 0.7789,
+      "output_text_mask_loss": 0.0061,
+      "predictor_learning_rate": 0.00044005164409509866,
+      "step": 2424
+    },
+    {
+      "epoch": 0.25,
+      "gumbel_tau": 0.5645,
+      "image_masks_loss": 0.0324,
+      "llm_learning_rate": 4.399977349728851e-06,
+      "loss": 0.7879,
+      "output_text_mask_loss": 0.0275,
+      "predictor_learning_rate": 0.0004399977349728851,
+      "step": 2425
+    },
+    {
+      "epoch": 0.25,
+      "gumbel_tau": 0.5643,
+      "image_masks_loss": 0.019,
+      "llm_learning_rate": 4.399438049276978e-06,
+      "loss": 0.7925,
+      "output_text_mask_loss": 0.0207,
+      "predictor_learning_rate": 0.00043994380492769783,
+      "step": 2426
+    },
+    {
+      "epoch": 0.25,
+      "gumbel_tau": 0.5642,
+      "image_masks_loss": 0.0132,
+      "llm_learning_rate": 4.398898539654756e-06,
+      "loss": 0.8135,
+      "output_text_mask_loss": 0.0058,
+      "predictor_learning_rate": 0.0004398898539654756,
+      "step": 2427
+    },
+    {
+      "epoch": 0.25,
+      "gumbel_tau": 0.5641,
+      "image_masks_loss": 0.0185,
+      "llm_learning_rate": 4.3983588209215975e-06,
+      "loss": 0.7174,
+      "output_text_mask_loss": 0.0149,
+      "predictor_learning_rate": 0.00043983588209215975,
+      "step": 2428
+    },
+    {
+      "epoch": 0.25,
+      "gumbel_tau": 0.5639,
+      "image_masks_loss": 0.0224,
+      "llm_learning_rate": 4.397818893136937e-06,
+      "loss": 0.7443,
+      "output_text_mask_loss": 0.0618,
+      "predictor_learning_rate": 0.0004397818893136937,
+      "step": 2429
+    },
+    {
+      "epoch": 0.25,
+      "gumbel_tau": 0.5638,
+      "image_masks_loss": 0.0301,
+      "llm_learning_rate": 4.397278756360233e-06,
+      "loss": 0.7897,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.0004397278756360232,
+      "step": 2430
+    },
+    {
+      "epoch": 0.25,
+      "gumbel_tau": 0.5637,
+      "image_masks_loss": 0.0184,
+      "llm_learning_rate": 4.396738410650966e-06,
+      "loss": 0.8815,
+      "output_text_mask_loss": 0.0399,
+      "predictor_learning_rate": 0.00043967384106509655,
+      "step": 2431
+    },
+    {
+      "epoch": 0.25,
+      "gumbel_tau": 0.5635,
+      "image_masks_loss": 0.0366,
+      "llm_learning_rate": 4.39619785606864e-06,
+      "loss": 0.8608,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.000439619785606864,
+      "step": 2432
+    },
+    {
+      "epoch": 0.25,
+      "gumbel_tau": 0.5634,
+      "image_masks_loss": 0.0103,
+      "llm_learning_rate": 4.395657092672783e-06,
+      "loss": 0.7839,
+      "output_text_mask_loss": 0.0467,
+      "predictor_learning_rate": 0.0004395657092672783,
+      "step": 2433
+    },
+    {
+      "epoch": 0.25,
+      "gumbel_tau": 0.5633,
+      "image_masks_loss": 0.0151,
+      "llm_learning_rate": 4.395116120522944e-06,
+      "loss": 0.7615,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.0004395116120522944,
+      "step": 2434
+    },
+    {
+      "epoch": 0.25,
+      "gumbel_tau": 0.5631,
+      "image_masks_loss": 0.0089,
+      "llm_learning_rate": 4.394574939678697e-06,
+      "loss": 0.7521,
+      "output_text_mask_loss": 0.0015,
+      "predictor_learning_rate": 0.00043945749396786973,
+      "step": 2435
+    },
+    {
+      "epoch": 0.25,
+      "gumbel_tau": 0.563,
+      "image_masks_loss": 0.0156,
+      "llm_learning_rate": 4.394033550199639e-06,
+      "loss": 0.7179,
+      "output_text_mask_loss": 0.0117,
+      "predictor_learning_rate": 0.00043940335501996383,
+      "step": 2436
+    },
+    {
+      "epoch": 0.25,
+      "gumbel_tau": 0.5629,
+      "image_masks_loss": 0.0193,
+      "llm_learning_rate": 4.3934919521453856e-06,
+      "loss": 0.7453,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00043934919521453856,
+      "step": 2437
+    },
+    {
+      "epoch": 0.25,
+      "gumbel_tau": 0.5627,
+      "image_masks_loss": 0.0347,
+      "llm_learning_rate": 4.392950145575582e-06,
+      "loss": 0.7432,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.00043929501455755815,
+      "step": 2438
+    },
+    {
+      "epoch": 0.25,
+      "gumbel_tau": 0.5626,
+      "image_masks_loss": 0.0156,
+      "llm_learning_rate": 4.392408130549891e-06,
+      "loss": 0.8733,
+      "output_text_mask_loss": 0.0012,
+      "predictor_learning_rate": 0.0004392408130549891,
+      "step": 2439
+    },
+    {
+      "epoch": 0.25,
+      "gumbel_tau": 0.5625,
+      "image_masks_loss": 0.0125,
+      "llm_learning_rate": 4.391865907128002e-06,
+      "loss": 0.6907,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.00043918659071280023,
+      "step": 2440
+    },
+    {
+      "epoch": 0.25,
+      "gumbel_tau": 0.5623,
+      "image_masks_loss": 0.0109,
+      "llm_learning_rate": 4.391323475369626e-06,
+      "loss": 0.8027,
+      "output_text_mask_loss": 0.2625,
+      "predictor_learning_rate": 0.0004391323475369626,
+      "step": 2441
+    },
+    {
+      "epoch": 0.25,
+      "gumbel_tau": 0.5622,
+      "image_masks_loss": 0.0203,
+      "llm_learning_rate": 4.390780835334495e-06,
+      "loss": 0.8123,
+      "output_text_mask_loss": 0.0141,
+      "predictor_learning_rate": 0.00043907808353344944,
+      "step": 2442
+    },
+    {
+      "epoch": 0.25,
+      "gumbel_tau": 0.5621,
+      "image_masks_loss": 0.0163,
+      "llm_learning_rate": 4.390237987082368e-06,
+      "loss": 0.7018,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.00043902379870823673,
+      "step": 2443
+    },
+    {
+      "epoch": 0.25,
+      "gumbel_tau": 0.5619,
+      "image_masks_loss": 0.0238,
+      "llm_learning_rate": 4.389694930673023e-06,
+      "loss": 0.7049,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 0.00043896949306730226,
+      "step": 2444
+    },
+    {
+      "epoch": 0.25,
+      "gumbel_tau": 0.5618,
+      "image_masks_loss": 0.0197,
+      "llm_learning_rate": 4.3891516661662625e-06,
+      "loss": 0.7388,
+      "output_text_mask_loss": 0.0305,
+      "predictor_learning_rate": 0.0004389151666166262,
+      "step": 2445
+    },
+    {
+      "epoch": 0.25,
+      "gumbel_tau": 0.5617,
+      "image_masks_loss": 0.0199,
+      "llm_learning_rate": 4.388608193621914e-06,
+      "loss": 0.7031,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.00043886081936219133,
+      "step": 2446
+    },
+    {
+      "epoch": 0.25,
+      "gumbel_tau": 0.5615,
+      "image_masks_loss": 0.0102,
+      "llm_learning_rate": 4.3880645130998236e-06,
+      "loss": 0.8056,
+      "output_text_mask_loss": 0.0166,
+      "predictor_learning_rate": 0.0004388064513099823,
+      "step": 2447
+    },
+    {
+      "epoch": 0.25,
+      "gumbel_tau": 0.5614,
+      "image_masks_loss": 0.0223,
+      "llm_learning_rate": 4.3875206246598634e-06,
+      "loss": 0.7424,
+      "output_text_mask_loss": 0.0467,
+      "predictor_learning_rate": 0.0004387520624659863,
+      "step": 2448
+    },
+    {
+      "epoch": 0.25,
+      "gumbel_tau": 0.5613,
+      "image_masks_loss": 0.0194,
+      "llm_learning_rate": 4.386976528361928e-06,
+      "loss": 0.8407,
+      "output_text_mask_loss": 0.0034,
+      "predictor_learning_rate": 0.0004386976528361927,
+      "step": 2449
+    },
+    {
+      "epoch": 0.25,
+      "gumbel_tau": 0.5611,
+      "image_masks_loss": 0.0271,
+      "llm_learning_rate": 4.386432224265934e-06,
+      "loss": 0.8358,
+      "output_text_mask_loss": 0.0954,
+      "predictor_learning_rate": 0.0004386432224265933,
+      "step": 2450
+    },
+    {
+      "epoch": 0.25,
+      "gumbel_tau": 0.561,
+      "image_masks_loss": 0.0126,
+      "llm_learning_rate": 4.385887712431821e-06,
+      "loss": 0.7291,
+      "output_text_mask_loss": 0.0158,
+      "predictor_learning_rate": 0.00043858877124318207,
+      "step": 2451
+    },
+    {
+      "epoch": 0.25,
+      "gumbel_tau": 0.5609,
+      "image_masks_loss": 0.0332,
+      "llm_learning_rate": 4.385342992919554e-06,
+      "loss": 0.7167,
+      "output_text_mask_loss": 0.0081,
+      "predictor_learning_rate": 0.00043853429929195534,
+      "step": 2452
+    },
+    {
+      "epoch": 0.25,
+      "gumbel_tau": 0.5608,
+      "image_masks_loss": 0.0076,
+      "llm_learning_rate": 4.384798065789116e-06,
+      "loss": 0.7954,
+      "output_text_mask_loss": 0.0275,
+      "predictor_learning_rate": 0.00043847980657891154,
+      "step": 2453
+    },
+    {
+      "epoch": 0.25,
+      "gumbel_tau": 0.5606,
+      "image_masks_loss": 0.0257,
+      "llm_learning_rate": 4.3842529311005175e-06,
+      "loss": 0.845,
+      "output_text_mask_loss": 0.0137,
+      "predictor_learning_rate": 0.0004384252931100517,
+      "step": 2454
+    },
+    {
+      "epoch": 0.25,
+      "gumbel_tau": 0.5605,
+      "image_masks_loss": 0.0133,
+      "llm_learning_rate": 4.3837075889137885e-06,
+      "loss": 0.7519,
+      "output_text_mask_loss": 0.001,
+      "predictor_learning_rate": 0.00043837075889137877,
+      "step": 2455
+    },
+    {
+      "epoch": 0.25,
+      "gumbel_tau": 0.5604,
+      "image_masks_loss": 0.0301,
+      "llm_learning_rate": 4.383162039288984e-06,
+      "loss": 0.7588,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 0.00043831620392889835,
+      "step": 2456
+    },
+    {
+      "epoch": 0.25,
+      "gumbel_tau": 0.5602,
+      "image_masks_loss": 0.0119,
+      "llm_learning_rate": 4.382616282286182e-06,
+      "loss": 0.8284,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00043826162822861817,
+      "step": 2457
+    },
+    {
+      "epoch": 0.25,
+      "gumbel_tau": 0.5601,
+      "image_masks_loss": 0.029,
+      "llm_learning_rate": 4.382070317965481e-06,
+      "loss": 0.7964,
+      "output_text_mask_loss": 0.0893,
+      "predictor_learning_rate": 0.00043820703179654806,
+      "step": 2458
+    },
+    {
+      "epoch": 0.25,
+      "gumbel_tau": 0.56,
+      "image_masks_loss": 0.0227,
+      "llm_learning_rate": 4.381524146387005e-06,
+      "loss": 0.7588,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004381524146387004,
+      "step": 2459
+    },
+    {
+      "epoch": 0.25,
+      "gumbel_tau": 0.5598,
+      "image_masks_loss": 0.0128,
+      "llm_learning_rate": 4.3809777676108975e-06,
+      "loss": 0.8576,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004380977767610897,
+      "step": 2460
+    },
+    {
+      "epoch": 0.25,
+      "gumbel_tau": 0.5597,
+      "image_masks_loss": 0.0093,
+      "llm_learning_rate": 4.380431181697331e-06,
+      "loss": 0.7995,
+      "output_text_mask_loss": 0.0267,
+      "predictor_learning_rate": 0.000438043118169733,
+      "step": 2461
+    },
+    {
+      "epoch": 0.25,
+      "gumbel_tau": 0.5596,
+      "image_masks_loss": 0.0328,
+      "llm_learning_rate": 4.379884388706492e-06,
+      "loss": 0.76,
+      "output_text_mask_loss": 0.0193,
+      "predictor_learning_rate": 0.00043798843887064924,
+      "step": 2462
+    },
+    {
+      "epoch": 0.25,
+      "gumbel_tau": 0.5594,
+      "image_masks_loss": 0.0112,
+      "llm_learning_rate": 4.379337388698598e-06,
+      "loss": 0.7309,
+      "output_text_mask_loss": 0.0032,
+      "predictor_learning_rate": 0.00043793373886985977,
+      "step": 2463
+    },
+    {
+      "epoch": 0.25,
+      "gumbel_tau": 0.5593,
+      "image_masks_loss": 0.0338,
+      "llm_learning_rate": 4.378790181733885e-06,
+      "loss": 0.7992,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.0004378790181733884,
+      "step": 2464
+    },
+    {
+      "epoch": 0.25,
+      "gumbel_tau": 0.5592,
+      "image_masks_loss": 0.021,
+      "llm_learning_rate": 4.378242767872611e-06,
+      "loss": 0.7674,
+      "output_text_mask_loss": 0.0717,
+      "predictor_learning_rate": 0.0004378242767872611,
+      "step": 2465
+    },
+    {
+      "epoch": 0.25,
+      "gumbel_tau": 0.559,
+      "image_masks_loss": 0.0137,
+      "llm_learning_rate": 4.3776951471750604e-06,
+      "loss": 0.7539,
+      "output_text_mask_loss": 0.0095,
+      "predictor_learning_rate": 0.00043776951471750605,
+      "step": 2466
+    },
+    {
+      "epoch": 0.25,
+      "gumbel_tau": 0.5589,
+      "image_masks_loss": 0.0127,
+      "llm_learning_rate": 4.377147319701538e-06,
+      "loss": 0.7992,
+      "output_text_mask_loss": 0.0185,
+      "predictor_learning_rate": 0.00043771473197015383,
+      "step": 2467
+    },
+    {
+      "epoch": 0.25,
+      "gumbel_tau": 0.5588,
+      "image_masks_loss": 0.0139,
+      "llm_learning_rate": 4.3765992855123715e-06,
+      "loss": 0.7822,
+      "output_text_mask_loss": 0.074,
+      "predictor_learning_rate": 0.00043765992855123717,
+      "step": 2468
+    },
+    {
+      "epoch": 0.25,
+      "gumbel_tau": 0.5586,
+      "image_masks_loss": 0.0112,
+      "llm_learning_rate": 4.376051044667913e-06,
+      "loss": 0.7155,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 0.00043760510446679126,
+      "step": 2469
+    },
+    {
+      "epoch": 0.25,
+      "gumbel_tau": 0.5585,
+      "image_masks_loss": 0.013,
+      "llm_learning_rate": 4.375502597228533e-06,
+      "loss": 0.7534,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.0004375502597228533,
+      "step": 2470
+    },
+    {
+      "epoch": 0.25,
+      "gumbel_tau": 0.5584,
+      "image_masks_loss": 0.0146,
+      "llm_learning_rate": 4.37495394325463e-06,
+      "loss": 0.7943,
+      "output_text_mask_loss": 0.0004,
+      "predictor_learning_rate": 0.00043749539432546295,
+      "step": 2471
+    },
+    {
+      "epoch": 0.25,
+      "gumbel_tau": 0.5582,
+      "image_masks_loss": 0.0183,
+      "llm_learning_rate": 4.374405082806622e-06,
+      "loss": 0.7843,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004374405082806622,
+      "step": 2472
+    },
+    {
+      "epoch": 0.25,
+      "gumbel_tau": 0.5581,
+      "image_masks_loss": 0.0088,
+      "llm_learning_rate": 4.3738560159449525e-06,
+      "loss": 0.6848,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.0004373856015944952,
+      "step": 2473
+    },
+    {
+      "epoch": 0.25,
+      "gumbel_tau": 0.558,
+      "image_masks_loss": 0.0121,
+      "llm_learning_rate": 4.373306742730083e-06,
+      "loss": 0.7079,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.00043733067427300835,
+      "step": 2474
+    },
+    {
+      "epoch": 0.25,
+      "gumbel_tau": 0.5578,
+      "image_masks_loss": 0.0154,
+      "llm_learning_rate": 4.372757263222503e-06,
+      "loss": 0.7936,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.0004372757263222503,
+      "step": 2475
+    },
+    {
+      "epoch": 0.25,
+      "gumbel_tau": 0.5577,
+      "image_masks_loss": 0.0079,
+      "llm_learning_rate": 4.3722075774827225e-06,
+      "loss": 0.7492,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.0004372207577482722,
+      "step": 2476
+    },
+    {
+      "epoch": 0.25,
+      "gumbel_tau": 0.5576,
+      "image_masks_loss": 0.0092,
+      "llm_learning_rate": 4.3716576855712735e-06,
+      "loss": 0.7082,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.00043716576855712733,
+      "step": 2477
+    },
+    {
+      "epoch": 0.25,
+      "gumbel_tau": 0.5575,
+      "image_masks_loss": 0.0067,
+      "llm_learning_rate": 4.37110758754871e-06,
+      "loss": 0.7752,
+      "output_text_mask_loss": 0.0225,
+      "predictor_learning_rate": 0.00043711075875487094,
+      "step": 2478
+    },
+    {
+      "epoch": 0.25,
+      "gumbel_tau": 0.5573,
+      "image_masks_loss": 0.0154,
+      "llm_learning_rate": 4.370557283475612e-06,
+      "loss": 0.711,
+      "output_text_mask_loss": 0.0311,
+      "predictor_learning_rate": 0.00043705572834756113,
+      "step": 2479
+    },
+    {
+      "epoch": 0.25,
+      "gumbel_tau": 0.5572,
+      "image_masks_loss": 0.0143,
+      "llm_learning_rate": 4.370006773412579e-06,
+      "loss": 0.8548,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.0004370006773412579,
+      "step": 2480
+    },
+    {
+      "epoch": 0.25,
+      "gumbel_tau": 0.5571,
+      "image_masks_loss": 0.0243,
+      "llm_learning_rate": 4.369456057420235e-06,
+      "loss": 0.7408,
+      "output_text_mask_loss": 0.0183,
+      "predictor_learning_rate": 0.0004369456057420235,
+      "step": 2481
+    },
+    {
+      "epoch": 0.25,
+      "gumbel_tau": 0.5569,
+      "image_masks_loss": 0.0481,
+      "llm_learning_rate": 4.368905135559226e-06,
+      "loss": 0.7424,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.00043689051355592256,
+      "step": 2482
+    },
+    {
+      "epoch": 0.25,
+      "gumbel_tau": 0.5568,
+      "image_masks_loss": 0.0174,
+      "llm_learning_rate": 4.368354007890221e-06,
+      "loss": 0.8385,
+      "output_text_mask_loss": 0.0401,
+      "predictor_learning_rate": 0.0004368354007890221,
+      "step": 2483
+    },
+    {
+      "epoch": 0.25,
+      "gumbel_tau": 0.5567,
+      "image_masks_loss": 0.0175,
+      "llm_learning_rate": 4.367802674473911e-06,
+      "loss": 0.7359,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.000436780267447391,
+      "step": 2484
+    },
+    {
+      "epoch": 0.25,
+      "gumbel_tau": 0.5565,
+      "image_masks_loss": 0.0294,
+      "llm_learning_rate": 4.367251135371009e-06,
+      "loss": 0.7742,
+      "output_text_mask_loss": 0.0004,
+      "predictor_learning_rate": 0.0004367251135371009,
+      "step": 2485
+    },
+    {
+      "epoch": 0.25,
+      "gumbel_tau": 0.5564,
+      "image_masks_loss": 0.0049,
+      "llm_learning_rate": 4.366699390642254e-06,
+      "loss": 0.8322,
+      "output_text_mask_loss": 0.0317,
+      "predictor_learning_rate": 0.00043666993906422535,
+      "step": 2486
+    },
+    {
+      "epoch": 0.25,
+      "gumbel_tau": 0.5563,
+      "image_masks_loss": 0.0044,
+      "llm_learning_rate": 4.366147440348404e-06,
+      "loss": 0.7886,
+      "output_text_mask_loss": 0.0235,
+      "predictor_learning_rate": 0.00043661474403484035,
+      "step": 2487
+    },
+    {
+      "epoch": 0.25,
+      "gumbel_tau": 0.5561,
+      "image_masks_loss": 0.0226,
+      "llm_learning_rate": 4.365595284550241e-06,
+      "loss": 0.7831,
+      "output_text_mask_loss": 0.0408,
+      "predictor_learning_rate": 0.00043655952845502404,
+      "step": 2488
+    },
+    {
+      "epoch": 0.26,
+      "gumbel_tau": 0.556,
+      "image_masks_loss": 0.0546,
+      "llm_learning_rate": 4.3650429233085695e-06,
+      "loss": 0.7646,
+      "output_text_mask_loss": 0.0652,
+      "predictor_learning_rate": 0.00043650429233085694,
+      "step": 2489
+    },
+    {
+      "epoch": 0.26,
+      "gumbel_tau": 0.5559,
+      "image_masks_loss": 0.0204,
+      "llm_learning_rate": 4.364490356684218e-06,
+      "loss": 0.8071,
+      "output_text_mask_loss": 0.024,
+      "predictor_learning_rate": 0.00043644903566842177,
+      "step": 2490
+    },
+    {
+      "epoch": 0.26,
+      "gumbel_tau": 0.5557,
+      "image_masks_loss": 0.0194,
+      "llm_learning_rate": 4.363937584738034e-06,
+      "loss": 0.7743,
+      "output_text_mask_loss": 0.0324,
+      "predictor_learning_rate": 0.00043639375847380344,
+      "step": 2491
+    },
+    {
+      "epoch": 0.26,
+      "gumbel_tau": 0.5556,
+      "image_masks_loss": 0.028,
+      "llm_learning_rate": 4.3633846075308925e-06,
+      "loss": 0.7688,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.00043633846075308926,
+      "step": 2492
+    },
+    {
+      "epoch": 0.26,
+      "gumbel_tau": 0.5555,
+      "image_masks_loss": 0.0433,
+      "llm_learning_rate": 4.362831425123688e-06,
+      "loss": 0.8633,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00043628314251236875,
+      "step": 2493
+    },
+    {
+      "epoch": 0.26,
+      "gumbel_tau": 0.5554,
+      "image_masks_loss": 0.0332,
+      "llm_learning_rate": 4.362278037577336e-06,
+      "loss": 0.7754,
+      "output_text_mask_loss": 0.0404,
+      "predictor_learning_rate": 0.0004362278037577336,
+      "step": 2494
+    },
+    {
+      "epoch": 0.26,
+      "gumbel_tau": 0.5552,
+      "image_masks_loss": 0.0254,
+      "llm_learning_rate": 4.361724444952779e-06,
+      "loss": 0.8497,
+      "output_text_mask_loss": 0.0282,
+      "predictor_learning_rate": 0.00043617244449527796,
+      "step": 2495
+    },
+    {
+      "epoch": 0.26,
+      "gumbel_tau": 0.5551,
+      "image_masks_loss": 0.0338,
+      "llm_learning_rate": 4.36117064731098e-06,
+      "loss": 0.7893,
+      "output_text_mask_loss": 0.0046,
+      "predictor_learning_rate": 0.000436117064731098,
+      "step": 2496
+    },
+    {
+      "epoch": 0.26,
+      "gumbel_tau": 0.555,
+      "image_masks_loss": 0.0109,
+      "llm_learning_rate": 4.360616644712924e-06,
+      "loss": 0.7678,
+      "output_text_mask_loss": 0.0141,
+      "predictor_learning_rate": 0.00043606166447129234,
+      "step": 2497
+    },
+    {
+      "epoch": 0.26,
+      "gumbel_tau": 0.5548,
+      "image_masks_loss": 0.0687,
+      "llm_learning_rate": 4.360062437219618e-06,
+      "loss": 0.8307,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.00043600624372196173,
+      "step": 2498
+    },
+    {
+      "epoch": 0.26,
+      "gumbel_tau": 0.5547,
+      "image_masks_loss": 0.0087,
+      "llm_learning_rate": 4.359508024892093e-06,
+      "loss": 0.8557,
+      "output_text_mask_loss": 0.0246,
+      "predictor_learning_rate": 0.00043595080248920925,
+      "step": 2499
+    },
+    {
+      "epoch": 0.26,
+      "gumbel_tau": 0.5546,
+      "image_masks_loss": 0.0277,
+      "llm_learning_rate": 4.358953407791403e-06,
+      "loss": 0.7481,
+      "output_text_mask_loss": 0.0313,
+      "predictor_learning_rate": 0.00043589534077914026,
+      "step": 2500
+    },
+    {
+      "epoch": 0.26,
+      "gumbel_tau": 0.5544,
+      "image_masks_loss": 0.0188,
+      "llm_learning_rate": 4.358398585978623e-06,
+      "loss": 0.8019,
+      "output_text_mask_loss": 0.0507,
+      "predictor_learning_rate": 0.0004358398585978622,
+      "step": 2501
+    },
+    {
+      "epoch": 0.26,
+      "gumbel_tau": 0.5543,
+      "image_masks_loss": 0.0259,
+      "llm_learning_rate": 4.3578435595148505e-06,
+      "loss": 0.7377,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.000435784355951485,
+      "step": 2502
+    },
+    {
+      "epoch": 0.26,
+      "gumbel_tau": 0.5542,
+      "image_masks_loss": 0.01,
+      "llm_learning_rate": 4.3572883284612075e-06,
+      "loss": 0.7419,
+      "output_text_mask_loss": 0.0077,
+      "predictor_learning_rate": 0.00043572883284612076,
+      "step": 2503
+    },
+    {
+      "epoch": 0.26,
+      "gumbel_tau": 0.554,
+      "image_masks_loss": 0.0123,
+      "llm_learning_rate": 4.356732892878838e-06,
+      "loss": 0.7071,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.00043567328928788377,
+      "step": 2504
+    },
+    {
+      "epoch": 0.26,
+      "gumbel_tau": 0.5539,
+      "image_masks_loss": 0.0059,
+      "llm_learning_rate": 4.3561772528289056e-06,
+      "loss": 0.7939,
+      "output_text_mask_loss": 0.0015,
+      "predictor_learning_rate": 0.0004356177252828905,
+      "step": 2505
+    },
+    {
+      "epoch": 0.26,
+      "gumbel_tau": 0.5538,
+      "image_masks_loss": 0.0368,
+      "llm_learning_rate": 4.3556214083726e-06,
+      "loss": 0.7759,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.00043556214083725995,
+      "step": 2506
+    },
+    {
+      "epoch": 0.26,
+      "gumbel_tau": 0.5537,
+      "image_masks_loss": 0.0088,
+      "llm_learning_rate": 4.355065359571131e-06,
+      "loss": 0.74,
+      "output_text_mask_loss": 0.0124,
+      "predictor_learning_rate": 0.00043550653595711306,
+      "step": 2507
+    },
+    {
+      "epoch": 0.26,
+      "gumbel_tau": 0.5535,
+      "image_masks_loss": 0.0317,
+      "llm_learning_rate": 4.354509106485733e-06,
+      "loss": 0.8083,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.00043545091064857325,
+      "step": 2508
+    },
+    {
+      "epoch": 0.26,
+      "gumbel_tau": 0.5534,
+      "image_masks_loss": 0.0124,
+      "llm_learning_rate": 4.353952649177661e-06,
+      "loss": 0.8672,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 0.00043539526491776604,
+      "step": 2509
+    },
+    {
+      "epoch": 0.26,
+      "gumbel_tau": 0.5533,
+      "image_masks_loss": 0.0138,
+      "llm_learning_rate": 4.3533959877081946e-06,
+      "loss": 0.7662,
+      "output_text_mask_loss": 0.0185,
+      "predictor_learning_rate": 0.0004353395987708194,
+      "step": 2510
+    },
+    {
+      "epoch": 0.26,
+      "gumbel_tau": 0.5531,
+      "image_masks_loss": 0.0123,
+      "llm_learning_rate": 4.3528391221386325e-06,
+      "loss": 0.7143,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00043528391221386327,
+      "step": 2511
+    },
+    {
+      "epoch": 0.26,
+      "gumbel_tau": 0.553,
+      "image_masks_loss": 0.0141,
+      "llm_learning_rate": 4.352282052530301e-06,
+      "loss": 0.7173,
+      "output_text_mask_loss": 0.0437,
+      "predictor_learning_rate": 0.00043522820525303003,
+      "step": 2512
+    },
+    {
+      "epoch": 0.26,
+      "gumbel_tau": 0.5529,
+      "image_masks_loss": 0.0038,
+      "llm_learning_rate": 4.351724778944543e-06,
+      "loss": 0.7265,
+      "output_text_mask_loss": 0.0101,
+      "predictor_learning_rate": 0.00043517247789445426,
+      "step": 2513
+    },
+    {
+      "epoch": 0.26,
+      "gumbel_tau": 0.5527,
+      "image_masks_loss": 0.0098,
+      "llm_learning_rate": 4.351167301442728e-06,
+      "loss": 0.7812,
+      "output_text_mask_loss": 0.071,
+      "predictor_learning_rate": 0.0004351167301442728,
+      "step": 2514
+    },
+    {
+      "epoch": 0.26,
+      "gumbel_tau": 0.5526,
+      "image_masks_loss": 0.0176,
+      "llm_learning_rate": 4.350609620086246e-06,
+      "loss": 0.7496,
+      "output_text_mask_loss": 0.0195,
+      "predictor_learning_rate": 0.00043506096200862453,
+      "step": 2515
+    },
+    {
+      "epoch": 0.26,
+      "gumbel_tau": 0.5525,
+      "image_masks_loss": 0.019,
+      "llm_learning_rate": 4.3500517349365115e-06,
+      "loss": 0.8082,
+      "output_text_mask_loss": 0.0195,
+      "predictor_learning_rate": 0.0004350051734936511,
+      "step": 2516
+    },
+    {
+      "epoch": 0.26,
+      "gumbel_tau": 0.5523,
+      "image_masks_loss": 0.0475,
+      "llm_learning_rate": 4.349493646054958e-06,
+      "loss": 0.7763,
+      "output_text_mask_loss": 0.0231,
+      "predictor_learning_rate": 0.0004349493646054957,
+      "step": 2517
+    },
+    {
+      "epoch": 0.26,
+      "gumbel_tau": 0.5522,
+      "image_masks_loss": 0.0267,
+      "llm_learning_rate": 4.348935353503045e-06,
+      "loss": 0.7442,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.00043489353535030445,
+      "step": 2518
+    },
+    {
+      "epoch": 0.26,
+      "gumbel_tau": 0.5521,
+      "image_masks_loss": 0.0146,
+      "llm_learning_rate": 4.348376857342252e-06,
+      "loss": 0.7091,
+      "output_text_mask_loss": 0.017,
+      "predictor_learning_rate": 0.0004348376857342252,
+      "step": 2519
+    },
+    {
+      "epoch": 0.26,
+      "gumbel_tau": 0.552,
+      "image_masks_loss": 0.0174,
+      "llm_learning_rate": 4.347818157634084e-06,
+      "loss": 0.7811,
+      "output_text_mask_loss": 0.0282,
+      "predictor_learning_rate": 0.00043478181576340835,
+      "step": 2520
+    },
+    {
+      "epoch": 0.26,
+      "gumbel_tau": 0.5518,
+      "image_masks_loss": 0.009,
+      "llm_learning_rate": 4.347259254440064e-06,
+      "loss": 0.7713,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004347259254440064,
+      "step": 2521
+    },
+    {
+      "epoch": 0.26,
+      "gumbel_tau": 0.5517,
+      "image_masks_loss": 0.0094,
+      "llm_learning_rate": 4.34670014782174e-06,
+      "loss": 0.6992,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004346700147821739,
+      "step": 2522
+    },
+    {
+      "epoch": 0.26,
+      "gumbel_tau": 0.5516,
+      "image_masks_loss": 0.0135,
+      "llm_learning_rate": 4.3461408378406815e-06,
+      "loss": 0.7213,
+      "output_text_mask_loss": 0.006,
+      "predictor_learning_rate": 0.0004346140837840681,
+      "step": 2523
+    },
+    {
+      "epoch": 0.26,
+      "gumbel_tau": 0.5514,
+      "image_masks_loss": 0.0187,
+      "llm_learning_rate": 4.3455813245584835e-06,
+      "loss": 0.7317,
+      "output_text_mask_loss": 0.0039,
+      "predictor_learning_rate": 0.0004345581324558483,
+      "step": 2524
+    },
+    {
+      "epoch": 0.26,
+      "gumbel_tau": 0.5513,
+      "image_masks_loss": 0.0168,
+      "llm_learning_rate": 4.345021608036758e-06,
+      "loss": 0.6921,
+      "output_text_mask_loss": 0.0269,
+      "predictor_learning_rate": 0.0004345021608036758,
+      "step": 2525
+    },
+    {
+      "epoch": 0.26,
+      "gumbel_tau": 0.5512,
+      "image_masks_loss": 0.0278,
+      "llm_learning_rate": 4.344461688337144e-06,
+      "loss": 0.8175,
+      "output_text_mask_loss": 0.0124,
+      "predictor_learning_rate": 0.0004344461688337143,
+      "step": 2526
+    },
+    {
+      "epoch": 0.26,
+      "gumbel_tau": 0.551,
+      "image_masks_loss": 0.0138,
+      "llm_learning_rate": 4.343901565521299e-06,
+      "loss": 0.8022,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.0004343901565521299,
+      "step": 2527
+    },
+    {
+      "epoch": 0.26,
+      "gumbel_tau": 0.5509,
+      "image_masks_loss": 0.0139,
+      "llm_learning_rate": 4.343341239650908e-06,
+      "loss": 0.7442,
+      "output_text_mask_loss": 0.0145,
+      "predictor_learning_rate": 0.00043433412396509083,
+      "step": 2528
+    },
+    {
+      "epoch": 0.26,
+      "gumbel_tau": 0.5508,
+      "image_masks_loss": 0.0146,
+      "llm_learning_rate": 4.3427807107876744e-06,
+      "loss": 0.7343,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.0004342780710787674,
+      "step": 2529
+    },
+    {
+      "epoch": 0.26,
+      "gumbel_tau": 0.5507,
+      "image_masks_loss": 0.0179,
+      "llm_learning_rate": 4.342219978993323e-06,
+      "loss": 0.9059,
+      "output_text_mask_loss": 0.1694,
+      "predictor_learning_rate": 0.0004342219978993323,
+      "step": 2530
+    },
+    {
+      "epoch": 0.26,
+      "gumbel_tau": 0.5505,
+      "image_masks_loss": 0.0326,
+      "llm_learning_rate": 4.341659044329605e-06,
+      "loss": 1.041,
+      "output_text_mask_loss": 0.1434,
+      "predictor_learning_rate": 0.00043416590443296046,
+      "step": 2531
+    },
+    {
+      "epoch": 0.26,
+      "gumbel_tau": 0.5504,
+      "image_masks_loss": 0.0221,
+      "llm_learning_rate": 4.34109790685829e-06,
+      "loss": 0.8887,
+      "output_text_mask_loss": 0.235,
+      "predictor_learning_rate": 0.00043410979068582895,
+      "step": 2532
+    },
+    {
+      "epoch": 0.26,
+      "gumbel_tau": 0.5503,
+      "image_masks_loss": 0.0201,
+      "llm_learning_rate": 4.340536566641173e-06,
+      "loss": 0.8829,
+      "output_text_mask_loss": 0.1755,
+      "predictor_learning_rate": 0.0004340536566641173,
+      "step": 2533
+    },
+    {
+      "epoch": 0.26,
+      "gumbel_tau": 0.5501,
+      "image_masks_loss": 0.0092,
+      "llm_learning_rate": 4.339975023740069e-06,
+      "loss": 1.0031,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004339975023740069,
+      "step": 2534
+    },
+    {
+      "epoch": 0.26,
+      "gumbel_tau": 0.55,
+      "image_masks_loss": 0.0326,
+      "llm_learning_rate": 4.339413278216819e-06,
+      "loss": 0.8549,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.0004339413278216818,
+      "step": 2535
+    },
+    {
+      "epoch": 0.26,
+      "gumbel_tau": 0.5499,
+      "image_masks_loss": 0.0546,
+      "llm_learning_rate": 4.338851330133279e-06,
+      "loss": 0.8026,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00043388513301332796,
+      "step": 2536
+    },
+    {
+      "epoch": 0.26,
+      "gumbel_tau": 0.5497,
+      "image_masks_loss": 0.0326,
+      "llm_learning_rate": 4.338289179551337e-06,
+      "loss": 0.7828,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.0004338289179551337,
+      "step": 2537
+    },
+    {
+      "epoch": 0.26,
+      "gumbel_tau": 0.5496,
+      "image_masks_loss": 0.012,
+      "llm_learning_rate": 4.337726826532894e-06,
+      "loss": 0.7723,
+      "output_text_mask_loss": 0.0206,
+      "predictor_learning_rate": 0.00043377268265328944,
+      "step": 2538
+    },
+    {
+      "epoch": 0.26,
+      "gumbel_tau": 0.5495,
+      "image_masks_loss": 0.0164,
+      "llm_learning_rate": 4.337164271139881e-06,
+      "loss": 0.8522,
+      "output_text_mask_loss": 0.0702,
+      "predictor_learning_rate": 0.00043371642711398806,
+      "step": 2539
+    },
+    {
+      "epoch": 0.26,
+      "gumbel_tau": 0.5494,
+      "image_masks_loss": 0.0089,
+      "llm_learning_rate": 4.336601513434246e-06,
+      "loss": 0.8921,
+      "output_text_mask_loss": 0.0977,
+      "predictor_learning_rate": 0.00043366015134342455,
+      "step": 2540
+    },
+    {
+      "epoch": 0.26,
+      "gumbel_tau": 0.5492,
+      "image_masks_loss": 0.0185,
+      "llm_learning_rate": 4.336038553477961e-06,
+      "loss": 0.8386,
+      "output_text_mask_loss": 0.2365,
+      "predictor_learning_rate": 0.00043360385534779613,
+      "step": 2541
+    },
+    {
+      "epoch": 0.26,
+      "gumbel_tau": 0.5491,
+      "image_masks_loss": 0.0284,
+      "llm_learning_rate": 4.335475391333022e-06,
+      "loss": 0.863,
+      "output_text_mask_loss": 0.2823,
+      "predictor_learning_rate": 0.00043354753913330213,
+      "step": 2542
+    },
+    {
+      "epoch": 0.26,
+      "gumbel_tau": 0.549,
+      "image_masks_loss": 0.0178,
+      "llm_learning_rate": 4.334912027061443e-06,
+      "loss": 0.8313,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004334912027061443,
+      "step": 2543
+    },
+    {
+      "epoch": 0.26,
+      "gumbel_tau": 0.5488,
+      "image_masks_loss": 0.0307,
+      "llm_learning_rate": 4.334348460725265e-06,
+      "loss": 0.9855,
+      "output_text_mask_loss": 0.1717,
+      "predictor_learning_rate": 0.0004334348460725265,
+      "step": 2544
+    },
+    {
+      "epoch": 0.26,
+      "gumbel_tau": 0.5487,
+      "image_masks_loss": 0.0198,
+      "llm_learning_rate": 4.333784692386549e-06,
+      "loss": 0.8405,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 0.0004333784692386549,
+      "step": 2545
+    },
+    {
+      "epoch": 0.26,
+      "gumbel_tau": 0.5486,
+      "image_masks_loss": 0.0267,
+      "llm_learning_rate": 4.333220722107377e-06,
+      "loss": 0.9644,
+      "output_text_mask_loss": 0.2258,
+      "predictor_learning_rate": 0.0004333220722107377,
+      "step": 2546
+    },
+    {
+      "epoch": 0.26,
+      "gumbel_tau": 0.5485,
+      "image_masks_loss": 0.0216,
+      "llm_learning_rate": 4.3326565499498575e-06,
+      "loss": 1.0044,
+      "output_text_mask_loss": 0.1694,
+      "predictor_learning_rate": 0.0004332656549949857,
+      "step": 2547
+    },
+    {
+      "epoch": 0.26,
+      "gumbel_tau": 0.5483,
+      "image_masks_loss": 0.0171,
+      "llm_learning_rate": 4.332092175976116e-06,
+      "loss": 0.8898,
+      "output_text_mask_loss": 0.267,
+      "predictor_learning_rate": 0.0004332092175976116,
+      "step": 2548
+    },
+    {
+      "epoch": 0.26,
+      "gumbel_tau": 0.5482,
+      "image_masks_loss": 0.0364,
+      "llm_learning_rate": 4.331527600248303e-06,
+      "loss": 0.8072,
+      "output_text_mask_loss": 0.0429,
+      "predictor_learning_rate": 0.00043315276002483025,
+      "step": 2549
+    },
+    {
+      "epoch": 0.26,
+      "gumbel_tau": 0.5481,
+      "image_masks_loss": 0.0326,
+      "llm_learning_rate": 4.33096282282859e-06,
+      "loss": 0.8779,
+      "output_text_mask_loss": 0.0429,
+      "predictor_learning_rate": 0.00043309628228285904,
+      "step": 2550
+    },
+    {
+      "epoch": 0.26,
+      "gumbel_tau": 0.5479,
+      "image_masks_loss": 0.0237,
+      "llm_learning_rate": 4.330397843779175e-06,
+      "loss": 0.8663,
+      "output_text_mask_loss": 0.2197,
+      "predictor_learning_rate": 0.0004330397843779174,
+      "step": 2551
+    },
+    {
+      "epoch": 0.26,
+      "gumbel_tau": 0.5478,
+      "image_masks_loss": 0.0079,
+      "llm_learning_rate": 4.32983266316227e-06,
+      "loss": 0.7546,
+      "output_text_mask_loss": 0.0238,
+      "predictor_learning_rate": 0.00043298326631622695,
+      "step": 2552
+    },
+    {
+      "epoch": 0.26,
+      "gumbel_tau": 0.5477,
+      "image_masks_loss": 0.0158,
+      "llm_learning_rate": 4.3292672810401184e-06,
+      "loss": 0.7684,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004329267281040118,
+      "step": 2553
+    },
+    {
+      "epoch": 0.26,
+      "gumbel_tau": 0.5475,
+      "image_masks_loss": 0.008,
+      "llm_learning_rate": 4.328701697474977e-06,
+      "loss": 0.7792,
+      "output_text_mask_loss": 0.0071,
+      "predictor_learning_rate": 0.0004328701697474977,
+      "step": 2554
+    },
+    {
+      "epoch": 0.26,
+      "gumbel_tau": 0.5474,
+      "image_masks_loss": 0.0206,
+      "llm_learning_rate": 4.328135912529133e-06,
+      "loss": 0.5763,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 0.0004328135912529133,
+      "step": 2555
+    },
+    {
+      "epoch": 0.26,
+      "gumbel_tau": 0.5473,
+      "image_masks_loss": 0.0176,
+      "llm_learning_rate": 4.32756992626489e-06,
+      "loss": 0.8104,
+      "output_text_mask_loss": 0.0275,
+      "predictor_learning_rate": 0.00043275699262648905,
+      "step": 2556
+    },
+    {
+      "epoch": 0.26,
+      "gumbel_tau": 0.5472,
+      "image_masks_loss": 0.0181,
+      "llm_learning_rate": 4.327003738744576e-06,
+      "loss": 0.5526,
+      "output_text_mask_loss": 0.0322,
+      "predictor_learning_rate": 0.0004327003738744576,
+      "step": 2557
+    },
+    {
+      "epoch": 0.26,
+      "gumbel_tau": 0.547,
+      "image_masks_loss": 0.0138,
+      "llm_learning_rate": 4.326437350030541e-06,
+      "loss": 0.7497,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.0004326437350030541,
+      "step": 2558
+    },
+    {
+      "epoch": 0.26,
+      "gumbel_tau": 0.5469,
+      "image_masks_loss": 0.0185,
+      "llm_learning_rate": 4.3258707601851565e-06,
+      "loss": 0.8357,
+      "output_text_mask_loss": 0.0095,
+      "predictor_learning_rate": 0.00043258707601851566,
+      "step": 2559
+    },
+    {
+      "epoch": 0.26,
+      "gumbel_tau": 0.5468,
+      "image_masks_loss": 0.0168,
+      "llm_learning_rate": 4.325303969270818e-06,
+      "loss": 0.7579,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.00043253039692708173,
+      "step": 2560
+    },
+    {
+      "epoch": 0.26,
+      "gumbel_tau": 0.5466,
+      "image_masks_loss": 0.0172,
+      "llm_learning_rate": 4.324736977349939e-06,
+      "loss": 0.8183,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.0004324736977349939,
+      "step": 2561
+    },
+    {
+      "epoch": 0.26,
+      "gumbel_tau": 0.5465,
+      "image_masks_loss": 0.0219,
+      "llm_learning_rate": 4.324169784484961e-06,
+      "loss": 0.8268,
+      "output_text_mask_loss": 0.0433,
+      "predictor_learning_rate": 0.00043241697844849603,
+      "step": 2562
+    },
+    {
+      "epoch": 0.26,
+      "gumbel_tau": 0.5464,
+      "image_masks_loss": 0.0313,
+      "llm_learning_rate": 4.323602390738342e-06,
+      "loss": 0.761,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.00043236023907383415,
+      "step": 2563
+    },
+    {
+      "epoch": 0.26,
+      "gumbel_tau": 0.5463,
+      "image_masks_loss": 0.0072,
+      "llm_learning_rate": 4.323034796172566e-06,
+      "loss": 0.7149,
+      "output_text_mask_loss": 0.0076,
+      "predictor_learning_rate": 0.0004323034796172566,
+      "step": 2564
+    },
+    {
+      "epoch": 0.26,
+      "gumbel_tau": 0.5461,
+      "image_masks_loss": 0.0059,
+      "llm_learning_rate": 4.322467000850138e-06,
+      "loss": 0.6931,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 0.0004322467000850138,
+      "step": 2565
+    },
+    {
+      "epoch": 0.26,
+      "gumbel_tau": 0.546,
+      "image_masks_loss": 0.0136,
+      "llm_learning_rate": 4.321899004833585e-06,
+      "loss": 0.6903,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00043218990048335843,
+      "step": 2566
+    },
+    {
+      "epoch": 0.26,
+      "gumbel_tau": 0.5459,
+      "image_masks_loss": 0.0113,
+      "llm_learning_rate": 4.321330808185454e-06,
+      "loss": 0.7312,
+      "output_text_mask_loss": 0.0395,
+      "predictor_learning_rate": 0.0004321330808185454,
+      "step": 2567
+    },
+    {
+      "epoch": 0.26,
+      "gumbel_tau": 0.5457,
+      "image_masks_loss": 0.0107,
+      "llm_learning_rate": 4.320762410968319e-06,
+      "loss": 0.7596,
+      "output_text_mask_loss": 0.0458,
+      "predictor_learning_rate": 0.00043207624109683183,
+      "step": 2568
+    },
+    {
+      "epoch": 0.26,
+      "gumbel_tau": 0.5456,
+      "image_masks_loss": 0.0292,
+      "llm_learning_rate": 4.3201938132447706e-06,
+      "loss": 0.7978,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.00043201938132447705,
+      "step": 2569
+    },
+    {
+      "epoch": 0.26,
+      "gumbel_tau": 0.5455,
+      "image_masks_loss": 0.0232,
+      "llm_learning_rate": 4.319625015077426e-06,
+      "loss": 0.7156,
+      "output_text_mask_loss": 0.0277,
+      "predictor_learning_rate": 0.0004319625015077426,
+      "step": 2570
+    },
+    {
+      "epoch": 0.26,
+      "gumbel_tau": 0.5454,
+      "image_masks_loss": 0.0145,
+      "llm_learning_rate": 4.319056016528922e-06,
+      "loss": 0.7882,
+      "output_text_mask_loss": 0.0231,
+      "predictor_learning_rate": 0.0004319056016528921,
+      "step": 2571
+    },
+    {
+      "epoch": 0.26,
+      "gumbel_tau": 0.5452,
+      "image_masks_loss": 0.0256,
+      "llm_learning_rate": 4.318486817661917e-06,
+      "loss": 0.7626,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00043184868176619165,
+      "step": 2572
+    },
+    {
+      "epoch": 0.26,
+      "gumbel_tau": 0.5451,
+      "image_masks_loss": 0.0163,
+      "llm_learning_rate": 4.3179174185390925e-06,
+      "loss": 0.7362,
+      "output_text_mask_loss": 0.0523,
+      "predictor_learning_rate": 0.00043179174185390925,
+      "step": 2573
+    },
+    {
+      "epoch": 0.26,
+      "gumbel_tau": 0.545,
+      "image_masks_loss": 0.0163,
+      "llm_learning_rate": 4.3173478192231535e-06,
+      "loss": 0.8393,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 0.0004317347819223153,
+      "step": 2574
+    },
+    {
+      "epoch": 0.26,
+      "gumbel_tau": 0.5448,
+      "image_masks_loss": 0.0227,
+      "llm_learning_rate": 4.316778019776824e-06,
+      "loss": 0.892,
+      "output_text_mask_loss": 0.0332,
+      "predictor_learning_rate": 0.00043167780197768234,
+      "step": 2575
+    },
+    {
+      "epoch": 0.26,
+      "gumbel_tau": 0.5447,
+      "image_masks_loss": 0.034,
+      "llm_learning_rate": 4.316208020262853e-06,
+      "loss": 0.7887,
+      "output_text_mask_loss": 0.0153,
+      "predictor_learning_rate": 0.0004316208020262853,
+      "step": 2576
+    },
+    {
+      "epoch": 0.26,
+      "gumbel_tau": 0.5446,
+      "image_masks_loss": 0.0137,
+      "llm_learning_rate": 4.3156378207440085e-06,
+      "loss": 0.7202,
+      "output_text_mask_loss": 0.0618,
+      "predictor_learning_rate": 0.00043156378207440083,
+      "step": 2577
+    },
+    {
+      "epoch": 0.26,
+      "gumbel_tau": 0.5445,
+      "image_masks_loss": 0.0177,
+      "llm_learning_rate": 4.315067421283083e-06,
+      "loss": 0.8912,
+      "output_text_mask_loss": 0.1221,
+      "predictor_learning_rate": 0.0004315067421283083,
+      "step": 2578
+    },
+    {
+      "epoch": 0.26,
+      "gumbel_tau": 0.5443,
+      "image_masks_loss": 0.0237,
+      "llm_learning_rate": 4.314496821942892e-06,
+      "loss": 0.906,
+      "output_text_mask_loss": 0.1389,
+      "predictor_learning_rate": 0.00043144968219428916,
+      "step": 2579
+    },
+    {
+      "epoch": 0.26,
+      "gumbel_tau": 0.5442,
+      "image_masks_loss": 0.0122,
+      "llm_learning_rate": 4.313926022786268e-06,
+      "loss": 0.8262,
+      "output_text_mask_loss": 0.1328,
+      "predictor_learning_rate": 0.00043139260227862677,
+      "step": 2580
+    },
+    {
+      "epoch": 0.26,
+      "gumbel_tau": 0.5441,
+      "image_masks_loss": 0.0083,
+      "llm_learning_rate": 4.313355023876069e-06,
+      "loss": 0.8549,
+      "output_text_mask_loss": 0.1167,
+      "predictor_learning_rate": 0.0004313355023876069,
+      "step": 2581
+    },
+    {
+      "epoch": 0.26,
+      "gumbel_tau": 0.5439,
+      "image_masks_loss": 0.0218,
+      "llm_learning_rate": 4.312783825275178e-06,
+      "loss": 0.8514,
+      "output_text_mask_loss": 0.1686,
+      "predictor_learning_rate": 0.00043127838252751775,
+      "step": 2582
+    },
+    {
+      "epoch": 0.26,
+      "gumbel_tau": 0.5438,
+      "image_masks_loss": 0.019,
+      "llm_learning_rate": 4.312212427046492e-06,
+      "loss": 0.7717,
+      "output_text_mask_loss": 0.1984,
+      "predictor_learning_rate": 0.0004312212427046492,
+      "step": 2583
+    },
+    {
+      "epoch": 0.26,
+      "gumbel_tau": 0.5437,
+      "image_masks_loss": 0.0094,
+      "llm_learning_rate": 4.3116408292529386e-06,
+      "loss": 0.8958,
+      "output_text_mask_loss": 0.1717,
+      "predictor_learning_rate": 0.00043116408292529385,
+      "step": 2584
+    },
+    {
+      "epoch": 0.26,
+      "gumbel_tau": 0.5436,
+      "image_masks_loss": 0.0395,
+      "llm_learning_rate": 4.311069031957461e-06,
+      "loss": 0.9672,
+      "output_text_mask_loss": 0.1411,
+      "predictor_learning_rate": 0.00043110690319574606,
+      "step": 2585
+    },
+    {
+      "epoch": 0.26,
+      "gumbel_tau": 0.5434,
+      "image_masks_loss": 0.0103,
+      "llm_learning_rate": 4.310497035223028e-06,
+      "loss": 0.7703,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 0.0004310497035223028,
+      "step": 2586
+    },
+    {
+      "epoch": 0.27,
+      "gumbel_tau": 0.5433,
+      "image_masks_loss": 0.014,
+      "llm_learning_rate": 4.309924839112628e-06,
+      "loss": 0.9151,
+      "output_text_mask_loss": 0.1793,
+      "predictor_learning_rate": 0.0004309924839112628,
+      "step": 2587
+    },
+    {
+      "epoch": 0.27,
+      "gumbel_tau": 0.5432,
+      "image_masks_loss": 0.0149,
+      "llm_learning_rate": 4.309352443689274e-06,
+      "loss": 0.7584,
+      "output_text_mask_loss": 0.0431,
+      "predictor_learning_rate": 0.0004309352443689274,
+      "step": 2588
+    },
+    {
+      "epoch": 0.27,
+      "gumbel_tau": 0.543,
+      "image_masks_loss": 0.0263,
+      "llm_learning_rate": 4.308779849015999e-06,
+      "loss": 0.8361,
+      "output_text_mask_loss": 0.0061,
+      "predictor_learning_rate": 0.0004308779849015999,
+      "step": 2589
+    },
+    {
+      "epoch": 0.27,
+      "gumbel_tau": 0.5429,
+      "image_masks_loss": 0.0341,
+      "llm_learning_rate": 4.308207055155858e-06,
+      "loss": 0.9085,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.00043082070551558573,
+      "step": 2590
+    },
+    {
+      "epoch": 0.27,
+      "gumbel_tau": 0.5428,
+      "image_masks_loss": 0.0215,
+      "llm_learning_rate": 4.307634062171928e-06,
+      "loss": 0.7426,
+      "output_text_mask_loss": 0.0345,
+      "predictor_learning_rate": 0.00043076340621719277,
+      "step": 2591
+    },
+    {
+      "epoch": 0.27,
+      "gumbel_tau": 0.5427,
+      "image_masks_loss": 0.018,
+      "llm_learning_rate": 4.307060870127309e-06,
+      "loss": 0.7875,
+      "output_text_mask_loss": 0.1312,
+      "predictor_learning_rate": 0.00043070608701273083,
+      "step": 2592
+    },
+    {
+      "epoch": 0.27,
+      "gumbel_tau": 0.5425,
+      "image_masks_loss": 0.0081,
+      "llm_learning_rate": 4.306487479085122e-06,
+      "loss": 0.7522,
+      "output_text_mask_loss": 0.0156,
+      "predictor_learning_rate": 0.00043064874790851214,
+      "step": 2593
+    },
+    {
+      "epoch": 0.27,
+      "gumbel_tau": 0.5424,
+      "image_masks_loss": 0.0422,
+      "llm_learning_rate": 4.30591388910851e-06,
+      "loss": 0.8338,
+      "output_text_mask_loss": 0.0893,
+      "predictor_learning_rate": 0.0004305913889108509,
+      "step": 2594
+    },
+    {
+      "epoch": 0.27,
+      "gumbel_tau": 0.5423,
+      "image_masks_loss": 0.0111,
+      "llm_learning_rate": 4.305340100260637e-06,
+      "loss": 0.8101,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 0.00043053401002606374,
+      "step": 2595
+    },
+    {
+      "epoch": 0.27,
+      "gumbel_tau": 0.5421,
+      "image_masks_loss": 0.0299,
+      "llm_learning_rate": 4.304766112604692e-06,
+      "loss": 0.6273,
+      "output_text_mask_loss": 0.013,
+      "predictor_learning_rate": 0.00043047661126046916,
+      "step": 2596
+    },
+    {
+      "epoch": 0.27,
+      "gumbel_tau": 0.542,
+      "image_masks_loss": 0.0155,
+      "llm_learning_rate": 4.3041919262038825e-06,
+      "loss": 0.7635,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.0004304191926203882,
+      "step": 2597
+    },
+    {
+      "epoch": 0.27,
+      "gumbel_tau": 0.5419,
+      "image_masks_loss": 0.0231,
+      "llm_learning_rate": 4.30361754112144e-06,
+      "loss": 0.7164,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00043036175411214397,
+      "step": 2598
+    },
+    {
+      "epoch": 0.27,
+      "gumbel_tau": 0.5418,
+      "image_masks_loss": 0.0148,
+      "llm_learning_rate": 4.303042957420615e-06,
+      "loss": 0.7008,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00043030429574206154,
+      "step": 2599
+    },
+    {
+      "epoch": 0.27,
+      "gumbel_tau": 0.5416,
+      "image_masks_loss": 0.0096,
+      "llm_learning_rate": 4.3024681751646845e-06,
+      "loss": 0.7961,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00043024681751646843,
+      "step": 2600
+    },
+    {
+      "epoch": 0.27,
+      "gumbel_tau": 0.5415,
+      "image_masks_loss": 0.0169,
+      "llm_learning_rate": 4.301893194416944e-06,
+      "loss": 0.6916,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.0004301893194416944,
+      "step": 2601
+    },
+    {
+      "epoch": 0.27,
+      "gumbel_tau": 0.5414,
+      "image_masks_loss": 0.0257,
+      "llm_learning_rate": 4.301318015240711e-06,
+      "loss": 0.8292,
+      "output_text_mask_loss": 0.005,
+      "predictor_learning_rate": 0.00043013180152407107,
+      "step": 2602
+    },
+    {
+      "epoch": 0.27,
+      "gumbel_tau": 0.5413,
+      "image_masks_loss": 0.0103,
+      "llm_learning_rate": 4.300742637699325e-06,
+      "loss": 0.7152,
+      "output_text_mask_loss": 0.09,
+      "predictor_learning_rate": 0.0004300742637699325,
+      "step": 2603
+    },
+    {
+      "epoch": 0.27,
+      "gumbel_tau": 0.5411,
+      "image_masks_loss": 0.017,
+      "llm_learning_rate": 4.300167061856149e-06,
+      "loss": 0.8069,
+      "output_text_mask_loss": 0.0492,
+      "predictor_learning_rate": 0.00043001670618561495,
+      "step": 2604
+    },
+    {
+      "epoch": 0.27,
+      "gumbel_tau": 0.541,
+      "image_masks_loss": 0.0089,
+      "llm_learning_rate": 4.299591287774567e-06,
+      "loss": 0.7541,
+      "output_text_mask_loss": 0.0238,
+      "predictor_learning_rate": 0.00042995912877745674,
+      "step": 2605
+    },
+    {
+      "epoch": 0.27,
+      "gumbel_tau": 0.5409,
+      "image_masks_loss": 0.0083,
+      "llm_learning_rate": 4.299015315517985e-06,
+      "loss": 0.8387,
+      "output_text_mask_loss": 0.0353,
+      "predictor_learning_rate": 0.00042990153155179846,
+      "step": 2606
+    },
+    {
+      "epoch": 0.27,
+      "gumbel_tau": 0.5407,
+      "image_masks_loss": 0.0167,
+      "llm_learning_rate": 4.298439145149828e-06,
+      "loss": 0.8723,
+      "output_text_mask_loss": 0.0292,
+      "predictor_learning_rate": 0.00042984391451498274,
+      "step": 2607
+    },
+    {
+      "epoch": 0.27,
+      "gumbel_tau": 0.5406,
+      "image_masks_loss": 0.0125,
+      "llm_learning_rate": 4.297862776733547e-06,
+      "loss": 0.8922,
+      "output_text_mask_loss": 0.0082,
+      "predictor_learning_rate": 0.00042978627767335467,
+      "step": 2608
+    },
+    {
+      "epoch": 0.27,
+      "gumbel_tau": 0.5405,
+      "image_masks_loss": 0.0205,
+      "llm_learning_rate": 4.297286210332612e-06,
+      "loss": 0.8274,
+      "output_text_mask_loss": 0.0364,
+      "predictor_learning_rate": 0.00042972862103326116,
+      "step": 2609
+    },
+    {
+      "epoch": 0.27,
+      "gumbel_tau": 0.5404,
+      "image_masks_loss": 0.0246,
+      "llm_learning_rate": 4.296709446010516e-06,
+      "loss": 0.717,
+      "output_text_mask_loss": 0.0185,
+      "predictor_learning_rate": 0.0004296709446010516,
+      "step": 2610
+    },
+    {
+      "epoch": 0.27,
+      "gumbel_tau": 0.5402,
+      "image_masks_loss": 0.0157,
+      "llm_learning_rate": 4.296132483830774e-06,
+      "loss": 0.8037,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.00042961324838307737,
+      "step": 2611
+    },
+    {
+      "epoch": 0.27,
+      "gumbel_tau": 0.5401,
+      "image_masks_loss": 0.0163,
+      "llm_learning_rate": 4.295555323856923e-06,
+      "loss": 0.6499,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004295555323856922,
+      "step": 2612
+    },
+    {
+      "epoch": 0.27,
+      "gumbel_tau": 0.54,
+      "image_masks_loss": 0.0429,
+      "llm_learning_rate": 4.29497796615252e-06,
+      "loss": 0.7195,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004294977966152519,
+      "step": 2613
+    },
+    {
+      "epoch": 0.27,
+      "gumbel_tau": 0.5399,
+      "image_masks_loss": 0.0121,
+      "llm_learning_rate": 4.2944004107811434e-06,
+      "loss": 0.8137,
+      "output_text_mask_loss": 0.0019,
+      "predictor_learning_rate": 0.0004294400410781143,
+      "step": 2614
+    },
+    {
+      "epoch": 0.27,
+      "gumbel_tau": 0.5397,
+      "image_masks_loss": 0.0139,
+      "llm_learning_rate": 4.293822657806398e-06,
+      "loss": 0.8361,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00042938226578063976,
+      "step": 2615
+    },
+    {
+      "epoch": 0.27,
+      "gumbel_tau": 0.5396,
+      "image_masks_loss": 0.008,
+      "llm_learning_rate": 4.2932447072919065e-06,
+      "loss": 0.6337,
+      "output_text_mask_loss": 0.0565,
+      "predictor_learning_rate": 0.0004293244707291906,
+      "step": 2616
+    },
+    {
+      "epoch": 0.27,
+      "gumbel_tau": 0.5395,
+      "image_masks_loss": 0.0225,
+      "llm_learning_rate": 4.292666559301312e-06,
+      "loss": 0.7763,
+      "output_text_mask_loss": 0.0065,
+      "predictor_learning_rate": 0.00042926665593013126,
+      "step": 2617
+    },
+    {
+      "epoch": 0.27,
+      "gumbel_tau": 0.5393,
+      "image_masks_loss": 0.0225,
+      "llm_learning_rate": 4.292088213898285e-06,
+      "loss": 0.8563,
+      "output_text_mask_loss": 0.0401,
+      "predictor_learning_rate": 0.0004292088213898285,
+      "step": 2618
+    },
+    {
+      "epoch": 0.27,
+      "gumbel_tau": 0.5392,
+      "image_masks_loss": 0.0217,
+      "llm_learning_rate": 4.291509671146511e-06,
+      "loss": 0.8729,
+      "output_text_mask_loss": 0.0286,
+      "predictor_learning_rate": 0.00042915096711465105,
+      "step": 2619
+    },
+    {
+      "epoch": 0.27,
+      "gumbel_tau": 0.5391,
+      "image_masks_loss": 0.0187,
+      "llm_learning_rate": 4.290930931109702e-06,
+      "loss": 0.8085,
+      "output_text_mask_loss": 0.0193,
+      "predictor_learning_rate": 0.0004290930931109702,
+      "step": 2620
+    },
+    {
+      "epoch": 0.27,
+      "gumbel_tau": 0.539,
+      "image_masks_loss": 0.0084,
+      "llm_learning_rate": 4.290351993851589e-06,
+      "loss": 0.7627,
+      "output_text_mask_loss": 0.0324,
+      "predictor_learning_rate": 0.0004290351993851589,
+      "step": 2621
+    },
+    {
+      "epoch": 0.27,
+      "gumbel_tau": 0.5388,
+      "image_masks_loss": 0.0315,
+      "llm_learning_rate": 4.289772859435927e-06,
+      "loss": 0.8133,
+      "output_text_mask_loss": 0.0326,
+      "predictor_learning_rate": 0.0004289772859435927,
+      "step": 2622
+    },
+    {
+      "epoch": 0.27,
+      "gumbel_tau": 0.5387,
+      "image_masks_loss": 0.0148,
+      "llm_learning_rate": 4.289193527926492e-06,
+      "loss": 0.735,
+      "output_text_mask_loss": 0.0782,
+      "predictor_learning_rate": 0.00042891935279264914,
+      "step": 2623
+    },
+    {
+      "epoch": 0.27,
+      "gumbel_tau": 0.5386,
+      "image_masks_loss": 0.0113,
+      "llm_learning_rate": 4.288613999387079e-06,
+      "loss": 0.7874,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00042886139993870786,
+      "step": 2624
+    },
+    {
+      "epoch": 0.27,
+      "gumbel_tau": 0.5385,
+      "image_masks_loss": 0.0307,
+      "llm_learning_rate": 4.288034273881509e-06,
+      "loss": 0.8065,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.0004288034273881509,
+      "step": 2625
+    },
+    {
+      "epoch": 0.27,
+      "gumbel_tau": 0.5383,
+      "image_masks_loss": 0.0175,
+      "llm_learning_rate": 4.287454351473622e-06,
+      "loss": 0.7706,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 0.00042874543514736217,
+      "step": 2626
+    },
+    {
+      "epoch": 0.27,
+      "gumbel_tau": 0.5382,
+      "image_masks_loss": 0.0135,
+      "llm_learning_rate": 4.286874232227281e-06,
+      "loss": 0.785,
+      "output_text_mask_loss": 0.1488,
+      "predictor_learning_rate": 0.0004286874232227281,
+      "step": 2627
+    },
+    {
+      "epoch": 0.27,
+      "gumbel_tau": 0.5381,
+      "image_masks_loss": 0.0128,
+      "llm_learning_rate": 4.286293916206369e-06,
+      "loss": 0.6772,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.00042862939162063683,
+      "step": 2628
+    },
+    {
+      "epoch": 0.27,
+      "gumbel_tau": 0.5379,
+      "image_masks_loss": 0.0146,
+      "llm_learning_rate": 4.285713403474792e-06,
+      "loss": 0.6233,
+      "output_text_mask_loss": 0.0068,
+      "predictor_learning_rate": 0.0004285713403474792,
+      "step": 2629
+    },
+    {
+      "epoch": 0.27,
+      "gumbel_tau": 0.5378,
+      "image_masks_loss": 0.0144,
+      "llm_learning_rate": 4.285132694096478e-06,
+      "loss": 0.7376,
+      "output_text_mask_loss": 0.0141,
+      "predictor_learning_rate": 0.00042851326940964774,
+      "step": 2630
+    },
+    {
+      "epoch": 0.27,
+      "gumbel_tau": 0.5377,
+      "image_masks_loss": 0.0146,
+      "llm_learning_rate": 4.284551788135375e-06,
+      "loss": 0.8157,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.00042845517881353744,
+      "step": 2631
+    },
+    {
+      "epoch": 0.27,
+      "gumbel_tau": 0.5376,
+      "image_masks_loss": 0.0202,
+      "llm_learning_rate": 4.283970685655454e-06,
+      "loss": 0.7392,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 0.0004283970685655454,
+      "step": 2632
+    },
+    {
+      "epoch": 0.27,
+      "gumbel_tau": 0.5374,
+      "image_masks_loss": 0.0248,
+      "llm_learning_rate": 4.283389386720708e-06,
+      "loss": 0.7285,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00042833893867207074,
+      "step": 2633
+    },
+    {
+      "epoch": 0.27,
+      "gumbel_tau": 0.5373,
+      "image_masks_loss": 0.0053,
+      "llm_learning_rate": 4.28280789139515e-06,
+      "loss": 0.793,
+      "output_text_mask_loss": 0.0397,
+      "predictor_learning_rate": 0.000428280789139515,
+      "step": 2634
+    },
+    {
+      "epoch": 0.27,
+      "gumbel_tau": 0.5372,
+      "image_masks_loss": 0.0423,
+      "llm_learning_rate": 4.282226199742816e-06,
+      "loss": 0.781,
+      "output_text_mask_loss": 0.0694,
+      "predictor_learning_rate": 0.00042822261997428157,
+      "step": 2635
+    },
+    {
+      "epoch": 0.27,
+      "gumbel_tau": 0.5371,
+      "image_masks_loss": 0.0066,
+      "llm_learning_rate": 4.281644311827763e-06,
+      "loss": 0.7957,
+      "output_text_mask_loss": 0.1015,
+      "predictor_learning_rate": 0.0004281644311827763,
+      "step": 2636
+    },
+    {
+      "epoch": 0.27,
+      "gumbel_tau": 0.5369,
+      "image_masks_loss": 0.009,
+      "llm_learning_rate": 4.281062227714071e-06,
+      "loss": 0.8104,
+      "output_text_mask_loss": 0.0195,
+      "predictor_learning_rate": 0.00042810622277140704,
+      "step": 2637
+    },
+    {
+      "epoch": 0.27,
+      "gumbel_tau": 0.5368,
+      "image_masks_loss": 0.0083,
+      "llm_learning_rate": 4.280479947465838e-06,
+      "loss": 0.6599,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.00042804799474658375,
+      "step": 2638
+    },
+    {
+      "epoch": 0.27,
+      "gumbel_tau": 0.5367,
+      "image_masks_loss": 0.0078,
+      "llm_learning_rate": 4.279897471147188e-06,
+      "loss": 0.7704,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.00042798974711471876,
+      "step": 2639
+    },
+    {
+      "epoch": 0.27,
+      "gumbel_tau": 0.5366,
+      "image_masks_loss": 0.0317,
+      "llm_learning_rate": 4.2793147988222635e-06,
+      "loss": 0.7916,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.00042793147988222634,
+      "step": 2640
+    },
+    {
+      "epoch": 0.27,
+      "gumbel_tau": 0.5364,
+      "image_masks_loss": 0.0194,
+      "llm_learning_rate": 4.27873193055523e-06,
+      "loss": 0.7053,
+      "output_text_mask_loss": 0.005,
+      "predictor_learning_rate": 0.00042787319305552296,
+      "step": 2641
+    },
+    {
+      "epoch": 0.27,
+      "gumbel_tau": 0.5363,
+      "image_masks_loss": 0.01,
+      "llm_learning_rate": 4.2781488664102746e-06,
+      "loss": 0.8525,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00042781488664102746,
+      "step": 2642
+    },
+    {
+      "epoch": 0.27,
+      "gumbel_tau": 0.5362,
+      "image_masks_loss": 0.0132,
+      "llm_learning_rate": 4.277565606451606e-06,
+      "loss": 0.6717,
+      "output_text_mask_loss": 0.0317,
+      "predictor_learning_rate": 0.00042775656064516056,
+      "step": 2643
+    },
+    {
+      "epoch": 0.27,
+      "gumbel_tau": 0.536,
+      "image_masks_loss": 0.0106,
+      "llm_learning_rate": 4.276982150743452e-06,
+      "loss": 0.7749,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004276982150743453,
+      "step": 2644
+    },
+    {
+      "epoch": 0.27,
+      "gumbel_tau": 0.5359,
+      "image_masks_loss": 0.0109,
+      "llm_learning_rate": 4.276398499350067e-06,
+      "loss": 0.723,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00042763984993500666,
+      "step": 2645
+    },
+    {
+      "epoch": 0.27,
+      "gumbel_tau": 0.5358,
+      "image_masks_loss": 0.0117,
+      "llm_learning_rate": 4.2758146523357215e-06,
+      "loss": 0.6887,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 0.0004275814652335721,
+      "step": 2646
+    },
+    {
+      "epoch": 0.27,
+      "gumbel_tau": 0.5357,
+      "image_masks_loss": 0.0068,
+      "llm_learning_rate": 4.275230609764711e-06,
+      "loss": 0.7745,
+      "output_text_mask_loss": 0.0252,
+      "predictor_learning_rate": 0.00042752306097647107,
+      "step": 2647
+    },
+    {
+      "epoch": 0.27,
+      "gumbel_tau": 0.5355,
+      "image_masks_loss": 0.0152,
+      "llm_learning_rate": 4.274646371701352e-06,
+      "loss": 0.7534,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.0004274646371701352,
+      "step": 2648
+    },
+    {
+      "epoch": 0.27,
+      "gumbel_tau": 0.5354,
+      "image_masks_loss": 0.0326,
+      "llm_learning_rate": 4.274061938209982e-06,
+      "loss": 0.7145,
+      "output_text_mask_loss": 0.0214,
+      "predictor_learning_rate": 0.00042740619382099817,
+      "step": 2649
+    },
+    {
+      "epoch": 0.27,
+      "gumbel_tau": 0.5353,
+      "image_masks_loss": 0.0244,
+      "llm_learning_rate": 4.273477309354959e-06,
+      "loss": 0.7621,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.0004273477309354959,
+      "step": 2650
+    },
+    {
+      "epoch": 0.27,
+      "gumbel_tau": 0.5352,
+      "image_masks_loss": 0.0263,
+      "llm_learning_rate": 4.272892485200665e-06,
+      "loss": 0.7968,
+      "output_text_mask_loss": 0.0572,
+      "predictor_learning_rate": 0.00042728924852006647,
+      "step": 2651
+    },
+    {
+      "epoch": 0.27,
+      "gumbel_tau": 0.535,
+      "image_masks_loss": 0.0286,
+      "llm_learning_rate": 4.272307465811502e-06,
+      "loss": 0.8034,
+      "output_text_mask_loss": 0.0298,
+      "predictor_learning_rate": 0.0004272307465811501,
+      "step": 2652
+    },
+    {
+      "epoch": 0.27,
+      "gumbel_tau": 0.5349,
+      "image_masks_loss": 0.0134,
+      "llm_learning_rate": 4.271722251251892e-06,
+      "loss": 0.8513,
+      "output_text_mask_loss": 0.1389,
+      "predictor_learning_rate": 0.0004271722251251892,
+      "step": 2653
+    },
+    {
+      "epoch": 0.27,
+      "gumbel_tau": 0.5348,
+      "image_masks_loss": 0.0081,
+      "llm_learning_rate": 4.271136841586282e-06,
+      "loss": 0.7666,
+      "output_text_mask_loss": 0.0231,
+      "predictor_learning_rate": 0.00042711368415862815,
+      "step": 2654
+    },
+    {
+      "epoch": 0.27,
+      "gumbel_tau": 0.5347,
+      "image_masks_loss": 0.0311,
+      "llm_learning_rate": 4.270551236879137e-06,
+      "loss": 0.8579,
+      "output_text_mask_loss": 0.0012,
+      "predictor_learning_rate": 0.0004270551236879137,
+      "step": 2655
+    },
+    {
+      "epoch": 0.27,
+      "gumbel_tau": 0.5345,
+      "image_masks_loss": 0.0101,
+      "llm_learning_rate": 4.269965437194948e-06,
+      "loss": 0.7721,
+      "output_text_mask_loss": 0.024,
+      "predictor_learning_rate": 0.0004269965437194947,
+      "step": 2656
+    },
+    {
+      "epoch": 0.27,
+      "gumbel_tau": 0.5344,
+      "image_masks_loss": 0.0458,
+      "llm_learning_rate": 4.26937944259822e-06,
+      "loss": 0.7768,
+      "output_text_mask_loss": 0.0244,
+      "predictor_learning_rate": 0.00042693794425982203,
+      "step": 2657
+    },
+    {
+      "epoch": 0.27,
+      "gumbel_tau": 0.5343,
+      "image_masks_loss": 0.005,
+      "llm_learning_rate": 4.26879325315349e-06,
+      "loss": 0.6837,
+      "output_text_mask_loss": 0.0622,
+      "predictor_learning_rate": 0.0004268793253153489,
+      "step": 2658
+    },
+    {
+      "epoch": 0.27,
+      "gumbel_tau": 0.5342,
+      "image_masks_loss": 0.0151,
+      "llm_learning_rate": 4.268206868925304e-06,
+      "loss": 0.7987,
+      "output_text_mask_loss": 0.0782,
+      "predictor_learning_rate": 0.0004268206868925304,
+      "step": 2659
+    },
+    {
+      "epoch": 0.27,
+      "gumbel_tau": 0.534,
+      "image_masks_loss": 0.0372,
+      "llm_learning_rate": 4.267620289978241e-06,
+      "loss": 0.8025,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.000426762028997824,
+      "step": 2660
+    },
+    {
+      "epoch": 0.27,
+      "gumbel_tau": 0.5339,
+      "image_masks_loss": 0.0201,
+      "llm_learning_rate": 4.267033516376893e-06,
+      "loss": 0.7047,
+      "output_text_mask_loss": 0.0078,
+      "predictor_learning_rate": 0.0004267033516376893,
+      "step": 2661
+    },
+    {
+      "epoch": 0.27,
+      "gumbel_tau": 0.5338,
+      "image_masks_loss": 0.0254,
+      "llm_learning_rate": 4.26644654818588e-06,
+      "loss": 0.8158,
+      "output_text_mask_loss": 0.0271,
+      "predictor_learning_rate": 0.0004266446548185879,
+      "step": 2662
+    },
+    {
+      "epoch": 0.27,
+      "gumbel_tau": 0.5336,
+      "image_masks_loss": 0.0045,
+      "llm_learning_rate": 4.265859385469837e-06,
+      "loss": 0.7492,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.00042658593854698367,
+      "step": 2663
+    },
+    {
+      "epoch": 0.27,
+      "gumbel_tau": 0.5335,
+      "image_masks_loss": 0.0081,
+      "llm_learning_rate": 4.265272028293426e-06,
+      "loss": 0.8354,
+      "output_text_mask_loss": 0.0118,
+      "predictor_learning_rate": 0.0004265272028293425,
+      "step": 2664
+    },
+    {
+      "epoch": 0.27,
+      "gumbel_tau": 0.5334,
+      "image_masks_loss": 0.0406,
+      "llm_learning_rate": 4.2646844767213265e-06,
+      "loss": 0.8599,
+      "output_text_mask_loss": 0.0664,
+      "predictor_learning_rate": 0.00042646844767213265,
+      "step": 2665
+    },
+    {
+      "epoch": 0.27,
+      "gumbel_tau": 0.5333,
+      "image_masks_loss": 0.0178,
+      "llm_learning_rate": 4.2640967308182425e-06,
+      "loss": 0.8356,
+      "output_text_mask_loss": 0.0322,
+      "predictor_learning_rate": 0.0004264096730818242,
+      "step": 2666
+    },
+    {
+      "epoch": 0.27,
+      "gumbel_tau": 0.5331,
+      "image_masks_loss": 0.0219,
+      "llm_learning_rate": 4.263508790648897e-06,
+      "loss": 0.776,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.00042635087906488966,
+      "step": 2667
+    },
+    {
+      "epoch": 0.27,
+      "gumbel_tau": 0.533,
+      "image_masks_loss": 0.0237,
+      "llm_learning_rate": 4.262920656278034e-06,
+      "loss": 0.7938,
+      "output_text_mask_loss": 0.0557,
+      "predictor_learning_rate": 0.0004262920656278034,
+      "step": 2668
+    },
+    {
+      "epoch": 0.27,
+      "gumbel_tau": 0.5329,
+      "image_masks_loss": 0.0082,
+      "llm_learning_rate": 4.262332327770424e-06,
+      "loss": 0.7398,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 0.00042623323277704236,
+      "step": 2669
+    },
+    {
+      "epoch": 0.27,
+      "gumbel_tau": 0.5328,
+      "image_masks_loss": 0.0155,
+      "llm_learning_rate": 4.2617438051908515e-06,
+      "loss": 0.8417,
+      "output_text_mask_loss": 0.0081,
+      "predictor_learning_rate": 0.00042617438051908513,
+      "step": 2670
+    },
+    {
+      "epoch": 0.27,
+      "gumbel_tau": 0.5326,
+      "image_masks_loss": 0.0372,
+      "llm_learning_rate": 4.261155088604127e-06,
+      "loss": 0.6814,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.00042611550886041265,
+      "step": 2671
+    },
+    {
+      "epoch": 0.27,
+      "gumbel_tau": 0.5325,
+      "image_masks_loss": 0.0441,
+      "llm_learning_rate": 4.260566178075081e-06,
+      "loss": 0.777,
+      "output_text_mask_loss": 0.0231,
+      "predictor_learning_rate": 0.0004260566178075081,
+      "step": 2672
+    },
+    {
+      "epoch": 0.27,
+      "gumbel_tau": 0.5324,
+      "image_masks_loss": 0.0154,
+      "llm_learning_rate": 4.2599770736685665e-06,
+      "loss": 0.7834,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004259977073668566,
+      "step": 2673
+    },
+    {
+      "epoch": 0.27,
+      "gumbel_tau": 0.5323,
+      "image_masks_loss": 0.0213,
+      "llm_learning_rate": 4.2593877754494565e-06,
+      "loss": 0.7854,
+      "output_text_mask_loss": 0.0801,
+      "predictor_learning_rate": 0.0004259387775449456,
+      "step": 2674
+    },
+    {
+      "epoch": 0.27,
+      "gumbel_tau": 0.5321,
+      "image_masks_loss": 0.0125,
+      "llm_learning_rate": 4.258798283482645e-06,
+      "loss": 0.8004,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004258798283482645,
+      "step": 2675
+    },
+    {
+      "epoch": 0.27,
+      "gumbel_tau": 0.532,
+      "image_masks_loss": 0.0254,
+      "llm_learning_rate": 4.25820859783305e-06,
+      "loss": 0.8199,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.000425820859783305,
+      "step": 2676
+    },
+    {
+      "epoch": 0.27,
+      "gumbel_tau": 0.5319,
+      "image_masks_loss": 0.0132,
+      "llm_learning_rate": 4.257618718565607e-06,
+      "loss": 0.7455,
+      "output_text_mask_loss": 0.0322,
+      "predictor_learning_rate": 0.00042576187185656067,
+      "step": 2677
+    },
+    {
+      "epoch": 0.27,
+      "gumbel_tau": 0.5318,
+      "image_masks_loss": 0.0294,
+      "llm_learning_rate": 4.257028645745276e-06,
+      "loss": 0.5997,
+      "output_text_mask_loss": 0.0055,
+      "predictor_learning_rate": 0.0004257028645745276,
+      "step": 2678
+    },
+    {
+      "epoch": 0.27,
+      "gumbel_tau": 0.5316,
+      "image_masks_loss": 0.0197,
+      "llm_learning_rate": 4.256438379437038e-06,
+      "loss": 0.8281,
+      "output_text_mask_loss": 0.1068,
+      "predictor_learning_rate": 0.00042564383794370375,
+      "step": 2679
+    },
+    {
+      "epoch": 0.27,
+      "gumbel_tau": 0.5315,
+      "image_masks_loss": 0.012,
+      "llm_learning_rate": 4.2558479197058925e-06,
+      "loss": 0.7981,
+      "output_text_mask_loss": 0.0322,
+      "predictor_learning_rate": 0.0004255847919705892,
+      "step": 2680
+    },
+    {
+      "epoch": 0.27,
+      "gumbel_tau": 0.5314,
+      "image_masks_loss": 0.0237,
+      "llm_learning_rate": 4.255257266616862e-06,
+      "loss": 0.7656,
+      "output_text_mask_loss": 0.0177,
+      "predictor_learning_rate": 0.00042552572666168623,
+      "step": 2681
+    },
+    {
+      "epoch": 0.27,
+      "gumbel_tau": 0.5313,
+      "image_masks_loss": 0.0187,
+      "llm_learning_rate": 4.254666420234993e-06,
+      "loss": 0.8086,
+      "output_text_mask_loss": 0.0687,
+      "predictor_learning_rate": 0.00042546664202349925,
+      "step": 2682
+    },
+    {
+      "epoch": 0.27,
+      "gumbel_tau": 0.5311,
+      "image_masks_loss": 0.0305,
+      "llm_learning_rate": 4.25407538062535e-06,
+      "loss": 0.6811,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004254075380625349,
+      "step": 2683
+    },
+    {
+      "epoch": 0.28,
+      "gumbel_tau": 0.531,
+      "image_masks_loss": 0.0169,
+      "llm_learning_rate": 4.2534841478530175e-06,
+      "loss": 0.8132,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.00042534841478530173,
+      "step": 2684
+    },
+    {
+      "epoch": 0.28,
+      "gumbel_tau": 0.5309,
+      "image_masks_loss": 0.0284,
+      "llm_learning_rate": 4.252892721983105e-06,
+      "loss": 0.8182,
+      "output_text_mask_loss": 0.0101,
+      "predictor_learning_rate": 0.00042528927219831047,
+      "step": 2685
+    },
+    {
+      "epoch": 0.28,
+      "gumbel_tau": 0.5308,
+      "image_masks_loss": 0.0246,
+      "llm_learning_rate": 4.252301103080743e-06,
+      "loss": 0.8233,
+      "output_text_mask_loss": 0.0173,
+      "predictor_learning_rate": 0.0004252301103080742,
+      "step": 2686
+    },
+    {
+      "epoch": 0.28,
+      "gumbel_tau": 0.5306,
+      "image_masks_loss": 0.0065,
+      "llm_learning_rate": 4.251709291211079e-06,
+      "loss": 0.8114,
+      "output_text_mask_loss": 0.1297,
+      "predictor_learning_rate": 0.0004251709291211079,
+      "step": 2687
+    },
+    {
+      "epoch": 0.28,
+      "gumbel_tau": 0.5305,
+      "image_masks_loss": 0.0034,
+      "llm_learning_rate": 4.251117286439287e-06,
+      "loss": 0.8305,
+      "output_text_mask_loss": 0.0145,
+      "predictor_learning_rate": 0.00042511172864392866,
+      "step": 2688
+    },
+    {
+      "epoch": 0.28,
+      "gumbel_tau": 0.5304,
+      "image_masks_loss": 0.0305,
+      "llm_learning_rate": 4.250525088830558e-06,
+      "loss": 0.7819,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.0004250525088830558,
+      "step": 2689
+    },
+    {
+      "epoch": 0.28,
+      "gumbel_tau": 0.5303,
+      "image_masks_loss": 0.029,
+      "llm_learning_rate": 4.249932698450108e-06,
+      "loss": 0.7582,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00042499326984501073,
+      "step": 2690
+    },
+    {
+      "epoch": 0.28,
+      "gumbel_tau": 0.5301,
+      "image_masks_loss": 0.0211,
+      "llm_learning_rate": 4.2493401153631705e-06,
+      "loss": 0.7398,
+      "output_text_mask_loss": 0.016,
+      "predictor_learning_rate": 0.00042493401153631706,
+      "step": 2691
+    },
+    {
+      "epoch": 0.28,
+      "gumbel_tau": 0.53,
+      "image_masks_loss": 0.0267,
+      "llm_learning_rate": 4.248747339635004e-06,
+      "loss": 0.7665,
+      "output_text_mask_loss": 0.1167,
+      "predictor_learning_rate": 0.0004248747339635003,
+      "step": 2692
+    },
+    {
+      "epoch": 0.28,
+      "gumbel_tau": 0.5299,
+      "image_masks_loss": 0.0355,
+      "llm_learning_rate": 4.248154371330885e-06,
+      "loss": 0.7613,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.00042481543713308845,
+      "step": 2693
+    },
+    {
+      "epoch": 0.28,
+      "gumbel_tau": 0.5298,
+      "image_masks_loss": 0.0132,
+      "llm_learning_rate": 4.247561210516113e-06,
+      "loss": 0.779,
+      "output_text_mask_loss": 0.017,
+      "predictor_learning_rate": 0.0004247561210516112,
+      "step": 2694
+    },
+    {
+      "epoch": 0.28,
+      "gumbel_tau": 0.5296,
+      "image_masks_loss": 0.0282,
+      "llm_learning_rate": 4.246967857256007e-06,
+      "loss": 0.8553,
+      "output_text_mask_loss": 0.0904,
+      "predictor_learning_rate": 0.00042469678572560065,
+      "step": 2695
+    },
+    {
+      "epoch": 0.28,
+      "gumbel_tau": 0.5295,
+      "image_masks_loss": 0.0126,
+      "llm_learning_rate": 4.246374311615909e-06,
+      "loss": 0.751,
+      "output_text_mask_loss": 0.0004,
+      "predictor_learning_rate": 0.0004246374311615909,
+      "step": 2696
+    },
+    {
+      "epoch": 0.28,
+      "gumbel_tau": 0.5294,
+      "image_masks_loss": 0.041,
+      "llm_learning_rate": 4.2457805736611836e-06,
+      "loss": 0.8122,
+      "output_text_mask_loss": 0.0675,
+      "predictor_learning_rate": 0.0004245780573661183,
+      "step": 2697
+    },
+    {
+      "epoch": 0.28,
+      "gumbel_tau": 0.5293,
+      "image_masks_loss": 0.0188,
+      "llm_learning_rate": 4.245186643457212e-06,
+      "loss": 0.7919,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00042451866434572115,
+      "step": 2698
+    },
+    {
+      "epoch": 0.28,
+      "gumbel_tau": 0.5291,
+      "image_masks_loss": 0.0492,
+      "llm_learning_rate": 4.2445925210694e-06,
+      "loss": 0.6999,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.00042445925210693995,
+      "step": 2699
+    },
+    {
+      "epoch": 0.28,
+      "gumbel_tau": 0.529,
+      "image_masks_loss": 0.0189,
+      "llm_learning_rate": 4.2439982065631736e-06,
+      "loss": 0.7215,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004243998206563173,
+      "step": 2700
+    },
+    {
+      "epoch": 0.28,
+      "gumbel_tau": 0.5289,
+      "image_masks_loss": 0.0176,
+      "llm_learning_rate": 4.24340370000398e-06,
+      "loss": 0.8383,
+      "output_text_mask_loss": 0.0973,
+      "predictor_learning_rate": 0.000424340370000398,
+      "step": 2701
+    },
+    {
+      "epoch": 0.28,
+      "gumbel_tau": 0.5288,
+      "image_masks_loss": 0.0248,
+      "llm_learning_rate": 4.242809001457288e-06,
+      "loss": 0.7833,
+      "output_text_mask_loss": 0.0134,
+      "predictor_learning_rate": 0.00042428090014572876,
+      "step": 2702
+    },
+    {
+      "epoch": 0.28,
+      "gumbel_tau": 0.5286,
+      "image_masks_loss": 0.014,
+      "llm_learning_rate": 4.242214110988586e-06,
+      "loss": 0.8243,
+      "output_text_mask_loss": 0.0488,
+      "predictor_learning_rate": 0.0004242214110988586,
+      "step": 2703
+    },
+    {
+      "epoch": 0.28,
+      "gumbel_tau": 0.5285,
+      "image_masks_loss": 0.0271,
+      "llm_learning_rate": 4.241619028663386e-06,
+      "loss": 0.6905,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 0.0004241619028663386,
+      "step": 2704
+    },
+    {
+      "epoch": 0.28,
+      "gumbel_tau": 0.5284,
+      "image_masks_loss": 0.0177,
+      "llm_learning_rate": 4.24102375454722e-06,
+      "loss": 0.7554,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.000424102375454722,
+      "step": 2705
+    },
+    {
+      "epoch": 0.28,
+      "gumbel_tau": 0.5283,
+      "image_masks_loss": 0.0067,
+      "llm_learning_rate": 4.240428288705639e-06,
+      "loss": 0.7758,
+      "output_text_mask_loss": 0.0851,
+      "predictor_learning_rate": 0.0004240428288705639,
+      "step": 2706
+    },
+    {
+      "epoch": 0.28,
+      "gumbel_tau": 0.5281,
+      "image_masks_loss": 0.037,
+      "llm_learning_rate": 4.23983263120422e-06,
+      "loss": 0.7661,
+      "output_text_mask_loss": 0.0269,
+      "predictor_learning_rate": 0.0004239832631204219,
+      "step": 2707
+    },
+    {
+      "epoch": 0.28,
+      "gumbel_tau": 0.528,
+      "image_masks_loss": 0.0265,
+      "llm_learning_rate": 4.239236782108555e-06,
+      "loss": 0.7824,
+      "output_text_mask_loss": 0.0324,
+      "predictor_learning_rate": 0.0004239236782108554,
+      "step": 2708
+    },
+    {
+      "epoch": 0.28,
+      "gumbel_tau": 0.5279,
+      "image_masks_loss": 0.0104,
+      "llm_learning_rate": 4.238640741484261e-06,
+      "loss": 0.8305,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 0.0004238640741484261,
+      "step": 2709
+    },
+    {
+      "epoch": 0.28,
+      "gumbel_tau": 0.5278,
+      "image_masks_loss": 0.018,
+      "llm_learning_rate": 4.238044509396977e-06,
+      "loss": 0.8352,
+      "output_text_mask_loss": 0.0645,
+      "predictor_learning_rate": 0.0004238044509396976,
+      "step": 2710
+    },
+    {
+      "epoch": 0.28,
+      "gumbel_tau": 0.5276,
+      "image_masks_loss": 0.0095,
+      "llm_learning_rate": 4.237448085912359e-06,
+      "loss": 0.7011,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004237448085912359,
+      "step": 2711
+    },
+    {
+      "epoch": 0.28,
+      "gumbel_tau": 0.5275,
+      "image_masks_loss": 0.0135,
+      "llm_learning_rate": 4.236851471096088e-06,
+      "loss": 0.7324,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 0.0004236851471096088,
+      "step": 2712
+    },
+    {
+      "epoch": 0.28,
+      "gumbel_tau": 0.5274,
+      "image_masks_loss": 0.0233,
+      "llm_learning_rate": 4.236254665013866e-06,
+      "loss": 0.7911,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004236254665013865,
+      "step": 2713
+    },
+    {
+      "epoch": 0.28,
+      "gumbel_tau": 0.5273,
+      "image_masks_loss": 0.0254,
+      "llm_learning_rate": 4.235657667731411e-06,
+      "loss": 0.7545,
+      "output_text_mask_loss": 0.0216,
+      "predictor_learning_rate": 0.00042356576677314106,
+      "step": 2714
+    },
+    {
+      "epoch": 0.28,
+      "gumbel_tau": 0.5271,
+      "image_masks_loss": 0.0296,
+      "llm_learning_rate": 4.235060479314468e-06,
+      "loss": 0.7545,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004235060479314468,
+      "step": 2715
+    },
+    {
+      "epoch": 0.28,
+      "gumbel_tau": 0.527,
+      "image_masks_loss": 0.0441,
+      "llm_learning_rate": 4.234463099828801e-06,
+      "loss": 0.8572,
+      "output_text_mask_loss": 0.1152,
+      "predictor_learning_rate": 0.00042344630998288014,
+      "step": 2716
+    },
+    {
+      "epoch": 0.28,
+      "gumbel_tau": 0.5269,
+      "image_masks_loss": 0.01,
+      "llm_learning_rate": 4.233865529340194e-06,
+      "loss": 0.6845,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004233865529340194,
+      "step": 2717
+    },
+    {
+      "epoch": 0.28,
+      "gumbel_tau": 0.5268,
+      "image_masks_loss": 0.008,
+      "llm_learning_rate": 4.233267767914453e-06,
+      "loss": 0.6664,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00042332677679144526,
+      "step": 2718
+    },
+    {
+      "epoch": 0.28,
+      "gumbel_tau": 0.5266,
+      "image_masks_loss": 0.0336,
+      "llm_learning_rate": 4.232669815617404e-06,
+      "loss": 0.8201,
+      "output_text_mask_loss": 0.0477,
+      "predictor_learning_rate": 0.00042326698156174036,
+      "step": 2719
+    },
+    {
+      "epoch": 0.28,
+      "gumbel_tau": 0.5265,
+      "image_masks_loss": 0.0206,
+      "llm_learning_rate": 4.2320716725148955e-06,
+      "loss": 0.7851,
+      "output_text_mask_loss": 0.0347,
+      "predictor_learning_rate": 0.00042320716725148956,
+      "step": 2720
+    },
+    {
+      "epoch": 0.28,
+      "gumbel_tau": 0.5264,
+      "image_masks_loss": 0.033,
+      "llm_learning_rate": 4.231473338672797e-06,
+      "loss": 0.7768,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004231473338672797,
+      "step": 2721
+    },
+    {
+      "epoch": 0.28,
+      "gumbel_tau": 0.5263,
+      "image_masks_loss": 0.0201,
+      "llm_learning_rate": 4.2308748141569975e-06,
+      "loss": 0.745,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.00042308748141569973,
+      "step": 2722
+    },
+    {
+      "epoch": 0.28,
+      "gumbel_tau": 0.5261,
+      "image_masks_loss": 0.0077,
+      "llm_learning_rate": 4.2302760990334084e-06,
+      "loss": 0.7062,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004230276099033408,
+      "step": 2723
+    },
+    {
+      "epoch": 0.28,
+      "gumbel_tau": 0.526,
+      "image_masks_loss": 0.0186,
+      "llm_learning_rate": 4.2296771933679605e-06,
+      "loss": 0.8187,
+      "output_text_mask_loss": 0.0019,
+      "predictor_learning_rate": 0.0004229677193367961,
+      "step": 2724
+    },
+    {
+      "epoch": 0.28,
+      "gumbel_tau": 0.5259,
+      "image_masks_loss": 0.0389,
+      "llm_learning_rate": 4.229078097226609e-06,
+      "loss": 0.779,
+      "output_text_mask_loss": 0.0488,
+      "predictor_learning_rate": 0.00042290780972266086,
+      "step": 2725
+    },
+    {
+      "epoch": 0.28,
+      "gumbel_tau": 0.5258,
+      "image_masks_loss": 0.015,
+      "llm_learning_rate": 4.228478810675324e-06,
+      "loss": 0.6457,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.0004228478810675324,
+      "step": 2726
+    },
+    {
+      "epoch": 0.28,
+      "gumbel_tau": 0.5257,
+      "image_masks_loss": 0.0226,
+      "llm_learning_rate": 4.227879333780103e-06,
+      "loss": 0.8346,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.0004227879333780103,
+      "step": 2727
+    },
+    {
+      "epoch": 0.28,
+      "gumbel_tau": 0.5255,
+      "image_masks_loss": 0.0159,
+      "llm_learning_rate": 4.227279666606962e-06,
+      "loss": 0.7183,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004227279666606961,
+      "step": 2728
+    },
+    {
+      "epoch": 0.28,
+      "gumbel_tau": 0.5254,
+      "image_masks_loss": 0.0156,
+      "llm_learning_rate": 4.226679809221936e-06,
+      "loss": 0.7651,
+      "output_text_mask_loss": 0.0231,
+      "predictor_learning_rate": 0.00042266798092219357,
+      "step": 2729
+    },
+    {
+      "epoch": 0.28,
+      "gumbel_tau": 0.5253,
+      "image_masks_loss": 0.0261,
+      "llm_learning_rate": 4.226079761691083e-06,
+      "loss": 0.8003,
+      "output_text_mask_loss": 0.0824,
+      "predictor_learning_rate": 0.0004226079761691083,
+      "step": 2730
+    },
+    {
+      "epoch": 0.28,
+      "gumbel_tau": 0.5252,
+      "image_masks_loss": 0.0338,
+      "llm_learning_rate": 4.225479524080483e-06,
+      "loss": 0.8298,
+      "output_text_mask_loss": 0.0599,
+      "predictor_learning_rate": 0.0004225479524080482,
+      "step": 2731
+    },
+    {
+      "epoch": 0.28,
+      "gumbel_tau": 0.525,
+      "image_masks_loss": 0.019,
+      "llm_learning_rate": 4.2248790964562335e-06,
+      "loss": 0.7641,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00042248790964562335,
+      "step": 2732
+    },
+    {
+      "epoch": 0.28,
+      "gumbel_tau": 0.5249,
+      "image_masks_loss": 0.0143,
+      "llm_learning_rate": 4.224278478884456e-06,
+      "loss": 0.7336,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00042242784788844556,
+      "step": 2733
+    },
+    {
+      "epoch": 0.28,
+      "gumbel_tau": 0.5248,
+      "image_masks_loss": 0.0063,
+      "llm_learning_rate": 4.223677671431292e-06,
+      "loss": 0.759,
+      "output_text_mask_loss": 0.0039,
+      "predictor_learning_rate": 0.0004223677671431292,
+      "step": 2734
+    },
+    {
+      "epoch": 0.28,
+      "gumbel_tau": 0.5247,
+      "image_masks_loss": 0.0148,
+      "llm_learning_rate": 4.223076674162903e-06,
+      "loss": 0.8812,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.0004223076674162903,
+      "step": 2735
+    },
+    {
+      "epoch": 0.28,
+      "gumbel_tau": 0.5245,
+      "image_masks_loss": 0.0191,
+      "llm_learning_rate": 4.222475487145473e-06,
+      "loss": 0.5929,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.00042224754871454727,
+      "step": 2736
+    },
+    {
+      "epoch": 0.28,
+      "gumbel_tau": 0.5244,
+      "image_masks_loss": 0.0239,
+      "llm_learning_rate": 4.221874110445206e-06,
+      "loss": 0.7863,
+      "output_text_mask_loss": 0.0021,
+      "predictor_learning_rate": 0.00042218741104452055,
+      "step": 2737
+    },
+    {
+      "epoch": 0.28,
+      "gumbel_tau": 0.5243,
+      "image_masks_loss": 0.0103,
+      "llm_learning_rate": 4.221272544128326e-06,
+      "loss": 0.8062,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.0004221272544128326,
+      "step": 2738
+    },
+    {
+      "epoch": 0.28,
+      "gumbel_tau": 0.5242,
+      "image_masks_loss": 0.0174,
+      "llm_learning_rate": 4.22067078826108e-06,
+      "loss": 0.7368,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 0.00042206707882610795,
+      "step": 2739
+    },
+    {
+      "epoch": 0.28,
+      "gumbel_tau": 0.524,
+      "image_masks_loss": 0.0088,
+      "llm_learning_rate": 4.2200688429097345e-06,
+      "loss": 0.7293,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004220068842909734,
+      "step": 2740
+    },
+    {
+      "epoch": 0.28,
+      "gumbel_tau": 0.5239,
+      "image_masks_loss": 0.0269,
+      "llm_learning_rate": 4.219466708140576e-06,
+      "loss": 0.7454,
+      "output_text_mask_loss": 0.0504,
+      "predictor_learning_rate": 0.0004219466708140576,
+      "step": 2741
+    },
+    {
+      "epoch": 0.28,
+      "gumbel_tau": 0.5238,
+      "image_masks_loss": 0.0227,
+      "llm_learning_rate": 4.218864384019914e-06,
+      "loss": 0.7729,
+      "output_text_mask_loss": 0.0175,
+      "predictor_learning_rate": 0.0004218864384019913,
+      "step": 2742
+    },
+    {
+      "epoch": 0.28,
+      "gumbel_tau": 0.5237,
+      "image_masks_loss": 0.0395,
+      "llm_learning_rate": 4.218261870614077e-06,
+      "loss": 0.7401,
+      "output_text_mask_loss": 0.0072,
+      "predictor_learning_rate": 0.00042182618706140763,
+      "step": 2743
+    },
+    {
+      "epoch": 0.28,
+      "gumbel_tau": 0.5235,
+      "image_masks_loss": 0.0193,
+      "llm_learning_rate": 4.217659167989416e-06,
+      "loss": 0.7595,
+      "output_text_mask_loss": 0.0107,
+      "predictor_learning_rate": 0.00042176591679894156,
+      "step": 2744
+    },
+    {
+      "epoch": 0.28,
+      "gumbel_tau": 0.5234,
+      "image_masks_loss": 0.0263,
+      "llm_learning_rate": 4.2170562762123e-06,
+      "loss": 0.74,
+      "output_text_mask_loss": 0.0322,
+      "predictor_learning_rate": 0.00042170562762123,
+      "step": 2745
+    },
+    {
+      "epoch": 0.28,
+      "gumbel_tau": 0.5233,
+      "image_masks_loss": 0.0144,
+      "llm_learning_rate": 4.216453195349124e-06,
+      "loss": 0.7493,
+      "output_text_mask_loss": 0.0172,
+      "predictor_learning_rate": 0.0004216453195349124,
+      "step": 2746
+    },
+    {
+      "epoch": 0.28,
+      "gumbel_tau": 0.5232,
+      "image_masks_loss": 0.0231,
+      "llm_learning_rate": 4.2158499254662985e-06,
+      "loss": 0.809,
+      "output_text_mask_loss": 0.0126,
+      "predictor_learning_rate": 0.00042158499254662987,
+      "step": 2747
+    },
+    {
+      "epoch": 0.28,
+      "gumbel_tau": 0.5231,
+      "image_masks_loss": 0.0103,
+      "llm_learning_rate": 4.215246466630257e-06,
+      "loss": 0.6335,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.0004215246466630257,
+      "step": 2748
+    },
+    {
+      "epoch": 0.28,
+      "gumbel_tau": 0.5229,
+      "image_masks_loss": 0.025,
+      "llm_learning_rate": 4.214642818907455e-06,
+      "loss": 0.8814,
+      "output_text_mask_loss": 0.0117,
+      "predictor_learning_rate": 0.0004214642818907454,
+      "step": 2749
+    },
+    {
+      "epoch": 0.28,
+      "gumbel_tau": 0.5228,
+      "image_masks_loss": 0.0326,
+      "llm_learning_rate": 4.214038982364365e-06,
+      "loss": 0.8012,
+      "output_text_mask_loss": 0.0111,
+      "predictor_learning_rate": 0.0004214038982364365,
+      "step": 2750
+    },
+    {
+      "epoch": 0.28,
+      "gumbel_tau": 0.5227,
+      "image_masks_loss": 0.0244,
+      "llm_learning_rate": 4.2134349570674855e-06,
+      "loss": 0.8327,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.0004213434957067485,
+      "step": 2751
+    },
+    {
+      "epoch": 0.28,
+      "gumbel_tau": 0.5226,
+      "image_masks_loss": 0.0181,
+      "llm_learning_rate": 4.212830743083332e-06,
+      "loss": 0.8339,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.00042128307430833314,
+      "step": 2752
+    },
+    {
+      "epoch": 0.28,
+      "gumbel_tau": 0.5224,
+      "image_masks_loss": 0.0603,
+      "llm_learning_rate": 4.212226340478441e-06,
+      "loss": 0.7954,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00042122263404784414,
+      "step": 2753
+    },
+    {
+      "epoch": 0.28,
+      "gumbel_tau": 0.5223,
+      "image_masks_loss": 0.032,
+      "llm_learning_rate": 4.2116217493193726e-06,
+      "loss": 0.7363,
+      "output_text_mask_loss": 0.0603,
+      "predictor_learning_rate": 0.0004211621749319372,
+      "step": 2754
+    },
+    {
+      "epoch": 0.28,
+      "gumbel_tau": 0.5222,
+      "image_masks_loss": 0.029,
+      "llm_learning_rate": 4.211016969672704e-06,
+      "loss": 0.8294,
+      "output_text_mask_loss": 0.1106,
+      "predictor_learning_rate": 0.0004211016969672704,
+      "step": 2755
+    },
+    {
+      "epoch": 0.28,
+      "gumbel_tau": 0.5221,
+      "image_masks_loss": 0.029,
+      "llm_learning_rate": 4.210412001605035e-06,
+      "loss": 0.6313,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004210412001605035,
+      "step": 2756
+    },
+    {
+      "epoch": 0.28,
+      "gumbel_tau": 0.5219,
+      "image_masks_loss": 0.0075,
+      "llm_learning_rate": 4.209806845182988e-06,
+      "loss": 0.7584,
+      "output_text_mask_loss": 0.0591,
+      "predictor_learning_rate": 0.00042098068451829876,
+      "step": 2757
+    },
+    {
+      "epoch": 0.28,
+      "gumbel_tau": 0.5218,
+      "image_masks_loss": 0.0019,
+      "llm_learning_rate": 4.209201500473202e-06,
+      "loss": 0.7811,
+      "output_text_mask_loss": 0.013,
+      "predictor_learning_rate": 0.0004209201500473202,
+      "step": 2758
+    },
+    {
+      "epoch": 0.28,
+      "gumbel_tau": 0.5217,
+      "image_masks_loss": 0.0124,
+      "llm_learning_rate": 4.20859596754234e-06,
+      "loss": 0.7194,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.00042085959675423396,
+      "step": 2759
+    },
+    {
+      "epoch": 0.28,
+      "gumbel_tau": 0.5216,
+      "image_masks_loss": 0.0267,
+      "llm_learning_rate": 4.207990246457084e-06,
+      "loss": 0.8282,
+      "output_text_mask_loss": 0.0444,
+      "predictor_learning_rate": 0.00042079902464570835,
+      "step": 2760
+    },
+    {
+      "epoch": 0.28,
+      "gumbel_tau": 0.5215,
+      "image_masks_loss": 0.0081,
+      "llm_learning_rate": 4.207384337284138e-06,
+      "loss": 0.8073,
+      "output_text_mask_loss": 0.0326,
+      "predictor_learning_rate": 0.00042073843372841374,
+      "step": 2761
+    },
+    {
+      "epoch": 0.28,
+      "gumbel_tau": 0.5213,
+      "image_masks_loss": 0.0224,
+      "llm_learning_rate": 4.206778240090226e-06,
+      "loss": 0.8124,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.0004206778240090225,
+      "step": 2762
+    },
+    {
+      "epoch": 0.28,
+      "gumbel_tau": 0.5212,
+      "image_masks_loss": 0.0066,
+      "llm_learning_rate": 4.2061719549420906e-06,
+      "loss": 0.7778,
+      "output_text_mask_loss": 0.0078,
+      "predictor_learning_rate": 0.00042061719549420907,
+      "step": 2763
+    },
+    {
+      "epoch": 0.28,
+      "gumbel_tau": 0.5211,
+      "image_masks_loss": 0.0178,
+      "llm_learning_rate": 4.205565481906501e-06,
+      "loss": 0.7538,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 0.0004205565481906501,
+      "step": 2764
+    },
+    {
+      "epoch": 0.28,
+      "gumbel_tau": 0.521,
+      "image_masks_loss": 0.0064,
+      "llm_learning_rate": 4.204958821050241e-06,
+      "loss": 0.8074,
+      "output_text_mask_loss": 0.1099,
+      "predictor_learning_rate": 0.0004204958821050241,
+      "step": 2765
+    },
+    {
+      "epoch": 0.28,
+      "gumbel_tau": 0.5208,
+      "image_masks_loss": 0.0206,
+      "llm_learning_rate": 4.204351972440118e-06,
+      "loss": 0.8712,
+      "output_text_mask_loss": 0.1312,
+      "predictor_learning_rate": 0.0004204351972440118,
+      "step": 2766
+    },
+    {
+      "epoch": 0.28,
+      "gumbel_tau": 0.5207,
+      "image_masks_loss": 0.0083,
+      "llm_learning_rate": 4.203744936142959e-06,
+      "loss": 0.7375,
+      "output_text_mask_loss": 0.021,
+      "predictor_learning_rate": 0.0004203744936142959,
+      "step": 2767
+    },
+    {
+      "epoch": 0.28,
+      "gumbel_tau": 0.5206,
+      "image_masks_loss": 0.0282,
+      "llm_learning_rate": 4.203137712225614e-06,
+      "loss": 0.736,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00042031377122256136,
+      "step": 2768
+    },
+    {
+      "epoch": 0.28,
+      "gumbel_tau": 0.5205,
+      "image_masks_loss": 0.0094,
+      "llm_learning_rate": 4.20253030075495e-06,
+      "loss": 0.8735,
+      "output_text_mask_loss": 0.1289,
+      "predictor_learning_rate": 0.000420253030075495,
+      "step": 2769
+    },
+    {
+      "epoch": 0.28,
+      "gumbel_tau": 0.5203,
+      "image_masks_loss": 0.0313,
+      "llm_learning_rate": 4.201922701797856e-06,
+      "loss": 0.7164,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004201922701797856,
+      "step": 2770
+    },
+    {
+      "epoch": 0.28,
+      "gumbel_tau": 0.5202,
+      "image_masks_loss": 0.0246,
+      "llm_learning_rate": 4.2013149154212455e-06,
+      "loss": 0.6592,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.0004201314915421245,
+      "step": 2771
+    },
+    {
+      "epoch": 0.28,
+      "gumbel_tau": 0.5201,
+      "image_masks_loss": 0.0111,
+      "llm_learning_rate": 4.2007069416920445e-06,
+      "loss": 0.7696,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.00042007069416920447,
+      "step": 2772
+    },
+    {
+      "epoch": 0.28,
+      "gumbel_tau": 0.52,
+      "image_masks_loss": 0.0174,
+      "llm_learning_rate": 4.200098780677208e-06,
+      "loss": 0.8141,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.0004200098780677208,
+      "step": 2773
+    },
+    {
+      "epoch": 0.28,
+      "gumbel_tau": 0.5199,
+      "image_masks_loss": 0.0153,
+      "llm_learning_rate": 4.1994904324437066e-06,
+      "loss": 0.711,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004199490432443707,
+      "step": 2774
+    },
+    {
+      "epoch": 0.28,
+      "gumbel_tau": 0.5197,
+      "image_masks_loss": 0.0155,
+      "llm_learning_rate": 4.198881897058535e-06,
+      "loss": 0.7172,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.0004198881897058534,
+      "step": 2775
+    },
+    {
+      "epoch": 0.28,
+      "gumbel_tau": 0.5196,
+      "image_masks_loss": 0.0099,
+      "llm_learning_rate": 4.198273174588703e-06,
+      "loss": 0.7424,
+      "output_text_mask_loss": 0.0284,
+      "predictor_learning_rate": 0.00041982731745887023,
+      "step": 2776
+    },
+    {
+      "epoch": 0.28,
+      "gumbel_tau": 0.5195,
+      "image_masks_loss": 0.0089,
+      "llm_learning_rate": 4.197664265101246e-06,
+      "loss": 0.7079,
+      "output_text_mask_loss": 0.1366,
+      "predictor_learning_rate": 0.00041976642651012455,
+      "step": 2777
+    },
+    {
+      "epoch": 0.28,
+      "gumbel_tau": 0.5194,
+      "image_masks_loss": 0.0123,
+      "llm_learning_rate": 4.197055168663219e-06,
+      "loss": 0.72,
+      "output_text_mask_loss": 0.0231,
+      "predictor_learning_rate": 0.00041970551686632186,
+      "step": 2778
+    },
+    {
+      "epoch": 0.28,
+      "gumbel_tau": 0.5192,
+      "image_masks_loss": 0.0234,
+      "llm_learning_rate": 4.196445885341697e-06,
+      "loss": 0.7521,
+      "output_text_mask_loss": 0.0233,
+      "predictor_learning_rate": 0.0004196445885341697,
+      "step": 2779
+    },
+    {
+      "epoch": 0.28,
+      "gumbel_tau": 0.5191,
+      "image_masks_loss": 0.0146,
+      "llm_learning_rate": 4.195836415203776e-06,
+      "loss": 0.8167,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.0004195836415203775,
+      "step": 2780
+    },
+    {
+      "epoch": 0.28,
+      "gumbel_tau": 0.519,
+      "image_masks_loss": 0.0083,
+      "llm_learning_rate": 4.195226758316571e-06,
+      "loss": 0.7032,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 0.0004195226758316571,
+      "step": 2781
+    },
+    {
+      "epoch": 0.29,
+      "gumbel_tau": 0.5189,
+      "image_masks_loss": 0.0134,
+      "llm_learning_rate": 4.19461691474722e-06,
+      "loss": 0.6845,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.000419461691474722,
+      "step": 2782
+    },
+    {
+      "epoch": 0.29,
+      "gumbel_tau": 0.5188,
+      "image_masks_loss": 0.0112,
+      "llm_learning_rate": 4.194006884562879e-06,
+      "loss": 0.7591,
+      "output_text_mask_loss": 0.0004,
+      "predictor_learning_rate": 0.00041940068845628786,
+      "step": 2783
+    },
+    {
+      "epoch": 0.29,
+      "gumbel_tau": 0.5186,
+      "image_masks_loss": 0.01,
+      "llm_learning_rate": 4.193396667830727e-06,
+      "loss": 0.8058,
+      "output_text_mask_loss": 0.0353,
+      "predictor_learning_rate": 0.0004193396667830727,
+      "step": 2784
+    },
+    {
+      "epoch": 0.29,
+      "gumbel_tau": 0.5185,
+      "image_masks_loss": 0.0031,
+      "llm_learning_rate": 4.192786264617961e-06,
+      "loss": 0.752,
+      "output_text_mask_loss": 0.0082,
+      "predictor_learning_rate": 0.00041927862646179615,
+      "step": 2785
+    },
+    {
+      "epoch": 0.29,
+      "gumbel_tau": 0.5184,
+      "image_masks_loss": 0.0234,
+      "llm_learning_rate": 4.192175674991803e-06,
+      "loss": 0.8123,
+      "output_text_mask_loss": 0.0771,
+      "predictor_learning_rate": 0.0004192175674991803,
+      "step": 2786
+    },
+    {
+      "epoch": 0.29,
+      "gumbel_tau": 0.5183,
+      "image_masks_loss": 0.0246,
+      "llm_learning_rate": 4.191564899019489e-06,
+      "loss": 0.7311,
+      "output_text_mask_loss": 0.0683,
+      "predictor_learning_rate": 0.0004191564899019489,
+      "step": 2787
+    },
+    {
+      "epoch": 0.29,
+      "gumbel_tau": 0.5181,
+      "image_masks_loss": 0.0406,
+      "llm_learning_rate": 4.190953936768282e-06,
+      "loss": 0.6822,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00041909539367682814,
+      "step": 2788
+    },
+    {
+      "epoch": 0.29,
+      "gumbel_tau": 0.518,
+      "image_masks_loss": 0.0112,
+      "llm_learning_rate": 4.190342788305459e-06,
+      "loss": 0.7814,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.0004190342788305459,
+      "step": 2789
+    },
+    {
+      "epoch": 0.29,
+      "gumbel_tau": 0.5179,
+      "image_masks_loss": 0.0237,
+      "llm_learning_rate": 4.189731453698325e-06,
+      "loss": 0.7637,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 0.00041897314536983244,
+      "step": 2790
+    },
+    {
+      "epoch": 0.29,
+      "gumbel_tau": 0.5178,
+      "image_masks_loss": 0.0273,
+      "llm_learning_rate": 4.1891199330141984e-06,
+      "loss": 0.8707,
+      "output_text_mask_loss": 0.0202,
+      "predictor_learning_rate": 0.00041891199330141985,
+      "step": 2791
+    },
+    {
+      "epoch": 0.29,
+      "gumbel_tau": 0.5177,
+      "image_masks_loss": 0.0145,
+      "llm_learning_rate": 4.188508226320422e-06,
+      "loss": 0.7202,
+      "output_text_mask_loss": 0.0118,
+      "predictor_learning_rate": 0.0004188508226320422,
+      "step": 2792
+    },
+    {
+      "epoch": 0.29,
+      "gumbel_tau": 0.5175,
+      "image_masks_loss": 0.0236,
+      "llm_learning_rate": 4.187896333684361e-06,
+      "loss": 0.8939,
+      "output_text_mask_loss": 0.0105,
+      "predictor_learning_rate": 0.000418789633368436,
+      "step": 2793
+    },
+    {
+      "epoch": 0.29,
+      "gumbel_tau": 0.5174,
+      "image_masks_loss": 0.0362,
+      "llm_learning_rate": 4.1872842551733934e-06,
+      "loss": 0.7807,
+      "output_text_mask_loss": 0.0488,
+      "predictor_learning_rate": 0.0004187284255173393,
+      "step": 2794
+    },
+    {
+      "epoch": 0.29,
+      "gumbel_tau": 0.5173,
+      "image_masks_loss": 0.0099,
+      "llm_learning_rate": 4.186671990854927e-06,
+      "loss": 0.7232,
+      "output_text_mask_loss": 0.0439,
+      "predictor_learning_rate": 0.0004186671990854926,
+      "step": 2795
+    },
+    {
+      "epoch": 0.29,
+      "gumbel_tau": 0.5172,
+      "image_masks_loss": 0.0099,
+      "llm_learning_rate": 4.186059540796383e-06,
+      "loss": 0.7326,
+      "output_text_mask_loss": 0.0039,
+      "predictor_learning_rate": 0.00041860595407963827,
+      "step": 2796
+    },
+    {
+      "epoch": 0.29,
+      "gumbel_tau": 0.517,
+      "image_masks_loss": 0.0246,
+      "llm_learning_rate": 4.185446905065207e-06,
+      "loss": 0.7083,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 0.00041854469050652064,
+      "step": 2797
+    },
+    {
+      "epoch": 0.29,
+      "gumbel_tau": 0.5169,
+      "image_masks_loss": 0.0259,
+      "llm_learning_rate": 4.184834083728864e-06,
+      "loss": 0.7084,
+      "output_text_mask_loss": NaN,
+      "predictor_learning_rate": 0.0004184834083728863,
+      "step": 2798
+    },
+    {
+      "epoch": 0.29,
+      "gumbel_tau": 0.5168,
+      "image_masks_loss": 0.0086,
+      "llm_learning_rate": 4.1842210768548395e-06,
+      "loss": 0.739,
+      "output_text_mask_loss": 0.0343,
+      "predictor_learning_rate": 0.0004184221076854839,
+      "step": 2799
+    },
+    {
+      "epoch": 0.29,
+      "gumbel_tau": 0.5167,
+      "image_masks_loss": 0.0038,
+      "llm_learning_rate": 4.1836078845106375e-06,
+      "loss": 0.8299,
+      "output_text_mask_loss": 0.0189,
+      "predictor_learning_rate": 0.0004183607884510637,
+      "step": 2800
+    },
+    {
+      "epoch": 0.29,
+      "gumbel_tau": 0.5166,
+      "image_masks_loss": 0.0191,
+      "llm_learning_rate": 4.1829945067637855e-06,
+      "loss": 0.6598,
+      "output_text_mask_loss": 0.0248,
+      "predictor_learning_rate": 0.00041829945067637855,
+      "step": 2801
+    },
+    {
+      "epoch": 0.29,
+      "gumbel_tau": 0.5164,
+      "image_masks_loss": 0.0284,
+      "llm_learning_rate": 4.182380943681829e-06,
+      "loss": 0.8442,
+      "output_text_mask_loss": 0.0603,
+      "predictor_learning_rate": 0.00041823809436818295,
+      "step": 2802
+    },
+    {
+      "epoch": 0.29,
+      "gumbel_tau": 0.5163,
+      "image_masks_loss": 0.0206,
+      "llm_learning_rate": 4.181767195332337e-06,
+      "loss": 0.7624,
+      "output_text_mask_loss": 0.0877,
+      "predictor_learning_rate": 0.00041817671953323364,
+      "step": 2803
+    },
+    {
+      "epoch": 0.29,
+      "gumbel_tau": 0.5162,
+      "image_masks_loss": 0.0216,
+      "llm_learning_rate": 4.181153261782894e-06,
+      "loss": 0.7971,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.0004181153261782894,
+      "step": 2804
+    },
+    {
+      "epoch": 0.29,
+      "gumbel_tau": 0.5161,
+      "image_masks_loss": 0.0284,
+      "llm_learning_rate": 4.1805391431011114e-06,
+      "loss": 0.7152,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.00041805391431011106,
+      "step": 2805
+    },
+    {
+      "epoch": 0.29,
+      "gumbel_tau": 0.5159,
+      "image_masks_loss": 0.0098,
+      "llm_learning_rate": 4.179924839354613e-06,
+      "loss": 0.7481,
+      "output_text_mask_loss": 0.029,
+      "predictor_learning_rate": 0.0004179924839354613,
+      "step": 2806
+    },
+    {
+      "epoch": 0.29,
+      "gumbel_tau": 0.5158,
+      "image_masks_loss": 0.0148,
+      "llm_learning_rate": 4.17931035061105e-06,
+      "loss": 0.7804,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.000417931035061105,
+      "step": 2807
+    },
+    {
+      "epoch": 0.29,
+      "gumbel_tau": 0.5157,
+      "image_masks_loss": 0.0081,
+      "llm_learning_rate": 4.178695676938091e-06,
+      "loss": 0.7765,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004178695676938091,
+      "step": 2808
+    },
+    {
+      "epoch": 0.29,
+      "gumbel_tau": 0.5156,
+      "image_masks_loss": 0.0082,
+      "llm_learning_rate": 4.178080818403425e-06,
+      "loss": 0.8477,
+      "output_text_mask_loss": 0.0248,
+      "predictor_learning_rate": 0.0004178080818403425,
+      "step": 2809
+    },
+    {
+      "epoch": 0.29,
+      "gumbel_tau": 0.5155,
+      "image_masks_loss": 0.0124,
+      "llm_learning_rate": 4.17746577507476e-06,
+      "loss": 0.7609,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.000417746577507476,
+      "step": 2810
+    },
+    {
+      "epoch": 0.29,
+      "gumbel_tau": 0.5153,
+      "image_masks_loss": 0.0122,
+      "llm_learning_rate": 4.176850547019829e-06,
+      "loss": 0.7108,
+      "output_text_mask_loss": 0.0105,
+      "predictor_learning_rate": 0.0004176850547019829,
+      "step": 2811
+    },
+    {
+      "epoch": 0.29,
+      "gumbel_tau": 0.5152,
+      "image_masks_loss": 0.0416,
+      "llm_learning_rate": 4.176235134306381e-06,
+      "loss": 0.7922,
+      "output_text_mask_loss": 0.066,
+      "predictor_learning_rate": 0.0004176235134306381,
+      "step": 2812
+    },
+    {
+      "epoch": 0.29,
+      "gumbel_tau": 0.5151,
+      "image_masks_loss": 0.0416,
+      "llm_learning_rate": 4.175619537002186e-06,
+      "loss": 0.7829,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00041756195370021857,
+      "step": 2813
+    },
+    {
+      "epoch": 0.29,
+      "gumbel_tau": 0.515,
+      "image_masks_loss": 0.0116,
+      "llm_learning_rate": 4.175003755175034e-06,
+      "loss": 0.6792,
+      "output_text_mask_loss": 0.0406,
+      "predictor_learning_rate": 0.00041750037551750344,
+      "step": 2814
+    },
+    {
+      "epoch": 0.29,
+      "gumbel_tau": 0.5148,
+      "image_masks_loss": 0.0012,
+      "llm_learning_rate": 4.174387788892738e-06,
+      "loss": 0.7314,
+      "output_text_mask_loss": 0.0374,
+      "predictor_learning_rate": 0.00041743877888927384,
+      "step": 2815
+    },
+    {
+      "epoch": 0.29,
+      "gumbel_tau": 0.5147,
+      "image_masks_loss": 0.0286,
+      "llm_learning_rate": 4.173771638223131e-06,
+      "loss": 0.7405,
+      "output_text_mask_loss": 0.001,
+      "predictor_learning_rate": 0.00041737716382231305,
+      "step": 2816
+    },
+    {
+      "epoch": 0.29,
+      "gumbel_tau": 0.5146,
+      "image_masks_loss": 0.0184,
+      "llm_learning_rate": 4.173155303234061e-06,
+      "loss": 0.763,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004173155303234061,
+      "step": 2817
+    },
+    {
+      "epoch": 0.29,
+      "gumbel_tau": 0.5145,
+      "image_masks_loss": 0.0162,
+      "llm_learning_rate": 4.172538783993404e-06,
+      "loss": 0.7929,
+      "output_text_mask_loss": 0.0683,
+      "predictor_learning_rate": 0.0004172538783993403,
+      "step": 2818
+    },
+    {
+      "epoch": 0.29,
+      "gumbel_tau": 0.5144,
+      "image_masks_loss": 0.0051,
+      "llm_learning_rate": 4.171922080569049e-06,
+      "loss": 0.6593,
+      "output_text_mask_loss": 0.0584,
+      "predictor_learning_rate": 0.00041719220805690486,
+      "step": 2819
+    },
+    {
+      "epoch": 0.29,
+      "gumbel_tau": 0.5142,
+      "image_masks_loss": 0.0332,
+      "llm_learning_rate": 4.171305193028911e-06,
+      "loss": 0.8455,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004171305193028911,
+      "step": 2820
+    },
+    {
+      "epoch": 0.29,
+      "gumbel_tau": 0.5141,
+      "image_masks_loss": 0.0138,
+      "llm_learning_rate": 4.170688121440922e-06,
+      "loss": 0.7763,
+      "output_text_mask_loss": 0.0225,
+      "predictor_learning_rate": 0.0004170688121440922,
+      "step": 2821
+    },
+    {
+      "epoch": 0.29,
+      "gumbel_tau": 0.514,
+      "image_masks_loss": 0.0159,
+      "llm_learning_rate": 4.170070865873037e-06,
+      "loss": 0.7239,
+      "output_text_mask_loss": 0.046,
+      "predictor_learning_rate": 0.0004170070865873036,
+      "step": 2822
+    },
+    {
+      "epoch": 0.29,
+      "gumbel_tau": 0.5139,
+      "image_masks_loss": 0.0151,
+      "llm_learning_rate": 4.169453426393227e-06,
+      "loss": 0.7981,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.00041694534263932273,
+      "step": 2823
+    },
+    {
+      "epoch": 0.29,
+      "gumbel_tau": 0.5138,
+      "image_masks_loss": 0.0188,
+      "llm_learning_rate": 4.168835803069487e-06,
+      "loss": 0.8019,
+      "output_text_mask_loss": 0.1068,
+      "predictor_learning_rate": 0.0004168835803069487,
+      "step": 2824
+    },
+    {
+      "epoch": 0.29,
+      "gumbel_tau": 0.5136,
+      "image_masks_loss": 0.0136,
+      "llm_learning_rate": 4.168217995969832e-06,
+      "loss": 0.7189,
+      "output_text_mask_loss": 0.0076,
+      "predictor_learning_rate": 0.0004168217995969832,
+      "step": 2825
+    },
+    {
+      "epoch": 0.29,
+      "gumbel_tau": 0.5135,
+      "image_masks_loss": 0.0439,
+      "llm_learning_rate": 4.167600005162295e-06,
+      "loss": 0.8403,
+      "output_text_mask_loss": 0.0019,
+      "predictor_learning_rate": 0.00041676000051622954,
+      "step": 2826
+    },
+    {
+      "epoch": 0.29,
+      "gumbel_tau": 0.5134,
+      "image_masks_loss": 0.0149,
+      "llm_learning_rate": 4.166981830714933e-06,
+      "loss": 0.7172,
+      "output_text_mask_loss": 0.0101,
+      "predictor_learning_rate": 0.0004166981830714932,
+      "step": 2827
+    },
+    {
+      "epoch": 0.29,
+      "gumbel_tau": 0.5133,
+      "image_masks_loss": 0.0336,
+      "llm_learning_rate": 4.1663634726958165e-06,
+      "loss": 0.8469,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.0004166363472695816,
+      "step": 2828
+    },
+    {
+      "epoch": 0.29,
+      "gumbel_tau": 0.5132,
+      "image_masks_loss": 0.0047,
+      "llm_learning_rate": 4.165744931173044e-06,
+      "loss": 0.7463,
+      "output_text_mask_loss": 0.0526,
+      "predictor_learning_rate": 0.00041657449311730433,
+      "step": 2829
+    },
+    {
+      "epoch": 0.29,
+      "gumbel_tau": 0.513,
+      "image_masks_loss": 0.0096,
+      "llm_learning_rate": 4.165126206214728e-06,
+      "loss": 0.8128,
+      "output_text_mask_loss": 0.1167,
+      "predictor_learning_rate": 0.0004165126206214728,
+      "step": 2830
+    },
+    {
+      "epoch": 0.29,
+      "gumbel_tau": 0.5129,
+      "image_masks_loss": 0.0088,
+      "llm_learning_rate": 4.164507297889005e-06,
+      "loss": 0.716,
+      "output_text_mask_loss": 0.0469,
+      "predictor_learning_rate": 0.00041645072978890055,
+      "step": 2831
+    },
+    {
+      "epoch": 0.29,
+      "gumbel_tau": 0.5128,
+      "image_masks_loss": 0.0216,
+      "llm_learning_rate": 4.163888206264033e-06,
+      "loss": 0.8344,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.0004163888206264032,
+      "step": 2832
+    },
+    {
+      "epoch": 0.29,
+      "gumbel_tau": 0.5127,
+      "image_masks_loss": 0.0196,
+      "llm_learning_rate": 4.163268931407983e-06,
+      "loss": 0.6796,
+      "output_text_mask_loss": 0.002,
+      "predictor_learning_rate": 0.0004163268931407983,
+      "step": 2833
+    },
+    {
+      "epoch": 0.29,
+      "gumbel_tau": 0.5125,
+      "image_masks_loss": 0.0133,
+      "llm_learning_rate": 4.162649473389056e-06,
+      "loss": 0.855,
+      "output_text_mask_loss": 0.058,
+      "predictor_learning_rate": 0.0004162649473389055,
+      "step": 2834
+    },
+    {
+      "epoch": 0.29,
+      "gumbel_tau": 0.5124,
+      "image_masks_loss": 0.0086,
+      "llm_learning_rate": 4.1620298322754635e-06,
+      "loss": 0.7722,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 0.0004162029832275464,
+      "step": 2835
+    },
+    {
+      "epoch": 0.29,
+      "gumbel_tau": 0.5123,
+      "image_masks_loss": 0.0063,
+      "llm_learning_rate": 4.161410008135444e-06,
+      "loss": 0.7173,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.00041614100081354446,
+      "step": 2836
+    },
+    {
+      "epoch": 0.29,
+      "gumbel_tau": 0.5122,
+      "image_masks_loss": 0.0122,
+      "llm_learning_rate": 4.1607900010372556e-06,
+      "loss": 0.8731,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 0.00041607900010372554,
+      "step": 2837
+    },
+    {
+      "epoch": 0.29,
+      "gumbel_tau": 0.5121,
+      "image_masks_loss": 0.0256,
+      "llm_learning_rate": 4.160169811049171e-06,
+      "loss": 0.8076,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00041601698110491715,
+      "step": 2838
+    },
+    {
+      "epoch": 0.29,
+      "gumbel_tau": 0.5119,
+      "image_masks_loss": 0.0087,
+      "llm_learning_rate": 4.15954943823949e-06,
+      "loss": 0.7759,
+      "output_text_mask_loss": 0.0326,
+      "predictor_learning_rate": 0.000415954943823949,
+      "step": 2839
+    },
+    {
+      "epoch": 0.29,
+      "gumbel_tau": 0.5118,
+      "image_masks_loss": 0.0175,
+      "llm_learning_rate": 4.158928882676528e-06,
+      "loss": 0.7983,
+      "output_text_mask_loss": 0.0202,
+      "predictor_learning_rate": 0.0004158928882676528,
+      "step": 2840
+    },
+    {
+      "epoch": 0.29,
+      "gumbel_tau": 0.5117,
+      "image_masks_loss": 0.0301,
+      "llm_learning_rate": 4.158308144428623e-06,
+      "loss": 0.7962,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004158308144428622,
+      "step": 2841
+    },
+    {
+      "epoch": 0.29,
+      "gumbel_tau": 0.5116,
+      "image_masks_loss": 0.0182,
+      "llm_learning_rate": 4.157687223564131e-06,
+      "loss": 0.7542,
+      "output_text_mask_loss": 0.0385,
+      "predictor_learning_rate": 0.00041576872235641303,
+      "step": 2842
+    },
+    {
+      "epoch": 0.29,
+      "gumbel_tau": 0.5115,
+      "image_masks_loss": 0.0196,
+      "llm_learning_rate": 4.157066120151429e-06,
+      "loss": 0.7252,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.0004157066120151429,
+      "step": 2843
+    },
+    {
+      "epoch": 0.29,
+      "gumbel_tau": 0.5113,
+      "image_masks_loss": 0.0175,
+      "llm_learning_rate": 4.1564448342589155e-06,
+      "loss": 0.7536,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004156444834258915,
+      "step": 2844
+    },
+    {
+      "epoch": 0.29,
+      "gumbel_tau": 0.5112,
+      "image_masks_loss": 0.036,
+      "llm_learning_rate": 4.155823365955007e-06,
+      "loss": 0.7816,
+      "output_text_mask_loss": 0.0078,
+      "predictor_learning_rate": 0.00041558233659550066,
+      "step": 2845
+    },
+    {
+      "epoch": 0.29,
+      "gumbel_tau": 0.5111,
+      "image_masks_loss": 0.0049,
+      "llm_learning_rate": 4.155201715308141e-06,
+      "loss": 0.6879,
+      "output_text_mask_loss": 0.0162,
+      "predictor_learning_rate": 0.00041552017153081414,
+      "step": 2846
+    },
+    {
+      "epoch": 0.29,
+      "gumbel_tau": 0.511,
+      "image_masks_loss": 0.0189,
+      "llm_learning_rate": 4.154579882386775e-06,
+      "loss": 0.7616,
+      "output_text_mask_loss": 0.0012,
+      "predictor_learning_rate": 0.0004154579882386775,
+      "step": 2847
+    },
+    {
+      "epoch": 0.29,
+      "gumbel_tau": 0.5109,
+      "image_masks_loss": 0.0185,
+      "llm_learning_rate": 4.153957867259389e-06,
+      "loss": 0.7813,
+      "output_text_mask_loss": 0.0299,
+      "predictor_learning_rate": 0.00041539578672593883,
+      "step": 2848
+    },
+    {
+      "epoch": 0.29,
+      "gumbel_tau": 0.5107,
+      "image_masks_loss": 0.009,
+      "llm_learning_rate": 4.153335669994476e-06,
+      "loss": 0.8448,
+      "output_text_mask_loss": 0.1038,
+      "predictor_learning_rate": 0.00041533356699944765,
+      "step": 2849
+    },
+    {
+      "epoch": 0.29,
+      "gumbel_tau": 0.5106,
+      "image_masks_loss": 0.0213,
+      "llm_learning_rate": 4.1527132906605584e-06,
+      "loss": 0.8151,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004152713290660558,
+      "step": 2850
+    },
+    {
+      "epoch": 0.29,
+      "gumbel_tau": 0.5105,
+      "image_masks_loss": 0.0127,
+      "llm_learning_rate": 4.1520907293261705e-06,
+      "loss": 0.8386,
+      "output_text_mask_loss": 0.0025,
+      "predictor_learning_rate": 0.000415209072932617,
+      "step": 2851
+    },
+    {
+      "epoch": 0.29,
+      "gumbel_tau": 0.5104,
+      "image_masks_loss": 0.0077,
+      "llm_learning_rate": 4.151467986059872e-06,
+      "loss": 0.7781,
+      "output_text_mask_loss": 0.0374,
+      "predictor_learning_rate": 0.0004151467986059872,
+      "step": 2852
+    },
+    {
+      "epoch": 0.29,
+      "gumbel_tau": 0.5103,
+      "image_masks_loss": 0.0106,
+      "llm_learning_rate": 4.150845060930241e-06,
+      "loss": 0.7134,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 0.00041508450609302407,
+      "step": 2853
+    },
+    {
+      "epoch": 0.29,
+      "gumbel_tau": 0.5101,
+      "image_masks_loss": 0.0063,
+      "llm_learning_rate": 4.1502219540058745e-06,
+      "loss": 0.8216,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004150221954005874,
+      "step": 2854
+    },
+    {
+      "epoch": 0.29,
+      "gumbel_tau": 0.51,
+      "image_masks_loss": 0.0201,
+      "llm_learning_rate": 4.14959866535539e-06,
+      "loss": 0.7717,
+      "output_text_mask_loss": 0.0034,
+      "predictor_learning_rate": 0.00041495986653553897,
+      "step": 2855
+    },
+    {
+      "epoch": 0.29,
+      "gumbel_tau": 0.5099,
+      "image_masks_loss": 0.0098,
+      "llm_learning_rate": 4.1489751950474266e-06,
+      "loss": 0.7844,
+      "output_text_mask_loss": 0.069,
+      "predictor_learning_rate": 0.0004148975195047426,
+      "step": 2856
+    },
+    {
+      "epoch": 0.29,
+      "gumbel_tau": 0.5098,
+      "image_masks_loss": 0.0341,
+      "llm_learning_rate": 4.148351543150642e-06,
+      "loss": 0.718,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.0004148351543150642,
+      "step": 2857
+    },
+    {
+      "epoch": 0.29,
+      "gumbel_tau": 0.5097,
+      "image_masks_loss": 0.0362,
+      "llm_learning_rate": 4.147727709733714e-06,
+      "loss": 0.8062,
+      "output_text_mask_loss": 0.0429,
+      "predictor_learning_rate": 0.00041477277097337135,
+      "step": 2858
+    },
+    {
+      "epoch": 0.29,
+      "gumbel_tau": 0.5095,
+      "image_masks_loss": 0.0114,
+      "llm_learning_rate": 4.14710369486534e-06,
+      "loss": 0.7165,
+      "output_text_mask_loss": 0.0322,
+      "predictor_learning_rate": 0.00041471036948653397,
+      "step": 2859
+    },
+    {
+      "epoch": 0.29,
+      "gumbel_tau": 0.5094,
+      "image_masks_loss": 0.0301,
+      "llm_learning_rate": 4.146479498614239e-06,
+      "loss": 0.6959,
+      "output_text_mask_loss": 0.0095,
+      "predictor_learning_rate": 0.00041464794986142387,
+      "step": 2860
+    },
+    {
+      "epoch": 0.29,
+      "gumbel_tau": 0.5093,
+      "image_masks_loss": 0.0156,
+      "llm_learning_rate": 4.145855121049147e-06,
+      "loss": 0.7874,
+      "output_text_mask_loss": 0.0042,
+      "predictor_learning_rate": 0.0004145855121049147,
+      "step": 2861
+    },
+    {
+      "epoch": 0.29,
+      "gumbel_tau": 0.5092,
+      "image_masks_loss": 0.0174,
+      "llm_learning_rate": 4.145230562238824e-06,
+      "loss": 0.8794,
+      "output_text_mask_loss": 0.0052,
+      "predictor_learning_rate": 0.0004145230562238823,
+      "step": 2862
+    },
+    {
+      "epoch": 0.29,
+      "gumbel_tau": 0.5091,
+      "image_masks_loss": 0.0269,
+      "llm_learning_rate": 4.144605822252047e-06,
+      "loss": 0.7607,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.00041446058222520464,
+      "step": 2863
+    },
+    {
+      "epoch": 0.29,
+      "gumbel_tau": 0.5089,
+      "image_masks_loss": 0.0181,
+      "llm_learning_rate": 4.143980901157613e-06,
+      "loss": 0.7534,
+      "output_text_mask_loss": 0.037,
+      "predictor_learning_rate": 0.0004143980901157613,
+      "step": 2864
+    },
+    {
+      "epoch": 0.29,
+      "gumbel_tau": 0.5088,
+      "image_masks_loss": 0.0183,
+      "llm_learning_rate": 4.143355799024341e-06,
+      "loss": 0.7361,
+      "output_text_mask_loss": 0.0095,
+      "predictor_learning_rate": 0.000414335579902434,
+      "step": 2865
+    },
+    {
+      "epoch": 0.29,
+      "gumbel_tau": 0.5087,
+      "image_masks_loss": 0.0706,
+      "llm_learning_rate": 4.142730515921067e-06,
+      "loss": 0.7764,
+      "output_text_mask_loss": 0.0162,
+      "predictor_learning_rate": 0.0004142730515921067,
+      "step": 2866
+    },
+    {
+      "epoch": 0.29,
+      "gumbel_tau": 0.5086,
+      "image_masks_loss": 0.0444,
+      "llm_learning_rate": 4.142105051916651e-06,
+      "loss": 0.7466,
+      "output_text_mask_loss": 0.0126,
+      "predictor_learning_rate": 0.000414210505191665,
+      "step": 2867
+    },
+    {
+      "epoch": 0.29,
+      "gumbel_tau": 0.5085,
+      "image_masks_loss": 0.0106,
+      "llm_learning_rate": 4.141479407079968e-06,
+      "loss": 0.7488,
+      "output_text_mask_loss": 0.0484,
+      "predictor_learning_rate": 0.00041414794070799676,
+      "step": 2868
+    },
+    {
+      "epoch": 0.29,
+      "gumbel_tau": 0.5083,
+      "image_masks_loss": 0.0107,
+      "llm_learning_rate": 4.140853581479916e-06,
+      "loss": 0.7015,
+      "output_text_mask_loss": 0.002,
+      "predictor_learning_rate": 0.0004140853581479916,
+      "step": 2869
+    },
+    {
+      "epoch": 0.29,
+      "gumbel_tau": 0.5082,
+      "image_masks_loss": 0.0105,
+      "llm_learning_rate": 4.140227575185413e-06,
+      "loss": 0.7482,
+      "output_text_mask_loss": 0.0132,
+      "predictor_learning_rate": 0.00041402275751854133,
+      "step": 2870
+    },
+    {
+      "epoch": 0.29,
+      "gumbel_tau": 0.5081,
+      "image_masks_loss": 0.0167,
+      "llm_learning_rate": 4.139601388265397e-06,
+      "loss": 0.7449,
+      "output_text_mask_loss": 0.0073,
+      "predictor_learning_rate": 0.00041396013882653973,
+      "step": 2871
+    },
+    {
+      "epoch": 0.29,
+      "gumbel_tau": 0.508,
+      "image_masks_loss": 0.0212,
+      "llm_learning_rate": 4.138975020788823e-06,
+      "loss": 0.7295,
+      "output_text_mask_loss": 0.0391,
+      "predictor_learning_rate": 0.0004138975020788823,
+      "step": 2872
+    },
+    {
+      "epoch": 0.29,
+      "gumbel_tau": 0.5079,
+      "image_masks_loss": 0.0425,
+      "llm_learning_rate": 4.138348472824671e-06,
+      "loss": 0.6973,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00041383484728246703,
+      "step": 2873
+    },
+    {
+      "epoch": 0.29,
+      "gumbel_tau": 0.5077,
+      "image_masks_loss": 0.0257,
+      "llm_learning_rate": 4.137721744441934e-06,
+      "loss": 0.7613,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 0.0004137721744441933,
+      "step": 2874
+    },
+    {
+      "epoch": 0.29,
+      "gumbel_tau": 0.5076,
+      "image_masks_loss": 0.0032,
+      "llm_learning_rate": 4.13709483570963e-06,
+      "loss": 0.7109,
+      "output_text_mask_loss": 0.0435,
+      "predictor_learning_rate": 0.000413709483570963,
+      "step": 2875
+    },
+    {
+      "epoch": 0.29,
+      "gumbel_tau": 0.5075,
+      "image_masks_loss": 0.0372,
+      "llm_learning_rate": 4.136467746696797e-06,
+      "loss": 0.8971,
+      "output_text_mask_loss": 0.0028,
+      "predictor_learning_rate": 0.0004136467746696797,
+      "step": 2876
+    },
+    {
+      "epoch": 0.29,
+      "gumbel_tau": 0.5074,
+      "image_masks_loss": 0.0277,
+      "llm_learning_rate": 4.1358404774724915e-06,
+      "loss": 0.8332,
+      "output_text_mask_loss": 0.0056,
+      "predictor_learning_rate": 0.0004135840477472491,
+      "step": 2877
+    },
+    {
+      "epoch": 0.29,
+      "gumbel_tau": 0.5073,
+      "image_masks_loss": 0.0362,
+      "llm_learning_rate": 4.135213028105788e-06,
+      "loss": 0.7741,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.0004135213028105788,
+      "step": 2878
+    },
+    {
+      "epoch": 0.29,
+      "gumbel_tau": 0.5071,
+      "image_masks_loss": 0.028,
+      "llm_learning_rate": 4.134585398665784e-06,
+      "loss": 0.8689,
+      "output_text_mask_loss": 0.0149,
+      "predictor_learning_rate": 0.0004134585398665783,
+      "step": 2879
+    },
+    {
+      "epoch": 0.3,
+      "gumbel_tau": 0.507,
+      "image_masks_loss": 0.0199,
+      "llm_learning_rate": 4.133957589221595e-06,
+      "loss": 0.8007,
+      "output_text_mask_loss": 0.0652,
+      "predictor_learning_rate": 0.00041339575892215943,
+      "step": 2880
+    },
+    {
+      "epoch": 0.3,
+      "gumbel_tau": 0.5069,
+      "image_masks_loss": 0.0359,
+      "llm_learning_rate": 4.133329599842356e-06,
+      "loss": 0.701,
+      "output_text_mask_loss": 0.0641,
+      "predictor_learning_rate": 0.0004133329599842356,
+      "step": 2881
+    },
+    {
+      "epoch": 0.3,
+      "gumbel_tau": 0.5068,
+      "image_masks_loss": 0.0248,
+      "llm_learning_rate": 4.132701430597223e-06,
+      "loss": 0.7707,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.0004132701430597223,
+      "step": 2882
+    },
+    {
+      "epoch": 0.3,
+      "gumbel_tau": 0.5067,
+      "image_masks_loss": 0.0133,
+      "llm_learning_rate": 4.132073081555374e-06,
+      "loss": 0.8136,
+      "output_text_mask_loss": 0.0496,
+      "predictor_learning_rate": 0.00041320730815553735,
+      "step": 2883
+    },
+    {
+      "epoch": 0.3,
+      "gumbel_tau": 0.5065,
+      "image_masks_loss": 0.0317,
+      "llm_learning_rate": 4.131444552786001e-06,
+      "loss": 0.7707,
+      "output_text_mask_loss": 0.0076,
+      "predictor_learning_rate": 0.00041314445527860003,
+      "step": 2884
+    },
+    {
+      "epoch": 0.3,
+      "gumbel_tau": 0.5064,
+      "image_masks_loss": 0.0261,
+      "llm_learning_rate": 4.130815844358319e-06,
+      "loss": 0.8275,
+      "output_text_mask_loss": 0.0374,
+      "predictor_learning_rate": 0.0004130815844358319,
+      "step": 2885
+    },
+    {
+      "epoch": 0.3,
+      "gumbel_tau": 0.5063,
+      "image_masks_loss": 0.0227,
+      "llm_learning_rate": 4.130186956341564e-06,
+      "loss": 0.7602,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004130186956341564,
+      "step": 2886
+    },
+    {
+      "epoch": 0.3,
+      "gumbel_tau": 0.5062,
+      "image_masks_loss": 0.0366,
+      "llm_learning_rate": 4.129557888804991e-06,
+      "loss": 0.873,
+      "output_text_mask_loss": 0.0401,
+      "predictor_learning_rate": 0.0004129557888804991,
+      "step": 2887
+    },
+    {
+      "epoch": 0.3,
+      "gumbel_tau": 0.5061,
+      "image_masks_loss": 0.0408,
+      "llm_learning_rate": 4.128928641817873e-06,
+      "loss": 0.7833,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 0.00041289286418178726,
+      "step": 2888
+    },
+    {
+      "epoch": 0.3,
+      "gumbel_tau": 0.5059,
+      "image_masks_loss": 0.0162,
+      "llm_learning_rate": 4.128299215449505e-06,
+      "loss": 0.8311,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.00041282992154495046,
+      "step": 2889
+    },
+    {
+      "epoch": 0.3,
+      "gumbel_tau": 0.5058,
+      "image_masks_loss": 0.0222,
+      "llm_learning_rate": 4.1276696097692004e-06,
+      "loss": 0.8136,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00041276696097692,
+      "step": 2890
+    },
+    {
+      "epoch": 0.3,
+      "gumbel_tau": 0.5057,
+      "image_masks_loss": 0.0435,
+      "llm_learning_rate": 4.127039824846293e-06,
+      "loss": 0.7718,
+      "output_text_mask_loss": 0.0458,
+      "predictor_learning_rate": 0.00041270398248462924,
+      "step": 2891
+    },
+    {
+      "epoch": 0.3,
+      "gumbel_tau": 0.5056,
+      "image_masks_loss": 0.0254,
+      "llm_learning_rate": 4.1264098607501355e-06,
+      "loss": 0.8238,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00041264098607501356,
+      "step": 2892
+    },
+    {
+      "epoch": 0.3,
+      "gumbel_tau": 0.5055,
+      "image_masks_loss": 0.0366,
+      "llm_learning_rate": 4.1257797175501025e-06,
+      "loss": 0.7712,
+      "output_text_mask_loss": 0.0248,
+      "predictor_learning_rate": 0.00041257797175501024,
+      "step": 2893
+    },
+    {
+      "epoch": 0.3,
+      "gumbel_tau": 0.5053,
+      "image_masks_loss": 0.0099,
+      "llm_learning_rate": 4.125149395315584e-06,
+      "loss": 0.7744,
+      "output_text_mask_loss": 0.0454,
+      "predictor_learning_rate": 0.00041251493953155846,
+      "step": 2894
+    },
+    {
+      "epoch": 0.3,
+      "gumbel_tau": 0.5052,
+      "image_masks_loss": 0.0164,
+      "llm_learning_rate": 4.1245188941159954e-06,
+      "loss": 0.7647,
+      "output_text_mask_loss": 0.0275,
+      "predictor_learning_rate": 0.00041245188941159953,
+      "step": 2895
+    },
+    {
+      "epoch": 0.3,
+      "gumbel_tau": 0.5051,
+      "image_masks_loss": 0.0153,
+      "llm_learning_rate": 4.123888214020768e-06,
+      "loss": 0.8077,
+      "output_text_mask_loss": 0.0549,
+      "predictor_learning_rate": 0.0004123888214020768,
+      "step": 2896
+    },
+    {
+      "epoch": 0.3,
+      "gumbel_tau": 0.505,
+      "image_masks_loss": 0.0462,
+      "llm_learning_rate": 4.123257355099352e-06,
+      "loss": 0.8073,
+      "output_text_mask_loss": 0.0965,
+      "predictor_learning_rate": 0.00041232573550993524,
+      "step": 2897
+    },
+    {
+      "epoch": 0.3,
+      "gumbel_tau": 0.5049,
+      "image_masks_loss": 0.0136,
+      "llm_learning_rate": 4.1226263174212225e-06,
+      "loss": 0.7983,
+      "output_text_mask_loss": 0.0195,
+      "predictor_learning_rate": 0.0004122626317421222,
+      "step": 2898
+    },
+    {
+      "epoch": 0.3,
+      "gumbel_tau": 0.5047,
+      "image_masks_loss": 0.0141,
+      "llm_learning_rate": 4.1219951010558665e-06,
+      "loss": 0.9071,
+      "output_text_mask_loss": 0.0486,
+      "predictor_learning_rate": 0.00041219951010558665,
+      "step": 2899
+    },
+    {
+      "epoch": 0.3,
+      "gumbel_tau": 0.5046,
+      "image_masks_loss": 0.0171,
+      "llm_learning_rate": 4.1213637060727986e-06,
+      "loss": 0.8286,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004121363706072798,
+      "step": 2900
+    },
+    {
+      "epoch": 0.3,
+      "gumbel_tau": 0.5045,
+      "image_masks_loss": 0.0294,
+      "llm_learning_rate": 4.120732132541547e-06,
+      "loss": 0.7765,
+      "output_text_mask_loss": 0.0032,
+      "predictor_learning_rate": 0.0004120732132541547,
+      "step": 2901
+    },
+    {
+      "epoch": 0.3,
+      "gumbel_tau": 0.5044,
+      "image_masks_loss": 0.0217,
+      "llm_learning_rate": 4.120100380531664e-06,
+      "loss": 0.907,
+      "output_text_mask_loss": 0.0019,
+      "predictor_learning_rate": 0.0004120100380531664,
+      "step": 2902
+    },
+    {
+      "epoch": 0.3,
+      "gumbel_tau": 0.5043,
+      "image_masks_loss": 0.0125,
+      "llm_learning_rate": 4.119468450112718e-06,
+      "loss": 0.7475,
+      "output_text_mask_loss": 0.0139,
+      "predictor_learning_rate": 0.0004119468450112718,
+      "step": 2903
+    },
+    {
+      "epoch": 0.3,
+      "gumbel_tau": 0.5042,
+      "image_masks_loss": 0.0199,
+      "llm_learning_rate": 4.1188363413542985e-06,
+      "loss": 0.7749,
+      "output_text_mask_loss": 0.0116,
+      "predictor_learning_rate": 0.00041188363413542986,
+      "step": 2904
+    },
+    {
+      "epoch": 0.3,
+      "gumbel_tau": 0.504,
+      "image_masks_loss": 0.0277,
+      "llm_learning_rate": 4.118204054326016e-06,
+      "loss": 0.7695,
+      "output_text_mask_loss": 0.0004,
+      "predictor_learning_rate": 0.0004118204054326016,
+      "step": 2905
+    },
+    {
+      "epoch": 0.3,
+      "gumbel_tau": 0.5039,
+      "image_masks_loss": 0.0259,
+      "llm_learning_rate": 4.117571589097497e-06,
+      "loss": 0.7999,
+      "output_text_mask_loss": 0.0368,
+      "predictor_learning_rate": 0.00041175715890974973,
+      "step": 2906
+    },
+    {
+      "epoch": 0.3,
+      "gumbel_tau": 0.5038,
+      "image_masks_loss": 0.0139,
+      "llm_learning_rate": 4.116938945738394e-06,
+      "loss": 0.776,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.0004116938945738393,
+      "step": 2907
+    },
+    {
+      "epoch": 0.3,
+      "gumbel_tau": 0.5037,
+      "image_masks_loss": 0.016,
+      "llm_learning_rate": 4.116306124318371e-06,
+      "loss": 0.7447,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.0004116306124318371,
+      "step": 2908
+    },
+    {
+      "epoch": 0.3,
+      "gumbel_tau": 0.5036,
+      "image_masks_loss": 0.0149,
+      "llm_learning_rate": 4.115673124907119e-06,
+      "loss": 0.7901,
+      "output_text_mask_loss": 0.0809,
+      "predictor_learning_rate": 0.00041156731249071185,
+      "step": 2909
+    },
+    {
+      "epoch": 0.3,
+      "gumbel_tau": 0.5034,
+      "image_masks_loss": 0.0089,
+      "llm_learning_rate": 4.115039947574342e-06,
+      "loss": 0.7224,
+      "output_text_mask_loss": 0.016,
+      "predictor_learning_rate": 0.0004115039947574342,
+      "step": 2910
+    },
+    {
+      "epoch": 0.3,
+      "gumbel_tau": 0.5033,
+      "image_masks_loss": 0.0073,
+      "llm_learning_rate": 4.11440659238977e-06,
+      "loss": 0.7665,
+      "output_text_mask_loss": 0.0162,
+      "predictor_learning_rate": 0.00041144065923897694,
+      "step": 2911
+    },
+    {
+      "epoch": 0.3,
+      "gumbel_tau": 0.5032,
+      "image_masks_loss": 0.0207,
+      "llm_learning_rate": 4.113773059423147e-06,
+      "loss": 0.7203,
+      "output_text_mask_loss": 0.0885,
+      "predictor_learning_rate": 0.0004113773059423147,
+      "step": 2912
+    },
+    {
+      "epoch": 0.3,
+      "gumbel_tau": 0.5031,
+      "image_masks_loss": 0.0062,
+      "llm_learning_rate": 4.113139348744241e-06,
+      "loss": 0.8283,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.00041131393487442403,
+      "step": 2913
+    },
+    {
+      "epoch": 0.3,
+      "gumbel_tau": 0.503,
+      "image_masks_loss": 0.0254,
+      "llm_learning_rate": 4.112505460422836e-06,
+      "loss": 0.7513,
+      "output_text_mask_loss": 0.0507,
+      "predictor_learning_rate": 0.0004112505460422836,
+      "step": 2914
+    },
+    {
+      "epoch": 0.3,
+      "gumbel_tau": 0.5028,
+      "image_masks_loss": 0.0077,
+      "llm_learning_rate": 4.111871394528738e-06,
+      "loss": 0.7353,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.00041118713945287376,
+      "step": 2915
+    },
+    {
+      "epoch": 0.3,
+      "gumbel_tau": 0.5027,
+      "image_masks_loss": 0.0226,
+      "llm_learning_rate": 4.111237151131772e-06,
+      "loss": 0.748,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004111237151131772,
+      "step": 2916
+    },
+    {
+      "epoch": 0.3,
+      "gumbel_tau": 0.5026,
+      "image_masks_loss": 0.0065,
+      "llm_learning_rate": 4.110602730301782e-06,
+      "loss": 0.7675,
+      "output_text_mask_loss": 0.0162,
+      "predictor_learning_rate": 0.00041106027303017826,
+      "step": 2917
+    },
+    {
+      "epoch": 0.3,
+      "gumbel_tau": 0.5025,
+      "image_masks_loss": 0.022,
+      "llm_learning_rate": 4.1099681321086325e-06,
+      "loss": 0.7108,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 0.00041099681321086324,
+      "step": 2918
+    },
+    {
+      "epoch": 0.3,
+      "gumbel_tau": 0.5024,
+      "image_masks_loss": 0.0091,
+      "llm_learning_rate": 4.109333356622205e-06,
+      "loss": 0.77,
+      "output_text_mask_loss": 0.024,
+      "predictor_learning_rate": 0.0004109333356622205,
+      "step": 2919
+    },
+    {
+      "epoch": 0.3,
+      "gumbel_tau": 0.5023,
+      "image_masks_loss": 0.0138,
+      "llm_learning_rate": 4.108698403912405e-06,
+      "loss": 0.8789,
+      "output_text_mask_loss": 0.0025,
+      "predictor_learning_rate": 0.0004108698403912405,
+      "step": 2920
+    },
+    {
+      "epoch": 0.3,
+      "gumbel_tau": 0.5021,
+      "image_masks_loss": 0.015,
+      "llm_learning_rate": 4.108063274049153e-06,
+      "loss": 0.7639,
+      "output_text_mask_loss": 0.0496,
+      "predictor_learning_rate": 0.0004108063274049153,
+      "step": 2921
+    },
+    {
+      "epoch": 0.3,
+      "gumbel_tau": 0.502,
+      "image_masks_loss": 0.0269,
+      "llm_learning_rate": 4.107427967102391e-06,
+      "loss": 0.7877,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004107427967102391,
+      "step": 2922
+    },
+    {
+      "epoch": 0.3,
+      "gumbel_tau": 0.5019,
+      "image_masks_loss": 0.0343,
+      "llm_learning_rate": 4.1067924831420815e-06,
+      "loss": 0.8412,
+      "output_text_mask_loss": 0.0204,
+      "predictor_learning_rate": 0.0004106792483142081,
+      "step": 2923
+    },
+    {
+      "epoch": 0.3,
+      "gumbel_tau": 0.5018,
+      "image_masks_loss": 0.0248,
+      "llm_learning_rate": 4.106156822238204e-06,
+      "loss": 0.7681,
+      "output_text_mask_loss": 0.005,
+      "predictor_learning_rate": 0.0004106156822238204,
+      "step": 2924
+    },
+    {
+      "epoch": 0.3,
+      "gumbel_tau": 0.5017,
+      "image_masks_loss": 0.0096,
+      "llm_learning_rate": 4.10552098446076e-06,
+      "loss": 0.7316,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.000410552098446076,
+      "step": 2925
+    },
+    {
+      "epoch": 0.3,
+      "gumbel_tau": 0.5015,
+      "image_masks_loss": 0.0201,
+      "llm_learning_rate": 4.104884969879769e-06,
+      "loss": 0.7375,
+      "output_text_mask_loss": 0.001,
+      "predictor_learning_rate": 0.0004104884969879769,
+      "step": 2926
+    },
+    {
+      "epoch": 0.3,
+      "gumbel_tau": 0.5014,
+      "image_masks_loss": 0.0334,
+      "llm_learning_rate": 4.104248778565271e-06,
+      "loss": 0.7938,
+      "output_text_mask_loss": 0.0208,
+      "predictor_learning_rate": 0.00041042487785652705,
+      "step": 2927
+    },
+    {
+      "epoch": 0.3,
+      "gumbel_tau": 0.5013,
+      "image_masks_loss": 0.013,
+      "llm_learning_rate": 4.103612410587323e-06,
+      "loss": 0.8398,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.0004103612410587323,
+      "step": 2928
+    },
+    {
+      "epoch": 0.3,
+      "gumbel_tau": 0.5012,
+      "image_masks_loss": 0.0136,
+      "llm_learning_rate": 4.102975866016005e-06,
+      "loss": 0.7789,
+      "output_text_mask_loss": 0.0078,
+      "predictor_learning_rate": 0.0004102975866016005,
+      "step": 2929
+    },
+    {
+      "epoch": 0.3,
+      "gumbel_tau": 0.5011,
+      "image_masks_loss": 0.0126,
+      "llm_learning_rate": 4.102339144921415e-06,
+      "loss": 0.7458,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.0004102339144921415,
+      "step": 2930
+    },
+    {
+      "epoch": 0.3,
+      "gumbel_tau": 0.501,
+      "image_masks_loss": 0.0267,
+      "llm_learning_rate": 4.101702247373669e-06,
+      "loss": 0.7565,
+      "output_text_mask_loss": 0.0147,
+      "predictor_learning_rate": 0.00041017022473736686,
+      "step": 2931
+    },
+    {
+      "epoch": 0.3,
+      "gumbel_tau": 0.5008,
+      "image_masks_loss": 0.0204,
+      "llm_learning_rate": 4.101065173442903e-06,
+      "loss": 0.7296,
+      "output_text_mask_loss": 0.0353,
+      "predictor_learning_rate": 0.00041010651734429027,
+      "step": 2932
+    },
+    {
+      "epoch": 0.3,
+      "gumbel_tau": 0.5007,
+      "image_masks_loss": 0.019,
+      "llm_learning_rate": 4.100427923199274e-06,
+      "loss": 0.7172,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00041004279231992733,
+      "step": 2933
+    },
+    {
+      "epoch": 0.3,
+      "gumbel_tau": 0.5006,
+      "image_masks_loss": 0.0231,
+      "llm_learning_rate": 4.099790496712957e-06,
+      "loss": 0.7491,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004099790496712957,
+      "step": 2934
+    },
+    {
+      "epoch": 0.3,
+      "gumbel_tau": 0.5005,
+      "image_masks_loss": 0.0241,
+      "llm_learning_rate": 4.099152894054148e-06,
+      "loss": 0.6918,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004099152894054148,
+      "step": 2935
+    },
+    {
+      "epoch": 0.3,
+      "gumbel_tau": 0.5004,
+      "image_masks_loss": 0.0148,
+      "llm_learning_rate": 4.09851511529306e-06,
+      "loss": 0.761,
+      "output_text_mask_loss": 0.0039,
+      "predictor_learning_rate": 0.000409851511529306,
+      "step": 2936
+    },
+    {
+      "epoch": 0.3,
+      "gumbel_tau": 0.5002,
+      "image_masks_loss": 0.0427,
+      "llm_learning_rate": 4.097877160499927e-06,
+      "loss": 0.7284,
+      "output_text_mask_loss": 0.0095,
+      "predictor_learning_rate": 0.0004097877160499926,
+      "step": 2937
+    },
+    {
+      "epoch": 0.3,
+      "gumbel_tau": 0.5001,
+      "image_masks_loss": 0.028,
+      "llm_learning_rate": 4.097239029745002e-06,
+      "loss": 0.8045,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 0.00040972390297450026,
+      "step": 2938
+    },
+    {
+      "epoch": 0.3,
+      "gumbel_tau": 0.5,
+      "image_masks_loss": 0.014,
+      "llm_learning_rate": 4.0966007230985574e-06,
+      "loss": 0.6666,
+      "output_text_mask_loss": 0.0058,
+      "predictor_learning_rate": 0.00040966007230985575,
+      "step": 2939
+    },
+    {
+      "epoch": 0.3,
+      "gumbel_tau": 0.4999,
+      "image_masks_loss": 0.0136,
+      "llm_learning_rate": 4.095962240630885e-06,
+      "loss": 0.6561,
+      "output_text_mask_loss": 0.0496,
+      "predictor_learning_rate": 0.0004095962240630885,
+      "step": 2940
+    },
+    {
+      "epoch": 0.3,
+      "gumbel_tau": 0.4998,
+      "image_masks_loss": 0.0151,
+      "llm_learning_rate": 4.095323582412296e-06,
+      "loss": 0.7337,
+      "output_text_mask_loss": 0.0467,
+      "predictor_learning_rate": 0.00040953235824122954,
+      "step": 2941
+    },
+    {
+      "epoch": 0.3,
+      "gumbel_tau": 0.4997,
+      "image_masks_loss": 0.0183,
+      "llm_learning_rate": 4.094684748513119e-06,
+      "loss": 0.818,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004094684748513119,
+      "step": 2942
+    },
+    {
+      "epoch": 0.3,
+      "gumbel_tau": 0.4995,
+      "image_masks_loss": 0.0144,
+      "llm_learning_rate": 4.094045739003707e-06,
+      "loss": 0.7201,
+      "output_text_mask_loss": 0.0172,
+      "predictor_learning_rate": 0.0004094045739003707,
+      "step": 2943
+    },
+    {
+      "epoch": 0.3,
+      "gumbel_tau": 0.4994,
+      "image_masks_loss": 0.0077,
+      "llm_learning_rate": 4.093406553954427e-06,
+      "loss": 0.7821,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00040934065539544273,
+      "step": 2944
+    },
+    {
+      "epoch": 0.3,
+      "gumbel_tau": 0.4993,
+      "image_masks_loss": 0.0223,
+      "llm_learning_rate": 4.092767193435668e-06,
+      "loss": 0.8725,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.0004092767193435668,
+      "step": 2945
+    },
+    {
+      "epoch": 0.3,
+      "gumbel_tau": 0.4992,
+      "image_masks_loss": 0.0095,
+      "llm_learning_rate": 4.0921276575178395e-06,
+      "loss": 0.6145,
+      "output_text_mask_loss": 0.0078,
+      "predictor_learning_rate": 0.0004092127657517839,
+      "step": 2946
+    },
+    {
+      "epoch": 0.3,
+      "gumbel_tau": 0.4991,
+      "image_masks_loss": 0.033,
+      "llm_learning_rate": 4.0914879462713655e-06,
+      "loss": 0.7352,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004091487946271365,
+      "step": 2947
+    },
+    {
+      "epoch": 0.3,
+      "gumbel_tau": 0.4989,
+      "image_masks_loss": 0.0261,
+      "llm_learning_rate": 4.090848059766694e-06,
+      "loss": 0.787,
+      "output_text_mask_loss": 0.0027,
+      "predictor_learning_rate": 0.0004090848059766694,
+      "step": 2948
+    },
+    {
+      "epoch": 0.3,
+      "gumbel_tau": 0.4988,
+      "image_masks_loss": 0.0166,
+      "llm_learning_rate": 4.090207998074291e-06,
+      "loss": 0.8162,
+      "output_text_mask_loss": 0.0195,
+      "predictor_learning_rate": 0.0004090207998074291,
+      "step": 2949
+    },
+    {
+      "epoch": 0.3,
+      "gumbel_tau": 0.4987,
+      "image_masks_loss": 0.0241,
+      "llm_learning_rate": 4.089567761264641e-06,
+      "loss": 0.7586,
+      "output_text_mask_loss": 0.0824,
+      "predictor_learning_rate": 0.0004089567761264641,
+      "step": 2950
+    },
+    {
+      "epoch": 0.3,
+      "gumbel_tau": 0.4986,
+      "image_masks_loss": 0.0167,
+      "llm_learning_rate": 4.088927349408249e-06,
+      "loss": 1.0179,
+      "output_text_mask_loss": 0.3021,
+      "predictor_learning_rate": 0.00040889273494082483,
+      "step": 2951
+    },
+    {
+      "epoch": 0.3,
+      "gumbel_tau": 0.4985,
+      "image_masks_loss": 0.0172,
+      "llm_learning_rate": 4.088286762575637e-06,
+      "loss": 1.0733,
+      "output_text_mask_loss": 0.1831,
+      "predictor_learning_rate": 0.0004088286762575637,
+      "step": 2952
+    },
+    {
+      "epoch": 0.3,
+      "gumbel_tau": 0.4984,
+      "image_masks_loss": 0.0294,
+      "llm_learning_rate": 4.08764600083735e-06,
+      "loss": 1.1992,
+      "output_text_mask_loss": 0.5798,
+      "predictor_learning_rate": 0.00040876460008373503,
+      "step": 2953
+    },
+    {
+      "epoch": 0.3,
+      "gumbel_tau": 0.4982,
+      "image_masks_loss": 0.0235,
+      "llm_learning_rate": 4.087005064263949e-06,
+      "loss": 1.3972,
+      "output_text_mask_loss": 0.7324,
+      "predictor_learning_rate": 0.00040870050642639487,
+      "step": 2954
+    },
+    {
+      "epoch": 0.3,
+      "gumbel_tau": 0.4981,
+      "image_masks_loss": 0.0205,
+      "llm_learning_rate": 4.0863639529260154e-06,
+      "loss": 1.4318,
+      "output_text_mask_loss": 0.9216,
+      "predictor_learning_rate": 0.0004086363952926015,
+      "step": 2955
+    },
+    {
+      "epoch": 0.3,
+      "gumbel_tau": 0.498,
+      "image_masks_loss": 0.0271,
+      "llm_learning_rate": 4.0857226668941495e-06,
+      "loss": 1.223,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004085722666894149,
+      "step": 2956
+    },
+    {
+      "epoch": 0.3,
+      "gumbel_tau": 0.4979,
+      "image_masks_loss": 0.0119,
+      "llm_learning_rate": 4.085081206238972e-06,
+      "loss": 1.1252,
+      "output_text_mask_loss": 0.2884,
+      "predictor_learning_rate": 0.00040850812062389713,
+      "step": 2957
+    },
+    {
+      "epoch": 0.3,
+      "gumbel_tau": 0.4978,
+      "image_masks_loss": 0.0275,
+      "llm_learning_rate": 4.08443957103112e-06,
+      "loss": 0.868,
+      "output_text_mask_loss": 0.1953,
+      "predictor_learning_rate": 0.000408443957103112,
+      "step": 2958
+    },
+    {
+      "epoch": 0.3,
+      "gumbel_tau": 0.4977,
+      "image_masks_loss": 0.0254,
+      "llm_learning_rate": 4.083797761341255e-06,
+      "loss": 0.8464,
+      "output_text_mask_loss": 0.0231,
+      "predictor_learning_rate": 0.00040837977613412545,
+      "step": 2959
+    },
+    {
+      "epoch": 0.3,
+      "gumbel_tau": 0.4975,
+      "image_masks_loss": 0.0176,
+      "llm_learning_rate": 4.083155777240052e-06,
+      "loss": 0.9128,
+      "output_text_mask_loss": 0.119,
+      "predictor_learning_rate": 0.0004083155777240051,
+      "step": 2960
+    },
+    {
+      "epoch": 0.3,
+      "gumbel_tau": 0.4974,
+      "image_masks_loss": 0.0244,
+      "llm_learning_rate": 4.082513618798208e-06,
+      "loss": 0.7675,
+      "output_text_mask_loss": 0.0021,
+      "predictor_learning_rate": 0.00040825136187982075,
+      "step": 2961
+    },
+    {
+      "epoch": 0.3,
+      "gumbel_tau": 0.4973,
+      "image_masks_loss": 0.0278,
+      "llm_learning_rate": 4.08187128608644e-06,
+      "loss": 0.7429,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.00040818712860864393,
+      "step": 2962
+    },
+    {
+      "epoch": 0.3,
+      "gumbel_tau": 0.4972,
+      "image_masks_loss": 0.0062,
+      "llm_learning_rate": 4.081228779175482e-06,
+      "loss": 0.7693,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.00040812287791754814,
+      "step": 2963
+    },
+    {
+      "epoch": 0.3,
+      "gumbel_tau": 0.4971,
+      "image_masks_loss": 0.0176,
+      "llm_learning_rate": 4.080586098136089e-06,
+      "loss": 0.8422,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.0004080586098136089,
+      "step": 2964
+    },
+    {
+      "epoch": 0.3,
+      "gumbel_tau": 0.4969,
+      "image_masks_loss": 0.0164,
+      "llm_learning_rate": 4.079943243039034e-06,
+      "loss": 0.882,
+      "output_text_mask_loss": 0.0416,
+      "predictor_learning_rate": 0.00040799432430390335,
+      "step": 2965
+    },
+    {
+      "epoch": 0.3,
+      "gumbel_tau": 0.4968,
+      "image_masks_loss": 0.0299,
+      "llm_learning_rate": 4.079300213955111e-06,
+      "loss": 0.8923,
+      "output_text_mask_loss": 0.0349,
+      "predictor_learning_rate": 0.00040793002139551105,
+      "step": 2966
+    },
+    {
+      "epoch": 0.3,
+      "gumbel_tau": 0.4967,
+      "image_masks_loss": 0.02,
+      "llm_learning_rate": 4.07865701095513e-06,
+      "loss": 0.8217,
+      "output_text_mask_loss": 0.0938,
+      "predictor_learning_rate": 0.0004078657010955129,
+      "step": 2967
+    },
+    {
+      "epoch": 0.3,
+      "gumbel_tau": 0.4966,
+      "image_masks_loss": 0.0204,
+      "llm_learning_rate": 4.078013634109922e-06,
+      "loss": 0.8351,
+      "output_text_mask_loss": 0.0397,
+      "predictor_learning_rate": 0.00040780136341099217,
+      "step": 2968
+    },
+    {
+      "epoch": 0.3,
+      "gumbel_tau": 0.4965,
+      "image_masks_loss": 0.0062,
+      "llm_learning_rate": 4.077370083490338e-06,
+      "loss": 0.7805,
+      "output_text_mask_loss": 0.0101,
+      "predictor_learning_rate": 0.00040773700834903384,
+      "step": 2969
+    },
+    {
+      "epoch": 0.3,
+      "gumbel_tau": 0.4964,
+      "image_masks_loss": 0.0087,
+      "llm_learning_rate": 4.076726359167248e-06,
+      "loss": 0.6241,
+      "output_text_mask_loss": 0.0081,
+      "predictor_learning_rate": 0.00040767263591672475,
+      "step": 2970
+    },
+    {
+      "epoch": 0.3,
+      "gumbel_tau": 0.4962,
+      "image_masks_loss": 0.0082,
+      "llm_learning_rate": 4.076082461211538e-06,
+      "loss": 0.8048,
+      "output_text_mask_loss": 0.01,
+      "predictor_learning_rate": 0.0004076082461211538,
+      "step": 2971
+    },
+    {
+      "epoch": 0.3,
+      "gumbel_tau": 0.4961,
+      "image_masks_loss": 0.0184,
+      "llm_learning_rate": 4.0754383896941175e-06,
+      "loss": 0.7999,
+      "output_text_mask_loss": 0.0828,
+      "predictor_learning_rate": 0.0004075438389694117,
+      "step": 2972
+    },
+    {
+      "epoch": 0.3,
+      "gumbel_tau": 0.496,
+      "image_masks_loss": 0.0161,
+      "llm_learning_rate": 4.0747941446859115e-06,
+      "loss": 0.8224,
+      "output_text_mask_loss": 0.005,
+      "predictor_learning_rate": 0.00040747941446859115,
+      "step": 2973
+    },
+    {
+      "epoch": 0.3,
+      "gumbel_tau": 0.4959,
+      "image_masks_loss": 0.0248,
+      "llm_learning_rate": 4.074149726257868e-06,
+      "loss": 0.8457,
+      "output_text_mask_loss": 0.0278,
+      "predictor_learning_rate": 0.00040741497262578675,
+      "step": 2974
+    },
+    {
+      "epoch": 0.3,
+      "gumbel_tau": 0.4958,
+      "image_masks_loss": 0.0078,
+      "llm_learning_rate": 4.07350513448095e-06,
+      "loss": 0.8509,
+      "output_text_mask_loss": 0.017,
+      "predictor_learning_rate": 0.0004073505134480949,
+      "step": 2975
+    },
+    {
+      "epoch": 0.3,
+      "gumbel_tau": 0.4957,
+      "image_masks_loss": 0.0058,
+      "llm_learning_rate": 4.0728603694261416e-06,
+      "loss": 0.7822,
+      "output_text_mask_loss": 0.0322,
+      "predictor_learning_rate": 0.00040728603694261415,
+      "step": 2976
+    },
+    {
+      "epoch": 0.31,
+      "gumbel_tau": 0.4955,
+      "image_masks_loss": 0.0155,
+      "llm_learning_rate": 4.072215431164445e-06,
+      "loss": 0.8262,
+      "output_text_mask_loss": 0.0322,
+      "predictor_learning_rate": 0.0004072215431164445,
+      "step": 2977
+    },
+    {
+      "epoch": 0.31,
+      "gumbel_tau": 0.4954,
+      "image_masks_loss": 0.0057,
+      "llm_learning_rate": 4.071570319766885e-06,
+      "loss": 0.7689,
+      "output_text_mask_loss": 0.074,
+      "predictor_learning_rate": 0.0004071570319766884,
+      "step": 2978
+    },
+    {
+      "epoch": 0.31,
+      "gumbel_tau": 0.4953,
+      "image_masks_loss": 0.023,
+      "llm_learning_rate": 4.070925035304498e-06,
+      "loss": 0.7892,
+      "output_text_mask_loss": 0.0288,
+      "predictor_learning_rate": 0.0004070925035304498,
+      "step": 2979
+    },
+    {
+      "epoch": 0.31,
+      "gumbel_tau": 0.4952,
+      "image_masks_loss": 0.0211,
+      "llm_learning_rate": 4.07027957784835e-06,
+      "loss": 0.7356,
+      "output_text_mask_loss": 0.0231,
+      "predictor_learning_rate": 0.0004070279577848349,
+      "step": 2980
+    },
+    {
+      "epoch": 0.31,
+      "gumbel_tau": 0.4951,
+      "image_masks_loss": 0.0088,
+      "llm_learning_rate": 4.069633947469515e-06,
+      "loss": 0.7631,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004069633947469515,
+      "step": 2981
+    },
+    {
+      "epoch": 0.31,
+      "gumbel_tau": 0.495,
+      "image_masks_loss": 0.0108,
+      "llm_learning_rate": 4.068988144239094e-06,
+      "loss": 0.8154,
+      "output_text_mask_loss": 0.0515,
+      "predictor_learning_rate": 0.00040689881442390944,
+      "step": 2982
+    },
+    {
+      "epoch": 0.31,
+      "gumbel_tau": 0.4948,
+      "image_masks_loss": 0.0116,
+      "llm_learning_rate": 4.068342168228204e-06,
+      "loss": 0.7679,
+      "output_text_mask_loss": 0.0969,
+      "predictor_learning_rate": 0.0004068342168228204,
+      "step": 2983
+    },
+    {
+      "epoch": 0.31,
+      "gumbel_tau": 0.4947,
+      "image_masks_loss": 0.0167,
+      "llm_learning_rate": 4.067696019507983e-06,
+      "loss": 0.7376,
+      "output_text_mask_loss": 0.0156,
+      "predictor_learning_rate": 0.0004067696019507982,
+      "step": 2984
+    },
+    {
+      "epoch": 0.31,
+      "gumbel_tau": 0.4946,
+      "image_masks_loss": 0.0179,
+      "llm_learning_rate": 4.0670496981495825e-06,
+      "loss": 0.7085,
+      "output_text_mask_loss": 0.116,
+      "predictor_learning_rate": 0.0004067049698149582,
+      "step": 2985
+    },
+    {
+      "epoch": 0.31,
+      "gumbel_tau": 0.4945,
+      "image_masks_loss": 0.0148,
+      "llm_learning_rate": 4.066403204224178e-06,
+      "loss": 0.8139,
+      "output_text_mask_loss": 0.116,
+      "predictor_learning_rate": 0.0004066403204224178,
+      "step": 2986
+    },
+    {
+      "epoch": 0.31,
+      "gumbel_tau": 0.4944,
+      "image_masks_loss": 0.0038,
+      "llm_learning_rate": 4.065756537802965e-06,
+      "loss": 0.7606,
+      "output_text_mask_loss": 0.0307,
+      "predictor_learning_rate": 0.0004065756537802965,
+      "step": 2987
+    },
+    {
+      "epoch": 0.31,
+      "gumbel_tau": 0.4943,
+      "image_masks_loss": 0.025,
+      "llm_learning_rate": 4.065109698957154e-06,
+      "loss": 0.8079,
+      "output_text_mask_loss": 0.0271,
+      "predictor_learning_rate": 0.00040651096989571535,
+      "step": 2988
+    },
+    {
+      "epoch": 0.31,
+      "gumbel_tau": 0.4941,
+      "image_masks_loss": 0.0103,
+      "llm_learning_rate": 4.064462687757976e-06,
+      "loss": 0.6593,
+      "output_text_mask_loss": NaN,
+      "predictor_learning_rate": 0.0004064462687757976,
+      "step": 2989
+    },
+    {
+      "epoch": 0.31,
+      "gumbel_tau": 0.494,
+      "image_masks_loss": 0.0127,
+      "llm_learning_rate": 4.063815504276684e-06,
+      "loss": 0.7638,
+      "output_text_mask_loss": 0.1251,
+      "predictor_learning_rate": 0.0004063815504276683,
+      "step": 2990
+    },
+    {
+      "epoch": 0.31,
+      "gumbel_tau": 0.4939,
+      "image_masks_loss": 0.0114,
+      "llm_learning_rate": 4.063168148584545e-06,
+      "loss": 0.7273,
+      "output_text_mask_loss": 0.0618,
+      "predictor_learning_rate": 0.0004063168148584544,
+      "step": 2991
+    },
+    {
+      "epoch": 0.31,
+      "gumbel_tau": 0.4938,
+      "image_masks_loss": 0.0088,
+      "llm_learning_rate": 4.062520620752846e-06,
+      "loss": 0.8113,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.00040625206207528454,
+      "step": 2992
+    },
+    {
+      "epoch": 0.31,
+      "gumbel_tau": 0.4937,
+      "image_masks_loss": 0.0169,
+      "llm_learning_rate": 4.061872920852897e-06,
+      "loss": 0.8413,
+      "output_text_mask_loss": 0.0156,
+      "predictor_learning_rate": 0.00040618729208528973,
+      "step": 2993
+    },
+    {
+      "epoch": 0.31,
+      "gumbel_tau": 0.4936,
+      "image_masks_loss": 0.0047,
+      "llm_learning_rate": 4.061225048956024e-06,
+      "loss": 0.7975,
+      "output_text_mask_loss": 0.0039,
+      "predictor_learning_rate": 0.0004061225048956023,
+      "step": 2994
+    },
+    {
+      "epoch": 0.31,
+      "gumbel_tau": 0.4934,
+      "image_masks_loss": 0.0078,
+      "llm_learning_rate": 4.0605770051335695e-06,
+      "loss": 0.7746,
+      "output_text_mask_loss": 0.0315,
+      "predictor_learning_rate": 0.0004060577005133569,
+      "step": 2995
+    },
+    {
+      "epoch": 0.31,
+      "gumbel_tau": 0.4933,
+      "image_masks_loss": 0.0196,
+      "llm_learning_rate": 4.0599287894569e-06,
+      "loss": 0.8365,
+      "output_text_mask_loss": 0.0103,
+      "predictor_learning_rate": 0.0004059928789456899,
+      "step": 2996
+    },
+    {
+      "epoch": 0.31,
+      "gumbel_tau": 0.4932,
+      "image_masks_loss": 0.0167,
+      "llm_learning_rate": 4.059280401997396e-06,
+      "loss": 0.6607,
+      "output_text_mask_loss": 0.0725,
+      "predictor_learning_rate": 0.0004059280401997396,
+      "step": 2997
+    },
+    {
+      "epoch": 0.31,
+      "gumbel_tau": 0.4931,
+      "image_masks_loss": 0.0191,
+      "llm_learning_rate": 4.058631842826462e-06,
+      "loss": 0.7515,
+      "output_text_mask_loss": 0.0078,
+      "predictor_learning_rate": 0.00040586318428264623,
+      "step": 2998
+    },
+    {
+      "epoch": 0.31,
+      "gumbel_tau": 0.493,
+      "image_masks_loss": 0.0187,
+      "llm_learning_rate": 4.057983112015519e-06,
+      "loss": 0.8684,
+      "output_text_mask_loss": 0.0282,
+      "predictor_learning_rate": 0.00040579831120155185,
+      "step": 2999
+    },
+    {
+      "epoch": 0.31,
+      "gumbel_tau": 0.4929,
+      "image_masks_loss": 0.0153,
+      "llm_learning_rate": 4.057334209636004e-06,
+      "loss": 0.7604,
+      "output_text_mask_loss": 0.0242,
+      "predictor_learning_rate": 0.0004057334209636004,
+      "step": 3000
+    },
+    {
+      "epoch": 0.31,
+      "gumbel_tau": 0.4927,
+      "image_masks_loss": 0.0239,
+      "llm_learning_rate": 4.056685135759377e-06,
+      "loss": 0.7398,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004056685135759377,
+      "step": 3001
+    },
+    {
+      "epoch": 0.31,
+      "gumbel_tau": 0.4926,
+      "image_masks_loss": 0.0176,
+      "llm_learning_rate": 4.056035890457116e-06,
+      "loss": 0.7541,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004056035890457116,
+      "step": 3002
+    },
+    {
+      "epoch": 0.31,
+      "gumbel_tau": 0.4925,
+      "image_masks_loss": 0.0259,
+      "llm_learning_rate": 4.055386473800718e-06,
+      "loss": 0.8098,
+      "output_text_mask_loss": 0.0237,
+      "predictor_learning_rate": 0.00040553864738007174,
+      "step": 3003
+    },
+    {
+      "epoch": 0.31,
+      "gumbel_tau": 0.4924,
+      "image_masks_loss": 0.0607,
+      "llm_learning_rate": 4.054736885861695e-06,
+      "loss": 0.6839,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00040547368858616953,
+      "step": 3004
+    },
+    {
+      "epoch": 0.31,
+      "gumbel_tau": 0.4923,
+      "image_masks_loss": 0.0248,
+      "llm_learning_rate": 4.054087126711585e-06,
+      "loss": 0.698,
+      "output_text_mask_loss": 0.0565,
+      "predictor_learning_rate": 0.0004054087126711584,
+      "step": 3005
+    },
+    {
+      "epoch": 0.31,
+      "gumbel_tau": 0.4922,
+      "image_masks_loss": 0.0196,
+      "llm_learning_rate": 4.053437196421937e-06,
+      "loss": 0.746,
+      "output_text_mask_loss": 0.0217,
+      "predictor_learning_rate": 0.0004053437196421936,
+      "step": 3006
+    },
+    {
+      "epoch": 0.31,
+      "gumbel_tau": 0.492,
+      "image_masks_loss": 0.0313,
+      "llm_learning_rate": 4.052787095064325e-06,
+      "loss": 0.7173,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 0.0004052787095064325,
+      "step": 3007
+    },
+    {
+      "epoch": 0.31,
+      "gumbel_tau": 0.4919,
+      "image_masks_loss": 0.0275,
+      "llm_learning_rate": 4.052136822710341e-06,
+      "loss": 0.7036,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 0.0004052136822710341,
+      "step": 3008
+    },
+    {
+      "epoch": 0.31,
+      "gumbel_tau": 0.4918,
+      "image_masks_loss": 0.0553,
+      "llm_learning_rate": 4.051486379431592e-06,
+      "loss": 0.8114,
+      "output_text_mask_loss": 0.0124,
+      "predictor_learning_rate": 0.00040514863794315916,
+      "step": 3009
+    },
+    {
+      "epoch": 0.31,
+      "gumbel_tau": 0.4917,
+      "image_masks_loss": 0.0237,
+      "llm_learning_rate": 4.050835765299707e-06,
+      "loss": 0.7776,
+      "output_text_mask_loss": 0.0023,
+      "predictor_learning_rate": 0.0004050835765299706,
+      "step": 3010
+    },
+    {
+      "epoch": 0.31,
+      "gumbel_tau": 0.4916,
+      "image_masks_loss": 0.0158,
+      "llm_learning_rate": 4.0501849803863325e-06,
+      "loss": 0.741,
+      "output_text_mask_loss": 0.0275,
+      "predictor_learning_rate": 0.0004050184980386332,
+      "step": 3011
+    },
+    {
+      "epoch": 0.31,
+      "gumbel_tau": 0.4915,
+      "image_masks_loss": 0.0177,
+      "llm_learning_rate": 4.049534024763136e-06,
+      "loss": 0.8568,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.00040495340247631356,
+      "step": 3012
+    },
+    {
+      "epoch": 0.31,
+      "gumbel_tau": 0.4914,
+      "image_masks_loss": 0.036,
+      "llm_learning_rate": 4.048882898501801e-06,
+      "loss": 0.8647,
+      "output_text_mask_loss": 0.1495,
+      "predictor_learning_rate": 0.00040488828985018005,
+      "step": 3013
+    },
+    {
+      "epoch": 0.31,
+      "gumbel_tau": 0.4912,
+      "image_masks_loss": 0.0242,
+      "llm_learning_rate": 4.048231601674031e-06,
+      "loss": 0.8323,
+      "output_text_mask_loss": 0.1129,
+      "predictor_learning_rate": 0.00040482316016740306,
+      "step": 3014
+    },
+    {
+      "epoch": 0.31,
+      "gumbel_tau": 0.4911,
+      "image_masks_loss": 0.0374,
+      "llm_learning_rate": 4.0475801343515485e-06,
+      "loss": 0.7294,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 0.00040475801343515486,
+      "step": 3015
+    },
+    {
+      "epoch": 0.31,
+      "gumbel_tau": 0.491,
+      "image_masks_loss": 0.0257,
+      "llm_learning_rate": 4.046928496606095e-06,
+      "loss": 0.7221,
+      "output_text_mask_loss": 0.0399,
+      "predictor_learning_rate": 0.00040469284966060946,
+      "step": 3016
+    },
+    {
+      "epoch": 0.31,
+      "gumbel_tau": 0.4909,
+      "image_masks_loss": 0.0374,
+      "llm_learning_rate": 4.046276688509429e-06,
+      "loss": 0.8928,
+      "output_text_mask_loss": 0.0584,
+      "predictor_learning_rate": 0.00040462766885094285,
+      "step": 3017
+    },
+    {
+      "epoch": 0.31,
+      "gumbel_tau": 0.4908,
+      "image_masks_loss": 0.0117,
+      "llm_learning_rate": 4.0456247101333305e-06,
+      "loss": 0.7116,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.000404562471013333,
+      "step": 3018
+    },
+    {
+      "epoch": 0.31,
+      "gumbel_tau": 0.4907,
+      "image_masks_loss": 0.0136,
+      "llm_learning_rate": 4.044972561549596e-06,
+      "loss": 0.7228,
+      "output_text_mask_loss": 0.0041,
+      "predictor_learning_rate": 0.0004044972561549596,
+      "step": 3019
+    },
+    {
+      "epoch": 0.31,
+      "gumbel_tau": 0.4905,
+      "image_masks_loss": 0.0085,
+      "llm_learning_rate": 4.044320242830041e-06,
+      "loss": 0.774,
+      "output_text_mask_loss": 0.0038,
+      "predictor_learning_rate": 0.0004044320242830041,
+      "step": 3020
+    },
+    {
+      "epoch": 0.31,
+      "gumbel_tau": 0.4904,
+      "image_masks_loss": 0.0191,
+      "llm_learning_rate": 4.043667754046502e-06,
+      "loss": 0.8234,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004043667754046502,
+      "step": 3021
+    },
+    {
+      "epoch": 0.31,
+      "gumbel_tau": 0.4903,
+      "image_masks_loss": 0.0194,
+      "llm_learning_rate": 4.0430150952708315e-06,
+      "loss": 0.6838,
+      "output_text_mask_loss": 0.0081,
+      "predictor_learning_rate": 0.0004043015095270831,
+      "step": 3022
+    },
+    {
+      "epoch": 0.31,
+      "gumbel_tau": 0.4902,
+      "image_masks_loss": 0.0193,
+      "llm_learning_rate": 4.042362266574902e-06,
+      "loss": 0.787,
+      "output_text_mask_loss": 0.0437,
+      "predictor_learning_rate": 0.0004042362266574902,
+      "step": 3023
+    },
+    {
+      "epoch": 0.31,
+      "gumbel_tau": 0.4901,
+      "image_masks_loss": 0.0259,
+      "llm_learning_rate": 4.041709268030603e-06,
+      "loss": 0.7802,
+      "output_text_mask_loss": 0.0702,
+      "predictor_learning_rate": 0.0004041709268030603,
+      "step": 3024
+    },
+    {
+      "epoch": 0.31,
+      "gumbel_tau": 0.49,
+      "image_masks_loss": 0.0198,
+      "llm_learning_rate": 4.041056099709847e-06,
+      "loss": 0.7831,
+      "output_text_mask_loss": 0.0278,
+      "predictor_learning_rate": 0.00040410560997098467,
+      "step": 3025
+    },
+    {
+      "epoch": 0.31,
+      "gumbel_tau": 0.4898,
+      "image_masks_loss": 0.0444,
+      "llm_learning_rate": 4.040402761684559e-06,
+      "loss": 0.788,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.0004040402761684559,
+      "step": 3026
+    },
+    {
+      "epoch": 0.31,
+      "gumbel_tau": 0.4897,
+      "image_masks_loss": 0.0142,
+      "llm_learning_rate": 4.039749254026689e-06,
+      "loss": 0.6978,
+      "output_text_mask_loss": 0.0322,
+      "predictor_learning_rate": 0.0004039749254026689,
+      "step": 3027
+    },
+    {
+      "epoch": 0.31,
+      "gumbel_tau": 0.4896,
+      "image_masks_loss": 0.0082,
+      "llm_learning_rate": 4.039095576808201e-06,
+      "loss": 0.7768,
+      "output_text_mask_loss": 0.0519,
+      "predictor_learning_rate": 0.0004039095576808201,
+      "step": 3028
+    },
+    {
+      "epoch": 0.31,
+      "gumbel_tau": 0.4895,
+      "image_masks_loss": 0.014,
+      "llm_learning_rate": 4.03844173010108e-06,
+      "loss": 0.7683,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.000403844173010108,
+      "step": 3029
+    },
+    {
+      "epoch": 0.31,
+      "gumbel_tau": 0.4894,
+      "image_masks_loss": 0.0201,
+      "llm_learning_rate": 4.037787713977329e-06,
+      "loss": 0.7796,
+      "output_text_mask_loss": 0.0126,
+      "predictor_learning_rate": 0.0004037787713977329,
+      "step": 3030
+    },
+    {
+      "epoch": 0.31,
+      "gumbel_tau": 0.4893,
+      "image_masks_loss": 0.0201,
+      "llm_learning_rate": 4.03713352850897e-06,
+      "loss": 0.8386,
+      "output_text_mask_loss": 0.0025,
+      "predictor_learning_rate": 0.0004037133528508969,
+      "step": 3031
+    },
+    {
+      "epoch": 0.31,
+      "gumbel_tau": 0.4892,
+      "image_masks_loss": 0.0073,
+      "llm_learning_rate": 4.036479173768043e-06,
+      "loss": 0.789,
+      "output_text_mask_loss": 0.0322,
+      "predictor_learning_rate": 0.0004036479173768042,
+      "step": 3032
+    },
+    {
+      "epoch": 0.31,
+      "gumbel_tau": 0.489,
+      "image_masks_loss": 0.0154,
+      "llm_learning_rate": 4.035824649826606e-06,
+      "loss": 0.7447,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.0004035824649826606,
+      "step": 3033
+    },
+    {
+      "epoch": 0.31,
+      "gumbel_tau": 0.4889,
+      "image_masks_loss": 0.0124,
+      "llm_learning_rate": 4.035169956756738e-06,
+      "loss": 0.733,
+      "output_text_mask_loss": 0.0385,
+      "predictor_learning_rate": 0.0004035169956756738,
+      "step": 3034
+    },
+    {
+      "epoch": 0.31,
+      "gumbel_tau": 0.4888,
+      "image_masks_loss": 0.0154,
+      "llm_learning_rate": 4.0345150946305356e-06,
+      "loss": 0.8075,
+      "output_text_mask_loss": 0.0374,
+      "predictor_learning_rate": 0.00040345150946305354,
+      "step": 3035
+    },
+    {
+      "epoch": 0.31,
+      "gumbel_tau": 0.4887,
+      "image_masks_loss": 0.024,
+      "llm_learning_rate": 4.033860063520112e-06,
+      "loss": 0.7278,
+      "output_text_mask_loss": 0.0322,
+      "predictor_learning_rate": 0.00040338600635201125,
+      "step": 3036
+    },
+    {
+      "epoch": 0.31,
+      "gumbel_tau": 0.4886,
+      "image_masks_loss": 0.0121,
+      "llm_learning_rate": 4.033204863497603e-06,
+      "loss": 0.7638,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00040332048634976025,
+      "step": 3037
+    },
+    {
+      "epoch": 0.31,
+      "gumbel_tau": 0.4885,
+      "image_masks_loss": 0.024,
+      "llm_learning_rate": 4.032549494635159e-06,
+      "loss": 0.8132,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 0.0004032549494635158,
+      "step": 3038
+    },
+    {
+      "epoch": 0.31,
+      "gumbel_tau": 0.4883,
+      "image_masks_loss": 0.0163,
+      "llm_learning_rate": 4.0318939570049495e-06,
+      "loss": 0.7195,
+      "output_text_mask_loss": 0.0126,
+      "predictor_learning_rate": 0.00040318939570049497,
+      "step": 3039
+    },
+    {
+      "epoch": 0.31,
+      "gumbel_tau": 0.4882,
+      "image_masks_loss": 0.014,
+      "llm_learning_rate": 4.031238250679167e-06,
+      "loss": 0.7879,
+      "output_text_mask_loss": 0.0309,
+      "predictor_learning_rate": 0.0004031238250679167,
+      "step": 3040
+    },
+    {
+      "epoch": 0.31,
+      "gumbel_tau": 0.4881,
+      "image_masks_loss": 0.018,
+      "llm_learning_rate": 4.030582375730016e-06,
+      "loss": 0.759,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 0.00040305823757300167,
+      "step": 3041
+    },
+    {
+      "epoch": 0.31,
+      "gumbel_tau": 0.488,
+      "image_masks_loss": 0.0135,
+      "llm_learning_rate": 4.029926332229727e-06,
+      "loss": 0.761,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 0.0004029926332229726,
+      "step": 3042
+    },
+    {
+      "epoch": 0.31,
+      "gumbel_tau": 0.4879,
+      "image_masks_loss": 0.0085,
+      "llm_learning_rate": 4.029270120250541e-06,
+      "loss": 0.7352,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00040292701202505405,
+      "step": 3043
+    },
+    {
+      "epoch": 0.31,
+      "gumbel_tau": 0.4878,
+      "image_masks_loss": 0.0225,
+      "llm_learning_rate": 4.028613739864723e-06,
+      "loss": 0.822,
+      "output_text_mask_loss": 0.0025,
+      "predictor_learning_rate": 0.0004028613739864723,
+      "step": 3044
+    },
+    {
+      "epoch": 0.31,
+      "gumbel_tau": 0.4877,
+      "image_masks_loss": 0.028,
+      "llm_learning_rate": 4.027957191144555e-06,
+      "loss": 0.8546,
+      "output_text_mask_loss": 0.0549,
+      "predictor_learning_rate": 0.00040279571911445544,
+      "step": 3045
+    },
+    {
+      "epoch": 0.31,
+      "gumbel_tau": 0.4875,
+      "image_masks_loss": 0.0171,
+      "llm_learning_rate": 4.0273004741623375e-06,
+      "loss": 0.8538,
+      "output_text_mask_loss": 0.1167,
+      "predictor_learning_rate": 0.00040273004741623377,
+      "step": 3046
+    },
+    {
+      "epoch": 0.31,
+      "gumbel_tau": 0.4874,
+      "image_masks_loss": 0.0199,
+      "llm_learning_rate": 4.0266435889903906e-06,
+      "loss": 0.7419,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004026643588990391,
+      "step": 3047
+    },
+    {
+      "epoch": 0.31,
+      "gumbel_tau": 0.4873,
+      "image_masks_loss": 0.0235,
+      "llm_learning_rate": 4.025986535701051e-06,
+      "loss": 0.7928,
+      "output_text_mask_loss": 0.0275,
+      "predictor_learning_rate": 0.00040259865357010504,
+      "step": 3048
+    },
+    {
+      "epoch": 0.31,
+      "gumbel_tau": 0.4872,
+      "image_masks_loss": 0.0156,
+      "llm_learning_rate": 4.025329314366675e-06,
+      "loss": 0.7904,
+      "output_text_mask_loss": 0.0338,
+      "predictor_learning_rate": 0.00040253293143666753,
+      "step": 3049
+    },
+    {
+      "epoch": 0.31,
+      "gumbel_tau": 0.4871,
+      "image_masks_loss": 0.0209,
+      "llm_learning_rate": 4.024671925059637e-06,
+      "loss": 0.7656,
+      "output_text_mask_loss": 0.0368,
+      "predictor_learning_rate": 0.00040246719250596375,
+      "step": 3050
+    },
+    {
+      "epoch": 0.31,
+      "gumbel_tau": 0.487,
+      "image_masks_loss": 0.0103,
+      "llm_learning_rate": 4.024014367852331e-06,
+      "loss": 0.7561,
+      "output_text_mask_loss": 0.0504,
+      "predictor_learning_rate": 0.00040240143678523313,
+      "step": 3051
+    },
+    {
+      "epoch": 0.31,
+      "gumbel_tau": 0.4869,
+      "image_masks_loss": 0.0075,
+      "llm_learning_rate": 4.023356642817169e-06,
+      "loss": 0.6975,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.00040233566428171685,
+      "step": 3052
+    },
+    {
+      "epoch": 0.31,
+      "gumbel_tau": 0.4867,
+      "image_masks_loss": 0.0226,
+      "llm_learning_rate": 4.0226987500265794e-06,
+      "loss": 0.7676,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.00040226987500265786,
+      "step": 3053
+    },
+    {
+      "epoch": 0.31,
+      "gumbel_tau": 0.4866,
+      "image_masks_loss": 0.038,
+      "llm_learning_rate": 4.022040689553012e-06,
+      "loss": 0.853,
+      "output_text_mask_loss": 0.0231,
+      "predictor_learning_rate": 0.00040220406895530113,
+      "step": 3054
+    },
+    {
+      "epoch": 0.31,
+      "gumbel_tau": 0.4865,
+      "image_masks_loss": 0.0196,
+      "llm_learning_rate": 4.0213824614689335e-06,
+      "loss": 0.6732,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.0004021382461468933,
+      "step": 3055
+    },
+    {
+      "epoch": 0.31,
+      "gumbel_tau": 0.4864,
+      "image_masks_loss": 0.0147,
+      "llm_learning_rate": 4.0207240658468295e-06,
+      "loss": 0.8548,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0004020724065846829,
+      "step": 3056
+    },
+    {
+      "epoch": 0.31,
+      "gumbel_tau": 0.4863,
+      "image_masks_loss": 0.0242,
+      "llm_learning_rate": 4.0200655027592035e-06,
+      "loss": 0.7781,
+      "output_text_mask_loss": 0.0483,
+      "predictor_learning_rate": 0.00040200655027592037,
+      "step": 3057
+    },
+    {
+      "epoch": 0.31,
+      "gumbel_tau": 0.4862,
+      "image_masks_loss": 0.0211,
+      "llm_learning_rate": 4.019406772278579e-06,
+      "loss": 0.7703,
+      "output_text_mask_loss": 0.0565,
+      "predictor_learning_rate": 0.00040194067722785793,
+      "step": 3058
+    },
+    {
+      "epoch": 0.31,
+      "gumbel_tau": 0.486,
+      "image_masks_loss": 0.0179,
+      "llm_learning_rate": 4.018747874477498e-06,
+      "loss": 0.7632,
+      "output_text_mask_loss": 0.0124,
+      "predictor_learning_rate": 0.00040187478744774973,
+      "step": 3059
+    },
+    {
+      "epoch": 0.31,
+      "gumbel_tau": 0.4859,
+      "image_masks_loss": 0.0165,
+      "llm_learning_rate": 4.018088809428517e-06,
+      "loss": 0.6639,
+      "output_text_mask_loss": 0.0082,
+      "predictor_learning_rate": 0.0004018088809428516,
+      "step": 3060
+    },
+    {
+      "epoch": 0.31,
+      "gumbel_tau": 0.4858,
+      "image_masks_loss": 0.005,
+      "llm_learning_rate": 4.017429577204214e-06,
+      "loss": 0.8486,
+      "output_text_mask_loss": 0.1411,
+      "predictor_learning_rate": 0.0004017429577204214,
+      "step": 3061
+    },
+    {
+      "epoch": 0.31,
+      "gumbel_tau": 0.4857,
+      "image_masks_loss": 0.0031,
+      "llm_learning_rate": 4.016770177877186e-06,
+      "loss": 0.7538,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00040167701778771864,
+      "step": 3062
+    },
+    {
+      "epoch": 0.31,
+      "gumbel_tau": 0.4856,
+      "image_masks_loss": 0.0138,
+      "llm_learning_rate": 4.016110611520049e-06,
+      "loss": 0.682,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00040161106115200486,
+      "step": 3063
+    },
+    {
+      "epoch": 0.31,
+      "gumbel_tau": 0.4855,
+      "image_masks_loss": 0.0412,
+      "llm_learning_rate": 4.015450878205433e-06,
+      "loss": 0.823,
+      "output_text_mask_loss": 0.0231,
+      "predictor_learning_rate": 0.0004015450878205433,
+      "step": 3064
+    },
+    {
+      "epoch": 0.31,
+      "gumbel_tau": 0.4854,
+      "image_masks_loss": 0.0083,
+      "llm_learning_rate": 4.014790978005992e-06,
+      "loss": 0.7918,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.00040147909780059916,
+      "step": 3065
+    },
+    {
+      "epoch": 0.31,
+      "gumbel_tau": 0.4852,
+      "image_masks_loss": 0.0068,
+      "llm_learning_rate": 4.014130910994394e-06,
+      "loss": 0.8187,
+      "output_text_mask_loss": 0.1335,
+      "predictor_learning_rate": 0.00040141309109943934,
+      "step": 3066
+    },
+    {
+      "epoch": 0.31,
+      "gumbel_tau": 0.4851,
+      "image_masks_loss": 0.0233,
+      "llm_learning_rate": 4.013470677243327e-06,
+      "loss": 0.6967,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.00040134706772433274,
+      "step": 3067
+    },
+    {
+      "epoch": 0.31,
+      "gumbel_tau": 0.485,
+      "image_masks_loss": 0.028,
+      "llm_learning_rate": 4.012810276825498e-06,
+      "loss": 0.81,
+      "output_text_mask_loss": 0.0172,
+      "predictor_learning_rate": 0.0004012810276825498,
+      "step": 3068
+    },
+    {
+      "epoch": 0.31,
+      "gumbel_tau": 0.4849,
+      "image_masks_loss": 0.0127,
+      "llm_learning_rate": 4.012149709813633e-06,
+      "loss": 0.7597,
+      "output_text_mask_loss": 0.0162,
+      "predictor_learning_rate": 0.0004012149709813633,
+      "step": 3069
+    },
+    {
+      "epoch": 0.31,
+      "gumbel_tau": 0.4848,
+      "image_masks_loss": 0.0199,
+      "llm_learning_rate": 4.011488976280472e-06,
+      "loss": 0.7995,
+      "output_text_mask_loss": 0.0565,
+      "predictor_learning_rate": 0.0004011488976280472,
+      "step": 3070
+    },
+    {
+      "epoch": 0.31,
+      "gumbel_tau": 0.4847,
+      "image_masks_loss": 0.0296,
+      "llm_learning_rate": 4.010828076298779e-06,
+      "loss": 0.8094,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00040108280762987794,
+      "step": 3071
+    },
+    {
+      "epoch": 0.31,
+      "gumbel_tau": 0.4846,
+      "image_masks_loss": 0.0154,
+      "llm_learning_rate": 4.010167009941333e-06,
+      "loss": 0.7742,
+      "output_text_mask_loss": 0.02,
+      "predictor_learning_rate": 0.00040101670099413326,
+      "step": 3072
+    },
+    {
+      "epoch": 0.31,
+      "gumbel_tau": 0.4844,
+      "image_masks_loss": 0.0055,
+      "llm_learning_rate": 4.0095057772809324e-06,
+      "loss": 0.6807,
+      "output_text_mask_loss": 0.0854,
+      "predictor_learning_rate": 0.0004009505777280932,
+      "step": 3073
+    },
+    {
+      "epoch": 0.31,
+      "gumbel_tau": 0.4843,
+      "image_masks_loss": 0.0542,
+      "llm_learning_rate": 4.008844378390393e-06,
+      "loss": 0.7523,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.0004008844378390393,
+      "step": 3074
+    },
+    {
+      "epoch": 0.32,
+      "gumbel_tau": 0.4842,
+      "image_masks_loss": 0.006,
+      "llm_learning_rate": 4.00818281334255e-06,
+      "loss": 0.8316,
+      "output_text_mask_loss": 0.0034,
+      "predictor_learning_rate": 0.00040081828133425493,
+      "step": 3075
+    },
+    {
+      "epoch": 0.32,
+      "gumbel_tau": 0.4841,
+      "image_masks_loss": 0.0179,
+      "llm_learning_rate": 4.007521082210256e-06,
+      "loss": 0.8149,
+      "output_text_mask_loss": 0.0414,
+      "predictor_learning_rate": 0.00040075210822102555,
+      "step": 3076
+    },
+    {
+      "epoch": 0.32,
+      "gumbel_tau": 0.484,
+      "image_masks_loss": 0.0128,
+      "llm_learning_rate": 4.006859185066382e-06,
+      "loss": 0.7895,
+      "output_text_mask_loss": 0.0488,
+      "predictor_learning_rate": 0.0004006859185066382,
+      "step": 3077
+    },
+    {
+      "epoch": 0.32,
+      "gumbel_tau": 0.4839,
+      "image_masks_loss": 0.0222,
+      "llm_learning_rate": 4.006197121983819e-06,
+      "loss": 0.8762,
+      "output_text_mask_loss": 0.0534,
+      "predictor_learning_rate": 0.0004006197121983819,
+      "step": 3078
+    },
+    {
+      "epoch": 0.32,
+      "gumbel_tau": 0.4838,
+      "image_masks_loss": 0.0418,
+      "llm_learning_rate": 4.005534893035474e-06,
+      "loss": 0.7557,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00040055348930354733,
+      "step": 3079
+    },
+    {
+      "epoch": 0.32,
+      "gumbel_tau": 0.4836,
+      "image_masks_loss": 0.0043,
+      "llm_learning_rate": 4.004872498294272e-06,
+      "loss": 0.7079,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00040048724982942726,
+      "step": 3080
+    },
+    {
+      "epoch": 0.32,
+      "gumbel_tau": 0.4835,
+      "image_masks_loss": 0.0181,
+      "llm_learning_rate": 4.00420993783316e-06,
+      "loss": 0.7916,
+      "output_text_mask_loss": 0.0237,
+      "predictor_learning_rate": 0.00040042099378331597,
+      "step": 3081
+    },
+    {
+      "epoch": 0.32,
+      "gumbel_tau": 0.4834,
+      "image_masks_loss": 0.0131,
+      "llm_learning_rate": 4.0035472117250985e-06,
+      "loss": 0.8149,
+      "output_text_mask_loss": 0.0261,
+      "predictor_learning_rate": 0.0004003547211725098,
+      "step": 3082
+    },
+    {
+      "epoch": 0.32,
+      "gumbel_tau": 0.4833,
+      "image_masks_loss": 0.0073,
+      "llm_learning_rate": 4.002884320043069e-06,
+      "loss": 0.7577,
+      "output_text_mask_loss": 0.095,
+      "predictor_learning_rate": 0.0004002884320043069,
+      "step": 3083
+    },
+    {
+      "epoch": 0.32,
+      "gumbel_tau": 0.4832,
+      "image_masks_loss": 0.0058,
+      "llm_learning_rate": 4.002221262860071e-06,
+      "loss": 0.7625,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.0004002221262860071,
+      "step": 3084
+    },
+    {
+      "epoch": 0.32,
+      "gumbel_tau": 0.4831,
+      "image_masks_loss": 0.0091,
+      "llm_learning_rate": 4.00155804024912e-06,
+      "loss": 0.7433,
+      "output_text_mask_loss": 0.0322,
+      "predictor_learning_rate": 0.00040015580402491204,
+      "step": 3085
+    },
+    {
+      "epoch": 0.32,
+      "gumbel_tau": 0.483,
+      "image_masks_loss": 0.0064,
+      "llm_learning_rate": 4.000894652283255e-06,
+      "loss": 0.7571,
+      "output_text_mask_loss": 0.0162,
+      "predictor_learning_rate": 0.0004000894652283255,
+      "step": 3086
+    },
+    {
+      "epoch": 0.32,
+      "gumbel_tau": 0.4828,
+      "image_masks_loss": 0.0296,
+      "llm_learning_rate": 4.000231099035527e-06,
+      "loss": 0.835,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.0004000231099035527,
+      "step": 3087
+    },
+    {
+      "epoch": 0.32,
+      "gumbel_tau": 0.4827,
+      "image_masks_loss": 0.0141,
+      "llm_learning_rate": 3.999567380579009e-06,
+      "loss": 0.7765,
+      "output_text_mask_loss": 0.0839,
+      "predictor_learning_rate": 0.0003999567380579009,
+      "step": 3088
+    },
+    {
+      "epoch": 0.32,
+      "gumbel_tau": 0.4826,
+      "image_masks_loss": 0.0165,
+      "llm_learning_rate": 3.99890349698679e-06,
+      "loss": 0.8467,
+      "output_text_mask_loss": 0.0618,
+      "predictor_learning_rate": 0.000399890349698679,
+      "step": 3089
+    },
+    {
+      "epoch": 0.32,
+      "gumbel_tau": 0.4825,
+      "image_masks_loss": 0.0193,
+      "llm_learning_rate": 3.99823944833198e-06,
+      "loss": 0.7106,
+      "output_text_mask_loss": 0.0032,
+      "predictor_learning_rate": 0.000399823944833198,
+      "step": 3090
+    },
+    {
+      "epoch": 0.32,
+      "gumbel_tau": 0.4824,
+      "image_masks_loss": 0.0146,
+      "llm_learning_rate": 3.9975752346877054e-06,
+      "loss": 0.6698,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0003997575234687705,
+      "step": 3091
+    },
+    {
+      "epoch": 0.32,
+      "gumbel_tau": 0.4823,
+      "image_masks_loss": 0.0256,
+      "llm_learning_rate": 3.9969108561271085e-06,
+      "loss": 0.8614,
+      "output_text_mask_loss": 0.0345,
+      "predictor_learning_rate": 0.00039969108561271086,
+      "step": 3092
+    },
+    {
+      "epoch": 0.32,
+      "gumbel_tau": 0.4822,
+      "image_masks_loss": 0.0292,
+      "llm_learning_rate": 3.996246312723355e-06,
+      "loss": 0.8958,
+      "output_text_mask_loss": 0.0324,
+      "predictor_learning_rate": 0.0003996246312723355,
+      "step": 3093
+    },
+    {
+      "epoch": 0.32,
+      "gumbel_tau": 0.4821,
+      "image_masks_loss": 0.0093,
+      "llm_learning_rate": 3.9955816045496245e-06,
+      "loss": 0.723,
+      "output_text_mask_loss": 0.0496,
+      "predictor_learning_rate": 0.0003995581604549624,
+      "step": 3094
+    },
+    {
+      "epoch": 0.32,
+      "gumbel_tau": 0.4819,
+      "image_masks_loss": 0.0192,
+      "llm_learning_rate": 3.994916731679117e-06,
+      "loss": 0.8243,
+      "output_text_mask_loss": 0.1289,
+      "predictor_learning_rate": 0.00039949167316791163,
+      "step": 3095
+    },
+    {
+      "epoch": 0.32,
+      "gumbel_tau": 0.4818,
+      "image_masks_loss": 0.0267,
+      "llm_learning_rate": 3.994251694185048e-06,
+      "loss": 0.6927,
+      "output_text_mask_loss": 0.0178,
+      "predictor_learning_rate": 0.0003994251694185048,
+      "step": 3096
+    },
+    {
+      "epoch": 0.32,
+      "gumbel_tau": 0.4817,
+      "image_masks_loss": 0.0188,
+      "llm_learning_rate": 3.993586492140656e-06,
+      "loss": 0.6686,
+      "output_text_mask_loss": 0.0931,
+      "predictor_learning_rate": 0.00039935864921406553,
+      "step": 3097
+    },
+    {
+      "epoch": 0.32,
+      "gumbel_tau": 0.4816,
+      "image_masks_loss": 0.0216,
+      "llm_learning_rate": 3.992921125619191e-06,
+      "loss": 0.7922,
+      "output_text_mask_loss": 0.0931,
+      "predictor_learning_rate": 0.00039929211256191913,
+      "step": 3098
+    },
+    {
+      "epoch": 0.32,
+      "gumbel_tau": 0.4815,
+      "image_masks_loss": 0.0234,
+      "llm_learning_rate": 3.992255594693927e-06,
+      "loss": 0.8405,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.00039922555946939265,
+      "step": 3099
+    },
+    {
+      "epoch": 0.32,
+      "gumbel_tau": 0.4814,
+      "image_masks_loss": 0.0222,
+      "llm_learning_rate": 3.991589899438152e-06,
+      "loss": 0.784,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.0003991589899438152,
+      "step": 3100
+    },
+    {
+      "epoch": 0.32,
+      "gumbel_tau": 0.4813,
+      "image_masks_loss": 0.0175,
+      "llm_learning_rate": 3.990924039925176e-06,
+      "loss": 0.7104,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 0.0003990924039925176,
+      "step": 3101
+    },
+    {
+      "epoch": 0.32,
+      "gumbel_tau": 0.4811,
+      "image_masks_loss": 0.0043,
+      "llm_learning_rate": 3.9902580162283236e-06,
+      "loss": 0.7456,
+      "output_text_mask_loss": 0.0223,
+      "predictor_learning_rate": 0.0003990258016228323,
+      "step": 3102
+    },
+    {
+      "epoch": 0.32,
+      "gumbel_tau": 0.481,
+      "image_masks_loss": 0.02,
+      "llm_learning_rate": 3.989591828420939e-06,
+      "loss": 0.7738,
+      "output_text_mask_loss": 0.0185,
+      "predictor_learning_rate": 0.0003989591828420939,
+      "step": 3103
+    },
+    {
+      "epoch": 0.32,
+      "gumbel_tau": 0.4809,
+      "image_masks_loss": 0.009,
+      "llm_learning_rate": 3.9889254765763845e-06,
+      "loss": 0.8026,
+      "output_text_mask_loss": 0.0313,
+      "predictor_learning_rate": 0.0003988925476576384,
+      "step": 3104
+    },
+    {
+      "epoch": 0.32,
+      "gumbel_tau": 0.4808,
+      "image_masks_loss": 0.0357,
+      "llm_learning_rate": 3.98825896076804e-06,
+      "loss": 0.766,
+      "output_text_mask_loss": 0.0324,
+      "predictor_learning_rate": 0.00039882589607680396,
+      "step": 3105
+    },
+    {
+      "epoch": 0.32,
+      "gumbel_tau": 0.4807,
+      "image_masks_loss": 0.0097,
+      "llm_learning_rate": 3.987592281069303e-06,
+      "loss": 0.7194,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0003987592281069303,
+      "step": 3106
+    },
+    {
+      "epoch": 0.32,
+      "gumbel_tau": 0.4806,
+      "image_masks_loss": 0.0239,
+      "llm_learning_rate": 3.986925437553592e-06,
+      "loss": 0.7484,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.0003986925437553591,
+      "step": 3107
+    },
+    {
+      "epoch": 0.32,
+      "gumbel_tau": 0.4805,
+      "image_masks_loss": 0.0117,
+      "llm_learning_rate": 3.9862584302943385e-06,
+      "loss": 0.841,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.00039862584302943386,
+      "step": 3108
+    },
+    {
+      "epoch": 0.32,
+      "gumbel_tau": 0.4803,
+      "image_masks_loss": 0.0111,
+      "llm_learning_rate": 3.985591259364997e-06,
+      "loss": 0.8207,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.00039855912593649966,
+      "step": 3109
+    },
+    {
+      "epoch": 0.32,
+      "gumbel_tau": 0.4802,
+      "image_masks_loss": 0.0201,
+      "llm_learning_rate": 3.984923924839036e-06,
+      "loss": 0.7208,
+      "output_text_mask_loss": 0.005,
+      "predictor_learning_rate": 0.0003984923924839036,
+      "step": 3110
+    },
+    {
+      "epoch": 0.32,
+      "gumbel_tau": 0.4801,
+      "image_masks_loss": 0.0042,
+      "llm_learning_rate": 3.984256426789946e-06,
+      "loss": 0.7689,
+      "output_text_mask_loss": 0.0179,
+      "predictor_learning_rate": 0.00039842564267899454,
+      "step": 3111
+    },
+    {
+      "epoch": 0.32,
+      "gumbel_tau": 0.48,
+      "image_masks_loss": 0.0199,
+      "llm_learning_rate": 3.983588765291231e-06,
+      "loss": 0.7806,
+      "output_text_mask_loss": 0.0477,
+      "predictor_learning_rate": 0.0003983588765291231,
+      "step": 3112
+    },
+    {
+      "epoch": 0.32,
+      "gumbel_tau": 0.4799,
+      "image_masks_loss": 0.0148,
+      "llm_learning_rate": 3.982920940416418e-06,
+      "loss": 0.7537,
+      "output_text_mask_loss": 0.005,
+      "predictor_learning_rate": 0.0003982920940416418,
+      "step": 3113
+    },
+    {
+      "epoch": 0.32,
+      "gumbel_tau": 0.4798,
+      "image_masks_loss": 0.0067,
+      "llm_learning_rate": 3.9822529522390475e-06,
+      "loss": 0.5976,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.0003982252952239047,
+      "step": 3114
+    },
+    {
+      "epoch": 0.32,
+      "gumbel_tau": 0.4797,
+      "image_masks_loss": 0.0236,
+      "llm_learning_rate": 3.98158480083268e-06,
+      "loss": 0.7962,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.0003981584800832679,
+      "step": 3115
+    },
+    {
+      "epoch": 0.32,
+      "gumbel_tau": 0.4796,
+      "image_masks_loss": 0.0259,
+      "llm_learning_rate": 3.980916486270894e-06,
+      "loss": 0.7776,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0003980916486270894,
+      "step": 3116
+    },
+    {
+      "epoch": 0.32,
+      "gumbel_tau": 0.4794,
+      "image_masks_loss": 0.0239,
+      "llm_learning_rate": 3.9802480086272865e-06,
+      "loss": 0.7674,
+      "output_text_mask_loss": 0.0484,
+      "predictor_learning_rate": 0.00039802480086272864,
+      "step": 3117
+    },
+    {
+      "epoch": 0.32,
+      "gumbel_tau": 0.4793,
+      "image_masks_loss": 0.0115,
+      "llm_learning_rate": 3.9795793679754704e-06,
+      "loss": 0.7596,
+      "output_text_mask_loss": 0.1297,
+      "predictor_learning_rate": 0.000397957936797547,
+      "step": 3118
+    },
+    {
+      "epoch": 0.32,
+      "gumbel_tau": 0.4792,
+      "image_masks_loss": 0.0108,
+      "llm_learning_rate": 3.978910564389079e-06,
+      "loss": 0.7703,
+      "output_text_mask_loss": 0.0029,
+      "predictor_learning_rate": 0.0003978910564389079,
+      "step": 3119
+    },
+    {
+      "epoch": 0.32,
+      "gumbel_tau": 0.4791,
+      "image_masks_loss": 0.0081,
+      "llm_learning_rate": 3.978241597941762e-06,
+      "loss": 0.7137,
+      "output_text_mask_loss": 0.1129,
+      "predictor_learning_rate": 0.0003978241597941762,
+      "step": 3120
+    },
+    {
+      "epoch": 0.32,
+      "gumbel_tau": 0.479,
+      "image_masks_loss": 0.0204,
+      "llm_learning_rate": 3.977572468707188e-06,
+      "loss": 0.7889,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.0003977572468707188,
+      "step": 3121
+    },
+    {
+      "epoch": 0.32,
+      "gumbel_tau": 0.4789,
+      "image_masks_loss": 0.0095,
+      "llm_learning_rate": 3.976903176759043e-06,
+      "loss": 0.8116,
+      "output_text_mask_loss": 0.0433,
+      "predictor_learning_rate": 0.00039769031767590427,
+      "step": 3122
+    },
+    {
+      "epoch": 0.32,
+      "gumbel_tau": 0.4788,
+      "image_masks_loss": 0.0097,
+      "llm_learning_rate": 3.97623372217103e-06,
+      "loss": 0.7898,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.00039762337221710295,
+      "step": 3123
+    },
+    {
+      "epoch": 0.32,
+      "gumbel_tau": 0.4787,
+      "image_masks_loss": 0.033,
+      "llm_learning_rate": 3.975564105016871e-06,
+      "loss": 0.6234,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.0003975564105016871,
+      "step": 3124
+    },
+    {
+      "epoch": 0.32,
+      "gumbel_tau": 0.4785,
+      "image_masks_loss": 0.0269,
+      "llm_learning_rate": 3.974894325370306e-06,
+      "loss": 0.7093,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.00039748943253703055,
+      "step": 3125
+    },
+    {
+      "epoch": 0.32,
+      "gumbel_tau": 0.4784,
+      "image_masks_loss": 0.006,
+      "llm_learning_rate": 3.974224383305093e-06,
+      "loss": 0.7824,
+      "output_text_mask_loss": 0.0429,
+      "predictor_learning_rate": 0.0003974224383305093,
+      "step": 3126
+    },
+    {
+      "epoch": 0.32,
+      "gumbel_tau": 0.4783,
+      "image_masks_loss": 0.0248,
+      "llm_learning_rate": 3.973554278895006e-06,
+      "loss": 0.8083,
+      "output_text_mask_loss": 0.0446,
+      "predictor_learning_rate": 0.00039735542788950064,
+      "step": 3127
+    },
+    {
+      "epoch": 0.32,
+      "gumbel_tau": 0.4782,
+      "image_masks_loss": 0.0383,
+      "llm_learning_rate": 3.972884012213842e-06,
+      "loss": 0.7506,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0003972884012213841,
+      "step": 3128
+    },
+    {
+      "epoch": 0.32,
+      "gumbel_tau": 0.4781,
+      "image_masks_loss": 0.0117,
+      "llm_learning_rate": 3.972213583335408e-06,
+      "loss": 0.7407,
+      "output_text_mask_loss": 0.0149,
+      "predictor_learning_rate": 0.0003972213583335408,
+      "step": 3129
+    },
+    {
+      "epoch": 0.32,
+      "gumbel_tau": 0.478,
+      "image_masks_loss": 0.012,
+      "llm_learning_rate": 3.971542992333535e-06,
+      "loss": 0.7092,
+      "output_text_mask_loss": 0.0374,
+      "predictor_learning_rate": 0.0003971542992333536,
+      "step": 3130
+    },
+    {
+      "epoch": 0.32,
+      "gumbel_tau": 0.4779,
+      "image_masks_loss": 0.0192,
+      "llm_learning_rate": 3.970872239282071e-06,
+      "loss": 0.731,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.00039708722392820705,
+      "step": 3131
+    },
+    {
+      "epoch": 0.32,
+      "gumbel_tau": 0.4778,
+      "image_masks_loss": 0.0198,
+      "llm_learning_rate": 3.970201324254879e-06,
+      "loss": 0.714,
+      "output_text_mask_loss": 0.0444,
+      "predictor_learning_rate": 0.0003970201324254879,
+      "step": 3132
+    },
+    {
+      "epoch": 0.32,
+      "gumbel_tau": 0.4776,
+      "image_masks_loss": 0.017,
+      "llm_learning_rate": 3.969530247325843e-06,
+      "loss": 0.7624,
+      "output_text_mask_loss": 0.0095,
+      "predictor_learning_rate": 0.0003969530247325843,
+      "step": 3133
+    },
+    {
+      "epoch": 0.32,
+      "gumbel_tau": 0.4775,
+      "image_masks_loss": 0.0121,
+      "llm_learning_rate": 3.968859008568863e-06,
+      "loss": 0.7704,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00039688590085688625,
+      "step": 3134
+    },
+    {
+      "epoch": 0.32,
+      "gumbel_tau": 0.4774,
+      "image_masks_loss": 0.0138,
+      "llm_learning_rate": 3.968187608057857e-06,
+      "loss": 0.7815,
+      "output_text_mask_loss": 0.0135,
+      "predictor_learning_rate": 0.0003968187608057857,
+      "step": 3135
+    },
+    {
+      "epoch": 0.32,
+      "gumbel_tau": 0.4773,
+      "image_masks_loss": 0.011,
+      "llm_learning_rate": 3.9675160458667606e-06,
+      "loss": 0.8053,
+      "output_text_mask_loss": 0.0071,
+      "predictor_learning_rate": 0.000396751604586676,
+      "step": 3136
+    },
+    {
+      "epoch": 0.32,
+      "gumbel_tau": 0.4772,
+      "image_masks_loss": 0.0368,
+      "llm_learning_rate": 3.96684432206953e-06,
+      "loss": 0.7513,
+      "output_text_mask_loss": 0.0193,
+      "predictor_learning_rate": 0.00039668443220695297,
+      "step": 3137
+    },
+    {
+      "epoch": 0.32,
+      "gumbel_tau": 0.4771,
+      "image_masks_loss": 0.0267,
+      "llm_learning_rate": 3.966172436740134e-06,
+      "loss": 0.733,
+      "output_text_mask_loss": 0.0012,
+      "predictor_learning_rate": 0.00039661724367401335,
+      "step": 3138
+    },
+    {
+      "epoch": 0.32,
+      "gumbel_tau": 0.477,
+      "image_masks_loss": 0.0132,
+      "llm_learning_rate": 3.965500389952563e-06,
+      "loss": 0.5588,
+      "output_text_mask_loss": 0.0107,
+      "predictor_learning_rate": 0.0003965500389952563,
+      "step": 3139
+    },
+    {
+      "epoch": 0.32,
+      "gumbel_tau": 0.4768,
+      "image_masks_loss": 0.0184,
+      "llm_learning_rate": 3.9648281817808265e-06,
+      "loss": 0.7566,
+      "output_text_mask_loss": 0.0671,
+      "predictor_learning_rate": 0.0003964828181780826,
+      "step": 3140
+    },
+    {
+      "epoch": 0.32,
+      "gumbel_tau": 0.4767,
+      "image_masks_loss": 0.0248,
+      "llm_learning_rate": 3.964155812298947e-06,
+      "loss": 0.809,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 0.0003964155812298947,
+      "step": 3141
+    },
+    {
+      "epoch": 0.32,
+      "gumbel_tau": 0.4766,
+      "image_masks_loss": 0.0162,
+      "llm_learning_rate": 3.963483281580969e-06,
+      "loss": 0.6745,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00039634832815809687,
+      "step": 3142
+    },
+    {
+      "epoch": 0.32,
+      "gumbel_tau": 0.4765,
+      "image_masks_loss": 0.0119,
+      "llm_learning_rate": 3.962810589700951e-06,
+      "loss": 0.7656,
+      "output_text_mask_loss": 0.0416,
+      "predictor_learning_rate": 0.00039628105897009504,
+      "step": 3143
+    },
+    {
+      "epoch": 0.32,
+      "gumbel_tau": 0.4764,
+      "image_masks_loss": 0.0234,
+      "llm_learning_rate": 3.962137736732973e-06,
+      "loss": 0.7099,
+      "output_text_mask_loss": 0.0015,
+      "predictor_learning_rate": 0.00039621377367329727,
+      "step": 3144
+    },
+    {
+      "epoch": 0.32,
+      "gumbel_tau": 0.4763,
+      "image_masks_loss": 0.0184,
+      "llm_learning_rate": 3.961464722751131e-06,
+      "loss": 0.7533,
+      "output_text_mask_loss": 0.0038,
+      "predictor_learning_rate": 0.000396146472275113,
+      "step": 3145
+    },
+    {
+      "epoch": 0.32,
+      "gumbel_tau": 0.4762,
+      "image_masks_loss": 0.0134,
+      "llm_learning_rate": 3.960791547829538e-06,
+      "loss": 0.8018,
+      "output_text_mask_loss": 0.1778,
+      "predictor_learning_rate": 0.00039607915478295373,
+      "step": 3146
+    },
+    {
+      "epoch": 0.32,
+      "gumbel_tau": 0.4761,
+      "image_masks_loss": 0.015,
+      "llm_learning_rate": 3.960118212042326e-06,
+      "loss": 0.7733,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0003960118212042326,
+      "step": 3147
+    },
+    {
+      "epoch": 0.32,
+      "gumbel_tau": 0.476,
+      "image_masks_loss": 0.0292,
+      "llm_learning_rate": 3.959444715463644e-06,
+      "loss": 0.9242,
+      "output_text_mask_loss": 0.0557,
+      "predictor_learning_rate": 0.00039594447154636436,
+      "step": 3148
+    },
+    {
+      "epoch": 0.32,
+      "gumbel_tau": 0.4758,
+      "image_masks_loss": 0.0275,
+      "llm_learning_rate": 3.95877105816766e-06,
+      "loss": 0.756,
+      "output_text_mask_loss": 0.0311,
+      "predictor_learning_rate": 0.00039587710581676593,
+      "step": 3149
+    },
+    {
+      "epoch": 0.32,
+      "gumbel_tau": 0.4757,
+      "image_masks_loss": 0.0261,
+      "llm_learning_rate": 3.958097240228557e-06,
+      "loss": 0.7335,
+      "output_text_mask_loss": 0.0269,
+      "predictor_learning_rate": 0.0003958097240228556,
+      "step": 3150
+    },
+    {
+      "epoch": 0.32,
+      "gumbel_tau": 0.4756,
+      "image_masks_loss": 0.0237,
+      "llm_learning_rate": 3.957423261720539e-06,
+      "loss": 0.6857,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.00039574232617205386,
+      "step": 3151
+    },
+    {
+      "epoch": 0.32,
+      "gumbel_tau": 0.4755,
+      "image_masks_loss": 0.0235,
+      "llm_learning_rate": 3.956749122717824e-06,
+      "loss": 0.7678,
+      "output_text_mask_loss": 0.0029,
+      "predictor_learning_rate": 0.00039567491227178236,
+      "step": 3152
+    },
+    {
+      "epoch": 0.32,
+      "gumbel_tau": 0.4754,
+      "image_masks_loss": 0.0322,
+      "llm_learning_rate": 3.956074823294651e-06,
+      "loss": 0.8525,
+      "output_text_mask_loss": 0.0202,
+      "predictor_learning_rate": 0.0003956074823294651,
+      "step": 3153
+    },
+    {
+      "epoch": 0.32,
+      "gumbel_tau": 0.4753,
+      "image_masks_loss": 0.0087,
+      "llm_learning_rate": 3.955400363525276e-06,
+      "loss": 0.7553,
+      "output_text_mask_loss": 0.0328,
+      "predictor_learning_rate": 0.00039554003635252754,
+      "step": 3154
+    },
+    {
+      "epoch": 0.32,
+      "gumbel_tau": 0.4752,
+      "image_masks_loss": 0.0111,
+      "llm_learning_rate": 3.95472574348397e-06,
+      "loss": 0.7135,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.00039547257434839703,
+      "step": 3155
+    },
+    {
+      "epoch": 0.32,
+      "gumbel_tau": 0.4751,
+      "image_masks_loss": 0.02,
+      "llm_learning_rate": 3.954050963245026e-06,
+      "loss": 0.6838,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0003954050963245026,
+      "step": 3156
+    },
+    {
+      "epoch": 0.32,
+      "gumbel_tau": 0.4749,
+      "image_masks_loss": 0.0345,
+      "llm_learning_rate": 3.953376022882751e-06,
+      "loss": 0.7586,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0003953376022882751,
+      "step": 3157
+    },
+    {
+      "epoch": 0.32,
+      "gumbel_tau": 0.4748,
+      "image_masks_loss": 0.0103,
+      "llm_learning_rate": 3.9527009224714704e-06,
+      "loss": 0.8063,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.00039527009224714707,
+      "step": 3158
+    },
+    {
+      "epoch": 0.32,
+      "gumbel_tau": 0.4747,
+      "image_masks_loss": 0.0343,
+      "llm_learning_rate": 3.952025662085529e-06,
+      "loss": 0.6713,
+      "output_text_mask_loss": 0.058,
+      "predictor_learning_rate": 0.0003952025662085529,
+      "step": 3159
+    },
+    {
+      "epoch": 0.32,
+      "gumbel_tau": 0.4746,
+      "image_masks_loss": 0.0052,
+      "llm_learning_rate": 3.951350241799288e-06,
+      "loss": 0.7171,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.00039513502417992875,
+      "step": 3160
+    },
+    {
+      "epoch": 0.32,
+      "gumbel_tau": 0.4745,
+      "image_masks_loss": 0.0073,
+      "llm_learning_rate": 3.950674661687125e-06,
+      "loss": 0.7715,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.00039506746616871244,
+      "step": 3161
+    },
+    {
+      "epoch": 0.32,
+      "gumbel_tau": 0.4744,
+      "image_masks_loss": 0.0204,
+      "llm_learning_rate": 3.949998921823436e-06,
+      "loss": 0.9248,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.0003949998921823436,
+      "step": 3162
+    },
+    {
+      "epoch": 0.32,
+      "gumbel_tau": 0.4743,
+      "image_masks_loss": 0.0105,
+      "llm_learning_rate": 3.9493230222826374e-06,
+      "loss": 0.738,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0003949323022282637,
+      "step": 3163
+    },
+    {
+      "epoch": 0.32,
+      "gumbel_tau": 0.4742,
+      "image_masks_loss": 0.0441,
+      "llm_learning_rate": 3.948646963139159e-06,
+      "loss": 0.8358,
+      "output_text_mask_loss": 0.0004,
+      "predictor_learning_rate": 0.00039486469631391585,
+      "step": 3164
+    },
+    {
+      "epoch": 0.32,
+      "gumbel_tau": 0.474,
+      "image_masks_loss": 0.0231,
+      "llm_learning_rate": 3.9479707444674494e-06,
+      "loss": 0.6352,
+      "output_text_mask_loss": 0.0309,
+      "predictor_learning_rate": 0.00039479707444674495,
+      "step": 3165
+    },
+    {
+      "epoch": 0.32,
+      "gumbel_tau": 0.4739,
+      "image_masks_loss": 0.0158,
+      "llm_learning_rate": 3.947294366341977e-06,
+      "loss": 0.7765,
+      "output_text_mask_loss": 0.0111,
+      "predictor_learning_rate": 0.0003947294366341977,
+      "step": 3166
+    },
+    {
+      "epoch": 0.32,
+      "gumbel_tau": 0.4738,
+      "image_masks_loss": 0.0177,
+      "llm_learning_rate": 3.9466178288372255e-06,
+      "loss": 0.7088,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0003946617828837225,
+      "step": 3167
+    },
+    {
+      "epoch": 0.32,
+      "gumbel_tau": 0.4737,
+      "image_masks_loss": 0.0048,
+      "llm_learning_rate": 3.945941132027697e-06,
+      "loss": 0.7223,
+      "output_text_mask_loss": 0.0097,
+      "predictor_learning_rate": 0.0003945941132027696,
+      "step": 3168
+    },
+    {
+      "epoch": 0.32,
+      "gumbel_tau": 0.4736,
+      "image_masks_loss": 0.0064,
+      "llm_learning_rate": 3.945264275987908e-06,
+      "loss": 0.7314,
+      "output_text_mask_loss": 0.0185,
+      "predictor_learning_rate": 0.0003945264275987908,
+      "step": 3169
+    },
+    {
+      "epoch": 0.32,
+      "gumbel_tau": 0.4735,
+      "image_masks_loss": 0.0366,
+      "llm_learning_rate": 3.9445872607924e-06,
+      "loss": 0.729,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.00039445872607923996,
+      "step": 3170
+    },
+    {
+      "epoch": 0.32,
+      "gumbel_tau": 0.4734,
+      "image_masks_loss": 0.0168,
+      "llm_learning_rate": 3.9439100865157245e-06,
+      "loss": 0.7761,
+      "output_text_mask_loss": 0.021,
+      "predictor_learning_rate": 0.0003943910086515724,
+      "step": 3171
+    },
+    {
+      "epoch": 0.33,
+      "gumbel_tau": 0.4733,
+      "image_masks_loss": 0.005,
+      "llm_learning_rate": 3.943232753232454e-06,
+      "loss": 0.757,
+      "output_text_mask_loss": 0.0141,
+      "predictor_learning_rate": 0.00039432327532324535,
+      "step": 3172
+    },
+    {
+      "epoch": 0.33,
+      "gumbel_tau": 0.4732,
+      "image_masks_loss": 0.019,
+      "llm_learning_rate": 3.942555261017178e-06,
+      "loss": 0.8314,
+      "output_text_mask_loss": 0.0614,
+      "predictor_learning_rate": 0.0003942555261017178,
+      "step": 3173
+    },
+    {
+      "epoch": 0.33,
+      "gumbel_tau": 0.473,
+      "image_masks_loss": 0.0271,
+      "llm_learning_rate": 3.9418776099445035e-06,
+      "loss": 0.7555,
+      "output_text_mask_loss": 0.0202,
+      "predictor_learning_rate": 0.0003941877609944503,
+      "step": 3174
+    },
+    {
+      "epoch": 0.33,
+      "gumbel_tau": 0.4729,
+      "image_masks_loss": 0.0067,
+      "llm_learning_rate": 3.941199800089054e-06,
+      "loss": 0.7815,
+      "output_text_mask_loss": 0.0034,
+      "predictor_learning_rate": 0.0003941199800089054,
+      "step": 3175
+    },
+    {
+      "epoch": 0.33,
+      "gumbel_tau": 0.4728,
+      "image_masks_loss": 0.0154,
+      "llm_learning_rate": 3.940521831525474e-06,
+      "loss": 0.7399,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00039405218315254737,
+      "step": 3176
+    },
+    {
+      "epoch": 0.33,
+      "gumbel_tau": 0.4727,
+      "image_masks_loss": 0.0115,
+      "llm_learning_rate": 3.93984370432842e-06,
+      "loss": 0.8393,
+      "output_text_mask_loss": 0.0071,
+      "predictor_learning_rate": 0.000393984370432842,
+      "step": 3177
+    },
+    {
+      "epoch": 0.33,
+      "gumbel_tau": 0.4726,
+      "image_masks_loss": 0.0126,
+      "llm_learning_rate": 3.9391654185725714e-06,
+      "loss": 0.655,
+      "output_text_mask_loss": 0.0055,
+      "predictor_learning_rate": 0.0003939165418572571,
+      "step": 3178
+    },
+    {
+      "epoch": 0.33,
+      "gumbel_tau": 0.4725,
+      "image_masks_loss": 0.01,
+      "llm_learning_rate": 3.9384869743326206e-06,
+      "loss": 0.7909,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.000393848697433262,
+      "step": 3179
+    },
+    {
+      "epoch": 0.33,
+      "gumbel_tau": 0.4724,
+      "image_masks_loss": 0.0173,
+      "llm_learning_rate": 3.93780837168328e-06,
+      "loss": 0.6907,
+      "output_text_mask_loss": 0.0645,
+      "predictor_learning_rate": 0.00039378083716832803,
+      "step": 3180
+    },
+    {
+      "epoch": 0.33,
+      "gumbel_tau": 0.4723,
+      "image_masks_loss": 0.012,
+      "llm_learning_rate": 3.937129610699279e-06,
+      "loss": 0.8001,
+      "output_text_mask_loss": 0.0153,
+      "predictor_learning_rate": 0.00039371296106992785,
+      "step": 3181
+    },
+    {
+      "epoch": 0.33,
+      "gumbel_tau": 0.4721,
+      "image_masks_loss": 0.0217,
+      "llm_learning_rate": 3.936450691455364e-06,
+      "loss": 0.6726,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.00039364506914553644,
+      "step": 3182
+    },
+    {
+      "epoch": 0.33,
+      "gumbel_tau": 0.472,
+      "image_masks_loss": 0.0166,
+      "llm_learning_rate": 3.9357716140263e-06,
+      "loss": 0.8091,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 0.00039357716140262997,
+      "step": 3183
+    },
+    {
+      "epoch": 0.33,
+      "gumbel_tau": 0.4719,
+      "image_masks_loss": 0.0048,
+      "llm_learning_rate": 3.935092378486867e-06,
+      "loss": 0.7683,
+      "output_text_mask_loss": 0.0101,
+      "predictor_learning_rate": 0.00039350923784868666,
+      "step": 3184
+    },
+    {
+      "epoch": 0.33,
+      "gumbel_tau": 0.4718,
+      "image_masks_loss": 0.0157,
+      "llm_learning_rate": 3.934412984911865e-06,
+      "loss": 0.8142,
+      "output_text_mask_loss": 0.0038,
+      "predictor_learning_rate": 0.0003934412984911865,
+      "step": 3185
+    },
+    {
+      "epoch": 0.33,
+      "gumbel_tau": 0.4717,
+      "image_masks_loss": 0.0192,
+      "llm_learning_rate": 3.93373343337611e-06,
+      "loss": 0.7685,
+      "output_text_mask_loss": 0.0294,
+      "predictor_learning_rate": 0.000393373343337611,
+      "step": 3186
+    },
+    {
+      "epoch": 0.33,
+      "gumbel_tau": 0.4716,
+      "image_masks_loss": 0.0452,
+      "llm_learning_rate": 3.9330537239544355e-06,
+      "loss": 0.8291,
+      "output_text_mask_loss": 0.0429,
+      "predictor_learning_rate": 0.00039330537239544357,
+      "step": 3187
+    },
+    {
+      "epoch": 0.33,
+      "gumbel_tau": 0.4715,
+      "image_masks_loss": 0.0056,
+      "llm_learning_rate": 3.932373856721694e-06,
+      "loss": 0.8697,
+      "output_text_mask_loss": 0.0004,
+      "predictor_learning_rate": 0.0003932373856721694,
+      "step": 3188
+    },
+    {
+      "epoch": 0.33,
+      "gumbel_tau": 0.4714,
+      "image_masks_loss": 0.0227,
+      "llm_learning_rate": 3.931693831752752e-06,
+      "loss": 0.7525,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.00039316938317527516,
+      "step": 3189
+    },
+    {
+      "epoch": 0.33,
+      "gumbel_tau": 0.4713,
+      "image_masks_loss": 0.0057,
+      "llm_learning_rate": 3.931013649122497e-06,
+      "loss": 0.7274,
+      "output_text_mask_loss": 0.0565,
+      "predictor_learning_rate": 0.0003931013649122497,
+      "step": 3190
+    },
+    {
+      "epoch": 0.33,
+      "gumbel_tau": 0.4711,
+      "image_masks_loss": 0.0076,
+      "llm_learning_rate": 3.93033330890583e-06,
+      "loss": 0.752,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.00039303333089058304,
+      "step": 3191
+    },
+    {
+      "epoch": 0.33,
+      "gumbel_tau": 0.471,
+      "image_masks_loss": 0.0259,
+      "llm_learning_rate": 3.929652811177674e-06,
+      "loss": 0.7447,
+      "output_text_mask_loss": 0.0832,
+      "predictor_learning_rate": 0.0003929652811177674,
+      "step": 3192
+    },
+    {
+      "epoch": 0.33,
+      "gumbel_tau": 0.4709,
+      "image_masks_loss": 0.0267,
+      "llm_learning_rate": 3.928972156012967e-06,
+      "loss": 0.7693,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.00039289721560129665,
+      "step": 3193
+    },
+    {
+      "epoch": 0.33,
+      "gumbel_tau": 0.4708,
+      "image_masks_loss": 0.0112,
+      "llm_learning_rate": 3.928291343486662e-06,
+      "loss": 0.6966,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.0003928291343486662,
+      "step": 3194
+    },
+    {
+      "epoch": 0.33,
+      "gumbel_tau": 0.4707,
+      "image_masks_loss": 0.0026,
+      "llm_learning_rate": 3.927610373673734e-06,
+      "loss": 0.7241,
+      "output_text_mask_loss": 0.0011,
+      "predictor_learning_rate": 0.00039276103736737334,
+      "step": 3195
+    },
+    {
+      "epoch": 0.33,
+      "gumbel_tau": 0.4706,
+      "image_masks_loss": 0.0075,
+      "llm_learning_rate": 3.92692924664917e-06,
+      "loss": 0.8373,
+      "output_text_mask_loss": 0.0275,
+      "predictor_learning_rate": 0.000392692924664917,
+      "step": 3196
+    },
+    {
+      "epoch": 0.33,
+      "gumbel_tau": 0.4705,
+      "image_masks_loss": 0.0261,
+      "llm_learning_rate": 3.92624796248798e-06,
+      "loss": 0.7692,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.00039262479624879796,
+      "step": 3197
+    },
+    {
+      "epoch": 0.33,
+      "gumbel_tau": 0.4704,
+      "image_masks_loss": 0.0062,
+      "llm_learning_rate": 3.925566521265188e-06,
+      "loss": 0.7927,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0003925566521265188,
+      "step": 3198
+    },
+    {
+      "epoch": 0.33,
+      "gumbel_tau": 0.4703,
+      "image_masks_loss": 0.0097,
+      "llm_learning_rate": 3.924884923055835e-06,
+      "loss": 0.7215,
+      "output_text_mask_loss": 0.0156,
+      "predictor_learning_rate": 0.0003924884923055835,
+      "step": 3199
+    },
+    {
+      "epoch": 0.33,
+      "gumbel_tau": 0.4701,
+      "image_masks_loss": 0.0214,
+      "llm_learning_rate": 3.92420316793498e-06,
+      "loss": 0.784,
+      "output_text_mask_loss": 0.0259,
+      "predictor_learning_rate": 0.000392420316793498,
+      "step": 3200
+    },
+    {
+      "epoch": 0.33,
+      "gumbel_tau": 0.47,
+      "image_masks_loss": 0.0222,
+      "llm_learning_rate": 3.9235212559777e-06,
+      "loss": 0.7842,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00039235212559777,
+      "step": 3201
+    },
+    {
+      "epoch": 0.33,
+      "gumbel_tau": 0.4699,
+      "image_masks_loss": 0.0183,
+      "llm_learning_rate": 3.922839187259089e-06,
+      "loss": 0.742,
+      "output_text_mask_loss": 0.0809,
+      "predictor_learning_rate": 0.0003922839187259088,
+      "step": 3202
+    },
+    {
+      "epoch": 0.33,
+      "gumbel_tau": 0.4698,
+      "image_masks_loss": 0.008,
+      "llm_learning_rate": 3.922156961854257e-06,
+      "loss": 0.718,
+      "output_text_mask_loss": 0.0124,
+      "predictor_learning_rate": 0.00039221569618542567,
+      "step": 3203
+    },
+    {
+      "epoch": 0.33,
+      "gumbel_tau": 0.4697,
+      "image_masks_loss": 0.0211,
+      "llm_learning_rate": 3.921474579838333e-06,
+      "loss": 0.6349,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.0003921474579838333,
+      "step": 3204
+    },
+    {
+      "epoch": 0.33,
+      "gumbel_tau": 0.4696,
+      "image_masks_loss": 0.0152,
+      "llm_learning_rate": 3.920792041286462e-06,
+      "loss": 0.7619,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.00039207920412864615,
+      "step": 3205
+    },
+    {
+      "epoch": 0.33,
+      "gumbel_tau": 0.4695,
+      "image_masks_loss": 0.0196,
+      "llm_learning_rate": 3.920109346273806e-06,
+      "loss": 0.8762,
+      "output_text_mask_loss": 0.0431,
+      "predictor_learning_rate": 0.00039201093462738057,
+      "step": 3206
+    },
+    {
+      "epoch": 0.33,
+      "gumbel_tau": 0.4694,
+      "image_masks_loss": 0.0305,
+      "llm_learning_rate": 3.9194264948755465e-06,
+      "loss": 0.6979,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.00039194264948755466,
+      "step": 3207
+    },
+    {
+      "epoch": 0.33,
+      "gumbel_tau": 0.4693,
+      "image_masks_loss": 0.0145,
+      "llm_learning_rate": 3.918743487166879e-06,
+      "loss": 0.6927,
+      "output_text_mask_loss": 0.0174,
+      "predictor_learning_rate": 0.0003918743487166879,
+      "step": 3208
+    },
+    {
+      "epoch": 0.33,
+      "gumbel_tau": 0.4691,
+      "image_masks_loss": 0.0309,
+      "llm_learning_rate": 3.9180603232230195e-06,
+      "loss": 0.7815,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.0003918060323223019,
+      "step": 3209
+    },
+    {
+      "epoch": 0.33,
+      "gumbel_tau": 0.469,
+      "image_masks_loss": 0.0141,
+      "llm_learning_rate": 3.917377003119198e-06,
+      "loss": 0.7927,
+      "output_text_mask_loss": 0.0185,
+      "predictor_learning_rate": 0.0003917377003119198,
+      "step": 3210
+    },
+    {
+      "epoch": 0.33,
+      "gumbel_tau": 0.4689,
+      "image_masks_loss": 0.0108,
+      "llm_learning_rate": 3.916693526930664e-06,
+      "loss": 0.732,
+      "output_text_mask_loss": 0.0319,
+      "predictor_learning_rate": 0.00039166935269306636,
+      "step": 3211
+    },
+    {
+      "epoch": 0.33,
+      "gumbel_tau": 0.4688,
+      "image_masks_loss": 0.007,
+      "llm_learning_rate": 3.9160098947326835e-06,
+      "loss": 0.786,
+      "output_text_mask_loss": 0.0035,
+      "predictor_learning_rate": 0.0003916009894732683,
+      "step": 3212
+    },
+    {
+      "epoch": 0.33,
+      "gumbel_tau": 0.4687,
+      "image_masks_loss": 0.0286,
+      "llm_learning_rate": 3.915326106600539e-06,
+      "loss": 0.7346,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.0003915326106600539,
+      "step": 3213
+    },
+    {
+      "epoch": 0.33,
+      "gumbel_tau": 0.4686,
+      "image_masks_loss": 0.0186,
+      "llm_learning_rate": 3.914642162609531e-06,
+      "loss": 0.7708,
+      "output_text_mask_loss": 0.0216,
+      "predictor_learning_rate": 0.00039146421626095305,
+      "step": 3214
+    },
+    {
+      "epoch": 0.33,
+      "gumbel_tau": 0.4685,
+      "image_masks_loss": 0.017,
+      "llm_learning_rate": 3.913958062834977e-06,
+      "loss": 0.6445,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.0003913958062834977,
+      "step": 3215
+    },
+    {
+      "epoch": 0.33,
+      "gumbel_tau": 0.4684,
+      "image_masks_loss": 0.0184,
+      "llm_learning_rate": 3.913273807352212e-06,
+      "loss": 0.8388,
+      "output_text_mask_loss": 0.0088,
+      "predictor_learning_rate": 0.00039132738073522125,
+      "step": 3216
+    },
+    {
+      "epoch": 0.33,
+      "gumbel_tau": 0.4683,
+      "image_masks_loss": 0.021,
+      "llm_learning_rate": 3.912589396236588e-06,
+      "loss": 0.7245,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.0003912589396236588,
+      "step": 3217
+    },
+    {
+      "epoch": 0.33,
+      "gumbel_tau": 0.4682,
+      "image_masks_loss": 0.0163,
+      "llm_learning_rate": 3.911904829563472e-06,
+      "loss": 0.759,
+      "output_text_mask_loss": 0.0023,
+      "predictor_learning_rate": 0.0003911904829563472,
+      "step": 3218
+    },
+    {
+      "epoch": 0.33,
+      "gumbel_tau": 0.468,
+      "image_masks_loss": 0.0137,
+      "llm_learning_rate": 3.911220107408252e-06,
+      "loss": 0.8144,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0003911220107408252,
+      "step": 3219
+    },
+    {
+      "epoch": 0.33,
+      "gumbel_tau": 0.4679,
+      "image_masks_loss": 0.0084,
+      "llm_learning_rate": 3.910535229846331e-06,
+      "loss": 0.8016,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.000391053522984633,
+      "step": 3220
+    },
+    {
+      "epoch": 0.33,
+      "gumbel_tau": 0.4678,
+      "image_masks_loss": 0.0082,
+      "llm_learning_rate": 3.909850196953127e-06,
+      "loss": 0.6685,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.00039098501969531265,
+      "step": 3221
+    },
+    {
+      "epoch": 0.33,
+      "gumbel_tau": 0.4677,
+      "image_masks_loss": 0.0286,
+      "llm_learning_rate": 3.90916500880408e-06,
+      "loss": 0.7106,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 0.000390916500880408,
+      "step": 3222
+    },
+    {
+      "epoch": 0.33,
+      "gumbel_tau": 0.4676,
+      "image_masks_loss": 0.0267,
+      "llm_learning_rate": 3.9084796654746436e-06,
+      "loss": 0.8388,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0003908479665474643,
+      "step": 3223
+    },
+    {
+      "epoch": 0.33,
+      "gumbel_tau": 0.4675,
+      "image_masks_loss": 0.0196,
+      "llm_learning_rate": 3.907794167040289e-06,
+      "loss": 0.742,
+      "output_text_mask_loss": 0.0187,
+      "predictor_learning_rate": 0.0003907794167040288,
+      "step": 3224
+    },
+    {
+      "epoch": 0.33,
+      "gumbel_tau": 0.4674,
+      "image_masks_loss": 0.0443,
+      "llm_learning_rate": 3.907108513576504e-06,
+      "loss": 0.7829,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0003907108513576504,
+      "step": 3225
+    },
+    {
+      "epoch": 0.33,
+      "gumbel_tau": 0.4673,
+      "image_masks_loss": 0.0073,
+      "llm_learning_rate": 3.906422705158797e-06,
+      "loss": 0.7766,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0003906422705158796,
+      "step": 3226
+    },
+    {
+      "epoch": 0.33,
+      "gumbel_tau": 0.4672,
+      "image_masks_loss": 0.0195,
+      "llm_learning_rate": 3.905736741862688e-06,
+      "loss": 0.8181,
+      "output_text_mask_loss": 0.0046,
+      "predictor_learning_rate": 0.0003905736741862688,
+      "step": 3227
+    },
+    {
+      "epoch": 0.33,
+      "gumbel_tau": 0.4671,
+      "image_masks_loss": 0.0129,
+      "llm_learning_rate": 3.905050623763717e-06,
+      "loss": 0.775,
+      "output_text_mask_loss": 0.0113,
+      "predictor_learning_rate": 0.00039050506237637173,
+      "step": 3228
+    },
+    {
+      "epoch": 0.33,
+      "gumbel_tau": 0.4669,
+      "image_masks_loss": 0.0079,
+      "llm_learning_rate": 3.904364350937443e-06,
+      "loss": 0.6821,
+      "output_text_mask_loss": 0.0462,
+      "predictor_learning_rate": 0.00039043643509374426,
+      "step": 3229
+    },
+    {
+      "epoch": 0.33,
+      "gumbel_tau": 0.4668,
+      "image_masks_loss": 0.0294,
+      "llm_learning_rate": 3.903677923459438e-06,
+      "loss": 0.8428,
+      "output_text_mask_loss": 0.0326,
+      "predictor_learning_rate": 0.00039036779234594376,
+      "step": 3230
+    },
+    {
+      "epoch": 0.33,
+      "gumbel_tau": 0.4667,
+      "image_masks_loss": 0.0219,
+      "llm_learning_rate": 3.902991341405295e-06,
+      "loss": 0.817,
+      "output_text_mask_loss": 0.0387,
+      "predictor_learning_rate": 0.0003902991341405294,
+      "step": 3231
+    },
+    {
+      "epoch": 0.33,
+      "gumbel_tau": 0.4666,
+      "image_masks_loss": 0.0171,
+      "llm_learning_rate": 3.9023046048506185e-06,
+      "loss": 0.6785,
+      "output_text_mask_loss": 0.0158,
+      "predictor_learning_rate": 0.00039023046048506185,
+      "step": 3232
+    },
+    {
+      "epoch": 0.33,
+      "gumbel_tau": 0.4665,
+      "image_masks_loss": 0.0142,
+      "llm_learning_rate": 3.9016177138710356e-06,
+      "loss": 0.7524,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.0003901617713871036,
+      "step": 3233
+    },
+    {
+      "epoch": 0.33,
+      "gumbel_tau": 0.4664,
+      "image_masks_loss": 0.0122,
+      "llm_learning_rate": 3.900930668542188e-06,
+      "loss": 0.7893,
+      "output_text_mask_loss": 0.0223,
+      "predictor_learning_rate": 0.0003900930668542189,
+      "step": 3234
+    },
+    {
+      "epoch": 0.33,
+      "gumbel_tau": 0.4663,
+      "image_masks_loss": 0.0131,
+      "llm_learning_rate": 3.9002434689397365e-06,
+      "loss": 0.6535,
+      "output_text_mask_loss": 0.0101,
+      "predictor_learning_rate": 0.0003900243468939736,
+      "step": 3235
+    },
+    {
+      "epoch": 0.33,
+      "gumbel_tau": 0.4662,
+      "image_masks_loss": 0.0448,
+      "llm_learning_rate": 3.899556115139354e-06,
+      "loss": 0.7383,
+      "output_text_mask_loss": 0.0056,
+      "predictor_learning_rate": 0.0003899556115139354,
+      "step": 3236
+    },
+    {
+      "epoch": 0.33,
+      "gumbel_tau": 0.4661,
+      "image_masks_loss": 0.0146,
+      "llm_learning_rate": 3.898868607216735e-06,
+      "loss": 0.7602,
+      "output_text_mask_loss": 0.0353,
+      "predictor_learning_rate": 0.00038988686072167345,
+      "step": 3237
+    },
+    {
+      "epoch": 0.33,
+      "gumbel_tau": 0.466,
+      "image_masks_loss": 0.0119,
+      "llm_learning_rate": 3.898180945247589e-06,
+      "loss": 0.7218,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0003898180945247588,
+      "step": 3238
+    },
+    {
+      "epoch": 0.33,
+      "gumbel_tau": 0.4658,
+      "image_masks_loss": 0.0433,
+      "llm_learning_rate": 3.897493129307643e-06,
+      "loss": 0.7847,
+      "output_text_mask_loss": 0.0504,
+      "predictor_learning_rate": 0.0003897493129307642,
+      "step": 3239
+    },
+    {
+      "epoch": 0.33,
+      "gumbel_tau": 0.4657,
+      "image_masks_loss": 0.0267,
+      "llm_learning_rate": 3.89680515947264e-06,
+      "loss": 0.7335,
+      "output_text_mask_loss": 0.0223,
+      "predictor_learning_rate": 0.00038968051594726397,
+      "step": 3240
+    },
+    {
+      "epoch": 0.33,
+      "gumbel_tau": 0.4656,
+      "image_masks_loss": 0.0383,
+      "llm_learning_rate": 3.896117035818342e-06,
+      "loss": 0.7757,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00038961170358183417,
+      "step": 3241
+    },
+    {
+      "epoch": 0.33,
+      "gumbel_tau": 0.4655,
+      "image_masks_loss": 0.0098,
+      "llm_learning_rate": 3.895428758420526e-06,
+      "loss": 0.7453,
+      "output_text_mask_loss": 0.0086,
+      "predictor_learning_rate": 0.00038954287584205263,
+      "step": 3242
+    },
+    {
+      "epoch": 0.33,
+      "gumbel_tau": 0.4654,
+      "image_masks_loss": 0.0185,
+      "llm_learning_rate": 3.8947403273549886e-06,
+      "loss": 0.8952,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00038947403273549883,
+      "step": 3243
+    },
+    {
+      "epoch": 0.33,
+      "gumbel_tau": 0.4653,
+      "image_masks_loss": 0.0175,
+      "llm_learning_rate": 3.894051742697539e-06,
+      "loss": 0.704,
+      "output_text_mask_loss": 0.0193,
+      "predictor_learning_rate": 0.00038940517426975385,
+      "step": 3244
+    },
+    {
+      "epoch": 0.33,
+      "gumbel_tau": 0.4652,
+      "image_masks_loss": 0.0507,
+      "llm_learning_rate": 3.893363004524006e-06,
+      "loss": 0.709,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00038933630045240054,
+      "step": 3245
+    },
+    {
+      "epoch": 0.33,
+      "gumbel_tau": 0.4651,
+      "image_masks_loss": 0.0307,
+      "llm_learning_rate": 3.8926741129102344e-06,
+      "loss": 0.774,
+      "output_text_mask_loss": 0.0019,
+      "predictor_learning_rate": 0.0003892674112910234,
+      "step": 3246
+    },
+    {
+      "epoch": 0.33,
+      "gumbel_tau": 0.465,
+      "image_masks_loss": 0.0248,
+      "llm_learning_rate": 3.891985067932089e-06,
+      "loss": 0.8542,
+      "output_text_mask_loss": 0.0124,
+      "predictor_learning_rate": 0.00038919850679320885,
+      "step": 3247
+    },
+    {
+      "epoch": 0.33,
+      "gumbel_tau": 0.4649,
+      "image_masks_loss": 0.0238,
+      "llm_learning_rate": 3.891295869665446e-06,
+      "loss": 0.8374,
+      "output_text_mask_loss": 0.0179,
+      "predictor_learning_rate": 0.00038912958696654454,
+      "step": 3248
+    },
+    {
+      "epoch": 0.33,
+      "gumbel_tau": 0.4647,
+      "image_masks_loss": 0.0298,
+      "llm_learning_rate": 3.890606518186204e-06,
+      "loss": 0.8076,
+      "output_text_mask_loss": 0.0622,
+      "predictor_learning_rate": 0.0003890606518186204,
+      "step": 3249
+    },
+    {
+      "epoch": 0.33,
+      "gumbel_tau": 0.4646,
+      "image_masks_loss": 0.0309,
+      "llm_learning_rate": 3.889917013570274e-06,
+      "loss": 0.8175,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00038899170135702735,
+      "step": 3250
+    },
+    {
+      "epoch": 0.33,
+      "gumbel_tau": 0.4645,
+      "image_masks_loss": 0.0284,
+      "llm_learning_rate": 3.889227355893587e-06,
+      "loss": 0.7376,
+      "output_text_mask_loss": 0.0023,
+      "predictor_learning_rate": 0.00038892273558935867,
+      "step": 3251
+    },
+    {
+      "epoch": 0.33,
+      "gumbel_tau": 0.4644,
+      "image_masks_loss": 0.0326,
+      "llm_learning_rate": 3.888537545232088e-06,
+      "loss": 0.7722,
+      "output_text_mask_loss": 0.0231,
+      "predictor_learning_rate": 0.0003888537545232088,
+      "step": 3252
+    },
+    {
+      "epoch": 0.33,
+      "gumbel_tau": 0.4643,
+      "image_masks_loss": 0.0237,
+      "llm_learning_rate": 3.887847581661744e-06,
+      "loss": 0.8317,
+      "output_text_mask_loss": 0.005,
+      "predictor_learning_rate": 0.0003887847581661743,
+      "step": 3253
+    },
+    {
+      "epoch": 0.33,
+      "gumbel_tau": 0.4642,
+      "image_masks_loss": 0.0475,
+      "llm_learning_rate": 3.887157465258531e-06,
+      "loss": 0.8334,
+      "output_text_mask_loss": 0.0481,
+      "predictor_learning_rate": 0.0003887157465258531,
+      "step": 3254
+    },
+    {
+      "epoch": 0.33,
+      "gumbel_tau": 0.4641,
+      "image_masks_loss": 0.0259,
+      "llm_learning_rate": 3.886467196098449e-06,
+      "loss": 0.7782,
+      "output_text_mask_loss": 0.0275,
+      "predictor_learning_rate": 0.0003886467196098449,
+      "step": 3255
+    },
+    {
+      "epoch": 0.33,
+      "gumbel_tau": 0.464,
+      "image_masks_loss": 0.0191,
+      "llm_learning_rate": 3.8857767742575115e-06,
+      "loss": 0.7285,
+      "output_text_mask_loss": 0.0056,
+      "predictor_learning_rate": 0.0003885776774257511,
+      "step": 3256
+    },
+    {
+      "epoch": 0.33,
+      "gumbel_tau": 0.4639,
+      "image_masks_loss": 0.0141,
+      "llm_learning_rate": 3.885086199811748e-06,
+      "loss": 0.7811,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.0003885086199811748,
+      "step": 3257
+    },
+    {
+      "epoch": 0.33,
+      "gumbel_tau": 0.4638,
+      "image_masks_loss": 0.0177,
+      "llm_learning_rate": 3.8843954728372075e-06,
+      "loss": 0.682,
+      "output_text_mask_loss": 0.0116,
+      "predictor_learning_rate": 0.0003884395472837207,
+      "step": 3258
+    },
+    {
+      "epoch": 0.33,
+      "gumbel_tau": 0.4636,
+      "image_masks_loss": 0.0614,
+      "llm_learning_rate": 3.883704593409954e-06,
+      "loss": 0.7403,
+      "output_text_mask_loss": 0.0052,
+      "predictor_learning_rate": 0.00038837045934099533,
+      "step": 3259
+    },
+    {
+      "epoch": 0.33,
+      "gumbel_tau": 0.4635,
+      "image_masks_loss": 0.0546,
+      "llm_learning_rate": 3.8830135616060685e-06,
+      "loss": 0.8104,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0003883013561606068,
+      "step": 3260
+    },
+    {
+      "epoch": 0.33,
+      "gumbel_tau": 0.4634,
+      "image_masks_loss": 0.0151,
+      "llm_learning_rate": 3.882322377501648e-06,
+      "loss": 0.7616,
+      "output_text_mask_loss": 0.0098,
+      "predictor_learning_rate": 0.0003882322377501648,
+      "step": 3261
+    },
+    {
+      "epoch": 0.33,
+      "gumbel_tau": 0.4633,
+      "image_masks_loss": 0.0374,
+      "llm_learning_rate": 3.8816310411728105e-06,
+      "loss": 0.8375,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 0.000388163104117281,
+      "step": 3262
+    },
+    {
+      "epoch": 0.33,
+      "gumbel_tau": 0.4632,
+      "image_masks_loss": 0.0257,
+      "llm_learning_rate": 3.880939552695684e-06,
+      "loss": 0.8071,
+      "output_text_mask_loss": 0.0071,
+      "predictor_learning_rate": 0.00038809395526956836,
+      "step": 3263
+    },
+    {
+      "epoch": 0.33,
+      "gumbel_tau": 0.4631,
+      "image_masks_loss": 0.0161,
+      "llm_learning_rate": 3.880247912146418e-06,
+      "loss": 0.6716,
+      "output_text_mask_loss": 0.0278,
+      "predictor_learning_rate": 0.00038802479121464176,
+      "step": 3264
+    },
+    {
+      "epoch": 0.33,
+      "gumbel_tau": 0.463,
+      "image_masks_loss": 0.0197,
+      "llm_learning_rate": 3.879556119601178e-06,
+      "loss": 0.8187,
+      "output_text_mask_loss": 0.0156,
+      "predictor_learning_rate": 0.0003879556119601178,
+      "step": 3265
+    },
+    {
+      "epoch": 0.33,
+      "gumbel_tau": 0.4629,
+      "image_masks_loss": 0.0149,
+      "llm_learning_rate": 3.878864175136144e-06,
+      "loss": 0.7819,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00038788641751361444,
+      "step": 3266
+    },
+    {
+      "epoch": 0.33,
+      "gumbel_tau": 0.4628,
+      "image_masks_loss": 0.037,
+      "llm_learning_rate": 3.878172078827517e-06,
+      "loss": 0.772,
+      "output_text_mask_loss": 0.0446,
+      "predictor_learning_rate": 0.00038781720788275164,
+      "step": 3267
+    },
+    {
+      "epoch": 0.33,
+      "gumbel_tau": 0.4627,
+      "image_masks_loss": 0.0492,
+      "llm_learning_rate": 3.87747983075151e-06,
+      "loss": 0.7787,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00038774798307515104,
+      "step": 3268
+    },
+    {
+      "epoch": 0.33,
+      "gumbel_tau": 0.4626,
+      "image_masks_loss": 0.0082,
+      "llm_learning_rate": 3.876787430984357e-06,
+      "loss": 0.7834,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.0003876787430984357,
+      "step": 3269
+    },
+    {
+      "epoch": 0.34,
+      "gumbel_tau": 0.4624,
+      "image_masks_loss": 0.0216,
+      "llm_learning_rate": 3.876094879602305e-06,
+      "loss": 0.798,
+      "output_text_mask_loss": 0.0422,
+      "predictor_learning_rate": 0.00038760948796023046,
+      "step": 3270
+    },
+    {
+      "epoch": 0.34,
+      "gumbel_tau": 0.4623,
+      "image_masks_loss": 0.0332,
+      "llm_learning_rate": 3.875402176681619e-06,
+      "loss": 0.7915,
+      "output_text_mask_loss": 0.0278,
+      "predictor_learning_rate": 0.00038754021766816184,
+      "step": 3271
+    },
+    {
+      "epoch": 0.34,
+      "gumbel_tau": 0.4622,
+      "image_masks_loss": 0.0273,
+      "llm_learning_rate": 3.874709322298582e-06,
+      "loss": 0.7643,
+      "output_text_mask_loss": 0.017,
+      "predictor_learning_rate": 0.0003874709322298582,
+      "step": 3272
+    },
+    {
+      "epoch": 0.34,
+      "gumbel_tau": 0.4621,
+      "image_masks_loss": 0.0192,
+      "llm_learning_rate": 3.874016316529493e-06,
+      "loss": 0.792,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.0003874016316529493,
+      "step": 3273
+    },
+    {
+      "epoch": 0.34,
+      "gumbel_tau": 0.462,
+      "image_masks_loss": 0.0203,
+      "llm_learning_rate": 3.873323159450666e-06,
+      "loss": 0.7811,
+      "output_text_mask_loss": 0.001,
+      "predictor_learning_rate": 0.00038733231594506655,
+      "step": 3274
+    },
+    {
+      "epoch": 0.34,
+      "gumbel_tau": 0.4619,
+      "image_masks_loss": 0.0326,
+      "llm_learning_rate": 3.872629851138434e-06,
+      "loss": 0.755,
+      "output_text_mask_loss": 0.0172,
+      "predictor_learning_rate": 0.0003872629851138434,
+      "step": 3275
+    },
+    {
+      "epoch": 0.34,
+      "gumbel_tau": 0.4618,
+      "image_masks_loss": 0.0142,
+      "llm_learning_rate": 3.8719363916691456e-06,
+      "loss": 0.769,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0003871936391669145,
+      "step": 3276
+    },
+    {
+      "epoch": 0.34,
+      "gumbel_tau": 0.4617,
+      "image_masks_loss": 0.0092,
+      "llm_learning_rate": 3.871242781119166e-06,
+      "loss": 0.7361,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0003871242781119166,
+      "step": 3277
+    },
+    {
+      "epoch": 0.34,
+      "gumbel_tau": 0.4616,
+      "image_masks_loss": 0.029,
+      "llm_learning_rate": 3.870549019564877e-06,
+      "loss": 0.8107,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00038705490195648764,
+      "step": 3278
+    },
+    {
+      "epoch": 0.34,
+      "gumbel_tau": 0.4615,
+      "image_masks_loss": 0.0114,
+      "llm_learning_rate": 3.869855107082676e-06,
+      "loss": 0.7958,
+      "output_text_mask_loss": 0.0296,
+      "predictor_learning_rate": 0.0003869855107082676,
+      "step": 3279
+    },
+    {
+      "epoch": 0.34,
+      "gumbel_tau": 0.4614,
+      "image_masks_loss": 0.0095,
+      "llm_learning_rate": 3.869161043748981e-06,
+      "loss": 0.8462,
+      "output_text_mask_loss": 0.0261,
+      "predictor_learning_rate": 0.00038691610437489805,
+      "step": 3280
+    },
+    {
+      "epoch": 0.34,
+      "gumbel_tau": 0.4612,
+      "image_masks_loss": 0.0355,
+      "llm_learning_rate": 3.8684668296402215e-06,
+      "loss": 0.7575,
+      "output_text_mask_loss": 0.002,
+      "predictor_learning_rate": 0.00038684668296402213,
+      "step": 3281
+    },
+    {
+      "epoch": 0.34,
+      "gumbel_tau": 0.4611,
+      "image_masks_loss": 0.0111,
+      "llm_learning_rate": 3.8677724648328465e-06,
+      "loss": 0.7639,
+      "output_text_mask_loss": 0.0469,
+      "predictor_learning_rate": 0.0003867772464832846,
+      "step": 3282
+    },
+    {
+      "epoch": 0.34,
+      "gumbel_tau": 0.461,
+      "image_masks_loss": 0.0237,
+      "llm_learning_rate": 3.867077949403322e-06,
+      "loss": 0.796,
+      "output_text_mask_loss": 0.0452,
+      "predictor_learning_rate": 0.00038670779494033217,
+      "step": 3283
+    },
+    {
+      "epoch": 0.34,
+      "gumbel_tau": 0.4609,
+      "image_masks_loss": 0.0515,
+      "llm_learning_rate": 3.8663832834281264e-06,
+      "loss": 0.793,
+      "output_text_mask_loss": 0.0111,
+      "predictor_learning_rate": 0.00038663832834281265,
+      "step": 3284
+    },
+    {
+      "epoch": 0.34,
+      "gumbel_tau": 0.4608,
+      "image_masks_loss": 0.0359,
+      "llm_learning_rate": 3.865688466983762e-06,
+      "loss": 0.7721,
+      "output_text_mask_loss": 0.0683,
+      "predictor_learning_rate": 0.00038656884669837624,
+      "step": 3285
+    },
+    {
+      "epoch": 0.34,
+      "gumbel_tau": 0.4607,
+      "image_masks_loss": 0.0166,
+      "llm_learning_rate": 3.864993500146741e-06,
+      "loss": 0.8239,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 0.0003864993500146741,
+      "step": 3286
+    },
+    {
+      "epoch": 0.34,
+      "gumbel_tau": 0.4606,
+      "image_masks_loss": 0.014,
+      "llm_learning_rate": 3.864298382993596e-06,
+      "loss": 0.7678,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.0003864298382993595,
+      "step": 3287
+    },
+    {
+      "epoch": 0.34,
+      "gumbel_tau": 0.4605,
+      "image_masks_loss": 0.0294,
+      "llm_learning_rate": 3.8636031156008734e-06,
+      "loss": 0.6941,
+      "output_text_mask_loss": 0.0233,
+      "predictor_learning_rate": 0.00038636031156008733,
+      "step": 3288
+    },
+    {
+      "epoch": 0.34,
+      "gumbel_tau": 0.4604,
+      "image_masks_loss": 0.0043,
+      "llm_learning_rate": 3.862907698045139e-06,
+      "loss": 0.8346,
+      "output_text_mask_loss": 0.0015,
+      "predictor_learning_rate": 0.0003862907698045138,
+      "step": 3289
+    },
+    {
+      "epoch": 0.34,
+      "gumbel_tau": 0.4603,
+      "image_masks_loss": 0.0313,
+      "llm_learning_rate": 3.862212130402973e-06,
+      "loss": 0.7873,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.0003862212130402973,
+      "step": 3290
+    },
+    {
+      "epoch": 0.34,
+      "gumbel_tau": 0.4602,
+      "image_masks_loss": 0.0131,
+      "llm_learning_rate": 3.861516412750973e-06,
+      "loss": 0.8727,
+      "output_text_mask_loss": 0.0164,
+      "predictor_learning_rate": 0.00038615164127509726,
+      "step": 3291
+    },
+    {
+      "epoch": 0.34,
+      "gumbel_tau": 0.4601,
+      "image_masks_loss": 0.0185,
+      "llm_learning_rate": 3.860820545165752e-06,
+      "loss": 0.706,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.0003860820545165752,
+      "step": 3292
+    },
+    {
+      "epoch": 0.34,
+      "gumbel_tau": 0.4599,
+      "image_masks_loss": 0.0353,
+      "llm_learning_rate": 3.860124527723943e-06,
+      "loss": 0.7918,
+      "output_text_mask_loss": 0.0082,
+      "predictor_learning_rate": 0.00038601245277239426,
+      "step": 3293
+    },
+    {
+      "epoch": 0.34,
+      "gumbel_tau": 0.4598,
+      "image_masks_loss": 0.0242,
+      "llm_learning_rate": 3.859428360502191e-06,
+      "loss": 0.7552,
+      "output_text_mask_loss": 0.0153,
+      "predictor_learning_rate": 0.0003859428360502191,
+      "step": 3294
+    },
+    {
+      "epoch": 0.34,
+      "gumbel_tau": 0.4597,
+      "image_masks_loss": 0.0173,
+      "llm_learning_rate": 3.858732043577159e-06,
+      "loss": 0.7958,
+      "output_text_mask_loss": 0.0372,
+      "predictor_learning_rate": 0.00038587320435771587,
+      "step": 3295
+    },
+    {
+      "epoch": 0.34,
+      "gumbel_tau": 0.4596,
+      "image_masks_loss": 0.013,
+      "llm_learning_rate": 3.858035577025529e-06,
+      "loss": 0.7937,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 0.00038580355770255286,
+      "step": 3296
+    },
+    {
+      "epoch": 0.34,
+      "gumbel_tau": 0.4595,
+      "image_masks_loss": 0.0098,
+      "llm_learning_rate": 3.857338960923996e-06,
+      "loss": 0.731,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.0003857338960923996,
+      "step": 3297
+    },
+    {
+      "epoch": 0.34,
+      "gumbel_tau": 0.4594,
+      "image_masks_loss": 0.0192,
+      "llm_learning_rate": 3.8566421953492735e-06,
+      "loss": 0.8017,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0003856642195349273,
+      "step": 3298
+    },
+    {
+      "epoch": 0.34,
+      "gumbel_tau": 0.4593,
+      "image_masks_loss": 0.0652,
+      "llm_learning_rate": 3.855945280378091e-06,
+      "loss": 0.7481,
+      "output_text_mask_loss": 0.0023,
+      "predictor_learning_rate": 0.000385594528037809,
+      "step": 3299
+    },
+    {
+      "epoch": 0.34,
+      "gumbel_tau": 0.4592,
+      "image_masks_loss": 0.0089,
+      "llm_learning_rate": 3.855248216087193e-06,
+      "loss": 0.7917,
+      "output_text_mask_loss": 0.0183,
+      "predictor_learning_rate": 0.00038552482160871925,
+      "step": 3300
+    },
+    {
+      "epoch": 0.34,
+      "gumbel_tau": 0.4591,
+      "image_masks_loss": 0.0192,
+      "llm_learning_rate": 3.854551002553343e-06,
+      "loss": 0.7835,
+      "output_text_mask_loss": 0.0086,
+      "predictor_learning_rate": 0.00038545510025533424,
+      "step": 3301
+    },
+    {
+      "epoch": 0.34,
+      "gumbel_tau": 0.459,
+      "image_masks_loss": 0.0488,
+      "llm_learning_rate": 3.853853639853319e-06,
+      "loss": 0.7654,
+      "output_text_mask_loss": 0.0366,
+      "predictor_learning_rate": 0.0003853853639853319,
+      "step": 3302
+    },
+    {
+      "epoch": 0.34,
+      "gumbel_tau": 0.4589,
+      "image_masks_loss": 0.0181,
+      "llm_learning_rate": 3.853156128063917e-06,
+      "loss": 0.7652,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.00038531561280639174,
+      "step": 3303
+    },
+    {
+      "epoch": 0.34,
+      "gumbel_tau": 0.4588,
+      "image_masks_loss": 0.0185,
+      "llm_learning_rate": 3.8524584672619485e-06,
+      "loss": 0.6812,
+      "output_text_mask_loss": 0.0086,
+      "predictor_learning_rate": 0.00038524584672619484,
+      "step": 3304
+    },
+    {
+      "epoch": 0.34,
+      "gumbel_tau": 0.4586,
+      "image_masks_loss": 0.0456,
+      "llm_learning_rate": 3.8517606575242406e-06,
+      "loss": 0.7948,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 0.00038517606575242405,
+      "step": 3305
+    },
+    {
+      "epoch": 0.34,
+      "gumbel_tau": 0.4585,
+      "image_masks_loss": 0.0127,
+      "llm_learning_rate": 3.8510626989276386e-06,
+      "loss": 0.7855,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00038510626989276387,
+      "step": 3306
+    },
+    {
+      "epoch": 0.34,
+      "gumbel_tau": 0.4584,
+      "image_masks_loss": 0.0031,
+      "llm_learning_rate": 3.850364591549003e-06,
+      "loss": 0.7208,
+      "output_text_mask_loss": 0.0217,
+      "predictor_learning_rate": 0.00038503645915490025,
+      "step": 3307
+    },
+    {
+      "epoch": 0.34,
+      "gumbel_tau": 0.4583,
+      "image_masks_loss": 0.025,
+      "llm_learning_rate": 3.84966633546521e-06,
+      "loss": 0.8093,
+      "output_text_mask_loss": 0.0139,
+      "predictor_learning_rate": 0.000384966633546521,
+      "step": 3308
+    },
+    {
+      "epoch": 0.34,
+      "gumbel_tau": 0.4582,
+      "image_masks_loss": 0.0209,
+      "llm_learning_rate": 3.848967930753155e-06,
+      "loss": 0.8231,
+      "output_text_mask_loss": 0.0238,
+      "predictor_learning_rate": 0.00038489679307531546,
+      "step": 3309
+    },
+    {
+      "epoch": 0.34,
+      "gumbel_tau": 0.4581,
+      "image_masks_loss": 0.0134,
+      "llm_learning_rate": 3.848269377489746e-06,
+      "loss": 0.6976,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.0003848269377489746,
+      "step": 3310
+    },
+    {
+      "epoch": 0.34,
+      "gumbel_tau": 0.458,
+      "image_masks_loss": 0.0103,
+      "llm_learning_rate": 3.84757067575191e-06,
+      "loss": 0.7009,
+      "output_text_mask_loss": 0.036,
+      "predictor_learning_rate": 0.000384757067575191,
+      "step": 3311
+    },
+    {
+      "epoch": 0.34,
+      "gumbel_tau": 0.4579,
+      "image_masks_loss": 0.0101,
+      "llm_learning_rate": 3.846871825616589e-06,
+      "loss": 0.7015,
+      "output_text_mask_loss": 0.0094,
+      "predictor_learning_rate": 0.0003846871825616589,
+      "step": 3312
+    },
+    {
+      "epoch": 0.34,
+      "gumbel_tau": 0.4578,
+      "image_masks_loss": 0.0093,
+      "llm_learning_rate": 3.8461728271607435e-06,
+      "loss": 0.7227,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.0003846172827160743,
+      "step": 3313
+    },
+    {
+      "epoch": 0.34,
+      "gumbel_tau": 0.4577,
+      "image_masks_loss": 0.0121,
+      "llm_learning_rate": 3.845473680461348e-06,
+      "loss": 0.7834,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00038454736804613475,
+      "step": 3314
+    },
+    {
+      "epoch": 0.34,
+      "gumbel_tau": 0.4576,
+      "image_masks_loss": 0.012,
+      "llm_learning_rate": 3.844774385595394e-06,
+      "loss": 0.7766,
+      "output_text_mask_loss": 0.0023,
+      "predictor_learning_rate": 0.0003844774385595393,
+      "step": 3315
+    },
+    {
+      "epoch": 0.34,
+      "gumbel_tau": 0.4575,
+      "image_masks_loss": 0.0315,
+      "llm_learning_rate": 3.844074942639887e-06,
+      "loss": 0.776,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.0003844074942639887,
+      "step": 3316
+    },
+    {
+      "epoch": 0.34,
+      "gumbel_tau": 0.4573,
+      "image_masks_loss": 0.0122,
+      "llm_learning_rate": 3.843375351671855e-06,
+      "loss": 0.7158,
+      "output_text_mask_loss": 0.0374,
+      "predictor_learning_rate": 0.0003843375351671855,
+      "step": 3317
+    },
+    {
+      "epoch": 0.34,
+      "gumbel_tau": 0.4572,
+      "image_masks_loss": 0.0189,
+      "llm_learning_rate": 3.842675612768338e-06,
+      "loss": 0.8294,
+      "output_text_mask_loss": 0.042,
+      "predictor_learning_rate": 0.0003842675612768338,
+      "step": 3318
+    },
+    {
+      "epoch": 0.34,
+      "gumbel_tau": 0.4571,
+      "image_masks_loss": 0.0067,
+      "llm_learning_rate": 3.841975726006391e-06,
+      "loss": 0.724,
+      "output_text_mask_loss": 0.0231,
+      "predictor_learning_rate": 0.00038419757260063907,
+      "step": 3319
+    },
+    {
+      "epoch": 0.34,
+      "gumbel_tau": 0.457,
+      "image_masks_loss": 0.0186,
+      "llm_learning_rate": 3.841275691463088e-06,
+      "loss": 0.7372,
+      "output_text_mask_loss": 0.071,
+      "predictor_learning_rate": 0.00038412756914630875,
+      "step": 3320
+    },
+    {
+      "epoch": 0.34,
+      "gumbel_tau": 0.4569,
+      "image_masks_loss": 0.0104,
+      "llm_learning_rate": 3.840575509215518e-06,
+      "loss": 0.6705,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0003840575509215518,
+      "step": 3321
+    },
+    {
+      "epoch": 0.34,
+      "gumbel_tau": 0.4568,
+      "image_masks_loss": 0.0116,
+      "llm_learning_rate": 3.839875179340788e-06,
+      "loss": 0.8183,
+      "output_text_mask_loss": 0.0081,
+      "predictor_learning_rate": 0.00038398751793407873,
+      "step": 3322
+    },
+    {
+      "epoch": 0.34,
+      "gumbel_tau": 0.4567,
+      "image_masks_loss": 0.0189,
+      "llm_learning_rate": 3.839174701916017e-06,
+      "loss": 0.7273,
+      "output_text_mask_loss": 0.0694,
+      "predictor_learning_rate": 0.0003839174701916017,
+      "step": 3323
+    },
+    {
+      "epoch": 0.34,
+      "gumbel_tau": 0.4566,
+      "image_masks_loss": 0.0165,
+      "llm_learning_rate": 3.838474077018346e-06,
+      "loss": 0.7751,
+      "output_text_mask_loss": 0.0126,
+      "predictor_learning_rate": 0.00038384740770183457,
+      "step": 3324
+    },
+    {
+      "epoch": 0.34,
+      "gumbel_tau": 0.4565,
+      "image_masks_loss": 0.0288,
+      "llm_learning_rate": 3.8377733047249285e-06,
+      "loss": 0.6772,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0003837773304724928,
+      "step": 3325
+    },
+    {
+      "epoch": 0.34,
+      "gumbel_tau": 0.4564,
+      "image_masks_loss": 0.0047,
+      "llm_learning_rate": 3.837072385112933e-06,
+      "loss": 0.7903,
+      "output_text_mask_loss": 0.036,
+      "predictor_learning_rate": 0.00038370723851129335,
+      "step": 3326
+    },
+    {
+      "epoch": 0.34,
+      "gumbel_tau": 0.4563,
+      "image_masks_loss": 0.0177,
+      "llm_learning_rate": 3.836371318259551e-06,
+      "loss": 0.6755,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0003836371318259551,
+      "step": 3327
+    },
+    {
+      "epoch": 0.34,
+      "gumbel_tau": 0.4562,
+      "image_masks_loss": 0.0115,
+      "llm_learning_rate": 3.835670104241982e-06,
+      "loss": 0.7907,
+      "output_text_mask_loss": 0.0067,
+      "predictor_learning_rate": 0.0003835670104241981,
+      "step": 3328
+    },
+    {
+      "epoch": 0.34,
+      "gumbel_tau": 0.4561,
+      "image_masks_loss": 0.0134,
+      "llm_learning_rate": 3.834968743137445e-06,
+      "loss": 0.6656,
+      "output_text_mask_loss": 0.0301,
+      "predictor_learning_rate": 0.0003834968743137445,
+      "step": 3329
+    },
+    {
+      "epoch": 0.34,
+      "gumbel_tau": 0.4559,
+      "image_masks_loss": 0.0067,
+      "llm_learning_rate": 3.834267235023177e-06,
+      "loss": 0.5927,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0003834267235023177,
+      "step": 3330
+    },
+    {
+      "epoch": 0.34,
+      "gumbel_tau": 0.4558,
+      "image_masks_loss": 0.0057,
+      "llm_learning_rate": 3.83356557997643e-06,
+      "loss": 0.6995,
+      "output_text_mask_loss": 0.0004,
+      "predictor_learning_rate": 0.0003833565579976429,
+      "step": 3331
+    },
+    {
+      "epoch": 0.34,
+      "gumbel_tau": 0.4557,
+      "image_masks_loss": 0.0135,
+      "llm_learning_rate": 3.832863778074469e-06,
+      "loss": 0.797,
+      "output_text_mask_loss": 0.1442,
+      "predictor_learning_rate": 0.0003832863778074469,
+      "step": 3332
+    },
+    {
+      "epoch": 0.34,
+      "gumbel_tau": 0.4556,
+      "image_masks_loss": 0.0107,
+      "llm_learning_rate": 3.832161829394582e-06,
+      "loss": 0.7807,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 0.00038321618293945815,
+      "step": 3333
+    },
+    {
+      "epoch": 0.34,
+      "gumbel_tau": 0.4555,
+      "image_masks_loss": 0.0343,
+      "llm_learning_rate": 3.831459734014066e-06,
+      "loss": 0.7533,
+      "output_text_mask_loss": 0.0244,
+      "predictor_learning_rate": 0.00038314597340140656,
+      "step": 3334
+    },
+    {
+      "epoch": 0.34,
+      "gumbel_tau": 0.4554,
+      "image_masks_loss": 0.0234,
+      "llm_learning_rate": 3.8307574920102374e-06,
+      "loss": 0.781,
+      "output_text_mask_loss": 0.0399,
+      "predictor_learning_rate": 0.00038307574920102373,
+      "step": 3335
+    },
+    {
+      "epoch": 0.34,
+      "gumbel_tau": 0.4553,
+      "image_masks_loss": 0.023,
+      "llm_learning_rate": 3.830055103460431e-06,
+      "loss": 0.6965,
+      "output_text_mask_loss": 0.0174,
+      "predictor_learning_rate": 0.00038300551034604316,
+      "step": 3336
+    },
+    {
+      "epoch": 0.34,
+      "gumbel_tau": 0.4552,
+      "image_masks_loss": 0.0144,
+      "llm_learning_rate": 3.8293525684419945e-06,
+      "loss": 0.8064,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00038293525684419944,
+      "step": 3337
+    },
+    {
+      "epoch": 0.34,
+      "gumbel_tau": 0.4551,
+      "image_masks_loss": 0.0134,
+      "llm_learning_rate": 3.828649887032291e-06,
+      "loss": 0.7773,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.00038286498870322906,
+      "step": 3338
+    },
+    {
+      "epoch": 0.34,
+      "gumbel_tau": 0.455,
+      "image_masks_loss": 0.0294,
+      "llm_learning_rate": 3.827947059308703e-06,
+      "loss": 0.7168,
+      "output_text_mask_loss": 0.0099,
+      "predictor_learning_rate": 0.0003827947059308703,
+      "step": 3339
+    },
+    {
+      "epoch": 0.34,
+      "gumbel_tau": 0.4549,
+      "image_masks_loss": 0.0261,
+      "llm_learning_rate": 3.827244085348627e-06,
+      "loss": 0.6419,
+      "output_text_mask_loss": 0.0565,
+      "predictor_learning_rate": 0.00038272440853486265,
+      "step": 3340
+    },
+    {
+      "epoch": 0.34,
+      "gumbel_tau": 0.4548,
+      "image_masks_loss": 0.0032,
+      "llm_learning_rate": 3.8265409652294764e-06,
+      "loss": 0.7544,
+      "output_text_mask_loss": 0.001,
+      "predictor_learning_rate": 0.0003826540965229476,
+      "step": 3341
+    },
+    {
+      "epoch": 0.34,
+      "gumbel_tau": 0.4547,
+      "image_masks_loss": 0.0162,
+      "llm_learning_rate": 3.825837699028678e-06,
+      "loss": 0.7278,
+      "output_text_mask_loss": 0.0248,
+      "predictor_learning_rate": 0.0003825837699028678,
+      "step": 3342
+    },
+    {
+      "epoch": 0.34,
+      "gumbel_tau": 0.4546,
+      "image_masks_loss": 0.0133,
+      "llm_learning_rate": 3.825134286823681e-06,
+      "loss": 0.7308,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 0.0003825134286823681,
+      "step": 3343
+    },
+    {
+      "epoch": 0.34,
+      "gumbel_tau": 0.4544,
+      "image_masks_loss": 0.0282,
+      "llm_learning_rate": 3.8244307286919444e-06,
+      "loss": 0.8074,
+      "output_text_mask_loss": 0.0488,
+      "predictor_learning_rate": 0.00038244307286919445,
+      "step": 3344
+    },
+    {
+      "epoch": 0.34,
+      "gumbel_tau": 0.4543,
+      "image_masks_loss": 0.0192,
+      "llm_learning_rate": 3.8237270247109445e-06,
+      "loss": 0.736,
+      "output_text_mask_loss": 0.0081,
+      "predictor_learning_rate": 0.00038237270247109444,
+      "step": 3345
+    },
+    {
+      "epoch": 0.34,
+      "gumbel_tau": 0.4542,
+      "image_masks_loss": 0.0292,
+      "llm_learning_rate": 3.823023174958178e-06,
+      "loss": 0.7676,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.0003823023174958178,
+      "step": 3346
+    },
+    {
+      "epoch": 0.34,
+      "gumbel_tau": 0.4541,
+      "image_masks_loss": 0.007,
+      "llm_learning_rate": 3.822319179511152e-06,
+      "loss": 0.7289,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0003822319179511152,
+      "step": 3347
+    },
+    {
+      "epoch": 0.34,
+      "gumbel_tau": 0.454,
+      "image_masks_loss": 0.0195,
+      "llm_learning_rate": 3.821615038447392e-06,
+      "loss": 0.7838,
+      "output_text_mask_loss": 0.0345,
+      "predictor_learning_rate": 0.00038216150384473926,
+      "step": 3348
+    },
+    {
+      "epoch": 0.34,
+      "gumbel_tau": 0.4539,
+      "image_masks_loss": 0.0133,
+      "llm_learning_rate": 3.820910751844441e-06,
+      "loss": 0.7862,
+      "output_text_mask_loss": 0.1015,
+      "predictor_learning_rate": 0.0003820910751844441,
+      "step": 3349
+    },
+    {
+      "epoch": 0.34,
+      "gumbel_tau": 0.4538,
+      "image_masks_loss": 0.0355,
+      "llm_learning_rate": 3.820206319779856e-06,
+      "loss": 0.6676,
+      "output_text_mask_loss": 0.0422,
+      "predictor_learning_rate": 0.0003820206319779855,
+      "step": 3350
+    },
+    {
+      "epoch": 0.34,
+      "gumbel_tau": 0.4537,
+      "image_masks_loss": 0.0216,
+      "llm_learning_rate": 3.8195017423312096e-06,
+      "loss": 0.6863,
+      "output_text_mask_loss": NaN,
+      "predictor_learning_rate": 0.0003819501742331209,
+      "step": 3351
+    },
+    {
+      "epoch": 0.34,
+      "gumbel_tau": 0.4536,
+      "image_masks_loss": 0.0437,
+      "llm_learning_rate": 3.818797019576092e-06,
+      "loss": 0.8495,
+      "output_text_mask_loss": 0.0004,
+      "predictor_learning_rate": 0.0003818797019576092,
+      "step": 3352
+    },
+    {
+      "epoch": 0.34,
+      "gumbel_tau": 0.4535,
+      "image_masks_loss": 0.028,
+      "llm_learning_rate": 3.81809215159211e-06,
+      "loss": 0.7471,
+      "output_text_mask_loss": 0.0235,
+      "predictor_learning_rate": 0.00038180921515921097,
+      "step": 3353
+    },
+    {
+      "epoch": 0.34,
+      "gumbel_tau": 0.4534,
+      "image_masks_loss": 0.029,
+      "llm_learning_rate": 3.817387138456884e-06,
+      "loss": 0.8796,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.00038173871384568835,
+      "step": 3354
+    },
+    {
+      "epoch": 0.34,
+      "gumbel_tau": 0.4533,
+      "image_masks_loss": 0.0319,
+      "llm_learning_rate": 3.8166819802480516e-06,
+      "loss": 0.7439,
+      "output_text_mask_loss": 0.0231,
+      "predictor_learning_rate": 0.00038166819802480515,
+      "step": 3355
+    },
+    {
+      "epoch": 0.34,
+      "gumbel_tau": 0.4532,
+      "image_masks_loss": 0.0282,
+      "llm_learning_rate": 3.815976677043267e-06,
+      "loss": 0.8158,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 0.0003815976677043267,
+      "step": 3356
+    },
+    {
+      "epoch": 0.34,
+      "gumbel_tau": 0.4531,
+      "image_masks_loss": 0.0136,
+      "llm_learning_rate": 3.8152712289202e-06,
+      "loss": 0.7425,
+      "output_text_mask_loss": 0.0172,
+      "predictor_learning_rate": 0.00038152712289201996,
+      "step": 3357
+    },
+    {
+      "epoch": 0.34,
+      "gumbel_tau": 0.4529,
+      "image_masks_loss": 0.0036,
+      "llm_learning_rate": 3.8145656359565342e-06,
+      "loss": 0.7561,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0003814565635956534,
+      "step": 3358
+    },
+    {
+      "epoch": 0.34,
+      "gumbel_tau": 0.4528,
+      "image_masks_loss": 0.0113,
+      "llm_learning_rate": 3.8138598982299733e-06,
+      "loss": 0.7989,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.0003813859898229973,
+      "step": 3359
+    },
+    {
+      "epoch": 0.34,
+      "gumbel_tau": 0.4527,
+      "image_masks_loss": 0.0128,
+      "llm_learning_rate": 3.8131540158182346e-06,
+      "loss": 0.631,
+      "output_text_mask_loss": 0.0004,
+      "predictor_learning_rate": 0.0003813154015818234,
+      "step": 3360
+    },
+    {
+      "epoch": 0.34,
+      "gumbel_tau": 0.4526,
+      "image_masks_loss": 0.0122,
+      "llm_learning_rate": 3.8124479887990495e-06,
+      "loss": 0.7391,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0003812447988799049,
+      "step": 3361
+    },
+    {
+      "epoch": 0.34,
+      "gumbel_tau": 0.4525,
+      "image_masks_loss": 0.0146,
+      "llm_learning_rate": 3.8117418172501697e-06,
+      "loss": 0.7413,
+      "output_text_mask_loss": 0.0039,
+      "predictor_learning_rate": 0.00038117418172501694,
+      "step": 3362
+    },
+    {
+      "epoch": 0.34,
+      "gumbel_tau": 0.4524,
+      "image_masks_loss": 0.0237,
+      "llm_learning_rate": 3.8110355012493584e-06,
+      "loss": 0.8635,
+      "output_text_mask_loss": 0.0282,
+      "predictor_learning_rate": 0.0003811035501249358,
+      "step": 3363
+    },
+    {
+      "epoch": 0.34,
+      "gumbel_tau": 0.4523,
+      "image_masks_loss": 0.0162,
+      "llm_learning_rate": 3.8103290408743977e-06,
+      "loss": 0.7988,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 0.00038103290408743974,
+      "step": 3364
+    },
+    {
+      "epoch": 0.34,
+      "gumbel_tau": 0.4522,
+      "image_masks_loss": 0.015,
+      "llm_learning_rate": 3.8096224362030846e-06,
+      "loss": 0.7335,
+      "output_text_mask_loss": 0.0133,
+      "predictor_learning_rate": 0.00038096224362030843,
+      "step": 3365
+    },
+    {
+      "epoch": 0.34,
+      "gumbel_tau": 0.4521,
+      "image_masks_loss": 0.0144,
+      "llm_learning_rate": 3.8089156873132314e-06,
+      "loss": 0.7394,
+      "output_text_mask_loss": 0.0181,
+      "predictor_learning_rate": 0.0003808915687313231,
+      "step": 3366
+    },
+    {
+      "epoch": 0.34,
+      "gumbel_tau": 0.452,
+      "image_masks_loss": 0.0094,
+      "llm_learning_rate": 3.8082087942826673e-06,
+      "loss": 0.7147,
+      "output_text_mask_loss": 0.0671,
+      "predictor_learning_rate": 0.00038082087942826673,
+      "step": 3367
+    },
+    {
+      "epoch": 0.35,
+      "gumbel_tau": 0.4519,
+      "image_masks_loss": 0.0185,
+      "llm_learning_rate": 3.8075017571892368e-06,
+      "loss": 0.8733,
+      "output_text_mask_loss": 0.0298,
+      "predictor_learning_rate": 0.0003807501757189237,
+      "step": 3368
+    },
+    {
+      "epoch": 0.35,
+      "gumbel_tau": 0.4518,
+      "image_masks_loss": 0.053,
+      "llm_learning_rate": 3.806794576110801e-06,
+      "loss": 0.7806,
+      "output_text_mask_loss": 0.0185,
+      "predictor_learning_rate": 0.0003806794576110801,
+      "step": 3369
+    },
+    {
+      "epoch": 0.35,
+      "gumbel_tau": 0.4517,
+      "image_masks_loss": 0.0165,
+      "llm_learning_rate": 3.8060872511252354e-06,
+      "loss": 0.8325,
+      "output_text_mask_loss": 0.0265,
+      "predictor_learning_rate": 0.0003806087251125235,
+      "step": 3370
+    },
+    {
+      "epoch": 0.35,
+      "gumbel_tau": 0.4516,
+      "image_masks_loss": 0.0118,
+      "llm_learning_rate": 3.8053797823104322e-06,
+      "loss": 0.7661,
+      "output_text_mask_loss": 0.0004,
+      "predictor_learning_rate": 0.00038053797823104317,
+      "step": 3371
+    },
+    {
+      "epoch": 0.35,
+      "gumbel_tau": 0.4515,
+      "image_masks_loss": 0.0189,
+      "llm_learning_rate": 3.8046721697443e-06,
+      "loss": 0.6721,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 0.00038046721697442996,
+      "step": 3372
+    },
+    {
+      "epoch": 0.35,
+      "gumbel_tau": 0.4513,
+      "image_masks_loss": 0.0138,
+      "llm_learning_rate": 3.8039644135047625e-06,
+      "loss": 0.7779,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00038039644135047626,
+      "step": 3373
+    },
+    {
+      "epoch": 0.35,
+      "gumbel_tau": 0.4512,
+      "image_masks_loss": 0.0155,
+      "llm_learning_rate": 3.8032565136697593e-06,
+      "loss": 0.7318,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.0003803256513669759,
+      "step": 3374
+    },
+    {
+      "epoch": 0.35,
+      "gumbel_tau": 0.4511,
+      "image_masks_loss": 0.0102,
+      "llm_learning_rate": 3.8025484703172455e-06,
+      "loss": 0.7275,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0003802548470317245,
+      "step": 3375
+    },
+    {
+      "epoch": 0.35,
+      "gumbel_tau": 0.451,
+      "image_masks_loss": 0.0058,
+      "llm_learning_rate": 3.8018402835251933e-06,
+      "loss": 0.7467,
+      "output_text_mask_loss": 0.0019,
+      "predictor_learning_rate": 0.0003801840283525193,
+      "step": 3376
+    },
+    {
+      "epoch": 0.35,
+      "gumbel_tau": 0.4509,
+      "image_masks_loss": 0.0237,
+      "llm_learning_rate": 3.8011319533715892e-06,
+      "loss": 0.7473,
+      "output_text_mask_loss": 0.0108,
+      "predictor_learning_rate": 0.0003801131953371589,
+      "step": 3377
+    },
+    {
+      "epoch": 0.35,
+      "gumbel_tau": 0.4508,
+      "image_masks_loss": 0.0127,
+      "llm_learning_rate": 3.8004234799344364e-06,
+      "loss": 0.8216,
+      "output_text_mask_loss": 0.0138,
+      "predictor_learning_rate": 0.0003800423479934436,
+      "step": 3378
+    },
+    {
+      "epoch": 0.35,
+      "gumbel_tau": 0.4507,
+      "image_masks_loss": 0.0145,
+      "llm_learning_rate": 3.7997148632917533e-06,
+      "loss": 0.6773,
+      "output_text_mask_loss": 0.0038,
+      "predictor_learning_rate": 0.00037997148632917533,
+      "step": 3379
+    },
+    {
+      "epoch": 0.35,
+      "gumbel_tau": 0.4506,
+      "image_masks_loss": 0.0064,
+      "llm_learning_rate": 3.7990061035215743e-06,
+      "loss": 0.7855,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.0003799006103521574,
+      "step": 3380
+    },
+    {
+      "epoch": 0.35,
+      "gumbel_tau": 0.4505,
+      "image_masks_loss": 0.014,
+      "llm_learning_rate": 3.798297200701949e-06,
+      "loss": 0.7232,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0003798297200701949,
+      "step": 3381
+    },
+    {
+      "epoch": 0.35,
+      "gumbel_tau": 0.4504,
+      "image_masks_loss": 0.034,
+      "llm_learning_rate": 3.7975881549109457e-06,
+      "loss": 0.7506,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0003797588154910945,
+      "step": 3382
+    },
+    {
+      "epoch": 0.35,
+      "gumbel_tau": 0.4503,
+      "image_masks_loss": 0.011,
+      "llm_learning_rate": 3.7968789662266437e-06,
+      "loss": 0.6723,
+      "output_text_mask_loss": 0.0198,
+      "predictor_learning_rate": 0.0003796878966226643,
+      "step": 3383
+    },
+    {
+      "epoch": 0.35,
+      "gumbel_tau": 0.4502,
+      "image_masks_loss": 0.0492,
+      "llm_learning_rate": 3.7961696347271417e-06,
+      "loss": 0.6875,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00037961696347271415,
+      "step": 3384
+    },
+    {
+      "epoch": 0.35,
+      "gumbel_tau": 0.4501,
+      "image_masks_loss": 0.024,
+      "llm_learning_rate": 3.7954601604905517e-06,
+      "loss": 0.6209,
+      "output_text_mask_loss": 0.0137,
+      "predictor_learning_rate": 0.0003795460160490551,
+      "step": 3385
+    },
+    {
+      "epoch": 0.35,
+      "gumbel_tau": 0.45,
+      "image_masks_loss": 0.0178,
+      "llm_learning_rate": 3.7947505435950037e-06,
+      "loss": 0.7783,
+      "output_text_mask_loss": 0.0046,
+      "predictor_learning_rate": 0.00037947505435950035,
+      "step": 3386
+    },
+    {
+      "epoch": 0.35,
+      "gumbel_tau": 0.4499,
+      "image_masks_loss": 0.0246,
+      "llm_learning_rate": 3.7940407841186424e-06,
+      "loss": 0.7363,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 0.0003794040784118642,
+      "step": 3387
+    },
+    {
+      "epoch": 0.35,
+      "gumbel_tau": 0.4498,
+      "image_masks_loss": 0.0225,
+      "llm_learning_rate": 3.7933308821396264e-06,
+      "loss": 0.7398,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.0003793330882139626,
+      "step": 3388
+    },
+    {
+      "epoch": 0.35,
+      "gumbel_tau": 0.4496,
+      "image_masks_loss": 0.0117,
+      "llm_learning_rate": 3.792620837736134e-06,
+      "loss": 0.7633,
+      "output_text_mask_loss": 0.0317,
+      "predictor_learning_rate": 0.0003792620837736134,
+      "step": 3389
+    },
+    {
+      "epoch": 0.35,
+      "gumbel_tau": 0.4495,
+      "image_masks_loss": 0.021,
+      "llm_learning_rate": 3.7919106509863546e-06,
+      "loss": 0.7689,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 0.00037919106509863543,
+      "step": 3390
+    },
+    {
+      "epoch": 0.35,
+      "gumbel_tau": 0.4494,
+      "image_masks_loss": 0.0257,
+      "llm_learning_rate": 3.791200321968498e-06,
+      "loss": 0.7293,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.0003791200321968498,
+      "step": 3391
+    },
+    {
+      "epoch": 0.35,
+      "gumbel_tau": 0.4493,
+      "image_masks_loss": 0.0261,
+      "llm_learning_rate": 3.7904898507607857e-06,
+      "loss": 0.7252,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.00037904898507607854,
+      "step": 3392
+    },
+    {
+      "epoch": 0.35,
+      "gumbel_tau": 0.4492,
+      "image_masks_loss": 0.0208,
+      "llm_learning_rate": 3.7897792374414554e-06,
+      "loss": 0.8009,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 0.00037897792374414553,
+      "step": 3393
+    },
+    {
+      "epoch": 0.35,
+      "gumbel_tau": 0.4491,
+      "image_masks_loss": 0.0296,
+      "llm_learning_rate": 3.7890684820887636e-06,
+      "loss": 0.706,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 0.0003789068482088763,
+      "step": 3394
+    },
+    {
+      "epoch": 0.35,
+      "gumbel_tau": 0.449,
+      "image_masks_loss": 0.0309,
+      "llm_learning_rate": 3.7883575847809795e-06,
+      "loss": 0.7264,
+      "output_text_mask_loss": 0.0275,
+      "predictor_learning_rate": 0.00037883575847809793,
+      "step": 3395
+    },
+    {
+      "epoch": 0.35,
+      "gumbel_tau": 0.4489,
+      "image_masks_loss": 0.0587,
+      "llm_learning_rate": 3.787646545596387e-06,
+      "loss": 0.8194,
+      "output_text_mask_loss": 0.0035,
+      "predictor_learning_rate": 0.0003787646545596387,
+      "step": 3396
+    },
+    {
+      "epoch": 0.35,
+      "gumbel_tau": 0.4488,
+      "image_masks_loss": 0.0454,
+      "llm_learning_rate": 3.78693536461329e-06,
+      "loss": 0.6897,
+      "output_text_mask_loss": 0.0041,
+      "predictor_learning_rate": 0.00037869353646132896,
+      "step": 3397
+    },
+    {
+      "epoch": 0.35,
+      "gumbel_tau": 0.4487,
+      "image_masks_loss": 0.0307,
+      "llm_learning_rate": 3.7862240419100036e-06,
+      "loss": 0.717,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00037862240419100034,
+      "step": 3398
+    },
+    {
+      "epoch": 0.35,
+      "gumbel_tau": 0.4486,
+      "image_masks_loss": 0.0223,
+      "llm_learning_rate": 3.7855125775648615e-06,
+      "loss": 0.7192,
+      "output_text_mask_loss": 0.0385,
+      "predictor_learning_rate": 0.00037855125775648616,
+      "step": 3399
+    },
+    {
+      "epoch": 0.35,
+      "gumbel_tau": 0.4485,
+      "image_masks_loss": 0.0345,
+      "llm_learning_rate": 3.784800971656211e-06,
+      "loss": 0.7897,
+      "output_text_mask_loss": 0.0164,
+      "predictor_learning_rate": 0.00037848009716562103,
+      "step": 3400
+    },
+    {
+      "epoch": 0.35,
+      "gumbel_tau": 0.4484,
+      "image_masks_loss": 0.0338,
+      "llm_learning_rate": 3.784089224262415e-06,
+      "loss": 0.7871,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.00037840892242624145,
+      "step": 3401
+    },
+    {
+      "epoch": 0.35,
+      "gumbel_tau": 0.4483,
+      "image_masks_loss": 0.0234,
+      "llm_learning_rate": 3.783377335461854e-06,
+      "loss": 0.7562,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0003783377335461854,
+      "step": 3402
+    },
+    {
+      "epoch": 0.35,
+      "gumbel_tau": 0.4482,
+      "image_masks_loss": 0.0109,
+      "llm_learning_rate": 3.782665305332921e-06,
+      "loss": 0.7605,
+      "output_text_mask_loss": 0.0847,
+      "predictor_learning_rate": 0.0003782665305332921,
+      "step": 3403
+    },
+    {
+      "epoch": 0.35,
+      "gumbel_tau": 0.4481,
+      "image_masks_loss": 0.0089,
+      "llm_learning_rate": 3.781953133954029e-06,
+      "loss": 0.6621,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00037819531339540284,
+      "step": 3404
+    },
+    {
+      "epoch": 0.35,
+      "gumbel_tau": 0.4479,
+      "image_masks_loss": 0.028,
+      "llm_learning_rate": 3.7812408214036015e-06,
+      "loss": 0.7953,
+      "output_text_mask_loss": 0.0679,
+      "predictor_learning_rate": 0.00037812408214036013,
+      "step": 3405
+    },
+    {
+      "epoch": 0.35,
+      "gumbel_tau": 0.4478,
+      "image_masks_loss": 0.0063,
+      "llm_learning_rate": 3.7805283677600814e-06,
+      "loss": 0.723,
+      "output_text_mask_loss": 0.017,
+      "predictor_learning_rate": 0.0003780528367760081,
+      "step": 3406
+    },
+    {
+      "epoch": 0.35,
+      "gumbel_tau": 0.4477,
+      "image_masks_loss": 0.0126,
+      "llm_learning_rate": 3.779815773101925e-06,
+      "loss": 0.7853,
+      "output_text_mask_loss": 0.0641,
+      "predictor_learning_rate": 0.00037798157731019243,
+      "step": 3407
+    },
+    {
+      "epoch": 0.35,
+      "gumbel_tau": 0.4476,
+      "image_masks_loss": 0.0317,
+      "llm_learning_rate": 3.7791030375076055e-06,
+      "loss": 0.772,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.0003779103037507605,
+      "step": 3408
+    },
+    {
+      "epoch": 0.35,
+      "gumbel_tau": 0.4475,
+      "image_masks_loss": 0.0068,
+      "llm_learning_rate": 3.7783901610556097e-06,
+      "loss": 0.7334,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00037783901610556093,
+      "step": 3409
+    },
+    {
+      "epoch": 0.35,
+      "gumbel_tau": 0.4474,
+      "image_masks_loss": 0.0043,
+      "llm_learning_rate": 3.7776771438244418e-06,
+      "loss": 0.7386,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.0003777677143824442,
+      "step": 3410
+    },
+    {
+      "epoch": 0.35,
+      "gumbel_tau": 0.4473,
+      "image_masks_loss": 0.0103,
+      "llm_learning_rate": 3.776963985892622e-06,
+      "loss": 0.7513,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 0.0003776963985892622,
+      "step": 3411
+    },
+    {
+      "epoch": 0.35,
+      "gumbel_tau": 0.4472,
+      "image_masks_loss": 0.0097,
+      "llm_learning_rate": 3.776250687338683e-06,
+      "loss": 0.8032,
+      "output_text_mask_loss": 0.0538,
+      "predictor_learning_rate": 0.00037762506873386824,
+      "step": 3412
+    },
+    {
+      "epoch": 0.35,
+      "gumbel_tau": 0.4471,
+      "image_masks_loss": 0.0104,
+      "llm_learning_rate": 3.775537248241175e-06,
+      "loss": 0.7453,
+      "output_text_mask_loss": 0.0277,
+      "predictor_learning_rate": 0.0003775537248241175,
+      "step": 3413
+    },
+    {
+      "epoch": 0.35,
+      "gumbel_tau": 0.447,
+      "image_masks_loss": 0.0117,
+      "llm_learning_rate": 3.7748236686786645e-06,
+      "loss": 0.78,
+      "output_text_mask_loss": 0.0103,
+      "predictor_learning_rate": 0.00037748236686786644,
+      "step": 3414
+    },
+    {
+      "epoch": 0.35,
+      "gumbel_tau": 0.4469,
+      "image_masks_loss": 0.0216,
+      "llm_learning_rate": 3.774109948729733e-06,
+      "loss": 0.6831,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.00037741099487297324,
+      "step": 3415
+    },
+    {
+      "epoch": 0.35,
+      "gumbel_tau": 0.4468,
+      "image_masks_loss": 0.0236,
+      "llm_learning_rate": 3.7733960884729742e-06,
+      "loss": 0.8147,
+      "output_text_mask_loss": 0.0774,
+      "predictor_learning_rate": 0.0003773396088472974,
+      "step": 3416
+    },
+    {
+      "epoch": 0.35,
+      "gumbel_tau": 0.4467,
+      "image_masks_loss": 0.0303,
+      "llm_learning_rate": 3.772682087987003e-06,
+      "loss": 0.7587,
+      "output_text_mask_loss": 0.0296,
+      "predictor_learning_rate": 0.0003772682087987003,
+      "step": 3417
+    },
+    {
+      "epoch": 0.35,
+      "gumbel_tau": 0.4466,
+      "image_masks_loss": 0.0217,
+      "llm_learning_rate": 3.771967947350445e-06,
+      "loss": 0.7679,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.0003771967947350445,
+      "step": 3418
+    },
+    {
+      "epoch": 0.35,
+      "gumbel_tau": 0.4465,
+      "image_masks_loss": 0.0122,
+      "llm_learning_rate": 3.7712536666419443e-06,
+      "loss": 0.7687,
+      "output_text_mask_loss": 0.037,
+      "predictor_learning_rate": 0.0003771253666641944,
+      "step": 3419
+    },
+    {
+      "epoch": 0.35,
+      "gumbel_tau": 0.4464,
+      "image_masks_loss": 0.0288,
+      "llm_learning_rate": 3.7705392459401567e-06,
+      "loss": 0.7487,
+      "output_text_mask_loss": 0.021,
+      "predictor_learning_rate": 0.0003770539245940157,
+      "step": 3420
+    },
+    {
+      "epoch": 0.35,
+      "gumbel_tau": 0.4463,
+      "image_masks_loss": 0.0153,
+      "llm_learning_rate": 3.7698246853237584e-06,
+      "loss": 0.7113,
+      "output_text_mask_loss": 0.0759,
+      "predictor_learning_rate": 0.0003769824685323758,
+      "step": 3421
+    },
+    {
+      "epoch": 0.35,
+      "gumbel_tau": 0.4462,
+      "image_masks_loss": 0.0271,
+      "llm_learning_rate": 3.7691099848714375e-06,
+      "loss": 0.7653,
+      "output_text_mask_loss": 0.021,
+      "predictor_learning_rate": 0.0003769109984871437,
+      "step": 3422
+    },
+    {
+      "epoch": 0.35,
+      "gumbel_tau": 0.4461,
+      "image_masks_loss": 0.0203,
+      "llm_learning_rate": 3.7683951446618973e-06,
+      "loss": 0.7372,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0003768395144661897,
+      "step": 3423
+    },
+    {
+      "epoch": 0.35,
+      "gumbel_tau": 0.4459,
+      "image_masks_loss": 0.0243,
+      "llm_learning_rate": 3.767680164773859e-06,
+      "loss": 0.7726,
+      "output_text_mask_loss": 0.0343,
+      "predictor_learning_rate": 0.0003767680164773859,
+      "step": 3424
+    },
+    {
+      "epoch": 0.35,
+      "gumbel_tau": 0.4458,
+      "image_masks_loss": 0.032,
+      "llm_learning_rate": 3.7669650452860577e-06,
+      "loss": 0.789,
+      "output_text_mask_loss": 0.0244,
+      "predictor_learning_rate": 0.0003766965045286058,
+      "step": 3425
+    },
+    {
+      "epoch": 0.35,
+      "gumbel_tau": 0.4457,
+      "image_masks_loss": 0.0237,
+      "llm_learning_rate": 3.7662497862772428e-06,
+      "loss": 0.8103,
+      "output_text_mask_loss": 0.0675,
+      "predictor_learning_rate": 0.00037662497862772426,
+      "step": 3426
+    },
+    {
+      "epoch": 0.35,
+      "gumbel_tau": 0.4456,
+      "image_masks_loss": 0.0471,
+      "llm_learning_rate": 3.765534387826182e-06,
+      "loss": 0.8728,
+      "output_text_mask_loss": 0.0725,
+      "predictor_learning_rate": 0.00037655343878261816,
+      "step": 3427
+    },
+    {
+      "epoch": 0.35,
+      "gumbel_tau": 0.4455,
+      "image_masks_loss": 0.0263,
+      "llm_learning_rate": 3.7648188500116547e-06,
+      "loss": 0.8281,
+      "output_text_mask_loss": 0.0139,
+      "predictor_learning_rate": 0.0003764818850011654,
+      "step": 3428
+    },
+    {
+      "epoch": 0.35,
+      "gumbel_tau": 0.4454,
+      "image_masks_loss": 0.0326,
+      "llm_learning_rate": 3.7641031729124588e-06,
+      "loss": 0.751,
+      "output_text_mask_loss": 0.0032,
+      "predictor_learning_rate": 0.0003764103172912459,
+      "step": 3429
+    },
+    {
+      "epoch": 0.35,
+      "gumbel_tau": 0.4453,
+      "image_masks_loss": 0.0108,
+      "llm_learning_rate": 3.7633873566074058e-06,
+      "loss": 0.7296,
+      "output_text_mask_loss": 0.0603,
+      "predictor_learning_rate": 0.00037633873566074054,
+      "step": 3430
+    },
+    {
+      "epoch": 0.35,
+      "gumbel_tau": 0.4452,
+      "image_masks_loss": 0.0236,
+      "llm_learning_rate": 3.762671401175323e-06,
+      "loss": 0.7397,
+      "output_text_mask_loss": 0.1938,
+      "predictor_learning_rate": 0.0003762671401175323,
+      "step": 3431
+    },
+    {
+      "epoch": 0.35,
+      "gumbel_tau": 0.4451,
+      "image_masks_loss": 0.0362,
+      "llm_learning_rate": 3.761955306695052e-06,
+      "loss": 0.7179,
+      "output_text_mask_loss": 0.0507,
+      "predictor_learning_rate": 0.0003761955306695052,
+      "step": 3432
+    },
+    {
+      "epoch": 0.35,
+      "gumbel_tau": 0.445,
+      "image_masks_loss": 0.0159,
+      "llm_learning_rate": 3.761239073245453e-06,
+      "loss": 0.798,
+      "output_text_mask_loss": 0.0025,
+      "predictor_learning_rate": 0.00037612390732454527,
+      "step": 3433
+    },
+    {
+      "epoch": 0.35,
+      "gumbel_tau": 0.4449,
+      "image_masks_loss": 0.0603,
+      "llm_learning_rate": 3.7605227009053978e-06,
+      "loss": 0.7874,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00037605227009053976,
+      "step": 3434
+    },
+    {
+      "epoch": 0.35,
+      "gumbel_tau": 0.4448,
+      "image_masks_loss": 0.0229,
+      "llm_learning_rate": 3.7598061897537753e-06,
+      "loss": 0.7539,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.0003759806189753775,
+      "step": 3435
+    },
+    {
+      "epoch": 0.35,
+      "gumbel_tau": 0.4447,
+      "image_masks_loss": 0.0063,
+      "llm_learning_rate": 3.7590895398694878e-06,
+      "loss": 0.836,
+      "output_text_mask_loss": 0.0626,
+      "predictor_learning_rate": 0.00037590895398694873,
+      "step": 3436
+    },
+    {
+      "epoch": 0.35,
+      "gumbel_tau": 0.4446,
+      "image_masks_loss": 0.0307,
+      "llm_learning_rate": 3.758372751331456e-06,
+      "loss": 0.7558,
+      "output_text_mask_loss": 0.0305,
+      "predictor_learning_rate": 0.0003758372751331456,
+      "step": 3437
+    },
+    {
+      "epoch": 0.35,
+      "gumbel_tau": 0.4445,
+      "image_masks_loss": 0.038,
+      "llm_learning_rate": 3.7576558242186157e-06,
+      "loss": 0.9015,
+      "output_text_mask_loss": 0.0058,
+      "predictor_learning_rate": 0.0003757655824218615,
+      "step": 3438
+    },
+    {
+      "epoch": 0.35,
+      "gumbel_tau": 0.4444,
+      "image_masks_loss": 0.0206,
+      "llm_learning_rate": 3.7569387586099126e-06,
+      "loss": 0.786,
+      "output_text_mask_loss": 0.0202,
+      "predictor_learning_rate": 0.00037569387586099123,
+      "step": 3439
+    },
+    {
+      "epoch": 0.35,
+      "gumbel_tau": 0.4443,
+      "image_masks_loss": 0.0217,
+      "llm_learning_rate": 3.756221554584315e-06,
+      "loss": 0.6784,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.0003756221554584315,
+      "step": 3440
+    },
+    {
+      "epoch": 0.35,
+      "gumbel_tau": 0.4442,
+      "image_masks_loss": 0.0296,
+      "llm_learning_rate": 3.7555042122208006e-06,
+      "loss": 0.9097,
+      "output_text_mask_loss": 0.0166,
+      "predictor_learning_rate": 0.00037555042122208005,
+      "step": 3441
+    },
+    {
+      "epoch": 0.35,
+      "gumbel_tau": 0.4441,
+      "image_masks_loss": 0.0706,
+      "llm_learning_rate": 3.754786731598367e-06,
+      "loss": 0.863,
+      "output_text_mask_loss": 0.0098,
+      "predictor_learning_rate": 0.00037547867315983663,
+      "step": 3442
+    },
+    {
+      "epoch": 0.35,
+      "gumbel_tau": 0.444,
+      "image_masks_loss": 0.0456,
+      "llm_learning_rate": 3.754069112796023e-06,
+      "loss": 0.7639,
+      "output_text_mask_loss": 0.0376,
+      "predictor_learning_rate": 0.00037540691127960224,
+      "step": 3443
+    },
+    {
+      "epoch": 0.35,
+      "gumbel_tau": 0.4438,
+      "image_masks_loss": 0.0102,
+      "llm_learning_rate": 3.753351355892794e-06,
+      "loss": 0.7978,
+      "output_text_mask_loss": 0.0038,
+      "predictor_learning_rate": 0.0003753351355892794,
+      "step": 3444
+    },
+    {
+      "epoch": 0.35,
+      "gumbel_tau": 0.4437,
+      "image_masks_loss": 0.0345,
+      "llm_learning_rate": 3.7526334609677234e-06,
+      "loss": 0.8692,
+      "output_text_mask_loss": 0.1007,
+      "predictor_learning_rate": 0.0003752633460967723,
+      "step": 3445
+    },
+    {
+      "epoch": 0.35,
+      "gumbel_tau": 0.4436,
+      "image_masks_loss": 0.0671,
+      "llm_learning_rate": 3.751915428099865e-06,
+      "loss": 0.7588,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00037519154280998645,
+      "step": 3446
+    },
+    {
+      "epoch": 0.35,
+      "gumbel_tau": 0.4435,
+      "image_masks_loss": 0.0071,
+      "llm_learning_rate": 3.751197257368292e-06,
+      "loss": 0.6977,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0003751197257368292,
+      "step": 3447
+    },
+    {
+      "epoch": 0.35,
+      "gumbel_tau": 0.4434,
+      "image_masks_loss": 0.0359,
+      "llm_learning_rate": 3.750478948852089e-06,
+      "loss": 0.8869,
+      "output_text_mask_loss": 0.0519,
+      "predictor_learning_rate": 0.0003750478948852089,
+      "step": 3448
+    },
+    {
+      "epoch": 0.35,
+      "gumbel_tau": 0.4433,
+      "image_masks_loss": 0.0219,
+      "llm_learning_rate": 3.7497605026303595e-06,
+      "loss": 0.8176,
+      "output_text_mask_loss": 0.0231,
+      "predictor_learning_rate": 0.00037497605026303593,
+      "step": 3449
+    },
+    {
+      "epoch": 0.35,
+      "gumbel_tau": 0.4432,
+      "image_masks_loss": 0.0393,
+      "llm_learning_rate": 3.7490419187822204e-06,
+      "loss": 0.6869,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.000374904191878222,
+      "step": 3450
+    },
+    {
+      "epoch": 0.35,
+      "gumbel_tau": 0.4431,
+      "image_masks_loss": 0.0385,
+      "llm_learning_rate": 3.748323197386802e-06,
+      "loss": 0.7762,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0003748323197386802,
+      "step": 3451
+    },
+    {
+      "epoch": 0.35,
+      "gumbel_tau": 0.443,
+      "image_masks_loss": 0.0243,
+      "llm_learning_rate": 3.747604338523253e-06,
+      "loss": 0.7595,
+      "output_text_mask_loss": 0.0107,
+      "predictor_learning_rate": 0.0003747604338523253,
+      "step": 3452
+    },
+    {
+      "epoch": 0.35,
+      "gumbel_tau": 0.4429,
+      "image_masks_loss": 0.0357,
+      "llm_learning_rate": 3.746885342270735e-06,
+      "loss": 0.7858,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0003746885342270735,
+      "step": 3453
+    },
+    {
+      "epoch": 0.35,
+      "gumbel_tau": 0.4428,
+      "image_masks_loss": 0.0084,
+      "llm_learning_rate": 3.7461662087084262e-06,
+      "loss": 0.7072,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.00037461662087084256,
+      "step": 3454
+    },
+    {
+      "epoch": 0.35,
+      "gumbel_tau": 0.4427,
+      "image_masks_loss": 0.0299,
+      "llm_learning_rate": 3.745446937915518e-06,
+      "loss": 0.7415,
+      "output_text_mask_loss": 0.001,
+      "predictor_learning_rate": 0.0003745446937915518,
+      "step": 3455
+    },
+    {
+      "epoch": 0.35,
+      "gumbel_tau": 0.4426,
+      "image_masks_loss": 0.0103,
+      "llm_learning_rate": 3.74472752997122e-06,
+      "loss": 0.7224,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00037447275299712195,
+      "step": 3456
+    },
+    {
+      "epoch": 0.35,
+      "gumbel_tau": 0.4425,
+      "image_masks_loss": 0.0157,
+      "llm_learning_rate": 3.7440079849547537e-06,
+      "loss": 0.8022,
+      "output_text_mask_loss": 0.0381,
+      "predictor_learning_rate": 0.00037440079849547535,
+      "step": 3457
+    },
+    {
+      "epoch": 0.35,
+      "gumbel_tau": 0.4424,
+      "image_masks_loss": 0.0086,
+      "llm_learning_rate": 3.7432883029453565e-06,
+      "loss": 0.7368,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00037432883029453565,
+      "step": 3458
+    },
+    {
+      "epoch": 0.35,
+      "gumbel_tau": 0.4423,
+      "image_masks_loss": 0.0349,
+      "llm_learning_rate": 3.742568484022282e-06,
+      "loss": 0.8211,
+      "output_text_mask_loss": 0.0075,
+      "predictor_learning_rate": 0.0003742568484022282,
+      "step": 3459
+    },
+    {
+      "epoch": 0.35,
+      "gumbel_tau": 0.4422,
+      "image_masks_loss": 0.0173,
+      "llm_learning_rate": 3.741848528264799e-06,
+      "loss": 0.6807,
+      "output_text_mask_loss": 0.0208,
+      "predictor_learning_rate": 0.0003741848528264799,
+      "step": 3460
+    },
+    {
+      "epoch": 0.35,
+      "gumbel_tau": 0.4421,
+      "image_masks_loss": 0.0062,
+      "llm_learning_rate": 3.74112843575219e-06,
+      "loss": 0.7298,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00037411284357521897,
+      "step": 3461
+    },
+    {
+      "epoch": 0.35,
+      "gumbel_tau": 0.442,
+      "image_masks_loss": 0.0261,
+      "llm_learning_rate": 3.7404082065637526e-06,
+      "loss": 0.6935,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.00037404082065637523,
+      "step": 3462
+    },
+    {
+      "epoch": 0.35,
+      "gumbel_tau": 0.4419,
+      "image_masks_loss": 0.028,
+      "llm_learning_rate": 3.7396878407788005e-06,
+      "loss": 0.7209,
+      "output_text_mask_loss": 0.0118,
+      "predictor_learning_rate": 0.00037396878407788006,
+      "step": 3463
+    },
+    {
+      "epoch": 0.35,
+      "gumbel_tau": 0.4418,
+      "image_masks_loss": 0.0307,
+      "llm_learning_rate": 3.7389673384766623e-06,
+      "loss": 0.7195,
+      "output_text_mask_loss": 0.0431,
+      "predictor_learning_rate": 0.0003738967338476662,
+      "step": 3464
+    },
+    {
+      "epoch": 0.36,
+      "gumbel_tau": 0.4417,
+      "image_masks_loss": 0.0345,
+      "llm_learning_rate": 3.7382466997366812e-06,
+      "loss": 0.7465,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 0.0003738246699736681,
+      "step": 3465
+    },
+    {
+      "epoch": 0.36,
+      "gumbel_tau": 0.4415,
+      "image_masks_loss": 0.0045,
+      "llm_learning_rate": 3.7375259246382154e-06,
+      "loss": 0.703,
+      "output_text_mask_loss": 0.0124,
+      "predictor_learning_rate": 0.00037375259246382153,
+      "step": 3466
+    },
+    {
+      "epoch": 0.36,
+      "gumbel_tau": 0.4414,
+      "image_masks_loss": 0.0074,
+      "llm_learning_rate": 3.7368050132606386e-06,
+      "loss": 0.7889,
+      "output_text_mask_loss": 0.0231,
+      "predictor_learning_rate": 0.0003736805013260639,
+      "step": 3467
+    },
+    {
+      "epoch": 0.36,
+      "gumbel_tau": 0.4413,
+      "image_masks_loss": 0.0225,
+      "llm_learning_rate": 3.7360839656833386e-06,
+      "loss": 0.7735,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.0003736083965683338,
+      "step": 3468
+    },
+    {
+      "epoch": 0.36,
+      "gumbel_tau": 0.4412,
+      "image_masks_loss": 0.0139,
+      "llm_learning_rate": 3.735362781985719e-06,
+      "loss": 0.6541,
+      "output_text_mask_loss": 0.0629,
+      "predictor_learning_rate": 0.00037353627819857187,
+      "step": 3469
+    },
+    {
+      "epoch": 0.36,
+      "gumbel_tau": 0.4411,
+      "image_masks_loss": 0.0406,
+      "llm_learning_rate": 3.734641462247198e-06,
+      "loss": 0.8288,
+      "output_text_mask_loss": 0.0429,
+      "predictor_learning_rate": 0.0003734641462247198,
+      "step": 3470
+    },
+    {
+      "epoch": 0.36,
+      "gumbel_tau": 0.441,
+      "image_masks_loss": 0.0169,
+      "llm_learning_rate": 3.733920006547209e-06,
+      "loss": 0.6493,
+      "output_text_mask_loss": 0.0215,
+      "predictor_learning_rate": 0.0003733920006547209,
+      "step": 3471
+    },
+    {
+      "epoch": 0.36,
+      "gumbel_tau": 0.4409,
+      "image_masks_loss": 0.0237,
+      "llm_learning_rate": 3.7331984149652013e-06,
+      "loss": 0.7241,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0003733198414965201,
+      "step": 3472
+    },
+    {
+      "epoch": 0.36,
+      "gumbel_tau": 0.4408,
+      "image_masks_loss": 0.0096,
+      "llm_learning_rate": 3.732476687580637e-06,
+      "loss": 0.8619,
+      "output_text_mask_loss": 0.0771,
+      "predictor_learning_rate": 0.00037324766875806363,
+      "step": 3473
+    },
+    {
+      "epoch": 0.36,
+      "gumbel_tau": 0.4407,
+      "image_masks_loss": 0.0205,
+      "llm_learning_rate": 3.731754824472994e-06,
+      "loss": 0.7626,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.0003731754824472994,
+      "step": 3474
+    },
+    {
+      "epoch": 0.36,
+      "gumbel_tau": 0.4406,
+      "image_masks_loss": 0.0158,
+      "llm_learning_rate": 3.7310328257217664e-06,
+      "loss": 0.8191,
+      "output_text_mask_loss": 0.005,
+      "predictor_learning_rate": 0.0003731032825721766,
+      "step": 3475
+    },
+    {
+      "epoch": 0.36,
+      "gumbel_tau": 0.4405,
+      "image_masks_loss": 0.0224,
+      "llm_learning_rate": 3.7303106914064623e-06,
+      "loss": 0.7081,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 0.00037303106914064624,
+      "step": 3476
+    },
+    {
+      "epoch": 0.36,
+      "gumbel_tau": 0.4404,
+      "image_masks_loss": 0.0278,
+      "llm_learning_rate": 3.729588421606605e-06,
+      "loss": 0.8143,
+      "output_text_mask_loss": 0.001,
+      "predictor_learning_rate": 0.00037295884216066044,
+      "step": 3477
+    },
+    {
+      "epoch": 0.36,
+      "gumbel_tau": 0.4403,
+      "image_masks_loss": 0.0179,
+      "llm_learning_rate": 3.7288660164017303e-06,
+      "loss": 0.7694,
+      "output_text_mask_loss": 0.0061,
+      "predictor_learning_rate": 0.000372886601640173,
+      "step": 3478
+    },
+    {
+      "epoch": 0.36,
+      "gumbel_tau": 0.4402,
+      "image_masks_loss": 0.0192,
+      "llm_learning_rate": 3.7281434758713934e-06,
+      "loss": 0.7509,
+      "output_text_mask_loss": 0.0298,
+      "predictor_learning_rate": 0.0003728143475871393,
+      "step": 3479
+    },
+    {
+      "epoch": 0.36,
+      "gumbel_tau": 0.4401,
+      "image_masks_loss": 0.0041,
+      "llm_learning_rate": 3.7274208000951618e-06,
+      "loss": 0.786,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.0003727420800095162,
+      "step": 3480
+    },
+    {
+      "epoch": 0.36,
+      "gumbel_tau": 0.44,
+      "image_masks_loss": 0.029,
+      "llm_learning_rate": 3.7266979891526168e-06,
+      "loss": 0.7319,
+      "output_text_mask_loss": 0.0116,
+      "predictor_learning_rate": 0.00037266979891526164,
+      "step": 3481
+    },
+    {
+      "epoch": 0.36,
+      "gumbel_tau": 0.4399,
+      "image_masks_loss": 0.038,
+      "llm_learning_rate": 3.725975043123358e-06,
+      "loss": 0.7664,
+      "output_text_mask_loss": 0.0099,
+      "predictor_learning_rate": 0.0003725975043123357,
+      "step": 3482
+    },
+    {
+      "epoch": 0.36,
+      "gumbel_tau": 0.4398,
+      "image_masks_loss": 0.0161,
+      "llm_learning_rate": 3.725251962086996e-06,
+      "loss": 0.818,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.0003725251962086996,
+      "step": 3483
+    },
+    {
+      "epoch": 0.36,
+      "gumbel_tau": 0.4397,
+      "image_masks_loss": 0.0129,
+      "llm_learning_rate": 3.7245287461231594e-06,
+      "loss": 0.7627,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.0003724528746123159,
+      "step": 3484
+    },
+    {
+      "epoch": 0.36,
+      "gumbel_tau": 0.4396,
+      "image_masks_loss": 0.0422,
+      "llm_learning_rate": 3.7238053953114896e-06,
+      "loss": 0.6725,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00037238053953114893,
+      "step": 3485
+    },
+    {
+      "epoch": 0.36,
+      "gumbel_tau": 0.4395,
+      "image_masks_loss": 0.0222,
+      "llm_learning_rate": 3.723081909731643e-06,
+      "loss": 0.7983,
+      "output_text_mask_loss": 0.0063,
+      "predictor_learning_rate": 0.00037230819097316427,
+      "step": 3486
+    },
+    {
+      "epoch": 0.36,
+      "gumbel_tau": 0.4394,
+      "image_masks_loss": 0.0393,
+      "llm_learning_rate": 3.7223582894632926e-06,
+      "loss": 0.7997,
+      "output_text_mask_loss": 0.0027,
+      "predictor_learning_rate": 0.0003722358289463292,
+      "step": 3487
+    },
+    {
+      "epoch": 0.36,
+      "gumbel_tau": 0.4393,
+      "image_masks_loss": 0.0166,
+      "llm_learning_rate": 3.721634534586125e-06,
+      "loss": 0.7187,
+      "output_text_mask_loss": 0.0231,
+      "predictor_learning_rate": 0.00037216345345861246,
+      "step": 3488
+    },
+    {
+      "epoch": 0.36,
+      "gumbel_tau": 0.4392,
+      "image_masks_loss": 0.0138,
+      "llm_learning_rate": 3.7209106451798405e-06,
+      "loss": 0.7341,
+      "output_text_mask_loss": 0.0156,
+      "predictor_learning_rate": 0.000372091064517984,
+      "step": 3489
+    },
+    {
+      "epoch": 0.36,
+      "gumbel_tau": 0.4391,
+      "image_masks_loss": 0.0154,
+      "llm_learning_rate": 3.720186621324157e-06,
+      "loss": 0.7852,
+      "output_text_mask_loss": 0.0088,
+      "predictor_learning_rate": 0.0003720186621324157,
+      "step": 3490
+    },
+    {
+      "epoch": 0.36,
+      "gumbel_tau": 0.439,
+      "image_masks_loss": 0.0119,
+      "llm_learning_rate": 3.7194624630988047e-06,
+      "loss": 0.7658,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.00037194624630988043,
+      "step": 3491
+    },
+    {
+      "epoch": 0.36,
+      "gumbel_tau": 0.4388,
+      "image_masks_loss": 0.0223,
+      "llm_learning_rate": 3.71873817058353e-06,
+      "loss": 0.7644,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.000371873817058353,
+      "step": 3492
+    },
+    {
+      "epoch": 0.36,
+      "gumbel_tau": 0.4387,
+      "image_masks_loss": 0.0221,
+      "llm_learning_rate": 3.7180137438580933e-06,
+      "loss": 0.8106,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.0003718013743858093,
+      "step": 3493
+    },
+    {
+      "epoch": 0.36,
+      "gumbel_tau": 0.4386,
+      "image_masks_loss": 0.021,
+      "llm_learning_rate": 3.7172891830022694e-06,
+      "loss": 0.8127,
+      "output_text_mask_loss": 0.0429,
+      "predictor_learning_rate": 0.0003717289183002269,
+      "step": 3494
+    },
+    {
+      "epoch": 0.36,
+      "gumbel_tau": 0.4385,
+      "image_masks_loss": 0.0106,
+      "llm_learning_rate": 3.7165644880958503e-06,
+      "loss": 0.8013,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.000371656448809585,
+      "step": 3495
+    },
+    {
+      "epoch": 0.36,
+      "gumbel_tau": 0.4384,
+      "image_masks_loss": 0.0062,
+      "llm_learning_rate": 3.7158396592186396e-06,
+      "loss": 0.7612,
+      "output_text_mask_loss": 0.0098,
+      "predictor_learning_rate": 0.00037158396592186396,
+      "step": 3496
+    },
+    {
+      "epoch": 0.36,
+      "gumbel_tau": 0.4383,
+      "image_masks_loss": 0.0117,
+      "llm_learning_rate": 3.715114696450457e-06,
+      "loss": 0.7502,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00037151146964504566,
+      "step": 3497
+    },
+    {
+      "epoch": 0.36,
+      "gumbel_tau": 0.4382,
+      "image_masks_loss": 0.0282,
+      "llm_learning_rate": 3.714389599871138e-06,
+      "loss": 0.7861,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00037143895998711376,
+      "step": 3498
+    },
+    {
+      "epoch": 0.36,
+      "gumbel_tau": 0.4381,
+      "image_masks_loss": 0.0111,
+      "llm_learning_rate": 3.713664369560531e-06,
+      "loss": 0.7706,
+      "output_text_mask_loss": 0.0254,
+      "predictor_learning_rate": 0.0003713664369560531,
+      "step": 3499
+    },
+    {
+      "epoch": 0.36,
+      "gumbel_tau": 0.438,
+      "image_masks_loss": 0.0254,
+      "llm_learning_rate": 3.7129390055985003e-06,
+      "loss": 0.8624,
+      "output_text_mask_loss": 0.0202,
+      "predictor_learning_rate": 0.00037129390055985,
+      "step": 3500
+    },
+    {
+      "epoch": 0.36,
+      "gumbel_tau": 0.4379,
+      "image_masks_loss": 0.0412,
+      "llm_learning_rate": 3.7122135080649247e-06,
+      "loss": 0.8156,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.00037122135080649245,
+      "step": 3501
+    },
+    {
+      "epoch": 0.36,
+      "gumbel_tau": 0.4378,
+      "image_masks_loss": 0.0156,
+      "llm_learning_rate": 3.7114878770396974e-06,
+      "loss": 0.8424,
+      "output_text_mask_loss": 0.0023,
+      "predictor_learning_rate": 0.0003711487877039697,
+      "step": 3502
+    },
+    {
+      "epoch": 0.36,
+      "gumbel_tau": 0.4377,
+      "image_masks_loss": 0.0167,
+      "llm_learning_rate": 3.7107621126027265e-06,
+      "loss": 0.7062,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.00037107621126027265,
+      "step": 3503
+    },
+    {
+      "epoch": 0.36,
+      "gumbel_tau": 0.4376,
+      "image_masks_loss": 0.0172,
+      "llm_learning_rate": 3.710036214833935e-06,
+      "loss": 0.7605,
+      "output_text_mask_loss": 0.0801,
+      "predictor_learning_rate": 0.0003710036214833935,
+      "step": 3504
+    },
+    {
+      "epoch": 0.36,
+      "gumbel_tau": 0.4375,
+      "image_masks_loss": 0.0117,
+      "llm_learning_rate": 3.7093101838132605e-06,
+      "loss": 0.7771,
+      "output_text_mask_loss": 0.0444,
+      "predictor_learning_rate": 0.00037093101838132605,
+      "step": 3505
+    },
+    {
+      "epoch": 0.36,
+      "gumbel_tau": 0.4374,
+      "image_masks_loss": 0.0284,
+      "llm_learning_rate": 3.708584019620655e-06,
+      "loss": 0.8078,
+      "output_text_mask_loss": 0.0172,
+      "predictor_learning_rate": 0.00037085840196206547,
+      "step": 3506
+    },
+    {
+      "epoch": 0.36,
+      "gumbel_tau": 0.4373,
+      "image_masks_loss": 0.0107,
+      "llm_learning_rate": 3.7078577223360847e-06,
+      "loss": 0.7358,
+      "output_text_mask_loss": 0.0345,
+      "predictor_learning_rate": 0.00037078577223360843,
+      "step": 3507
+    },
+    {
+      "epoch": 0.36,
+      "gumbel_tau": 0.4372,
+      "image_masks_loss": 0.0357,
+      "llm_learning_rate": 3.707131292039532e-06,
+      "loss": 0.7829,
+      "output_text_mask_loss": 0.1495,
+      "predictor_learning_rate": 0.00037071312920395315,
+      "step": 3508
+    },
+    {
+      "epoch": 0.36,
+      "gumbel_tau": 0.4371,
+      "image_masks_loss": 0.0246,
+      "llm_learning_rate": 3.706404728810993e-06,
+      "loss": 0.7564,
+      "output_text_mask_loss": 0.1045,
+      "predictor_learning_rate": 0.00037064047288109923,
+      "step": 3509
+    },
+    {
+      "epoch": 0.36,
+      "gumbel_tau": 0.437,
+      "image_masks_loss": 0.0238,
+      "llm_learning_rate": 3.705678032730477e-06,
+      "loss": 0.8268,
+      "output_text_mask_loss": 0.0595,
+      "predictor_learning_rate": 0.00037056780327304767,
+      "step": 3510
+    },
+    {
+      "epoch": 0.36,
+      "gumbel_tau": 0.4369,
+      "image_masks_loss": 0.028,
+      "llm_learning_rate": 3.7049512038780113e-06,
+      "loss": 0.8372,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0003704951203878011,
+      "step": 3511
+    },
+    {
+      "epoch": 0.36,
+      "gumbel_tau": 0.4368,
+      "image_masks_loss": 0.0074,
+      "llm_learning_rate": 3.7042242423336354e-06,
+      "loss": 0.8696,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0003704224242333635,
+      "step": 3512
+    },
+    {
+      "epoch": 0.36,
+      "gumbel_tau": 0.4367,
+      "image_masks_loss": 0.0111,
+      "llm_learning_rate": 3.703497148177403e-06,
+      "loss": 0.7884,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 0.0003703497148177403,
+      "step": 3513
+    },
+    {
+      "epoch": 0.36,
+      "gumbel_tau": 0.4366,
+      "image_masks_loss": 0.0091,
+      "llm_learning_rate": 3.7027699214893843e-06,
+      "loss": 0.8171,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 0.0003702769921489384,
+      "step": 3514
+    },
+    {
+      "epoch": 0.36,
+      "gumbel_tau": 0.4365,
+      "image_masks_loss": 0.0166,
+      "llm_learning_rate": 3.7020425623496627e-06,
+      "loss": 0.7217,
+      "output_text_mask_loss": 0.2075,
+      "predictor_learning_rate": 0.00037020425623496625,
+      "step": 3515
+    },
+    {
+      "epoch": 0.36,
+      "gumbel_tau": 0.4364,
+      "image_masks_loss": 0.0282,
+      "llm_learning_rate": 3.7013150708383367e-06,
+      "loss": 0.9118,
+      "output_text_mask_loss": 0.0694,
+      "predictor_learning_rate": 0.00037013150708383363,
+      "step": 3516
+    },
+    {
+      "epoch": 0.36,
+      "gumbel_tau": 0.4363,
+      "image_masks_loss": 0.0086,
+      "llm_learning_rate": 3.7005874470355187e-06,
+      "loss": 0.8786,
+      "output_text_mask_loss": 0.0041,
+      "predictor_learning_rate": 0.00037005874470355185,
+      "step": 3517
+    },
+    {
+      "epoch": 0.36,
+      "gumbel_tau": 0.4362,
+      "image_masks_loss": 0.0236,
+      "llm_learning_rate": 3.6998596910213375e-06,
+      "loss": 0.7746,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.00036998596910213375,
+      "step": 3518
+    },
+    {
+      "epoch": 0.36,
+      "gumbel_tau": 0.4361,
+      "image_masks_loss": 0.0183,
+      "llm_learning_rate": 3.6991318028759344e-06,
+      "loss": 0.7708,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.00036991318028759343,
+      "step": 3519
+    },
+    {
+      "epoch": 0.36,
+      "gumbel_tau": 0.436,
+      "image_masks_loss": 0.0138,
+      "llm_learning_rate": 3.6984037826794657e-06,
+      "loss": 0.7691,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.00036984037826794656,
+      "step": 3520
+    },
+    {
+      "epoch": 0.36,
+      "gumbel_tau": 0.4359,
+      "image_masks_loss": 0.0121,
+      "llm_learning_rate": 3.6976756305121035e-06,
+      "loss": 0.8158,
+      "output_text_mask_loss": 0.0309,
+      "predictor_learning_rate": 0.0003697675630512103,
+      "step": 3521
+    },
+    {
+      "epoch": 0.36,
+      "gumbel_tau": 0.4358,
+      "image_masks_loss": 0.0036,
+      "llm_learning_rate": 3.696947346454034e-06,
+      "loss": 0.8446,
+      "output_text_mask_loss": 0.1495,
+      "predictor_learning_rate": 0.0003696947346454034,
+      "step": 3522
+    },
+    {
+      "epoch": 0.36,
+      "gumbel_tau": 0.4357,
+      "image_masks_loss": 0.028,
+      "llm_learning_rate": 3.6962189305854558e-06,
+      "loss": 0.723,
+      "output_text_mask_loss": 0.0984,
+      "predictor_learning_rate": 0.00036962189305854556,
+      "step": 3523
+    },
+    {
+      "epoch": 0.36,
+      "gumbel_tau": 0.4355,
+      "image_masks_loss": 0.0265,
+      "llm_learning_rate": 3.695490382986584e-06,
+      "loss": 0.8009,
+      "output_text_mask_loss": 0.0012,
+      "predictor_learning_rate": 0.0003695490382986584,
+      "step": 3524
+    },
+    {
+      "epoch": 0.36,
+      "gumbel_tau": 0.4354,
+      "image_masks_loss": 0.0086,
+      "llm_learning_rate": 3.6947617037376495e-06,
+      "loss": 0.7448,
+      "output_text_mask_loss": 0.1076,
+      "predictor_learning_rate": 0.0003694761703737649,
+      "step": 3525
+    },
+    {
+      "epoch": 0.36,
+      "gumbel_tau": 0.4353,
+      "image_masks_loss": 0.016,
+      "llm_learning_rate": 3.694032892918894e-06,
+      "loss": 0.8113,
+      "output_text_mask_loss": 0.0969,
+      "predictor_learning_rate": 0.00036940328929188935,
+      "step": 3526
+    },
+    {
+      "epoch": 0.36,
+      "gumbel_tau": 0.4352,
+      "image_masks_loss": 0.0233,
+      "llm_learning_rate": 3.6933039506105773e-06,
+      "loss": 0.7139,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0003693303950610577,
+      "step": 3527
+    },
+    {
+      "epoch": 0.36,
+      "gumbel_tau": 0.4351,
+      "image_masks_loss": 0.0161,
+      "llm_learning_rate": 3.692574876892971e-06,
+      "loss": 0.8413,
+      "output_text_mask_loss": 0.0075,
+      "predictor_learning_rate": 0.00036925748768929705,
+      "step": 3528
+    },
+    {
+      "epoch": 0.36,
+      "gumbel_tau": 0.435,
+      "image_masks_loss": 0.0246,
+      "llm_learning_rate": 3.6918456718463635e-06,
+      "loss": 0.7698,
+      "output_text_mask_loss": 0.0004,
+      "predictor_learning_rate": 0.00036918456718463633,
+      "step": 3529
+    },
+    {
+      "epoch": 0.36,
+      "gumbel_tau": 0.4349,
+      "image_masks_loss": 0.0309,
+      "llm_learning_rate": 3.691116335551055e-06,
+      "loss": 0.7088,
+      "output_text_mask_loss": 0.0137,
+      "predictor_learning_rate": 0.00036911163355510544,
+      "step": 3530
+    },
+    {
+      "epoch": 0.36,
+      "gumbel_tau": 0.4348,
+      "image_masks_loss": 0.0114,
+      "llm_learning_rate": 3.690386868087364e-06,
+      "loss": 0.6891,
+      "output_text_mask_loss": 0.0248,
+      "predictor_learning_rate": 0.0003690386868087364,
+      "step": 3531
+    },
+    {
+      "epoch": 0.36,
+      "gumbel_tau": 0.4347,
+      "image_masks_loss": 0.0305,
+      "llm_learning_rate": 3.6896572695356186e-06,
+      "loss": 0.6969,
+      "output_text_mask_loss": 0.005,
+      "predictor_learning_rate": 0.0003689657269535618,
+      "step": 3532
+    },
+    {
+      "epoch": 0.36,
+      "gumbel_tau": 0.4346,
+      "image_masks_loss": 0.0096,
+      "llm_learning_rate": 3.688927539976165e-06,
+      "loss": 0.722,
+      "output_text_mask_loss": 0.0725,
+      "predictor_learning_rate": 0.00036889275399761646,
+      "step": 3533
+    },
+    {
+      "epoch": 0.36,
+      "gumbel_tau": 0.4345,
+      "image_masks_loss": 0.0149,
+      "llm_learning_rate": 3.6881976794893625e-06,
+      "loss": 0.838,
+      "output_text_mask_loss": 0.0193,
+      "predictor_learning_rate": 0.00036881976794893626,
+      "step": 3534
+    },
+    {
+      "epoch": 0.36,
+      "gumbel_tau": 0.4344,
+      "image_masks_loss": 0.0217,
+      "llm_learning_rate": 3.687467688155585e-06,
+      "loss": 0.7262,
+      "output_text_mask_loss": 0.0058,
+      "predictor_learning_rate": 0.00036874676881555845,
+      "step": 3535
+    },
+    {
+      "epoch": 0.36,
+      "gumbel_tau": 0.4343,
+      "image_masks_loss": 0.0296,
+      "llm_learning_rate": 3.68673756605522e-06,
+      "loss": 0.7934,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 0.00036867375660552197,
+      "step": 3536
+    },
+    {
+      "epoch": 0.36,
+      "gumbel_tau": 0.4342,
+      "image_masks_loss": 0.0298,
+      "llm_learning_rate": 3.6860073132686718e-06,
+      "loss": 0.7547,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.0003686007313268672,
+      "step": 3537
+    },
+    {
+      "epoch": 0.36,
+      "gumbel_tau": 0.4341,
+      "image_masks_loss": 0.0486,
+      "llm_learning_rate": 3.685276929876356e-06,
+      "loss": 0.7239,
+      "output_text_mask_loss": 0.0141,
+      "predictor_learning_rate": 0.0003685276929876356,
+      "step": 3538
+    },
+    {
+      "epoch": 0.36,
+      "gumbel_tau": 0.434,
+      "image_masks_loss": 0.0162,
+      "llm_learning_rate": 3.6845464159587053e-06,
+      "loss": 0.7529,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0003684546415958705,
+      "step": 3539
+    },
+    {
+      "epoch": 0.36,
+      "gumbel_tau": 0.4339,
+      "image_masks_loss": 0.0224,
+      "llm_learning_rate": 3.683815771596164e-06,
+      "loss": 0.8519,
+      "output_text_mask_loss": 0.0267,
+      "predictor_learning_rate": 0.0003683815771596164,
+      "step": 3540
+    },
+    {
+      "epoch": 0.36,
+      "gumbel_tau": 0.4338,
+      "image_masks_loss": 0.0326,
+      "llm_learning_rate": 3.683084996869194e-06,
+      "loss": 0.788,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 0.00036830849968691936,
+      "step": 3541
+    },
+    {
+      "epoch": 0.36,
+      "gumbel_tau": 0.4337,
+      "image_masks_loss": 0.0197,
+      "llm_learning_rate": 3.6823540918582677e-06,
+      "loss": 0.7969,
+      "output_text_mask_loss": 0.013,
+      "predictor_learning_rate": 0.0003682354091858268,
+      "step": 3542
+    },
+    {
+      "epoch": 0.36,
+      "gumbel_tau": 0.4336,
+      "image_masks_loss": 0.011,
+      "llm_learning_rate": 3.681623056643876e-06,
+      "loss": 0.7283,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.0003681623056643876,
+      "step": 3543
+    },
+    {
+      "epoch": 0.36,
+      "gumbel_tau": 0.4335,
+      "image_masks_loss": 0.0334,
+      "llm_learning_rate": 3.680891891306521e-06,
+      "loss": 0.7915,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 0.0003680891891306521,
+      "step": 3544
+    },
+    {
+      "epoch": 0.36,
+      "gumbel_tau": 0.4334,
+      "image_masks_loss": 0.022,
+      "llm_learning_rate": 3.6801605959267208e-06,
+      "loss": 0.6564,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00036801605959267203,
+      "step": 3545
+    },
+    {
+      "epoch": 0.36,
+      "gumbel_tau": 0.4333,
+      "image_masks_loss": 0.0198,
+      "llm_learning_rate": 3.6794291705850067e-06,
+      "loss": 0.816,
+      "output_text_mask_loss": 0.0441,
+      "predictor_learning_rate": 0.00036794291705850065,
+      "step": 3546
+    },
+    {
+      "epoch": 0.36,
+      "gumbel_tau": 0.4332,
+      "image_masks_loss": 0.0267,
+      "llm_learning_rate": 3.6786976153619257e-06,
+      "loss": 0.7741,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.00036786976153619256,
+      "step": 3547
+    },
+    {
+      "epoch": 0.36,
+      "gumbel_tau": 0.4331,
+      "image_masks_loss": 0.032,
+      "llm_learning_rate": 3.6779659303380378e-06,
+      "loss": 0.8178,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00036779659303380375,
+      "step": 3548
+    },
+    {
+      "epoch": 0.36,
+      "gumbel_tau": 0.433,
+      "image_masks_loss": 0.0152,
+      "llm_learning_rate": 3.677234115593918e-06,
+      "loss": 0.7781,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.00036772341155939174,
+      "step": 3549
+    },
+    {
+      "epoch": 0.36,
+      "gumbel_tau": 0.4329,
+      "image_masks_loss": 0.0182,
+      "llm_learning_rate": 3.676502171210155e-06,
+      "loss": 0.7195,
+      "output_text_mask_loss": 0.02,
+      "predictor_learning_rate": 0.00036765021712101547,
+      "step": 3550
+    },
+    {
+      "epoch": 0.36,
+      "gumbel_tau": 0.4328,
+      "image_masks_loss": 0.0135,
+      "llm_learning_rate": 3.6757700972673523e-06,
+      "loss": 0.813,
+      "output_text_mask_loss": 0.0454,
+      "predictor_learning_rate": 0.0003675770097267352,
+      "step": 3551
+    },
+    {
+      "epoch": 0.36,
+      "gumbel_tau": 0.4327,
+      "image_masks_loss": 0.0134,
+      "llm_learning_rate": 3.675037893846127e-06,
+      "loss": 0.8345,
+      "output_text_mask_loss": 0.0004,
+      "predictor_learning_rate": 0.0003675037893846127,
+      "step": 3552
+    },
+    {
+      "epoch": 0.36,
+      "gumbel_tau": 0.4326,
+      "image_masks_loss": 0.0212,
+      "llm_learning_rate": 3.674305561027112e-06,
+      "loss": 0.8039,
+      "output_text_mask_loss": 0.0231,
+      "predictor_learning_rate": 0.0003674305561027112,
+      "step": 3553
+    },
+    {
+      "epoch": 0.36,
+      "gumbel_tau": 0.4325,
+      "image_masks_loss": 0.0169,
+      "llm_learning_rate": 3.6735730988909536e-06,
+      "loss": 0.7646,
+      "output_text_mask_loss": 0.0641,
+      "predictor_learning_rate": 0.0003673573098890953,
+      "step": 3554
+    },
+    {
+      "epoch": 0.36,
+      "gumbel_tau": 0.4324,
+      "image_masks_loss": 0.004,
+      "llm_learning_rate": 3.672840507518311e-06,
+      "loss": 0.8061,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 0.00036728405075183107,
+      "step": 3555
+    },
+    {
+      "epoch": 0.36,
+      "gumbel_tau": 0.4323,
+      "image_masks_loss": 0.015,
+      "llm_learning_rate": 3.6721077869898593e-06,
+      "loss": 0.6927,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.0003672107786989859,
+      "step": 3556
+    },
+    {
+      "epoch": 0.36,
+      "gumbel_tau": 0.4322,
+      "image_masks_loss": 0.0052,
+      "llm_learning_rate": 3.6713749373862874e-06,
+      "loss": 0.7816,
+      "output_text_mask_loss": 0.0078,
+      "predictor_learning_rate": 0.0003671374937386287,
+      "step": 3557
+    },
+    {
+      "epoch": 0.36,
+      "gumbel_tau": 0.4321,
+      "image_masks_loss": 0.0232,
+      "llm_learning_rate": 3.670641958788298e-06,
+      "loss": 0.789,
+      "output_text_mask_loss": 0.0077,
+      "predictor_learning_rate": 0.00036706419587882977,
+      "step": 3558
+    },
+    {
+      "epoch": 0.36,
+      "gumbel_tau": 0.432,
+      "image_masks_loss": 0.015,
+      "llm_learning_rate": 3.669908851276609e-06,
+      "loss": 0.7426,
+      "output_text_mask_loss": 0.0195,
+      "predictor_learning_rate": 0.00036699088512766087,
+      "step": 3559
+    },
+    {
+      "epoch": 0.36,
+      "gumbel_tau": 0.4319,
+      "image_masks_loss": 0.0064,
+      "llm_learning_rate": 3.669175614931951e-06,
+      "loss": 0.8531,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 0.0003669175614931951,
+      "step": 3560
+    },
+    {
+      "epoch": 0.36,
+      "gumbel_tau": 0.4318,
+      "image_masks_loss": 0.0117,
+      "llm_learning_rate": 3.6684422498350705e-06,
+      "loss": 0.7627,
+      "output_text_mask_loss": 0.0231,
+      "predictor_learning_rate": 0.000366844224983507,
+      "step": 3561
+    },
+    {
+      "epoch": 0.36,
+      "gumbel_tau": 0.4317,
+      "image_masks_loss": 0.0113,
+      "llm_learning_rate": 3.6677087560667263e-06,
+      "loss": 0.8378,
+      "output_text_mask_loss": 0.0095,
+      "predictor_learning_rate": 0.0003667708756066726,
+      "step": 3562
+    },
+    {
+      "epoch": 0.37,
+      "gumbel_tau": 0.4316,
+      "image_masks_loss": 0.0152,
+      "llm_learning_rate": 3.666975133707693e-06,
+      "loss": 0.7167,
+      "output_text_mask_loss": 0.0324,
+      "predictor_learning_rate": 0.00036669751337076926,
+      "step": 3563
+    },
+    {
+      "epoch": 0.37,
+      "gumbel_tau": 0.4315,
+      "image_masks_loss": 0.0093,
+      "llm_learning_rate": 3.6662413828387593e-06,
+      "loss": 0.7151,
+      "output_text_mask_loss": 0.0477,
+      "predictor_learning_rate": 0.0003666241382838759,
+      "step": 3564
+    },
+    {
+      "epoch": 0.37,
+      "gumbel_tau": 0.4314,
+      "image_masks_loss": 0.0162,
+      "llm_learning_rate": 3.665507503540726e-06,
+      "loss": 0.7982,
+      "output_text_mask_loss": 0.0172,
+      "predictor_learning_rate": 0.0003665507503540726,
+      "step": 3565
+    },
+    {
+      "epoch": 0.37,
+      "gumbel_tau": 0.4313,
+      "image_masks_loss": 0.0099,
+      "llm_learning_rate": 3.664773495894411e-06,
+      "loss": 0.7494,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 0.0003664773495894411,
+      "step": 3566
+    },
+    {
+      "epoch": 0.37,
+      "gumbel_tau": 0.4312,
+      "image_masks_loss": 0.0165,
+      "llm_learning_rate": 3.664039359980644e-06,
+      "loss": 0.726,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 0.0003664039359980644,
+      "step": 3567
+    },
+    {
+      "epoch": 0.37,
+      "gumbel_tau": 0.4311,
+      "image_masks_loss": 0.0204,
+      "llm_learning_rate": 3.6633050958802698e-06,
+      "loss": 0.6774,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00036633050958802694,
+      "step": 3568
+    },
+    {
+      "epoch": 0.37,
+      "gumbel_tau": 0.4309,
+      "image_masks_loss": 0.0122,
+      "llm_learning_rate": 3.6625707036741476e-06,
+      "loss": 0.7485,
+      "output_text_mask_loss": 0.0235,
+      "predictor_learning_rate": 0.00036625707036741474,
+      "step": 3569
+    },
+    {
+      "epoch": 0.37,
+      "gumbel_tau": 0.4308,
+      "image_masks_loss": 0.0237,
+      "llm_learning_rate": 3.6618361834431493e-06,
+      "loss": 0.7625,
+      "output_text_mask_loss": 0.0292,
+      "predictor_learning_rate": 0.0003661836183443149,
+      "step": 3570
+    },
+    {
+      "epoch": 0.37,
+      "gumbel_tau": 0.4307,
+      "image_masks_loss": 0.0204,
+      "llm_learning_rate": 3.661101535268163e-06,
+      "loss": 0.7212,
+      "output_text_mask_loss": 0.0065,
+      "predictor_learning_rate": 0.0003661101535268163,
+      "step": 3571
+    },
+    {
+      "epoch": 0.37,
+      "gumbel_tau": 0.4306,
+      "image_masks_loss": 0.0052,
+      "llm_learning_rate": 3.660366759230089e-06,
+      "loss": 0.6156,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0003660366759230089,
+      "step": 3572
+    },
+    {
+      "epoch": 0.37,
+      "gumbel_tau": 0.4305,
+      "image_masks_loss": 0.0147,
+      "llm_learning_rate": 3.6596318554098438e-06,
+      "loss": 0.7531,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00036596318554098434,
+      "step": 3573
+    },
+    {
+      "epoch": 0.37,
+      "gumbel_tau": 0.4304,
+      "image_masks_loss": 0.0118,
+      "llm_learning_rate": 3.658896823888355e-06,
+      "loss": 0.8038,
+      "output_text_mask_loss": 0.0023,
+      "predictor_learning_rate": 0.0003658896823888355,
+      "step": 3574
+    },
+    {
+      "epoch": 0.37,
+      "gumbel_tau": 0.4303,
+      "image_masks_loss": 0.0334,
+      "llm_learning_rate": 3.658161664746566e-06,
+      "loss": 0.8304,
+      "output_text_mask_loss": 0.0277,
+      "predictor_learning_rate": 0.0003658161664746566,
+      "step": 3575
+    },
+    {
+      "epoch": 0.37,
+      "gumbel_tau": 0.4302,
+      "image_masks_loss": 0.028,
+      "llm_learning_rate": 3.6574263780654355e-06,
+      "loss": 0.7983,
+      "output_text_mask_loss": 0.0523,
+      "predictor_learning_rate": 0.00036574263780654354,
+      "step": 3576
+    },
+    {
+      "epoch": 0.37,
+      "gumbel_tau": 0.4301,
+      "image_masks_loss": 0.0169,
+      "llm_learning_rate": 3.6566909639259353e-06,
+      "loss": 0.7824,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0003656690963925935,
+      "step": 3577
+    },
+    {
+      "epoch": 0.37,
+      "gumbel_tau": 0.43,
+      "image_masks_loss": 0.0298,
+      "llm_learning_rate": 3.6559554224090476e-06,
+      "loss": 0.7863,
+      "output_text_mask_loss": 0.0004,
+      "predictor_learning_rate": 0.00036559554224090476,
+      "step": 3578
+    },
+    {
+      "epoch": 0.37,
+      "gumbel_tau": 0.4299,
+      "image_masks_loss": 0.0125,
+      "llm_learning_rate": 3.6552197535957756e-06,
+      "loss": 0.8078,
+      "output_text_mask_loss": 0.0175,
+      "predictor_learning_rate": 0.00036552197535957755,
+      "step": 3579
+    },
+    {
+      "epoch": 0.37,
+      "gumbel_tau": 0.4298,
+      "image_masks_loss": 0.0246,
+      "llm_learning_rate": 3.6544839575671305e-06,
+      "loss": 0.7677,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.00036544839575671303,
+      "step": 3580
+    },
+    {
+      "epoch": 0.37,
+      "gumbel_tau": 0.4297,
+      "image_masks_loss": 0.0125,
+      "llm_learning_rate": 3.653748034404141e-06,
+      "loss": 0.7014,
+      "output_text_mask_loss": 0.0488,
+      "predictor_learning_rate": 0.0003653748034404141,
+      "step": 3581
+    },
+    {
+      "epoch": 0.37,
+      "gumbel_tau": 0.4296,
+      "image_masks_loss": 0.0051,
+      "llm_learning_rate": 3.6530119841878468e-06,
+      "loss": 0.7368,
+      "output_text_mask_loss": 0.0086,
+      "predictor_learning_rate": 0.00036530119841878466,
+      "step": 3582
+    },
+    {
+      "epoch": 0.37,
+      "gumbel_tau": 0.4295,
+      "image_masks_loss": 0.0145,
+      "llm_learning_rate": 3.652275806999306e-06,
+      "loss": 0.7663,
+      "output_text_mask_loss": 0.0032,
+      "predictor_learning_rate": 0.00036522758069993056,
+      "step": 3583
+    },
+    {
+      "epoch": 0.37,
+      "gumbel_tau": 0.4294,
+      "image_masks_loss": 0.01,
+      "llm_learning_rate": 3.651539502919586e-06,
+      "loss": 0.8425,
+      "output_text_mask_loss": 0.0141,
+      "predictor_learning_rate": 0.0003651539502919586,
+      "step": 3584
+    },
+    {
+      "epoch": 0.37,
+      "gumbel_tau": 0.4293,
+      "image_masks_loss": 0.0232,
+      "llm_learning_rate": 3.650803072029771e-06,
+      "loss": 0.7181,
+      "output_text_mask_loss": 0.0557,
+      "predictor_learning_rate": 0.0003650803072029771,
+      "step": 3585
+    },
+    {
+      "epoch": 0.37,
+      "gumbel_tau": 0.4292,
+      "image_masks_loss": 0.0077,
+      "llm_learning_rate": 3.6500665144109587e-06,
+      "loss": 0.542,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00036500665144109586,
+      "step": 3586
+    },
+    {
+      "epoch": 0.37,
+      "gumbel_tau": 0.4291,
+      "image_masks_loss": 0.0109,
+      "llm_learning_rate": 3.6493298301442602e-06,
+      "loss": 0.7258,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.00036493298301442604,
+      "step": 3587
+    },
+    {
+      "epoch": 0.37,
+      "gumbel_tau": 0.429,
+      "image_masks_loss": 0.0178,
+      "llm_learning_rate": 3.648593019310801e-06,
+      "loss": 0.7035,
+      "output_text_mask_loss": 0.0534,
+      "predictor_learning_rate": 0.00036485930193108006,
+      "step": 3588
+    },
+    {
+      "epoch": 0.37,
+      "gumbel_tau": 0.4289,
+      "image_masks_loss": 0.011,
+      "llm_learning_rate": 3.64785608199172e-06,
+      "loss": 0.7848,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.00036478560819917197,
+      "step": 3589
+    },
+    {
+      "epoch": 0.37,
+      "gumbel_tau": 0.4288,
+      "image_masks_loss": 0.0196,
+      "llm_learning_rate": 3.6471190182681697e-06,
+      "loss": 0.7605,
+      "output_text_mask_loss": 0.0153,
+      "predictor_learning_rate": 0.000364711901826817,
+      "step": 3590
+    },
+    {
+      "epoch": 0.37,
+      "gumbel_tau": 0.4287,
+      "image_masks_loss": 0.0311,
+      "llm_learning_rate": 3.646381828221318e-06,
+      "loss": 0.7462,
+      "output_text_mask_loss": 0.0004,
+      "predictor_learning_rate": 0.0003646381828221318,
+      "step": 3591
+    },
+    {
+      "epoch": 0.37,
+      "gumbel_tau": 0.4286,
+      "image_masks_loss": 0.0319,
+      "llm_learning_rate": 3.645644511932347e-06,
+      "loss": 0.7431,
+      "output_text_mask_loss": 0.1129,
+      "predictor_learning_rate": 0.0003645644511932347,
+      "step": 3592
+    },
+    {
+      "epoch": 0.37,
+      "gumbel_tau": 0.4285,
+      "image_masks_loss": 0.0093,
+      "llm_learning_rate": 3.6449070694824502e-06,
+      "loss": 0.7446,
+      "output_text_mask_loss": 0.0023,
+      "predictor_learning_rate": 0.000364490706948245,
+      "step": 3593
+    },
+    {
+      "epoch": 0.37,
+      "gumbel_tau": 0.4284,
+      "image_masks_loss": 0.0183,
+      "llm_learning_rate": 3.6441695009528368e-06,
+      "loss": 0.7969,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.00036441695009528365,
+      "step": 3594
+    },
+    {
+      "epoch": 0.37,
+      "gumbel_tau": 0.4283,
+      "image_masks_loss": 0.0155,
+      "llm_learning_rate": 3.6434318064247287e-06,
+      "loss": 0.7665,
+      "output_text_mask_loss": 0.0132,
+      "predictor_learning_rate": 0.00036434318064247287,
+      "step": 3595
+    },
+    {
+      "epoch": 0.37,
+      "gumbel_tau": 0.4282,
+      "image_masks_loss": 0.0381,
+      "llm_learning_rate": 3.642693985979364e-06,
+      "loss": 0.7401,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 0.0003642693985979364,
+      "step": 3596
+    },
+    {
+      "epoch": 0.37,
+      "gumbel_tau": 0.4281,
+      "image_masks_loss": 0.007,
+      "llm_learning_rate": 3.6419560396979923e-06,
+      "loss": 0.8152,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.00036419560396979924,
+      "step": 3597
+    },
+    {
+      "epoch": 0.37,
+      "gumbel_tau": 0.428,
+      "image_masks_loss": 0.0231,
+      "llm_learning_rate": 3.641217967661877e-06,
+      "loss": 0.7543,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.00036412179676618766,
+      "step": 3598
+    },
+    {
+      "epoch": 0.37,
+      "gumbel_tau": 0.4279,
+      "image_masks_loss": 0.0062,
+      "llm_learning_rate": 3.6404797699522974e-06,
+      "loss": 0.7493,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 0.0003640479769952297,
+      "step": 3599
+    },
+    {
+      "epoch": 0.37,
+      "gumbel_tau": 0.4278,
+      "image_masks_loss": 0.0207,
+      "llm_learning_rate": 3.639741446650546e-06,
+      "loss": 0.778,
+      "output_text_mask_loss": 0.0071,
+      "predictor_learning_rate": 0.00036397414466505455,
+      "step": 3600
+    },
+    {
+      "epoch": 0.37,
+      "gumbel_tau": 0.4277,
+      "image_masks_loss": 0.0167,
+      "llm_learning_rate": 3.639002997837927e-06,
+      "loss": 0.8163,
+      "output_text_mask_loss": 0.0458,
+      "predictor_learning_rate": 0.0003639002997837927,
+      "step": 3601
+    },
+    {
+      "epoch": 0.37,
+      "gumbel_tau": 0.4276,
+      "image_masks_loss": 0.0134,
+      "llm_learning_rate": 3.638264423595761e-06,
+      "loss": 0.7428,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.0003638264423595761,
+      "step": 3602
+    },
+    {
+      "epoch": 0.37,
+      "gumbel_tau": 0.4275,
+      "image_masks_loss": 0.0137,
+      "llm_learning_rate": 3.6375257240053817e-06,
+      "loss": 0.8582,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.00036375257240053816,
+      "step": 3603
+    },
+    {
+      "epoch": 0.37,
+      "gumbel_tau": 0.4274,
+      "image_masks_loss": 0.0154,
+      "llm_learning_rate": 3.636786899148136e-06,
+      "loss": 0.6419,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0003636786899148136,
+      "step": 3604
+    },
+    {
+      "epoch": 0.37,
+      "gumbel_tau": 0.4273,
+      "image_masks_loss": 0.0111,
+      "llm_learning_rate": 3.6360479491053846e-06,
+      "loss": 0.607,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00036360479491053845,
+      "step": 3605
+    },
+    {
+      "epoch": 0.37,
+      "gumbel_tau": 0.4272,
+      "image_masks_loss": 0.0083,
+      "llm_learning_rate": 3.635308873958503e-06,
+      "loss": 0.7534,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.00036353088739585025,
+      "step": 3606
+    },
+    {
+      "epoch": 0.37,
+      "gumbel_tau": 0.4271,
+      "image_masks_loss": 0.0273,
+      "llm_learning_rate": 3.634569673788878e-06,
+      "loss": 0.7915,
+      "output_text_mask_loss": 0.0261,
+      "predictor_learning_rate": 0.0003634569673788878,
+      "step": 3607
+    },
+    {
+      "epoch": 0.37,
+      "gumbel_tau": 0.427,
+      "image_masks_loss": 0.0402,
+      "llm_learning_rate": 3.6338303486779146e-06,
+      "loss": 0.7486,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00036338303486779143,
+      "step": 3608
+    },
+    {
+      "epoch": 0.37,
+      "gumbel_tau": 0.4269,
+      "image_masks_loss": 0.0172,
+      "llm_learning_rate": 3.633090898707028e-06,
+      "loss": 0.7275,
+      "output_text_mask_loss": 0.0107,
+      "predictor_learning_rate": 0.00036330908987070276,
+      "step": 3609
+    },
+    {
+      "epoch": 0.37,
+      "gumbel_tau": 0.4268,
+      "image_masks_loss": 0.0137,
+      "llm_learning_rate": 3.632351323957647e-06,
+      "loss": 0.6722,
+      "output_text_mask_loss": 0.0416,
+      "predictor_learning_rate": 0.00036323513239576465,
+      "step": 3610
+    },
+    {
+      "epoch": 0.37,
+      "gumbel_tau": 0.4267,
+      "image_masks_loss": 0.0334,
+      "llm_learning_rate": 3.631611624511215e-06,
+      "loss": 0.7263,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.00036316116245112154,
+      "step": 3611
+    },
+    {
+      "epoch": 0.37,
+      "gumbel_tau": 0.4266,
+      "image_masks_loss": 0.0092,
+      "llm_learning_rate": 3.630871800449192e-06,
+      "loss": 0.7842,
+      "output_text_mask_loss": 0.0931,
+      "predictor_learning_rate": 0.0003630871800449192,
+      "step": 3612
+    },
+    {
+      "epoch": 0.37,
+      "gumbel_tau": 0.4265,
+      "image_masks_loss": 0.0026,
+      "llm_learning_rate": 3.6301318518530475e-06,
+      "loss": 0.8012,
+      "output_text_mask_loss": 0.0679,
+      "predictor_learning_rate": 0.00036301318518530473,
+      "step": 3613
+    },
+    {
+      "epoch": 0.37,
+      "gumbel_tau": 0.4264,
+      "image_masks_loss": 0.0111,
+      "llm_learning_rate": 3.6293917788042654e-06,
+      "loss": 0.7594,
+      "output_text_mask_loss": 0.0039,
+      "predictor_learning_rate": 0.00036293917788042654,
+      "step": 3614
+    },
+    {
+      "epoch": 0.37,
+      "gumbel_tau": 0.4263,
+      "image_masks_loss": 0.0182,
+      "llm_learning_rate": 3.628651581384346e-06,
+      "loss": 0.8917,
+      "output_text_mask_loss": 0.0217,
+      "predictor_learning_rate": 0.00036286515813843456,
+      "step": 3615
+    },
+    {
+      "epoch": 0.37,
+      "gumbel_tau": 0.4262,
+      "image_masks_loss": 0.0105,
+      "llm_learning_rate": 3.6279112596747997e-06,
+      "loss": 0.7585,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.00036279112596747994,
+      "step": 3616
+    },
+    {
+      "epoch": 0.37,
+      "gumbel_tau": 0.4261,
+      "image_masks_loss": 0.0087,
+      "llm_learning_rate": 3.6271708137571527e-06,
+      "loss": 0.8208,
+      "output_text_mask_loss": 0.0366,
+      "predictor_learning_rate": 0.0003627170813757152,
+      "step": 3617
+    },
+    {
+      "epoch": 0.37,
+      "gumbel_tau": 0.426,
+      "image_masks_loss": 0.0121,
+      "llm_learning_rate": 3.6264302437129458e-06,
+      "loss": 0.701,
+      "output_text_mask_loss": 0.0189,
+      "predictor_learning_rate": 0.00036264302437129454,
+      "step": 3618
+    },
+    {
+      "epoch": 0.37,
+      "gumbel_tau": 0.4259,
+      "image_masks_loss": 0.0103,
+      "llm_learning_rate": 3.625689549623731e-06,
+      "loss": 0.6756,
+      "output_text_mask_loss": 0.0127,
+      "predictor_learning_rate": 0.0003625689549623731,
+      "step": 3619
+    },
+    {
+      "epoch": 0.37,
+      "gumbel_tau": 0.4258,
+      "image_masks_loss": 0.0128,
+      "llm_learning_rate": 3.6249487315710756e-06,
+      "loss": 0.7005,
+      "output_text_mask_loss": 0.0126,
+      "predictor_learning_rate": 0.00036249487315710753,
+      "step": 3620
+    },
+    {
+      "epoch": 0.37,
+      "gumbel_tau": 0.4257,
+      "image_masks_loss": 0.0229,
+      "llm_learning_rate": 3.6242077896365602e-06,
+      "loss": 0.7407,
+      "output_text_mask_loss": 0.0393,
+      "predictor_learning_rate": 0.000362420778963656,
+      "step": 3621
+    },
+    {
+      "epoch": 0.37,
+      "gumbel_tau": 0.4256,
+      "image_masks_loss": 0.0134,
+      "llm_learning_rate": 3.623466723901779e-06,
+      "loss": 0.6424,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.00036234667239017787,
+      "step": 3622
+    },
+    {
+      "epoch": 0.37,
+      "gumbel_tau": 0.4255,
+      "image_masks_loss": 0.0153,
+      "llm_learning_rate": 3.62272553444834e-06,
+      "loss": 0.757,
+      "output_text_mask_loss": 0.0004,
+      "predictor_learning_rate": 0.00036227255344483403,
+      "step": 3623
+    },
+    {
+      "epoch": 0.37,
+      "gumbel_tau": 0.4254,
+      "image_masks_loss": 0.0243,
+      "llm_learning_rate": 3.6219842213578643e-06,
+      "loss": 0.7458,
+      "output_text_mask_loss": 0.017,
+      "predictor_learning_rate": 0.0003621984221357864,
+      "step": 3624
+    },
+    {
+      "epoch": 0.37,
+      "gumbel_tau": 0.4253,
+      "image_masks_loss": 0.0102,
+      "llm_learning_rate": 3.621242784711987e-06,
+      "loss": 0.7704,
+      "output_text_mask_loss": 0.0014,
+      "predictor_learning_rate": 0.00036212427847119865,
+      "step": 3625
+    },
+    {
+      "epoch": 0.37,
+      "gumbel_tau": 0.4252,
+      "image_masks_loss": 0.0174,
+      "llm_learning_rate": 3.6205012245923565e-06,
+      "loss": 0.7322,
+      "output_text_mask_loss": 0.0349,
+      "predictor_learning_rate": 0.0003620501224592356,
+      "step": 3626
+    },
+    {
+      "epoch": 0.37,
+      "gumbel_tau": 0.4251,
+      "image_masks_loss": 0.0326,
+      "llm_learning_rate": 3.619759541080635e-06,
+      "loss": 0.7625,
+      "output_text_mask_loss": 0.0599,
+      "predictor_learning_rate": 0.0003619759541080635,
+      "step": 3627
+    },
+    {
+      "epoch": 0.37,
+      "gumbel_tau": 0.425,
+      "image_masks_loss": 0.0259,
+      "llm_learning_rate": 3.6190177342585e-06,
+      "loss": 0.7384,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.00036190177342585,
+      "step": 3628
+    },
+    {
+      "epoch": 0.37,
+      "gumbel_tau": 0.4249,
+      "image_masks_loss": 0.033,
+      "llm_learning_rate": 3.6182758042076383e-06,
+      "loss": 0.7286,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.0003618275804207638,
+      "step": 3629
+    },
+    {
+      "epoch": 0.37,
+      "gumbel_tau": 0.4248,
+      "image_masks_loss": 0.0137,
+      "llm_learning_rate": 3.6175337510097542e-06,
+      "loss": 0.7492,
+      "output_text_mask_loss": 0.02,
+      "predictor_learning_rate": 0.0003617533751009754,
+      "step": 3630
+    },
+    {
+      "epoch": 0.37,
+      "gumbel_tau": 0.4247,
+      "image_masks_loss": 0.0713,
+      "llm_learning_rate": 3.6167915747465646e-06,
+      "loss": 0.781,
+      "output_text_mask_loss": 0.0071,
+      "predictor_learning_rate": 0.0003616791574746565,
+      "step": 3631
+    },
+    {
+      "epoch": 0.37,
+      "gumbel_tau": 0.4246,
+      "image_masks_loss": 0.0248,
+      "llm_learning_rate": 3.6160492754998e-06,
+      "loss": 0.8246,
+      "output_text_mask_loss": 0.0012,
+      "predictor_learning_rate": 0.00036160492754998,
+      "step": 3632
+    },
+    {
+      "epoch": 0.37,
+      "gumbel_tau": 0.4245,
+      "image_masks_loss": 0.0205,
+      "llm_learning_rate": 3.6153068533512024e-06,
+      "loss": 0.8042,
+      "output_text_mask_loss": 0.0099,
+      "predictor_learning_rate": 0.00036153068533512024,
+      "step": 3633
+    },
+    {
+      "epoch": 0.37,
+      "gumbel_tau": 0.4244,
+      "image_masks_loss": 0.0507,
+      "llm_learning_rate": 3.61456430838253e-06,
+      "loss": 0.7556,
+      "output_text_mask_loss": 0.0219,
+      "predictor_learning_rate": 0.00036145643083825295,
+      "step": 3634
+    },
+    {
+      "epoch": 0.37,
+      "gumbel_tau": 0.4243,
+      "image_masks_loss": 0.0133,
+      "llm_learning_rate": 3.613821640675554e-06,
+      "loss": 0.6943,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0003613821640675554,
+      "step": 3635
+    },
+    {
+      "epoch": 0.37,
+      "gumbel_tau": 0.4242,
+      "image_masks_loss": 0.0088,
+      "llm_learning_rate": 3.6130788503120574e-06,
+      "loss": 0.8245,
+      "output_text_mask_loss": 0.0225,
+      "predictor_learning_rate": 0.0003613078850312057,
+      "step": 3636
+    },
+    {
+      "epoch": 0.37,
+      "gumbel_tau": 0.4241,
+      "image_masks_loss": 0.0056,
+      "llm_learning_rate": 3.6123359373738375e-06,
+      "loss": 0.7114,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.0003612335937373837,
+      "step": 3637
+    },
+    {
+      "epoch": 0.37,
+      "gumbel_tau": 0.424,
+      "image_masks_loss": 0.0176,
+      "llm_learning_rate": 3.611592901942708e-06,
+      "loss": 0.7782,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0003611592901942708,
+      "step": 3638
+    },
+    {
+      "epoch": 0.37,
+      "gumbel_tau": 0.4239,
+      "image_masks_loss": 0.0218,
+      "llm_learning_rate": 3.610849744100492e-06,
+      "loss": 0.8309,
+      "output_text_mask_loss": 0.0004,
+      "predictor_learning_rate": 0.0003610849744100492,
+      "step": 3639
+    },
+    {
+      "epoch": 0.37,
+      "gumbel_tau": 0.4238,
+      "image_masks_loss": 0.0164,
+      "llm_learning_rate": 3.6101064639290272e-06,
+      "loss": 0.8635,
+      "output_text_mask_loss": 0.0736,
+      "predictor_learning_rate": 0.00036101064639290273,
+      "step": 3640
+    },
+    {
+      "epoch": 0.37,
+      "gumbel_tau": 0.4237,
+      "image_masks_loss": 0.0118,
+      "llm_learning_rate": 3.609363061510166e-06,
+      "loss": 0.7996,
+      "output_text_mask_loss": 0.0465,
+      "predictor_learning_rate": 0.00036093630615101655,
+      "step": 3641
+    },
+    {
+      "epoch": 0.37,
+      "gumbel_tau": 0.4236,
+      "image_masks_loss": 0.019,
+      "llm_learning_rate": 3.608619536925775e-06,
+      "loss": 0.7808,
+      "output_text_mask_loss": 0.0023,
+      "predictor_learning_rate": 0.00036086195369257747,
+      "step": 3642
+    },
+    {
+      "epoch": 0.37,
+      "gumbel_tau": 0.4235,
+      "image_masks_loss": 0.0184,
+      "llm_learning_rate": 3.6078758902577295e-06,
+      "loss": 0.7522,
+      "output_text_mask_loss": 0.0301,
+      "predictor_learning_rate": 0.0003607875890257729,
+      "step": 3643
+    },
+    {
+      "epoch": 0.37,
+      "gumbel_tau": 0.4234,
+      "image_masks_loss": 0.0038,
+      "llm_learning_rate": 3.6071321215879236e-06,
+      "loss": 0.7953,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.00036071321215879235,
+      "step": 3644
+    },
+    {
+      "epoch": 0.37,
+      "gumbel_tau": 0.4233,
+      "image_masks_loss": 0.0101,
+      "llm_learning_rate": 3.606388230998263e-06,
+      "loss": 0.7544,
+      "output_text_mask_loss": 0.0341,
+      "predictor_learning_rate": 0.00036063882309982624,
+      "step": 3645
+    },
+    {
+      "epoch": 0.37,
+      "gumbel_tau": 0.4232,
+      "image_masks_loss": 0.0084,
+      "llm_learning_rate": 3.605644218570665e-06,
+      "loss": 0.6976,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0003605644218570665,
+      "step": 3646
+    },
+    {
+      "epoch": 0.37,
+      "gumbel_tau": 0.4231,
+      "image_masks_loss": 0.0305,
+      "llm_learning_rate": 3.6049000843870634e-06,
+      "loss": 0.7625,
+      "output_text_mask_loss": 0.0233,
+      "predictor_learning_rate": 0.0003604900084387063,
+      "step": 3647
+    },
+    {
+      "epoch": 0.37,
+      "gumbel_tau": 0.423,
+      "image_masks_loss": 0.0234,
+      "llm_learning_rate": 3.6041558285294036e-06,
+      "loss": 0.8257,
+      "output_text_mask_loss": 0.0374,
+      "predictor_learning_rate": 0.00036041558285294033,
+      "step": 3648
+    },
+    {
+      "epoch": 0.37,
+      "gumbel_tau": 0.4229,
+      "image_masks_loss": 0.0237,
+      "llm_learning_rate": 3.6034114510796446e-06,
+      "loss": 0.7181,
+      "output_text_mask_loss": 0.0137,
+      "predictor_learning_rate": 0.00036034114510796446,
+      "step": 3649
+    },
+    {
+      "epoch": 0.37,
+      "gumbel_tau": 0.4228,
+      "image_masks_loss": 0.0077,
+      "llm_learning_rate": 3.602666952119759e-06,
+      "loss": 0.8297,
+      "output_text_mask_loss": 0.0134,
+      "predictor_learning_rate": 0.00036026669521197586,
+      "step": 3650
+    },
+    {
+      "epoch": 0.37,
+      "gumbel_tau": 0.4227,
+      "image_masks_loss": 0.0066,
+      "llm_learning_rate": 3.6019223317317327e-06,
+      "loss": 0.7077,
+      "output_text_mask_loss": 0.0012,
+      "predictor_learning_rate": 0.0003601922331731732,
+      "step": 3651
+    },
+    {
+      "epoch": 0.37,
+      "gumbel_tau": 0.4226,
+      "image_masks_loss": 0.0139,
+      "llm_learning_rate": 3.601177589997564e-06,
+      "loss": 0.8359,
+      "output_text_mask_loss": 0.0645,
+      "predictor_learning_rate": 0.0003601177589997564,
+      "step": 3652
+    },
+    {
+      "epoch": 0.37,
+      "gumbel_tau": 0.4225,
+      "image_masks_loss": 0.0195,
+      "llm_learning_rate": 3.6004327269992666e-06,
+      "loss": 0.6916,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.00036004327269992665,
+      "step": 3653
+    },
+    {
+      "epoch": 0.37,
+      "gumbel_tau": 0.4224,
+      "image_masks_loss": 0.0207,
+      "llm_learning_rate": 3.5996877428188664e-06,
+      "loss": 0.7465,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.0003599687742818866,
+      "step": 3654
+    },
+    {
+      "epoch": 0.37,
+      "gumbel_tau": 0.4223,
+      "image_masks_loss": 0.0072,
+      "llm_learning_rate": 3.598942637538403e-06,
+      "loss": 0.7754,
+      "output_text_mask_loss": 0.0357,
+      "predictor_learning_rate": 0.00035989426375384025,
+      "step": 3655
+    },
+    {
+      "epoch": 0.37,
+      "gumbel_tau": 0.4222,
+      "image_masks_loss": 0.0116,
+      "llm_learning_rate": 3.5981974112399277e-06,
+      "loss": 0.7909,
+      "output_text_mask_loss": 0.0319,
+      "predictor_learning_rate": 0.00035981974112399273,
+      "step": 3656
+    },
+    {
+      "epoch": 0.37,
+      "gumbel_tau": 0.4221,
+      "image_masks_loss": 0.0267,
+      "llm_learning_rate": 3.597452064005507e-06,
+      "loss": 0.7766,
+      "output_text_mask_loss": 0.0025,
+      "predictor_learning_rate": 0.0003597452064005507,
+      "step": 3657
+    },
+    {
+      "epoch": 0.37,
+      "gumbel_tau": 0.422,
+      "image_masks_loss": 0.0128,
+      "llm_learning_rate": 3.5967065959172216e-06,
+      "loss": 0.6028,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00035967065959172215,
+      "step": 3658
+    },
+    {
+      "epoch": 0.37,
+      "gumbel_tau": 0.4219,
+      "image_masks_loss": 0.0213,
+      "llm_learning_rate": 3.5959610070571625e-06,
+      "loss": 0.7929,
+      "output_text_mask_loss": 0.0744,
+      "predictor_learning_rate": 0.00035959610070571623,
+      "step": 3659
+    },
+    {
+      "epoch": 0.38,
+      "gumbel_tau": 0.4218,
+      "image_masks_loss": 0.0148,
+      "llm_learning_rate": 3.595215297507436e-06,
+      "loss": 0.7804,
+      "output_text_mask_loss": 0.0261,
+      "predictor_learning_rate": 0.0003595215297507436,
+      "step": 3660
+    },
+    {
+      "epoch": 0.38,
+      "gumbel_tau": 0.4217,
+      "image_masks_loss": 0.0149,
+      "llm_learning_rate": 3.5944694673501625e-06,
+      "loss": 0.6381,
+      "output_text_mask_loss": 0.006,
+      "predictor_learning_rate": 0.0003594469467350162,
+      "step": 3661
+    },
+    {
+      "epoch": 0.38,
+      "gumbel_tau": 0.4216,
+      "image_masks_loss": 0.0269,
+      "llm_learning_rate": 3.5937235166674723e-06,
+      "loss": 0.7431,
+      "output_text_mask_loss": 0.0038,
+      "predictor_learning_rate": 0.0003593723516667472,
+      "step": 3662
+    },
+    {
+      "epoch": 0.38,
+      "gumbel_tau": 0.4215,
+      "image_masks_loss": 0.0146,
+      "llm_learning_rate": 3.5929774455415124e-06,
+      "loss": 0.7804,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0003592977445541512,
+      "step": 3663
+    },
+    {
+      "epoch": 0.38,
+      "gumbel_tau": 0.4214,
+      "image_masks_loss": 0.0207,
+      "llm_learning_rate": 3.592231254054442e-06,
+      "loss": 0.7842,
+      "output_text_mask_loss": 0.0687,
+      "predictor_learning_rate": 0.00035922312540544415,
+      "step": 3664
+    },
+    {
+      "epoch": 0.38,
+      "gumbel_tau": 0.4213,
+      "image_masks_loss": 0.0257,
+      "llm_learning_rate": 3.591484942288433e-06,
+      "loss": 0.6504,
+      "output_text_mask_loss": NaN,
+      "predictor_learning_rate": 0.0003591484942288433,
+      "step": 3665
+    },
+    {
+      "epoch": 0.38,
+      "gumbel_tau": 0.4212,
+      "image_masks_loss": 0.0227,
+      "llm_learning_rate": 3.5907385103256705e-06,
+      "loss": 0.7534,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.000359073851032567,
+      "step": 3666
+    },
+    {
+      "epoch": 0.38,
+      "gumbel_tau": 0.4211,
+      "image_masks_loss": 0.0176,
+      "llm_learning_rate": 3.589991958248354e-06,
+      "loss": 0.8607,
+      "output_text_mask_loss": 0.1167,
+      "predictor_learning_rate": 0.0003589991958248354,
+      "step": 3667
+    },
+    {
+      "epoch": 0.38,
+      "gumbel_tau": 0.421,
+      "image_masks_loss": 0.0113,
+      "llm_learning_rate": 3.589245286138695e-06,
+      "loss": 0.6879,
+      "output_text_mask_loss": 0.0073,
+      "predictor_learning_rate": 0.00035892452861386947,
+      "step": 3668
+    },
+    {
+      "epoch": 0.38,
+      "gumbel_tau": 0.4209,
+      "image_masks_loss": 0.0037,
+      "llm_learning_rate": 3.5884984940789196e-06,
+      "loss": 0.7627,
+      "output_text_mask_loss": 0.0256,
+      "predictor_learning_rate": 0.00035884984940789196,
+      "step": 3669
+    },
+    {
+      "epoch": 0.38,
+      "gumbel_tau": 0.4208,
+      "image_masks_loss": 0.0062,
+      "llm_learning_rate": 3.5877515821512654e-06,
+      "loss": 0.7371,
+      "output_text_mask_loss": 0.0174,
+      "predictor_learning_rate": 0.0003587751582151265,
+      "step": 3670
+    },
+    {
+      "epoch": 0.38,
+      "gumbel_tau": 0.4207,
+      "image_masks_loss": 0.021,
+      "llm_learning_rate": 3.5870045504379837e-06,
+      "loss": 0.7896,
+      "output_text_mask_loss": 0.0103,
+      "predictor_learning_rate": 0.0003587004550437983,
+      "step": 3671
+    },
+    {
+      "epoch": 0.38,
+      "gumbel_tau": 0.4206,
+      "image_masks_loss": 0.0376,
+      "llm_learning_rate": 3.5862573990213393e-06,
+      "loss": 0.7874,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.0003586257399021339,
+      "step": 3672
+    },
+    {
+      "epoch": 0.38,
+      "gumbel_tau": 0.4205,
+      "image_masks_loss": 0.0139,
+      "llm_learning_rate": 3.5855101279836113e-06,
+      "loss": 0.7652,
+      "output_text_mask_loss": 0.0032,
+      "predictor_learning_rate": 0.0003585510127983611,
+      "step": 3673
+    },
+    {
+      "epoch": 0.38,
+      "gumbel_tau": 0.4204,
+      "image_masks_loss": 0.0236,
+      "llm_learning_rate": 3.5847627374070898e-06,
+      "loss": 0.8212,
+      "output_text_mask_loss": 0.0168,
+      "predictor_learning_rate": 0.00035847627374070896,
+      "step": 3674
+    },
+    {
+      "epoch": 0.38,
+      "gumbel_tau": 0.4203,
+      "image_masks_loss": 0.0173,
+      "llm_learning_rate": 3.5840152273740796e-06,
+      "loss": 0.7548,
+      "output_text_mask_loss": 0.0011,
+      "predictor_learning_rate": 0.00035840152273740794,
+      "step": 3675
+    },
+    {
+      "epoch": 0.38,
+      "gumbel_tau": 0.4202,
+      "image_masks_loss": 0.0152,
+      "llm_learning_rate": 3.5832675979668976e-06,
+      "loss": 0.8701,
+      "output_text_mask_loss": 0.0469,
+      "predictor_learning_rate": 0.0003583267597966897,
+      "step": 3676
+    },
+    {
+      "epoch": 0.38,
+      "gumbel_tau": 0.4201,
+      "image_masks_loss": 0.009,
+      "llm_learning_rate": 3.5825198492678753e-06,
+      "loss": 0.7861,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.0003582519849267875,
+      "step": 3677
+    },
+    {
+      "epoch": 0.38,
+      "gumbel_tau": 0.42,
+      "image_masks_loss": 0.017,
+      "llm_learning_rate": 3.581771981359356e-06,
+      "loss": 0.7601,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 0.00035817719813593555,
+      "step": 3678
+    },
+    {
+      "epoch": 0.38,
+      "gumbel_tau": 0.4199,
+      "image_masks_loss": 0.0027,
+      "llm_learning_rate": 3.5810239943236963e-06,
+      "loss": 0.7576,
+      "output_text_mask_loss": 0.0141,
+      "predictor_learning_rate": 0.0003581023994323696,
+      "step": 3679
+    },
+    {
+      "epoch": 0.38,
+      "gumbel_tau": 0.4198,
+      "image_masks_loss": 0.0212,
+      "llm_learning_rate": 3.5802758882432666e-06,
+      "loss": 0.6348,
+      "output_text_mask_loss": 0.0292,
+      "predictor_learning_rate": 0.0003580275888243266,
+      "step": 3680
+    },
+    {
+      "epoch": 0.38,
+      "gumbel_tau": 0.4197,
+      "image_masks_loss": 0.0217,
+      "llm_learning_rate": 3.5795276632004498e-06,
+      "loss": 0.7964,
+      "output_text_mask_loss": 0.0076,
+      "predictor_learning_rate": 0.000357952766320045,
+      "step": 3681
+    },
+    {
+      "epoch": 0.38,
+      "gumbel_tau": 0.4196,
+      "image_masks_loss": 0.0155,
+      "llm_learning_rate": 3.578779319277642e-06,
+      "loss": 0.7613,
+      "output_text_mask_loss": 0.0202,
+      "predictor_learning_rate": 0.00035787793192776417,
+      "step": 3682
+    },
+    {
+      "epoch": 0.38,
+      "gumbel_tau": 0.4195,
+      "image_masks_loss": 0.0134,
+      "llm_learning_rate": 3.578030856557253e-06,
+      "loss": 0.7526,
+      "output_text_mask_loss": 0.0463,
+      "predictor_learning_rate": 0.00035780308565572526,
+      "step": 3683
+    },
+    {
+      "epoch": 0.38,
+      "gumbel_tau": 0.4194,
+      "image_masks_loss": 0.018,
+      "llm_learning_rate": 3.5772822751217053e-06,
+      "loss": 0.7921,
+      "output_text_mask_loss": 0.0179,
+      "predictor_learning_rate": 0.0003577282275121705,
+      "step": 3684
+    },
+    {
+      "epoch": 0.38,
+      "gumbel_tau": 0.4193,
+      "image_masks_loss": 0.0172,
+      "llm_learning_rate": 3.576533575053433e-06,
+      "loss": 0.8067,
+      "output_text_mask_loss": 0.0429,
+      "predictor_learning_rate": 0.0003576533575053433,
+      "step": 3685
+    },
+    {
+      "epoch": 0.38,
+      "gumbel_tau": 0.4192,
+      "image_masks_loss": 0.0091,
+      "llm_learning_rate": 3.5757847564348857e-06,
+      "loss": 0.8091,
+      "output_text_mask_loss": 0.0098,
+      "predictor_learning_rate": 0.00035757847564348857,
+      "step": 3686
+    },
+    {
+      "epoch": 0.38,
+      "gumbel_tau": 0.4191,
+      "image_masks_loss": 0.0134,
+      "llm_learning_rate": 3.5750358193485245e-06,
+      "loss": 0.8265,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.0003575035819348524,
+      "step": 3687
+    },
+    {
+      "epoch": 0.38,
+      "gumbel_tau": 0.419,
+      "image_masks_loss": 0.0163,
+      "llm_learning_rate": 3.5742867638768254e-06,
+      "loss": 0.7659,
+      "output_text_mask_loss": 0.0324,
+      "predictor_learning_rate": 0.0003574286763876825,
+      "step": 3688
+    },
+    {
+      "epoch": 0.38,
+      "gumbel_tau": 0.4189,
+      "image_masks_loss": 0.0108,
+      "llm_learning_rate": 3.5735375901022745e-06,
+      "loss": 0.673,
+      "output_text_mask_loss": 0.0021,
+      "predictor_learning_rate": 0.0003573537590102274,
+      "step": 3689
+    },
+    {
+      "epoch": 0.38,
+      "gumbel_tau": 0.4188,
+      "image_masks_loss": 0.0237,
+      "llm_learning_rate": 3.5727882981073725e-06,
+      "loss": 0.7604,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.0003572788298107372,
+      "step": 3690
+    },
+    {
+      "epoch": 0.38,
+      "gumbel_tau": 0.4187,
+      "image_masks_loss": 0.0127,
+      "llm_learning_rate": 3.5720388879746342e-06,
+      "loss": 0.84,
+      "output_text_mask_loss": 0.0018,
+      "predictor_learning_rate": 0.0003572038887974634,
+      "step": 3691
+    },
+    {
+      "epoch": 0.38,
+      "gumbel_tau": 0.4186,
+      "image_masks_loss": 0.0206,
+      "llm_learning_rate": 3.571289359786585e-06,
+      "loss": 0.8518,
+      "output_text_mask_loss": 0.0021,
+      "predictor_learning_rate": 0.00035712893597865847,
+      "step": 3692
+    },
+    {
+      "epoch": 0.38,
+      "gumbel_tau": 0.4185,
+      "image_masks_loss": 0.0184,
+      "llm_learning_rate": 3.570539713625766e-06,
+      "loss": 0.7536,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.00035705397136257655,
+      "step": 3693
+    },
+    {
+      "epoch": 0.38,
+      "gumbel_tau": 0.4184,
+      "image_masks_loss": 0.0139,
+      "llm_learning_rate": 3.569789949574729e-06,
+      "loss": 0.7347,
+      "output_text_mask_loss": 0.0149,
+      "predictor_learning_rate": 0.0003569789949574729,
+      "step": 3694
+    },
+    {
+      "epoch": 0.38,
+      "gumbel_tau": 0.4183,
+      "image_masks_loss": 0.0149,
+      "llm_learning_rate": 3.569040067716039e-06,
+      "loss": 0.717,
+      "output_text_mask_loss": 0.0058,
+      "predictor_learning_rate": 0.0003569040067716039,
+      "step": 3695
+    },
+    {
+      "epoch": 0.38,
+      "gumbel_tau": 0.4182,
+      "image_masks_loss": 0.0259,
+      "llm_learning_rate": 3.5682900681322767e-06,
+      "loss": 0.8391,
+      "output_text_mask_loss": 0.0458,
+      "predictor_learning_rate": 0.0003568290068132277,
+      "step": 3696
+    },
+    {
+      "epoch": 0.38,
+      "gumbel_tau": 0.4181,
+      "image_masks_loss": 0.0137,
+      "llm_learning_rate": 3.5675399509060326e-06,
+      "loss": 0.7936,
+      "output_text_mask_loss": 0.0147,
+      "predictor_learning_rate": 0.00035675399509060325,
+      "step": 3697
+    },
+    {
+      "epoch": 0.38,
+      "gumbel_tau": 0.418,
+      "image_masks_loss": 0.0175,
+      "llm_learning_rate": 3.5667897161199105e-06,
+      "loss": 0.6995,
+      "output_text_mask_loss": 0.0162,
+      "predictor_learning_rate": 0.00035667897161199103,
+      "step": 3698
+    },
+    {
+      "epoch": 0.38,
+      "gumbel_tau": 0.4179,
+      "image_masks_loss": 0.0228,
+      "llm_learning_rate": 3.5660393638565295e-06,
+      "loss": 0.8359,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.0003566039363856529,
+      "step": 3699
+    },
+    {
+      "epoch": 0.38,
+      "gumbel_tau": 0.4178,
+      "image_masks_loss": 0.0058,
+      "llm_learning_rate": 3.565288894198519e-06,
+      "loss": 0.8504,
+      "output_text_mask_loss": 0.0553,
+      "predictor_learning_rate": 0.00035652888941985184,
+      "step": 3700
+    },
+    {
+      "epoch": 0.38,
+      "gumbel_tau": 0.4177,
+      "image_masks_loss": 0.0106,
+      "llm_learning_rate": 3.5645383072285224e-06,
+      "loss": 0.7984,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.0003564538307228522,
+      "step": 3701
+    },
+    {
+      "epoch": 0.38,
+      "gumbel_tau": 0.4176,
+      "image_masks_loss": 0.0176,
+      "llm_learning_rate": 3.5637876030291956e-06,
+      "loss": 0.731,
+      "output_text_mask_loss": 0.0284,
+      "predictor_learning_rate": 0.00035637876030291955,
+      "step": 3702
+    },
+    {
+      "epoch": 0.38,
+      "gumbel_tau": 0.4175,
+      "image_masks_loss": 0.0187,
+      "llm_learning_rate": 3.5630367816832095e-06,
+      "loss": 0.6351,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 0.00035630367816832096,
+      "step": 3703
+    },
+    {
+      "epoch": 0.38,
+      "gumbel_tau": 0.4174,
+      "image_masks_loss": 0.0257,
+      "llm_learning_rate": 3.562285843273245e-06,
+      "loss": 0.7142,
+      "output_text_mask_loss": 0.0187,
+      "predictor_learning_rate": 0.0003562285843273245,
+      "step": 3704
+    },
+    {
+      "epoch": 0.38,
+      "gumbel_tau": 0.4173,
+      "image_masks_loss": 0.0292,
+      "llm_learning_rate": 3.561534787881996e-06,
+      "loss": 0.838,
+      "output_text_mask_loss": 0.006,
+      "predictor_learning_rate": 0.0003561534787881996,
+      "step": 3705
+    },
+    {
+      "epoch": 0.38,
+      "gumbel_tau": 0.4172,
+      "image_masks_loss": 0.0097,
+      "llm_learning_rate": 3.5607836155921733e-06,
+      "loss": 0.7405,
+      "output_text_mask_loss": 0.0225,
+      "predictor_learning_rate": 0.0003560783615592173,
+      "step": 3706
+    },
+    {
+      "epoch": 0.38,
+      "gumbel_tau": 0.4171,
+      "image_masks_loss": 0.0269,
+      "llm_learning_rate": 3.5600323264864963e-06,
+      "loss": 0.8595,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0003560032326486496,
+      "step": 3707
+    },
+    {
+      "epoch": 0.38,
+      "gumbel_tau": 0.417,
+      "image_masks_loss": 0.0111,
+      "llm_learning_rate": 3.559280920647697e-06,
+      "loss": 0.7468,
+      "output_text_mask_loss": 0.0025,
+      "predictor_learning_rate": 0.0003559280920647697,
+      "step": 3708
+    },
+    {
+      "epoch": 0.38,
+      "gumbel_tau": 0.4169,
+      "image_masks_loss": 0.0179,
+      "llm_learning_rate": 3.558529398158524e-06,
+      "loss": 0.8,
+      "output_text_mask_loss": 0.0305,
+      "predictor_learning_rate": 0.00035585293981585235,
+      "step": 3709
+    },
+    {
+      "epoch": 0.38,
+      "gumbel_tau": 0.4168,
+      "image_masks_loss": 0.0228,
+      "llm_learning_rate": 3.557777759101736e-06,
+      "loss": 0.6795,
+      "output_text_mask_loss": 0.0275,
+      "predictor_learning_rate": 0.00035577777591017356,
+      "step": 3710
+    },
+    {
+      "epoch": 0.38,
+      "gumbel_tau": 0.4168,
+      "image_masks_loss": 0.0353,
+      "llm_learning_rate": 3.5570260035601043e-06,
+      "loss": 0.844,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.0003557026003560104,
+      "step": 3711
+    },
+    {
+      "epoch": 0.38,
+      "gumbel_tau": 0.4167,
+      "image_masks_loss": 0.0166,
+      "llm_learning_rate": 3.556274131616415e-06,
+      "loss": 0.7401,
+      "output_text_mask_loss": 0.0117,
+      "predictor_learning_rate": 0.00035562741316164147,
+      "step": 3712
+    },
+    {
+      "epoch": 0.38,
+      "gumbel_tau": 0.4166,
+      "image_masks_loss": 0.0084,
+      "llm_learning_rate": 3.555522143353466e-06,
+      "loss": 0.7176,
+      "output_text_mask_loss": 0.0099,
+      "predictor_learning_rate": 0.0003555522143353466,
+      "step": 3713
+    },
+    {
+      "epoch": 0.38,
+      "gumbel_tau": 0.4165,
+      "image_masks_loss": 0.0086,
+      "llm_learning_rate": 3.5547700388540673e-06,
+      "loss": 0.7156,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00035547700388540674,
+      "step": 3714
+    },
+    {
+      "epoch": 0.38,
+      "gumbel_tau": 0.4164,
+      "image_masks_loss": 0.01,
+      "llm_learning_rate": 3.5540178182010422e-06,
+      "loss": 0.7567,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.0003554017818201042,
+      "step": 3715
+    },
+    {
+      "epoch": 0.38,
+      "gumbel_tau": 0.4163,
+      "image_masks_loss": 0.0256,
+      "llm_learning_rate": 3.553265481477228e-06,
+      "loss": 0.7871,
+      "output_text_mask_loss": 0.006,
+      "predictor_learning_rate": 0.0003553265481477228,
+      "step": 3716
+    },
+    {
+      "epoch": 0.38,
+      "gumbel_tau": 0.4162,
+      "image_masks_loss": 0.0427,
+      "llm_learning_rate": 3.5525130287654723e-06,
+      "loss": 0.7093,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0003552513028765472,
+      "step": 3717
+    },
+    {
+      "epoch": 0.38,
+      "gumbel_tau": 0.4161,
+      "image_masks_loss": 0.0196,
+      "llm_learning_rate": 3.551760460148638e-06,
+      "loss": 0.8311,
+      "output_text_mask_loss": 0.038,
+      "predictor_learning_rate": 0.00035517604601486376,
+      "step": 3718
+    },
+    {
+      "epoch": 0.38,
+      "gumbel_tau": 0.416,
+      "image_masks_loss": 0.0071,
+      "llm_learning_rate": 3.5510077757095984e-06,
+      "loss": 0.8459,
+      "output_text_mask_loss": 0.0231,
+      "predictor_learning_rate": 0.00035510077757095983,
+      "step": 3719
+    },
+    {
+      "epoch": 0.38,
+      "gumbel_tau": 0.4159,
+      "image_masks_loss": 0.0158,
+      "llm_learning_rate": 3.550254975531242e-06,
+      "loss": 0.8998,
+      "output_text_mask_loss": 0.0839,
+      "predictor_learning_rate": 0.0003550254975531242,
+      "step": 3720
+    },
+    {
+      "epoch": 0.38,
+      "gumbel_tau": 0.4158,
+      "image_masks_loss": 0.0196,
+      "llm_learning_rate": 3.5495020596964674e-06,
+      "loss": 0.7234,
+      "output_text_mask_loss": 0.135,
+      "predictor_learning_rate": 0.0003549502059696467,
+      "step": 3721
+    },
+    {
+      "epoch": 0.38,
+      "gumbel_tau": 0.4157,
+      "image_masks_loss": 0.015,
+      "llm_learning_rate": 3.5487490282881894e-06,
+      "loss": 0.7884,
+      "output_text_mask_loss": 0.0809,
+      "predictor_learning_rate": 0.0003548749028288189,
+      "step": 3722
+    },
+    {
+      "epoch": 0.38,
+      "gumbel_tau": 0.4156,
+      "image_masks_loss": 0.0154,
+      "llm_learning_rate": 3.5479958813893323e-06,
+      "loss": 0.8541,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 0.0003547995881389332,
+      "step": 3723
+    },
+    {
+      "epoch": 0.38,
+      "gumbel_tau": 0.4155,
+      "image_masks_loss": 0.0071,
+      "llm_learning_rate": 3.547242619082834e-06,
+      "loss": 1.0547,
+      "output_text_mask_loss": 0.6042,
+      "predictor_learning_rate": 0.00035472426190828336,
+      "step": 3724
+    },
+    {
+      "epoch": 0.38,
+      "gumbel_tau": 0.4154,
+      "image_masks_loss": 0.0237,
+      "llm_learning_rate": 3.5464892414516457e-06,
+      "loss": 1.059,
+      "output_text_mask_loss": 0.1007,
+      "predictor_learning_rate": 0.0003546489241451645,
+      "step": 3725
+    },
+    {
+      "epoch": 0.38,
+      "gumbel_tau": 0.4153,
+      "image_masks_loss": 0.0147,
+      "llm_learning_rate": 3.5457357485787326e-06,
+      "loss": 0.9864,
+      "output_text_mask_loss": 0.0801,
+      "predictor_learning_rate": 0.0003545735748578732,
+      "step": 3726
+    },
+    {
+      "epoch": 0.38,
+      "gumbel_tau": 0.4152,
+      "image_masks_loss": 0.0139,
+      "llm_learning_rate": 3.5449821405470676e-06,
+      "loss": 0.9396,
+      "output_text_mask_loss": 0.1541,
+      "predictor_learning_rate": 0.0003544982140547067,
+      "step": 3727
+    },
+    {
+      "epoch": 0.38,
+      "gumbel_tau": 0.4151,
+      "image_masks_loss": 0.0186,
+      "llm_learning_rate": 3.5442284174396422e-06,
+      "loss": 0.7654,
+      "output_text_mask_loss": 0.074,
+      "predictor_learning_rate": 0.0003544228417439642,
+      "step": 3728
+    },
+    {
+      "epoch": 0.38,
+      "gumbel_tau": 0.415,
+      "image_masks_loss": 0.0098,
+      "llm_learning_rate": 3.543474579339459e-06,
+      "loss": 0.7605,
+      "output_text_mask_loss": 0.0771,
+      "predictor_learning_rate": 0.00035434745793394584,
+      "step": 3729
+    },
+    {
+      "epoch": 0.38,
+      "gumbel_tau": 0.4149,
+      "image_masks_loss": 0.0313,
+      "llm_learning_rate": 3.5427206263295293e-06,
+      "loss": 0.7533,
+      "output_text_mask_loss": 0.0097,
+      "predictor_learning_rate": 0.0003542720626329529,
+      "step": 3730
+    },
+    {
+      "epoch": 0.38,
+      "gumbel_tau": 0.4148,
+      "image_masks_loss": 0.0198,
+      "llm_learning_rate": 3.541966558492882e-06,
+      "loss": 0.7098,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.00035419665584928816,
+      "step": 3731
+    },
+    {
+      "epoch": 0.38,
+      "gumbel_tau": 0.4147,
+      "image_masks_loss": 0.0063,
+      "llm_learning_rate": 3.5412123759125566e-06,
+      "loss": 0.8001,
+      "output_text_mask_loss": 0.0061,
+      "predictor_learning_rate": 0.00035412123759125564,
+      "step": 3732
+    },
+    {
+      "epoch": 0.38,
+      "gumbel_tau": 0.4146,
+      "image_masks_loss": 0.0328,
+      "llm_learning_rate": 3.5404580786716052e-06,
+      "loss": 0.7651,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00035404580786716053,
+      "step": 3733
+    },
+    {
+      "epoch": 0.38,
+      "gumbel_tau": 0.4145,
+      "image_masks_loss": 0.0347,
+      "llm_learning_rate": 3.539703666853092e-06,
+      "loss": 0.7801,
+      "output_text_mask_loss": 0.0736,
+      "predictor_learning_rate": 0.0003539703666853092,
+      "step": 3734
+    },
+    {
+      "epoch": 0.38,
+      "gumbel_tau": 0.4144,
+      "image_masks_loss": 0.0222,
+      "llm_learning_rate": 3.5389491405400967e-06,
+      "loss": 0.7356,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0003538949140540096,
+      "step": 3735
+    },
+    {
+      "epoch": 0.38,
+      "gumbel_tau": 0.4143,
+      "image_masks_loss": 0.0206,
+      "llm_learning_rate": 3.538194499815707e-06,
+      "loss": 0.8912,
+      "output_text_mask_loss": 0.0641,
+      "predictor_learning_rate": 0.0003538194499815707,
+      "step": 3736
+    },
+    {
+      "epoch": 0.38,
+      "gumbel_tau": 0.4142,
+      "image_masks_loss": 0.0188,
+      "llm_learning_rate": 3.5374397447630255e-06,
+      "loss": 0.7942,
+      "output_text_mask_loss": 0.0278,
+      "predictor_learning_rate": 0.0003537439744763025,
+      "step": 3737
+    },
+    {
+      "epoch": 0.38,
+      "gumbel_tau": 0.4141,
+      "image_masks_loss": 0.0124,
+      "llm_learning_rate": 3.5366848754651696e-06,
+      "loss": 0.7638,
+      "output_text_mask_loss": 0.0341,
+      "predictor_learning_rate": 0.0003536684875465169,
+      "step": 3738
+    },
+    {
+      "epoch": 0.38,
+      "gumbel_tau": 0.414,
+      "image_masks_loss": 0.0154,
+      "llm_learning_rate": 3.535929892005266e-06,
+      "loss": 0.6979,
+      "output_text_mask_loss": 0.0301,
+      "predictor_learning_rate": 0.00035359298920052657,
+      "step": 3739
+    },
+    {
+      "epoch": 0.38,
+      "gumbel_tau": 0.4139,
+      "image_masks_loss": 0.0052,
+      "llm_learning_rate": 3.535174794466455e-06,
+      "loss": 0.7433,
+      "output_text_mask_loss": 0.0076,
+      "predictor_learning_rate": 0.0003535174794466455,
+      "step": 3740
+    },
+    {
+      "epoch": 0.38,
+      "gumbel_tau": 0.4138,
+      "image_masks_loss": 0.0212,
+      "llm_learning_rate": 3.5344195829318896e-06,
+      "loss": 0.7593,
+      "output_text_mask_loss": 0.0261,
+      "predictor_learning_rate": 0.0003534419582931889,
+      "step": 3741
+    },
+    {
+      "epoch": 0.38,
+      "gumbel_tau": 0.4137,
+      "image_masks_loss": 0.0119,
+      "llm_learning_rate": 3.533664257484736e-06,
+      "loss": 0.7038,
+      "output_text_mask_loss": 0.1083,
+      "predictor_learning_rate": 0.0003533664257484736,
+      "step": 3742
+    },
+    {
+      "epoch": 0.38,
+      "gumbel_tau": 0.4136,
+      "image_masks_loss": 0.0092,
+      "llm_learning_rate": 3.5329088182081724e-06,
+      "loss": 0.75,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0003532908818208172,
+      "step": 3743
+    },
+    {
+      "epoch": 0.38,
+      "gumbel_tau": 0.4135,
+      "image_masks_loss": 0.0303,
+      "llm_learning_rate": 3.5321532651853884e-06,
+      "loss": 0.7663,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0003532153265185388,
+      "step": 3744
+    },
+    {
+      "epoch": 0.38,
+      "gumbel_tau": 0.4134,
+      "image_masks_loss": 0.0222,
+      "llm_learning_rate": 3.5313975984995887e-06,
+      "loss": 0.7567,
+      "output_text_mask_loss": 0.0584,
+      "predictor_learning_rate": 0.0003531397598499588,
+      "step": 3745
+    },
+    {
+      "epoch": 0.38,
+      "gumbel_tau": 0.4133,
+      "image_masks_loss": 0.0157,
+      "llm_learning_rate": 3.5306418182339873e-06,
+      "loss": 0.695,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0003530641818233987,
+      "step": 3746
+    },
+    {
+      "epoch": 0.38,
+      "gumbel_tau": 0.4132,
+      "image_masks_loss": 0.0107,
+      "llm_learning_rate": 3.529885924471813e-06,
+      "loss": 0.7026,
+      "output_text_mask_loss": 0.0025,
+      "predictor_learning_rate": 0.0003529885924471813,
+      "step": 3747
+    },
+    {
+      "epoch": 0.38,
+      "gumbel_tau": 0.4131,
+      "image_masks_loss": 0.0131,
+      "llm_learning_rate": 3.5291299172963068e-06,
+      "loss": 0.7602,
+      "output_text_mask_loss": 0.0687,
+      "predictor_learning_rate": 0.00035291299172963064,
+      "step": 3748
+    },
+    {
+      "epoch": 0.38,
+      "gumbel_tau": 0.413,
+      "image_masks_loss": 0.0103,
+      "llm_learning_rate": 3.5283737967907224e-06,
+      "loss": 0.7204,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.0003528373796790722,
+      "step": 3749
+    },
+    {
+      "epoch": 0.38,
+      "gumbel_tau": 0.4129,
+      "image_masks_loss": 0.0152,
+      "llm_learning_rate": 3.527617563038324e-06,
+      "loss": 0.7455,
+      "output_text_mask_loss": 0.0056,
+      "predictor_learning_rate": 0.0003527617563038324,
+      "step": 3750
+    },
+    {
+      "epoch": 0.38,
+      "gumbel_tau": 0.4128,
+      "image_masks_loss": 0.0068,
+      "llm_learning_rate": 3.5268612161223916e-06,
+      "loss": 0.7854,
+      "output_text_mask_loss": 0.0288,
+      "predictor_learning_rate": 0.0003526861216122391,
+      "step": 3751
+    },
+    {
+      "epoch": 0.38,
+      "gumbel_tau": 0.4127,
+      "image_masks_loss": 0.0164,
+      "llm_learning_rate": 3.526104756126214e-06,
+      "loss": 0.7367,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0003526104756126214,
+      "step": 3752
+    },
+    {
+      "epoch": 0.38,
+      "gumbel_tau": 0.4126,
+      "image_masks_loss": 0.0141,
+      "llm_learning_rate": 3.525348183133096e-06,
+      "loss": 0.8882,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.0003525348183133096,
+      "step": 3753
+    },
+    {
+      "epoch": 0.38,
+      "gumbel_tau": 0.4125,
+      "image_masks_loss": 0.0126,
+      "llm_learning_rate": 3.524591497226352e-06,
+      "loss": 0.8447,
+      "output_text_mask_loss": 0.0679,
+      "predictor_learning_rate": 0.0003524591497226352,
+      "step": 3754
+    },
+    {
+      "epoch": 0.38,
+      "gumbel_tau": 0.4124,
+      "image_masks_loss": 0.0092,
+      "llm_learning_rate": 3.5238346984893097e-06,
+      "loss": 0.6788,
+      "output_text_mask_loss": 0.0294,
+      "predictor_learning_rate": 0.0003523834698489309,
+      "step": 3755
+    },
+    {
+      "epoch": 0.38,
+      "gumbel_tau": 0.4124,
+      "image_masks_loss": 0.0105,
+      "llm_learning_rate": 3.5230777870053102e-06,
+      "loss": 0.7731,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.000352307778700531,
+      "step": 3756
+    },
+    {
+      "epoch": 0.38,
+      "gumbel_tau": 0.4123,
+      "image_masks_loss": 0.0261,
+      "llm_learning_rate": 3.5223207628577054e-06,
+      "loss": 0.8039,
+      "output_text_mask_loss": 0.0134,
+      "predictor_learning_rate": 0.0003522320762857705,
+      "step": 3757
+    },
+    {
+      "epoch": 0.39,
+      "gumbel_tau": 0.4122,
+      "image_masks_loss": 0.0305,
+      "llm_learning_rate": 3.5215636261298614e-06,
+      "loss": 0.8329,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.00035215636261298614,
+      "step": 3758
+    },
+    {
+      "epoch": 0.39,
+      "gumbel_tau": 0.4121,
+      "image_masks_loss": 0.0152,
+      "llm_learning_rate": 3.520806376905156e-06,
+      "loss": 0.6653,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0003520806376905156,
+      "step": 3759
+    },
+    {
+      "epoch": 0.39,
+      "gumbel_tau": 0.412,
+      "image_masks_loss": 0.018,
+      "llm_learning_rate": 3.5200490152669774e-06,
+      "loss": 0.771,
+      "output_text_mask_loss": 0.0025,
+      "predictor_learning_rate": 0.00035200490152669775,
+      "step": 3760
+    },
+    {
+      "epoch": 0.39,
+      "gumbel_tau": 0.4119,
+      "image_masks_loss": 0.023,
+      "llm_learning_rate": 3.5192915412987304e-06,
+      "loss": 0.8074,
+      "output_text_mask_loss": 0.0359,
+      "predictor_learning_rate": 0.00035192915412987305,
+      "step": 3761
+    },
+    {
+      "epoch": 0.39,
+      "gumbel_tau": 0.4118,
+      "image_masks_loss": 0.0109,
+      "llm_learning_rate": 3.518533955083828e-06,
+      "loss": 0.7399,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.0003518533955083828,
+      "step": 3762
+    },
+    {
+      "epoch": 0.39,
+      "gumbel_tau": 0.4117,
+      "image_masks_loss": 0.034,
+      "llm_learning_rate": 3.5177762567056977e-06,
+      "loss": 0.7268,
+      "output_text_mask_loss": 0.0257,
+      "predictor_learning_rate": 0.0003517776256705698,
+      "step": 3763
+    },
+    {
+      "epoch": 0.39,
+      "gumbel_tau": 0.4116,
+      "image_masks_loss": 0.0243,
+      "llm_learning_rate": 3.5170184462477797e-06,
+      "loss": 0.7393,
+      "output_text_mask_loss": 0.0065,
+      "predictor_learning_rate": 0.00035170184462477794,
+      "step": 3764
+    },
+    {
+      "epoch": 0.39,
+      "gumbel_tau": 0.4115,
+      "image_masks_loss": 0.0072,
+      "llm_learning_rate": 3.5162605237935243e-06,
+      "loss": 0.7314,
+      "output_text_mask_loss": 0.0193,
+      "predictor_learning_rate": 0.0003516260523793524,
+      "step": 3765
+    },
+    {
+      "epoch": 0.39,
+      "gumbel_tau": 0.4114,
+      "image_masks_loss": 0.0296,
+      "llm_learning_rate": 3.515502489426397e-06,
+      "loss": 0.824,
+      "output_text_mask_loss": 0.0645,
+      "predictor_learning_rate": 0.0003515502489426397,
+      "step": 3766
+    },
+    {
+      "epoch": 0.39,
+      "gumbel_tau": 0.4113,
+      "image_masks_loss": 0.0118,
+      "llm_learning_rate": 3.514744343229872e-06,
+      "loss": 0.7804,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0003514744343229872,
+      "step": 3767
+    },
+    {
+      "epoch": 0.39,
+      "gumbel_tau": 0.4112,
+      "image_masks_loss": 0.019,
+      "llm_learning_rate": 3.5139860852874414e-06,
+      "loss": 0.8015,
+      "output_text_mask_loss": 0.0444,
+      "predictor_learning_rate": 0.0003513986085287441,
+      "step": 3768
+    },
+    {
+      "epoch": 0.39,
+      "gumbel_tau": 0.4111,
+      "image_masks_loss": 0.0233,
+      "llm_learning_rate": 3.513227715682604e-06,
+      "loss": 0.8389,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.0003513227715682604,
+      "step": 3769
+    },
+    {
+      "epoch": 0.39,
+      "gumbel_tau": 0.411,
+      "image_masks_loss": 0.0209,
+      "llm_learning_rate": 3.512469234498873e-06,
+      "loss": 0.8022,
+      "output_text_mask_loss": 0.0381,
+      "predictor_learning_rate": 0.0003512469234498873,
+      "step": 3770
+    },
+    {
+      "epoch": 0.39,
+      "gumbel_tau": 0.4109,
+      "image_masks_loss": 0.0134,
+      "llm_learning_rate": 3.5117106418197756e-06,
+      "loss": 0.7431,
+      "output_text_mask_loss": 0.0717,
+      "predictor_learning_rate": 0.00035117106418197755,
+      "step": 3771
+    },
+    {
+      "epoch": 0.39,
+      "gumbel_tau": 0.4108,
+      "image_masks_loss": 0.0362,
+      "llm_learning_rate": 3.5109519377288494e-06,
+      "loss": 0.7349,
+      "output_text_mask_loss": 0.0332,
+      "predictor_learning_rate": 0.0003510951937728849,
+      "step": 3772
+    },
+    {
+      "epoch": 0.39,
+      "gumbel_tau": 0.4107,
+      "image_masks_loss": 0.0319,
+      "llm_learning_rate": 3.5101931223096423e-06,
+      "loss": 0.776,
+      "output_text_mask_loss": 0.0139,
+      "predictor_learning_rate": 0.0003510193122309642,
+      "step": 3773
+    },
+    {
+      "epoch": 0.39,
+      "gumbel_tau": 0.4106,
+      "image_masks_loss": 0.0161,
+      "llm_learning_rate": 3.50943419564572e-06,
+      "loss": 0.7359,
+      "output_text_mask_loss": 0.0124,
+      "predictor_learning_rate": 0.000350943419564572,
+      "step": 3774
+    },
+    {
+      "epoch": 0.39,
+      "gumbel_tau": 0.4105,
+      "image_masks_loss": 0.023,
+      "llm_learning_rate": 3.5086751578206546e-06,
+      "loss": 0.7876,
+      "output_text_mask_loss": 0.0479,
+      "predictor_learning_rate": 0.00035086751578206543,
+      "step": 3775
+    },
+    {
+      "epoch": 0.39,
+      "gumbel_tau": 0.4104,
+      "image_masks_loss": 0.0157,
+      "llm_learning_rate": 3.507916008918034e-06,
+      "loss": 0.7658,
+      "output_text_mask_loss": 0.0076,
+      "predictor_learning_rate": 0.0003507916008918034,
+      "step": 3776
+    },
+    {
+      "epoch": 0.39,
+      "gumbel_tau": 0.4103,
+      "image_masks_loss": 0.0168,
+      "llm_learning_rate": 3.5071567490214585e-06,
+      "loss": 0.7827,
+      "output_text_mask_loss": 0.0534,
+      "predictor_learning_rate": 0.00035071567490214587,
+      "step": 3777
+    },
+    {
+      "epoch": 0.39,
+      "gumbel_tau": 0.4102,
+      "image_masks_loss": 0.0399,
+      "llm_learning_rate": 3.506397378214538e-06,
+      "loss": 0.8083,
+      "output_text_mask_loss": 0.0858,
+      "predictor_learning_rate": 0.00035063973782145376,
+      "step": 3778
+    },
+    {
+      "epoch": 0.39,
+      "gumbel_tau": 0.4101,
+      "image_masks_loss": 0.0294,
+      "llm_learning_rate": 3.5056378965808962e-06,
+      "loss": 0.7606,
+      "output_text_mask_loss": 0.0374,
+      "predictor_learning_rate": 0.0003505637896580896,
+      "step": 3779
+    },
+    {
+      "epoch": 0.39,
+      "gumbel_tau": 0.41,
+      "image_masks_loss": 0.0163,
+      "llm_learning_rate": 3.504878304204169e-06,
+      "loss": 0.873,
+      "output_text_mask_loss": 0.0332,
+      "predictor_learning_rate": 0.00035048783042041686,
+      "step": 3780
+    },
+    {
+      "epoch": 0.39,
+      "gumbel_tau": 0.4099,
+      "image_masks_loss": 0.0374,
+      "llm_learning_rate": 3.504118601168006e-06,
+      "loss": 0.8646,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 0.0003504118601168006,
+      "step": 3781
+    },
+    {
+      "epoch": 0.39,
+      "gumbel_tau": 0.4098,
+      "image_masks_loss": 0.0084,
+      "llm_learning_rate": 3.5033587875560656e-06,
+      "loss": 0.6354,
+      "output_text_mask_loss": 0.0702,
+      "predictor_learning_rate": 0.00035033587875560654,
+      "step": 3782
+    },
+    {
+      "epoch": 0.39,
+      "gumbel_tau": 0.4097,
+      "image_masks_loss": 0.0145,
+      "llm_learning_rate": 3.50259886345202e-06,
+      "loss": 0.7669,
+      "output_text_mask_loss": 0.013,
+      "predictor_learning_rate": 0.00035025988634520194,
+      "step": 3783
+    },
+    {
+      "epoch": 0.39,
+      "gumbel_tau": 0.4096,
+      "image_masks_loss": 0.0113,
+      "llm_learning_rate": 3.5018388289395543e-06,
+      "loss": 0.7822,
+      "output_text_mask_loss": 0.0244,
+      "predictor_learning_rate": 0.00035018388289395543,
+      "step": 3784
+    },
+    {
+      "epoch": 0.39,
+      "gumbel_tau": 0.4095,
+      "image_masks_loss": 0.0143,
+      "llm_learning_rate": 3.501078684102366e-06,
+      "loss": 0.7535,
+      "output_text_mask_loss": 0.0172,
+      "predictor_learning_rate": 0.00035010786841023657,
+      "step": 3785
+    },
+    {
+      "epoch": 0.39,
+      "gumbel_tau": 0.4094,
+      "image_masks_loss": 0.0261,
+      "llm_learning_rate": 3.500318429024162e-06,
+      "loss": 0.7927,
+      "output_text_mask_loss": 0.0896,
+      "predictor_learning_rate": 0.0003500318429024162,
+      "step": 3786
+    },
+    {
+      "epoch": 0.39,
+      "gumbel_tau": 0.4093,
+      "image_masks_loss": 0.0067,
+      "llm_learning_rate": 3.499558063788665e-06,
+      "loss": 0.7563,
+      "output_text_mask_loss": 0.0679,
+      "predictor_learning_rate": 0.0003499558063788665,
+      "step": 3787
+    },
+    {
+      "epoch": 0.39,
+      "gumbel_tau": 0.4092,
+      "image_masks_loss": 0.0118,
+      "llm_learning_rate": 3.498797588479608e-06,
+      "loss": 0.7332,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 0.00034987975884796075,
+      "step": 3788
+    },
+    {
+      "epoch": 0.39,
+      "gumbel_tau": 0.4092,
+      "image_masks_loss": 0.0116,
+      "llm_learning_rate": 3.498037003180736e-06,
+      "loss": 0.7542,
+      "output_text_mask_loss": 0.0322,
+      "predictor_learning_rate": 0.0003498037003180736,
+      "step": 3789
+    },
+    {
+      "epoch": 0.39,
+      "gumbel_tau": 0.4091,
+      "image_masks_loss": 0.0093,
+      "llm_learning_rate": 3.4972763079758053e-06,
+      "loss": 0.769,
+      "output_text_mask_loss": 0.1038,
+      "predictor_learning_rate": 0.00034972763079758055,
+      "step": 3790
+    },
+    {
+      "epoch": 0.39,
+      "gumbel_tau": 0.409,
+      "image_masks_loss": 0.034,
+      "llm_learning_rate": 3.496515502948587e-06,
+      "loss": 0.7794,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0003496515502948587,
+      "step": 3791
+    },
+    {
+      "epoch": 0.39,
+      "gumbel_tau": 0.4089,
+      "image_masks_loss": 0.0221,
+      "llm_learning_rate": 3.4957545881828614e-06,
+      "loss": 0.7594,
+      "output_text_mask_loss": 0.0763,
+      "predictor_learning_rate": 0.0003495754588182861,
+      "step": 3792
+    },
+    {
+      "epoch": 0.39,
+      "gumbel_tau": 0.4088,
+      "image_masks_loss": 0.0114,
+      "llm_learning_rate": 3.494993563762423e-06,
+      "loss": 0.8093,
+      "output_text_mask_loss": 0.0134,
+      "predictor_learning_rate": 0.00034949935637624223,
+      "step": 3793
+    },
+    {
+      "epoch": 0.39,
+      "gumbel_tau": 0.4087,
+      "image_masks_loss": 0.0209,
+      "llm_learning_rate": 3.4942324297710767e-06,
+      "loss": 0.8021,
+      "output_text_mask_loss": 0.0023,
+      "predictor_learning_rate": 0.00034942324297710764,
+      "step": 3794
+    },
+    {
+      "epoch": 0.39,
+      "gumbel_tau": 0.4086,
+      "image_masks_loss": 0.0117,
+      "llm_learning_rate": 3.4934711862926405e-06,
+      "loss": 0.7286,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.00034934711862926403,
+      "step": 3795
+    },
+    {
+      "epoch": 0.39,
+      "gumbel_tau": 0.4085,
+      "image_masks_loss": 0.034,
+      "llm_learning_rate": 3.4927098334109444e-06,
+      "loss": 0.7568,
+      "output_text_mask_loss": 0.0202,
+      "predictor_learning_rate": 0.0003492709833410944,
+      "step": 3796
+    },
+    {
+      "epoch": 0.39,
+      "gumbel_tau": 0.4084,
+      "image_masks_loss": 0.0051,
+      "llm_learning_rate": 3.491948371209831e-06,
+      "loss": 0.8261,
+      "output_text_mask_loss": 0.0401,
+      "predictor_learning_rate": 0.0003491948371209831,
+      "step": 3797
+    },
+    {
+      "epoch": 0.39,
+      "gumbel_tau": 0.4083,
+      "image_masks_loss": 0.0088,
+      "llm_learning_rate": 3.491186799773153e-06,
+      "loss": 0.7479,
+      "output_text_mask_loss": 0.0038,
+      "predictor_learning_rate": 0.00034911867997731527,
+      "step": 3798
+    },
+    {
+      "epoch": 0.39,
+      "gumbel_tau": 0.4082,
+      "image_masks_loss": 0.01,
+      "llm_learning_rate": 3.4904251191847767e-06,
+      "loss": 0.6794,
+      "output_text_mask_loss": 0.0137,
+      "predictor_learning_rate": 0.00034904251191847764,
+      "step": 3799
+    },
+    {
+      "epoch": 0.39,
+      "gumbel_tau": 0.4081,
+      "image_masks_loss": 0.0129,
+      "llm_learning_rate": 3.489663329528581e-06,
+      "loss": 0.8087,
+      "output_text_mask_loss": 0.0607,
+      "predictor_learning_rate": 0.0003489663329528581,
+      "step": 3800
+    },
+    {
+      "epoch": 0.39,
+      "gumbel_tau": 0.408,
+      "image_masks_loss": 0.016,
+      "llm_learning_rate": 3.4889014308884546e-06,
+      "loss": 0.9232,
+      "output_text_mask_loss": 0.0065,
+      "predictor_learning_rate": 0.00034889014308884543,
+      "step": 3801
+    },
+    {
+      "epoch": 0.39,
+      "gumbel_tau": 0.4079,
+      "image_masks_loss": 0.0313,
+      "llm_learning_rate": 3.488139423348299e-06,
+      "loss": 0.6948,
+      "output_text_mask_loss": 0.0011,
+      "predictor_learning_rate": 0.00034881394233482986,
+      "step": 3802
+    },
+    {
+      "epoch": 0.39,
+      "gumbel_tau": 0.4078,
+      "image_masks_loss": 0.018,
+      "llm_learning_rate": 3.48737730699203e-06,
+      "loss": 0.8553,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.000348737730699203,
+      "step": 3803
+    },
+    {
+      "epoch": 0.39,
+      "gumbel_tau": 0.4077,
+      "image_masks_loss": 0.0378,
+      "llm_learning_rate": 3.4866150819035726e-06,
+      "loss": 0.7279,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.00034866150819035726,
+      "step": 3804
+    },
+    {
+      "epoch": 0.39,
+      "gumbel_tau": 0.4076,
+      "image_masks_loss": 0.0193,
+      "llm_learning_rate": 3.4858527481668647e-06,
+      "loss": 0.7175,
+      "output_text_mask_loss": 0.0278,
+      "predictor_learning_rate": 0.00034858527481668644,
+      "step": 3805
+    },
+    {
+      "epoch": 0.39,
+      "gumbel_tau": 0.4075,
+      "image_masks_loss": 0.008,
+      "llm_learning_rate": 3.485090305865856e-06,
+      "loss": 0.7734,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.00034850903058658556,
+      "step": 3806
+    },
+    {
+      "epoch": 0.39,
+      "gumbel_tau": 0.4074,
+      "image_masks_loss": 0.0225,
+      "llm_learning_rate": 3.484327755084509e-06,
+      "loss": 0.7801,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0003484327755084509,
+      "step": 3807
+    },
+    {
+      "epoch": 0.39,
+      "gumbel_tau": 0.4073,
+      "image_masks_loss": 0.0191,
+      "llm_learning_rate": 3.483565095906797e-06,
+      "loss": 0.7182,
+      "output_text_mask_loss": 0.0023,
+      "predictor_learning_rate": 0.00034835650959067974,
+      "step": 3808
+    },
+    {
+      "epoch": 0.39,
+      "gumbel_tau": 0.4072,
+      "image_masks_loss": 0.0214,
+      "llm_learning_rate": 3.4828023284167062e-06,
+      "loss": 0.7768,
+      "output_text_mask_loss": 0.0078,
+      "predictor_learning_rate": 0.0003482802328416706,
+      "step": 3809
+    },
+    {
+      "epoch": 0.39,
+      "gumbel_tau": 0.4071,
+      "image_masks_loss": 0.0324,
+      "llm_learning_rate": 3.4820394526982337e-06,
+      "loss": 0.8328,
+      "output_text_mask_loss": 0.0004,
+      "predictor_learning_rate": 0.00034820394526982335,
+      "step": 3810
+    },
+    {
+      "epoch": 0.39,
+      "gumbel_tau": 0.407,
+      "image_masks_loss": 0.0133,
+      "llm_learning_rate": 3.4812764688353894e-06,
+      "loss": 0.6878,
+      "output_text_mask_loss": 0.05,
+      "predictor_learning_rate": 0.0003481276468835389,
+      "step": 3811
+    },
+    {
+      "epoch": 0.39,
+      "gumbel_tau": 0.4069,
+      "image_masks_loss": 0.0191,
+      "llm_learning_rate": 3.4805133769121933e-06,
+      "loss": 0.767,
+      "output_text_mask_loss": 0.0496,
+      "predictor_learning_rate": 0.0003480513376912193,
+      "step": 3812
+    },
+    {
+      "epoch": 0.39,
+      "gumbel_tau": 0.4068,
+      "image_masks_loss": 0.0269,
+      "llm_learning_rate": 3.479750177012682e-06,
+      "loss": 0.7432,
+      "output_text_mask_loss": 0.0277,
+      "predictor_learning_rate": 0.0003479750177012682,
+      "step": 3813
+    },
+    {
+      "epoch": 0.39,
+      "gumbel_tau": 0.4067,
+      "image_masks_loss": 0.0294,
+      "llm_learning_rate": 3.478986869220898e-06,
+      "loss": 0.7778,
+      "output_text_mask_loss": 0.0213,
+      "predictor_learning_rate": 0.00034789868692208976,
+      "step": 3814
+    },
+    {
+      "epoch": 0.39,
+      "gumbel_tau": 0.4067,
+      "image_masks_loss": 0.0484,
+      "llm_learning_rate": 3.478223453620899e-06,
+      "loss": 0.7653,
+      "output_text_mask_loss": 0.0431,
+      "predictor_learning_rate": 0.0003478223453620899,
+      "step": 3815
+    },
+    {
+      "epoch": 0.39,
+      "gumbel_tau": 0.4066,
+      "image_masks_loss": 0.0301,
+      "llm_learning_rate": 3.477459930296755e-06,
+      "loss": 0.7401,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.00034774599302967547,
+      "step": 3816
+    },
+    {
+      "epoch": 0.39,
+      "gumbel_tau": 0.4065,
+      "image_masks_loss": 0.0172,
+      "llm_learning_rate": 3.476696299332546e-06,
+      "loss": 0.8032,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.0003476696299332546,
+      "step": 3817
+    },
+    {
+      "epoch": 0.39,
+      "gumbel_tau": 0.4064,
+      "image_masks_loss": 0.0076,
+      "llm_learning_rate": 3.475932560812365e-06,
+      "loss": 0.6792,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0003475932560812365,
+      "step": 3818
+    },
+    {
+      "epoch": 0.39,
+      "gumbel_tau": 0.4063,
+      "image_masks_loss": 0.005,
+      "llm_learning_rate": 3.4751687148203167e-06,
+      "loss": 0.7987,
+      "output_text_mask_loss": 0.0145,
+      "predictor_learning_rate": 0.00034751687148203164,
+      "step": 3819
+    },
+    {
+      "epoch": 0.39,
+      "gumbel_tau": 0.4062,
+      "image_masks_loss": 0.0062,
+      "llm_learning_rate": 3.4744047614405172e-06,
+      "loss": 0.7494,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.0003474404761440517,
+      "step": 3820
+    },
+    {
+      "epoch": 0.39,
+      "gumbel_tau": 0.4061,
+      "image_masks_loss": 0.0155,
+      "llm_learning_rate": 3.473640700757096e-06,
+      "loss": 0.693,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.0003473640700757096,
+      "step": 3821
+    },
+    {
+      "epoch": 0.39,
+      "gumbel_tau": 0.406,
+      "image_masks_loss": 0.0128,
+      "llm_learning_rate": 3.4728765328541906e-06,
+      "loss": 0.7766,
+      "output_text_mask_loss": 0.0198,
+      "predictor_learning_rate": 0.0003472876532854191,
+      "step": 3822
+    },
+    {
+      "epoch": 0.39,
+      "gumbel_tau": 0.4059,
+      "image_masks_loss": 0.0081,
+      "llm_learning_rate": 3.472112257815955e-06,
+      "loss": 0.802,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.00034721122578159545,
+      "step": 3823
+    },
+    {
+      "epoch": 0.39,
+      "gumbel_tau": 0.4058,
+      "image_masks_loss": 0.0324,
+      "llm_learning_rate": 3.471347875726553e-06,
+      "loss": 0.7675,
+      "output_text_mask_loss": 0.0025,
+      "predictor_learning_rate": 0.0003471347875726553,
+      "step": 3824
+    },
+    {
+      "epoch": 0.39,
+      "gumbel_tau": 0.4057,
+      "image_masks_loss": 0.0115,
+      "llm_learning_rate": 3.4705833866701584e-06,
+      "loss": 0.7678,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.0003470583386670158,
+      "step": 3825
+    },
+    {
+      "epoch": 0.39,
+      "gumbel_tau": 0.4056,
+      "image_masks_loss": 0.0179,
+      "llm_learning_rate": 3.4698187907309593e-06,
+      "loss": 0.8418,
+      "output_text_mask_loss": 0.1076,
+      "predictor_learning_rate": 0.0003469818790730959,
+      "step": 3826
+    },
+    {
+      "epoch": 0.39,
+      "gumbel_tau": 0.4055,
+      "image_masks_loss": 0.0254,
+      "llm_learning_rate": 3.4690540879931562e-06,
+      "loss": 0.8149,
+      "output_text_mask_loss": 0.046,
+      "predictor_learning_rate": 0.0003469054087993156,
+      "step": 3827
+    },
+    {
+      "epoch": 0.39,
+      "gumbel_tau": 0.4054,
+      "image_masks_loss": 0.0052,
+      "llm_learning_rate": 3.4682892785409565e-06,
+      "loss": 0.7407,
+      "output_text_mask_loss": 0.0299,
+      "predictor_learning_rate": 0.0003468289278540956,
+      "step": 3828
+    },
+    {
+      "epoch": 0.39,
+      "gumbel_tau": 0.4053,
+      "image_masks_loss": 0.0204,
+      "llm_learning_rate": 3.4675243624585857e-06,
+      "loss": 0.8174,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.0003467524362458586,
+      "step": 3829
+    },
+    {
+      "epoch": 0.39,
+      "gumbel_tau": 0.4052,
+      "image_masks_loss": 0.0477,
+      "llm_learning_rate": 3.4667593398302774e-06,
+      "loss": 0.7303,
+      "output_text_mask_loss": 0.0216,
+      "predictor_learning_rate": 0.0003466759339830277,
+      "step": 3830
+    },
+    {
+      "epoch": 0.39,
+      "gumbel_tau": 0.4051,
+      "image_masks_loss": 0.0349,
+      "llm_learning_rate": 3.465994210740276e-06,
+      "loss": 0.5987,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0003465994210740276,
+      "step": 3831
+    },
+    {
+      "epoch": 0.39,
+      "gumbel_tau": 0.405,
+      "image_masks_loss": 0.0175,
+      "llm_learning_rate": 3.4652289752728412e-06,
+      "loss": 0.7515,
+      "output_text_mask_loss": 0.0299,
+      "predictor_learning_rate": 0.0003465228975272841,
+      "step": 3832
+    },
+    {
+      "epoch": 0.39,
+      "gumbel_tau": 0.4049,
+      "image_masks_loss": 0.0139,
+      "llm_learning_rate": 3.4644636335122416e-06,
+      "loss": 0.7608,
+      "output_text_mask_loss": 0.0263,
+      "predictor_learning_rate": 0.00034644636335122415,
+      "step": 3833
+    },
+    {
+      "epoch": 0.39,
+      "gumbel_tau": 0.4048,
+      "image_masks_loss": 0.0406,
+      "llm_learning_rate": 3.463698185542759e-06,
+      "loss": 0.8,
+      "output_text_mask_loss": 0.0098,
+      "predictor_learning_rate": 0.00034636981855427586,
+      "step": 3834
+    },
+    {
+      "epoch": 0.39,
+      "gumbel_tau": 0.4047,
+      "image_masks_loss": 0.0191,
+      "llm_learning_rate": 3.4629326314486846e-06,
+      "loss": 0.743,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.00034629326314486843,
+      "step": 3835
+    },
+    {
+      "epoch": 0.39,
+      "gumbel_tau": 0.4046,
+      "image_masks_loss": 0.0296,
+      "llm_learning_rate": 3.462166971314325e-06,
+      "loss": 0.838,
+      "output_text_mask_loss": 0.0137,
+      "predictor_learning_rate": 0.00034621669713143244,
+      "step": 3836
+    },
+    {
+      "epoch": 0.39,
+      "gumbel_tau": 0.4045,
+      "image_masks_loss": 0.0301,
+      "llm_learning_rate": 3.4614012052239954e-06,
+      "loss": 0.7211,
+      "output_text_mask_loss": 0.0055,
+      "predictor_learning_rate": 0.0003461401205223995,
+      "step": 3837
+    },
+    {
+      "epoch": 0.39,
+      "gumbel_tau": 0.4044,
+      "image_masks_loss": 0.0098,
+      "llm_learning_rate": 3.4606353332620226e-06,
+      "loss": 0.844,
+      "output_text_mask_loss": 0.0675,
+      "predictor_learning_rate": 0.00034606353332620224,
+      "step": 3838
+    },
+    {
+      "epoch": 0.39,
+      "gumbel_tau": 0.4044,
+      "image_masks_loss": 0.0186,
+      "llm_learning_rate": 3.4598693555127476e-06,
+      "loss": 0.8138,
+      "output_text_mask_loss": 0.0012,
+      "predictor_learning_rate": 0.00034598693555127475,
+      "step": 3839
+    },
+    {
+      "epoch": 0.39,
+      "gumbel_tau": 0.4043,
+      "image_masks_loss": 0.0542,
+      "llm_learning_rate": 3.459103272060521e-06,
+      "loss": 0.7851,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 0.0003459103272060521,
+      "step": 3840
+    },
+    {
+      "epoch": 0.39,
+      "gumbel_tau": 0.4042,
+      "image_masks_loss": 0.0404,
+      "llm_learning_rate": 3.4583370829897064e-06,
+      "loss": 0.8317,
+      "output_text_mask_loss": 0.0208,
+      "predictor_learning_rate": 0.0003458337082989706,
+      "step": 3841
+    },
+    {
+      "epoch": 0.39,
+      "gumbel_tau": 0.4041,
+      "image_masks_loss": 0.0196,
+      "llm_learning_rate": 3.457570788384677e-06,
+      "loss": 0.7292,
+      "output_text_mask_loss": 0.0038,
+      "predictor_learning_rate": 0.0003457570788384677,
+      "step": 3842
+    },
+    {
+      "epoch": 0.39,
+      "gumbel_tau": 0.404,
+      "image_masks_loss": 0.0387,
+      "llm_learning_rate": 3.4568043883298197e-06,
+      "loss": 0.812,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 0.00034568043883298195,
+      "step": 3843
+    },
+    {
+      "epoch": 0.39,
+      "gumbel_tau": 0.4039,
+      "image_masks_loss": 0.0313,
+      "llm_learning_rate": 3.456037882909532e-06,
+      "loss": 0.7252,
+      "output_text_mask_loss": 0.0523,
+      "predictor_learning_rate": 0.0003456037882909532,
+      "step": 3844
+    },
+    {
+      "epoch": 0.39,
+      "gumbel_tau": 0.4038,
+      "image_masks_loss": 0.061,
+      "llm_learning_rate": 3.4552712722082225e-06,
+      "loss": 0.7177,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.00034552712722082224,
+      "step": 3845
+    },
+    {
+      "epoch": 0.39,
+      "gumbel_tau": 0.4037,
+      "image_masks_loss": 0.015,
+      "llm_learning_rate": 3.454504556310314e-06,
+      "loss": 0.7483,
+      "output_text_mask_loss": 0.0025,
+      "predictor_learning_rate": 0.0003454504556310314,
+      "step": 3846
+    },
+    {
+      "epoch": 0.39,
+      "gumbel_tau": 0.4036,
+      "image_masks_loss": 0.0149,
+      "llm_learning_rate": 3.4537377353002366e-06,
+      "loss": 0.8395,
+      "output_text_mask_loss": 0.0721,
+      "predictor_learning_rate": 0.0003453737735300236,
+      "step": 3847
+    },
+    {
+      "epoch": 0.39,
+      "gumbel_tau": 0.4035,
+      "image_masks_loss": 0.029,
+      "llm_learning_rate": 3.452970809262436e-06,
+      "loss": 0.8051,
+      "output_text_mask_loss": 0.0145,
+      "predictor_learning_rate": 0.00034529708092624357,
+      "step": 3848
+    },
+    {
+      "epoch": 0.39,
+      "gumbel_tau": 0.4034,
+      "image_masks_loss": 0.0267,
+      "llm_learning_rate": 3.452203778281367e-06,
+      "loss": 0.6978,
+      "output_text_mask_loss": 0.0429,
+      "predictor_learning_rate": 0.0003452203778281367,
+      "step": 3849
+    },
+    {
+      "epoch": 0.39,
+      "gumbel_tau": 0.4033,
+      "image_masks_loss": 0.0108,
+      "llm_learning_rate": 3.4514366424414975e-06,
+      "loss": 0.785,
+      "output_text_mask_loss": 0.0227,
+      "predictor_learning_rate": 0.0003451436642441497,
+      "step": 3850
+    },
+    {
+      "epoch": 0.39,
+      "gumbel_tau": 0.4032,
+      "image_masks_loss": 0.0204,
+      "llm_learning_rate": 3.4506694018273044e-06,
+      "loss": 0.8672,
+      "output_text_mask_loss": 0.0992,
+      "predictor_learning_rate": 0.00034506694018273044,
+      "step": 3851
+    },
+    {
+      "epoch": 0.39,
+      "gumbel_tau": 0.4031,
+      "image_masks_loss": 0.0147,
+      "llm_learning_rate": 3.4499020565232805e-06,
+      "loss": 0.7173,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.000344990205652328,
+      "step": 3852
+    },
+    {
+      "epoch": 0.39,
+      "gumbel_tau": 0.403,
+      "image_masks_loss": 0.0286,
+      "llm_learning_rate": 3.4491346066139258e-06,
+      "loss": 0.7462,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.00034491346066139257,
+      "step": 3853
+    },
+    {
+      "epoch": 0.39,
+      "gumbel_tau": 0.4029,
+      "image_masks_loss": 0.0292,
+      "llm_learning_rate": 3.448367052183754e-06,
+      "loss": 0.7691,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00034483670521837534,
+      "step": 3854
+    },
+    {
+      "epoch": 0.39,
+      "gumbel_tau": 0.4028,
+      "image_masks_loss": 0.0094,
+      "llm_learning_rate": 3.447599393317291e-06,
+      "loss": 0.7286,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.0003447599393317291,
+      "step": 3855
+    },
+    {
+      "epoch": 0.4,
+      "gumbel_tau": 0.4027,
+      "image_masks_loss": 0.0261,
+      "llm_learning_rate": 3.446831630099071e-06,
+      "loss": 0.8368,
+      "output_text_mask_loss": 0.0298,
+      "predictor_learning_rate": 0.0003446831630099071,
+      "step": 3856
+    },
+    {
+      "epoch": 0.4,
+      "gumbel_tau": 0.4026,
+      "image_masks_loss": 0.025,
+      "llm_learning_rate": 3.4460637626136427e-06,
+      "loss": 0.6642,
+      "output_text_mask_loss": 0.058,
+      "predictor_learning_rate": 0.00034460637626136426,
+      "step": 3857
+    },
+    {
+      "epoch": 0.4,
+      "gumbel_tau": 0.4025,
+      "image_masks_loss": 0.0165,
+      "llm_learning_rate": 3.4452957909455662e-06,
+      "loss": 0.8161,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 0.00034452957909455664,
+      "step": 3858
+    },
+    {
+      "epoch": 0.4,
+      "gumbel_tau": 0.4025,
+      "image_masks_loss": 0.0309,
+      "llm_learning_rate": 3.4445277151794122e-06,
+      "loss": 0.8605,
+      "output_text_mask_loss": 0.0847,
+      "predictor_learning_rate": 0.0003444527715179412,
+      "step": 3859
+    },
+    {
+      "epoch": 0.4,
+      "gumbel_tau": 0.4024,
+      "image_masks_loss": 0.0046,
+      "llm_learning_rate": 3.4437595353997613e-06,
+      "loss": 0.7246,
+      "output_text_mask_loss": 0.0317,
+      "predictor_learning_rate": 0.0003443759535399761,
+      "step": 3860
+    },
+    {
+      "epoch": 0.4,
+      "gumbel_tau": 0.4023,
+      "image_masks_loss": 0.0099,
+      "llm_learning_rate": 3.4429912516912086e-06,
+      "loss": 0.7442,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.00034429912516912084,
+      "step": 3861
+    },
+    {
+      "epoch": 0.4,
+      "gumbel_tau": 0.4022,
+      "image_masks_loss": 0.0233,
+      "llm_learning_rate": 3.442222864138359e-06,
+      "loss": 0.7708,
+      "output_text_mask_loss": 0.0042,
+      "predictor_learning_rate": 0.00034422228641383583,
+      "step": 3862
+    },
+    {
+      "epoch": 0.4,
+      "gumbel_tau": 0.4021,
+      "image_masks_loss": 0.0217,
+      "llm_learning_rate": 3.441454372825829e-06,
+      "loss": 0.7842,
+      "output_text_mask_loss": 0.0687,
+      "predictor_learning_rate": 0.00034414543728258283,
+      "step": 3863
+    },
+    {
+      "epoch": 0.4,
+      "gumbel_tau": 0.402,
+      "image_masks_loss": 0.0244,
+      "llm_learning_rate": 3.4406857778382457e-06,
+      "loss": 0.8318,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 0.00034406857778382453,
+      "step": 3864
+    },
+    {
+      "epoch": 0.4,
+      "gumbel_tau": 0.4019,
+      "image_masks_loss": 0.0062,
+      "llm_learning_rate": 3.4399170792602497e-06,
+      "loss": 0.8531,
+      "output_text_mask_loss": 0.0309,
+      "predictor_learning_rate": 0.00034399170792602496,
+      "step": 3865
+    },
+    {
+      "epoch": 0.4,
+      "gumbel_tau": 0.4018,
+      "image_masks_loss": 0.0265,
+      "llm_learning_rate": 3.4391482771764915e-06,
+      "loss": 0.8366,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 0.00034391482771764916,
+      "step": 3866
+    },
+    {
+      "epoch": 0.4,
+      "gumbel_tau": 0.4017,
+      "image_masks_loss": 0.0223,
+      "llm_learning_rate": 3.438379371671633e-06,
+      "loss": 0.7364,
+      "output_text_mask_loss": 0.0081,
+      "predictor_learning_rate": 0.00034383793716716326,
+      "step": 3867
+    },
+    {
+      "epoch": 0.4,
+      "gumbel_tau": 0.4016,
+      "image_masks_loss": 0.0176,
+      "llm_learning_rate": 3.4376103628303477e-06,
+      "loss": 0.7806,
+      "output_text_mask_loss": 0.0054,
+      "predictor_learning_rate": 0.0003437610362830347,
+      "step": 3868
+    },
+    {
+      "epoch": 0.4,
+      "gumbel_tau": 0.4015,
+      "image_masks_loss": 0.0157,
+      "llm_learning_rate": 3.436841250737321e-06,
+      "loss": 0.7208,
+      "output_text_mask_loss": 0.0012,
+      "predictor_learning_rate": 0.0003436841250737321,
+      "step": 3869
+    },
+    {
+      "epoch": 0.4,
+      "gumbel_tau": 0.4014,
+      "image_masks_loss": 0.0115,
+      "llm_learning_rate": 3.4360720354772488e-06,
+      "loss": 0.7903,
+      "output_text_mask_loss": 0.0137,
+      "predictor_learning_rate": 0.00034360720354772485,
+      "step": 3870
+    },
+    {
+      "epoch": 0.4,
+      "gumbel_tau": 0.4013,
+      "image_masks_loss": 0.0175,
+      "llm_learning_rate": 3.4353027171348396e-06,
+      "loss": 0.8456,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 0.00034353027171348395,
+      "step": 3871
+    },
+    {
+      "epoch": 0.4,
+      "gumbel_tau": 0.4012,
+      "image_masks_loss": 0.0256,
+      "llm_learning_rate": 3.4345332957948118e-06,
+      "loss": 0.718,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.0003434533295794812,
+      "step": 3872
+    },
+    {
+      "epoch": 0.4,
+      "gumbel_tau": 0.4011,
+      "image_masks_loss": 0.0162,
+      "llm_learning_rate": 3.4337637715418957e-06,
+      "loss": 0.6282,
+      "output_text_mask_loss": 0.0271,
+      "predictor_learning_rate": 0.00034337637715418955,
+      "step": 3873
+    },
+    {
+      "epoch": 0.4,
+      "gumbel_tau": 0.401,
+      "image_masks_loss": 0.0114,
+      "llm_learning_rate": 3.4329941444608333e-06,
+      "loss": 0.7268,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00034329941444608334,
+      "step": 3874
+    },
+    {
+      "epoch": 0.4,
+      "gumbel_tau": 0.4009,
+      "image_masks_loss": 0.0229,
+      "llm_learning_rate": 3.4322244146363777e-06,
+      "loss": 0.759,
+      "output_text_mask_loss": 0.0504,
+      "predictor_learning_rate": 0.00034322244146363774,
+      "step": 3875
+    },
+    {
+      "epoch": 0.4,
+      "gumbel_tau": 0.4008,
+      "image_masks_loss": 0.0027,
+      "llm_learning_rate": 3.431454582153293e-06,
+      "loss": 0.7581,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.0003431454582153293,
+      "step": 3876
+    },
+    {
+      "epoch": 0.4,
+      "gumbel_tau": 0.4007,
+      "image_masks_loss": 0.0064,
+      "llm_learning_rate": 3.430684647096354e-06,
+      "loss": 0.7871,
+      "output_text_mask_loss": 0.0687,
+      "predictor_learning_rate": 0.0003430684647096354,
+      "step": 3877
+    },
+    {
+      "epoch": 0.4,
+      "gumbel_tau": 0.4007,
+      "image_masks_loss": 0.0336,
+      "llm_learning_rate": 3.42991460955035e-06,
+      "loss": 0.9153,
+      "output_text_mask_loss": 0.0071,
+      "predictor_learning_rate": 0.00034299146095503495,
+      "step": 3878
+    },
+    {
+      "epoch": 0.4,
+      "gumbel_tau": 0.4006,
+      "image_masks_loss": 0.0066,
+      "llm_learning_rate": 3.429144469600077e-06,
+      "loss": 0.6789,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0003429144469600077,
+      "step": 3879
+    },
+    {
+      "epoch": 0.4,
+      "gumbel_tau": 0.4005,
+      "image_masks_loss": 0.016,
+      "llm_learning_rate": 3.4283742273303457e-06,
+      "loss": 0.7944,
+      "output_text_mask_loss": 0.0118,
+      "predictor_learning_rate": 0.00034283742273303453,
+      "step": 3880
+    },
+    {
+      "epoch": 0.4,
+      "gumbel_tau": 0.4004,
+      "image_masks_loss": 0.0181,
+      "llm_learning_rate": 3.4276038828259765e-06,
+      "loss": 0.8414,
+      "output_text_mask_loss": 0.0077,
+      "predictor_learning_rate": 0.00034276038828259765,
+      "step": 3881
+    },
+    {
+      "epoch": 0.4,
+      "gumbel_tau": 0.4003,
+      "image_masks_loss": 0.0116,
+      "llm_learning_rate": 3.426833436171802e-06,
+      "loss": 0.7449,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.00034268334361718016,
+      "step": 3882
+    },
+    {
+      "epoch": 0.4,
+      "gumbel_tau": 0.4002,
+      "image_masks_loss": 0.0143,
+      "llm_learning_rate": 3.426062887452665e-06,
+      "loss": 0.7068,
+      "output_text_mask_loss": 0.021,
+      "predictor_learning_rate": 0.00034260628874526643,
+      "step": 3883
+    },
+    {
+      "epoch": 0.4,
+      "gumbel_tau": 0.4001,
+      "image_masks_loss": 0.0202,
+      "llm_learning_rate": 3.4252922367534196e-06,
+      "loss": 0.7492,
+      "output_text_mask_loss": 0.006,
+      "predictor_learning_rate": 0.000342529223675342,
+      "step": 3884
+    },
+    {
+      "epoch": 0.4,
+      "gumbel_tau": 0.4,
+      "image_masks_loss": 0.0322,
+      "llm_learning_rate": 3.4245214841589326e-06,
+      "loss": 0.7758,
+      "output_text_mask_loss": 0.0645,
+      "predictor_learning_rate": 0.0003424521484158932,
+      "step": 3885
+    },
+    {
+      "epoch": 0.4,
+      "gumbel_tau": 0.3999,
+      "image_masks_loss": 0.01,
+      "llm_learning_rate": 3.4237506297540805e-06,
+      "loss": 0.6522,
+      "output_text_mask_loss": 0.0159,
+      "predictor_learning_rate": 0.00034237506297540805,
+      "step": 3886
+    },
+    {
+      "epoch": 0.4,
+      "gumbel_tau": 0.3998,
+      "image_masks_loss": 0.0042,
+      "llm_learning_rate": 3.4229796736237508e-06,
+      "loss": 0.7517,
+      "output_text_mask_loss": 0.0141,
+      "predictor_learning_rate": 0.0003422979673623751,
+      "step": 3887
+    },
+    {
+      "epoch": 0.4,
+      "gumbel_tau": 0.3997,
+      "image_masks_loss": 0.0051,
+      "llm_learning_rate": 3.4222086158528445e-06,
+      "loss": 0.7257,
+      "output_text_mask_loss": 0.0015,
+      "predictor_learning_rate": 0.0003422208615852844,
+      "step": 3888
+    },
+    {
+      "epoch": 0.4,
+      "gumbel_tau": 0.3996,
+      "image_masks_loss": 0.0132,
+      "llm_learning_rate": 3.4214374565262708e-06,
+      "loss": 0.6968,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.0003421437456526271,
+      "step": 3889
+    },
+    {
+      "epoch": 0.4,
+      "gumbel_tau": 0.3995,
+      "image_masks_loss": 0.0328,
+      "llm_learning_rate": 3.4206661957289515e-06,
+      "loss": 0.722,
+      "output_text_mask_loss": 0.0012,
+      "predictor_learning_rate": 0.00034206661957289513,
+      "step": 3890
+    },
+    {
+      "epoch": 0.4,
+      "gumbel_tau": 0.3994,
+      "image_masks_loss": 0.0141,
+      "llm_learning_rate": 3.4198948335458203e-06,
+      "loss": 0.7467,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.000341989483354582,
+      "step": 3891
+    },
+    {
+      "epoch": 0.4,
+      "gumbel_tau": 0.3993,
+      "image_masks_loss": 0.0107,
+      "llm_learning_rate": 3.419123370061822e-06,
+      "loss": 0.7329,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00034191233700618217,
+      "step": 3892
+    },
+    {
+      "epoch": 0.4,
+      "gumbel_tau": 0.3992,
+      "image_masks_loss": 0.0144,
+      "llm_learning_rate": 3.4183518053619102e-06,
+      "loss": 0.7801,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.000341835180536191,
+      "step": 3893
+    },
+    {
+      "epoch": 0.4,
+      "gumbel_tau": 0.3991,
+      "image_masks_loss": 0.0132,
+      "llm_learning_rate": 3.417580139531052e-06,
+      "loss": 0.7129,
+      "output_text_mask_loss": 0.0504,
+      "predictor_learning_rate": 0.00034175801395310517,
+      "step": 3894
+    },
+    {
+      "epoch": 0.4,
+      "gumbel_tau": 0.399,
+      "image_masks_loss": 0.0085,
+      "llm_learning_rate": 3.416808372654225e-06,
+      "loss": 0.8314,
+      "output_text_mask_loss": 0.0041,
+      "predictor_learning_rate": 0.0003416808372654225,
+      "step": 3895
+    },
+    {
+      "epoch": 0.4,
+      "gumbel_tau": 0.399,
+      "image_masks_loss": 0.0236,
+      "llm_learning_rate": 3.416036504816417e-06,
+      "loss": 0.6543,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0003416036504816417,
+      "step": 3896
+    },
+    {
+      "epoch": 0.4,
+      "gumbel_tau": 0.3989,
+      "image_masks_loss": 0.0124,
+      "llm_learning_rate": 3.4152645361026306e-06,
+      "loss": 0.9119,
+      "output_text_mask_loss": 0.0061,
+      "predictor_learning_rate": 0.000341526453610263,
+      "step": 3897
+    },
+    {
+      "epoch": 0.4,
+      "gumbel_tau": 0.3988,
+      "image_masks_loss": 0.0195,
+      "llm_learning_rate": 3.4144924665978736e-06,
+      "loss": 0.7172,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00034144924665978734,
+      "step": 3898
+    },
+    {
+      "epoch": 0.4,
+      "gumbel_tau": 0.3987,
+      "image_masks_loss": 0.0151,
+      "llm_learning_rate": 3.41372029638717e-06,
+      "loss": 0.8102,
+      "output_text_mask_loss": 0.0444,
+      "predictor_learning_rate": 0.00034137202963871693,
+      "step": 3899
+    },
+    {
+      "epoch": 0.4,
+      "gumbel_tau": 0.3986,
+      "image_masks_loss": 0.019,
+      "llm_learning_rate": 3.4129480255555507e-06,
+      "loss": 0.8218,
+      "output_text_mask_loss": 0.0071,
+      "predictor_learning_rate": 0.00034129480255555507,
+      "step": 3900
+    },
+    {
+      "epoch": 0.4,
+      "gumbel_tau": 0.3985,
+      "image_masks_loss": 0.0357,
+      "llm_learning_rate": 3.4121756541880626e-06,
+      "loss": 0.7238,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.0003412175654188062,
+      "step": 3901
+    },
+    {
+      "epoch": 0.4,
+      "gumbel_tau": 0.3984,
+      "image_masks_loss": 0.0152,
+      "llm_learning_rate": 3.4114031823697595e-06,
+      "loss": 0.762,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.0003411403182369759,
+      "step": 3902
+    },
+    {
+      "epoch": 0.4,
+      "gumbel_tau": 0.3983,
+      "image_masks_loss": 0.0381,
+      "llm_learning_rate": 3.410630610185707e-06,
+      "loss": 0.7542,
+      "output_text_mask_loss": 0.0039,
+      "predictor_learning_rate": 0.0003410630610185707,
+      "step": 3903
+    },
+    {
+      "epoch": 0.4,
+      "gumbel_tau": 0.3982,
+      "image_masks_loss": 0.0202,
+      "llm_learning_rate": 3.4098579377209846e-06,
+      "loss": 0.7471,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.00034098579377209843,
+      "step": 3904
+    },
+    {
+      "epoch": 0.4,
+      "gumbel_tau": 0.3981,
+      "image_masks_loss": 0.0261,
+      "llm_learning_rate": 3.409085165060679e-06,
+      "loss": 0.8302,
+      "output_text_mask_loss": 0.058,
+      "predictor_learning_rate": 0.00034090851650606785,
+      "step": 3905
+    },
+    {
+      "epoch": 0.4,
+      "gumbel_tau": 0.398,
+      "image_masks_loss": 0.0028,
+      "llm_learning_rate": 3.408312292289889e-06,
+      "loss": 0.7836,
+      "output_text_mask_loss": 0.005,
+      "predictor_learning_rate": 0.0003408312292289889,
+      "step": 3906
+    },
+    {
+      "epoch": 0.4,
+      "gumbel_tau": 0.3979,
+      "image_masks_loss": 0.0088,
+      "llm_learning_rate": 3.407539319493728e-06,
+      "loss": 0.7483,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00034075393194937277,
+      "step": 3907
+    },
+    {
+      "epoch": 0.4,
+      "gumbel_tau": 0.3978,
+      "image_masks_loss": 0.0033,
+      "llm_learning_rate": 3.4067662467573147e-06,
+      "loss": 0.8146,
+      "output_text_mask_loss": 0.0086,
+      "predictor_learning_rate": 0.0003406766246757315,
+      "step": 3908
+    },
+    {
+      "epoch": 0.4,
+      "gumbel_tau": 0.3977,
+      "image_masks_loss": 0.0185,
+      "llm_learning_rate": 3.405993074165783e-06,
+      "loss": 0.8586,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.0003405993074165783,
+      "step": 3909
+    },
+    {
+      "epoch": 0.4,
+      "gumbel_tau": 0.3976,
+      "image_masks_loss": 0.0195,
+      "llm_learning_rate": 3.4052198018042753e-06,
+      "loss": 0.7152,
+      "output_text_mask_loss": 0.0437,
+      "predictor_learning_rate": 0.00034052198018042753,
+      "step": 3910
+    },
+    {
+      "epoch": 0.4,
+      "gumbel_tau": 0.3975,
+      "image_masks_loss": 0.0172,
+      "llm_learning_rate": 3.4044464297579477e-06,
+      "loss": 0.8004,
+      "output_text_mask_loss": 0.0039,
+      "predictor_learning_rate": 0.00034044464297579475,
+      "step": 3911
+    },
+    {
+      "epoch": 0.4,
+      "gumbel_tau": 0.3975,
+      "image_masks_loss": 0.0186,
+      "llm_learning_rate": 3.403672958111964e-06,
+      "loss": 0.7941,
+      "output_text_mask_loss": 0.005,
+      "predictor_learning_rate": 0.00034036729581119634,
+      "step": 3912
+    },
+    {
+      "epoch": 0.4,
+      "gumbel_tau": 0.3974,
+      "image_masks_loss": 0.0082,
+      "llm_learning_rate": 3.402899386951501e-06,
+      "loss": 0.7606,
+      "output_text_mask_loss": 0.0576,
+      "predictor_learning_rate": 0.00034028993869515006,
+      "step": 3913
+    },
+    {
+      "epoch": 0.4,
+      "gumbel_tau": 0.3973,
+      "image_masks_loss": 0.0097,
+      "llm_learning_rate": 3.402125716361747e-06,
+      "loss": 0.7734,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.0003402125716361747,
+      "step": 3914
+    },
+    {
+      "epoch": 0.4,
+      "gumbel_tau": 0.3972,
+      "image_masks_loss": 0.0142,
+      "llm_learning_rate": 3.401351946427899e-06,
+      "loss": 0.7872,
+      "output_text_mask_loss": 0.0107,
+      "predictor_learning_rate": 0.00034013519464278993,
+      "step": 3915
+    },
+    {
+      "epoch": 0.4,
+      "gumbel_tau": 0.3971,
+      "image_masks_loss": 0.021,
+      "llm_learning_rate": 3.4005780772351678e-06,
+      "loss": 0.7494,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00034005780772351676,
+      "step": 3916
+    },
+    {
+      "epoch": 0.4,
+      "gumbel_tau": 0.397,
+      "image_masks_loss": 0.0148,
+      "llm_learning_rate": 3.399804108868772e-06,
+      "loss": 0.7669,
+      "output_text_mask_loss": 0.029,
+      "predictor_learning_rate": 0.0003399804108868772,
+      "step": 3917
+    },
+    {
+      "epoch": 0.4,
+      "gumbel_tau": 0.3969,
+      "image_masks_loss": 0.0203,
+      "llm_learning_rate": 3.399030041413944e-06,
+      "loss": 0.7826,
+      "output_text_mask_loss": 0.0193,
+      "predictor_learning_rate": 0.00033990300414139444,
+      "step": 3918
+    },
+    {
+      "epoch": 0.4,
+      "gumbel_tau": 0.3968,
+      "image_masks_loss": 0.021,
+      "llm_learning_rate": 3.3982558749559253e-06,
+      "loss": 0.8138,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.0003398255874955925,
+      "step": 3919
+    },
+    {
+      "epoch": 0.4,
+      "gumbel_tau": 0.3967,
+      "image_masks_loss": 0.0408,
+      "llm_learning_rate": 3.397481609579969e-06,
+      "loss": 0.7946,
+      "output_text_mask_loss": 0.1373,
+      "predictor_learning_rate": 0.0003397481609579969,
+      "step": 3920
+    },
+    {
+      "epoch": 0.4,
+      "gumbel_tau": 0.3966,
+      "image_masks_loss": 0.0193,
+      "llm_learning_rate": 3.396707245371339e-06,
+      "loss": 0.7262,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0003396707245371339,
+      "step": 3921
+    },
+    {
+      "epoch": 0.4,
+      "gumbel_tau": 0.3965,
+      "image_masks_loss": 0.0099,
+      "llm_learning_rate": 3.395932782415309e-06,
+      "loss": 0.7608,
+      "output_text_mask_loss": 0.0198,
+      "predictor_learning_rate": 0.00033959327824153085,
+      "step": 3922
+    },
+    {
+      "epoch": 0.4,
+      "gumbel_tau": 0.3964,
+      "image_masks_loss": 0.0121,
+      "llm_learning_rate": 3.3951582207971658e-06,
+      "loss": 0.8268,
+      "output_text_mask_loss": 0.1129,
+      "predictor_learning_rate": 0.00033951582207971656,
+      "step": 3923
+    },
+    {
+      "epoch": 0.4,
+      "gumbel_tau": 0.3963,
+      "image_masks_loss": 0.0298,
+      "llm_learning_rate": 3.394383560602206e-06,
+      "loss": 0.7666,
+      "output_text_mask_loss": 0.0058,
+      "predictor_learning_rate": 0.00033943835606022054,
+      "step": 3924
+    },
+    {
+      "epoch": 0.4,
+      "gumbel_tau": 0.3962,
+      "image_masks_loss": 0.0234,
+      "llm_learning_rate": 3.3936088019157354e-06,
+      "loss": 0.7395,
+      "output_text_mask_loss": 0.1404,
+      "predictor_learning_rate": 0.0003393608801915735,
+      "step": 3925
+    },
+    {
+      "epoch": 0.4,
+      "gumbel_tau": 0.3961,
+      "image_masks_loss": 0.012,
+      "llm_learning_rate": 3.3928339448230728e-06,
+      "loss": 0.6879,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0003392833944823073,
+      "step": 3926
+    },
+    {
+      "epoch": 0.4,
+      "gumbel_tau": 0.396,
+      "image_masks_loss": 0.0039,
+      "llm_learning_rate": 3.3920589894095484e-06,
+      "loss": 0.7598,
+      "output_text_mask_loss": 0.0576,
+      "predictor_learning_rate": 0.00033920589894095485,
+      "step": 3927
+    },
+    {
+      "epoch": 0.4,
+      "gumbel_tau": 0.396,
+      "image_masks_loss": 0.0183,
+      "llm_learning_rate": 3.391283935760501e-06,
+      "loss": 0.7344,
+      "output_text_mask_loss": 0.0039,
+      "predictor_learning_rate": 0.00033912839357605007,
+      "step": 3928
+    },
+    {
+      "epoch": 0.4,
+      "gumbel_tau": 0.3959,
+      "image_masks_loss": 0.0201,
+      "llm_learning_rate": 3.3905087839612818e-06,
+      "loss": 0.7825,
+      "output_text_mask_loss": 0.0065,
+      "predictor_learning_rate": 0.00033905087839612816,
+      "step": 3929
+    },
+    {
+      "epoch": 0.4,
+      "gumbel_tau": 0.3958,
+      "image_masks_loss": 0.0205,
+      "llm_learning_rate": 3.389733534097251e-06,
+      "loss": 0.7483,
+      "output_text_mask_loss": 0.0078,
+      "predictor_learning_rate": 0.00033897335340972505,
+      "step": 3930
+    },
+    {
+      "epoch": 0.4,
+      "gumbel_tau": 0.3957,
+      "image_masks_loss": 0.0119,
+      "llm_learning_rate": 3.388958186253782e-06,
+      "loss": 0.802,
+      "output_text_mask_loss": 0.0004,
+      "predictor_learning_rate": 0.0003388958186253782,
+      "step": 3931
+    },
+    {
+      "epoch": 0.4,
+      "gumbel_tau": 0.3956,
+      "image_masks_loss": 0.0241,
+      "llm_learning_rate": 3.388182740516257e-06,
+      "loss": 0.6775,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 0.0003388182740516257,
+      "step": 3932
+    },
+    {
+      "epoch": 0.4,
+      "gumbel_tau": 0.3955,
+      "image_masks_loss": 0.0105,
+      "llm_learning_rate": 3.3874071969700706e-06,
+      "loss": 0.8027,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 0.00033874071969700704,
+      "step": 3933
+    },
+    {
+      "epoch": 0.4,
+      "gumbel_tau": 0.3954,
+      "image_masks_loss": 0.0271,
+      "llm_learning_rate": 3.386631555700627e-06,
+      "loss": 0.7364,
+      "output_text_mask_loss": 0.0271,
+      "predictor_learning_rate": 0.00033866315557006264,
+      "step": 3934
+    },
+    {
+      "epoch": 0.4,
+      "gumbel_tau": 0.3953,
+      "image_masks_loss": 0.0113,
+      "llm_learning_rate": 3.3858558167933407e-06,
+      "loss": 0.7711,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 0.00033858558167933407,
+      "step": 3935
+    },
+    {
+      "epoch": 0.4,
+      "gumbel_tau": 0.3952,
+      "image_masks_loss": 0.0154,
+      "llm_learning_rate": 3.3850799803336397e-06,
+      "loss": 0.8036,
+      "output_text_mask_loss": 0.0618,
+      "predictor_learning_rate": 0.00033850799803336395,
+      "step": 3936
+    },
+    {
+      "epoch": 0.4,
+      "gumbel_tau": 0.3951,
+      "image_masks_loss": 0.0435,
+      "llm_learning_rate": 3.38430404640696e-06,
+      "loss": 0.813,
+      "output_text_mask_loss": 0.0444,
+      "predictor_learning_rate": 0.000338430404640696,
+      "step": 3937
+    },
+    {
+      "epoch": 0.4,
+      "gumbel_tau": 0.395,
+      "image_masks_loss": 0.0278,
+      "llm_learning_rate": 3.3835280150987488e-06,
+      "loss": 0.7831,
+      "output_text_mask_loss": 0.0137,
+      "predictor_learning_rate": 0.00033835280150987484,
+      "step": 3938
+    },
+    {
+      "epoch": 0.4,
+      "gumbel_tau": 0.3949,
+      "image_masks_loss": 0.0425,
+      "llm_learning_rate": 3.3827518864944646e-06,
+      "loss": 0.7381,
+      "output_text_mask_loss": 0.0004,
+      "predictor_learning_rate": 0.00033827518864944646,
+      "step": 3939
+    },
+    {
+      "epoch": 0.4,
+      "gumbel_tau": 0.3948,
+      "image_masks_loss": 0.0133,
+      "llm_learning_rate": 3.381975660679576e-06,
+      "loss": 0.7294,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.0003381975660679576,
+      "step": 3940
+    },
+    {
+      "epoch": 0.4,
+      "gumbel_tau": 0.3947,
+      "image_masks_loss": 0.0195,
+      "llm_learning_rate": 3.381199337739564e-06,
+      "loss": 0.6871,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.0003381199337739564,
+      "step": 3941
+    },
+    {
+      "epoch": 0.4,
+      "gumbel_tau": 0.3946,
+      "image_masks_loss": 0.0186,
+      "llm_learning_rate": 3.3804229177599168e-06,
+      "loss": 0.7441,
+      "output_text_mask_loss": 0.0071,
+      "predictor_learning_rate": 0.00033804229177599165,
+      "step": 3942
+    },
+    {
+      "epoch": 0.4,
+      "gumbel_tau": 0.3946,
+      "image_masks_loss": 0.0185,
+      "llm_learning_rate": 3.3796464008261375e-06,
+      "loss": 0.7543,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.00033796464008261376,
+      "step": 3943
+    },
+    {
+      "epoch": 0.4,
+      "gumbel_tau": 0.3945,
+      "image_masks_loss": 0.0332,
+      "llm_learning_rate": 3.3788697870237374e-06,
+      "loss": 0.722,
+      "output_text_mask_loss": 0.0193,
+      "predictor_learning_rate": 0.0003378869787023737,
+      "step": 3944
+    },
+    {
+      "epoch": 0.4,
+      "gumbel_tau": 0.3944,
+      "image_masks_loss": 0.034,
+      "llm_learning_rate": 3.3780930764382385e-06,
+      "loss": 0.7747,
+      "output_text_mask_loss": 0.0194,
+      "predictor_learning_rate": 0.00033780930764382383,
+      "step": 3945
+    },
+    {
+      "epoch": 0.4,
+      "gumbel_tau": 0.3943,
+      "image_masks_loss": 0.0126,
+      "llm_learning_rate": 3.377316269155174e-06,
+      "loss": 0.749,
+      "output_text_mask_loss": 0.006,
+      "predictor_learning_rate": 0.0003377316269155174,
+      "step": 3946
+    },
+    {
+      "epoch": 0.4,
+      "gumbel_tau": 0.3942,
+      "image_masks_loss": 0.0124,
+      "llm_learning_rate": 3.376539365260089e-06,
+      "loss": 0.7873,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.00033765393652600886,
+      "step": 3947
+    },
+    {
+      "epoch": 0.4,
+      "gumbel_tau": 0.3941,
+      "image_masks_loss": 0.0061,
+      "llm_learning_rate": 3.3757623648385353e-06,
+      "loss": 0.8444,
+      "output_text_mask_loss": 0.0286,
+      "predictor_learning_rate": 0.0003375762364838535,
+      "step": 3948
+    },
+    {
+      "epoch": 0.4,
+      "gumbel_tau": 0.394,
+      "image_masks_loss": 0.0049,
+      "llm_learning_rate": 3.3749852679760796e-06,
+      "loss": 0.7814,
+      "output_text_mask_loss": 0.0238,
+      "predictor_learning_rate": 0.00033749852679760796,
+      "step": 3949
+    },
+    {
+      "epoch": 0.4,
+      "gumbel_tau": 0.3939,
+      "image_masks_loss": 0.0252,
+      "llm_learning_rate": 3.374208074758298e-06,
+      "loss": 0.8011,
+      "output_text_mask_loss": 0.0301,
+      "predictor_learning_rate": 0.00033742080747582975,
+      "step": 3950
+    },
+    {
+      "epoch": 0.4,
+      "gumbel_tau": 0.3938,
+      "image_masks_loss": 0.0288,
+      "llm_learning_rate": 3.3734307852707748e-06,
+      "loss": 0.8915,
+      "output_text_mask_loss": 0.0145,
+      "predictor_learning_rate": 0.00033734307852707746,
+      "step": 3951
+    },
+    {
+      "epoch": 0.4,
+      "gumbel_tau": 0.3937,
+      "image_masks_loss": 0.017,
+      "llm_learning_rate": 3.3726533995991085e-06,
+      "loss": 0.7105,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0003372653399599108,
+      "step": 3952
+    },
+    {
+      "epoch": 0.41,
+      "gumbel_tau": 0.3936,
+      "image_masks_loss": 0.0149,
+      "llm_learning_rate": 3.371875917828906e-06,
+      "loss": 0.7791,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0003371875917828906,
+      "step": 3953
+    },
+    {
+      "epoch": 0.41,
+      "gumbel_tau": 0.3935,
+      "image_masks_loss": 0.0467,
+      "llm_learning_rate": 3.3710983400457865e-06,
+      "loss": 0.8051,
+      "output_text_mask_loss": 0.0422,
+      "predictor_learning_rate": 0.0003371098340045786,
+      "step": 3954
+    },
+    {
+      "epoch": 0.41,
+      "gumbel_tau": 0.3934,
+      "image_masks_loss": 0.0094,
+      "llm_learning_rate": 3.370320666335376e-06,
+      "loss": 0.7149,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.0003370320666335376,
+      "step": 3955
+    },
+    {
+      "epoch": 0.41,
+      "gumbel_tau": 0.3933,
+      "image_masks_loss": 0.0164,
+      "llm_learning_rate": 3.3695428967833156e-06,
+      "loss": 0.8157,
+      "output_text_mask_loss": 0.0162,
+      "predictor_learning_rate": 0.00033695428967833155,
+      "step": 3956
+    },
+    {
+      "epoch": 0.41,
+      "gumbel_tau": 0.3933,
+      "image_masks_loss": 0.0127,
+      "llm_learning_rate": 3.368765031475255e-06,
+      "loss": 0.8274,
+      "output_text_mask_loss": 0.0085,
+      "predictor_learning_rate": 0.0003368765031475255,
+      "step": 3957
+    },
+    {
+      "epoch": 0.41,
+      "gumbel_tau": 0.3932,
+      "image_masks_loss": 0.0058,
+      "llm_learning_rate": 3.3679870704968535e-06,
+      "loss": 0.8209,
+      "output_text_mask_loss": 0.0261,
+      "predictor_learning_rate": 0.00033679870704968537,
+      "step": 3958
+    },
+    {
+      "epoch": 0.41,
+      "gumbel_tau": 0.3931,
+      "image_masks_loss": 0.0123,
+      "llm_learning_rate": 3.367209013933783e-06,
+      "loss": 0.7306,
+      "output_text_mask_loss": 0.2319,
+      "predictor_learning_rate": 0.0003367209013933783,
+      "step": 3959
+    },
+    {
+      "epoch": 0.41,
+      "gumbel_tau": 0.393,
+      "image_masks_loss": 0.0078,
+      "llm_learning_rate": 3.366430861871724e-06,
+      "loss": 0.8033,
+      "output_text_mask_loss": 0.0038,
+      "predictor_learning_rate": 0.0003366430861871724,
+      "step": 3960
+    },
+    {
+      "epoch": 0.41,
+      "gumbel_tau": 0.3929,
+      "image_masks_loss": 0.0507,
+      "llm_learning_rate": 3.3656526143963687e-06,
+      "loss": 0.7341,
+      "output_text_mask_loss": 0.0039,
+      "predictor_learning_rate": 0.00033656526143963685,
+      "step": 3961
+    },
+    {
+      "epoch": 0.41,
+      "gumbel_tau": 0.3928,
+      "image_masks_loss": 0.0123,
+      "llm_learning_rate": 3.364874271593419e-06,
+      "loss": 0.8566,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.00033648742715934187,
+      "step": 3962
+    },
+    {
+      "epoch": 0.41,
+      "gumbel_tau": 0.3927,
+      "image_masks_loss": 0.0081,
+      "llm_learning_rate": 3.3640958335485887e-06,
+      "loss": 0.6772,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00033640958335485885,
+      "step": 3963
+    },
+    {
+      "epoch": 0.41,
+      "gumbel_tau": 0.3926,
+      "image_masks_loss": 0.0319,
+      "llm_learning_rate": 3.3633173003475996e-06,
+      "loss": 0.6831,
+      "output_text_mask_loss": 0.0217,
+      "predictor_learning_rate": 0.0003363317300347599,
+      "step": 3964
+    },
+    {
+      "epoch": 0.41,
+      "gumbel_tau": 0.3925,
+      "image_masks_loss": 0.0298,
+      "llm_learning_rate": 3.3625386720761865e-06,
+      "loss": 0.7874,
+      "output_text_mask_loss": 0.001,
+      "predictor_learning_rate": 0.00033625386720761866,
+      "step": 3965
+    },
+    {
+      "epoch": 0.41,
+      "gumbel_tau": 0.3924,
+      "image_masks_loss": 0.0152,
+      "llm_learning_rate": 3.361759948820095e-06,
+      "loss": 0.8712,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 0.00033617599488200943,
+      "step": 3966
+    },
+    {
+      "epoch": 0.41,
+      "gumbel_tau": 0.3923,
+      "image_masks_loss": 0.0181,
+      "llm_learning_rate": 3.3609811306650763e-06,
+      "loss": 0.7461,
+      "output_text_mask_loss": 0.0231,
+      "predictor_learning_rate": 0.0003360981130665076,
+      "step": 3967
+    },
+    {
+      "epoch": 0.41,
+      "gumbel_tau": 0.3922,
+      "image_masks_loss": 0.0359,
+      "llm_learning_rate": 3.360202217696899e-06,
+      "loss": 0.7724,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0003360202217696899,
+      "step": 3968
+    },
+    {
+      "epoch": 0.41,
+      "gumbel_tau": 0.3921,
+      "image_masks_loss": 0.014,
+      "llm_learning_rate": 3.359423210001338e-06,
+      "loss": 0.7091,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0003359423210001338,
+      "step": 3969
+    },
+    {
+      "epoch": 0.41,
+      "gumbel_tau": 0.392,
+      "image_masks_loss": 0.0199,
+      "llm_learning_rate": 3.358644107664178e-06,
+      "loss": 0.7931,
+      "output_text_mask_loss": 0.0071,
+      "predictor_learning_rate": 0.0003358644107664178,
+      "step": 3970
+    },
+    {
+      "epoch": 0.41,
+      "gumbel_tau": 0.392,
+      "image_masks_loss": 0.0299,
+      "llm_learning_rate": 3.357864910771216e-06,
+      "loss": 0.7271,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00033578649107712154,
+      "step": 3971
+    },
+    {
+      "epoch": 0.41,
+      "gumbel_tau": 0.3919,
+      "image_masks_loss": 0.019,
+      "llm_learning_rate": 3.3570856194082598e-06,
+      "loss": 0.7573,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.00033570856194082595,
+      "step": 3972
+    },
+    {
+      "epoch": 0.41,
+      "gumbel_tau": 0.3918,
+      "image_masks_loss": 0.0119,
+      "llm_learning_rate": 3.3563062336611257e-06,
+      "loss": 0.6972,
+      "output_text_mask_loss": 0.071,
+      "predictor_learning_rate": 0.00033563062336611254,
+      "step": 3973
+    },
+    {
+      "epoch": 0.41,
+      "gumbel_tau": 0.3917,
+      "image_masks_loss": 0.0196,
+      "llm_learning_rate": 3.355526753615642e-06,
+      "loss": 0.7853,
+      "output_text_mask_loss": 0.0534,
+      "predictor_learning_rate": 0.0003355526753615642,
+      "step": 3974
+    },
+    {
+      "epoch": 0.41,
+      "gumbel_tau": 0.3916,
+      "image_masks_loss": 0.0154,
+      "llm_learning_rate": 3.354747179357647e-06,
+      "loss": 0.8101,
+      "output_text_mask_loss": 0.0771,
+      "predictor_learning_rate": 0.00033547471793576465,
+      "step": 3975
+    },
+    {
+      "epoch": 0.41,
+      "gumbel_tau": 0.3915,
+      "image_masks_loss": 0.0246,
+      "llm_learning_rate": 3.3539675109729885e-06,
+      "loss": 0.732,
+      "output_text_mask_loss": 0.0278,
+      "predictor_learning_rate": 0.00033539675109729884,
+      "step": 3976
+    },
+    {
+      "epoch": 0.41,
+      "gumbel_tau": 0.3914,
+      "image_masks_loss": 0.0385,
+      "llm_learning_rate": 3.353187748547525e-06,
+      "loss": 0.7016,
+      "output_text_mask_loss": 0.0717,
+      "predictor_learning_rate": 0.0003353187748547525,
+      "step": 3977
+    },
+    {
+      "epoch": 0.41,
+      "gumbel_tau": 0.3913,
+      "image_masks_loss": 0.0104,
+      "llm_learning_rate": 3.3524078921671265e-06,
+      "loss": 0.7882,
+      "output_text_mask_loss": 0.0116,
+      "predictor_learning_rate": 0.00033524078921671265,
+      "step": 3978
+    },
+    {
+      "epoch": 0.41,
+      "gumbel_tau": 0.3912,
+      "image_masks_loss": 0.0188,
+      "llm_learning_rate": 3.351627941917673e-06,
+      "loss": 0.8265,
+      "output_text_mask_loss": 0.0023,
+      "predictor_learning_rate": 0.00033516279419176727,
+      "step": 3979
+    },
+    {
+      "epoch": 0.41,
+      "gumbel_tau": 0.3911,
+      "image_masks_loss": 0.0164,
+      "llm_learning_rate": 3.3508478978850526e-06,
+      "loss": 0.7446,
+      "output_text_mask_loss": 0.0376,
+      "predictor_learning_rate": 0.00033508478978850524,
+      "step": 3980
+    },
+    {
+      "epoch": 0.41,
+      "gumbel_tau": 0.391,
+      "image_masks_loss": 0.0093,
+      "llm_learning_rate": 3.350067760155167e-06,
+      "loss": 0.7382,
+      "output_text_mask_loss": 0.0156,
+      "predictor_learning_rate": 0.00033500677601551665,
+      "step": 3981
+    },
+    {
+      "epoch": 0.41,
+      "gumbel_tau": 0.3909,
+      "image_masks_loss": 0.0212,
+      "llm_learning_rate": 3.3492875288139257e-06,
+      "loss": 0.702,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.00033492875288139255,
+      "step": 3982
+    },
+    {
+      "epoch": 0.41,
+      "gumbel_tau": 0.3908,
+      "image_masks_loss": 0.0121,
+      "llm_learning_rate": 3.348507203947251e-06,
+      "loss": 0.7996,
+      "output_text_mask_loss": 0.0732,
+      "predictor_learning_rate": 0.0003348507203947251,
+      "step": 3983
+    },
+    {
+      "epoch": 0.41,
+      "gumbel_tau": 0.3908,
+      "image_masks_loss": 0.0093,
+      "llm_learning_rate": 3.347726785641072e-06,
+      "loss": 0.7312,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.0003347726785641072,
+      "step": 3984
+    },
+    {
+      "epoch": 0.41,
+      "gumbel_tau": 0.3907,
+      "image_masks_loss": 0.0061,
+      "llm_learning_rate": 3.346946273981332e-06,
+      "loss": 0.7196,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.0003346946273981332,
+      "step": 3985
+    },
+    {
+      "epoch": 0.41,
+      "gumbel_tau": 0.3906,
+      "image_masks_loss": 0.0168,
+      "llm_learning_rate": 3.346165669053982e-06,
+      "loss": 0.8441,
+      "output_text_mask_loss": 0.0402,
+      "predictor_learning_rate": 0.00033461656690539815,
+      "step": 3986
+    },
+    {
+      "epoch": 0.41,
+      "gumbel_tau": 0.3905,
+      "image_masks_loss": 0.0178,
+      "llm_learning_rate": 3.3453849709449827e-06,
+      "loss": 0.802,
+      "output_text_mask_loss": 0.0116,
+      "predictor_learning_rate": 0.0003345384970944982,
+      "step": 3987
+    },
+    {
+      "epoch": 0.41,
+      "gumbel_tau": 0.3904,
+      "image_masks_loss": 0.0357,
+      "llm_learning_rate": 3.344604179740308e-06,
+      "loss": 0.756,
+      "output_text_mask_loss": 0.021,
+      "predictor_learning_rate": 0.0003344604179740308,
+      "step": 3988
+    },
+    {
+      "epoch": 0.41,
+      "gumbel_tau": 0.3903,
+      "image_masks_loss": 0.0209,
+      "llm_learning_rate": 3.3438232955259392e-06,
+      "loss": 0.7403,
+      "output_text_mask_loss": 0.0081,
+      "predictor_learning_rate": 0.0003343823295525939,
+      "step": 3989
+    },
+    {
+      "epoch": 0.41,
+      "gumbel_tau": 0.3902,
+      "image_masks_loss": 0.0158,
+      "llm_learning_rate": 3.3430423183878704e-06,
+      "loss": 0.7762,
+      "output_text_mask_loss": 0.0294,
+      "predictor_learning_rate": 0.000334304231838787,
+      "step": 3990
+    },
+    {
+      "epoch": 0.41,
+      "gumbel_tau": 0.3901,
+      "image_masks_loss": 0.0173,
+      "llm_learning_rate": 3.3422612484121025e-06,
+      "loss": 0.7875,
+      "output_text_mask_loss": 0.0137,
+      "predictor_learning_rate": 0.00033422612484121025,
+      "step": 3991
+    },
+    {
+      "epoch": 0.41,
+      "gumbel_tau": 0.39,
+      "image_masks_loss": 0.0184,
+      "llm_learning_rate": 3.3414800856846506e-06,
+      "loss": 0.7584,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.00033414800856846504,
+      "step": 3992
+    },
+    {
+      "epoch": 0.41,
+      "gumbel_tau": 0.3899,
+      "image_masks_loss": 0.0194,
+      "llm_learning_rate": 3.3406988302915376e-06,
+      "loss": 0.8329,
+      "output_text_mask_loss": 0.0322,
+      "predictor_learning_rate": 0.00033406988302915375,
+      "step": 3993
+    },
+    {
+      "epoch": 0.41,
+      "gumbel_tau": 0.3898,
+      "image_masks_loss": 0.0257,
+      "llm_learning_rate": 3.339917482318797e-06,
+      "loss": 0.7787,
+      "output_text_mask_loss": 0.0092,
+      "predictor_learning_rate": 0.0003339917482318797,
+      "step": 3994
+    },
+    {
+      "epoch": 0.41,
+      "gumbel_tau": 0.3897,
+      "image_masks_loss": 0.0151,
+      "llm_learning_rate": 3.3391360418524714e-06,
+      "loss": 0.759,
+      "output_text_mask_loss": 0.0019,
+      "predictor_learning_rate": 0.0003339136041852471,
+      "step": 3995
+    },
+    {
+      "epoch": 0.41,
+      "gumbel_tau": 0.3897,
+      "image_masks_loss": 0.0241,
+      "llm_learning_rate": 3.338354508978617e-06,
+      "loss": 0.7123,
+      "output_text_mask_loss": 0.0488,
+      "predictor_learning_rate": 0.0003338354508978617,
+      "step": 3996
+    },
+    {
+      "epoch": 0.41,
+      "gumbel_tau": 0.3896,
+      "image_masks_loss": 0.0244,
+      "llm_learning_rate": 3.3375728837832954e-06,
+      "loss": 0.7752,
+      "output_text_mask_loss": 0.0744,
+      "predictor_learning_rate": 0.0003337572883783295,
+      "step": 3997
+    },
+    {
+      "epoch": 0.41,
+      "gumbel_tau": 0.3895,
+      "image_masks_loss": 0.0277,
+      "llm_learning_rate": 3.3367911663525833e-06,
+      "loss": 0.7113,
+      "output_text_mask_loss": 0.0254,
+      "predictor_learning_rate": 0.0003336791166352583,
+      "step": 3998
+    },
+    {
+      "epoch": 0.41,
+      "gumbel_tau": 0.3894,
+      "image_masks_loss": 0.021,
+      "llm_learning_rate": 3.336009356772563e-06,
+      "loss": 0.844,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.0003336009356772563,
+      "step": 3999
+    },
+    {
+      "epoch": 0.41,
+      "gumbel_tau": 0.3893,
+      "image_masks_loss": 0.0161,
+      "llm_learning_rate": 3.3352274551293305e-06,
+      "loss": 0.7744,
+      "output_text_mask_loss": 0.0012,
+      "predictor_learning_rate": 0.00033352274551293307,
+      "step": 4000
+    },
+    {
+      "epoch": 0.41,
+      "gumbel_tau": 0.3892,
+      "image_masks_loss": 0.0188,
+      "llm_learning_rate": 3.3344454615089913e-06,
+      "loss": 0.7844,
+      "output_text_mask_loss": 0.0315,
+      "predictor_learning_rate": 0.0003334445461508991,
+      "step": 4001
+    },
+    {
+      "epoch": 0.41,
+      "gumbel_tau": 0.3891,
+      "image_masks_loss": 0.038,
+      "llm_learning_rate": 3.3336633759976585e-06,
+      "loss": 0.7135,
+      "output_text_mask_loss": 0.0299,
+      "predictor_learning_rate": 0.00033336633759976583,
+      "step": 4002
+    },
+    {
+      "epoch": 0.41,
+      "gumbel_tau": 0.389,
+      "image_masks_loss": 0.0448,
+      "llm_learning_rate": 3.332881198681458e-06,
+      "loss": 0.7091,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0003332881198681458,
+      "step": 4003
+    },
+    {
+      "epoch": 0.41,
+      "gumbel_tau": 0.3889,
+      "image_masks_loss": 0.0094,
+      "llm_learning_rate": 3.3320989296465246e-06,
+      "loss": 0.7904,
+      "output_text_mask_loss": 0.0454,
+      "predictor_learning_rate": 0.0003332098929646524,
+      "step": 4004
+    },
+    {
+      "epoch": 0.41,
+      "gumbel_tau": 0.3888,
+      "image_masks_loss": 0.0237,
+      "llm_learning_rate": 3.3313165689790034e-06,
+      "loss": 0.7108,
+      "output_text_mask_loss": 0.0164,
+      "predictor_learning_rate": 0.0003331316568979003,
+      "step": 4005
+    },
+    {
+      "epoch": 0.41,
+      "gumbel_tau": 0.3887,
+      "image_masks_loss": 0.0273,
+      "llm_learning_rate": 3.33053411676505e-06,
+      "loss": 0.7751,
+      "output_text_mask_loss": 0.036,
+      "predictor_learning_rate": 0.000333053411676505,
+      "step": 4006
+    },
+    {
+      "epoch": 0.41,
+      "gumbel_tau": 0.3886,
+      "image_masks_loss": 0.0183,
+      "llm_learning_rate": 3.3297515730908297e-06,
+      "loss": 0.7153,
+      "output_text_mask_loss": 0.0179,
+      "predictor_learning_rate": 0.00033297515730908297,
+      "step": 4007
+    },
+    {
+      "epoch": 0.41,
+      "gumbel_tau": 0.3885,
+      "image_masks_loss": 0.0051,
+      "llm_learning_rate": 3.3289689380425185e-06,
+      "loss": 0.8115,
+      "output_text_mask_loss": 0.006,
+      "predictor_learning_rate": 0.0003328968938042518,
+      "step": 4008
+    },
+    {
+      "epoch": 0.41,
+      "gumbel_tau": 0.3885,
+      "image_masks_loss": 0.0142,
+      "llm_learning_rate": 3.3281862117063008e-06,
+      "loss": 0.705,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00033281862117063004,
+      "step": 4009
+    },
+    {
+      "epoch": 0.41,
+      "gumbel_tau": 0.3884,
+      "image_masks_loss": 0.042,
+      "llm_learning_rate": 3.3274033941683727e-06,
+      "loss": 0.699,
+      "output_text_mask_loss": 0.005,
+      "predictor_learning_rate": 0.00033274033941683724,
+      "step": 4010
+    },
+    {
+      "epoch": 0.41,
+      "gumbel_tau": 0.3883,
+      "image_masks_loss": 0.0213,
+      "llm_learning_rate": 3.3266204855149403e-06,
+      "loss": 0.7126,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.000332662048551494,
+      "step": 4011
+    },
+    {
+      "epoch": 0.41,
+      "gumbel_tau": 0.3882,
+      "image_masks_loss": 0.0108,
+      "llm_learning_rate": 3.3258374858322183e-06,
+      "loss": 0.8065,
+      "output_text_mask_loss": 0.0641,
+      "predictor_learning_rate": 0.00033258374858322184,
+      "step": 4012
+    },
+    {
+      "epoch": 0.41,
+      "gumbel_tau": 0.3881,
+      "image_masks_loss": 0.0238,
+      "llm_learning_rate": 3.3250543952064334e-06,
+      "loss": 0.7853,
+      "output_text_mask_loss": 0.032,
+      "predictor_learning_rate": 0.0003325054395206433,
+      "step": 4013
+    },
+    {
+      "epoch": 0.41,
+      "gumbel_tau": 0.388,
+      "image_masks_loss": 0.0406,
+      "llm_learning_rate": 3.3242712137238207e-06,
+      "loss": 0.6075,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0003324271213723821,
+      "step": 4014
+    },
+    {
+      "epoch": 0.41,
+      "gumbel_tau": 0.3879,
+      "image_masks_loss": 0.0103,
+      "llm_learning_rate": 3.323487941470626e-06,
+      "loss": 0.7898,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.0003323487941470626,
+      "step": 4015
+    },
+    {
+      "epoch": 0.41,
+      "gumbel_tau": 0.3878,
+      "image_masks_loss": 0.0097,
+      "llm_learning_rate": 3.3227045785331046e-06,
+      "loss": 0.7321,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.00033227045785331046,
+      "step": 4016
+    },
+    {
+      "epoch": 0.41,
+      "gumbel_tau": 0.3877,
+      "image_masks_loss": 0.0196,
+      "llm_learning_rate": 3.3219211249975226e-06,
+      "loss": 0.7888,
+      "output_text_mask_loss": 0.006,
+      "predictor_learning_rate": 0.0003321921124997522,
+      "step": 4017
+    },
+    {
+      "epoch": 0.41,
+      "gumbel_tau": 0.3876,
+      "image_masks_loss": 0.0252,
+      "llm_learning_rate": 3.3211375809501556e-06,
+      "loss": 0.8208,
+      "output_text_mask_loss": 0.0204,
+      "predictor_learning_rate": 0.00033211375809501557,
+      "step": 4018
+    },
+    {
+      "epoch": 0.41,
+      "gumbel_tau": 0.3875,
+      "image_masks_loss": 0.0328,
+      "llm_learning_rate": 3.3203539464772895e-06,
+      "loss": 0.7738,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.00033203539464772894,
+      "step": 4019
+    },
+    {
+      "epoch": 0.41,
+      "gumbel_tau": 0.3875,
+      "image_masks_loss": 0.0257,
+      "llm_learning_rate": 3.319570221665219e-06,
+      "loss": 0.6955,
+      "output_text_mask_loss": 0.0429,
+      "predictor_learning_rate": 0.00033195702216652184,
+      "step": 4020
+    },
+    {
+      "epoch": 0.41,
+      "gumbel_tau": 0.3874,
+      "image_masks_loss": 0.0083,
+      "llm_learning_rate": 3.3187864066002507e-06,
+      "loss": 0.7739,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00033187864066002505,
+      "step": 4021
+    },
+    {
+      "epoch": 0.41,
+      "gumbel_tau": 0.3873,
+      "image_masks_loss": 0.0254,
+      "llm_learning_rate": 3.3180025013686994e-06,
+      "loss": 0.8789,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.0003318002501368699,
+      "step": 4022
+    },
+    {
+      "epoch": 0.41,
+      "gumbel_tau": 0.3872,
+      "image_masks_loss": 0.0114,
+      "llm_learning_rate": 3.3172185060568903e-06,
+      "loss": 0.782,
+      "output_text_mask_loss": 0.0216,
+      "predictor_learning_rate": 0.000331721850605689,
+      "step": 4023
+    },
+    {
+      "epoch": 0.41,
+      "gumbel_tau": 0.3871,
+      "image_masks_loss": 0.0156,
+      "llm_learning_rate": 3.316434420751159e-06,
+      "loss": 0.8132,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.0003316434420751159,
+      "step": 4024
+    },
+    {
+      "epoch": 0.41,
+      "gumbel_tau": 0.387,
+      "image_masks_loss": 0.0192,
+      "llm_learning_rate": 3.3156502455378512e-06,
+      "loss": 0.8241,
+      "output_text_mask_loss": 0.0103,
+      "predictor_learning_rate": 0.00033156502455378514,
+      "step": 4025
+    },
+    {
+      "epoch": 0.41,
+      "gumbel_tau": 0.3869,
+      "image_masks_loss": 0.022,
+      "llm_learning_rate": 3.314865980503321e-06,
+      "loss": 0.8327,
+      "output_text_mask_loss": 0.033,
+      "predictor_learning_rate": 0.0003314865980503321,
+      "step": 4026
+    },
+    {
+      "epoch": 0.41,
+      "gumbel_tau": 0.3868,
+      "image_masks_loss": 0.023,
+      "llm_learning_rate": 3.314081625733935e-06,
+      "loss": 0.7222,
+      "output_text_mask_loss": 0.0134,
+      "predictor_learning_rate": 0.00033140816257339344,
+      "step": 4027
+    },
+    {
+      "epoch": 0.41,
+      "gumbel_tau": 0.3867,
+      "image_masks_loss": 0.0092,
+      "llm_learning_rate": 3.3132971813160656e-06,
+      "loss": 0.7795,
+      "output_text_mask_loss": 0.0004,
+      "predictor_learning_rate": 0.0003313297181316066,
+      "step": 4028
+    },
+    {
+      "epoch": 0.41,
+      "gumbel_tau": 0.3866,
+      "image_masks_loss": 0.0122,
+      "llm_learning_rate": 3.3125126473361004e-06,
+      "loss": 0.7676,
+      "output_text_mask_loss": 0.0431,
+      "predictor_learning_rate": 0.00033125126473361,
+      "step": 4029
+    },
+    {
+      "epoch": 0.41,
+      "gumbel_tau": 0.3865,
+      "image_masks_loss": 0.0241,
+      "llm_learning_rate": 3.3117280238804316e-06,
+      "loss": 0.701,
+      "output_text_mask_loss": 0.0081,
+      "predictor_learning_rate": 0.00033117280238804313,
+      "step": 4030
+    },
+    {
+      "epoch": 0.41,
+      "gumbel_tau": 0.3864,
+      "image_masks_loss": 0.0246,
+      "llm_learning_rate": 3.310943311035466e-06,
+      "loss": 0.6908,
+      "output_text_mask_loss": 0.1389,
+      "predictor_learning_rate": 0.0003310943311035466,
+      "step": 4031
+    },
+    {
+      "epoch": 0.41,
+      "gumbel_tau": 0.3864,
+      "image_masks_loss": 0.0317,
+      "llm_learning_rate": 3.3101585088876155e-06,
+      "loss": 0.6691,
+      "output_text_mask_loss": 0.0391,
+      "predictor_learning_rate": 0.0003310158508887615,
+      "step": 4032
+    },
+    {
+      "epoch": 0.41,
+      "gumbel_tau": 0.3863,
+      "image_masks_loss": 0.0158,
+      "llm_learning_rate": 3.3093736175233058e-06,
+      "loss": 0.7072,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00033093736175233057,
+      "step": 4033
+    },
+    {
+      "epoch": 0.41,
+      "gumbel_tau": 0.3862,
+      "image_masks_loss": 0.0359,
+      "llm_learning_rate": 3.3085886370289712e-06,
+      "loss": 0.851,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0003308588637028971,
+      "step": 4034
+    },
+    {
+      "epoch": 0.41,
+      "gumbel_tau": 0.3861,
+      "image_masks_loss": 0.0142,
+      "llm_learning_rate": 3.307803567491054e-06,
+      "loss": 0.639,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0003307803567491054,
+      "step": 4035
+    },
+    {
+      "epoch": 0.41,
+      "gumbel_tau": 0.386,
+      "image_masks_loss": 0.0052,
+      "llm_learning_rate": 3.3070184089960084e-06,
+      "loss": 0.7966,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0003307018408996008,
+      "step": 4036
+    },
+    {
+      "epoch": 0.41,
+      "gumbel_tau": 0.3859,
+      "image_masks_loss": 0.0187,
+      "llm_learning_rate": 3.306233161630299e-06,
+      "loss": 0.7828,
+      "output_text_mask_loss": 0.004,
+      "predictor_learning_rate": 0.0003306233161630299,
+      "step": 4037
+    },
+    {
+      "epoch": 0.41,
+      "gumbel_tau": 0.3858,
+      "image_masks_loss": 0.0121,
+      "llm_learning_rate": 3.305447825480397e-06,
+      "loss": 0.7276,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.0003305447825480397,
+      "step": 4038
+    },
+    {
+      "epoch": 0.41,
+      "gumbel_tau": 0.3857,
+      "image_masks_loss": 0.0118,
+      "llm_learning_rate": 3.3046624006327873e-06,
+      "loss": 0.8302,
+      "output_text_mask_loss": 0.0269,
+      "predictor_learning_rate": 0.0003304662400632787,
+      "step": 4039
+    },
+    {
+      "epoch": 0.41,
+      "gumbel_tau": 0.3856,
+      "image_masks_loss": 0.0252,
+      "llm_learning_rate": 3.303876887173962e-06,
+      "loss": 0.6973,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.0003303876887173962,
+      "step": 4040
+    },
+    {
+      "epoch": 0.41,
+      "gumbel_tau": 0.3855,
+      "image_masks_loss": 0.0184,
+      "llm_learning_rate": 3.303091285190422e-06,
+      "loss": 0.7775,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.00033030912851904224,
+      "step": 4041
+    },
+    {
+      "epoch": 0.41,
+      "gumbel_tau": 0.3854,
+      "image_masks_loss": 0.0123,
+      "llm_learning_rate": 3.3023055947686814e-06,
+      "loss": 0.7657,
+      "output_text_mask_loss": 0.0334,
+      "predictor_learning_rate": 0.00033023055947686814,
+      "step": 4042
+    },
+    {
+      "epoch": 0.41,
+      "gumbel_tau": 0.3854,
+      "image_masks_loss": 0.0108,
+      "llm_learning_rate": 3.301519815995262e-06,
+      "loss": 0.6783,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0003301519815995262,
+      "step": 4043
+    },
+    {
+      "epoch": 0.41,
+      "gumbel_tau": 0.3853,
+      "image_masks_loss": 0.0437,
+      "llm_learning_rate": 3.3007339489566953e-06,
+      "loss": 0.7504,
+      "output_text_mask_loss": 0.0198,
+      "predictor_learning_rate": 0.0003300733948956695,
+      "step": 4044
+    },
+    {
+      "epoch": 0.41,
+      "gumbel_tau": 0.3852,
+      "image_masks_loss": 0.0208,
+      "llm_learning_rate": 3.299947993739522e-06,
+      "loss": 0.723,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00032999479937395213,
+      "step": 4045
+    },
+    {
+      "epoch": 0.41,
+      "gumbel_tau": 0.3851,
+      "image_masks_loss": 0.0425,
+      "llm_learning_rate": 3.299161950430294e-06,
+      "loss": 0.8118,
+      "output_text_mask_loss": 0.0248,
+      "predictor_learning_rate": 0.0003299161950430294,
+      "step": 4046
+    },
+    {
+      "epoch": 0.41,
+      "gumbel_tau": 0.385,
+      "image_masks_loss": 0.0088,
+      "llm_learning_rate": 3.298375819115572e-06,
+      "loss": 0.6151,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00032983758191155717,
+      "step": 4047
+    },
+    {
+      "epoch": 0.41,
+      "gumbel_tau": 0.3849,
+      "image_masks_loss": 0.0267,
+      "llm_learning_rate": 3.297589599881926e-06,
+      "loss": 0.798,
+      "output_text_mask_loss": 0.0429,
+      "predictor_learning_rate": 0.0003297589599881926,
+      "step": 4048
+    },
+    {
+      "epoch": 0.41,
+      "gumbel_tau": 0.3848,
+      "image_masks_loss": 0.0158,
+      "llm_learning_rate": 3.296803292815937e-06,
+      "loss": 0.734,
+      "output_text_mask_loss": 0.0088,
+      "predictor_learning_rate": 0.0003296803292815937,
+      "step": 4049
+    },
+    {
+      "epoch": 0.41,
+      "gumbel_tau": 0.3847,
+      "image_masks_loss": 0.0178,
+      "llm_learning_rate": 3.2960168980041957e-06,
+      "loss": 0.6988,
+      "output_text_mask_loss": 0.0137,
+      "predictor_learning_rate": 0.00032960168980041953,
+      "step": 4050
+    },
+    {
+      "epoch": 0.42,
+      "gumbel_tau": 0.3846,
+      "image_masks_loss": 0.0256,
+      "llm_learning_rate": 3.2952304155332993e-06,
+      "loss": 0.7232,
+      "output_text_mask_loss": 0.0032,
+      "predictor_learning_rate": 0.0003295230415533299,
+      "step": 4051
+    },
+    {
+      "epoch": 0.42,
+      "gumbel_tau": 0.3845,
+      "image_masks_loss": 0.015,
+      "llm_learning_rate": 3.2944438454898585e-06,
+      "loss": 0.7335,
+      "output_text_mask_loss": 0.0086,
+      "predictor_learning_rate": 0.00032944438454898583,
+      "step": 4052
+    },
+    {
+      "epoch": 0.42,
+      "gumbel_tau": 0.3844,
+      "image_masks_loss": 0.0088,
+      "llm_learning_rate": 3.293657187960492e-06,
+      "loss": 0.773,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00032936571879604913,
+      "step": 4053
+    },
+    {
+      "epoch": 0.42,
+      "gumbel_tau": 0.3844,
+      "image_masks_loss": 0.0206,
+      "llm_learning_rate": 3.292870443031828e-06,
+      "loss": 0.8794,
+      "output_text_mask_loss": 0.0139,
+      "predictor_learning_rate": 0.0003292870443031828,
+      "step": 4054
+    },
+    {
+      "epoch": 0.42,
+      "gumbel_tau": 0.3843,
+      "image_masks_loss": 0.0238,
+      "llm_learning_rate": 3.2920836107905048e-06,
+      "loss": 0.7421,
+      "output_text_mask_loss": 0.069,
+      "predictor_learning_rate": 0.00032920836107905044,
+      "step": 4055
+    },
+    {
+      "epoch": 0.42,
+      "gumbel_tau": 0.3842,
+      "image_masks_loss": 0.0271,
+      "llm_learning_rate": 3.29129669132317e-06,
+      "loss": 0.8874,
+      "output_text_mask_loss": 0.0793,
+      "predictor_learning_rate": 0.00032912966913231694,
+      "step": 4056
+    },
+    {
+      "epoch": 0.42,
+      "gumbel_tau": 0.3841,
+      "image_masks_loss": 0.0132,
+      "llm_learning_rate": 3.2905096847164807e-06,
+      "loss": 0.8603,
+      "output_text_mask_loss": 0.0603,
+      "predictor_learning_rate": 0.00032905096847164804,
+      "step": 4057
+    },
+    {
+      "epoch": 0.42,
+      "gumbel_tau": 0.384,
+      "image_masks_loss": 0.0292,
+      "llm_learning_rate": 3.289722591057105e-06,
+      "loss": 0.7081,
+      "output_text_mask_loss": 0.0086,
+      "predictor_learning_rate": 0.0003289722591057105,
+      "step": 4058
+    },
+    {
+      "epoch": 0.42,
+      "gumbel_tau": 0.3839,
+      "image_masks_loss": 0.0162,
+      "llm_learning_rate": 3.288935410431718e-06,
+      "loss": 0.7407,
+      "output_text_mask_loss": 0.0223,
+      "predictor_learning_rate": 0.0003288935410431718,
+      "step": 4059
+    },
+    {
+      "epoch": 0.42,
+      "gumbel_tau": 0.3838,
+      "image_masks_loss": 0.0234,
+      "llm_learning_rate": 3.2881481429270067e-06,
+      "loss": 0.7536,
+      "output_text_mask_loss": 0.0679,
+      "predictor_learning_rate": 0.0003288148142927007,
+      "step": 4060
+    },
+    {
+      "epoch": 0.42,
+      "gumbel_tau": 0.3837,
+      "image_masks_loss": 0.0212,
+      "llm_learning_rate": 3.287360788629666e-06,
+      "loss": 0.7421,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0003287360788629666,
+      "step": 4061
+    },
+    {
+      "epoch": 0.42,
+      "gumbel_tau": 0.3836,
+      "image_masks_loss": 0.0179,
+      "llm_learning_rate": 3.2865733476264016e-06,
+      "loss": 0.66,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00032865733476264013,
+      "step": 4062
+    },
+    {
+      "epoch": 0.42,
+      "gumbel_tau": 0.3835,
+      "image_masks_loss": 0.0332,
+      "llm_learning_rate": 3.285785820003928e-06,
+      "loss": 0.834,
+      "output_text_mask_loss": 0.0076,
+      "predictor_learning_rate": 0.00032857858200039275,
+      "step": 4063
+    },
+    {
+      "epoch": 0.42,
+      "gumbel_tau": 0.3835,
+      "image_masks_loss": 0.0211,
+      "llm_learning_rate": 3.2849982058489697e-06,
+      "loss": 0.7382,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00032849982058489696,
+      "step": 4064
+    },
+    {
+      "epoch": 0.42,
+      "gumbel_tau": 0.3834,
+      "image_masks_loss": 0.0473,
+      "llm_learning_rate": 3.28421050524826e-06,
+      "loss": 0.7885,
+      "output_text_mask_loss": 0.0018,
+      "predictor_learning_rate": 0.00032842105052482593,
+      "step": 4065
+    },
+    {
+      "epoch": 0.42,
+      "gumbel_tau": 0.3833,
+      "image_masks_loss": 0.0168,
+      "llm_learning_rate": 3.2834227182885437e-06,
+      "loss": 0.7999,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 0.00032834227182885437,
+      "step": 4066
+    },
+    {
+      "epoch": 0.42,
+      "gumbel_tau": 0.3832,
+      "image_masks_loss": 0.004,
+      "llm_learning_rate": 3.282634845056572e-06,
+      "loss": 0.7786,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 0.0003282634845056572,
+      "step": 4067
+    },
+    {
+      "epoch": 0.42,
+      "gumbel_tau": 0.3831,
+      "image_masks_loss": 0.0119,
+      "llm_learning_rate": 3.281846885639108e-06,
+      "loss": 0.8075,
+      "output_text_mask_loss": 0.0004,
+      "predictor_learning_rate": 0.00032818468856391083,
+      "step": 4068
+    },
+    {
+      "epoch": 0.42,
+      "gumbel_tau": 0.383,
+      "image_masks_loss": 0.0267,
+      "llm_learning_rate": 3.2810588401229235e-06,
+      "loss": 0.7797,
+      "output_text_mask_loss": 0.0977,
+      "predictor_learning_rate": 0.0003281058840122924,
+      "step": 4069
+    },
+    {
+      "epoch": 0.42,
+      "gumbel_tau": 0.3829,
+      "image_masks_loss": 0.0179,
+      "llm_learning_rate": 3.2802707085948e-06,
+      "loss": 0.7445,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.00032802707085947995,
+      "step": 4070
+    },
+    {
+      "epoch": 0.42,
+      "gumbel_tau": 0.3828,
+      "image_masks_loss": 0.0212,
+      "llm_learning_rate": 3.2794824911415276e-06,
+      "loss": 0.8472,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0003279482491141527,
+      "step": 4071
+    },
+    {
+      "epoch": 0.42,
+      "gumbel_tau": 0.3827,
+      "image_masks_loss": 0.028,
+      "llm_learning_rate": 3.278694187849907e-06,
+      "loss": 0.7035,
+      "output_text_mask_loss": 0.0078,
+      "predictor_learning_rate": 0.00032786941878499074,
+      "step": 4072
+    },
+    {
+      "epoch": 0.42,
+      "gumbel_tau": 0.3826,
+      "image_masks_loss": 0.0263,
+      "llm_learning_rate": 3.2779057988067485e-06,
+      "loss": 0.6627,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00032779057988067484,
+      "step": 4073
+    },
+    {
+      "epoch": 0.42,
+      "gumbel_tau": 0.3825,
+      "image_masks_loss": 0.0364,
+      "llm_learning_rate": 3.277117324098871e-06,
+      "loss": 0.7931,
+      "output_text_mask_loss": 0.0132,
+      "predictor_learning_rate": 0.00032771173240988703,
+      "step": 4074
+    },
+    {
+      "epoch": 0.42,
+      "gumbel_tau": 0.3825,
+      "image_masks_loss": 0.0134,
+      "llm_learning_rate": 3.2763287638131014e-06,
+      "loss": 0.7871,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00032763287638131014,
+      "step": 4075
+    },
+    {
+      "epoch": 0.42,
+      "gumbel_tau": 0.3824,
+      "image_masks_loss": 0.0182,
+      "llm_learning_rate": 3.2755401180362804e-06,
+      "loss": 0.8087,
+      "output_text_mask_loss": 0.1007,
+      "predictor_learning_rate": 0.00032755401180362804,
+      "step": 4076
+    },
+    {
+      "epoch": 0.42,
+      "gumbel_tau": 0.3823,
+      "image_masks_loss": 0.0082,
+      "llm_learning_rate": 3.274751386855254e-06,
+      "loss": 0.8873,
+      "output_text_mask_loss": 0.0185,
+      "predictor_learning_rate": 0.0003274751386855254,
+      "step": 4077
+    },
+    {
+      "epoch": 0.42,
+      "gumbel_tau": 0.3822,
+      "image_masks_loss": 0.0082,
+      "llm_learning_rate": 3.273962570356879e-06,
+      "loss": 0.7209,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00032739625703568785,
+      "step": 4078
+    },
+    {
+      "epoch": 0.42,
+      "gumbel_tau": 0.3821,
+      "image_masks_loss": 0.0087,
+      "llm_learning_rate": 3.273173668628022e-06,
+      "loss": 0.7616,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00032731736686280216,
+      "step": 4079
+    },
+    {
+      "epoch": 0.42,
+      "gumbel_tau": 0.382,
+      "image_masks_loss": 0.0242,
+      "llm_learning_rate": 3.272384681755558e-06,
+      "loss": 0.8182,
+      "output_text_mask_loss": 0.0404,
+      "predictor_learning_rate": 0.0003272384681755558,
+      "step": 4080
+    },
+    {
+      "epoch": 0.42,
+      "gumbel_tau": 0.3819,
+      "image_masks_loss": 0.0364,
+      "llm_learning_rate": 3.2715956098263723e-06,
+      "loss": 0.7136,
+      "output_text_mask_loss": 0.0231,
+      "predictor_learning_rate": 0.00032715956098263723,
+      "step": 4081
+    },
+    {
+      "epoch": 0.42,
+      "gumbel_tau": 0.3818,
+      "image_masks_loss": 0.0179,
+      "llm_learning_rate": 3.27080645292736e-06,
+      "loss": 0.8113,
+      "output_text_mask_loss": 0.0238,
+      "predictor_learning_rate": 0.00032708064529273596,
+      "step": 4082
+    },
+    {
+      "epoch": 0.42,
+      "gumbel_tau": 0.3817,
+      "image_masks_loss": 0.0083,
+      "llm_learning_rate": 3.2700172111454244e-06,
+      "loss": 0.7618,
+      "output_text_mask_loss": 0.0774,
+      "predictor_learning_rate": 0.0003270017211145424,
+      "step": 4083
+    },
+    {
+      "epoch": 0.42,
+      "gumbel_tau": 0.3816,
+      "image_masks_loss": 0.0149,
+      "llm_learning_rate": 3.2692278845674772e-06,
+      "loss": 0.7581,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 0.0003269227884567477,
+      "step": 4084
+    },
+    {
+      "epoch": 0.42,
+      "gumbel_tau": 0.3816,
+      "image_masks_loss": 0.0088,
+      "llm_learning_rate": 3.2684384732804425e-06,
+      "loss": 0.7661,
+      "output_text_mask_loss": 0.0618,
+      "predictor_learning_rate": 0.0003268438473280442,
+      "step": 4085
+    },
+    {
+      "epoch": 0.42,
+      "gumbel_tau": 0.3815,
+      "image_masks_loss": 0.0101,
+      "llm_learning_rate": 3.2676489773712515e-06,
+      "loss": 0.8268,
+      "output_text_mask_loss": 0.0033,
+      "predictor_learning_rate": 0.00032676489773712515,
+      "step": 4086
+    },
+    {
+      "epoch": 0.42,
+      "gumbel_tau": 0.3814,
+      "image_masks_loss": 0.0083,
+      "llm_learning_rate": 3.266859396926846e-06,
+      "loss": 0.6408,
+      "output_text_mask_loss": NaN,
+      "predictor_learning_rate": 0.00032668593969268457,
+      "step": 4087
+    },
+    {
+      "epoch": 0.42,
+      "gumbel_tau": 0.3813,
+      "image_masks_loss": 0.0047,
+      "llm_learning_rate": 3.266069732034175e-06,
+      "loss": 0.6802,
+      "output_text_mask_loss": 0.0345,
+      "predictor_learning_rate": 0.0003266069732034175,
+      "step": 4088
+    },
+    {
+      "epoch": 0.42,
+      "gumbel_tau": 0.3812,
+      "image_masks_loss": 0.0223,
+      "llm_learning_rate": 3.265279982780199e-06,
+      "loss": 0.7712,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.0003265279982780199,
+      "step": 4089
+    },
+    {
+      "epoch": 0.42,
+      "gumbel_tau": 0.3811,
+      "image_masks_loss": 0.0097,
+      "llm_learning_rate": 3.2644901492518864e-06,
+      "loss": 0.7638,
+      "output_text_mask_loss": 0.161,
+      "predictor_learning_rate": 0.0003264490149251886,
+      "step": 4090
+    },
+    {
+      "epoch": 0.42,
+      "gumbel_tau": 0.381,
+      "image_masks_loss": 0.0146,
+      "llm_learning_rate": 3.2637002315362153e-06,
+      "loss": 0.7272,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0003263700231536215,
+      "step": 4091
+    },
+    {
+      "epoch": 0.42,
+      "gumbel_tau": 0.3809,
+      "image_masks_loss": 0.0039,
+      "llm_learning_rate": 3.2629102297201744e-06,
+      "loss": 0.7501,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 0.0003262910229720174,
+      "step": 4092
+    },
+    {
+      "epoch": 0.42,
+      "gumbel_tau": 0.3808,
+      "image_masks_loss": 0.011,
+      "llm_learning_rate": 3.2621201438907597e-06,
+      "loss": 0.7695,
+      "output_text_mask_loss": 0.0286,
+      "predictor_learning_rate": 0.00032621201438907596,
+      "step": 4093
+    },
+    {
+      "epoch": 0.42,
+      "gumbel_tau": 0.3807,
+      "image_masks_loss": 0.0216,
+      "llm_learning_rate": 3.2613299741349765e-06,
+      "loss": 0.714,
+      "output_text_mask_loss": 0.0278,
+      "predictor_learning_rate": 0.00032613299741349763,
+      "step": 4094
+    },
+    {
+      "epoch": 0.42,
+      "gumbel_tau": 0.3807,
+      "image_masks_loss": 0.0189,
+      "llm_learning_rate": 3.260539720539841e-06,
+      "loss": 0.7717,
+      "output_text_mask_loss": 0.0012,
+      "predictor_learning_rate": 0.00032605397205398405,
+      "step": 4095
+    },
+    {
+      "epoch": 0.42,
+      "gumbel_tau": 0.3806,
+      "image_masks_loss": 0.0185,
+      "llm_learning_rate": 3.259749383192378e-06,
+      "loss": 0.724,
+      "output_text_mask_loss": 0.0402,
+      "predictor_learning_rate": 0.0003259749383192378,
+      "step": 4096
+    },
+    {
+      "epoch": 0.42,
+      "gumbel_tau": 0.3805,
+      "image_masks_loss": 0.0307,
+      "llm_learning_rate": 3.25895896217962e-06,
+      "loss": 0.8119,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.000325895896217962,
+      "step": 4097
+    },
+    {
+      "epoch": 0.42,
+      "gumbel_tau": 0.3804,
+      "image_masks_loss": 0.0087,
+      "llm_learning_rate": 3.2581684575886106e-06,
+      "loss": 0.6557,
+      "output_text_mask_loss": 0.087,
+      "predictor_learning_rate": 0.00032581684575886104,
+      "step": 4098
+    },
+    {
+      "epoch": 0.42,
+      "gumbel_tau": 0.3803,
+      "image_masks_loss": 0.0307,
+      "llm_learning_rate": 3.2573778695064025e-06,
+      "loss": 0.771,
+      "output_text_mask_loss": 0.02,
+      "predictor_learning_rate": 0.0003257377869506402,
+      "step": 4099
+    },
+    {
+      "epoch": 0.42,
+      "gumbel_tau": 0.3802,
+      "image_masks_loss": 0.0252,
+      "llm_learning_rate": 3.256587198020056e-06,
+      "loss": 0.6952,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 0.0003256587198020056,
+      "step": 4100
+    },
+    {
+      "epoch": 0.42,
+      "gumbel_tau": 0.3801,
+      "image_masks_loss": 0.0171,
+      "llm_learning_rate": 3.2557964432166415e-06,
+      "loss": 0.7696,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.00032557964432166416,
+      "step": 4101
+    },
+    {
+      "epoch": 0.42,
+      "gumbel_tau": 0.38,
+      "image_masks_loss": 0.0113,
+      "llm_learning_rate": 3.25500560518324e-06,
+      "loss": 0.725,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00032550056051832397,
+      "step": 4102
+    },
+    {
+      "epoch": 0.42,
+      "gumbel_tau": 0.3799,
+      "image_masks_loss": 0.0163,
+      "llm_learning_rate": 3.2542146840069393e-06,
+      "loss": 0.774,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 0.00032542146840069394,
+      "step": 4103
+    },
+    {
+      "epoch": 0.42,
+      "gumbel_tau": 0.3798,
+      "image_masks_loss": 0.0282,
+      "llm_learning_rate": 3.2534236797748374e-06,
+      "loss": 0.7858,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.0003253423679774837,
+      "step": 4104
+    },
+    {
+      "epoch": 0.42,
+      "gumbel_tau": 0.3798,
+      "image_masks_loss": 0.0187,
+      "llm_learning_rate": 3.2526325925740428e-06,
+      "loss": 0.6989,
+      "output_text_mask_loss": 0.0467,
+      "predictor_learning_rate": 0.00032526325925740423,
+      "step": 4105
+    },
+    {
+      "epoch": 0.42,
+      "gumbel_tau": 0.3797,
+      "image_masks_loss": 0.0124,
+      "llm_learning_rate": 3.2518414224916696e-06,
+      "loss": 0.7133,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.00032518414224916694,
+      "step": 4106
+    },
+    {
+      "epoch": 0.42,
+      "gumbel_tau": 0.3796,
+      "image_masks_loss": 0.0174,
+      "llm_learning_rate": 3.251050169614844e-06,
+      "loss": 0.8092,
+      "output_text_mask_loss": 0.0011,
+      "predictor_learning_rate": 0.00032510501696148437,
+      "step": 4107
+    },
+    {
+      "epoch": 0.42,
+      "gumbel_tau": 0.3795,
+      "image_masks_loss": 0.0294,
+      "llm_learning_rate": 3.2502588340307016e-06,
+      "loss": 0.8585,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.00032502588340307015,
+      "step": 4108
+    },
+    {
+      "epoch": 0.42,
+      "gumbel_tau": 0.3794,
+      "image_masks_loss": 0.0208,
+      "llm_learning_rate": 3.249467415826385e-06,
+      "loss": 0.7436,
+      "output_text_mask_loss": 0.0032,
+      "predictor_learning_rate": 0.0003249467415826385,
+      "step": 4109
+    },
+    {
+      "epoch": 0.42,
+      "gumbel_tau": 0.3793,
+      "image_masks_loss": 0.011,
+      "llm_learning_rate": 3.248675915089047e-06,
+      "loss": 0.7927,
+      "output_text_mask_loss": 0.0248,
+      "predictor_learning_rate": 0.0003248675915089047,
+      "step": 4110
+    },
+    {
+      "epoch": 0.42,
+      "gumbel_tau": 0.3792,
+      "image_masks_loss": 0.0353,
+      "llm_learning_rate": 3.2478843319058495e-06,
+      "loss": 0.7539,
+      "output_text_mask_loss": 0.0423,
+      "predictor_learning_rate": 0.0003247884331905849,
+      "step": 4111
+    },
+    {
+      "epoch": 0.42,
+      "gumbel_tau": 0.3791,
+      "image_masks_loss": 0.0162,
+      "llm_learning_rate": 3.247092666363964e-06,
+      "loss": 0.684,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 0.0003247092666363964,
+      "step": 4112
+    },
+    {
+      "epoch": 0.42,
+      "gumbel_tau": 0.379,
+      "image_masks_loss": 0.0069,
+      "llm_learning_rate": 3.24630091855057e-06,
+      "loss": 0.718,
+      "output_text_mask_loss": 0.0322,
+      "predictor_learning_rate": 0.000324630091855057,
+      "step": 4113
+    },
+    {
+      "epoch": 0.42,
+      "gumbel_tau": 0.379,
+      "image_masks_loss": 0.0257,
+      "llm_learning_rate": 3.2455090885528563e-06,
+      "loss": 0.7381,
+      "output_text_mask_loss": 0.0004,
+      "predictor_learning_rate": 0.0003245509088552856,
+      "step": 4114
+    },
+    {
+      "epoch": 0.42,
+      "gumbel_tau": 0.3789,
+      "image_masks_loss": 0.0234,
+      "llm_learning_rate": 3.2447171764580217e-06,
+      "loss": 0.7205,
+      "output_text_mask_loss": 0.0101,
+      "predictor_learning_rate": 0.0003244717176458021,
+      "step": 4115
+    },
+    {
+      "epoch": 0.42,
+      "gumbel_tau": 0.3788,
+      "image_masks_loss": 0.0102,
+      "llm_learning_rate": 3.2439251823532725e-06,
+      "loss": 0.7248,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.0003243925182353272,
+      "step": 4116
+    },
+    {
+      "epoch": 0.42,
+      "gumbel_tau": 0.3787,
+      "image_masks_loss": 0.0296,
+      "llm_learning_rate": 3.2431331063258243e-06,
+      "loss": 0.7467,
+      "output_text_mask_loss": 0.0107,
+      "predictor_learning_rate": 0.0003243133106325824,
+      "step": 4117
+    },
+    {
+      "epoch": 0.42,
+      "gumbel_tau": 0.3786,
+      "image_masks_loss": 0.0114,
+      "llm_learning_rate": 3.242340948462904e-06,
+      "loss": 0.8015,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.0003242340948462904,
+      "step": 4118
+    },
+    {
+      "epoch": 0.42,
+      "gumbel_tau": 0.3785,
+      "image_masks_loss": 0.0097,
+      "llm_learning_rate": 3.2415487088517457e-06,
+      "loss": 0.7583,
+      "output_text_mask_loss": 0.0263,
+      "predictor_learning_rate": 0.0003241548708851745,
+      "step": 4119
+    },
+    {
+      "epoch": 0.42,
+      "gumbel_tau": 0.3784,
+      "image_masks_loss": 0.0256,
+      "llm_learning_rate": 3.24075638757959e-06,
+      "loss": 0.8019,
+      "output_text_mask_loss": 0.0124,
+      "predictor_learning_rate": 0.00032407563875795896,
+      "step": 4120
+    },
+    {
+      "epoch": 0.42,
+      "gumbel_tau": 0.3783,
+      "image_masks_loss": 0.0171,
+      "llm_learning_rate": 3.2399639847336924e-06,
+      "loss": 0.7036,
+      "output_text_mask_loss": 0.0196,
+      "predictor_learning_rate": 0.0003239963984733692,
+      "step": 4121
+    },
+    {
+      "epoch": 0.42,
+      "gumbel_tau": 0.3782,
+      "image_masks_loss": 0.0208,
+      "llm_learning_rate": 3.2391715004013126e-06,
+      "loss": 0.7604,
+      "output_text_mask_loss": 0.0144,
+      "predictor_learning_rate": 0.0003239171500401312,
+      "step": 4122
+    },
+    {
+      "epoch": 0.42,
+      "gumbel_tau": 0.3781,
+      "image_masks_loss": 0.0098,
+      "llm_learning_rate": 3.2383789346697204e-06,
+      "loss": 0.8525,
+      "output_text_mask_loss": 0.01,
+      "predictor_learning_rate": 0.000323837893466972,
+      "step": 4123
+    },
+    {
+      "epoch": 0.42,
+      "gumbel_tau": 0.3781,
+      "image_masks_loss": 0.0134,
+      "llm_learning_rate": 3.237586287626195e-06,
+      "loss": 0.6503,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 0.00032375862876261944,
+      "step": 4124
+    },
+    {
+      "epoch": 0.42,
+      "gumbel_tau": 0.378,
+      "image_masks_loss": 0.0092,
+      "llm_learning_rate": 3.236793559358025e-06,
+      "loss": 0.8363,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.0003236793559358024,
+      "step": 4125
+    },
+    {
+      "epoch": 0.42,
+      "gumbel_tau": 0.3779,
+      "image_masks_loss": 0.013,
+      "llm_learning_rate": 3.236000749952507e-06,
+      "loss": 0.8325,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.0003236000749952507,
+      "step": 4126
+    },
+    {
+      "epoch": 0.42,
+      "gumbel_tau": 0.3778,
+      "image_masks_loss": 0.0155,
+      "llm_learning_rate": 3.235207859496947e-06,
+      "loss": 0.7111,
+      "output_text_mask_loss": 0.09,
+      "predictor_learning_rate": 0.00032352078594969464,
+      "step": 4127
+    },
+    {
+      "epoch": 0.42,
+      "gumbel_tau": 0.3777,
+      "image_masks_loss": 0.0267,
+      "llm_learning_rate": 3.23441488807866e-06,
+      "loss": 0.762,
+      "output_text_mask_loss": 0.0816,
+      "predictor_learning_rate": 0.0003234414888078659,
+      "step": 4128
+    },
+    {
+      "epoch": 0.42,
+      "gumbel_tau": 0.3776,
+      "image_masks_loss": 0.029,
+      "llm_learning_rate": 3.233621835784969e-06,
+      "loss": 0.8444,
+      "output_text_mask_loss": 0.0444,
+      "predictor_learning_rate": 0.0003233621835784969,
+      "step": 4129
+    },
+    {
+      "epoch": 0.42,
+      "gumbel_tau": 0.3775,
+      "image_masks_loss": 0.0118,
+      "llm_learning_rate": 3.232828702703208e-06,
+      "loss": 0.7063,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.0003232828702703208,
+      "step": 4130
+    },
+    {
+      "epoch": 0.42,
+      "gumbel_tau": 0.3774,
+      "image_masks_loss": 0.0157,
+      "llm_learning_rate": 3.2320354889207177e-06,
+      "loss": 0.719,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.00032320354889207176,
+      "step": 4131
+    },
+    {
+      "epoch": 0.42,
+      "gumbel_tau": 0.3773,
+      "image_masks_loss": 0.0128,
+      "llm_learning_rate": 3.2312421945248496e-06,
+      "loss": 0.7719,
+      "output_text_mask_loss": 0.013,
+      "predictor_learning_rate": 0.00032312421945248493,
+      "step": 4132
+    },
+    {
+      "epoch": 0.42,
+      "gumbel_tau": 0.3773,
+      "image_masks_loss": 0.0074,
+      "llm_learning_rate": 3.2304488196029614e-06,
+      "loss": 0.7683,
+      "output_text_mask_loss": 0.017,
+      "predictor_learning_rate": 0.0003230448819602961,
+      "step": 4133
+    },
+    {
+      "epoch": 0.42,
+      "gumbel_tau": 0.3772,
+      "image_masks_loss": 0.0071,
+      "llm_learning_rate": 3.2296553642424226e-06,
+      "loss": 0.7349,
+      "output_text_mask_loss": 0.0098,
+      "predictor_learning_rate": 0.00032296553642424225,
+      "step": 4134
+    },
+    {
+      "epoch": 0.42,
+      "gumbel_tau": 0.3771,
+      "image_masks_loss": 0.0389,
+      "llm_learning_rate": 3.2288618285306095e-06,
+      "loss": 0.8412,
+      "output_text_mask_loss": 0.0174,
+      "predictor_learning_rate": 0.0003228861828530609,
+      "step": 4135
+    },
+    {
+      "epoch": 0.42,
+      "gumbel_tau": 0.377,
+      "image_masks_loss": 0.0216,
+      "llm_learning_rate": 3.2280682125549086e-06,
+      "loss": 0.7178,
+      "output_text_mask_loss": 0.0261,
+      "predictor_learning_rate": 0.00032280682125549083,
+      "step": 4136
+    },
+    {
+      "epoch": 0.42,
+      "gumbel_tau": 0.3769,
+      "image_masks_loss": 0.0218,
+      "llm_learning_rate": 3.227274516402714e-06,
+      "loss": 0.7527,
+      "output_text_mask_loss": 0.0422,
+      "predictor_learning_rate": 0.0003227274516402714,
+      "step": 4137
+    },
+    {
+      "epoch": 0.42,
+      "gumbel_tau": 0.3768,
+      "image_masks_loss": 0.0175,
+      "llm_learning_rate": 3.2264807401614305e-06,
+      "loss": 0.7184,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.000322648074016143,
+      "step": 4138
+    },
+    {
+      "epoch": 0.42,
+      "gumbel_tau": 0.3767,
+      "image_masks_loss": 0.0176,
+      "llm_learning_rate": 3.2256868839184694e-06,
+      "loss": 0.717,
+      "output_text_mask_loss": 0.0056,
+      "predictor_learning_rate": 0.00032256868839184696,
+      "step": 4139
+    },
+    {
+      "epoch": 0.42,
+      "gumbel_tau": 0.3766,
+      "image_masks_loss": 0.0215,
+      "llm_learning_rate": 3.2248929477612516e-06,
+      "loss": 0.6872,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.00032248929477612514,
+      "step": 4140
+    },
+    {
+      "epoch": 0.42,
+      "gumbel_tau": 0.3765,
+      "image_masks_loss": 0.0168,
+      "llm_learning_rate": 3.224098931777209e-06,
+      "loss": 0.6804,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00032240989317772087,
+      "step": 4141
+    },
+    {
+      "epoch": 0.42,
+      "gumbel_tau": 0.3765,
+      "image_masks_loss": 0.0286,
+      "llm_learning_rate": 3.223304836053779e-06,
+      "loss": 0.7368,
+      "output_text_mask_loss": 0.005,
+      "predictor_learning_rate": 0.00032233048360537787,
+      "step": 4142
+    },
+    {
+      "epoch": 0.42,
+      "gumbel_tau": 0.3764,
+      "image_masks_loss": 0.0132,
+      "llm_learning_rate": 3.222510660678408e-06,
+      "loss": 0.6928,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0003222510660678408,
+      "step": 4143
+    },
+    {
+      "epoch": 0.42,
+      "gumbel_tau": 0.3763,
+      "image_masks_loss": 0.0099,
+      "llm_learning_rate": 3.221716405738555e-06,
+      "loss": 0.7124,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.00032217164057385546,
+      "step": 4144
+    },
+    {
+      "epoch": 0.42,
+      "gumbel_tau": 0.3762,
+      "image_masks_loss": 0.0105,
+      "llm_learning_rate": 3.2209220713216833e-06,
+      "loss": 0.7578,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0003220922071321683,
+      "step": 4145
+    },
+    {
+      "epoch": 0.42,
+      "gumbel_tau": 0.3761,
+      "image_masks_loss": 0.0105,
+      "llm_learning_rate": 3.220127657515267e-06,
+      "loss": 0.764,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.00032201276575152663,
+      "step": 4146
+    },
+    {
+      "epoch": 0.42,
+      "gumbel_tau": 0.376,
+      "image_masks_loss": 0.0055,
+      "llm_learning_rate": 3.2193331644067895e-06,
+      "loss": 0.7011,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.00032193331644067894,
+      "step": 4147
+    },
+    {
+      "epoch": 0.42,
+      "gumbel_tau": 0.3759,
+      "image_masks_loss": 0.0353,
+      "llm_learning_rate": 3.2185385920837415e-06,
+      "loss": 0.8934,
+      "output_text_mask_loss": 0.0109,
+      "predictor_learning_rate": 0.00032185385920837414,
+      "step": 4148
+    },
+    {
+      "epoch": 0.43,
+      "gumbel_tau": 0.3758,
+      "image_masks_loss": 0.0515,
+      "llm_learning_rate": 3.2177439406336237e-06,
+      "loss": 0.8111,
+      "output_text_mask_loss": 0.0359,
+      "predictor_learning_rate": 0.00032177439406336233,
+      "step": 4149
+    },
+    {
+      "epoch": 0.43,
+      "gumbel_tau": 0.3757,
+      "image_masks_loss": 0.0056,
+      "llm_learning_rate": 3.2169492101439436e-06,
+      "loss": 0.7129,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.00032169492101439434,
+      "step": 4150
+    },
+    {
+      "epoch": 0.43,
+      "gumbel_tau": 0.3757,
+      "image_masks_loss": 0.0195,
+      "llm_learning_rate": 3.21615440070222e-06,
+      "loss": 0.759,
+      "output_text_mask_loss": 0.0172,
+      "predictor_learning_rate": 0.000321615440070222,
+      "step": 4151
+    },
+    {
+      "epoch": 0.43,
+      "gumbel_tau": 0.3756,
+      "image_masks_loss": 0.0484,
+      "llm_learning_rate": 3.2153595123959793e-06,
+      "loss": 0.7745,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 0.0003215359512395979,
+      "step": 4152
+    },
+    {
+      "epoch": 0.43,
+      "gumbel_tau": 0.3755,
+      "image_masks_loss": 0.0292,
+      "llm_learning_rate": 3.2145645453127548e-06,
+      "loss": 0.846,
+      "output_text_mask_loss": 0.1015,
+      "predictor_learning_rate": 0.0003214564545312755,
+      "step": 4153
+    },
+    {
+      "epoch": 0.43,
+      "gumbel_tau": 0.3754,
+      "image_masks_loss": 0.0313,
+      "llm_learning_rate": 3.2137694995400915e-06,
+      "loss": 0.7665,
+      "output_text_mask_loss": 0.0162,
+      "predictor_learning_rate": 0.0003213769499540091,
+      "step": 4154
+    },
+    {
+      "epoch": 0.43,
+      "gumbel_tau": 0.3753,
+      "image_masks_loss": 0.0166,
+      "llm_learning_rate": 3.2129743751655406e-06,
+      "loss": 0.5189,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00032129743751655403,
+      "step": 4155
+    },
+    {
+      "epoch": 0.43,
+      "gumbel_tau": 0.3752,
+      "image_masks_loss": 0.0233,
+      "llm_learning_rate": 3.212179172276663e-06,
+      "loss": 0.7979,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.0003212179172276663,
+      "step": 4156
+    },
+    {
+      "epoch": 0.43,
+      "gumbel_tau": 0.3751,
+      "image_masks_loss": 0.0163,
+      "llm_learning_rate": 3.2113838909610297e-06,
+      "loss": 0.7669,
+      "output_text_mask_loss": 0.0313,
+      "predictor_learning_rate": 0.00032113838909610294,
+      "step": 4157
+    },
+    {
+      "epoch": 0.43,
+      "gumbel_tau": 0.375,
+      "image_masks_loss": 0.0203,
+      "llm_learning_rate": 3.2105885313062175e-06,
+      "loss": 0.7874,
+      "output_text_mask_loss": 0.0584,
+      "predictor_learning_rate": 0.00032105885313062175,
+      "step": 4158
+    },
+    {
+      "epoch": 0.43,
+      "gumbel_tau": 0.375,
+      "image_masks_loss": 0.016,
+      "llm_learning_rate": 3.2097930933998127e-06,
+      "loss": 0.6347,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00032097930933998125,
+      "step": 4159
+    },
+    {
+      "epoch": 0.43,
+      "gumbel_tau": 0.3749,
+      "image_masks_loss": 0.0167,
+      "llm_learning_rate": 3.208997577329412e-06,
+      "loss": 0.7236,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00032089975773294115,
+      "step": 4160
+    },
+    {
+      "epoch": 0.43,
+      "gumbel_tau": 0.3748,
+      "image_masks_loss": 0.0298,
+      "llm_learning_rate": 3.2082019831826196e-06,
+      "loss": 0.7598,
+      "output_text_mask_loss": 0.0275,
+      "predictor_learning_rate": 0.0003208201983182619,
+      "step": 4161
+    },
+    {
+      "epoch": 0.43,
+      "gumbel_tau": 0.3747,
+      "image_masks_loss": 0.0087,
+      "llm_learning_rate": 3.207406311047046e-06,
+      "loss": 0.7453,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.0003207406311047046,
+      "step": 4162
+    },
+    {
+      "epoch": 0.43,
+      "gumbel_tau": 0.3746,
+      "image_masks_loss": 0.0243,
+      "llm_learning_rate": 3.2066105610103148e-06,
+      "loss": 0.7099,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 0.00032066105610103145,
+      "step": 4163
+    },
+    {
+      "epoch": 0.43,
+      "gumbel_tau": 0.3745,
+      "image_masks_loss": 0.0275,
+      "llm_learning_rate": 3.2058147331600543e-06,
+      "loss": 0.7576,
+      "output_text_mask_loss": 0.0231,
+      "predictor_learning_rate": 0.0003205814733160054,
+      "step": 4164
+    },
+    {
+      "epoch": 0.43,
+      "gumbel_tau": 0.3744,
+      "image_masks_loss": 0.0353,
+      "llm_learning_rate": 3.2050188275839035e-06,
+      "loss": 0.7942,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0003205018827583903,
+      "step": 4165
+    },
+    {
+      "epoch": 0.43,
+      "gumbel_tau": 0.3743,
+      "image_masks_loss": 0.0257,
+      "llm_learning_rate": 3.204222844369508e-06,
+      "loss": 0.7444,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.0003204222844369508,
+      "step": 4166
+    },
+    {
+      "epoch": 0.43,
+      "gumbel_tau": 0.3742,
+      "image_masks_loss": 0.0136,
+      "llm_learning_rate": 3.2034267836045253e-06,
+      "loss": 0.762,
+      "output_text_mask_loss": 0.0488,
+      "predictor_learning_rate": 0.0003203426783604525,
+      "step": 4167
+    },
+    {
+      "epoch": 0.43,
+      "gumbel_tau": 0.3742,
+      "image_masks_loss": 0.012,
+      "llm_learning_rate": 3.202630645376618e-06,
+      "loss": 0.8061,
+      "output_text_mask_loss": 0.0286,
+      "predictor_learning_rate": 0.0003202630645376618,
+      "step": 4168
+    },
+    {
+      "epoch": 0.43,
+      "gumbel_tau": 0.3741,
+      "image_masks_loss": 0.0139,
+      "llm_learning_rate": 3.2018344297734587e-06,
+      "loss": 0.744,
+      "output_text_mask_loss": 0.0118,
+      "predictor_learning_rate": 0.00032018344297734587,
+      "step": 4169
+    },
+    {
+      "epoch": 0.43,
+      "gumbel_tau": 0.374,
+      "image_masks_loss": 0.0097,
+      "llm_learning_rate": 3.2010381368827297e-06,
+      "loss": 0.7068,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.000320103813688273,
+      "step": 4170
+    },
+    {
+      "epoch": 0.43,
+      "gumbel_tau": 0.3739,
+      "image_masks_loss": 0.0155,
+      "llm_learning_rate": 3.2002417667921193e-06,
+      "loss": 0.8267,
+      "output_text_mask_loss": 0.001,
+      "predictor_learning_rate": 0.0003200241766792119,
+      "step": 4171
+    },
+    {
+      "epoch": 0.43,
+      "gumbel_tau": 0.3738,
+      "image_masks_loss": 0.0223,
+      "llm_learning_rate": 3.199445319589325e-06,
+      "loss": 0.7852,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.0003199445319589325,
+      "step": 4172
+    },
+    {
+      "epoch": 0.43,
+      "gumbel_tau": 0.3737,
+      "image_masks_loss": 0.0174,
+      "llm_learning_rate": 3.1986487953620547e-06,
+      "loss": 0.6973,
+      "output_text_mask_loss": 0.0275,
+      "predictor_learning_rate": 0.0003198648795362055,
+      "step": 4173
+    },
+    {
+      "epoch": 0.43,
+      "gumbel_tau": 0.3736,
+      "image_masks_loss": 0.0395,
+      "llm_learning_rate": 3.1978521941980227e-06,
+      "loss": 0.7784,
+      "output_text_mask_loss": 0.0425,
+      "predictor_learning_rate": 0.00031978521941980226,
+      "step": 4174
+    },
+    {
+      "epoch": 0.43,
+      "gumbel_tau": 0.3735,
+      "image_masks_loss": 0.014,
+      "llm_learning_rate": 3.1970555161849525e-06,
+      "loss": 0.7253,
+      "output_text_mask_loss": 0.0124,
+      "predictor_learning_rate": 0.00031970555161849523,
+      "step": 4175
+    },
+    {
+      "epoch": 0.43,
+      "gumbel_tau": 0.3735,
+      "image_masks_loss": 0.0294,
+      "llm_learning_rate": 3.196258761410577e-06,
+      "loss": 0.8016,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00031962587614105763,
+      "step": 4176
+    },
+    {
+      "epoch": 0.43,
+      "gumbel_tau": 0.3734,
+      "image_masks_loss": 0.033,
+      "llm_learning_rate": 3.1954619299626347e-06,
+      "loss": 0.7906,
+      "output_text_mask_loss": 0.0065,
+      "predictor_learning_rate": 0.0003195461929962635,
+      "step": 4177
+    },
+    {
+      "epoch": 0.43,
+      "gumbel_tau": 0.3733,
+      "image_masks_loss": 0.0115,
+      "llm_learning_rate": 3.194665021928877e-06,
+      "loss": 0.8298,
+      "output_text_mask_loss": 0.1312,
+      "predictor_learning_rate": 0.00031946650219288766,
+      "step": 4178
+    },
+    {
+      "epoch": 0.43,
+      "gumbel_tau": 0.3732,
+      "image_masks_loss": 0.0207,
+      "llm_learning_rate": 3.1938680373970586e-06,
+      "loss": 0.7939,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.0003193868037397059,
+      "step": 4179
+    },
+    {
+      "epoch": 0.43,
+      "gumbel_tau": 0.3731,
+      "image_masks_loss": 0.0122,
+      "llm_learning_rate": 3.193070976454947e-06,
+      "loss": 0.8148,
+      "output_text_mask_loss": 0.0158,
+      "predictor_learning_rate": 0.00031930709764549473,
+      "step": 4180
+    },
+    {
+      "epoch": 0.43,
+      "gumbel_tau": 0.373,
+      "image_masks_loss": 0.0175,
+      "llm_learning_rate": 3.1922738391903156e-06,
+      "loss": 0.7706,
+      "output_text_mask_loss": 0.01,
+      "predictor_learning_rate": 0.00031922738391903155,
+      "step": 4181
+    },
+    {
+      "epoch": 0.43,
+      "gumbel_tau": 0.3729,
+      "image_masks_loss": 0.0324,
+      "llm_learning_rate": 3.191476625690947e-06,
+      "loss": 0.7753,
+      "output_text_mask_loss": 0.0648,
+      "predictor_learning_rate": 0.0003191476625690947,
+      "step": 4182
+    },
+    {
+      "epoch": 0.43,
+      "gumbel_tau": 0.3728,
+      "image_masks_loss": 0.0141,
+      "llm_learning_rate": 3.190679336044632e-06,
+      "loss": 0.6874,
+      "output_text_mask_loss": 0.0081,
+      "predictor_learning_rate": 0.0003190679336044632,
+      "step": 4183
+    },
+    {
+      "epoch": 0.43,
+      "gumbel_tau": 0.3727,
+      "image_masks_loss": 0.0226,
+      "llm_learning_rate": 3.1898819703391704e-06,
+      "loss": 0.7049,
+      "output_text_mask_loss": 0.0231,
+      "predictor_learning_rate": 0.000318988197033917,
+      "step": 4184
+    },
+    {
+      "epoch": 0.43,
+      "gumbel_tau": 0.3727,
+      "image_masks_loss": 0.0136,
+      "llm_learning_rate": 3.189084528662369e-06,
+      "loss": 0.82,
+      "output_text_mask_loss": 0.0038,
+      "predictor_learning_rate": 0.0003189084528662369,
+      "step": 4185
+    },
+    {
+      "epoch": 0.43,
+      "gumbel_tau": 0.3726,
+      "image_masks_loss": 0.0052,
+      "llm_learning_rate": 3.1882870111020446e-06,
+      "loss": 0.7919,
+      "output_text_mask_loss": 0.0374,
+      "predictor_learning_rate": 0.00031882870111020446,
+      "step": 4186
+    },
+    {
+      "epoch": 0.43,
+      "gumbel_tau": 0.3725,
+      "image_masks_loss": 0.0218,
+      "llm_learning_rate": 3.1874894177460217e-06,
+      "loss": 0.8143,
+      "output_text_mask_loss": 0.0736,
+      "predictor_learning_rate": 0.0003187489417746022,
+      "step": 4187
+    },
+    {
+      "epoch": 0.43,
+      "gumbel_tau": 0.3724,
+      "image_masks_loss": 0.0214,
+      "llm_learning_rate": 3.186691748682133e-06,
+      "loss": 0.5729,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0003186691748682133,
+      "step": 4188
+    },
+    {
+      "epoch": 0.43,
+      "gumbel_tau": 0.3723,
+      "image_masks_loss": 0.009,
+      "llm_learning_rate": 3.1858940039982197e-06,
+      "loss": 0.7653,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00031858940039982195,
+      "step": 4189
+    },
+    {
+      "epoch": 0.43,
+      "gumbel_tau": 0.3722,
+      "image_masks_loss": 0.0214,
+      "llm_learning_rate": 3.1850961837821303e-06,
+      "loss": 0.7582,
+      "output_text_mask_loss": 0.0145,
+      "predictor_learning_rate": 0.000318509618378213,
+      "step": 4190
+    },
+    {
+      "epoch": 0.43,
+      "gumbel_tau": 0.3721,
+      "image_masks_loss": 0.0181,
+      "llm_learning_rate": 3.1842982881217227e-06,
+      "loss": 0.7744,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.0003184298288121722,
+      "step": 4191
+    },
+    {
+      "epoch": 0.43,
+      "gumbel_tau": 0.372,
+      "image_masks_loss": 0.019,
+      "llm_learning_rate": 3.183500317104864e-06,
+      "loss": 0.812,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 0.00031835003171048635,
+      "step": 4192
+    },
+    {
+      "epoch": 0.43,
+      "gumbel_tau": 0.372,
+      "image_masks_loss": 0.0042,
+      "llm_learning_rate": 3.1827022708194277e-06,
+      "loss": 0.7486,
+      "output_text_mask_loss": 0.0523,
+      "predictor_learning_rate": 0.00031827022708194277,
+      "step": 4193
+    },
+    {
+      "epoch": 0.43,
+      "gumbel_tau": 0.3719,
+      "image_masks_loss": 0.0109,
+      "llm_learning_rate": 3.1819041493532975e-06,
+      "loss": 0.7285,
+      "output_text_mask_loss": 0.0376,
+      "predictor_learning_rate": 0.00031819041493532974,
+      "step": 4194
+    },
+    {
+      "epoch": 0.43,
+      "gumbel_tau": 0.3718,
+      "image_masks_loss": 0.0093,
+      "llm_learning_rate": 3.1811059527943622e-06,
+      "loss": 0.7916,
+      "output_text_mask_loss": 0.0039,
+      "predictor_learning_rate": 0.0003181105952794362,
+      "step": 4195
+    },
+    {
+      "epoch": 0.43,
+      "gumbel_tau": 0.3717,
+      "image_masks_loss": 0.0078,
+      "llm_learning_rate": 3.1803076812305235e-06,
+      "loss": 0.8759,
+      "output_text_mask_loss": 0.0147,
+      "predictor_learning_rate": 0.0003180307681230523,
+      "step": 4196
+    },
+    {
+      "epoch": 0.43,
+      "gumbel_tau": 0.3716,
+      "image_masks_loss": 0.0163,
+      "llm_learning_rate": 3.1795093347496864e-06,
+      "loss": 0.7923,
+      "output_text_mask_loss": 0.0278,
+      "predictor_learning_rate": 0.0003179509334749686,
+      "step": 4197
+    },
+    {
+      "epoch": 0.43,
+      "gumbel_tau": 0.3715,
+      "image_masks_loss": 0.0189,
+      "llm_learning_rate": 3.1787109134397687e-06,
+      "loss": 0.7538,
+      "output_text_mask_loss": 0.0098,
+      "predictor_learning_rate": 0.0003178710913439768,
+      "step": 4198
+    },
+    {
+      "epoch": 0.43,
+      "gumbel_tau": 0.3714,
+      "image_masks_loss": 0.0103,
+      "llm_learning_rate": 3.1779124173886934e-06,
+      "loss": 0.7777,
+      "output_text_mask_loss": 0.0156,
+      "predictor_learning_rate": 0.0003177912417388693,
+      "step": 4199
+    },
+    {
+      "epoch": 0.43,
+      "gumbel_tau": 0.3713,
+      "image_masks_loss": 0.0298,
+      "llm_learning_rate": 3.1771138466843925e-06,
+      "loss": 0.6173,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.0003177113846684392,
+      "step": 4200
+    },
+    {
+      "epoch": 0.43,
+      "gumbel_tau": 0.3713,
+      "image_masks_loss": 0.0393,
+      "llm_learning_rate": 3.1763152014148064e-06,
+      "loss": 0.7315,
+      "output_text_mask_loss": 0.0025,
+      "predictor_learning_rate": 0.00031763152014148064,
+      "step": 4201
+    },
+    {
+      "epoch": 0.43,
+      "gumbel_tau": 0.3712,
+      "image_masks_loss": 0.0236,
+      "llm_learning_rate": 3.175516481667884e-06,
+      "loss": 0.7975,
+      "output_text_mask_loss": 0.0072,
+      "predictor_learning_rate": 0.00031755164816678837,
+      "step": 4202
+    },
+    {
+      "epoch": 0.43,
+      "gumbel_tau": 0.3711,
+      "image_masks_loss": 0.0259,
+      "llm_learning_rate": 3.174717687531582e-06,
+      "loss": 0.7327,
+      "output_text_mask_loss": 0.013,
+      "predictor_learning_rate": 0.0003174717687531582,
+      "step": 4203
+    },
+    {
+      "epoch": 0.43,
+      "gumbel_tau": 0.371,
+      "image_masks_loss": 0.0093,
+      "llm_learning_rate": 3.1739188190938656e-06,
+      "loss": 0.8015,
+      "output_text_mask_loss": 0.0223,
+      "predictor_learning_rate": 0.00031739188190938654,
+      "step": 4204
+    },
+    {
+      "epoch": 0.43,
+      "gumbel_tau": 0.3709,
+      "image_masks_loss": 0.0194,
+      "llm_learning_rate": 3.1731198764427075e-06,
+      "loss": 0.7324,
+      "output_text_mask_loss": 0.0381,
+      "predictor_learning_rate": 0.0003173119876442707,
+      "step": 4205
+    },
+    {
+      "epoch": 0.43,
+      "gumbel_tau": 0.3708,
+      "image_masks_loss": 0.0235,
+      "llm_learning_rate": 3.1723208596660894e-06,
+      "loss": 0.7982,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.0003172320859666089,
+      "step": 4206
+    },
+    {
+      "epoch": 0.43,
+      "gumbel_tau": 0.3707,
+      "image_masks_loss": 0.0225,
+      "llm_learning_rate": 3.1715217688520006e-06,
+      "loss": 0.6514,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.0003171521768852,
+      "step": 4207
+    },
+    {
+      "epoch": 0.43,
+      "gumbel_tau": 0.3706,
+      "image_masks_loss": 0.0133,
+      "llm_learning_rate": 3.1707226040884385e-06,
+      "loss": 0.8324,
+      "output_text_mask_loss": 0.0125,
+      "predictor_learning_rate": 0.00031707226040884385,
+      "step": 4208
+    },
+    {
+      "epoch": 0.43,
+      "gumbel_tau": 0.3706,
+      "image_masks_loss": 0.0193,
+      "llm_learning_rate": 3.16992336546341e-06,
+      "loss": 0.7028,
+      "output_text_mask_loss": 0.0082,
+      "predictor_learning_rate": 0.00031699233654634097,
+      "step": 4209
+    },
+    {
+      "epoch": 0.43,
+      "gumbel_tau": 0.3705,
+      "image_masks_loss": 0.0211,
+      "llm_learning_rate": 3.1691240530649282e-06,
+      "loss": 0.8469,
+      "output_text_mask_loss": 0.02,
+      "predictor_learning_rate": 0.00031691240530649277,
+      "step": 4210
+    },
+    {
+      "epoch": 0.43,
+      "gumbel_tau": 0.3704,
+      "image_masks_loss": 0.0089,
+      "llm_learning_rate": 3.1683246669810146e-06,
+      "loss": 0.7661,
+      "output_text_mask_loss": 0.0326,
+      "predictor_learning_rate": 0.00031683246669810146,
+      "step": 4211
+    },
+    {
+      "epoch": 0.43,
+      "gumbel_tau": 0.3703,
+      "image_masks_loss": 0.0054,
+      "llm_learning_rate": 3.167525207299701e-06,
+      "loss": 0.7208,
+      "output_text_mask_loss": 0.0046,
+      "predictor_learning_rate": 0.00031675252072997005,
+      "step": 4212
+    },
+    {
+      "epoch": 0.43,
+      "gumbel_tau": 0.3702,
+      "image_masks_loss": 0.0191,
+      "llm_learning_rate": 3.166725674109024e-06,
+      "loss": 0.8633,
+      "output_text_mask_loss": 0.0219,
+      "predictor_learning_rate": 0.0003166725674109024,
+      "step": 4213
+    },
+    {
+      "epoch": 0.43,
+      "gumbel_tau": 0.3701,
+      "image_masks_loss": 0.0163,
+      "llm_learning_rate": 3.1659260674970303e-06,
+      "loss": 0.7412,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 0.00031659260674970303,
+      "step": 4214
+    },
+    {
+      "epoch": 0.43,
+      "gumbel_tau": 0.37,
+      "image_masks_loss": 0.0137,
+      "llm_learning_rate": 3.1651263875517756e-06,
+      "loss": 0.7125,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.00031651263875517757,
+      "step": 4215
+    },
+    {
+      "epoch": 0.43,
+      "gumbel_tau": 0.3699,
+      "image_masks_loss": 0.0113,
+      "llm_learning_rate": 3.164326634361322e-06,
+      "loss": 0.8138,
+      "output_text_mask_loss": 0.0729,
+      "predictor_learning_rate": 0.0003164326634361322,
+      "step": 4216
+    },
+    {
+      "epoch": 0.43,
+      "gumbel_tau": 0.3699,
+      "image_masks_loss": 0.0134,
+      "llm_learning_rate": 3.163526808013738e-06,
+      "loss": 0.8264,
+      "output_text_mask_loss": 0.0229,
+      "predictor_learning_rate": 0.0003163526808013738,
+      "step": 4217
+    },
+    {
+      "epoch": 0.43,
+      "gumbel_tau": 0.3698,
+      "image_masks_loss": 0.0278,
+      "llm_learning_rate": 3.1627269085971057e-06,
+      "loss": 0.8216,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00031627269085971057,
+      "step": 4218
+    },
+    {
+      "epoch": 0.43,
+      "gumbel_tau": 0.3697,
+      "image_masks_loss": 0.0107,
+      "llm_learning_rate": 3.1619269361995093e-06,
+      "loss": 0.6966,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0003161926936199509,
+      "step": 4219
+    },
+    {
+      "epoch": 0.43,
+      "gumbel_tau": 0.3696,
+      "image_masks_loss": 0.0213,
+      "llm_learning_rate": 3.1611268909090443e-06,
+      "loss": 0.7643,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.0003161126890909044,
+      "step": 4220
+    },
+    {
+      "epoch": 0.43,
+      "gumbel_tau": 0.3695,
+      "image_masks_loss": 0.0317,
+      "llm_learning_rate": 3.1603267728138133e-06,
+      "loss": 0.7475,
+      "output_text_mask_loss": 0.0231,
+      "predictor_learning_rate": 0.0003160326772813813,
+      "step": 4221
+    },
+    {
+      "epoch": 0.43,
+      "gumbel_tau": 0.3694,
+      "image_masks_loss": 0.0091,
+      "llm_learning_rate": 3.159526582001927e-06,
+      "loss": 0.8026,
+      "output_text_mask_loss": 0.0012,
+      "predictor_learning_rate": 0.0003159526582001927,
+      "step": 4222
+    },
+    {
+      "epoch": 0.43,
+      "gumbel_tau": 0.3693,
+      "image_masks_loss": 0.0077,
+      "llm_learning_rate": 3.1587263185615057e-06,
+      "loss": 0.7746,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00031587263185615055,
+      "step": 4223
+    },
+    {
+      "epoch": 0.43,
+      "gumbel_tau": 0.3692,
+      "image_masks_loss": 0.012,
+      "llm_learning_rate": 3.1579259825806736e-06,
+      "loss": 0.7574,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00031579259825806734,
+      "step": 4224
+    },
+    {
+      "epoch": 0.43,
+      "gumbel_tau": 0.3692,
+      "image_masks_loss": 0.0288,
+      "llm_learning_rate": 3.1571255741475676e-06,
+      "loss": 0.7073,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0003157125574147567,
+      "step": 4225
+    },
+    {
+      "epoch": 0.43,
+      "gumbel_tau": 0.3691,
+      "image_masks_loss": 0.0257,
+      "llm_learning_rate": 3.15632509335033e-06,
+      "loss": 0.7524,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.00031563250933503296,
+      "step": 4226
+    },
+    {
+      "epoch": 0.43,
+      "gumbel_tau": 0.369,
+      "image_masks_loss": 0.0135,
+      "llm_learning_rate": 3.15552454027711e-06,
+      "loss": 0.7303,
+      "output_text_mask_loss": 0.021,
+      "predictor_learning_rate": 0.00031555245402771097,
+      "step": 4227
+    },
+    {
+      "epoch": 0.43,
+      "gumbel_tau": 0.3689,
+      "image_masks_loss": 0.0221,
+      "llm_learning_rate": 3.1547239150160687e-06,
+      "loss": 0.7655,
+      "output_text_mask_loss": 0.0378,
+      "predictor_learning_rate": 0.0003154723915016068,
+      "step": 4228
+    },
+    {
+      "epoch": 0.43,
+      "gumbel_tau": 0.3688,
+      "image_masks_loss": 0.0068,
+      "llm_learning_rate": 3.153923217655371e-06,
+      "loss": 0.7507,
+      "output_text_mask_loss": 0.0355,
+      "predictor_learning_rate": 0.00031539232176553707,
+      "step": 4229
+    },
+    {
+      "epoch": 0.43,
+      "gumbel_tau": 0.3687,
+      "image_masks_loss": 0.0234,
+      "llm_learning_rate": 3.1531224482831924e-06,
+      "loss": 0.7209,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.00031531224482831923,
+      "step": 4230
+    },
+    {
+      "epoch": 0.43,
+      "gumbel_tau": 0.3686,
+      "image_masks_loss": 0.0233,
+      "llm_learning_rate": 3.152321606987714e-06,
+      "loss": 0.7496,
+      "output_text_mask_loss": 0.0221,
+      "predictor_learning_rate": 0.0003152321606987714,
+      "step": 4231
+    },
+    {
+      "epoch": 0.43,
+      "gumbel_tau": 0.3685,
+      "image_masks_loss": 0.0108,
+      "llm_learning_rate": 3.151520693857128e-06,
+      "loss": 0.7813,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0003151520693857128,
+      "step": 4232
+    },
+    {
+      "epoch": 0.43,
+      "gumbel_tau": 0.3685,
+      "image_masks_loss": 0.0218,
+      "llm_learning_rate": 3.1507197089796322e-06,
+      "loss": 0.8743,
+      "output_text_mask_loss": 0.0275,
+      "predictor_learning_rate": 0.0003150719708979632,
+      "step": 4233
+    },
+    {
+      "epoch": 0.43,
+      "gumbel_tau": 0.3684,
+      "image_masks_loss": 0.0294,
+      "llm_learning_rate": 3.1499186524434315e-06,
+      "loss": 0.8088,
+      "output_text_mask_loss": 0.0025,
+      "predictor_learning_rate": 0.00031499186524434317,
+      "step": 4234
+    },
+    {
+      "epoch": 0.43,
+      "gumbel_tau": 0.3683,
+      "image_masks_loss": 0.0115,
+      "llm_learning_rate": 3.1491175243367427e-06,
+      "loss": 0.787,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.00031491175243367425,
+      "step": 4235
+    },
+    {
+      "epoch": 0.43,
+      "gumbel_tau": 0.3682,
+      "image_masks_loss": 0.0294,
+      "llm_learning_rate": 3.1483163247477856e-06,
+      "loss": 0.8617,
+      "output_text_mask_loss": 0.0568,
+      "predictor_learning_rate": 0.0003148316324747785,
+      "step": 4236
+    },
+    {
+      "epoch": 0.43,
+      "gumbel_tau": 0.3681,
+      "image_masks_loss": 0.0203,
+      "llm_learning_rate": 3.14751505376479e-06,
+      "loss": 0.7406,
+      "output_text_mask_loss": 0.0324,
+      "predictor_learning_rate": 0.000314751505376479,
+      "step": 4237
+    },
+    {
+      "epoch": 0.43,
+      "gumbel_tau": 0.368,
+      "image_masks_loss": 0.0183,
+      "llm_learning_rate": 3.1467137114759954e-06,
+      "loss": 0.7884,
+      "output_text_mask_loss": 0.0317,
+      "predictor_learning_rate": 0.0003146713711475995,
+      "step": 4238
+    },
+    {
+      "epoch": 0.43,
+      "gumbel_tau": 0.3679,
+      "image_masks_loss": 0.022,
+      "llm_learning_rate": 3.1459122979696467e-06,
+      "loss": 0.7805,
+      "output_text_mask_loss": 0.024,
+      "predictor_learning_rate": 0.00031459122979696463,
+      "step": 4239
+    },
+    {
+      "epoch": 0.43,
+      "gumbel_tau": 0.3679,
+      "image_masks_loss": 0.0246,
+      "llm_learning_rate": 3.145110813333996e-06,
+      "loss": 0.7748,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00031451108133339955,
+      "step": 4240
+    },
+    {
+      "epoch": 0.43,
+      "gumbel_tau": 0.3678,
+      "image_masks_loss": 0.0084,
+      "llm_learning_rate": 3.1443092576573057e-06,
+      "loss": 0.7469,
+      "output_text_mask_loss": 0.0278,
+      "predictor_learning_rate": 0.00031443092576573055,
+      "step": 4241
+    },
+    {
+      "epoch": 0.43,
+      "gumbel_tau": 0.3677,
+      "image_masks_loss": 0.0133,
+      "llm_learning_rate": 3.1435076310278457e-06,
+      "loss": 0.7197,
+      "output_text_mask_loss": 0.0196,
+      "predictor_learning_rate": 0.00031435076310278455,
+      "step": 4242
+    },
+    {
+      "epoch": 0.43,
+      "gumbel_tau": 0.3676,
+      "image_masks_loss": 0.0232,
+      "llm_learning_rate": 3.142705933533891e-06,
+      "loss": 0.7194,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0003142705933533891,
+      "step": 4243
+    },
+    {
+      "epoch": 0.43,
+      "gumbel_tau": 0.3675,
+      "image_masks_loss": 0.0175,
+      "llm_learning_rate": 3.1419041652637285e-06,
+      "loss": 0.7421,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.0003141904165263728,
+      "step": 4244
+    },
+    {
+      "epoch": 0.43,
+      "gumbel_tau": 0.3674,
+      "image_masks_loss": 0.0126,
+      "llm_learning_rate": 3.1411023263056485e-06,
+      "loss": 0.6631,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.0003141102326305648,
+      "step": 4245
+    },
+    {
+      "epoch": 0.44,
+      "gumbel_tau": 0.3673,
+      "image_masks_loss": 0.0128,
+      "llm_learning_rate": 3.140300416747953e-06,
+      "loss": 0.7676,
+      "output_text_mask_loss": 0.0526,
+      "predictor_learning_rate": 0.00031403004167479526,
+      "step": 4246
+    },
+    {
+      "epoch": 0.44,
+      "gumbel_tau": 0.3672,
+      "image_masks_loss": 0.0197,
+      "llm_learning_rate": 3.139498436678948e-06,
+      "loss": 0.7517,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.0003139498436678948,
+      "step": 4247
+    },
+    {
+      "epoch": 0.44,
+      "gumbel_tau": 0.3672,
+      "image_masks_loss": 0.0032,
+      "llm_learning_rate": 3.1386963861869524e-06,
+      "loss": 0.6232,
+      "output_text_mask_loss": 0.0004,
+      "predictor_learning_rate": 0.00031386963861869524,
+      "step": 4248
+    },
+    {
+      "epoch": 0.44,
+      "gumbel_tau": 0.3671,
+      "image_masks_loss": 0.0098,
+      "llm_learning_rate": 3.1378942653602875e-06,
+      "loss": 0.7112,
+      "output_text_mask_loss": 0.0101,
+      "predictor_learning_rate": 0.0003137894265360287,
+      "step": 4249
+    },
+    {
+      "epoch": 0.44,
+      "gumbel_tau": 0.367,
+      "image_masks_loss": 0.0216,
+      "llm_learning_rate": 3.1370920742872846e-06,
+      "loss": 0.6511,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00031370920742872845,
+      "step": 4250
+    },
+    {
+      "epoch": 0.44,
+      "gumbel_tau": 0.3669,
+      "image_masks_loss": 0.0153,
+      "llm_learning_rate": 3.1362898130562842e-06,
+      "loss": 0.6554,
+      "output_text_mask_loss": NaN,
+      "predictor_learning_rate": 0.0003136289813056284,
+      "step": 4251
+    },
+    {
+      "epoch": 0.44,
+      "gumbel_tau": 0.3668,
+      "image_masks_loss": 0.0172,
+      "llm_learning_rate": 3.1354874817556323e-06,
+      "loss": 0.797,
+      "output_text_mask_loss": 0.0698,
+      "predictor_learning_rate": 0.00031354874817556325,
+      "step": 4252
+    },
+    {
+      "epoch": 0.44,
+      "gumbel_tau": 0.3667,
+      "image_masks_loss": 0.0048,
+      "llm_learning_rate": 3.134685080473683e-06,
+      "loss": 0.7917,
+      "output_text_mask_loss": 0.0372,
+      "predictor_learning_rate": 0.0003134685080473683,
+      "step": 4253
+    },
+    {
+      "epoch": 0.44,
+      "gumbel_tau": 0.3666,
+      "image_masks_loss": 0.0223,
+      "llm_learning_rate": 3.1338826092987994e-06,
+      "loss": 0.7723,
+      "output_text_mask_loss": 0.0488,
+      "predictor_learning_rate": 0.0003133882609298799,
+      "step": 4254
+    },
+    {
+      "epoch": 0.44,
+      "gumbel_tau": 0.3666,
+      "image_masks_loss": 0.0116,
+      "llm_learning_rate": 3.1330800683193507e-06,
+      "loss": 0.791,
+      "output_text_mask_loss": 0.0029,
+      "predictor_learning_rate": 0.00031330800683193507,
+      "step": 4255
+    },
+    {
+      "epoch": 0.44,
+      "gumbel_tau": 0.3665,
+      "image_masks_loss": 0.0151,
+      "llm_learning_rate": 3.1322774576237147e-06,
+      "loss": 0.7055,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.00031322774576237147,
+      "step": 4256
+    },
+    {
+      "epoch": 0.44,
+      "gumbel_tau": 0.3664,
+      "image_masks_loss": 0.0076,
+      "llm_learning_rate": 3.131474777300277e-06,
+      "loss": 0.7499,
+      "output_text_mask_loss": 0.0376,
+      "predictor_learning_rate": 0.0003131474777300277,
+      "step": 4257
+    },
+    {
+      "epoch": 0.44,
+      "gumbel_tau": 0.3663,
+      "image_masks_loss": 0.021,
+      "llm_learning_rate": 3.13067202743743e-06,
+      "loss": 0.7831,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.000313067202743743,
+      "step": 4258
+    },
+    {
+      "epoch": 0.44,
+      "gumbel_tau": 0.3662,
+      "image_masks_loss": 0.0083,
+      "llm_learning_rate": 3.129869208123575e-06,
+      "loss": 0.6869,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.0003129869208123575,
+      "step": 4259
+    },
+    {
+      "epoch": 0.44,
+      "gumbel_tau": 0.3661,
+      "image_masks_loss": 0.0058,
+      "llm_learning_rate": 3.1290663194471188e-06,
+      "loss": 0.7105,
+      "output_text_mask_loss": 0.0324,
+      "predictor_learning_rate": 0.00031290663194471184,
+      "step": 4260
+    },
+    {
+      "epoch": 0.44,
+      "gumbel_tau": 0.366,
+      "image_masks_loss": 0.0199,
+      "llm_learning_rate": 3.128263361496479e-06,
+      "loss": 0.744,
+      "output_text_mask_loss": 0.0089,
+      "predictor_learning_rate": 0.0003128263361496479,
+      "step": 4261
+    },
+    {
+      "epoch": 0.44,
+      "gumbel_tau": 0.366,
+      "image_masks_loss": 0.0248,
+      "llm_learning_rate": 3.1274603343600784e-06,
+      "loss": 0.7307,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00031274603343600785,
+      "step": 4262
+    },
+    {
+      "epoch": 0.44,
+      "gumbel_tau": 0.3659,
+      "image_masks_loss": 0.0064,
+      "llm_learning_rate": 3.1266572381263483e-06,
+      "loss": 0.8116,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.00031266572381263483,
+      "step": 4263
+    },
+    {
+      "epoch": 0.44,
+      "gumbel_tau": 0.3658,
+      "image_masks_loss": 0.0095,
+      "llm_learning_rate": 3.125854072883727e-06,
+      "loss": 0.8096,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.0003125854072883727,
+      "step": 4264
+    },
+    {
+      "epoch": 0.44,
+      "gumbel_tau": 0.3657,
+      "image_masks_loss": 0.0104,
+      "llm_learning_rate": 3.1250508387206617e-06,
+      "loss": 0.7688,
+      "output_text_mask_loss": 0.0763,
+      "predictor_learning_rate": 0.00031250508387206617,
+      "step": 4265
+    },
+    {
+      "epoch": 0.44,
+      "gumbel_tau": 0.3656,
+      "image_masks_loss": 0.0211,
+      "llm_learning_rate": 3.124247535725605e-06,
+      "loss": 0.7392,
+      "output_text_mask_loss": 0.0004,
+      "predictor_learning_rate": 0.00031242475357256053,
+      "step": 4266
+    },
+    {
+      "epoch": 0.44,
+      "gumbel_tau": 0.3655,
+      "image_masks_loss": 0.0136,
+      "llm_learning_rate": 3.12344416398702e-06,
+      "loss": 0.7253,
+      "output_text_mask_loss": 0.0033,
+      "predictor_learning_rate": 0.00031234441639870197,
+      "step": 4267
+    },
+    {
+      "epoch": 0.44,
+      "gumbel_tau": 0.3654,
+      "image_masks_loss": 0.0067,
+      "llm_learning_rate": 3.122640723593375e-06,
+      "loss": 0.7151,
+      "output_text_mask_loss": 0.0298,
+      "predictor_learning_rate": 0.00031226407235933747,
+      "step": 4268
+    },
+    {
+      "epoch": 0.44,
+      "gumbel_tau": 0.3653,
+      "image_masks_loss": 0.0198,
+      "llm_learning_rate": 3.121837214633146e-06,
+      "loss": 0.7725,
+      "output_text_mask_loss": 0.0172,
+      "predictor_learning_rate": 0.0003121837214633146,
+      "step": 4269
+    },
+    {
+      "epoch": 0.44,
+      "gumbel_tau": 0.3653,
+      "image_masks_loss": 0.0124,
+      "llm_learning_rate": 3.121033637194818e-06,
+      "loss": 0.7504,
+      "output_text_mask_loss": 0.0193,
+      "predictor_learning_rate": 0.00031210336371948177,
+      "step": 4270
+    },
+    {
+      "epoch": 0.44,
+      "gumbel_tau": 0.3652,
+      "image_masks_loss": 0.01,
+      "llm_learning_rate": 3.120229991366883e-06,
+      "loss": 0.7285,
+      "output_text_mask_loss": 0.0057,
+      "predictor_learning_rate": 0.0003120229991366882,
+      "step": 4271
+    },
+    {
+      "epoch": 0.44,
+      "gumbel_tau": 0.3651,
+      "image_masks_loss": 0.016,
+      "llm_learning_rate": 3.1194262772378393e-06,
+      "loss": 0.7648,
+      "output_text_mask_loss": 0.0504,
+      "predictor_learning_rate": 0.0003119426277237839,
+      "step": 4272
+    },
+    {
+      "epoch": 0.44,
+      "gumbel_tau": 0.365,
+      "image_masks_loss": 0.0144,
+      "llm_learning_rate": 3.1186224948961945e-06,
+      "loss": 0.7165,
+      "output_text_mask_loss": 0.0096,
+      "predictor_learning_rate": 0.00031186224948961944,
+      "step": 4273
+    },
+    {
+      "epoch": 0.44,
+      "gumbel_tau": 0.3649,
+      "image_masks_loss": 0.0095,
+      "llm_learning_rate": 3.1178186444304625e-06,
+      "loss": 0.8297,
+      "output_text_mask_loss": 0.0416,
+      "predictor_learning_rate": 0.00031178186444304624,
+      "step": 4274
+    },
+    {
+      "epoch": 0.44,
+      "gumbel_tau": 0.3648,
+      "image_masks_loss": 0.0246,
+      "llm_learning_rate": 3.1170147259291654e-06,
+      "loss": 0.7434,
+      "output_text_mask_loss": 0.0027,
+      "predictor_learning_rate": 0.0003117014725929165,
+      "step": 4275
+    },
+    {
+      "epoch": 0.44,
+      "gumbel_tau": 0.3647,
+      "image_masks_loss": 0.0141,
+      "llm_learning_rate": 3.1162107394808322e-06,
+      "loss": 0.6913,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0003116210739480832,
+      "step": 4276
+    },
+    {
+      "epoch": 0.44,
+      "gumbel_tau": 0.3647,
+      "image_masks_loss": 0.0292,
+      "llm_learning_rate": 3.115406685174e-06,
+      "loss": 0.6446,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.0003115406685174,
+      "step": 4277
+    },
+    {
+      "epoch": 0.44,
+      "gumbel_tau": 0.3646,
+      "image_masks_loss": 0.0242,
+      "llm_learning_rate": 3.1146025630972126e-06,
+      "loss": 0.653,
+      "output_text_mask_loss": 0.021,
+      "predictor_learning_rate": 0.0003114602563097212,
+      "step": 4278
+    },
+    {
+      "epoch": 0.44,
+      "gumbel_tau": 0.3645,
+      "image_masks_loss": 0.0246,
+      "llm_learning_rate": 3.113798373339022e-06,
+      "loss": 0.7299,
+      "output_text_mask_loss": 0.0102,
+      "predictor_learning_rate": 0.00031137983733390217,
+      "step": 4279
+    },
+    {
+      "epoch": 0.44,
+      "gumbel_tau": 0.3644,
+      "image_masks_loss": 0.0217,
+      "llm_learning_rate": 3.1129941159879873e-06,
+      "loss": 0.7027,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00031129941159879873,
+      "step": 4280
+    },
+    {
+      "epoch": 0.44,
+      "gumbel_tau": 0.3643,
+      "image_masks_loss": 0.0139,
+      "llm_learning_rate": 3.112189791132676e-06,
+      "loss": 0.7157,
+      "output_text_mask_loss": 0.0193,
+      "predictor_learning_rate": 0.00031121897911326756,
+      "step": 4281
+    },
+    {
+      "epoch": 0.44,
+      "gumbel_tau": 0.3642,
+      "image_masks_loss": 0.021,
+      "llm_learning_rate": 3.1113853988616605e-06,
+      "loss": 0.649,
+      "output_text_mask_loss": 0.0023,
+      "predictor_learning_rate": 0.000311138539886166,
+      "step": 4282
+    },
+    {
+      "epoch": 0.44,
+      "gumbel_tau": 0.3641,
+      "image_masks_loss": 0.0177,
+      "llm_learning_rate": 3.1105809392635233e-06,
+      "loss": 0.7475,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.00031105809392635233,
+      "step": 4283
+    },
+    {
+      "epoch": 0.44,
+      "gumbel_tau": 0.3641,
+      "image_masks_loss": 0.0189,
+      "llm_learning_rate": 3.1097764124268528e-06,
+      "loss": 0.8393,
+      "output_text_mask_loss": 0.0248,
+      "predictor_learning_rate": 0.00031097764124268524,
+      "step": 4284
+    },
+    {
+      "epoch": 0.44,
+      "gumbel_tau": 0.364,
+      "image_masks_loss": 0.0146,
+      "llm_learning_rate": 3.1089718184402456e-06,
+      "loss": 0.8063,
+      "output_text_mask_loss": 0.0189,
+      "predictor_learning_rate": 0.00031089718184402457,
+      "step": 4285
+    },
+    {
+      "epoch": 0.44,
+      "gumbel_tau": 0.3639,
+      "image_masks_loss": 0.0155,
+      "llm_learning_rate": 3.108167157392305e-06,
+      "loss": 0.685,
+      "output_text_mask_loss": 0.0374,
+      "predictor_learning_rate": 0.00031081671573923047,
+      "step": 4286
+    },
+    {
+      "epoch": 0.44,
+      "gumbel_tau": 0.3638,
+      "image_masks_loss": 0.0076,
+      "llm_learning_rate": 3.107362429371643e-06,
+      "loss": 0.8567,
+      "output_text_mask_loss": 0.0496,
+      "predictor_learning_rate": 0.0003107362429371643,
+      "step": 4287
+    },
+    {
+      "epoch": 0.44,
+      "gumbel_tau": 0.3637,
+      "image_masks_loss": 0.0187,
+      "llm_learning_rate": 3.1065576344668772e-06,
+      "loss": 0.8162,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0003106557634466877,
+      "step": 4288
+    },
+    {
+      "epoch": 0.44,
+      "gumbel_tau": 0.3636,
+      "image_masks_loss": 0.0176,
+      "llm_learning_rate": 3.1057527727666336e-06,
+      "loss": 0.9168,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0003105752772766633,
+      "step": 4289
+    },
+    {
+      "epoch": 0.44,
+      "gumbel_tau": 0.3635,
+      "image_masks_loss": 0.0088,
+      "llm_learning_rate": 3.1049478443595456e-06,
+      "loss": 0.7128,
+      "output_text_mask_loss": 0.0139,
+      "predictor_learning_rate": 0.00031049478443595453,
+      "step": 4290
+    },
+    {
+      "epoch": 0.44,
+      "gumbel_tau": 0.3635,
+      "image_masks_loss": 0.0202,
+      "llm_learning_rate": 3.104142849334254e-06,
+      "loss": 0.7494,
+      "output_text_mask_loss": 0.0162,
+      "predictor_learning_rate": 0.0003104142849334254,
+      "step": 4291
+    },
+    {
+      "epoch": 0.44,
+      "gumbel_tau": 0.3634,
+      "image_masks_loss": 0.0077,
+      "llm_learning_rate": 3.1033377877794054e-06,
+      "loss": 0.6732,
+      "output_text_mask_loss": 0.0809,
+      "predictor_learning_rate": 0.0003103337787779405,
+      "step": 4292
+    },
+    {
+      "epoch": 0.44,
+      "gumbel_tau": 0.3633,
+      "image_masks_loss": 0.0182,
+      "llm_learning_rate": 3.1025326597836558e-06,
+      "loss": 0.6172,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00031025326597836555,
+      "step": 4293
+    },
+    {
+      "epoch": 0.44,
+      "gumbel_tau": 0.3632,
+      "image_masks_loss": 0.02,
+      "llm_learning_rate": 3.101727465435668e-06,
+      "loss": 0.7013,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 0.00031017274654356676,
+      "step": 4294
+    },
+    {
+      "epoch": 0.44,
+      "gumbel_tau": 0.3631,
+      "image_masks_loss": 0.0169,
+      "llm_learning_rate": 3.1009222048241104e-06,
+      "loss": 0.7179,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.000310092220482411,
+      "step": 4295
+    },
+    {
+      "epoch": 0.44,
+      "gumbel_tau": 0.363,
+      "image_masks_loss": 0.0093,
+      "llm_learning_rate": 3.100116878037662e-06,
+      "loss": 0.7137,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00031001168780376615,
+      "step": 4296
+    },
+    {
+      "epoch": 0.44,
+      "gumbel_tau": 0.3629,
+      "image_masks_loss": 0.0311,
+      "llm_learning_rate": 3.099311485165006e-06,
+      "loss": 0.7796,
+      "output_text_mask_loss": 0.005,
+      "predictor_learning_rate": 0.0003099311485165006,
+      "step": 4297
+    },
+    {
+      "epoch": 0.44,
+      "gumbel_tau": 0.3629,
+      "image_masks_loss": 0.0181,
+      "llm_learning_rate": 3.0985060262948345e-06,
+      "loss": 0.5632,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00030985060262948345,
+      "step": 4298
+    },
+    {
+      "epoch": 0.44,
+      "gumbel_tau": 0.3628,
+      "image_masks_loss": 0.01,
+      "llm_learning_rate": 3.0977005015158456e-06,
+      "loss": 0.7172,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00030977005015158456,
+      "step": 4299
+    },
+    {
+      "epoch": 0.44,
+      "gumbel_tau": 0.3627,
+      "image_masks_loss": 0.0228,
+      "llm_learning_rate": 3.0968949109167475e-06,
+      "loss": 0.741,
+      "output_text_mask_loss": 0.0019,
+      "predictor_learning_rate": 0.0003096894910916747,
+      "step": 4300
+    },
+    {
+      "epoch": 0.44,
+      "gumbel_tau": 0.3626,
+      "image_masks_loss": 0.0241,
+      "llm_learning_rate": 3.0960892545862515e-06,
+      "loss": 0.7339,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0003096089254586251,
+      "step": 4301
+    },
+    {
+      "epoch": 0.44,
+      "gumbel_tau": 0.3625,
+      "image_masks_loss": 0.0201,
+      "llm_learning_rate": 3.095283532613079e-06,
+      "loss": 0.8081,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00030952835326130784,
+      "step": 4302
+    },
+    {
+      "epoch": 0.44,
+      "gumbel_tau": 0.3624,
+      "image_masks_loss": 0.0085,
+      "llm_learning_rate": 3.0944777450859576e-06,
+      "loss": 0.799,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.00030944777450859573,
+      "step": 4303
+    },
+    {
+      "epoch": 0.44,
+      "gumbel_tau": 0.3623,
+      "image_masks_loss": 0.0104,
+      "llm_learning_rate": 3.0936718920936236e-06,
+      "loss": 0.8304,
+      "output_text_mask_loss": 0.0265,
+      "predictor_learning_rate": 0.0003093671892093624,
+      "step": 4304
+    },
+    {
+      "epoch": 0.44,
+      "gumbel_tau": 0.3623,
+      "image_masks_loss": 0.0332,
+      "llm_learning_rate": 3.0928659737248173e-06,
+      "loss": 0.7974,
+      "output_text_mask_loss": 0.006,
+      "predictor_learning_rate": 0.0003092865973724817,
+      "step": 4305
+    },
+    {
+      "epoch": 0.44,
+      "gumbel_tau": 0.3622,
+      "image_masks_loss": 0.0357,
+      "llm_learning_rate": 3.0920599900682907e-06,
+      "loss": 0.7057,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00030920599900682903,
+      "step": 4306
+    },
+    {
+      "epoch": 0.44,
+      "gumbel_tau": 0.3621,
+      "image_masks_loss": 0.0108,
+      "llm_learning_rate": 3.0912539412127984e-06,
+      "loss": 0.7282,
+      "output_text_mask_loss": 0.0431,
+      "predictor_learning_rate": 0.0003091253941212798,
+      "step": 4307
+    },
+    {
+      "epoch": 0.44,
+      "gumbel_tau": 0.362,
+      "image_masks_loss": 0.0059,
+      "llm_learning_rate": 3.0904478272471054e-06,
+      "loss": 0.7441,
+      "output_text_mask_loss": 0.0257,
+      "predictor_learning_rate": 0.00030904478272471053,
+      "step": 4308
+    },
+    {
+      "epoch": 0.44,
+      "gumbel_tau": 0.3619,
+      "image_masks_loss": 0.0053,
+      "llm_learning_rate": 3.0896416482599827e-06,
+      "loss": 0.797,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.0003089641648259983,
+      "step": 4309
+    },
+    {
+      "epoch": 0.44,
+      "gumbel_tau": 0.3618,
+      "image_masks_loss": 0.0452,
+      "llm_learning_rate": 3.0888354043402087e-06,
+      "loss": 0.786,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.00030888354043402086,
+      "step": 4310
+    },
+    {
+      "epoch": 0.44,
+      "gumbel_tau": 0.3617,
+      "image_masks_loss": 0.0082,
+      "llm_learning_rate": 3.0880290955765674e-06,
+      "loss": 0.5917,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.00030880290955765673,
+      "step": 4311
+    },
+    {
+      "epoch": 0.44,
+      "gumbel_tau": 0.3617,
+      "image_masks_loss": 0.0322,
+      "llm_learning_rate": 3.0872227220578527e-06,
+      "loss": 0.8419,
+      "output_text_mask_loss": 0.0309,
+      "predictor_learning_rate": 0.00030872227220578525,
+      "step": 4312
+    },
+    {
+      "epoch": 0.44,
+      "gumbel_tau": 0.3616,
+      "image_masks_loss": 0.0416,
+      "llm_learning_rate": 3.086416283872864e-06,
+      "loss": 0.7479,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.0003086416283872864,
+      "step": 4313
+    },
+    {
+      "epoch": 0.44,
+      "gumbel_tau": 0.3615,
+      "image_masks_loss": 0.0175,
+      "llm_learning_rate": 3.085609781110408e-06,
+      "loss": 0.8034,
+      "output_text_mask_loss": 0.0479,
+      "predictor_learning_rate": 0.00030856097811104077,
+      "step": 4314
+    },
+    {
+      "epoch": 0.44,
+      "gumbel_tau": 0.3614,
+      "image_masks_loss": 0.0093,
+      "llm_learning_rate": 3.084803213859298e-06,
+      "loss": 0.6756,
+      "output_text_mask_loss": 0.0322,
+      "predictor_learning_rate": 0.0003084803213859298,
+      "step": 4315
+    },
+    {
+      "epoch": 0.44,
+      "gumbel_tau": 0.3613,
+      "image_masks_loss": 0.0083,
+      "llm_learning_rate": 3.083996582208356e-06,
+      "loss": 0.8208,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 0.00030839965822083553,
+      "step": 4316
+    },
+    {
+      "epoch": 0.44,
+      "gumbel_tau": 0.3612,
+      "image_masks_loss": 0.0198,
+      "llm_learning_rate": 3.0831898862464094e-06,
+      "loss": 0.8201,
+      "output_text_mask_loss": 0.0198,
+      "predictor_learning_rate": 0.0003083189886246409,
+      "step": 4317
+    },
+    {
+      "epoch": 0.44,
+      "gumbel_tau": 0.3611,
+      "image_masks_loss": 0.0248,
+      "llm_learning_rate": 3.0823831260622937e-06,
+      "loss": 0.7264,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.0003082383126062294,
+      "step": 4318
+    },
+    {
+      "epoch": 0.44,
+      "gumbel_tau": 0.3611,
+      "image_masks_loss": 0.0222,
+      "llm_learning_rate": 3.081576301744851e-06,
+      "loss": 0.7634,
+      "output_text_mask_loss": 0.0168,
+      "predictor_learning_rate": 0.00030815763017448506,
+      "step": 4319
+    },
+    {
+      "epoch": 0.44,
+      "gumbel_tau": 0.361,
+      "image_masks_loss": 0.0075,
+      "llm_learning_rate": 3.08076941338293e-06,
+      "loss": 0.7527,
+      "output_text_mask_loss": 0.0467,
+      "predictor_learning_rate": 0.000308076941338293,
+      "step": 4320
+    },
+    {
+      "epoch": 0.44,
+      "gumbel_tau": 0.3609,
+      "image_masks_loss": 0.0128,
+      "llm_learning_rate": 3.079962461065388e-06,
+      "loss": 0.7707,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00030799624610653874,
+      "step": 4321
+    },
+    {
+      "epoch": 0.44,
+      "gumbel_tau": 0.3608,
+      "image_masks_loss": 0.0112,
+      "llm_learning_rate": 3.0791554448810877e-06,
+      "loss": 0.7461,
+      "output_text_mask_loss": 0.0198,
+      "predictor_learning_rate": 0.00030791554448810876,
+      "step": 4322
+    },
+    {
+      "epoch": 0.44,
+      "gumbel_tau": 0.3607,
+      "image_masks_loss": 0.0099,
+      "llm_learning_rate": 3.0783483649189e-06,
+      "loss": 0.746,
+      "output_text_mask_loss": 0.0231,
+      "predictor_learning_rate": 0.00030783483649188996,
+      "step": 4323
+    },
+    {
+      "epoch": 0.44,
+      "gumbel_tau": 0.3606,
+      "image_masks_loss": 0.037,
+      "llm_learning_rate": 3.0775412212677023e-06,
+      "loss": 0.797,
+      "output_text_mask_loss": 0.0086,
+      "predictor_learning_rate": 0.0003077541221267702,
+      "step": 4324
+    },
+    {
+      "epoch": 0.44,
+      "gumbel_tau": 0.3606,
+      "image_masks_loss": 0.0115,
+      "llm_learning_rate": 3.0767340140163783e-06,
+      "loss": 0.66,
+      "output_text_mask_loss": 0.0618,
+      "predictor_learning_rate": 0.0003076734014016378,
+      "step": 4325
+    },
+    {
+      "epoch": 0.44,
+      "gumbel_tau": 0.3605,
+      "image_masks_loss": 0.0244,
+      "llm_learning_rate": 3.0759267432538204e-06,
+      "loss": 0.6802,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.000307592674325382,
+      "step": 4326
+    },
+    {
+      "epoch": 0.44,
+      "gumbel_tau": 0.3604,
+      "image_masks_loss": 0.0039,
+      "llm_learning_rate": 3.075119409068927e-06,
+      "loss": 0.7637,
+      "output_text_mask_loss": 0.0004,
+      "predictor_learning_rate": 0.00030751194090689266,
+      "step": 4327
+    },
+    {
+      "epoch": 0.44,
+      "gumbel_tau": 0.3603,
+      "image_masks_loss": 0.0054,
+      "llm_learning_rate": 3.074312011550603e-06,
+      "loss": 0.7561,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.00030743120115506026,
+      "step": 4328
+    },
+    {
+      "epoch": 0.44,
+      "gumbel_tau": 0.3602,
+      "image_masks_loss": 0.0292,
+      "llm_learning_rate": 3.0735045507877615e-06,
+      "loss": 0.7636,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.0003073504550787761,
+      "step": 4329
+    },
+    {
+      "epoch": 0.44,
+      "gumbel_tau": 0.3601,
+      "image_masks_loss": 0.0058,
+      "llm_learning_rate": 3.0726970268693208e-06,
+      "loss": 0.7633,
+      "output_text_mask_loss": 0.0322,
+      "predictor_learning_rate": 0.00030726970268693207,
+      "step": 4330
+    },
+    {
+      "epoch": 0.44,
+      "gumbel_tau": 0.36,
+      "image_masks_loss": 0.0193,
+      "llm_learning_rate": 3.0718894398842076e-06,
+      "loss": 0.6796,
+      "output_text_mask_loss": 0.0025,
+      "predictor_learning_rate": 0.00030718894398842076,
+      "step": 4331
+    },
+    {
+      "epoch": 0.44,
+      "gumbel_tau": 0.36,
+      "image_masks_loss": 0.0349,
+      "llm_learning_rate": 3.0710817899213563e-06,
+      "loss": 0.797,
+      "output_text_mask_loss": 0.0139,
+      "predictor_learning_rate": 0.0003071081789921356,
+      "step": 4332
+    },
+    {
+      "epoch": 0.44,
+      "gumbel_tau": 0.3599,
+      "image_masks_loss": 0.0143,
+      "llm_learning_rate": 3.070274077069706e-06,
+      "loss": 0.6456,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.0003070274077069706,
+      "step": 4333
+    },
+    {
+      "epoch": 0.44,
+      "gumbel_tau": 0.3598,
+      "image_masks_loss": 0.0106,
+      "llm_learning_rate": 3.0694663014182035e-06,
+      "loss": 0.7148,
+      "output_text_mask_loss": 0.0012,
+      "predictor_learning_rate": 0.00030694663014182034,
+      "step": 4334
+    },
+    {
+      "epoch": 0.44,
+      "gumbel_tau": 0.3597,
+      "image_masks_loss": 0.0408,
+      "llm_learning_rate": 3.0686584630558035e-06,
+      "loss": 0.6977,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0003068658463055803,
+      "step": 4335
+    },
+    {
+      "epoch": 0.44,
+      "gumbel_tau": 0.3596,
+      "image_masks_loss": 0.0303,
+      "llm_learning_rate": 3.067850562071468e-06,
+      "loss": 0.6969,
+      "output_text_mask_loss": 0.002,
+      "predictor_learning_rate": 0.00030678505620714675,
+      "step": 4336
+    },
+    {
+      "epoch": 0.44,
+      "gumbel_tau": 0.3595,
+      "image_masks_loss": 0.0162,
+      "llm_learning_rate": 3.067042598554163e-06,
+      "loss": 0.7659,
+      "output_text_mask_loss": 0.0422,
+      "predictor_learning_rate": 0.0003067042598554163,
+      "step": 4337
+    },
+    {
+      "epoch": 0.44,
+      "gumbel_tau": 0.3594,
+      "image_masks_loss": 0.0343,
+      "llm_learning_rate": 3.066234572592864e-06,
+      "loss": 0.7125,
+      "output_text_mask_loss": 0.0425,
+      "predictor_learning_rate": 0.00030662345725928643,
+      "step": 4338
+    },
+    {
+      "epoch": 0.44,
+      "gumbel_tau": 0.3594,
+      "image_masks_loss": 0.0372,
+      "llm_learning_rate": 3.065426484276553e-06,
+      "loss": 0.7865,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.0003065426484276553,
+      "step": 4339
+    },
+    {
+      "epoch": 0.44,
+      "gumbel_tau": 0.3593,
+      "image_masks_loss": 0.0288,
+      "llm_learning_rate": 3.0646183336942182e-06,
+      "loss": 0.7854,
+      "output_text_mask_loss": 0.0242,
+      "predictor_learning_rate": 0.0003064618333694218,
+      "step": 4340
+    },
+    {
+      "epoch": 0.44,
+      "gumbel_tau": 0.3592,
+      "image_masks_loss": 0.0097,
+      "llm_learning_rate": 3.063810120934854e-06,
+      "loss": 0.6674,
+      "output_text_mask_loss": 0.0338,
+      "predictor_learning_rate": 0.00030638101209348533,
+      "step": 4341
+    },
+    {
+      "epoch": 0.44,
+      "gumbel_tau": 0.3591,
+      "image_masks_loss": 0.0066,
+      "llm_learning_rate": 3.063001846087464e-06,
+      "loss": 0.7147,
+      "output_text_mask_loss": 0.0385,
+      "predictor_learning_rate": 0.0003063001846087464,
+      "step": 4342
+    },
+    {
+      "epoch": 0.44,
+      "gumbel_tau": 0.359,
+      "image_masks_loss": 0.0045,
+      "llm_learning_rate": 3.0621935092410564e-06,
+      "loss": 0.7448,
+      "output_text_mask_loss": 0.006,
+      "predictor_learning_rate": 0.00030621935092410563,
+      "step": 4343
+    },
+    {
+      "epoch": 0.45,
+      "gumbel_tau": 0.3589,
+      "image_masks_loss": 0.0305,
+      "llm_learning_rate": 3.061385110484647e-06,
+      "loss": 0.7565,
+      "output_text_mask_loss": 0.0322,
+      "predictor_learning_rate": 0.0003061385110484647,
+      "step": 4344
+    },
+    {
+      "epoch": 0.45,
+      "gumbel_tau": 0.3589,
+      "image_masks_loss": 0.0103,
+      "llm_learning_rate": 3.06057664990726e-06,
+      "loss": 0.7246,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.00030605766499072594,
+      "step": 4345
+    },
+    {
+      "epoch": 0.45,
+      "gumbel_tau": 0.3588,
+      "image_masks_loss": 0.0051,
+      "llm_learning_rate": 3.0597681275979235e-06,
+      "loss": 0.7551,
+      "output_text_mask_loss": 0.0051,
+      "predictor_learning_rate": 0.0003059768127597923,
+      "step": 4346
+    },
+    {
+      "epoch": 0.45,
+      "gumbel_tau": 0.3587,
+      "image_masks_loss": 0.0131,
+      "llm_learning_rate": 3.058959543645673e-06,
+      "loss": 0.7529,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0003058959543645673,
+      "step": 4347
+    },
+    {
+      "epoch": 0.45,
+      "gumbel_tau": 0.3586,
+      "image_masks_loss": 0.0252,
+      "llm_learning_rate": 3.058150898139553e-06,
+      "loss": 0.8545,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0003058150898139553,
+      "step": 4348
+    },
+    {
+      "epoch": 0.45,
+      "gumbel_tau": 0.3585,
+      "image_masks_loss": 0.0187,
+      "llm_learning_rate": 3.057342191168613e-06,
+      "loss": 0.705,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00030573421911686126,
+      "step": 4349
+    },
+    {
+      "epoch": 0.45,
+      "gumbel_tau": 0.3584,
+      "image_masks_loss": 0.0051,
+      "llm_learning_rate": 3.0565334228219097e-06,
+      "loss": 0.8272,
+      "output_text_mask_loss": 0.0732,
+      "predictor_learning_rate": 0.00030565334228219095,
+      "step": 4350
+    },
+    {
+      "epoch": 0.45,
+      "gumbel_tau": 0.3583,
+      "image_masks_loss": 0.0191,
+      "llm_learning_rate": 3.055724593188505e-06,
+      "loss": 0.8331,
+      "output_text_mask_loss": 0.0141,
+      "predictor_learning_rate": 0.00030557245931885047,
+      "step": 4351
+    },
+    {
+      "epoch": 0.45,
+      "gumbel_tau": 0.3583,
+      "image_masks_loss": 0.0062,
+      "llm_learning_rate": 3.054915702357472e-06,
+      "loss": 0.7986,
+      "output_text_mask_loss": 0.0328,
+      "predictor_learning_rate": 0.00030549157023574715,
+      "step": 4352
+    },
+    {
+      "epoch": 0.45,
+      "gumbel_tau": 0.3582,
+      "image_masks_loss": 0.0065,
+      "llm_learning_rate": 3.0541067504178844e-06,
+      "loss": 0.7957,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0003054106750417884,
+      "step": 4353
+    },
+    {
+      "epoch": 0.45,
+      "gumbel_tau": 0.3581,
+      "image_masks_loss": 0.0138,
+      "llm_learning_rate": 3.0532977374588274e-06,
+      "loss": 0.5948,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0003053297737458827,
+      "step": 4354
+    },
+    {
+      "epoch": 0.45,
+      "gumbel_tau": 0.358,
+      "image_masks_loss": 0.0124,
+      "llm_learning_rate": 3.052488663569392e-06,
+      "loss": 0.7532,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0003052488663569392,
+      "step": 4355
+    },
+    {
+      "epoch": 0.45,
+      "gumbel_tau": 0.3579,
+      "image_masks_loss": 0.0257,
+      "llm_learning_rate": 3.0516795288386742e-06,
+      "loss": 0.732,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.0003051679528838674,
+      "step": 4356
+    },
+    {
+      "epoch": 0.45,
+      "gumbel_tau": 0.3578,
+      "image_masks_loss": 0.0149,
+      "llm_learning_rate": 3.0508703333557773e-06,
+      "loss": 0.7875,
+      "output_text_mask_loss": 0.0231,
+      "predictor_learning_rate": 0.0003050870333355777,
+      "step": 4357
+    },
+    {
+      "epoch": 0.45,
+      "gumbel_tau": 0.3578,
+      "image_masks_loss": 0.0113,
+      "llm_learning_rate": 3.050061077209812e-06,
+      "loss": 0.7991,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 0.0003050061077209812,
+      "step": 4358
+    },
+    {
+      "epoch": 0.45,
+      "gumbel_tau": 0.3577,
+      "image_masks_loss": 0.0057,
+      "llm_learning_rate": 3.049251760489897e-06,
+      "loss": 0.7028,
+      "output_text_mask_loss": 0.0107,
+      "predictor_learning_rate": 0.0003049251760489897,
+      "step": 4359
+    },
+    {
+      "epoch": 0.45,
+      "gumbel_tau": 0.3576,
+      "image_masks_loss": 0.013,
+      "llm_learning_rate": 3.048442383285153e-06,
+      "loss": 0.7132,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 0.00030484423832851526,
+      "step": 4360
+    },
+    {
+      "epoch": 0.45,
+      "gumbel_tau": 0.3575,
+      "image_masks_loss": 0.0078,
+      "llm_learning_rate": 3.0476329456847133e-06,
+      "loss": 0.7537,
+      "output_text_mask_loss": 0.0126,
+      "predictor_learning_rate": 0.0003047632945684713,
+      "step": 4361
+    },
+    {
+      "epoch": 0.45,
+      "gumbel_tau": 0.3574,
+      "image_masks_loss": 0.015,
+      "llm_learning_rate": 3.0468234477777137e-06,
+      "loss": 0.6633,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.00030468234477777135,
+      "step": 4362
+    },
+    {
+      "epoch": 0.45,
+      "gumbel_tau": 0.3573,
+      "image_masks_loss": 0.0134,
+      "llm_learning_rate": 3.0460138896532983e-06,
+      "loss": 0.7848,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.0003046013889653298,
+      "step": 4363
+    },
+    {
+      "epoch": 0.45,
+      "gumbel_tau": 0.3572,
+      "image_masks_loss": 0.0081,
+      "llm_learning_rate": 3.0452042714006157e-06,
+      "loss": 0.6794,
+      "output_text_mask_loss": 0.0146,
+      "predictor_learning_rate": 0.00030452042714006157,
+      "step": 4364
+    },
+    {
+      "epoch": 0.45,
+      "gumbel_tau": 0.3572,
+      "image_masks_loss": 0.0081,
+      "llm_learning_rate": 3.0443945931088266e-06,
+      "loss": 0.7001,
+      "output_text_mask_loss": 0.0095,
+      "predictor_learning_rate": 0.0003044394593108826,
+      "step": 4365
+    },
+    {
+      "epoch": 0.45,
+      "gumbel_tau": 0.3571,
+      "image_masks_loss": 0.0364,
+      "llm_learning_rate": 3.043584854867091e-06,
+      "loss": 0.7572,
+      "output_text_mask_loss": 0.0147,
+      "predictor_learning_rate": 0.00030435848548670906,
+      "step": 4366
+    },
+    {
+      "epoch": 0.45,
+      "gumbel_tau": 0.357,
+      "image_masks_loss": 0.0086,
+      "llm_learning_rate": 3.0427750567645803e-06,
+      "loss": 0.7621,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.000304277505676458,
+      "step": 4367
+    },
+    {
+      "epoch": 0.45,
+      "gumbel_tau": 0.3569,
+      "image_masks_loss": 0.0165,
+      "llm_learning_rate": 3.0419651988904714e-06,
+      "loss": 0.8045,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.00030419651988904714,
+      "step": 4368
+    },
+    {
+      "epoch": 0.45,
+      "gumbel_tau": 0.3568,
+      "image_masks_loss": 0.0168,
+      "llm_learning_rate": 3.041155281333948e-06,
+      "loss": 0.6613,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.00030411552813339475,
+      "step": 4369
+    },
+    {
+      "epoch": 0.45,
+      "gumbel_tau": 0.3567,
+      "image_masks_loss": 0.0051,
+      "llm_learning_rate": 3.040345304184199e-06,
+      "loss": 0.7537,
+      "output_text_mask_loss": 0.0343,
+      "predictor_learning_rate": 0.00030403453041841986,
+      "step": 4370
+    },
+    {
+      "epoch": 0.45,
+      "gumbel_tau": 0.3567,
+      "image_masks_loss": 0.0118,
+      "llm_learning_rate": 3.0395352675304218e-06,
+      "loss": 0.7944,
+      "output_text_mask_loss": 0.0242,
+      "predictor_learning_rate": 0.00030395352675304215,
+      "step": 4371
+    },
+    {
+      "epoch": 0.45,
+      "gumbel_tau": 0.3566,
+      "image_masks_loss": 0.0215,
+      "llm_learning_rate": 3.03872517146182e-06,
+      "loss": 0.7861,
+      "output_text_mask_loss": 0.0217,
+      "predictor_learning_rate": 0.00030387251714618195,
+      "step": 4372
+    },
+    {
+      "epoch": 0.45,
+      "gumbel_tau": 0.3565,
+      "image_masks_loss": 0.0214,
+      "llm_learning_rate": 3.0379150160676013e-06,
+      "loss": 0.7592,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 0.0003037915016067601,
+      "step": 4373
+    },
+    {
+      "epoch": 0.45,
+      "gumbel_tau": 0.3564,
+      "image_masks_loss": 0.0188,
+      "llm_learning_rate": 3.0371048014369835e-06,
+      "loss": 0.7257,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.00030371048014369835,
+      "step": 4374
+    },
+    {
+      "epoch": 0.45,
+      "gumbel_tau": 0.3563,
+      "image_masks_loss": 0.0087,
+      "llm_learning_rate": 3.0362945276591888e-06,
+      "loss": 0.7816,
+      "output_text_mask_loss": 0.0162,
+      "predictor_learning_rate": 0.00030362945276591886,
+      "step": 4375
+    },
+    {
+      "epoch": 0.45,
+      "gumbel_tau": 0.3562,
+      "image_masks_loss": 0.0113,
+      "llm_learning_rate": 3.035484194823445e-06,
+      "loss": 0.7812,
+      "output_text_mask_loss": 0.0992,
+      "predictor_learning_rate": 0.0003035484194823445,
+      "step": 4376
+    },
+    {
+      "epoch": 0.45,
+      "gumbel_tau": 0.3562,
+      "image_masks_loss": 0.0107,
+      "llm_learning_rate": 3.03467380301899e-06,
+      "loss": 0.8404,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.00030346738030189894,
+      "step": 4377
+    },
+    {
+      "epoch": 0.45,
+      "gumbel_tau": 0.3561,
+      "image_masks_loss": 0.0271,
+      "llm_learning_rate": 3.0338633523350645e-06,
+      "loss": 0.8279,
+      "output_text_mask_loss": 0.0032,
+      "predictor_learning_rate": 0.0003033863352335064,
+      "step": 4378
+    },
+    {
+      "epoch": 0.45,
+      "gumbel_tau": 0.356,
+      "image_masks_loss": 0.0206,
+      "llm_learning_rate": 3.0330528428609176e-06,
+      "loss": 0.7846,
+      "output_text_mask_loss": 0.0126,
+      "predictor_learning_rate": 0.0003033052842860917,
+      "step": 4379
+    },
+    {
+      "epoch": 0.45,
+      "gumbel_tau": 0.3559,
+      "image_masks_loss": 0.0092,
+      "llm_learning_rate": 3.032242274685804e-06,
+      "loss": 0.7528,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.00030322422746858037,
+      "step": 4380
+    },
+    {
+      "epoch": 0.45,
+      "gumbel_tau": 0.3558,
+      "image_masks_loss": 0.017,
+      "llm_learning_rate": 3.0314316478989857e-06,
+      "loss": 0.7597,
+      "output_text_mask_loss": 0.0098,
+      "predictor_learning_rate": 0.00030314316478989855,
+      "step": 4381
+    },
+    {
+      "epoch": 0.45,
+      "gumbel_tau": 0.3557,
+      "image_masks_loss": 0.0149,
+      "llm_learning_rate": 3.030620962589731e-06,
+      "loss": 0.7565,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 0.00030306209625897306,
+      "step": 4382
+    },
+    {
+      "epoch": 0.45,
+      "gumbel_tau": 0.3557,
+      "image_masks_loss": 0.0179,
+      "llm_learning_rate": 3.029810218847314e-06,
+      "loss": 0.5753,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00030298102188473135,
+      "step": 4383
+    },
+    {
+      "epoch": 0.45,
+      "gumbel_tau": 0.3556,
+      "image_masks_loss": 0.0326,
+      "llm_learning_rate": 3.028999416761015e-06,
+      "loss": 0.7371,
+      "output_text_mask_loss": 0.0004,
+      "predictor_learning_rate": 0.00030289994167610154,
+      "step": 4384
+    },
+    {
+      "epoch": 0.45,
+      "gumbel_tau": 0.3555,
+      "image_masks_loss": 0.016,
+      "llm_learning_rate": 3.028188556420123e-06,
+      "loss": 0.7107,
+      "output_text_mask_loss": 0.045,
+      "predictor_learning_rate": 0.00030281885564201227,
+      "step": 4385
+    },
+    {
+      "epoch": 0.45,
+      "gumbel_tau": 0.3554,
+      "image_masks_loss": 0.0237,
+      "llm_learning_rate": 3.027377637913929e-06,
+      "loss": 0.7339,
+      "output_text_mask_loss": 0.0313,
+      "predictor_learning_rate": 0.0003027377637913929,
+      "step": 4386
+    },
+    {
+      "epoch": 0.45,
+      "gumbel_tau": 0.3553,
+      "image_masks_loss": 0.0157,
+      "llm_learning_rate": 3.0265666613317363e-06,
+      "loss": 0.727,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.00030265666613317363,
+      "step": 4387
+    },
+    {
+      "epoch": 0.45,
+      "gumbel_tau": 0.3552,
+      "image_masks_loss": 0.0173,
+      "llm_learning_rate": 3.02575562676285e-06,
+      "loss": 0.647,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00030257556267628493,
+      "step": 4388
+    },
+    {
+      "epoch": 0.45,
+      "gumbel_tau": 0.3551,
+      "image_masks_loss": 0.0233,
+      "llm_learning_rate": 3.0249445342965822e-06,
+      "loss": 0.8359,
+      "output_text_mask_loss": 0.0075,
+      "predictor_learning_rate": 0.0003024944534296582,
+      "step": 4389
+    },
+    {
+      "epoch": 0.45,
+      "gumbel_tau": 0.3551,
+      "image_masks_loss": 0.01,
+      "llm_learning_rate": 3.0241333840222524e-06,
+      "loss": 0.8164,
+      "output_text_mask_loss": 0.046,
+      "predictor_learning_rate": 0.00030241333840222526,
+      "step": 4390
+    },
+    {
+      "epoch": 0.45,
+      "gumbel_tau": 0.355,
+      "image_masks_loss": 0.0202,
+      "llm_learning_rate": 3.0233221760291882e-06,
+      "loss": 0.7945,
+      "output_text_mask_loss": 0.0299,
+      "predictor_learning_rate": 0.0003023322176029188,
+      "step": 4391
+    },
+    {
+      "epoch": 0.45,
+      "gumbel_tau": 0.3549,
+      "image_masks_loss": 0.0144,
+      "llm_learning_rate": 3.0225109104067196e-06,
+      "loss": 0.6961,
+      "output_text_mask_loss": 0.0202,
+      "predictor_learning_rate": 0.0003022510910406719,
+      "step": 4392
+    },
+    {
+      "epoch": 0.45,
+      "gumbel_tau": 0.3548,
+      "image_masks_loss": 0.0376,
+      "llm_learning_rate": 3.021699587244186e-06,
+      "loss": 0.7222,
+      "output_text_mask_loss": 0.0278,
+      "predictor_learning_rate": 0.0003021699587244186,
+      "step": 4393
+    },
+    {
+      "epoch": 0.45,
+      "gumbel_tau": 0.3547,
+      "image_masks_loss": 0.0103,
+      "llm_learning_rate": 3.0208882066309323e-06,
+      "loss": 0.7952,
+      "output_text_mask_loss": 0.0075,
+      "predictor_learning_rate": 0.00030208882066309325,
+      "step": 4394
+    },
+    {
+      "epoch": 0.45,
+      "gumbel_tau": 0.3546,
+      "image_masks_loss": 0.0097,
+      "llm_learning_rate": 3.0200767686563088e-06,
+      "loss": 0.729,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.00030200767686563084,
+      "step": 4395
+    },
+    {
+      "epoch": 0.45,
+      "gumbel_tau": 0.3546,
+      "image_masks_loss": 0.016,
+      "llm_learning_rate": 3.0192652734096723e-06,
+      "loss": 0.723,
+      "output_text_mask_loss": 0.0126,
+      "predictor_learning_rate": 0.0003019265273409672,
+      "step": 4396
+    },
+    {
+      "epoch": 0.45,
+      "gumbel_tau": 0.3545,
+      "image_masks_loss": 0.0345,
+      "llm_learning_rate": 3.0184537209803875e-06,
+      "loss": 0.8351,
+      "output_text_mask_loss": 0.0041,
+      "predictor_learning_rate": 0.00030184537209803876,
+      "step": 4397
+    },
+    {
+      "epoch": 0.45,
+      "gumbel_tau": 0.3544,
+      "image_masks_loss": 0.0208,
+      "llm_learning_rate": 3.017642111457825e-06,
+      "loss": 0.7189,
+      "output_text_mask_loss": 0.0023,
+      "predictor_learning_rate": 0.00030176421114578244,
+      "step": 4398
+    },
+    {
+      "epoch": 0.45,
+      "gumbel_tau": 0.3543,
+      "image_masks_loss": 0.0275,
+      "llm_learning_rate": 3.016830444931359e-06,
+      "loss": 0.7675,
+      "output_text_mask_loss": 0.0156,
+      "predictor_learning_rate": 0.00030168304449313587,
+      "step": 4399
+    },
+    {
+      "epoch": 0.45,
+      "gumbel_tau": 0.3542,
+      "image_masks_loss": 0.0198,
+      "llm_learning_rate": 3.016018721490374e-06,
+      "loss": 0.6837,
+      "output_text_mask_loss": 0.0185,
+      "predictor_learning_rate": 0.00030160187214903737,
+      "step": 4400
+    },
+    {
+      "epoch": 0.45,
+      "gumbel_tau": 0.3541,
+      "image_masks_loss": 0.0319,
+      "llm_learning_rate": 3.015206941224258e-06,
+      "loss": 0.7135,
+      "output_text_mask_loss": 0.0208,
+      "predictor_learning_rate": 0.00030152069412242577,
+      "step": 4401
+    },
+    {
+      "epoch": 0.45,
+      "gumbel_tau": 0.3541,
+      "image_masks_loss": 0.0069,
+      "llm_learning_rate": 3.014395104222406e-06,
+      "loss": 0.7125,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.0003014395104222406,
+      "step": 4402
+    },
+    {
+      "epoch": 0.45,
+      "gumbel_tau": 0.354,
+      "image_masks_loss": 0.0133,
+      "llm_learning_rate": 3.01358321057422e-06,
+      "loss": 0.7535,
+      "output_text_mask_loss": 0.0319,
+      "predictor_learning_rate": 0.00030135832105742196,
+      "step": 4403
+    },
+    {
+      "epoch": 0.45,
+      "gumbel_tau": 0.3539,
+      "image_masks_loss": 0.0073,
+      "llm_learning_rate": 3.0127712603691057e-06,
+      "loss": 0.8222,
+      "output_text_mask_loss": 0.0145,
+      "predictor_learning_rate": 0.00030127712603691053,
+      "step": 4404
+    },
+    {
+      "epoch": 0.45,
+      "gumbel_tau": 0.3538,
+      "image_masks_loss": 0.0187,
+      "llm_learning_rate": 3.011959253696479e-06,
+      "loss": 0.802,
+      "output_text_mask_loss": 0.0305,
+      "predictor_learning_rate": 0.00030119592536964786,
+      "step": 4405
+    },
+    {
+      "epoch": 0.45,
+      "gumbel_tau": 0.3537,
+      "image_masks_loss": 0.0238,
+      "llm_learning_rate": 3.0111471906457577e-06,
+      "loss": 0.7901,
+      "output_text_mask_loss": 0.0081,
+      "predictor_learning_rate": 0.0003011147190645758,
+      "step": 4406
+    },
+    {
+      "epoch": 0.45,
+      "gumbel_tau": 0.3536,
+      "image_masks_loss": 0.0172,
+      "llm_learning_rate": 3.01033507130637e-06,
+      "loss": 0.7282,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 0.000301033507130637,
+      "step": 4407
+    },
+    {
+      "epoch": 0.45,
+      "gumbel_tau": 0.3536,
+      "image_masks_loss": 0.0195,
+      "llm_learning_rate": 3.009522895767748e-06,
+      "loss": 0.7432,
+      "output_text_mask_loss": 0.0454,
+      "predictor_learning_rate": 0.0003009522895767748,
+      "step": 4408
+    },
+    {
+      "epoch": 0.45,
+      "gumbel_tau": 0.3535,
+      "image_masks_loss": 0.0309,
+      "llm_learning_rate": 3.0087106641193285e-06,
+      "loss": 0.7569,
+      "output_text_mask_loss": 0.0652,
+      "predictor_learning_rate": 0.00030087106641193284,
+      "step": 4409
+    },
+    {
+      "epoch": 0.45,
+      "gumbel_tau": 0.3534,
+      "image_masks_loss": 0.0045,
+      "llm_learning_rate": 3.007898376450559e-06,
+      "loss": 0.7933,
+      "output_text_mask_loss": 0.0294,
+      "predictor_learning_rate": 0.00030078983764505584,
+      "step": 4410
+    },
+    {
+      "epoch": 0.45,
+      "gumbel_tau": 0.3533,
+      "image_masks_loss": 0.0133,
+      "llm_learning_rate": 3.0070860328508888e-06,
+      "loss": 0.7567,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0003007086032850889,
+      "step": 4411
+    },
+    {
+      "epoch": 0.45,
+      "gumbel_tau": 0.3532,
+      "image_masks_loss": 0.0125,
+      "llm_learning_rate": 3.006273633409774e-06,
+      "loss": 0.7288,
+      "output_text_mask_loss": 0.0032,
+      "predictor_learning_rate": 0.0003006273633409774,
+      "step": 4412
+    },
+    {
+      "epoch": 0.45,
+      "gumbel_tau": 0.3531,
+      "image_masks_loss": 0.0141,
+      "llm_learning_rate": 3.0054611782166798e-06,
+      "loss": 0.7909,
+      "output_text_mask_loss": 0.0252,
+      "predictor_learning_rate": 0.00030054611782166795,
+      "step": 4413
+    },
+    {
+      "epoch": 0.45,
+      "gumbel_tau": 0.3531,
+      "image_masks_loss": 0.0154,
+      "llm_learning_rate": 3.004648667361075e-06,
+      "loss": 0.7493,
+      "output_text_mask_loss": 0.0298,
+      "predictor_learning_rate": 0.00030046486673610743,
+      "step": 4414
+    },
+    {
+      "epoch": 0.45,
+      "gumbel_tau": 0.353,
+      "image_masks_loss": 0.0176,
+      "llm_learning_rate": 3.0038361009324342e-06,
+      "loss": 0.7455,
+      "output_text_mask_loss": 0.032,
+      "predictor_learning_rate": 0.0003003836100932434,
+      "step": 4415
+    },
+    {
+      "epoch": 0.45,
+      "gumbel_tau": 0.3529,
+      "image_masks_loss": 0.0066,
+      "llm_learning_rate": 3.00302347902024e-06,
+      "loss": 0.7749,
+      "output_text_mask_loss": 0.0025,
+      "predictor_learning_rate": 0.000300302347902024,
+      "step": 4416
+    },
+    {
+      "epoch": 0.45,
+      "gumbel_tau": 0.3528,
+      "image_masks_loss": 0.0054,
+      "llm_learning_rate": 3.00221080171398e-06,
+      "loss": 0.6709,
+      "output_text_mask_loss": 0.0107,
+      "predictor_learning_rate": 0.000300221080171398,
+      "step": 4417
+    },
+    {
+      "epoch": 0.45,
+      "gumbel_tau": 0.3527,
+      "image_masks_loss": 0.0162,
+      "llm_learning_rate": 3.001398069103148e-06,
+      "loss": 0.7569,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.00030013980691031476,
+      "step": 4418
+    },
+    {
+      "epoch": 0.45,
+      "gumbel_tau": 0.3526,
+      "image_masks_loss": 0.0206,
+      "llm_learning_rate": 3.0005852812772436e-06,
+      "loss": 0.6923,
+      "output_text_mask_loss": 0.0086,
+      "predictor_learning_rate": 0.00030005852812772435,
+      "step": 4419
+    },
+    {
+      "epoch": 0.45,
+      "gumbel_tau": 0.3526,
+      "image_masks_loss": 0.0186,
+      "llm_learning_rate": 2.9997724383257738e-06,
+      "loss": 0.7143,
+      "output_text_mask_loss": 0.0019,
+      "predictor_learning_rate": 0.00029997724383257733,
+      "step": 4420
+    },
+    {
+      "epoch": 0.45,
+      "gumbel_tau": 0.3525,
+      "image_masks_loss": 0.0109,
+      "llm_learning_rate": 2.998959540338249e-06,
+      "loss": 0.7505,
+      "output_text_mask_loss": 0.002,
+      "predictor_learning_rate": 0.00029989595403382487,
+      "step": 4421
+    },
+    {
+      "epoch": 0.45,
+      "gumbel_tau": 0.3524,
+      "image_masks_loss": 0.0224,
+      "llm_learning_rate": 2.9981465874041883e-06,
+      "loss": 0.7333,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0002998146587404188,
+      "step": 4422
+    },
+    {
+      "epoch": 0.45,
+      "gumbel_tau": 0.3523,
+      "image_masks_loss": 0.013,
+      "llm_learning_rate": 2.9973335796131155e-06,
+      "loss": 0.7311,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.0002997333579613115,
+      "step": 4423
+    },
+    {
+      "epoch": 0.45,
+      "gumbel_tau": 0.3522,
+      "image_masks_loss": 0.0154,
+      "llm_learning_rate": 2.996520517054562e-06,
+      "loss": 0.8387,
+      "output_text_mask_loss": 0.0391,
+      "predictor_learning_rate": 0.0002996520517054562,
+      "step": 4424
+    },
+    {
+      "epoch": 0.45,
+      "gumbel_tau": 0.3521,
+      "image_masks_loss": 0.025,
+      "llm_learning_rate": 2.995707399818062e-06,
+      "loss": 0.8336,
+      "output_text_mask_loss": 0.0134,
+      "predictor_learning_rate": 0.00029957073998180617,
+      "step": 4425
+    },
+    {
+      "epoch": 0.45,
+      "gumbel_tau": 0.3521,
+      "image_masks_loss": 0.0028,
+      "llm_learning_rate": 2.99489422799316e-06,
+      "loss": 0.6894,
+      "output_text_mask_loss": 0.0511,
+      "predictor_learning_rate": 0.00029948942279931596,
+      "step": 4426
+    },
+    {
+      "epoch": 0.45,
+      "gumbel_tau": 0.352,
+      "image_masks_loss": 0.0034,
+      "llm_learning_rate": 2.9940810016694032e-06,
+      "loss": 0.7164,
+      "output_text_mask_loss": 0.053,
+      "predictor_learning_rate": 0.0002994081001669403,
+      "step": 4427
+    },
+    {
+      "epoch": 0.45,
+      "gumbel_tau": 0.3519,
+      "image_masks_loss": 0.0128,
+      "llm_learning_rate": 2.993267720936346e-06,
+      "loss": 0.7854,
+      "output_text_mask_loss": 0.0347,
+      "predictor_learning_rate": 0.00029932677209363456,
+      "step": 4428
+    },
+    {
+      "epoch": 0.45,
+      "gumbel_tau": 0.3518,
+      "image_masks_loss": 0.0101,
+      "llm_learning_rate": 2.9924543858835483e-06,
+      "loss": 0.8059,
+      "output_text_mask_loss": 0.0177,
+      "predictor_learning_rate": 0.0002992454385883548,
+      "step": 4429
+    },
+    {
+      "epoch": 0.45,
+      "gumbel_tau": 0.3517,
+      "image_masks_loss": 0.025,
+      "llm_learning_rate": 2.991640996600577e-06,
+      "loss": 0.7776,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.0002991640996600577,
+      "step": 4430
+    },
+    {
+      "epoch": 0.45,
+      "gumbel_tau": 0.3516,
+      "image_masks_loss": 0.0113,
+      "llm_learning_rate": 2.9908275531770038e-06,
+      "loss": 0.7472,
+      "output_text_mask_loss": 0.0391,
+      "predictor_learning_rate": 0.00029908275531770036,
+      "step": 4431
+    },
+    {
+      "epoch": 0.45,
+      "gumbel_tau": 0.3516,
+      "image_masks_loss": 0.0113,
+      "llm_learning_rate": 2.990014055702407e-06,
+      "loss": 0.7207,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0002990014055702407,
+      "step": 4432
+    },
+    {
+      "epoch": 0.45,
+      "gumbel_tau": 0.3515,
+      "image_masks_loss": 0.0156,
+      "llm_learning_rate": 2.9892005042663718e-06,
+      "loss": 0.752,
+      "output_text_mask_loss": 0.0137,
+      "predictor_learning_rate": 0.0002989200504266372,
+      "step": 4433
+    },
+    {
+      "epoch": 0.45,
+      "gumbel_tau": 0.3514,
+      "image_masks_loss": 0.0127,
+      "llm_learning_rate": 2.988386898958487e-06,
+      "loss": 0.7708,
+      "output_text_mask_loss": 0.0496,
+      "predictor_learning_rate": 0.0002988386898958487,
+      "step": 4434
+    },
+    {
+      "epoch": 0.45,
+      "gumbel_tau": 0.3513,
+      "image_masks_loss": 0.0179,
+      "llm_learning_rate": 2.987573239868348e-06,
+      "loss": 0.7648,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0002987573239868348,
+      "step": 4435
+    },
+    {
+      "epoch": 0.45,
+      "gumbel_tau": 0.3512,
+      "image_masks_loss": 0.0121,
+      "llm_learning_rate": 2.9867595270855594e-06,
+      "loss": 0.8302,
+      "output_text_mask_loss": 0.033,
+      "predictor_learning_rate": 0.0002986759527085559,
+      "step": 4436
+    },
+    {
+      "epoch": 0.45,
+      "gumbel_tau": 0.3511,
+      "image_masks_loss": 0.0081,
+      "llm_learning_rate": 2.9859457606997268e-06,
+      "loss": 0.8432,
+      "output_text_mask_loss": 0.0399,
+      "predictor_learning_rate": 0.00029859457606997266,
+      "step": 4437
+    },
+    {
+      "epoch": 0.45,
+      "gumbel_tau": 0.3511,
+      "image_masks_loss": 0.0036,
+      "llm_learning_rate": 2.985131940800464e-06,
+      "loss": 0.7176,
+      "output_text_mask_loss": 0.0277,
+      "predictor_learning_rate": 0.0002985131940800464,
+      "step": 4438
+    },
+    {
+      "epoch": 0.45,
+      "gumbel_tau": 0.351,
+      "image_masks_loss": 0.0067,
+      "llm_learning_rate": 2.9843180674773924e-06,
+      "loss": 0.7529,
+      "output_text_mask_loss": 0.0488,
+      "predictor_learning_rate": 0.00029843180674773924,
+      "step": 4439
+    },
+    {
+      "epoch": 0.45,
+      "gumbel_tau": 0.3509,
+      "image_masks_loss": 0.025,
+      "llm_learning_rate": 2.983504140820136e-06,
+      "loss": 0.8867,
+      "output_text_mask_loss": 0.0107,
+      "predictor_learning_rate": 0.0002983504140820136,
+      "step": 4440
+    },
+    {
+      "epoch": 0.46,
+      "gumbel_tau": 0.3508,
+      "image_masks_loss": 0.0162,
+      "llm_learning_rate": 2.982690160918326e-06,
+      "loss": 0.7797,
+      "output_text_mask_loss": 0.0006,
+      "predictor_learning_rate": 0.0002982690160918326,
+      "step": 4441
+    },
+    {
+      "epoch": 0.46,
+      "gumbel_tau": 0.3507,
+      "image_masks_loss": 0.0217,
+      "llm_learning_rate": 2.9818761278616016e-06,
+      "loss": 0.8017,
+      "output_text_mask_loss": 0.0004,
+      "predictor_learning_rate": 0.00029818761278616014,
+      "step": 4442
+    },
+    {
+      "epoch": 0.46,
+      "gumbel_tau": 0.3507,
+      "image_masks_loss": 0.002,
+      "llm_learning_rate": 2.981062041739604e-06,
+      "loss": 0.7943,
+      "output_text_mask_loss": 0.0702,
+      "predictor_learning_rate": 0.00029810620417396037,
+      "step": 4443
+    },
+    {
+      "epoch": 0.46,
+      "gumbel_tau": 0.3506,
+      "image_masks_loss": 0.0119,
+      "llm_learning_rate": 2.980247902641983e-06,
+      "loss": 0.7241,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.00029802479026419824,
+      "step": 4444
+    },
+    {
+      "epoch": 0.46,
+      "gumbel_tau": 0.3505,
+      "image_masks_loss": 0.0082,
+      "llm_learning_rate": 2.9794337106583927e-06,
+      "loss": 0.7718,
+      "output_text_mask_loss": 0.0862,
+      "predictor_learning_rate": 0.00029794337106583925,
+      "step": 4445
+    },
+    {
+      "epoch": 0.46,
+      "gumbel_tau": 0.3504,
+      "image_masks_loss": 0.006,
+      "llm_learning_rate": 2.9786194658784947e-06,
+      "loss": 0.72,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0002978619465878494,
+      "step": 4446
+    },
+    {
+      "epoch": 0.46,
+      "gumbel_tau": 0.3503,
+      "image_masks_loss": 0.0191,
+      "llm_learning_rate": 2.977805168391955e-06,
+      "loss": 0.7532,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 0.0002977805168391955,
+      "step": 4447
+    },
+    {
+      "epoch": 0.46,
+      "gumbel_tau": 0.3502,
+      "image_masks_loss": 0.0108,
+      "llm_learning_rate": 2.9769908182884456e-06,
+      "loss": 0.6727,
+      "output_text_mask_loss": 0.0124,
+      "predictor_learning_rate": 0.00029769908182884455,
+      "step": 4448
+    },
+    {
+      "epoch": 0.46,
+      "gumbel_tau": 0.3502,
+      "image_masks_loss": 0.0101,
+      "llm_learning_rate": 2.9761764156576458e-06,
+      "loss": 0.6767,
+      "output_text_mask_loss": 0.0202,
+      "predictor_learning_rate": 0.00029761764156576456,
+      "step": 4449
+    },
+    {
+      "epoch": 0.46,
+      "gumbel_tau": 0.3501,
+      "image_masks_loss": 0.0196,
+      "llm_learning_rate": 2.9753619605892376e-06,
+      "loss": 0.7651,
+      "output_text_mask_loss": 0.0012,
+      "predictor_learning_rate": 0.00029753619605892375,
+      "step": 4450
+    },
+    {
+      "epoch": 0.46,
+      "gumbel_tau": 0.35,
+      "image_masks_loss": 0.0113,
+      "llm_learning_rate": 2.9745474531729113e-06,
+      "loss": 0.7859,
+      "output_text_mask_loss": 0.0256,
+      "predictor_learning_rate": 0.0002974547453172911,
+      "step": 4451
+    },
+    {
+      "epoch": 0.46,
+      "gumbel_tau": 0.3499,
+      "image_masks_loss": 0.0246,
+      "llm_learning_rate": 2.9737328934983626e-06,
+      "loss": 0.753,
+      "output_text_mask_loss": 0.0368,
+      "predictor_learning_rate": 0.0002973732893498363,
+      "step": 4452
+    },
+    {
+      "epoch": 0.46,
+      "gumbel_tau": 0.3498,
+      "image_masks_loss": 0.0138,
+      "llm_learning_rate": 2.9729182816552927e-06,
+      "loss": 0.7571,
+      "output_text_mask_loss": 0.002,
+      "predictor_learning_rate": 0.00029729182816552925,
+      "step": 4453
+    },
+    {
+      "epoch": 0.46,
+      "gumbel_tau": 0.3497,
+      "image_masks_loss": 0.0112,
+      "llm_learning_rate": 2.9721036177334083e-06,
+      "loss": 0.8085,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.00029721036177334084,
+      "step": 4454
+    },
+    {
+      "epoch": 0.46,
+      "gumbel_tau": 0.3497,
+      "image_masks_loss": 0.0103,
+      "llm_learning_rate": 2.9712889018224217e-06,
+      "loss": 0.8006,
+      "output_text_mask_loss": 0.0687,
+      "predictor_learning_rate": 0.0002971288901822422,
+      "step": 4455
+    },
+    {
+      "epoch": 0.46,
+      "gumbel_tau": 0.3496,
+      "image_masks_loss": 0.011,
+      "llm_learning_rate": 2.9704741340120513e-06,
+      "loss": 0.796,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00029704741340120513,
+      "step": 4456
+    },
+    {
+      "epoch": 0.46,
+      "gumbel_tau": 0.3495,
+      "image_masks_loss": 0.0072,
+      "llm_learning_rate": 2.9696593143920222e-06,
+      "loss": 0.7152,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.0002969659314392022,
+      "step": 4457
+    },
+    {
+      "epoch": 0.46,
+      "gumbel_tau": 0.3494,
+      "image_masks_loss": 0.0267,
+      "llm_learning_rate": 2.9688444430520623e-06,
+      "loss": 0.6663,
+      "output_text_mask_loss": 0.017,
+      "predictor_learning_rate": 0.0002968844443052062,
+      "step": 4458
+    },
+    {
+      "epoch": 0.46,
+      "gumbel_tau": 0.3493,
+      "image_masks_loss": 0.0102,
+      "llm_learning_rate": 2.9680295200819087e-06,
+      "loss": 0.7733,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 0.00029680295200819085,
+      "step": 4459
+    },
+    {
+      "epoch": 0.46,
+      "gumbel_tau": 0.3492,
+      "image_masks_loss": 0.0085,
+      "llm_learning_rate": 2.967214545571302e-06,
+      "loss": 0.7275,
+      "output_text_mask_loss": 0.017,
+      "predictor_learning_rate": 0.00029672145455713017,
+      "step": 4460
+    },
+    {
+      "epoch": 0.46,
+      "gumbel_tau": 0.3492,
+      "image_masks_loss": 0.0137,
+      "llm_learning_rate": 2.9663995196099884e-06,
+      "loss": 0.6321,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0002966399519609988,
+      "step": 4461
+    },
+    {
+      "epoch": 0.46,
+      "gumbel_tau": 0.3491,
+      "image_masks_loss": 0.009,
+      "llm_learning_rate": 2.9655844422877207e-06,
+      "loss": 0.7773,
+      "output_text_mask_loss": 0.0034,
+      "predictor_learning_rate": 0.00029655844422877203,
+      "step": 4462
+    },
+    {
+      "epoch": 0.46,
+      "gumbel_tau": 0.349,
+      "image_masks_loss": 0.0439,
+      "llm_learning_rate": 2.9647693136942574e-06,
+      "loss": 0.7186,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0002964769313694257,
+      "step": 4463
+    },
+    {
+      "epoch": 0.46,
+      "gumbel_tau": 0.3489,
+      "image_masks_loss": 0.0119,
+      "llm_learning_rate": 2.963954133919361e-06,
+      "loss": 0.7896,
+      "output_text_mask_loss": 0.0298,
+      "predictor_learning_rate": 0.0002963954133919361,
+      "step": 4464
+    },
+    {
+      "epoch": 0.46,
+      "gumbel_tau": 0.3488,
+      "image_masks_loss": 0.021,
+      "llm_learning_rate": 2.9631389030528034e-06,
+      "loss": 0.7904,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0002963138903052803,
+      "step": 4465
+    },
+    {
+      "epoch": 0.46,
+      "gumbel_tau": 0.3488,
+      "image_masks_loss": 0.0215,
+      "llm_learning_rate": 2.962323621184358e-06,
+      "loss": 0.8161,
+      "output_text_mask_loss": 0.058,
+      "predictor_learning_rate": 0.0002962323621184358,
+      "step": 4466
+    },
+    {
+      "epoch": 0.46,
+      "gumbel_tau": 0.3487,
+      "image_masks_loss": 0.0162,
+      "llm_learning_rate": 2.9615082884038055e-06,
+      "loss": 0.787,
+      "output_text_mask_loss": 0.0179,
+      "predictor_learning_rate": 0.0002961508288403805,
+      "step": 4467
+    },
+    {
+      "epoch": 0.46,
+      "gumbel_tau": 0.3486,
+      "image_masks_loss": 0.0156,
+      "llm_learning_rate": 2.9606929048009327e-06,
+      "loss": 0.7226,
+      "output_text_mask_loss": 0.0319,
+      "predictor_learning_rate": 0.00029606929048009325,
+      "step": 4468
+    },
+    {
+      "epoch": 0.46,
+      "gumbel_tau": 0.3485,
+      "image_masks_loss": 0.0116,
+      "llm_learning_rate": 2.95987747046553e-06,
+      "loss": 0.7443,
+      "output_text_mask_loss": 0.0088,
+      "predictor_learning_rate": 0.00029598774704655303,
+      "step": 4469
+    },
+    {
+      "epoch": 0.46,
+      "gumbel_tau": 0.3484,
+      "image_masks_loss": 0.0134,
+      "llm_learning_rate": 2.959061985487397e-06,
+      "loss": 0.7372,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.0002959061985487397,
+      "step": 4470
+    },
+    {
+      "epoch": 0.46,
+      "gumbel_tau": 0.3483,
+      "image_masks_loss": 0.0161,
+      "llm_learning_rate": 2.9582464499563347e-06,
+      "loss": 0.6823,
+      "output_text_mask_loss": 0.0004,
+      "predictor_learning_rate": 0.00029582464499563343,
+      "step": 4471
+    },
+    {
+      "epoch": 0.46,
+      "gumbel_tau": 0.3483,
+      "image_masks_loss": 0.0213,
+      "llm_learning_rate": 2.957430863962153e-06,
+      "loss": 0.7267,
+      "output_text_mask_loss": 0.0553,
+      "predictor_learning_rate": 0.0002957430863962153,
+      "step": 4472
+    },
+    {
+      "epoch": 0.46,
+      "gumbel_tau": 0.3482,
+      "image_masks_loss": 0.0169,
+      "llm_learning_rate": 2.956615227594666e-06,
+      "loss": 0.7876,
+      "output_text_mask_loss": 0.0372,
+      "predictor_learning_rate": 0.0002956615227594666,
+      "step": 4473
+    },
+    {
+      "epoch": 0.46,
+      "gumbel_tau": 0.3481,
+      "image_masks_loss": 0.0056,
+      "llm_learning_rate": 2.9557995409436923e-06,
+      "loss": 0.7383,
+      "output_text_mask_loss": 0.0629,
+      "predictor_learning_rate": 0.00029557995409436923,
+      "step": 4474
+    },
+    {
+      "epoch": 0.46,
+      "gumbel_tau": 0.348,
+      "image_masks_loss": 0.0226,
+      "llm_learning_rate": 2.9549838040990586e-06,
+      "loss": 0.8401,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.00029549838040990583,
+      "step": 4475
+    },
+    {
+      "epoch": 0.46,
+      "gumbel_tau": 0.3479,
+      "image_masks_loss": 0.0117,
+      "llm_learning_rate": 2.9541680171505956e-06,
+      "loss": 0.7329,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.0002954168017150595,
+      "step": 4476
+    },
+    {
+      "epoch": 0.46,
+      "gumbel_tau": 0.3479,
+      "image_masks_loss": 0.0198,
+      "llm_learning_rate": 2.953352180188137e-06,
+      "loss": 0.7476,
+      "output_text_mask_loss": 0.0132,
+      "predictor_learning_rate": 0.0002953352180188137,
+      "step": 4477
+    },
+    {
+      "epoch": 0.46,
+      "gumbel_tau": 0.3478,
+      "image_masks_loss": 0.012,
+      "llm_learning_rate": 2.952536293301528e-06,
+      "loss": 0.8194,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.00029525362933015276,
+      "step": 4478
+    },
+    {
+      "epoch": 0.46,
+      "gumbel_tau": 0.3477,
+      "image_masks_loss": 0.0057,
+      "llm_learning_rate": 2.951720356580614e-06,
+      "loss": 0.717,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 0.00029517203565806137,
+      "step": 4479
+    },
+    {
+      "epoch": 0.46,
+      "gumbel_tau": 0.3476,
+      "image_masks_loss": 0.0118,
+      "llm_learning_rate": 2.9509043701152477e-06,
+      "loss": 0.767,
+      "output_text_mask_loss": 0.0198,
+      "predictor_learning_rate": 0.00029509043701152475,
+      "step": 4480
+    },
+    {
+      "epoch": 0.46,
+      "gumbel_tau": 0.3475,
+      "image_masks_loss": 0.022,
+      "llm_learning_rate": 2.9500883339952886e-06,
+      "loss": 0.7589,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00029500883339952884,
+      "step": 4481
+    },
+    {
+      "epoch": 0.46,
+      "gumbel_tau": 0.3474,
+      "image_masks_loss": 0.0273,
+      "llm_learning_rate": 2.9492722483105994e-06,
+      "loss": 0.7873,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0002949272248310599,
+      "step": 4482
+    },
+    {
+      "epoch": 0.46,
+      "gumbel_tau": 0.3474,
+      "image_masks_loss": 0.0082,
+      "llm_learning_rate": 2.948456113151049e-06,
+      "loss": 0.7643,
+      "output_text_mask_loss": 0.0465,
+      "predictor_learning_rate": 0.0002948456113151049,
+      "step": 4483
+    },
+    {
+      "epoch": 0.46,
+      "gumbel_tau": 0.3473,
+      "image_masks_loss": 0.0205,
+      "llm_learning_rate": 2.947639928606513e-06,
+      "loss": 0.703,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.00029476399286065125,
+      "step": 4484
+    },
+    {
+      "epoch": 0.46,
+      "gumbel_tau": 0.3472,
+      "image_masks_loss": 0.0147,
+      "llm_learning_rate": 2.9468236947668715e-06,
+      "loss": 0.8551,
+      "output_text_mask_loss": 0.0095,
+      "predictor_learning_rate": 0.00029468236947668714,
+      "step": 4485
+    },
+    {
+      "epoch": 0.46,
+      "gumbel_tau": 0.3471,
+      "image_masks_loss": 0.0192,
+      "llm_learning_rate": 2.946007411722009e-06,
+      "loss": 0.6628,
+      "output_text_mask_loss": 0.0179,
+      "predictor_learning_rate": 0.0002946007411722009,
+      "step": 4486
+    },
+    {
+      "epoch": 0.46,
+      "gumbel_tau": 0.347,
+      "image_masks_loss": 0.0425,
+      "llm_learning_rate": 2.9451910795618168e-06,
+      "loss": 0.7779,
+      "output_text_mask_loss": 0.0101,
+      "predictor_learning_rate": 0.0002945191079561816,
+      "step": 4487
+    },
+    {
+      "epoch": 0.46,
+      "gumbel_tau": 0.3469,
+      "image_masks_loss": 0.024,
+      "llm_learning_rate": 2.944374698376191e-06,
+      "loss": 0.8544,
+      "output_text_mask_loss": 0.0032,
+      "predictor_learning_rate": 0.0002944374698376191,
+      "step": 4488
+    },
+    {
+      "epoch": 0.46,
+      "gumbel_tau": 0.3469,
+      "image_masks_loss": 0.0182,
+      "llm_learning_rate": 2.943558268255034e-06,
+      "loss": 0.7779,
+      "output_text_mask_loss": 0.1183,
+      "predictor_learning_rate": 0.0002943558268255034,
+      "step": 4489
+    },
+    {
+      "epoch": 0.46,
+      "gumbel_tau": 0.3468,
+      "image_masks_loss": 0.0077,
+      "llm_learning_rate": 2.942741789288252e-06,
+      "loss": 0.7772,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 0.0002942741789288252,
+      "step": 4490
+    },
+    {
+      "epoch": 0.46,
+      "gumbel_tau": 0.3467,
+      "image_masks_loss": 0.0238,
+      "llm_learning_rate": 2.9419252615657585e-06,
+      "loss": 0.7569,
+      "output_text_mask_loss": 0.1144,
+      "predictor_learning_rate": 0.0002941925261565758,
+      "step": 4491
+    },
+    {
+      "epoch": 0.46,
+      "gumbel_tau": 0.3466,
+      "image_masks_loss": 0.0353,
+      "llm_learning_rate": 2.9411086851774705e-06,
+      "loss": 0.7407,
+      "output_text_mask_loss": 0.0179,
+      "predictor_learning_rate": 0.000294110868517747,
+      "step": 4492
+    },
+    {
+      "epoch": 0.46,
+      "gumbel_tau": 0.3465,
+      "image_masks_loss": 0.0222,
+      "llm_learning_rate": 2.940292060213312e-06,
+      "loss": 0.7218,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0002940292060213312,
+      "step": 4493
+    },
+    {
+      "epoch": 0.46,
+      "gumbel_tau": 0.3465,
+      "image_masks_loss": 0.0242,
+      "llm_learning_rate": 2.939475386763211e-06,
+      "loss": 0.7724,
+      "output_text_mask_loss": 0.0278,
+      "predictor_learning_rate": 0.0002939475386763211,
+      "step": 4494
+    },
+    {
+      "epoch": 0.46,
+      "gumbel_tau": 0.3464,
+      "image_masks_loss": 0.0114,
+      "llm_learning_rate": 2.9386586649171015e-06,
+      "loss": 0.7546,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.0002938658664917102,
+      "step": 4495
+    },
+    {
+      "epoch": 0.46,
+      "gumbel_tau": 0.3463,
+      "image_masks_loss": 0.0202,
+      "llm_learning_rate": 2.937841894764922e-06,
+      "loss": 0.6407,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0002937841894764922,
+      "step": 4496
+    },
+    {
+      "epoch": 0.46,
+      "gumbel_tau": 0.3462,
+      "image_masks_loss": 0.0234,
+      "llm_learning_rate": 2.937025076396619e-06,
+      "loss": 0.7921,
+      "output_text_mask_loss": 0.0231,
+      "predictor_learning_rate": 0.0002937025076396619,
+      "step": 4497
+    },
+    {
+      "epoch": 0.46,
+      "gumbel_tau": 0.3461,
+      "image_masks_loss": 0.0139,
+      "llm_learning_rate": 2.9362082099021405e-06,
+      "loss": 0.7629,
+      "output_text_mask_loss": 0.0389,
+      "predictor_learning_rate": 0.00029362082099021404,
+      "step": 4498
+    },
+    {
+      "epoch": 0.46,
+      "gumbel_tau": 0.3461,
+      "image_masks_loss": 0.0118,
+      "llm_learning_rate": 2.9353912953714423e-06,
+      "loss": 0.7252,
+      "output_text_mask_loss": 0.0077,
+      "predictor_learning_rate": 0.00029353912953714425,
+      "step": 4499
+    },
+    {
+      "epoch": 0.46,
+      "gumbel_tau": 0.346,
+      "image_masks_loss": 0.0163,
+      "llm_learning_rate": 2.934574332894485e-06,
+      "loss": 0.7298,
+      "output_text_mask_loss": 0.0137,
+      "predictor_learning_rate": 0.00029345743328944843,
+      "step": 4500
+    },
+    {
+      "epoch": 0.46,
+      "gumbel_tau": 0.3459,
+      "image_masks_loss": 0.017,
+      "llm_learning_rate": 2.9337573225612345e-06,
+      "loss": 0.7638,
+      "output_text_mask_loss": 0.0095,
+      "predictor_learning_rate": 0.00029337573225612343,
+      "step": 4501
+    },
+    {
+      "epoch": 0.46,
+      "gumbel_tau": 0.3458,
+      "image_masks_loss": 0.0226,
+      "llm_learning_rate": 2.9329402644616608e-06,
+      "loss": 0.8472,
+      "output_text_mask_loss": 0.0273,
+      "predictor_learning_rate": 0.00029329402644616605,
+      "step": 4502
+    },
+    {
+      "epoch": 0.46,
+      "gumbel_tau": 0.3457,
+      "image_masks_loss": 0.019,
+      "llm_learning_rate": 2.9321231586857407e-06,
+      "loss": 0.7014,
+      "output_text_mask_loss": 0.0362,
+      "predictor_learning_rate": 0.0002932123158685741,
+      "step": 4503
+    },
+    {
+      "epoch": 0.46,
+      "gumbel_tau": 0.3456,
+      "image_masks_loss": 0.0179,
+      "llm_learning_rate": 2.931306005323457e-06,
+      "loss": 0.7303,
+      "output_text_mask_loss": 0.0721,
+      "predictor_learning_rate": 0.0002931306005323457,
+      "step": 4504
+    },
+    {
+      "epoch": 0.46,
+      "gumbel_tau": 0.3456,
+      "image_masks_loss": 0.0147,
+      "llm_learning_rate": 2.930488804464795e-06,
+      "loss": 0.6983,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.00029304888044647947,
+      "step": 4505
+    },
+    {
+      "epoch": 0.46,
+      "gumbel_tau": 0.3455,
+      "image_masks_loss": 0.0133,
+      "llm_learning_rate": 2.929671556199746e-06,
+      "loss": 0.7518,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.0002929671556199746,
+      "step": 4506
+    },
+    {
+      "epoch": 0.46,
+      "gumbel_tau": 0.3454,
+      "image_masks_loss": 0.0156,
+      "llm_learning_rate": 2.928854260618309e-06,
+      "loss": 0.712,
+      "output_text_mask_loss": 0.005,
+      "predictor_learning_rate": 0.0002928854260618309,
+      "step": 4507
+    },
+    {
+      "epoch": 0.46,
+      "gumbel_tau": 0.3453,
+      "image_masks_loss": 0.0011,
+      "llm_learning_rate": 2.928036917810485e-06,
+      "loss": 0.8191,
+      "output_text_mask_loss": 0.002,
+      "predictor_learning_rate": 0.00029280369178104847,
+      "step": 4508
+    },
+    {
+      "epoch": 0.46,
+      "gumbel_tau": 0.3452,
+      "image_masks_loss": 0.0095,
+      "llm_learning_rate": 2.927219527866283e-06,
+      "loss": 0.8052,
+      "output_text_mask_loss": 0.0141,
+      "predictor_learning_rate": 0.00029272195278662826,
+      "step": 4509
+    },
+    {
+      "epoch": 0.46,
+      "gumbel_tau": 0.3452,
+      "image_masks_loss": 0.0198,
+      "llm_learning_rate": 2.926402090875714e-06,
+      "loss": 0.8122,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.0002926402090875714,
+      "step": 4510
+    },
+    {
+      "epoch": 0.46,
+      "gumbel_tau": 0.3451,
+      "image_masks_loss": 0.0235,
+      "llm_learning_rate": 2.9255846069287973e-06,
+      "loss": 0.7648,
+      "output_text_mask_loss": 0.0124,
+      "predictor_learning_rate": 0.00029255846069287973,
+      "step": 4511
+    },
+    {
+      "epoch": 0.46,
+      "gumbel_tau": 0.345,
+      "image_masks_loss": 0.0058,
+      "llm_learning_rate": 2.924767076115557e-06,
+      "loss": 0.682,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 0.0002924767076115556,
+      "step": 4512
+    },
+    {
+      "epoch": 0.46,
+      "gumbel_tau": 0.3449,
+      "image_masks_loss": 0.0189,
+      "llm_learning_rate": 2.9239494985260187e-06,
+      "loss": 0.7793,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.00029239494985260187,
+      "step": 4513
+    },
+    {
+      "epoch": 0.46,
+      "gumbel_tau": 0.3448,
+      "image_masks_loss": 0.0134,
+      "llm_learning_rate": 2.923131874250218e-06,
+      "loss": 0.6598,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.0002923131874250218,
+      "step": 4514
+    },
+    {
+      "epoch": 0.46,
+      "gumbel_tau": 0.3447,
+      "image_masks_loss": 0.0052,
+      "llm_learning_rate": 2.9223142033781925e-06,
+      "loss": 0.84,
+      "output_text_mask_loss": 0.0058,
+      "predictor_learning_rate": 0.00029223142033781924,
+      "step": 4515
+    },
+    {
+      "epoch": 0.46,
+      "gumbel_tau": 0.3447,
+      "image_masks_loss": 0.0093,
+      "llm_learning_rate": 2.9214964859999857e-06,
+      "loss": 0.7833,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.00029214964859999854,
+      "step": 4516
+    },
+    {
+      "epoch": 0.46,
+      "gumbel_tau": 0.3446,
+      "image_masks_loss": 0.0241,
+      "llm_learning_rate": 2.9206787222056477e-06,
+      "loss": 0.747,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.00029206787222056473,
+      "step": 4517
+    },
+    {
+      "epoch": 0.46,
+      "gumbel_tau": 0.3445,
+      "image_masks_loss": 0.0098,
+      "llm_learning_rate": 2.9198609120852317e-06,
+      "loss": 0.7409,
+      "output_text_mask_loss": 0.0332,
+      "predictor_learning_rate": 0.00029198609120852315,
+      "step": 4518
+    },
+    {
+      "epoch": 0.46,
+      "gumbel_tau": 0.3444,
+      "image_masks_loss": 0.0066,
+      "llm_learning_rate": 2.9190430557287964e-06,
+      "loss": 0.7823,
+      "output_text_mask_loss": 0.0092,
+      "predictor_learning_rate": 0.0002919043055728796,
+      "step": 4519
+    },
+    {
+      "epoch": 0.46,
+      "gumbel_tau": 0.3443,
+      "image_masks_loss": 0.0242,
+      "llm_learning_rate": 2.918225153226407e-06,
+      "loss": 0.7093,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0002918225153226407,
+      "step": 4520
+    },
+    {
+      "epoch": 0.46,
+      "gumbel_tau": 0.3443,
+      "image_masks_loss": 0.0137,
+      "llm_learning_rate": 2.9174072046681325e-06,
+      "loss": 0.7908,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00029174072046681324,
+      "step": 4521
+    },
+    {
+      "epoch": 0.46,
+      "gumbel_tau": 0.3442,
+      "image_masks_loss": 0.0229,
+      "llm_learning_rate": 2.916589210144046e-06,
+      "loss": 0.6544,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.0002916589210144046,
+      "step": 4522
+    },
+    {
+      "epoch": 0.46,
+      "gumbel_tau": 0.3441,
+      "image_masks_loss": 0.0311,
+      "llm_learning_rate": 2.9157711697442286e-06,
+      "loss": 0.6596,
+      "output_text_mask_loss": 0.0071,
+      "predictor_learning_rate": 0.00029157711697442283,
+      "step": 4523
+    },
+    {
+      "epoch": 0.46,
+      "gumbel_tau": 0.344,
+      "image_masks_loss": 0.0278,
+      "llm_learning_rate": 2.9149530835587637e-06,
+      "loss": 0.669,
+      "output_text_mask_loss": 0.0204,
+      "predictor_learning_rate": 0.00029149530835587633,
+      "step": 4524
+    },
+    {
+      "epoch": 0.46,
+      "gumbel_tau": 0.3439,
+      "image_masks_loss": 0.0381,
+      "llm_learning_rate": 2.9141349516777407e-06,
+      "loss": 0.841,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 0.00029141349516777405,
+      "step": 4525
+    },
+    {
+      "epoch": 0.46,
+      "gumbel_tau": 0.3439,
+      "image_masks_loss": 0.0271,
+      "llm_learning_rate": 2.9133167741912543e-06,
+      "loss": 0.7471,
+      "output_text_mask_loss": 0.0282,
+      "predictor_learning_rate": 0.00029133167741912544,
+      "step": 4526
+    },
+    {
+      "epoch": 0.46,
+      "gumbel_tau": 0.3438,
+      "image_masks_loss": 0.0172,
+      "llm_learning_rate": 2.912498551189404e-06,
+      "loss": 0.765,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.00029124985511894044,
+      "step": 4527
+    },
+    {
+      "epoch": 0.46,
+      "gumbel_tau": 0.3437,
+      "image_masks_loss": 0.0073,
+      "llm_learning_rate": 2.9116802827622948e-06,
+      "loss": 0.7618,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 0.00029116802827622946,
+      "step": 4528
+    },
+    {
+      "epoch": 0.46,
+      "gumbel_tau": 0.3436,
+      "image_masks_loss": 0.0263,
+      "llm_learning_rate": 2.9108619690000355e-06,
+      "loss": 0.7403,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00029108619690000354,
+      "step": 4529
+    },
+    {
+      "epoch": 0.46,
+      "gumbel_tau": 0.3435,
+      "image_masks_loss": 0.0133,
+      "llm_learning_rate": 2.9100436099927415e-06,
+      "loss": 0.7471,
+      "output_text_mask_loss": 0.0546,
+      "predictor_learning_rate": 0.00029100436099927413,
+      "step": 4530
+    },
+    {
+      "epoch": 0.46,
+      "gumbel_tau": 0.3434,
+      "image_masks_loss": 0.0169,
+      "llm_learning_rate": 2.909225205830532e-06,
+      "loss": 0.7112,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.00029092252058305316,
+      "step": 4531
+    },
+    {
+      "epoch": 0.46,
+      "gumbel_tau": 0.3434,
+      "image_masks_loss": 0.0233,
+      "llm_learning_rate": 2.908406756603531e-06,
+      "loss": 0.8222,
+      "output_text_mask_loss": 0.0391,
+      "predictor_learning_rate": 0.000290840675660353,
+      "step": 4532
+    },
+    {
+      "epoch": 0.46,
+      "gumbel_tau": 0.3433,
+      "image_masks_loss": 0.0181,
+      "llm_learning_rate": 2.907588262401868e-06,
+      "loss": 0.663,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.0002907588262401868,
+      "step": 4533
+    },
+    {
+      "epoch": 0.46,
+      "gumbel_tau": 0.3432,
+      "image_masks_loss": 0.0233,
+      "llm_learning_rate": 2.9067697233156776e-06,
+      "loss": 0.7499,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00029067697233156777,
+      "step": 4534
+    },
+    {
+      "epoch": 0.46,
+      "gumbel_tau": 0.3431,
+      "image_masks_loss": 0.0269,
+      "llm_learning_rate": 2.9059511394350993e-06,
+      "loss": 0.7107,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.0002905951139435099,
+      "step": 4535
+    },
+    {
+      "epoch": 0.46,
+      "gumbel_tau": 0.343,
+      "image_masks_loss": 0.0278,
+      "llm_learning_rate": 2.9051325108502774e-06,
+      "loss": 0.7268,
+      "output_text_mask_loss": 0.0116,
+      "predictor_learning_rate": 0.0002905132510850277,
+      "step": 4536
+    },
+    {
+      "epoch": 0.46,
+      "gumbel_tau": 0.343,
+      "image_masks_loss": 0.0069,
+      "llm_learning_rate": 2.904313837651361e-06,
+      "loss": 0.7561,
+      "output_text_mask_loss": 0.0603,
+      "predictor_learning_rate": 0.0002904313837651361,
+      "step": 4537
+    },
+    {
+      "epoch": 0.46,
+      "gumbel_tau": 0.3429,
+      "image_masks_loss": 0.0127,
+      "llm_learning_rate": 2.9034951199285044e-06,
+      "loss": 0.673,
+      "output_text_mask_loss": 0.0088,
+      "predictor_learning_rate": 0.0002903495119928504,
+      "step": 4538
+    },
+    {
+      "epoch": 0.47,
+      "gumbel_tau": 0.3428,
+      "image_masks_loss": 0.003,
+      "llm_learning_rate": 2.9026763577718663e-06,
+      "loss": 0.7386,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0002902676357771866,
+      "step": 4539
+    },
+    {
+      "epoch": 0.47,
+      "gumbel_tau": 0.3427,
+      "image_masks_loss": 0.0257,
+      "llm_learning_rate": 2.901857551271611e-06,
+      "loss": 0.7712,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.00029018575512716104,
+      "step": 4540
+    },
+    {
+      "epoch": 0.47,
+      "gumbel_tau": 0.3426,
+      "image_masks_loss": 0.0195,
+      "llm_learning_rate": 2.9010387005179073e-06,
+      "loss": 0.7091,
+      "output_text_mask_loss": 0.0238,
+      "predictor_learning_rate": 0.00029010387005179076,
+      "step": 4541
+    },
+    {
+      "epoch": 0.47,
+      "gumbel_tau": 0.3426,
+      "image_masks_loss": 0.0282,
+      "llm_learning_rate": 2.900219805600928e-06,
+      "loss": 0.673,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00029002198056009277,
+      "step": 4542
+    },
+    {
+      "epoch": 0.47,
+      "gumbel_tau": 0.3425,
+      "image_masks_loss": 0.0106,
+      "llm_learning_rate": 2.8994008666108523e-06,
+      "loss": 0.6853,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.00028994008666108524,
+      "step": 4543
+    },
+    {
+      "epoch": 0.47,
+      "gumbel_tau": 0.3424,
+      "image_masks_loss": 0.0271,
+      "llm_learning_rate": 2.8985818836378638e-06,
+      "loss": 0.7874,
+      "output_text_mask_loss": 0.0446,
+      "predictor_learning_rate": 0.00028985818836378636,
+      "step": 4544
+    },
+    {
+      "epoch": 0.47,
+      "gumbel_tau": 0.3423,
+      "image_masks_loss": 0.0202,
+      "llm_learning_rate": 2.89776285677215e-06,
+      "loss": 0.7914,
+      "output_text_mask_loss": 0.069,
+      "predictor_learning_rate": 0.00028977628567721496,
+      "step": 4545
+    },
+    {
+      "epoch": 0.47,
+      "gumbel_tau": 0.3422,
+      "image_masks_loss": 0.017,
+      "llm_learning_rate": 2.896943786103905e-06,
+      "loss": 0.7105,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.0002896943786103905,
+      "step": 4546
+    },
+    {
+      "epoch": 0.47,
+      "gumbel_tau": 0.3422,
+      "image_masks_loss": 0.011,
+      "llm_learning_rate": 2.896124671723326e-06,
+      "loss": 0.7672,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.0002896124671723326,
+      "step": 4547
+    },
+    {
+      "epoch": 0.47,
+      "gumbel_tau": 0.3421,
+      "image_masks_loss": 0.012,
+      "llm_learning_rate": 2.8953055137206154e-06,
+      "loss": 0.7409,
+      "output_text_mask_loss": 0.0156,
+      "predictor_learning_rate": 0.00028953055137206153,
+      "step": 4548
+    },
+    {
+      "epoch": 0.47,
+      "gumbel_tau": 0.342,
+      "image_masks_loss": 0.0114,
+      "llm_learning_rate": 2.8944863121859812e-06,
+      "loss": 0.6821,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.0002894486312185981,
+      "step": 4549
+    },
+    {
+      "epoch": 0.47,
+      "gumbel_tau": 0.3419,
+      "image_masks_loss": 0.0105,
+      "llm_learning_rate": 2.893667067209637e-06,
+      "loss": 0.7046,
+      "output_text_mask_loss": 0.0402,
+      "predictor_learning_rate": 0.00028936670672096365,
+      "step": 4550
+    },
+    {
+      "epoch": 0.47,
+      "gumbel_tau": 0.3418,
+      "image_masks_loss": 0.0053,
+      "llm_learning_rate": 2.8928477788817964e-06,
+      "loss": 0.772,
+      "output_text_mask_loss": 0.0717,
+      "predictor_learning_rate": 0.00028928477788817966,
+      "step": 4551
+    },
+    {
+      "epoch": 0.47,
+      "gumbel_tau": 0.3418,
+      "image_masks_loss": 0.0299,
+      "llm_learning_rate": 2.8920284472926846e-06,
+      "loss": 0.6649,
+      "output_text_mask_loss": 0.1602,
+      "predictor_learning_rate": 0.00028920284472926843,
+      "step": 4552
+    },
+    {
+      "epoch": 0.47,
+      "gumbel_tau": 0.3417,
+      "image_masks_loss": 0.0036,
+      "llm_learning_rate": 2.8912090725325265e-06,
+      "loss": 0.713,
+      "output_text_mask_loss": 0.0353,
+      "predictor_learning_rate": 0.00028912090725325264,
+      "step": 4553
+    },
+    {
+      "epoch": 0.47,
+      "gumbel_tau": 0.3416,
+      "image_masks_loss": 0.0155,
+      "llm_learning_rate": 2.8903896546915544e-06,
+      "loss": 0.7334,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00028903896546915545,
+      "step": 4554
+    },
+    {
+      "epoch": 0.47,
+      "gumbel_tau": 0.3415,
+      "image_masks_loss": 0.006,
+      "llm_learning_rate": 2.8895701938600036e-06,
+      "loss": 0.7752,
+      "output_text_mask_loss": 0.0473,
+      "predictor_learning_rate": 0.00028895701938600035,
+      "step": 4555
+    },
+    {
+      "epoch": 0.47,
+      "gumbel_tau": 0.3414,
+      "image_masks_loss": 0.0201,
+      "llm_learning_rate": 2.8887506901281153e-06,
+      "loss": 0.7589,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.00028887506901281154,
+      "step": 4556
+    },
+    {
+      "epoch": 0.47,
+      "gumbel_tau": 0.3413,
+      "image_masks_loss": 0.0126,
+      "llm_learning_rate": 2.8879311435861353e-06,
+      "loss": 0.7432,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.0002887931143586135,
+      "step": 4557
+    },
+    {
+      "epoch": 0.47,
+      "gumbel_tau": 0.3413,
+      "image_masks_loss": 0.0093,
+      "llm_learning_rate": 2.8871115543243133e-06,
+      "loss": 0.7497,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 0.00028871115543243134,
+      "step": 4558
+    },
+    {
+      "epoch": 0.47,
+      "gumbel_tau": 0.3412,
+      "image_masks_loss": 0.0237,
+      "llm_learning_rate": 2.8862919224329058e-06,
+      "loss": 0.7966,
+      "output_text_mask_loss": 0.0278,
+      "predictor_learning_rate": 0.00028862919224329056,
+      "step": 4559
+    },
+    {
+      "epoch": 0.47,
+      "gumbel_tau": 0.3411,
+      "image_masks_loss": 0.0374,
+      "llm_learning_rate": 2.8854722480021704e-06,
+      "loss": 0.6649,
+      "output_text_mask_loss": 0.0151,
+      "predictor_learning_rate": 0.000288547224800217,
+      "step": 4560
+    },
+    {
+      "epoch": 0.47,
+      "gumbel_tau": 0.341,
+      "image_masks_loss": 0.0145,
+      "llm_learning_rate": 2.884652531122372e-06,
+      "loss": 0.7466,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0002884652531122372,
+      "step": 4561
+    },
+    {
+      "epoch": 0.47,
+      "gumbel_tau": 0.3409,
+      "image_masks_loss": 0.0056,
+      "llm_learning_rate": 2.8838327718837804e-06,
+      "loss": 0.7254,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.000288383277188378,
+      "step": 4562
+    },
+    {
+      "epoch": 0.47,
+      "gumbel_tau": 0.3409,
+      "image_masks_loss": 0.0092,
+      "llm_learning_rate": 2.8830129703766696e-06,
+      "loss": 0.7372,
+      "output_text_mask_loss": 0.0035,
+      "predictor_learning_rate": 0.00028830129703766695,
+      "step": 4563
+    },
+    {
+      "epoch": 0.47,
+      "gumbel_tau": 0.3408,
+      "image_masks_loss": 0.0042,
+      "llm_learning_rate": 2.8821931266913173e-06,
+      "loss": 0.7082,
+      "output_text_mask_loss": 0.0322,
+      "predictor_learning_rate": 0.00028821931266913173,
+      "step": 4564
+    },
+    {
+      "epoch": 0.47,
+      "gumbel_tau": 0.3407,
+      "image_masks_loss": 0.0179,
+      "llm_learning_rate": 2.8813732409180062e-06,
+      "loss": 0.6671,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.0002881373240918006,
+      "step": 4565
+    },
+    {
+      "epoch": 0.47,
+      "gumbel_tau": 0.3406,
+      "image_masks_loss": 0.0137,
+      "llm_learning_rate": 2.880553313147024e-06,
+      "loss": 0.7157,
+      "output_text_mask_loss": 0.0702,
+      "predictor_learning_rate": 0.0002880553313147024,
+      "step": 4566
+    },
+    {
+      "epoch": 0.47,
+      "gumbel_tau": 0.3405,
+      "image_masks_loss": 0.0214,
+      "llm_learning_rate": 2.8797333434686643e-06,
+      "loss": 0.7097,
+      "output_text_mask_loss": 0.0012,
+      "predictor_learning_rate": 0.0002879733343468664,
+      "step": 4567
+    },
+    {
+      "epoch": 0.47,
+      "gumbel_tau": 0.3405,
+      "image_masks_loss": 0.0284,
+      "llm_learning_rate": 2.8789133319732225e-06,
+      "loss": 0.6953,
+      "output_text_mask_loss": 0.0027,
+      "predictor_learning_rate": 0.00028789133319732224,
+      "step": 4568
+    },
+    {
+      "epoch": 0.47,
+      "gumbel_tau": 0.3404,
+      "image_masks_loss": 0.0161,
+      "llm_learning_rate": 2.8780932787510006e-06,
+      "loss": 0.7975,
+      "output_text_mask_loss": 0.0071,
+      "predictor_learning_rate": 0.00028780932787510004,
+      "step": 4569
+    },
+    {
+      "epoch": 0.47,
+      "gumbel_tau": 0.3403,
+      "image_masks_loss": 0.0176,
+      "llm_learning_rate": 2.877273183892305e-06,
+      "loss": 0.7327,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0002877273183892305,
+      "step": 4570
+    },
+    {
+      "epoch": 0.47,
+      "gumbel_tau": 0.3402,
+      "image_masks_loss": 0.0248,
+      "llm_learning_rate": 2.8764530474874454e-06,
+      "loss": 0.7,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.0002876453047487445,
+      "step": 4571
+    },
+    {
+      "epoch": 0.47,
+      "gumbel_tau": 0.3401,
+      "image_masks_loss": 0.0166,
+      "llm_learning_rate": 2.8756328696267376e-06,
+      "loss": 0.7351,
+      "output_text_mask_loss": 0.0269,
+      "predictor_learning_rate": 0.00028756328696267374,
+      "step": 4572
+    },
+    {
+      "epoch": 0.47,
+      "gumbel_tau": 0.3401,
+      "image_masks_loss": 0.0052,
+      "llm_learning_rate": 2.874812650400502e-06,
+      "loss": 0.5748,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.0002874812650400502,
+      "step": 4573
+    },
+    {
+      "epoch": 0.47,
+      "gumbel_tau": 0.34,
+      "image_masks_loss": 0.0125,
+      "llm_learning_rate": 2.873992389899063e-06,
+      "loss": 0.8247,
+      "output_text_mask_loss": 0.0793,
+      "predictor_learning_rate": 0.0002873992389899063,
+      "step": 4574
+    },
+    {
+      "epoch": 0.47,
+      "gumbel_tau": 0.3399,
+      "image_masks_loss": 0.0071,
+      "llm_learning_rate": 2.873172088212748e-06,
+      "loss": 0.6425,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00028731720882127477,
+      "step": 4575
+    },
+    {
+      "epoch": 0.47,
+      "gumbel_tau": 0.3398,
+      "image_masks_loss": 0.0154,
+      "llm_learning_rate": 2.8723517454318915e-06,
+      "loss": 0.7477,
+      "output_text_mask_loss": 0.0479,
+      "predictor_learning_rate": 0.00028723517454318914,
+      "step": 4576
+    },
+    {
+      "epoch": 0.47,
+      "gumbel_tau": 0.3397,
+      "image_masks_loss": 0.025,
+      "llm_learning_rate": 2.8715313616468317e-06,
+      "loss": 0.8411,
+      "output_text_mask_loss": 0.0284,
+      "predictor_learning_rate": 0.00028715313616468316,
+      "step": 4577
+    },
+    {
+      "epoch": 0.47,
+      "gumbel_tau": 0.3397,
+      "image_masks_loss": 0.0343,
+      "llm_learning_rate": 2.870710936947911e-06,
+      "loss": 0.7172,
+      "output_text_mask_loss": 0.0202,
+      "predictor_learning_rate": 0.00028707109369479104,
+      "step": 4578
+    },
+    {
+      "epoch": 0.47,
+      "gumbel_tau": 0.3396,
+      "image_masks_loss": 0.0132,
+      "llm_learning_rate": 2.869890471425476e-06,
+      "loss": 0.7158,
+      "output_text_mask_loss": 0.0755,
+      "predictor_learning_rate": 0.0002869890471425476,
+      "step": 4579
+    },
+    {
+      "epoch": 0.47,
+      "gumbel_tau": 0.3395,
+      "image_masks_loss": 0.0102,
+      "llm_learning_rate": 2.8690699651698784e-06,
+      "loss": 0.7887,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.0002869069965169878,
+      "step": 4580
+    },
+    {
+      "epoch": 0.47,
+      "gumbel_tau": 0.3394,
+      "image_masks_loss": 0.0114,
+      "llm_learning_rate": 2.868249418271474e-06,
+      "loss": 0.8377,
+      "output_text_mask_loss": 0.0204,
+      "predictor_learning_rate": 0.00028682494182714734,
+      "step": 4581
+    },
+    {
+      "epoch": 0.47,
+      "gumbel_tau": 0.3393,
+      "image_masks_loss": 0.0115,
+      "llm_learning_rate": 2.8674288308206238e-06,
+      "loss": 0.7265,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.0002867428830820624,
+      "step": 4582
+    },
+    {
+      "epoch": 0.47,
+      "gumbel_tau": 0.3393,
+      "image_masks_loss": 0.0194,
+      "llm_learning_rate": 2.866608202907692e-06,
+      "loss": 0.7745,
+      "output_text_mask_loss": 0.0164,
+      "predictor_learning_rate": 0.00028666082029076917,
+      "step": 4583
+    },
+    {
+      "epoch": 0.47,
+      "gumbel_tau": 0.3392,
+      "image_masks_loss": 0.0234,
+      "llm_learning_rate": 2.865787534623048e-06,
+      "loss": 0.7523,
+      "output_text_mask_loss": 0.0277,
+      "predictor_learning_rate": 0.00028657875346230476,
+      "step": 4584
+    },
+    {
+      "epoch": 0.47,
+      "gumbel_tau": 0.3391,
+      "image_masks_loss": 0.0163,
+      "llm_learning_rate": 2.864966826057066e-06,
+      "loss": 0.6899,
+      "output_text_mask_loss": 0.0237,
+      "predictor_learning_rate": 0.0002864966826057066,
+      "step": 4585
+    },
+    {
+      "epoch": 0.47,
+      "gumbel_tau": 0.339,
+      "image_masks_loss": 0.0087,
+      "llm_learning_rate": 2.864146077300125e-06,
+      "loss": 0.8246,
+      "output_text_mask_loss": 0.006,
+      "predictor_learning_rate": 0.0002864146077300125,
+      "step": 4586
+    },
+    {
+      "epoch": 0.47,
+      "gumbel_tau": 0.3389,
+      "image_masks_loss": 0.0034,
+      "llm_learning_rate": 2.8633252884426065e-06,
+      "loss": 0.7844,
+      "output_text_mask_loss": 0.0219,
+      "predictor_learning_rate": 0.00028633252884426063,
+      "step": 4587
+    },
+    {
+      "epoch": 0.47,
+      "gumbel_tau": 0.3389,
+      "image_masks_loss": 0.0072,
+      "llm_learning_rate": 2.862504459574898e-06,
+      "loss": 0.7808,
+      "output_text_mask_loss": 0.0012,
+      "predictor_learning_rate": 0.0002862504459574898,
+      "step": 4588
+    },
+    {
+      "epoch": 0.47,
+      "gumbel_tau": 0.3388,
+      "image_masks_loss": 0.0109,
+      "llm_learning_rate": 2.861683590787391e-06,
+      "loss": 0.7882,
+      "output_text_mask_loss": 0.0171,
+      "predictor_learning_rate": 0.0002861683590787391,
+      "step": 4589
+    },
+    {
+      "epoch": 0.47,
+      "gumbel_tau": 0.3387,
+      "image_masks_loss": 0.0178,
+      "llm_learning_rate": 2.860862682170481e-06,
+      "loss": 0.7328,
+      "output_text_mask_loss": 0.0012,
+      "predictor_learning_rate": 0.0002860862682170481,
+      "step": 4590
+    },
+    {
+      "epoch": 0.47,
+      "gumbel_tau": 0.3386,
+      "image_masks_loss": 0.0098,
+      "llm_learning_rate": 2.8600417338145685e-06,
+      "loss": 0.7636,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00028600417338145684,
+      "step": 4591
+    },
+    {
+      "epoch": 0.47,
+      "gumbel_tau": 0.3385,
+      "image_masks_loss": 0.0096,
+      "llm_learning_rate": 2.859220745810059e-06,
+      "loss": 0.7037,
+      "output_text_mask_loss": 0.0023,
+      "predictor_learning_rate": 0.0002859220745810059,
+      "step": 4592
+    },
+    {
+      "epoch": 0.47,
+      "gumbel_tau": 0.3385,
+      "image_masks_loss": 0.0182,
+      "llm_learning_rate": 2.8583997182473604e-06,
+      "loss": 0.8234,
+      "output_text_mask_loss": 0.0208,
+      "predictor_learning_rate": 0.00028583997182473605,
+      "step": 4593
+    },
+    {
+      "epoch": 0.47,
+      "gumbel_tau": 0.3384,
+      "image_masks_loss": 0.0039,
+      "llm_learning_rate": 2.8575786512168863e-06,
+      "loss": 0.7354,
+      "output_text_mask_loss": 0.1198,
+      "predictor_learning_rate": 0.0002857578651216886,
+      "step": 4594
+    },
+    {
+      "epoch": 0.47,
+      "gumbel_tau": 0.3383,
+      "image_masks_loss": 0.0126,
+      "llm_learning_rate": 2.856757544809055e-06,
+      "loss": 0.6342,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.0002856757544809055,
+      "step": 4595
+    },
+    {
+      "epoch": 0.47,
+      "gumbel_tau": 0.3382,
+      "image_masks_loss": 0.0216,
+      "llm_learning_rate": 2.8559363991142885e-06,
+      "loss": 0.7085,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.00028559363991142886,
+      "step": 4596
+    },
+    {
+      "epoch": 0.47,
+      "gumbel_tau": 0.3381,
+      "image_masks_loss": 0.0208,
+      "llm_learning_rate": 2.8551152142230127e-06,
+      "loss": 0.7419,
+      "output_text_mask_loss": 0.1312,
+      "predictor_learning_rate": 0.00028551152142230124,
+      "step": 4597
+    },
+    {
+      "epoch": 0.47,
+      "gumbel_tau": 0.3381,
+      "image_masks_loss": 0.0146,
+      "llm_learning_rate": 2.8542939902256584e-06,
+      "loss": 0.7569,
+      "output_text_mask_loss": 0.0086,
+      "predictor_learning_rate": 0.0002854293990225658,
+      "step": 4598
+    },
+    {
+      "epoch": 0.47,
+      "gumbel_tau": 0.338,
+      "image_masks_loss": 0.0156,
+      "llm_learning_rate": 2.8534727272126605e-06,
+      "loss": 0.6782,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.000285347272721266,
+      "step": 4599
+    },
+    {
+      "epoch": 0.47,
+      "gumbel_tau": 0.3379,
+      "image_masks_loss": 0.0101,
+      "llm_learning_rate": 2.8526514252744587e-06,
+      "loss": 0.8082,
+      "output_text_mask_loss": 0.0158,
+      "predictor_learning_rate": 0.00028526514252744583,
+      "step": 4600
+    },
+    {
+      "epoch": 0.47,
+      "gumbel_tau": 0.3378,
+      "image_masks_loss": 0.0069,
+      "llm_learning_rate": 2.8518300845014967e-06,
+      "loss": 0.7714,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00028518300845014965,
+      "step": 4601
+    },
+    {
+      "epoch": 0.47,
+      "gumbel_tau": 0.3377,
+      "image_masks_loss": 0.0126,
+      "llm_learning_rate": 2.8510087049842215e-06,
+      "loss": 0.7248,
+      "output_text_mask_loss": 0.0086,
+      "predictor_learning_rate": 0.00028510087049842214,
+      "step": 4602
+    },
+    {
+      "epoch": 0.47,
+      "gumbel_tau": 0.3377,
+      "image_masks_loss": 0.0296,
+      "llm_learning_rate": 2.850187286813087e-06,
+      "loss": 0.7497,
+      "output_text_mask_loss": 0.0137,
+      "predictor_learning_rate": 0.0002850187286813087,
+      "step": 4603
+    },
+    {
+      "epoch": 0.47,
+      "gumbel_tau": 0.3376,
+      "image_masks_loss": 0.0186,
+      "llm_learning_rate": 2.849365830078547e-06,
+      "loss": 0.8237,
+      "output_text_mask_loss": 0.0028,
+      "predictor_learning_rate": 0.0002849365830078547,
+      "step": 4604
+    },
+    {
+      "epoch": 0.47,
+      "gumbel_tau": 0.3375,
+      "image_masks_loss": 0.0179,
+      "llm_learning_rate": 2.8485443348710646e-06,
+      "loss": 0.7958,
+      "output_text_mask_loss": 0.0248,
+      "predictor_learning_rate": 0.0002848544334871064,
+      "step": 4605
+    },
+    {
+      "epoch": 0.47,
+      "gumbel_tau": 0.3374,
+      "image_masks_loss": 0.0165,
+      "llm_learning_rate": 2.8477228012811043e-06,
+      "loss": 0.7231,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.0002847722801281104,
+      "step": 4606
+    },
+    {
+      "epoch": 0.47,
+      "gumbel_tau": 0.3373,
+      "image_masks_loss": 0.0174,
+      "llm_learning_rate": 2.8469012293991333e-06,
+      "loss": 0.7221,
+      "output_text_mask_loss": 0.0099,
+      "predictor_learning_rate": 0.00028469012293991334,
+      "step": 4607
+    },
+    {
+      "epoch": 0.47,
+      "gumbel_tau": 0.3373,
+      "image_masks_loss": 0.0301,
+      "llm_learning_rate": 2.846079619315627e-06,
+      "loss": 0.7788,
+      "output_text_mask_loss": 0.0076,
+      "predictor_learning_rate": 0.0002846079619315627,
+      "step": 4608
+    },
+    {
+      "epoch": 0.47,
+      "gumbel_tau": 0.3372,
+      "image_masks_loss": 0.0259,
+      "llm_learning_rate": 2.845257971121062e-06,
+      "loss": 0.7703,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0002845257971121062,
+      "step": 4609
+    },
+    {
+      "epoch": 0.47,
+      "gumbel_tau": 0.3371,
+      "image_masks_loss": 0.0254,
+      "llm_learning_rate": 2.8444362849059193e-06,
+      "loss": 0.7104,
+      "output_text_mask_loss": 0.0139,
+      "predictor_learning_rate": 0.0002844436284905919,
+      "step": 4610
+    },
+    {
+      "epoch": 0.47,
+      "gumbel_tau": 0.337,
+      "image_masks_loss": 0.0083,
+      "llm_learning_rate": 2.843614560760687e-06,
+      "loss": 0.814,
+      "output_text_mask_loss": 0.0078,
+      "predictor_learning_rate": 0.0002843614560760687,
+      "step": 4611
+    },
+    {
+      "epoch": 0.47,
+      "gumbel_tau": 0.3369,
+      "image_masks_loss": 0.0085,
+      "llm_learning_rate": 2.842792798775853e-06,
+      "loss": 0.86,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00028427927987758527,
+      "step": 4612
+    },
+    {
+      "epoch": 0.47,
+      "gumbel_tau": 0.3369,
+      "image_masks_loss": 0.0311,
+      "llm_learning_rate": 2.841970999041913e-06,
+      "loss": 0.8412,
+      "output_text_mask_loss": 0.0076,
+      "predictor_learning_rate": 0.0002841970999041913,
+      "step": 4613
+    },
+    {
+      "epoch": 0.47,
+      "gumbel_tau": 0.3368,
+      "image_masks_loss": 0.0153,
+      "llm_learning_rate": 2.841149161649365e-06,
+      "loss": 0.7989,
+      "output_text_mask_loss": 0.037,
+      "predictor_learning_rate": 0.0002841149161649365,
+      "step": 4614
+    },
+    {
+      "epoch": 0.47,
+      "gumbel_tau": 0.3367,
+      "image_masks_loss": 0.0127,
+      "llm_learning_rate": 2.840327286688711e-06,
+      "loss": 0.7478,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0002840327286688711,
+      "step": 4615
+    },
+    {
+      "epoch": 0.47,
+      "gumbel_tau": 0.3366,
+      "image_masks_loss": 0.0227,
+      "llm_learning_rate": 2.8395053742504582e-06,
+      "loss": 0.6661,
+      "output_text_mask_loss": 0.0092,
+      "predictor_learning_rate": 0.0002839505374250458,
+      "step": 4616
+    },
+    {
+      "epoch": 0.47,
+      "gumbel_tau": 0.3365,
+      "image_masks_loss": 0.0186,
+      "llm_learning_rate": 2.838683424425118e-06,
+      "loss": 0.7118,
+      "output_text_mask_loss": 0.0231,
+      "predictor_learning_rate": 0.00028386834244251174,
+      "step": 4617
+    },
+    {
+      "epoch": 0.47,
+      "gumbel_tau": 0.3365,
+      "image_masks_loss": 0.0141,
+      "llm_learning_rate": 2.8378614373032036e-06,
+      "loss": 0.7677,
+      "output_text_mask_loss": 0.01,
+      "predictor_learning_rate": 0.00028378614373032035,
+      "step": 4618
+    },
+    {
+      "epoch": 0.47,
+      "gumbel_tau": 0.3364,
+      "image_masks_loss": 0.0303,
+      "llm_learning_rate": 2.837039412975236e-06,
+      "loss": 0.6697,
+      "output_text_mask_loss": 0.0177,
+      "predictor_learning_rate": 0.0002837039412975236,
+      "step": 4619
+    },
+    {
+      "epoch": 0.47,
+      "gumbel_tau": 0.3363,
+      "image_masks_loss": 0.004,
+      "llm_learning_rate": 2.836217351531737e-06,
+      "loss": 0.7679,
+      "output_text_mask_loss": 0.0465,
+      "predictor_learning_rate": 0.0002836217351531737,
+      "step": 4620
+    },
+    {
+      "epoch": 0.47,
+      "gumbel_tau": 0.3362,
+      "image_masks_loss": 0.0218,
+      "llm_learning_rate": 2.8353952530632344e-06,
+      "loss": 0.8292,
+      "output_text_mask_loss": 0.038,
+      "predictor_learning_rate": 0.0002835395253063234,
+      "step": 4621
+    },
+    {
+      "epoch": 0.47,
+      "gumbel_tau": 0.3362,
+      "image_masks_loss": 0.0187,
+      "llm_learning_rate": 2.8345731176602598e-06,
+      "loss": 0.7166,
+      "output_text_mask_loss": 0.0534,
+      "predictor_learning_rate": 0.00028345731176602595,
+      "step": 4622
+    },
+    {
+      "epoch": 0.47,
+      "gumbel_tau": 0.3361,
+      "image_masks_loss": 0.014,
+      "llm_learning_rate": 2.8337509454133475e-06,
+      "loss": 0.7431,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.00028337509454133477,
+      "step": 4623
+    },
+    {
+      "epoch": 0.47,
+      "gumbel_tau": 0.336,
+      "image_masks_loss": 0.0078,
+      "llm_learning_rate": 2.832928736413039e-06,
+      "loss": 0.7291,
+      "output_text_mask_loss": 0.0065,
+      "predictor_learning_rate": 0.0002832928736413039,
+      "step": 4624
+    },
+    {
+      "epoch": 0.47,
+      "gumbel_tau": 0.3359,
+      "image_masks_loss": 0.0094,
+      "llm_learning_rate": 2.832106490749876e-06,
+      "loss": 0.7019,
+      "output_text_mask_loss": 0.0275,
+      "predictor_learning_rate": 0.0002832106490749876,
+      "step": 4625
+    },
+    {
+      "epoch": 0.47,
+      "gumbel_tau": 0.3358,
+      "image_masks_loss": 0.0152,
+      "llm_learning_rate": 2.831284208514406e-06,
+      "loss": 0.7449,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.00028312842085144056,
+      "step": 4626
+    },
+    {
+      "epoch": 0.47,
+      "gumbel_tau": 0.3358,
+      "image_masks_loss": 0.0103,
+      "llm_learning_rate": 2.8304618897971812e-06,
+      "loss": 0.751,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.0002830461889797181,
+      "step": 4627
+    },
+    {
+      "epoch": 0.47,
+      "gumbel_tau": 0.3357,
+      "image_masks_loss": 0.016,
+      "llm_learning_rate": 2.8296395346887577e-06,
+      "loss": 0.7358,
+      "output_text_mask_loss": 0.005,
+      "predictor_learning_rate": 0.00028296395346887575,
+      "step": 4628
+    },
+    {
+      "epoch": 0.47,
+      "gumbel_tau": 0.3356,
+      "image_masks_loss": 0.0239,
+      "llm_learning_rate": 2.828817143279694e-06,
+      "loss": 0.8215,
+      "output_text_mask_loss": 0.0216,
+      "predictor_learning_rate": 0.0002828817143279694,
+      "step": 4629
+    },
+    {
+      "epoch": 0.47,
+      "gumbel_tau": 0.3355,
+      "image_masks_loss": 0.0095,
+      "llm_learning_rate": 2.8279947156605534e-06,
+      "loss": 0.738,
+      "output_text_mask_loss": 0.0071,
+      "predictor_learning_rate": 0.00028279947156605533,
+      "step": 4630
+    },
+    {
+      "epoch": 0.47,
+      "gumbel_tau": 0.3354,
+      "image_masks_loss": 0.02,
+      "llm_learning_rate": 2.8271722519219048e-06,
+      "loss": 0.7654,
+      "output_text_mask_loss": 0.0187,
+      "predictor_learning_rate": 0.00028271722519219045,
+      "step": 4631
+    },
+    {
+      "epoch": 0.47,
+      "gumbel_tau": 0.3354,
+      "image_masks_loss": 0.0149,
+      "llm_learning_rate": 2.8263497521543192e-06,
+      "loss": 0.7724,
+      "output_text_mask_loss": 0.0153,
+      "predictor_learning_rate": 0.00028263497521543193,
+      "step": 4632
+    },
+    {
+      "epoch": 0.47,
+      "gumbel_tau": 0.3353,
+      "image_masks_loss": 0.014,
+      "llm_learning_rate": 2.8255272164483717e-06,
+      "loss": 0.7812,
+      "output_text_mask_loss": 0.0004,
+      "predictor_learning_rate": 0.00028255272164483714,
+      "step": 4633
+    },
+    {
+      "epoch": 0.47,
+      "gumbel_tau": 0.3352,
+      "image_masks_loss": 0.0218,
+      "llm_learning_rate": 2.824704644894642e-06,
+      "loss": 0.7626,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0002824704644894642,
+      "step": 4634
+    },
+    {
+      "epoch": 0.47,
+      "gumbel_tau": 0.3351,
+      "image_masks_loss": 0.0207,
+      "llm_learning_rate": 2.823882037583714e-06,
+      "loss": 0.7894,
+      "output_text_mask_loss": 0.0526,
+      "predictor_learning_rate": 0.0002823882037583714,
+      "step": 4635
+    },
+    {
+      "epoch": 0.47,
+      "gumbel_tau": 0.335,
+      "image_masks_loss": 0.0025,
+      "llm_learning_rate": 2.8230593946061736e-06,
+      "loss": 0.735,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0002823059394606173,
+      "step": 4636
+    },
+    {
+      "epoch": 0.48,
+      "gumbel_tau": 0.335,
+      "image_masks_loss": 0.0161,
+      "llm_learning_rate": 2.822236716052614e-06,
+      "loss": 0.7326,
+      "output_text_mask_loss": 0.017,
+      "predictor_learning_rate": 0.0002822236716052614,
+      "step": 4637
+    },
+    {
+      "epoch": 0.48,
+      "gumbel_tau": 0.3349,
+      "image_masks_loss": 0.0294,
+      "llm_learning_rate": 2.821414002013629e-06,
+      "loss": 0.6791,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0002821414002013629,
+      "step": 4638
+    },
+    {
+      "epoch": 0.48,
+      "gumbel_tau": 0.3348,
+      "image_masks_loss": 0.0275,
+      "llm_learning_rate": 2.820591252579818e-06,
+      "loss": 0.8269,
+      "output_text_mask_loss": 0.0055,
+      "predictor_learning_rate": 0.00028205912525798175,
+      "step": 4639
+    },
+    {
+      "epoch": 0.48,
+      "gumbel_tau": 0.3347,
+      "image_masks_loss": 0.0102,
+      "llm_learning_rate": 2.819768467841785e-06,
+      "loss": 0.778,
+      "output_text_mask_loss": 0.0488,
+      "predictor_learning_rate": 0.00028197684678417845,
+      "step": 4640
+    },
+    {
+      "epoch": 0.48,
+      "gumbel_tau": 0.3346,
+      "image_masks_loss": 0.0237,
+      "llm_learning_rate": 2.8189456478901354e-06,
+      "loss": 0.7267,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00028189456478901355,
+      "step": 4641
+    },
+    {
+      "epoch": 0.48,
+      "gumbel_tau": 0.3346,
+      "image_masks_loss": 0.0124,
+      "llm_learning_rate": 2.818122792815481e-06,
+      "loss": 0.6915,
+      "output_text_mask_loss": 0.0269,
+      "predictor_learning_rate": 0.0002818122792815481,
+      "step": 4642
+    },
+    {
+      "epoch": 0.48,
+      "gumbel_tau": 0.3345,
+      "image_masks_loss": 0.0277,
+      "llm_learning_rate": 2.817299902708436e-06,
+      "loss": 0.7013,
+      "output_text_mask_loss": 0.0124,
+      "predictor_learning_rate": 0.00028172999027084356,
+      "step": 4643
+    },
+    {
+      "epoch": 0.48,
+      "gumbel_tau": 0.3344,
+      "image_masks_loss": 0.0187,
+      "llm_learning_rate": 2.816476977659619e-06,
+      "loss": 0.7123,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.00028164769776596187,
+      "step": 4644
+    },
+    {
+      "epoch": 0.48,
+      "gumbel_tau": 0.3343,
+      "image_masks_loss": 0.0153,
+      "llm_learning_rate": 2.815654017759652e-06,
+      "loss": 0.7271,
+      "output_text_mask_loss": 0.0391,
+      "predictor_learning_rate": 0.0002815654017759652,
+      "step": 4645
+    },
+    {
+      "epoch": 0.48,
+      "gumbel_tau": 0.3343,
+      "image_masks_loss": 0.0168,
+      "llm_learning_rate": 2.814831023099161e-06,
+      "loss": 0.7905,
+      "output_text_mask_loss": 0.0032,
+      "predictor_learning_rate": 0.0002814831023099161,
+      "step": 4646
+    },
+    {
+      "epoch": 0.48,
+      "gumbel_tau": 0.3342,
+      "image_masks_loss": 0.0186,
+      "llm_learning_rate": 2.8140079937687774e-06,
+      "loss": 0.8236,
+      "output_text_mask_loss": 0.0322,
+      "predictor_learning_rate": 0.00028140079937687774,
+      "step": 4647
+    },
+    {
+      "epoch": 0.48,
+      "gumbel_tau": 0.3341,
+      "image_masks_loss": 0.017,
+      "llm_learning_rate": 2.813184929859134e-06,
+      "loss": 0.7651,
+      "output_text_mask_loss": 0.0671,
+      "predictor_learning_rate": 0.0002813184929859134,
+      "step": 4648
+    },
+    {
+      "epoch": 0.48,
+      "gumbel_tau": 0.334,
+      "image_masks_loss": 0.0271,
+      "llm_learning_rate": 2.8123618314608673e-06,
+      "loss": 0.7793,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.00028123618314608676,
+      "step": 4649
+    },
+    {
+      "epoch": 0.48,
+      "gumbel_tau": 0.3339,
+      "image_masks_loss": 0.0226,
+      "llm_learning_rate": 2.8115386986646215e-06,
+      "loss": 0.6559,
+      "output_text_mask_loss": 0.0324,
+      "predictor_learning_rate": 0.00028115386986646214,
+      "step": 4650
+    },
+    {
+      "epoch": 0.48,
+      "gumbel_tau": 0.3339,
+      "image_masks_loss": 0.0071,
+      "llm_learning_rate": 2.81071553156104e-06,
+      "loss": 0.7741,
+      "output_text_mask_loss": 0.0025,
+      "predictor_learning_rate": 0.000281071553156104,
+      "step": 4651
+    },
+    {
+      "epoch": 0.48,
+      "gumbel_tau": 0.3338,
+      "image_masks_loss": 0.0041,
+      "llm_learning_rate": 2.8098923302407716e-06,
+      "loss": 0.7991,
+      "output_text_mask_loss": 0.0633,
+      "predictor_learning_rate": 0.00028098923302407713,
+      "step": 4652
+    },
+    {
+      "epoch": 0.48,
+      "gumbel_tau": 0.3337,
+      "image_masks_loss": 0.0132,
+      "llm_learning_rate": 2.8090690947944697e-06,
+      "loss": 0.7276,
+      "output_text_mask_loss": 0.0305,
+      "predictor_learning_rate": 0.000280906909479447,
+      "step": 4653
+    },
+    {
+      "epoch": 0.48,
+      "gumbel_tau": 0.3336,
+      "image_masks_loss": 0.0208,
+      "llm_learning_rate": 2.8082458253127917e-06,
+      "loss": 0.752,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0002808245825312791,
+      "step": 4654
+    },
+    {
+      "epoch": 0.48,
+      "gumbel_tau": 0.3335,
+      "image_masks_loss": 0.0151,
+      "llm_learning_rate": 2.807422521886395e-06,
+      "loss": 0.756,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0002807422521886395,
+      "step": 4655
+    },
+    {
+      "epoch": 0.48,
+      "gumbel_tau": 0.3335,
+      "image_masks_loss": 0.0058,
+      "llm_learning_rate": 2.8065991846059466e-06,
+      "loss": 0.8496,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 0.00028065991846059463,
+      "step": 4656
+    },
+    {
+      "epoch": 0.48,
+      "gumbel_tau": 0.3334,
+      "image_masks_loss": 0.0104,
+      "llm_learning_rate": 2.805775813562113e-06,
+      "loss": 0.7367,
+      "output_text_mask_loss": 0.0085,
+      "predictor_learning_rate": 0.0002805775813562113,
+      "step": 4657
+    },
+    {
+      "epoch": 0.48,
+      "gumbel_tau": 0.3333,
+      "image_masks_loss": 0.0217,
+      "llm_learning_rate": 2.804952408845565e-06,
+      "loss": 0.8456,
+      "output_text_mask_loss": 0.0984,
+      "predictor_learning_rate": 0.00028049524088455646,
+      "step": 4658
+    },
+    {
+      "epoch": 0.48,
+      "gumbel_tau": 0.3332,
+      "image_masks_loss": 0.0103,
+      "llm_learning_rate": 2.8041289705469785e-06,
+      "loss": 0.7597,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 0.00028041289705469784,
+      "step": 4659
+    },
+    {
+      "epoch": 0.48,
+      "gumbel_tau": 0.3332,
+      "image_masks_loss": 0.0587,
+      "llm_learning_rate": 2.803305498757033e-06,
+      "loss": 0.7808,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00028033054987570326,
+      "step": 4660
+    },
+    {
+      "epoch": 0.48,
+      "gumbel_tau": 0.3331,
+      "image_masks_loss": 0.0051,
+      "llm_learning_rate": 2.802481993566411e-06,
+      "loss": 0.681,
+      "output_text_mask_loss": 0.0075,
+      "predictor_learning_rate": 0.0002802481993566411,
+      "step": 4661
+    },
+    {
+      "epoch": 0.48,
+      "gumbel_tau": 0.333,
+      "image_masks_loss": 0.0086,
+      "llm_learning_rate": 2.8016584550657966e-06,
+      "loss": 0.7473,
+      "output_text_mask_loss": 0.021,
+      "predictor_learning_rate": 0.00028016584550657965,
+      "step": 4662
+    },
+    {
+      "epoch": 0.48,
+      "gumbel_tau": 0.3329,
+      "image_masks_loss": 0.0199,
+      "llm_learning_rate": 2.800834883345882e-06,
+      "loss": 0.7394,
+      "output_text_mask_loss": 0.0843,
+      "predictor_learning_rate": 0.0002800834883345882,
+      "step": 4663
+    },
+    {
+      "epoch": 0.48,
+      "gumbel_tau": 0.3328,
+      "image_masks_loss": 0.0077,
+      "llm_learning_rate": 2.80001127849736e-06,
+      "loss": 0.7702,
+      "output_text_mask_loss": 0.0121,
+      "predictor_learning_rate": 0.00028000112784973594,
+      "step": 4664
+    },
+    {
+      "epoch": 0.48,
+      "gumbel_tau": 0.3328,
+      "image_masks_loss": 0.0433,
+      "llm_learning_rate": 2.7991876406109275e-06,
+      "loss": 0.7185,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.00027991876406109276,
+      "step": 4665
+    },
+    {
+      "epoch": 0.48,
+      "gumbel_tau": 0.3327,
+      "image_masks_loss": 0.0025,
+      "llm_learning_rate": 2.798363969777286e-06,
+      "loss": 0.7685,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.00027983639697772857,
+      "step": 4666
+    },
+    {
+      "epoch": 0.48,
+      "gumbel_tau": 0.3326,
+      "image_masks_loss": 0.0107,
+      "llm_learning_rate": 2.7975402660871394e-06,
+      "loss": 0.8352,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.0002797540266087139,
+      "step": 4667
+    },
+    {
+      "epoch": 0.48,
+      "gumbel_tau": 0.3325,
+      "image_masks_loss": 0.0106,
+      "llm_learning_rate": 2.7967165296311954e-06,
+      "loss": 0.7741,
+      "output_text_mask_loss": 0.0023,
+      "predictor_learning_rate": 0.0002796716529631195,
+      "step": 4668
+    },
+    {
+      "epoch": 0.48,
+      "gumbel_tau": 0.3324,
+      "image_masks_loss": 0.0248,
+      "llm_learning_rate": 2.7958927605001664e-06,
+      "loss": 0.7205,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0002795892760500166,
+      "step": 4669
+    },
+    {
+      "epoch": 0.48,
+      "gumbel_tau": 0.3324,
+      "image_masks_loss": 0.0081,
+      "llm_learning_rate": 2.7950689587847684e-06,
+      "loss": 0.7762,
+      "output_text_mask_loss": 0.037,
+      "predictor_learning_rate": 0.0002795068958784768,
+      "step": 4670
+    },
+    {
+      "epoch": 0.48,
+      "gumbel_tau": 0.3323,
+      "image_masks_loss": 0.0252,
+      "llm_learning_rate": 2.7942451245757186e-06,
+      "loss": 0.7024,
+      "output_text_mask_loss": 0.033,
+      "predictor_learning_rate": 0.00027942451245757183,
+      "step": 4671
+    },
+    {
+      "epoch": 0.48,
+      "gumbel_tau": 0.3322,
+      "image_masks_loss": 0.0156,
+      "llm_learning_rate": 2.7934212579637403e-06,
+      "loss": 0.7287,
+      "output_text_mask_loss": 0.0305,
+      "predictor_learning_rate": 0.00027934212579637405,
+      "step": 4672
+    },
+    {
+      "epoch": 0.48,
+      "gumbel_tau": 0.3321,
+      "image_masks_loss": 0.0093,
+      "llm_learning_rate": 2.79259735903956e-06,
+      "loss": 0.66,
+      "output_text_mask_loss": 0.0038,
+      "predictor_learning_rate": 0.00027925973590395595,
+      "step": 4673
+    },
+    {
+      "epoch": 0.48,
+      "gumbel_tau": 0.3321,
+      "image_masks_loss": 0.0148,
+      "llm_learning_rate": 2.791773427893907e-06,
+      "loss": 0.7726,
+      "output_text_mask_loss": 0.0134,
+      "predictor_learning_rate": 0.00027917734278939065,
+      "step": 4674
+    },
+    {
+      "epoch": 0.48,
+      "gumbel_tau": 0.332,
+      "image_masks_loss": 0.0074,
+      "llm_learning_rate": 2.7909494646175134e-06,
+      "loss": 0.8356,
+      "output_text_mask_loss": 0.0294,
+      "predictor_learning_rate": 0.00027909494646175133,
+      "step": 4675
+    },
+    {
+      "epoch": 0.48,
+      "gumbel_tau": 0.3319,
+      "image_masks_loss": 0.0106,
+      "llm_learning_rate": 2.790125469301117e-06,
+      "loss": 0.7516,
+      "output_text_mask_loss": 0.0175,
+      "predictor_learning_rate": 0.0002790125469301117,
+      "step": 4676
+    },
+    {
+      "epoch": 0.48,
+      "gumbel_tau": 0.3318,
+      "image_masks_loss": 0.0061,
+      "llm_learning_rate": 2.7893014420354573e-06,
+      "loss": 0.7027,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00027893014420354574,
+      "step": 4677
+    },
+    {
+      "epoch": 0.48,
+      "gumbel_tau": 0.3317,
+      "image_masks_loss": 0.0244,
+      "llm_learning_rate": 2.7884773829112787e-06,
+      "loss": 0.7257,
+      "output_text_mask_loss": 0.0645,
+      "predictor_learning_rate": 0.00027884773829112787,
+      "step": 4678
+    },
+    {
+      "epoch": 0.48,
+      "gumbel_tau": 0.3317,
+      "image_masks_loss": 0.013,
+      "llm_learning_rate": 2.7876532920193284e-06,
+      "loss": 0.709,
+      "output_text_mask_loss": 0.0004,
+      "predictor_learning_rate": 0.0002787653292019328,
+      "step": 4679
+    },
+    {
+      "epoch": 0.48,
+      "gumbel_tau": 0.3316,
+      "image_masks_loss": 0.0127,
+      "llm_learning_rate": 2.7868291694503572e-06,
+      "loss": 0.648,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.00027868291694503573,
+      "step": 4680
+    },
+    {
+      "epoch": 0.48,
+      "gumbel_tau": 0.3315,
+      "image_masks_loss": 0.0069,
+      "llm_learning_rate": 2.7860050152951183e-06,
+      "loss": 0.8044,
+      "output_text_mask_loss": 0.0244,
+      "predictor_learning_rate": 0.0002786005015295118,
+      "step": 4681
+    },
+    {
+      "epoch": 0.48,
+      "gumbel_tau": 0.3314,
+      "image_masks_loss": 0.0138,
+      "llm_learning_rate": 2.7851808296443704e-06,
+      "loss": 0.737,
+      "output_text_mask_loss": 0.0492,
+      "predictor_learning_rate": 0.00027851808296443703,
+      "step": 4682
+    },
+    {
+      "epoch": 0.48,
+      "gumbel_tau": 0.3313,
+      "image_masks_loss": 0.0121,
+      "llm_learning_rate": 2.7843566125888744e-06,
+      "loss": 0.776,
+      "output_text_mask_loss": 0.0275,
+      "predictor_learning_rate": 0.00027843566125888744,
+      "step": 4683
+    },
+    {
+      "epoch": 0.48,
+      "gumbel_tau": 0.3313,
+      "image_masks_loss": 0.0261,
+      "llm_learning_rate": 2.7835323642193944e-06,
+      "loss": 0.6959,
+      "output_text_mask_loss": NaN,
+      "predictor_learning_rate": 0.00027835323642193944,
+      "step": 4684
+    },
+    {
+      "epoch": 0.48,
+      "gumbel_tau": 0.3312,
+      "image_masks_loss": 0.0225,
+      "llm_learning_rate": 2.782708084626699e-06,
+      "loss": 0.8391,
+      "output_text_mask_loss": 0.0383,
+      "predictor_learning_rate": 0.0002782708084626699,
+      "step": 4685
+    },
+    {
+      "epoch": 0.48,
+      "gumbel_tau": 0.3311,
+      "image_masks_loss": 0.0226,
+      "llm_learning_rate": 2.7818837739015597e-06,
+      "loss": 0.7118,
+      "output_text_mask_loss": 0.0068,
+      "predictor_learning_rate": 0.000278188377390156,
+      "step": 4686
+    },
+    {
+      "epoch": 0.48,
+      "gumbel_tau": 0.331,
+      "image_masks_loss": 0.0237,
+      "llm_learning_rate": 2.7810594321347517e-06,
+      "loss": 0.6501,
+      "output_text_mask_loss": 0.0286,
+      "predictor_learning_rate": 0.00027810594321347516,
+      "step": 4687
+    },
+    {
+      "epoch": 0.48,
+      "gumbel_tau": 0.331,
+      "image_masks_loss": 0.0109,
+      "llm_learning_rate": 2.7802350594170522e-06,
+      "loss": 0.6938,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0002780235059417052,
+      "step": 4688
+    },
+    {
+      "epoch": 0.48,
+      "gumbel_tau": 0.3309,
+      "image_masks_loss": 0.0515,
+      "llm_learning_rate": 2.779410655839245e-06,
+      "loss": 0.7604,
+      "output_text_mask_loss": 0.0221,
+      "predictor_learning_rate": 0.0002779410655839245,
+      "step": 4689
+    },
+    {
+      "epoch": 0.48,
+      "gumbel_tau": 0.3308,
+      "image_masks_loss": 0.0168,
+      "llm_learning_rate": 2.7785862214921133e-06,
+      "loss": 0.7344,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0002778586221492113,
+      "step": 4690
+    },
+    {
+      "epoch": 0.48,
+      "gumbel_tau": 0.3307,
+      "image_masks_loss": 0.0087,
+      "llm_learning_rate": 2.777761756466446e-06,
+      "loss": 0.7865,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.0002777761756466446,
+      "step": 4691
+    },
+    {
+      "epoch": 0.48,
+      "gumbel_tau": 0.3306,
+      "image_masks_loss": 0.009,
+      "llm_learning_rate": 2.776937260853036e-06,
+      "loss": 0.7467,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.00027769372608530356,
+      "step": 4692
+    },
+    {
+      "epoch": 0.48,
+      "gumbel_tau": 0.3306,
+      "image_masks_loss": 0.0202,
+      "llm_learning_rate": 2.7761127347426774e-06,
+      "loss": 0.6829,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.00027761127347426774,
+      "step": 4693
+    },
+    {
+      "epoch": 0.48,
+      "gumbel_tau": 0.3305,
+      "image_masks_loss": 0.0259,
+      "llm_learning_rate": 2.7752881782261705e-06,
+      "loss": 0.7591,
+      "output_text_mask_loss": 0.0025,
+      "predictor_learning_rate": 0.000277528817822617,
+      "step": 4694
+    },
+    {
+      "epoch": 0.48,
+      "gumbel_tau": 0.3304,
+      "image_masks_loss": 0.0243,
+      "llm_learning_rate": 2.7744635913943146e-06,
+      "loss": 0.7648,
+      "output_text_mask_loss": 0.0111,
+      "predictor_learning_rate": 0.00027744635913943145,
+      "step": 4695
+    },
+    {
+      "epoch": 0.48,
+      "gumbel_tau": 0.3303,
+      "image_masks_loss": 0.0347,
+      "llm_learning_rate": 2.7736389743379176e-06,
+      "loss": 0.7317,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00027736389743379176,
+      "step": 4696
+    },
+    {
+      "epoch": 0.48,
+      "gumbel_tau": 0.3303,
+      "image_masks_loss": 0.0182,
+      "llm_learning_rate": 2.7728143271477877e-06,
+      "loss": 0.8177,
+      "output_text_mask_loss": 0.0004,
+      "predictor_learning_rate": 0.0002772814327147787,
+      "step": 4697
+    },
+    {
+      "epoch": 0.48,
+      "gumbel_tau": 0.3302,
+      "image_masks_loss": 0.0246,
+      "llm_learning_rate": 2.771989649914736e-06,
+      "loss": 0.7775,
+      "output_text_mask_loss": 0.0221,
+      "predictor_learning_rate": 0.0002771989649914736,
+      "step": 4698
+    },
+    {
+      "epoch": 0.48,
+      "gumbel_tau": 0.3301,
+      "image_masks_loss": 0.0035,
+      "llm_learning_rate": 2.7711649427295778e-06,
+      "loss": 0.6848,
+      "output_text_mask_loss": 0.0023,
+      "predictor_learning_rate": 0.0002771164942729578,
+      "step": 4699
+    },
+    {
+      "epoch": 0.48,
+      "gumbel_tau": 0.33,
+      "image_masks_loss": 0.0131,
+      "llm_learning_rate": 2.770340205683133e-06,
+      "loss": 0.6868,
+      "output_text_mask_loss": 0.0397,
+      "predictor_learning_rate": 0.0002770340205683133,
+      "step": 4700
+    },
+    {
+      "epoch": 0.48,
+      "gumbel_tau": 0.3299,
+      "image_masks_loss": 0.0102,
+      "llm_learning_rate": 2.7695154388662214e-06,
+      "loss": 0.6637,
+      "output_text_mask_loss": 0.0408,
+      "predictor_learning_rate": 0.00027695154388662215,
+      "step": 4701
+    },
+    {
+      "epoch": 0.48,
+      "gumbel_tau": 0.3299,
+      "image_masks_loss": 0.0058,
+      "llm_learning_rate": 2.76869064236967e-06,
+      "loss": 0.8241,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.000276869064236967,
+      "step": 4702
+    },
+    {
+      "epoch": 0.48,
+      "gumbel_tau": 0.3298,
+      "image_masks_loss": 0.0275,
+      "llm_learning_rate": 2.767865816284307e-06,
+      "loss": 0.7128,
+      "output_text_mask_loss": 0.0292,
+      "predictor_learning_rate": 0.0002767865816284307,
+      "step": 4703
+    },
+    {
+      "epoch": 0.48,
+      "gumbel_tau": 0.3297,
+      "image_masks_loss": 0.0136,
+      "llm_learning_rate": 2.7670409607009624e-06,
+      "loss": 0.739,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 0.0002767040960700962,
+      "step": 4704
+    },
+    {
+      "epoch": 0.48,
+      "gumbel_tau": 0.3296,
+      "image_masks_loss": 0.0095,
+      "llm_learning_rate": 2.7662160757104733e-06,
+      "loss": 0.7004,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.0002766216075710473,
+      "step": 4705
+    },
+    {
+      "epoch": 0.48,
+      "gumbel_tau": 0.3296,
+      "image_masks_loss": 0.0099,
+      "llm_learning_rate": 2.765391161403677e-06,
+      "loss": 0.796,
+      "output_text_mask_loss": 0.0023,
+      "predictor_learning_rate": 0.0002765391161403677,
+      "step": 4706
+    },
+    {
+      "epoch": 0.48,
+      "gumbel_tau": 0.3295,
+      "image_masks_loss": 0.0193,
+      "llm_learning_rate": 2.764566217871415e-06,
+      "loss": 0.808,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0002764566217871415,
+      "step": 4707
+    },
+    {
+      "epoch": 0.48,
+      "gumbel_tau": 0.3294,
+      "image_masks_loss": 0.0237,
+      "llm_learning_rate": 2.7637412452045313e-06,
+      "loss": 0.7883,
+      "output_text_mask_loss": 0.0166,
+      "predictor_learning_rate": 0.0002763741245204531,
+      "step": 4708
+    },
+    {
+      "epoch": 0.48,
+      "gumbel_tau": 0.3293,
+      "image_masks_loss": 0.0153,
+      "llm_learning_rate": 2.762916243493875e-06,
+      "loss": 0.831,
+      "output_text_mask_loss": 0.0416,
+      "predictor_learning_rate": 0.00027629162434938747,
+      "step": 4709
+    },
+    {
+      "epoch": 0.48,
+      "gumbel_tau": 0.3292,
+      "image_masks_loss": 0.0173,
+      "llm_learning_rate": 2.7620912128302957e-06,
+      "loss": 0.7623,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.00027620912128302954,
+      "step": 4710
+    },
+    {
+      "epoch": 0.48,
+      "gumbel_tau": 0.3292,
+      "image_masks_loss": 0.0104,
+      "llm_learning_rate": 2.761266153304648e-06,
+      "loss": 0.7383,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.00027612661533046477,
+      "step": 4711
+    },
+    {
+      "epoch": 0.48,
+      "gumbel_tau": 0.3291,
+      "image_masks_loss": 0.0036,
+      "llm_learning_rate": 2.7604410650077902e-06,
+      "loss": 0.7734,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.000276044106500779,
+      "step": 4712
+    },
+    {
+      "epoch": 0.48,
+      "gumbel_tau": 0.329,
+      "image_masks_loss": 0.023,
+      "llm_learning_rate": 2.7596159480305818e-06,
+      "loss": 0.8091,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.0002759615948030582,
+      "step": 4713
+    },
+    {
+      "epoch": 0.48,
+      "gumbel_tau": 0.3289,
+      "image_masks_loss": 0.0048,
+      "llm_learning_rate": 2.7587908024638865e-06,
+      "loss": 0.6507,
+      "output_text_mask_loss": 0.006,
+      "predictor_learning_rate": 0.0002758790802463886,
+      "step": 4714
+    },
+    {
+      "epoch": 0.48,
+      "gumbel_tau": 0.3289,
+      "image_masks_loss": 0.0278,
+      "llm_learning_rate": 2.757965628398572e-06,
+      "loss": 0.7207,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0002757965628398572,
+      "step": 4715
+    },
+    {
+      "epoch": 0.48,
+      "gumbel_tau": 0.3288,
+      "image_masks_loss": 0.0107,
+      "llm_learning_rate": 2.7571404259255087e-06,
+      "loss": 0.7216,
+      "output_text_mask_loss": 0.0089,
+      "predictor_learning_rate": 0.00027571404259255083,
+      "step": 4716
+    },
+    {
+      "epoch": 0.48,
+      "gumbel_tau": 0.3287,
+      "image_masks_loss": 0.0083,
+      "llm_learning_rate": 2.7563151951355684e-06,
+      "loss": 0.8974,
+      "output_text_mask_loss": 0.0277,
+      "predictor_learning_rate": 0.0002756315195135568,
+      "step": 4717
+    },
+    {
+      "epoch": 0.48,
+      "gumbel_tau": 0.3286,
+      "image_masks_loss": 0.0178,
+      "llm_learning_rate": 2.755489936119628e-06,
+      "loss": 0.8344,
+      "output_text_mask_loss": 0.0511,
+      "predictor_learning_rate": 0.0002755489936119628,
+      "step": 4718
+    },
+    {
+      "epoch": 0.48,
+      "gumbel_tau": 0.3285,
+      "image_masks_loss": 0.0138,
+      "llm_learning_rate": 2.7546646489685662e-06,
+      "loss": 0.7416,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.0002754664648968566,
+      "step": 4719
+    },
+    {
+      "epoch": 0.48,
+      "gumbel_tau": 0.3285,
+      "image_masks_loss": 0.0071,
+      "llm_learning_rate": 2.753839333773266e-06,
+      "loss": 0.6943,
+      "output_text_mask_loss": 0.0565,
+      "predictor_learning_rate": 0.0002753839333773266,
+      "step": 4720
+    },
+    {
+      "epoch": 0.48,
+      "gumbel_tau": 0.3284,
+      "image_masks_loss": 0.0145,
+      "llm_learning_rate": 2.753013990624614e-06,
+      "loss": 0.7417,
+      "output_text_mask_loss": 0.017,
+      "predictor_learning_rate": 0.0002753013990624614,
+      "step": 4721
+    },
+    {
+      "epoch": 0.48,
+      "gumbel_tau": 0.3283,
+      "image_masks_loss": 0.017,
+      "llm_learning_rate": 2.7521886196134966e-06,
+      "loss": 0.7714,
+      "output_text_mask_loss": 0.074,
+      "predictor_learning_rate": 0.00027521886196134966,
+      "step": 4722
+    },
+    {
+      "epoch": 0.48,
+      "gumbel_tau": 0.3282,
+      "image_masks_loss": 0.0072,
+      "llm_learning_rate": 2.7513632208308076e-06,
+      "loss": 0.7372,
+      "output_text_mask_loss": 0.002,
+      "predictor_learning_rate": 0.00027513632208308074,
+      "step": 4723
+    },
+    {
+      "epoch": 0.48,
+      "gumbel_tau": 0.3282,
+      "image_masks_loss": 0.014,
+      "llm_learning_rate": 2.7505377943674398e-06,
+      "loss": 0.8147,
+      "output_text_mask_loss": 0.0607,
+      "predictor_learning_rate": 0.00027505377943674397,
+      "step": 4724
+    },
+    {
+      "epoch": 0.48,
+      "gumbel_tau": 0.3281,
+      "image_masks_loss": 0.015,
+      "llm_learning_rate": 2.749712340314293e-06,
+      "loss": 0.782,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.00027497123403142927,
+      "step": 4725
+    },
+    {
+      "epoch": 0.48,
+      "gumbel_tau": 0.328,
+      "image_masks_loss": 0.0212,
+      "llm_learning_rate": 2.7488868587622674e-06,
+      "loss": 0.7584,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.00027488868587622673,
+      "step": 4726
+    },
+    {
+      "epoch": 0.48,
+      "gumbel_tau": 0.3279,
+      "image_masks_loss": 0.018,
+      "llm_learning_rate": 2.748061349802266e-06,
+      "loss": 0.7452,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.00027480613498022654,
+      "step": 4727
+    },
+    {
+      "epoch": 0.48,
+      "gumbel_tau": 0.3279,
+      "image_masks_loss": 0.0173,
+      "llm_learning_rate": 2.747235813525196e-06,
+      "loss": 0.8079,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00027472358135251956,
+      "step": 4728
+    },
+    {
+      "epoch": 0.48,
+      "gumbel_tau": 0.3278,
+      "image_masks_loss": 0.0271,
+      "llm_learning_rate": 2.7464102500219676e-06,
+      "loss": 0.8042,
+      "output_text_mask_loss": 0.0111,
+      "predictor_learning_rate": 0.0002746410250021967,
+      "step": 4729
+    },
+    {
+      "epoch": 0.48,
+      "gumbel_tau": 0.3277,
+      "image_masks_loss": 0.0203,
+      "llm_learning_rate": 2.745584659383493e-06,
+      "loss": 0.7777,
+      "output_text_mask_loss": 0.0101,
+      "predictor_learning_rate": 0.00027455846593834926,
+      "step": 4730
+    },
+    {
+      "epoch": 0.48,
+      "gumbel_tau": 0.3276,
+      "image_masks_loss": 0.0124,
+      "llm_learning_rate": 2.7447590417006893e-06,
+      "loss": 0.7651,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.0002744759041700689,
+      "step": 4731
+    },
+    {
+      "epoch": 0.48,
+      "gumbel_tau": 0.3275,
+      "image_masks_loss": 0.0161,
+      "llm_learning_rate": 2.743933397064475e-06,
+      "loss": 0.7282,
+      "output_text_mask_loss": 0.0095,
+      "predictor_learning_rate": 0.0002743933397064475,
+      "step": 4732
+    },
+    {
+      "epoch": 0.48,
+      "gumbel_tau": 0.3275,
+      "image_masks_loss": 0.0039,
+      "llm_learning_rate": 2.743107725565771e-06,
+      "loss": 0.7512,
+      "output_text_mask_loss": 0.0172,
+      "predictor_learning_rate": 0.0002743107725565771,
+      "step": 4733
+    },
+    {
+      "epoch": 0.49,
+      "gumbel_tau": 0.3274,
+      "image_masks_loss": 0.036,
+      "llm_learning_rate": 2.7422820272955032e-06,
+      "loss": 0.8009,
+      "output_text_mask_loss": 0.0065,
+      "predictor_learning_rate": 0.0002742282027295503,
+      "step": 4734
+    },
+    {
+      "epoch": 0.49,
+      "gumbel_tau": 0.3273,
+      "image_masks_loss": 0.0104,
+      "llm_learning_rate": 2.7414563023445993e-06,
+      "loss": 0.7217,
+      "output_text_mask_loss": 0.0137,
+      "predictor_learning_rate": 0.00027414563023445994,
+      "step": 4735
+    },
+    {
+      "epoch": 0.49,
+      "gumbel_tau": 0.3272,
+      "image_masks_loss": 0.0332,
+      "llm_learning_rate": 2.7406305508039886e-06,
+      "loss": 0.781,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00027406305508039885,
+      "step": 4736
+    },
+    {
+      "epoch": 0.49,
+      "gumbel_tau": 0.3272,
+      "image_masks_loss": 0.0068,
+      "llm_learning_rate": 2.7398047727646064e-06,
+      "loss": 0.768,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00027398047727646065,
+      "step": 4737
+    },
+    {
+      "epoch": 0.49,
+      "gumbel_tau": 0.3271,
+      "image_masks_loss": 0.012,
+      "llm_learning_rate": 2.7389789683173882e-06,
+      "loss": 0.6796,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.0002738978968317388,
+      "step": 4738
+    },
+    {
+      "epoch": 0.49,
+      "gumbel_tau": 0.327,
+      "image_masks_loss": 0.0153,
+      "llm_learning_rate": 2.7381531375532743e-06,
+      "loss": 0.7432,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 0.00027381531375532744,
+      "step": 4739
+    },
+    {
+      "epoch": 0.49,
+      "gumbel_tau": 0.3269,
+      "image_masks_loss": 0.0182,
+      "llm_learning_rate": 2.7373272805632056e-06,
+      "loss": 0.596,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.00027373272805632053,
+      "step": 4740
+    },
+    {
+      "epoch": 0.49,
+      "gumbel_tau": 0.3268,
+      "image_masks_loss": 0.0052,
+      "llm_learning_rate": 2.736501397438128e-06,
+      "loss": 0.7829,
+      "output_text_mask_loss": 0.0124,
+      "predictor_learning_rate": 0.00027365013974381283,
+      "step": 4741
+    },
+    {
+      "epoch": 0.49,
+      "gumbel_tau": 0.3268,
+      "image_masks_loss": 0.0166,
+      "llm_learning_rate": 2.7356754882689905e-06,
+      "loss": 0.6647,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.000273567548826899,
+      "step": 4742
+    },
+    {
+      "epoch": 0.49,
+      "gumbel_tau": 0.3267,
+      "image_masks_loss": 0.008,
+      "llm_learning_rate": 2.7348495531467427e-06,
+      "loss": 0.7371,
+      "output_text_mask_loss": 0.0054,
+      "predictor_learning_rate": 0.00027348495531467426,
+      "step": 4743
+    },
+    {
+      "epoch": 0.49,
+      "gumbel_tau": 0.3266,
+      "image_masks_loss": 0.0252,
+      "llm_learning_rate": 2.73402359216234e-06,
+      "loss": 0.8094,
+      "output_text_mask_loss": 0.0793,
+      "predictor_learning_rate": 0.00027340235921623397,
+      "step": 4744
+    },
+    {
+      "epoch": 0.49,
+      "gumbel_tau": 0.3265,
+      "image_masks_loss": 0.0122,
+      "llm_learning_rate": 2.7331976054067376e-06,
+      "loss": 0.7391,
+      "output_text_mask_loss": 0.0444,
+      "predictor_learning_rate": 0.00027331976054067374,
+      "step": 4745
+    },
+    {
+      "epoch": 0.49,
+      "gumbel_tau": 0.3265,
+      "image_masks_loss": 0.0143,
+      "llm_learning_rate": 2.7323715929708947e-06,
+      "loss": 0.7677,
+      "output_text_mask_loss": 0.0025,
+      "predictor_learning_rate": 0.0002732371592970895,
+      "step": 4746
+    },
+    {
+      "epoch": 0.49,
+      "gumbel_tau": 0.3264,
+      "image_masks_loss": 0.0077,
+      "llm_learning_rate": 2.731545554945775e-06,
+      "loss": 0.6811,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0002731545554945775,
+      "step": 4747
+    },
+    {
+      "epoch": 0.49,
+      "gumbel_tau": 0.3263,
+      "image_masks_loss": 0.0125,
+      "llm_learning_rate": 2.730719491422343e-06,
+      "loss": 0.6475,
+      "output_text_mask_loss": 0.0032,
+      "predictor_learning_rate": 0.00027307194914223426,
+      "step": 4748
+    },
+    {
+      "epoch": 0.49,
+      "gumbel_tau": 0.3262,
+      "image_masks_loss": 0.0107,
+      "llm_learning_rate": 2.729893402491567e-06,
+      "loss": 0.7486,
+      "output_text_mask_loss": 0.0124,
+      "predictor_learning_rate": 0.00027298934024915666,
+      "step": 4749
+    },
+    {
+      "epoch": 0.49,
+      "gumbel_tau": 0.3262,
+      "image_masks_loss": 0.0263,
+      "llm_learning_rate": 2.729067288244416e-06,
+      "loss": 0.7528,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 0.0002729067288244416,
+      "step": 4750
+    },
+    {
+      "epoch": 0.49,
+      "gumbel_tau": 0.3261,
+      "image_masks_loss": 0.0122,
+      "llm_learning_rate": 2.728241148771866e-06,
+      "loss": 0.6953,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 0.0002728241148771866,
+      "step": 4751
+    },
+    {
+      "epoch": 0.49,
+      "gumbel_tau": 0.326,
+      "image_masks_loss": 0.0381,
+      "llm_learning_rate": 2.7274149841648932e-06,
+      "loss": 0.8011,
+      "output_text_mask_loss": 0.0095,
+      "predictor_learning_rate": 0.0002727414984164893,
+      "step": 4752
+    },
+    {
+      "epoch": 0.49,
+      "gumbel_tau": 0.3259,
+      "image_masks_loss": 0.0234,
+      "llm_learning_rate": 2.7265887945144743e-06,
+      "loss": 0.7242,
+      "output_text_mask_loss": 0.0039,
+      "predictor_learning_rate": 0.0002726588794514474,
+      "step": 4753
+    },
+    {
+      "epoch": 0.49,
+      "gumbel_tau": 0.3258,
+      "image_masks_loss": 0.0022,
+      "llm_learning_rate": 2.725762579911594e-06,
+      "loss": 0.7772,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0002725762579911594,
+      "step": 4754
+    },
+    {
+      "epoch": 0.49,
+      "gumbel_tau": 0.3258,
+      "image_masks_loss": 0.0113,
+      "llm_learning_rate": 2.724936340447235e-06,
+      "loss": 0.7976,
+      "output_text_mask_loss": 0.0446,
+      "predictor_learning_rate": 0.00027249363404472347,
+      "step": 4755
+    },
+    {
+      "epoch": 0.49,
+      "gumbel_tau": 0.3257,
+      "image_masks_loss": 0.0126,
+      "llm_learning_rate": 2.7241100762123846e-06,
+      "loss": 0.8366,
+      "output_text_mask_loss": 0.0156,
+      "predictor_learning_rate": 0.00027241100762123846,
+      "step": 4756
+    },
+    {
+      "epoch": 0.49,
+      "gumbel_tau": 0.3256,
+      "image_masks_loss": 0.0345,
+      "llm_learning_rate": 2.723283787298034e-06,
+      "loss": 0.7398,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.0002723283787298034,
+      "step": 4757
+    },
+    {
+      "epoch": 0.49,
+      "gumbel_tau": 0.3255,
+      "image_masks_loss": 0.0116,
+      "llm_learning_rate": 2.7224574737951755e-06,
+      "loss": 0.729,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0002722457473795175,
+      "step": 4758
+    },
+    {
+      "epoch": 0.49,
+      "gumbel_tau": 0.3255,
+      "image_masks_loss": 0.0129,
+      "llm_learning_rate": 2.721631135794804e-06,
+      "loss": 0.8005,
+      "output_text_mask_loss": 0.0181,
+      "predictor_learning_rate": 0.00027216311357948036,
+      "step": 4759
+    },
+    {
+      "epoch": 0.49,
+      "gumbel_tau": 0.3254,
+      "image_masks_loss": 0.0067,
+      "llm_learning_rate": 2.7208047733879183e-06,
+      "loss": 0.7033,
+      "output_text_mask_loss": 0.0252,
+      "predictor_learning_rate": 0.0002720804773387918,
+      "step": 4760
+    },
+    {
+      "epoch": 0.49,
+      "gumbel_tau": 0.3253,
+      "image_masks_loss": 0.0122,
+      "llm_learning_rate": 2.7199783866655192e-06,
+      "loss": 0.7272,
+      "output_text_mask_loss": 0.0353,
+      "predictor_learning_rate": 0.0002719978386665519,
+      "step": 4761
+    },
+    {
+      "epoch": 0.49,
+      "gumbel_tau": 0.3252,
+      "image_masks_loss": 0.0257,
+      "llm_learning_rate": 2.71915197571861e-06,
+      "loss": 0.7759,
+      "output_text_mask_loss": 0.0086,
+      "predictor_learning_rate": 0.000271915197571861,
+      "step": 4762
+    },
+    {
+      "epoch": 0.49,
+      "gumbel_tau": 0.3252,
+      "image_masks_loss": 0.014,
+      "llm_learning_rate": 2.7183255406381976e-06,
+      "loss": 0.6961,
+      "output_text_mask_loss": 0.0496,
+      "predictor_learning_rate": 0.00027183255406381974,
+      "step": 4763
+    },
+    {
+      "epoch": 0.49,
+      "gumbel_tau": 0.3251,
+      "image_masks_loss": 0.0089,
+      "llm_learning_rate": 2.71749908151529e-06,
+      "loss": 0.6817,
+      "output_text_mask_loss": 0.0107,
+      "predictor_learning_rate": 0.000271749908151529,
+      "step": 4764
+    },
+    {
+      "epoch": 0.49,
+      "gumbel_tau": 0.325,
+      "image_masks_loss": 0.0153,
+      "llm_learning_rate": 2.7166725984408994e-06,
+      "loss": 0.7634,
+      "output_text_mask_loss": 0.0095,
+      "predictor_learning_rate": 0.0002716672598440899,
+      "step": 4765
+    },
+    {
+      "epoch": 0.49,
+      "gumbel_tau": 0.3249,
+      "image_masks_loss": 0.018,
+      "llm_learning_rate": 2.715846091506039e-06,
+      "loss": 0.7247,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.0002715846091506039,
+      "step": 4766
+    },
+    {
+      "epoch": 0.49,
+      "gumbel_tau": 0.3248,
+      "image_masks_loss": 0.01,
+      "llm_learning_rate": 2.7150195608017265e-06,
+      "loss": 0.7584,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00027150195608017266,
+      "step": 4767
+    },
+    {
+      "epoch": 0.49,
+      "gumbel_tau": 0.3248,
+      "image_masks_loss": 0.0073,
+      "llm_learning_rate": 2.7141930064189815e-06,
+      "loss": 0.7463,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.0002714193006418981,
+      "step": 4768
+    },
+    {
+      "epoch": 0.49,
+      "gumbel_tau": 0.3247,
+      "image_masks_loss": 0.0199,
+      "llm_learning_rate": 2.713366428448825e-06,
+      "loss": 0.7252,
+      "output_text_mask_loss": 0.0143,
+      "predictor_learning_rate": 0.0002713366428448825,
+      "step": 4769
+    },
+    {
+      "epoch": 0.49,
+      "gumbel_tau": 0.3246,
+      "image_masks_loss": 0.0275,
+      "llm_learning_rate": 2.712539826982282e-06,
+      "loss": 0.7975,
+      "output_text_mask_loss": 0.0098,
+      "predictor_learning_rate": 0.0002712539826982282,
+      "step": 4770
+    },
+    {
+      "epoch": 0.49,
+      "gumbel_tau": 0.3245,
+      "image_masks_loss": 0.0135,
+      "llm_learning_rate": 2.7117132021103803e-06,
+      "loss": 0.7559,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.00027117132021103803,
+      "step": 4771
+    },
+    {
+      "epoch": 0.49,
+      "gumbel_tau": 0.3245,
+      "image_masks_loss": 0.0109,
+      "llm_learning_rate": 2.710886553924149e-06,
+      "loss": 0.7782,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0002710886553924149,
+      "step": 4772
+    },
+    {
+      "epoch": 0.49,
+      "gumbel_tau": 0.3244,
+      "image_masks_loss": 0.0146,
+      "llm_learning_rate": 2.710059882514621e-06,
+      "loss": 0.7217,
+      "output_text_mask_loss": 0.0116,
+      "predictor_learning_rate": 0.00027100598825146205,
+      "step": 4773
+    },
+    {
+      "epoch": 0.49,
+      "gumbel_tau": 0.3243,
+      "image_masks_loss": 0.0271,
+      "llm_learning_rate": 2.7092331879728305e-06,
+      "loss": 0.8022,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.00027092331879728307,
+      "step": 4774
+    },
+    {
+      "epoch": 0.49,
+      "gumbel_tau": 0.3242,
+      "image_masks_loss": 0.0034,
+      "llm_learning_rate": 2.708406470389815e-06,
+      "loss": 0.7449,
+      "output_text_mask_loss": 0.0039,
+      "predictor_learning_rate": 0.0002708406470389815,
+      "step": 4775
+    },
+    {
+      "epoch": 0.49,
+      "gumbel_tau": 0.3242,
+      "image_masks_loss": 0.0108,
+      "llm_learning_rate": 2.707579729856615e-06,
+      "loss": 0.7291,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0002707579729856615,
+      "step": 4776
+    },
+    {
+      "epoch": 0.49,
+      "gumbel_tau": 0.3241,
+      "image_masks_loss": 0.0296,
+      "llm_learning_rate": 2.7067529664642735e-06,
+      "loss": 0.7165,
+      "output_text_mask_loss": 0.0032,
+      "predictor_learning_rate": 0.0002706752966464273,
+      "step": 4777
+    },
+    {
+      "epoch": 0.49,
+      "gumbel_tau": 0.324,
+      "image_masks_loss": 0.0095,
+      "llm_learning_rate": 2.705926180303834e-06,
+      "loss": 0.7262,
+      "output_text_mask_loss": 0.0078,
+      "predictor_learning_rate": 0.0002705926180303834,
+      "step": 4778
+    },
+    {
+      "epoch": 0.49,
+      "gumbel_tau": 0.3239,
+      "image_masks_loss": 0.0043,
+      "llm_learning_rate": 2.7050993714663447e-06,
+      "loss": 0.691,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0002705099371466345,
+      "step": 4779
+    },
+    {
+      "epoch": 0.49,
+      "gumbel_tau": 0.3239,
+      "image_masks_loss": 0.0183,
+      "llm_learning_rate": 2.704272540042856e-06,
+      "loss": 0.7468,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.0002704272540042856,
+      "step": 4780
+    },
+    {
+      "epoch": 0.49,
+      "gumbel_tau": 0.3238,
+      "image_masks_loss": 0.0055,
+      "llm_learning_rate": 2.703445686124421e-06,
+      "loss": 0.7352,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00027034456861244207,
+      "step": 4781
+    },
+    {
+      "epoch": 0.49,
+      "gumbel_tau": 0.3237,
+      "image_masks_loss": 0.0075,
+      "llm_learning_rate": 2.702618809802093e-06,
+      "loss": 0.8826,
+      "output_text_mask_loss": 0.1022,
+      "predictor_learning_rate": 0.00027026188098020924,
+      "step": 4782
+    },
+    {
+      "epoch": 0.49,
+      "gumbel_tau": 0.3236,
+      "image_masks_loss": 0.0062,
+      "llm_learning_rate": 2.7017919111669304e-06,
+      "loss": 0.7913,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.00027017919111669306,
+      "step": 4783
+    },
+    {
+      "epoch": 0.49,
+      "gumbel_tau": 0.3235,
+      "image_masks_loss": 0.0145,
+      "llm_learning_rate": 2.7009649903099936e-06,
+      "loss": 0.7486,
+      "output_text_mask_loss": 0.001,
+      "predictor_learning_rate": 0.00027009649903099935,
+      "step": 4784
+    },
+    {
+      "epoch": 0.49,
+      "gumbel_tau": 0.3235,
+      "image_masks_loss": 0.0145,
+      "llm_learning_rate": 2.7001380473223437e-06,
+      "loss": 0.787,
+      "output_text_mask_loss": 0.0132,
+      "predictor_learning_rate": 0.00027001380473223435,
+      "step": 4785
+    },
+    {
+      "epoch": 0.49,
+      "gumbel_tau": 0.3234,
+      "image_masks_loss": 0.0222,
+      "llm_learning_rate": 2.6993110822950465e-06,
+      "loss": 0.7158,
+      "output_text_mask_loss": 0.0322,
+      "predictor_learning_rate": 0.0002699311082295046,
+      "step": 4786
+    },
+    {
+      "epoch": 0.49,
+      "gumbel_tau": 0.3233,
+      "image_masks_loss": 0.0135,
+      "llm_learning_rate": 2.6984840953191693e-06,
+      "loss": 0.7199,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.0002698484095319169,
+      "step": 4787
+    },
+    {
+      "epoch": 0.49,
+      "gumbel_tau": 0.3232,
+      "image_masks_loss": 0.0069,
+      "llm_learning_rate": 2.6976570864857814e-06,
+      "loss": 0.7422,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.0002697657086485781,
+      "step": 4788
+    },
+    {
+      "epoch": 0.49,
+      "gumbel_tau": 0.3232,
+      "image_masks_loss": 0.0152,
+      "llm_learning_rate": 2.6968300558859537e-06,
+      "loss": 0.6922,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 0.0002696830055885954,
+      "step": 4789
+    },
+    {
+      "epoch": 0.49,
+      "gumbel_tau": 0.3231,
+      "image_masks_loss": 0.0101,
+      "llm_learning_rate": 2.6960030036107634e-06,
+      "loss": 0.6855,
+      "output_text_mask_loss": 0.0075,
+      "predictor_learning_rate": 0.0002696003003610763,
+      "step": 4790
+    },
+    {
+      "epoch": 0.49,
+      "gumbel_tau": 0.323,
+      "image_masks_loss": 0.01,
+      "llm_learning_rate": 2.6951759297512855e-06,
+      "loss": 0.7056,
+      "output_text_mask_loss": 0.0278,
+      "predictor_learning_rate": 0.0002695175929751285,
+      "step": 4791
+    },
+    {
+      "epoch": 0.49,
+      "gumbel_tau": 0.3229,
+      "image_masks_loss": 0.0168,
+      "llm_learning_rate": 2.694348834398599e-06,
+      "loss": 0.644,
+      "output_text_mask_loss": 0.0055,
+      "predictor_learning_rate": 0.0002694348834398599,
+      "step": 4792
+    },
+    {
+      "epoch": 0.49,
+      "gumbel_tau": 0.3229,
+      "image_masks_loss": 0.019,
+      "llm_learning_rate": 2.693521717643786e-06,
+      "loss": 0.6462,
+      "output_text_mask_loss": 0.0004,
+      "predictor_learning_rate": 0.0002693521717643786,
+      "step": 4793
+    },
+    {
+      "epoch": 0.49,
+      "gumbel_tau": 0.3228,
+      "image_masks_loss": 0.0145,
+      "llm_learning_rate": 2.692694579577931e-06,
+      "loss": 0.7411,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.00026926945795779304,
+      "step": 4794
+    },
+    {
+      "epoch": 0.49,
+      "gumbel_tau": 0.3227,
+      "image_masks_loss": 0.0176,
+      "llm_learning_rate": 2.6918674202921184e-06,
+      "loss": 0.6891,
+      "output_text_mask_loss": 0.0618,
+      "predictor_learning_rate": 0.0002691867420292118,
+      "step": 4795
+    },
+    {
+      "epoch": 0.49,
+      "gumbel_tau": 0.3226,
+      "image_masks_loss": 0.0095,
+      "llm_learning_rate": 2.6910402398774386e-06,
+      "loss": 0.85,
+      "output_text_mask_loss": 0.0149,
+      "predictor_learning_rate": 0.00026910402398774386,
+      "step": 4796
+    },
+    {
+      "epoch": 0.49,
+      "gumbel_tau": 0.3226,
+      "image_masks_loss": 0.005,
+      "llm_learning_rate": 2.690213038424983e-06,
+      "loss": 0.7755,
+      "output_text_mask_loss": 0.0393,
+      "predictor_learning_rate": 0.0002690213038424983,
+      "step": 4797
+    },
+    {
+      "epoch": 0.49,
+      "gumbel_tau": 0.3225,
+      "image_masks_loss": 0.0178,
+      "llm_learning_rate": 2.689385816025843e-06,
+      "loss": 0.7313,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.00026893858160258427,
+      "step": 4798
+    },
+    {
+      "epoch": 0.49,
+      "gumbel_tau": 0.3224,
+      "image_masks_loss": 0.0074,
+      "llm_learning_rate": 2.688558572771115e-06,
+      "loss": 0.7622,
+      "output_text_mask_loss": 0.0044,
+      "predictor_learning_rate": 0.00026885585727711147,
+      "step": 4799
+    },
+    {
+      "epoch": 0.49,
+      "gumbel_tau": 0.3223,
+      "image_masks_loss": 0.0105,
+      "llm_learning_rate": 2.687731308751898e-06,
+      "loss": 0.6936,
+      "output_text_mask_loss": 0.0041,
+      "predictor_learning_rate": 0.0002687731308751898,
+      "step": 4800
+    },
+    {
+      "epoch": 0.49,
+      "gumbel_tau": 0.3223,
+      "image_masks_loss": 0.0177,
+      "llm_learning_rate": 2.6869040240592904e-06,
+      "loss": 0.7323,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.00026869040240592903,
+      "step": 4801
+    },
+    {
+      "epoch": 0.49,
+      "gumbel_tau": 0.3222,
+      "image_masks_loss": 0.0066,
+      "llm_learning_rate": 2.686076718784395e-06,
+      "loss": 0.8062,
+      "output_text_mask_loss": 0.042,
+      "predictor_learning_rate": 0.00026860767187843947,
+      "step": 4802
+    },
+    {
+      "epoch": 0.49,
+      "gumbel_tau": 0.3221,
+      "image_masks_loss": 0.0228,
+      "llm_learning_rate": 2.6852493930183176e-06,
+      "loss": 0.7459,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00026852493930183176,
+      "step": 4803
+    },
+    {
+      "epoch": 0.49,
+      "gumbel_tau": 0.322,
+      "image_masks_loss": 0.0311,
+      "llm_learning_rate": 2.684422046852164e-06,
+      "loss": 0.7265,
+      "output_text_mask_loss": 0.0124,
+      "predictor_learning_rate": 0.0002684422046852164,
+      "step": 4804
+    },
+    {
+      "epoch": 0.49,
+      "gumbel_tau": 0.3219,
+      "image_masks_loss": 0.0099,
+      "llm_learning_rate": 2.6835946803770434e-06,
+      "loss": 0.771,
+      "output_text_mask_loss": 0.0179,
+      "predictor_learning_rate": 0.0002683594680377043,
+      "step": 4805
+    },
+    {
+      "epoch": 0.49,
+      "gumbel_tau": 0.3219,
+      "image_masks_loss": 0.012,
+      "llm_learning_rate": 2.6827672936840687e-06,
+      "loss": 0.8172,
+      "output_text_mask_loss": 0.0252,
+      "predictor_learning_rate": 0.00026827672936840683,
+      "step": 4806
+    },
+    {
+      "epoch": 0.49,
+      "gumbel_tau": 0.3218,
+      "image_masks_loss": 0.0248,
+      "llm_learning_rate": 2.681939886864352e-06,
+      "loss": 0.8018,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0002681939886864352,
+      "step": 4807
+    },
+    {
+      "epoch": 0.49,
+      "gumbel_tau": 0.3217,
+      "image_masks_loss": 0.0077,
+      "llm_learning_rate": 2.6811124600090094e-06,
+      "loss": 0.7869,
+      "output_text_mask_loss": 0.0572,
+      "predictor_learning_rate": 0.00026811124600090093,
+      "step": 4808
+    },
+    {
+      "epoch": 0.49,
+      "gumbel_tau": 0.3216,
+      "image_masks_loss": 0.0301,
+      "llm_learning_rate": 2.68028501320916e-06,
+      "loss": 0.7647,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 0.000268028501320916,
+      "step": 4809
+    },
+    {
+      "epoch": 0.49,
+      "gumbel_tau": 0.3216,
+      "image_masks_loss": 0.016,
+      "llm_learning_rate": 2.679457546555923e-06,
+      "loss": 0.7727,
+      "output_text_mask_loss": 0.0202,
+      "predictor_learning_rate": 0.0002679457546555923,
+      "step": 4810
+    },
+    {
+      "epoch": 0.49,
+      "gumbel_tau": 0.3215,
+      "image_masks_loss": 0.016,
+      "llm_learning_rate": 2.678630060140421e-06,
+      "loss": 0.7102,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00026786300601404204,
+      "step": 4811
+    },
+    {
+      "epoch": 0.49,
+      "gumbel_tau": 0.3214,
+      "image_masks_loss": 0.0113,
+      "llm_learning_rate": 2.6778025540537787e-06,
+      "loss": 0.7179,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00026778025540537785,
+      "step": 4812
+    },
+    {
+      "epoch": 0.49,
+      "gumbel_tau": 0.3213,
+      "image_masks_loss": 0.0179,
+      "llm_learning_rate": 2.6769750283871237e-06,
+      "loss": 0.7715,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.00026769750283871235,
+      "step": 4813
+    },
+    {
+      "epoch": 0.49,
+      "gumbel_tau": 0.3213,
+      "image_masks_loss": 0.0072,
+      "llm_learning_rate": 2.6761474832315844e-06,
+      "loss": 0.7512,
+      "output_text_mask_loss": 0.0322,
+      "predictor_learning_rate": 0.00026761474832315843,
+      "step": 4814
+    },
+    {
+      "epoch": 0.49,
+      "gumbel_tau": 0.3212,
+      "image_masks_loss": 0.0067,
+      "llm_learning_rate": 2.6753199186782914e-06,
+      "loss": 0.7724,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 0.00026753199186782913,
+      "step": 4815
+    },
+    {
+      "epoch": 0.49,
+      "gumbel_tau": 0.3211,
+      "image_masks_loss": 0.0239,
+      "llm_learning_rate": 2.6744923348183786e-06,
+      "loss": 0.745,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0002674492334818379,
+      "step": 4816
+    },
+    {
+      "epoch": 0.49,
+      "gumbel_tau": 0.321,
+      "image_masks_loss": 0.0057,
+      "llm_learning_rate": 2.6736647317429813e-06,
+      "loss": 0.7685,
+      "output_text_mask_loss": 0.0012,
+      "predictor_learning_rate": 0.00026736647317429814,
+      "step": 4817
+    },
+    {
+      "epoch": 0.49,
+      "gumbel_tau": 0.321,
+      "image_masks_loss": 0.0128,
+      "llm_learning_rate": 2.672837109543237e-06,
+      "loss": 0.762,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.0002672837109543237,
+      "step": 4818
+    },
+    {
+      "epoch": 0.49,
+      "gumbel_tau": 0.3209,
+      "image_masks_loss": 0.0108,
+      "llm_learning_rate": 2.6720094683102856e-06,
+      "loss": 0.6896,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0002672009468310285,
+      "step": 4819
+    },
+    {
+      "epoch": 0.49,
+      "gumbel_tau": 0.3208,
+      "image_masks_loss": 0.0065,
+      "llm_learning_rate": 2.6711818081352682e-06,
+      "loss": 0.7794,
+      "output_text_mask_loss": 0.02,
+      "predictor_learning_rate": 0.0002671181808135268,
+      "step": 4820
+    },
+    {
+      "epoch": 0.49,
+      "gumbel_tau": 0.3207,
+      "image_masks_loss": 0.0132,
+      "llm_learning_rate": 2.670354129109328e-06,
+      "loss": 0.8117,
+      "output_text_mask_loss": 0.0139,
+      "predictor_learning_rate": 0.0002670354129109328,
+      "step": 4821
+    },
+    {
+      "epoch": 0.49,
+      "gumbel_tau": 0.3207,
+      "image_masks_loss": 0.0128,
+      "llm_learning_rate": 2.6695264313236126e-06,
+      "loss": 0.8191,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.0002669526431323612,
+      "step": 4822
+    },
+    {
+      "epoch": 0.49,
+      "gumbel_tau": 0.3206,
+      "image_masks_loss": 0.0102,
+      "llm_learning_rate": 2.668698714869268e-06,
+      "loss": 0.7693,
+      "output_text_mask_loss": 0.0058,
+      "predictor_learning_rate": 0.0002668698714869268,
+      "step": 4823
+    },
+    {
+      "epoch": 0.49,
+      "gumbel_tau": 0.3205,
+      "image_masks_loss": 0.0134,
+      "llm_learning_rate": 2.667870979837446e-06,
+      "loss": 0.6813,
+      "output_text_mask_loss": 0.006,
+      "predictor_learning_rate": 0.0002667870979837446,
+      "step": 4824
+    },
+    {
+      "epoch": 0.49,
+      "gumbel_tau": 0.3204,
+      "image_masks_loss": 0.0094,
+      "llm_learning_rate": 2.6670432263192975e-06,
+      "loss": 0.7893,
+      "output_text_mask_loss": 0.0109,
+      "predictor_learning_rate": 0.00026670432263192973,
+      "step": 4825
+    },
+    {
+      "epoch": 0.49,
+      "gumbel_tau": 0.3204,
+      "image_masks_loss": 0.0164,
+      "llm_learning_rate": 2.666215454405977e-06,
+      "loss": 0.6863,
+      "output_text_mask_loss": 0.0185,
+      "predictor_learning_rate": 0.00026662154544059767,
+      "step": 4826
+    },
+    {
+      "epoch": 0.49,
+      "gumbel_tau": 0.3203,
+      "image_masks_loss": 0.0123,
+      "llm_learning_rate": 2.6653876641886407e-06,
+      "loss": 0.8567,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 0.00026653876641886406,
+      "step": 4827
+    },
+    {
+      "epoch": 0.49,
+      "gumbel_tau": 0.3202,
+      "image_masks_loss": 0.0311,
+      "llm_learning_rate": 2.664559855758446e-06,
+      "loss": 0.7639,
+      "output_text_mask_loss": 0.0221,
+      "predictor_learning_rate": 0.0002664559855758446,
+      "step": 4828
+    },
+    {
+      "epoch": 0.49,
+      "gumbel_tau": 0.3201,
+      "image_masks_loss": 0.0062,
+      "llm_learning_rate": 2.6637320292065537e-06,
+      "loss": 0.706,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 0.00026637320292065536,
+      "step": 4829
+    },
+    {
+      "epoch": 0.49,
+      "gumbel_tau": 0.3201,
+      "image_masks_loss": 0.0292,
+      "llm_learning_rate": 2.6629041846241255e-06,
+      "loss": 0.7814,
+      "output_text_mask_loss": 0.0671,
+      "predictor_learning_rate": 0.00026629041846241253,
+      "step": 4830
+    },
+    {
+      "epoch": 0.49,
+      "gumbel_tau": 0.32,
+      "image_masks_loss": 0.0038,
+      "llm_learning_rate": 2.6620763221023247e-06,
+      "loss": 0.7437,
+      "output_text_mask_loss": 0.0416,
+      "predictor_learning_rate": 0.0002662076322102325,
+      "step": 4831
+    },
+    {
+      "epoch": 0.5,
+      "gumbel_tau": 0.3199,
+      "image_masks_loss": 0.0151,
+      "llm_learning_rate": 2.6612484417323196e-06,
+      "loss": 0.7663,
+      "output_text_mask_loss": 0.0389,
+      "predictor_learning_rate": 0.000266124844173232,
+      "step": 4832
+    },
+    {
+      "epoch": 0.5,
+      "gumbel_tau": 0.3198,
+      "image_masks_loss": 0.0104,
+      "llm_learning_rate": 2.6604205436052767e-06,
+      "loss": 0.7476,
+      "output_text_mask_loss": 0.002,
+      "predictor_learning_rate": 0.00026604205436052766,
+      "step": 4833
+    },
+    {
+      "epoch": 0.5,
+      "gumbel_tau": 0.3198,
+      "image_masks_loss": 0.0084,
+      "llm_learning_rate": 2.6595926278123655e-06,
+      "loss": 0.8104,
+      "output_text_mask_loss": 0.0229,
+      "predictor_learning_rate": 0.00026595926278123657,
+      "step": 4834
+    },
+    {
+      "epoch": 0.5,
+      "gumbel_tau": 0.3197,
+      "image_masks_loss": 0.0248,
+      "llm_learning_rate": 2.658764694444759e-06,
+      "loss": 0.809,
+      "output_text_mask_loss": 0.0038,
+      "predictor_learning_rate": 0.00026587646944447593,
+      "step": 4835
+    },
+    {
+      "epoch": 0.5,
+      "gumbel_tau": 0.3196,
+      "image_masks_loss": 0.0036,
+      "llm_learning_rate": 2.6579367435936316e-06,
+      "loss": 0.7057,
+      "output_text_mask_loss": 0.0729,
+      "predictor_learning_rate": 0.0002657936743593631,
+      "step": 4836
+    },
+    {
+      "epoch": 0.5,
+      "gumbel_tau": 0.3195,
+      "image_masks_loss": 0.0053,
+      "llm_learning_rate": 2.657108775350157e-06,
+      "loss": 0.8102,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.00026571087753501574,
+      "step": 4837
+    },
+    {
+      "epoch": 0.5,
+      "gumbel_tau": 0.3195,
+      "image_masks_loss": 0.0069,
+      "llm_learning_rate": 2.6562807898055147e-06,
+      "loss": 0.6871,
+      "output_text_mask_loss": 0.0174,
+      "predictor_learning_rate": 0.0002656280789805515,
+      "step": 4838
+    },
+    {
+      "epoch": 0.5,
+      "gumbel_tau": 0.3194,
+      "image_masks_loss": 0.0157,
+      "llm_learning_rate": 2.655452787050884e-06,
+      "loss": 0.7287,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.00026554527870508837,
+      "step": 4839
+    },
+    {
+      "epoch": 0.5,
+      "gumbel_tau": 0.3193,
+      "image_masks_loss": 0.0173,
+      "llm_learning_rate": 2.6546247671774455e-06,
+      "loss": 0.7198,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00026546247671774454,
+      "step": 4840
+    },
+    {
+      "epoch": 0.5,
+      "gumbel_tau": 0.3192,
+      "image_masks_loss": 0.0202,
+      "llm_learning_rate": 2.6537967302763836e-06,
+      "loss": 0.7193,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.00026537967302763834,
+      "step": 4841
+    },
+    {
+      "epoch": 0.5,
+      "gumbel_tau": 0.3192,
+      "image_masks_loss": 0.0175,
+      "llm_learning_rate": 2.652968676438883e-06,
+      "loss": 0.8208,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.00026529686764388826,
+      "step": 4842
+    },
+    {
+      "epoch": 0.5,
+      "gumbel_tau": 0.3191,
+      "image_masks_loss": 0.0026,
+      "llm_learning_rate": 2.6521406057561313e-06,
+      "loss": 0.7389,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 0.0002652140605756131,
+      "step": 4843
+    },
+    {
+      "epoch": 0.5,
+      "gumbel_tau": 0.319,
+      "image_masks_loss": 0.0193,
+      "llm_learning_rate": 2.651312518319317e-06,
+      "loss": 0.7683,
+      "output_text_mask_loss": 0.0078,
+      "predictor_learning_rate": 0.00026513125183193166,
+      "step": 4844
+    },
+    {
+      "epoch": 0.5,
+      "gumbel_tau": 0.3189,
+      "image_masks_loss": 0.0114,
+      "llm_learning_rate": 2.650484414219631e-06,
+      "loss": 0.7467,
+      "output_text_mask_loss": 0.0016,
+      "predictor_learning_rate": 0.00026504844142196304,
+      "step": 4845
+    },
+    {
+      "epoch": 0.5,
+      "gumbel_tau": 0.3188,
+      "image_masks_loss": 0.0214,
+      "llm_learning_rate": 2.6496562935482666e-06,
+      "loss": 0.737,
+      "output_text_mask_loss": 0.0039,
+      "predictor_learning_rate": 0.00026496562935482664,
+      "step": 4846
+    },
+    {
+      "epoch": 0.5,
+      "gumbel_tau": 0.3188,
+      "image_masks_loss": 0.0278,
+      "llm_learning_rate": 2.6488281563964174e-06,
+      "loss": 0.7458,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 0.00026488281563964174,
+      "step": 4847
+    },
+    {
+      "epoch": 0.5,
+      "gumbel_tau": 0.3187,
+      "image_masks_loss": 0.0092,
+      "llm_learning_rate": 2.6480000028552798e-06,
+      "loss": 0.7985,
+      "output_text_mask_loss": 0.0055,
+      "predictor_learning_rate": 0.000264800000285528,
+      "step": 4848
+    },
+    {
+      "epoch": 0.5,
+      "gumbel_tau": 0.3186,
+      "image_masks_loss": 0.0117,
+      "llm_learning_rate": 2.647171833016053e-06,
+      "loss": 0.7701,
+      "output_text_mask_loss": 0.0182,
+      "predictor_learning_rate": 0.00026471718330160527,
+      "step": 4849
+    },
+    {
+      "epoch": 0.5,
+      "gumbel_tau": 0.3185,
+      "image_masks_loss": 0.0176,
+      "llm_learning_rate": 2.6463436469699354e-06,
+      "loss": 0.7609,
+      "output_text_mask_loss": 0.0303,
+      "predictor_learning_rate": 0.00026463436469699355,
+      "step": 4850
+    },
+    {
+      "epoch": 0.5,
+      "gumbel_tau": 0.3185,
+      "image_masks_loss": 0.0107,
+      "llm_learning_rate": 2.64551544480813e-06,
+      "loss": 0.7213,
+      "output_text_mask_loss": 0.036,
+      "predictor_learning_rate": 0.000264551544480813,
+      "step": 4851
+    },
+    {
+      "epoch": 0.5,
+      "gumbel_tau": 0.3184,
+      "image_masks_loss": 0.0103,
+      "llm_learning_rate": 2.6446872266218394e-06,
+      "loss": 0.7757,
+      "output_text_mask_loss": 0.0401,
+      "predictor_learning_rate": 0.0002644687226621839,
+      "step": 4852
+    },
+    {
+      "epoch": 0.5,
+      "gumbel_tau": 0.3183,
+      "image_masks_loss": 0.0059,
+      "llm_learning_rate": 2.6438589925022695e-06,
+      "loss": 0.6357,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0002643858992502269,
+      "step": 4853
+    },
+    {
+      "epoch": 0.5,
+      "gumbel_tau": 0.3182,
+      "image_masks_loss": 0.0168,
+      "llm_learning_rate": 2.643030742540626e-06,
+      "loss": 0.6698,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 0.00026430307425406255,
+      "step": 4854
+    },
+    {
+      "epoch": 0.5,
+      "gumbel_tau": 0.3182,
+      "image_masks_loss": 0.0305,
+      "llm_learning_rate": 2.6422024768281195e-06,
+      "loss": 0.7835,
+      "output_text_mask_loss": 0.0065,
+      "predictor_learning_rate": 0.00026422024768281194,
+      "step": 4855
+    },
+    {
+      "epoch": 0.5,
+      "gumbel_tau": 0.3181,
+      "image_masks_loss": 0.0393,
+      "llm_learning_rate": 2.641374195455959e-06,
+      "loss": 0.7614,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.0002641374195455959,
+      "step": 4856
+    },
+    {
+      "epoch": 0.5,
+      "gumbel_tau": 0.318,
+      "image_masks_loss": 0.005,
+      "llm_learning_rate": 2.6405458985153575e-06,
+      "loss": 0.6973,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.00026405458985153573,
+      "step": 4857
+    },
+    {
+      "epoch": 0.5,
+      "gumbel_tau": 0.3179,
+      "image_masks_loss": 0.0227,
+      "llm_learning_rate": 2.639717586097529e-06,
+      "loss": 0.6982,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.00026397175860975287,
+      "step": 4858
+    },
+    {
+      "epoch": 0.5,
+      "gumbel_tau": 0.3179,
+      "image_masks_loss": 0.0085,
+      "llm_learning_rate": 2.638889258293688e-06,
+      "loss": 0.7897,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.0002638889258293688,
+      "step": 4859
+    },
+    {
+      "epoch": 0.5,
+      "gumbel_tau": 0.3178,
+      "image_masks_loss": 0.013,
+      "llm_learning_rate": 2.638060915195052e-06,
+      "loss": 0.7427,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0002638060915195052,
+      "step": 4860
+    },
+    {
+      "epoch": 0.5,
+      "gumbel_tau": 0.3177,
+      "image_masks_loss": 0.014,
+      "llm_learning_rate": 2.6372325568928416e-06,
+      "loss": 0.776,
+      "output_text_mask_loss": 0.0246,
+      "predictor_learning_rate": 0.00026372325568928414,
+      "step": 4861
+    },
+    {
+      "epoch": 0.5,
+      "gumbel_tau": 0.3176,
+      "image_masks_loss": 0.0244,
+      "llm_learning_rate": 2.636404183478276e-06,
+      "loss": 0.8821,
+      "output_text_mask_loss": 0.0107,
+      "predictor_learning_rate": 0.0002636404183478276,
+      "step": 4862
+    },
+    {
+      "epoch": 0.5,
+      "gumbel_tau": 0.3176,
+      "image_masks_loss": 0.0104,
+      "llm_learning_rate": 2.6355757950425776e-06,
+      "loss": 0.7024,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.0002635575795042577,
+      "step": 4863
+    },
+    {
+      "epoch": 0.5,
+      "gumbel_tau": 0.3175,
+      "image_masks_loss": 0.0077,
+      "llm_learning_rate": 2.634747391676971e-06,
+      "loss": 0.7741,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.0002634747391676971,
+      "step": 4864
+    },
+    {
+      "epoch": 0.5,
+      "gumbel_tau": 0.3174,
+      "image_masks_loss": 0.0072,
+      "llm_learning_rate": 2.633918973472682e-06,
+      "loss": 0.7759,
+      "output_text_mask_loss": 0.1213,
+      "predictor_learning_rate": 0.0002633918973472682,
+      "step": 4865
+    },
+    {
+      "epoch": 0.5,
+      "gumbel_tau": 0.3173,
+      "image_masks_loss": 0.0064,
+      "llm_learning_rate": 2.633090540520936e-06,
+      "loss": 0.6847,
+      "output_text_mask_loss": 0.0381,
+      "predictor_learning_rate": 0.0002633090540520936,
+      "step": 4866
+    },
+    {
+      "epoch": 0.5,
+      "gumbel_tau": 0.3173,
+      "image_masks_loss": 0.0077,
+      "llm_learning_rate": 2.6322620929129646e-06,
+      "loss": 0.7084,
+      "output_text_mask_loss": 0.0275,
+      "predictor_learning_rate": 0.0002632262092912964,
+      "step": 4867
+    },
+    {
+      "epoch": 0.5,
+      "gumbel_tau": 0.3172,
+      "image_masks_loss": 0.0074,
+      "llm_learning_rate": 2.631433630739996e-06,
+      "loss": 0.7603,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.0002631433630739996,
+      "step": 4868
+    },
+    {
+      "epoch": 0.5,
+      "gumbel_tau": 0.3171,
+      "image_masks_loss": 0.007,
+      "llm_learning_rate": 2.6306051540932635e-06,
+      "loss": 0.8066,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00026306051540932635,
+      "step": 4869
+    },
+    {
+      "epoch": 0.5,
+      "gumbel_tau": 0.317,
+      "image_masks_loss": 0.014,
+      "llm_learning_rate": 2.6297766630640005e-06,
+      "loss": 0.6912,
+      "output_text_mask_loss": 0.0023,
+      "predictor_learning_rate": 0.00026297766630640004,
+      "step": 4870
+    },
+    {
+      "epoch": 0.5,
+      "gumbel_tau": 0.317,
+      "image_masks_loss": 0.0163,
+      "llm_learning_rate": 2.628948157743442e-06,
+      "loss": 0.7274,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0002628948157743442,
+      "step": 4871
+    },
+    {
+      "epoch": 0.5,
+      "gumbel_tau": 0.3169,
+      "image_masks_loss": 0.01,
+      "llm_learning_rate": 2.628119638222826e-06,
+      "loss": 0.7144,
+      "output_text_mask_loss": 0.001,
+      "predictor_learning_rate": 0.0002628119638222826,
+      "step": 4872
+    },
+    {
+      "epoch": 0.5,
+      "gumbel_tau": 0.3168,
+      "image_masks_loss": 0.0087,
+      "llm_learning_rate": 2.627291104593389e-06,
+      "loss": 0.7113,
+      "output_text_mask_loss": 0.0034,
+      "predictor_learning_rate": 0.0002627291104593389,
+      "step": 4873
+    },
+    {
+      "epoch": 0.5,
+      "gumbel_tau": 0.3168,
+      "image_masks_loss": 0.013,
+      "llm_learning_rate": 2.6264625569463737e-06,
+      "loss": 0.7589,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.00026264625569463734,
+      "step": 4874
+    },
+    {
+      "epoch": 0.5,
+      "gumbel_tau": 0.3167,
+      "image_masks_loss": 0.0292,
+      "llm_learning_rate": 2.6256339953730186e-06,
+      "loss": 0.6613,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.00026256339953730184,
+      "step": 4875
+    },
+    {
+      "epoch": 0.5,
+      "gumbel_tau": 0.3166,
+      "image_masks_loss": 0.0204,
+      "llm_learning_rate": 2.6248054199645678e-06,
+      "loss": 0.7888,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.00026248054199645677,
+      "step": 4876
+    },
+    {
+      "epoch": 0.5,
+      "gumbel_tau": 0.3165,
+      "image_masks_loss": 0.0152,
+      "llm_learning_rate": 2.623976830812267e-06,
+      "loss": 0.8009,
+      "output_text_mask_loss": 0.0023,
+      "predictor_learning_rate": 0.0002623976830812267,
+      "step": 4877
+    },
+    {
+      "epoch": 0.5,
+      "gumbel_tau": 0.3165,
+      "image_masks_loss": 0.0073,
+      "llm_learning_rate": 2.623148228007361e-06,
+      "loss": 0.7248,
+      "output_text_mask_loss": 0.0679,
+      "predictor_learning_rate": 0.0002623148228007361,
+      "step": 4878
+    },
+    {
+      "epoch": 0.5,
+      "gumbel_tau": 0.3164,
+      "image_masks_loss": 0.0315,
+      "llm_learning_rate": 2.6223196116410976e-06,
+      "loss": 0.7506,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00026223196116410976,
+      "step": 4879
+    },
+    {
+      "epoch": 0.5,
+      "gumbel_tau": 0.3163,
+      "image_masks_loss": 0.034,
+      "llm_learning_rate": 2.6214909818047266e-06,
+      "loss": 0.8272,
+      "output_text_mask_loss": 0.0124,
+      "predictor_learning_rate": 0.0002621490981804727,
+      "step": 4880
+    },
+    {
+      "epoch": 0.5,
+      "gumbel_tau": 0.3162,
+      "image_masks_loss": 0.0114,
+      "llm_learning_rate": 2.6206623385894982e-06,
+      "loss": 0.7691,
+      "output_text_mask_loss": 0.0355,
+      "predictor_learning_rate": 0.00026206623385894976,
+      "step": 4881
+    },
+    {
+      "epoch": 0.5,
+      "gumbel_tau": 0.3162,
+      "image_masks_loss": 0.0081,
+      "llm_learning_rate": 2.6198336820866645e-06,
+      "loss": 0.7419,
+      "output_text_mask_loss": 0.058,
+      "predictor_learning_rate": 0.00026198336820866644,
+      "step": 4882
+    },
+    {
+      "epoch": 0.5,
+      "gumbel_tau": 0.3161,
+      "image_masks_loss": 0.0178,
+      "llm_learning_rate": 2.6190050123874784e-06,
+      "loss": 0.7194,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.0002619005012387478,
+      "step": 4883
+    },
+    {
+      "epoch": 0.5,
+      "gumbel_tau": 0.316,
+      "image_masks_loss": 0.0244,
+      "llm_learning_rate": 2.618176329583195e-06,
+      "loss": 0.7312,
+      "output_text_mask_loss": 0.0406,
+      "predictor_learning_rate": 0.0002618176329583195,
+      "step": 4884
+    },
+    {
+      "epoch": 0.5,
+      "gumbel_tau": 0.3159,
+      "image_masks_loss": 0.0134,
+      "llm_learning_rate": 2.6173476337650717e-06,
+      "loss": 0.7259,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.00026173476337650715,
+      "step": 4885
+    },
+    {
+      "epoch": 0.5,
+      "gumbel_tau": 0.3159,
+      "image_masks_loss": 0.016,
+      "llm_learning_rate": 2.6165189250243643e-06,
+      "loss": 0.6597,
+      "output_text_mask_loss": 0.0162,
+      "predictor_learning_rate": 0.00026165189250243644,
+      "step": 4886
+    },
+    {
+      "epoch": 0.5,
+      "gumbel_tau": 0.3158,
+      "image_masks_loss": 0.0125,
+      "llm_learning_rate": 2.6156902034523346e-06,
+      "loss": 0.7664,
+      "output_text_mask_loss": 0.0095,
+      "predictor_learning_rate": 0.00026156902034523343,
+      "step": 4887
+    },
+    {
+      "epoch": 0.5,
+      "gumbel_tau": 0.3157,
+      "image_masks_loss": 0.0145,
+      "llm_learning_rate": 2.614861469140241e-06,
+      "loss": 0.8126,
+      "output_text_mask_loss": 0.0374,
+      "predictor_learning_rate": 0.0002614861469140241,
+      "step": 4888
+    },
+    {
+      "epoch": 0.5,
+      "gumbel_tau": 0.3156,
+      "image_masks_loss": 0.0092,
+      "llm_learning_rate": 2.614032722179347e-06,
+      "loss": 0.7723,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.0002614032722179347,
+      "step": 4889
+    },
+    {
+      "epoch": 0.5,
+      "gumbel_tau": 0.3156,
+      "image_masks_loss": 0.0334,
+      "llm_learning_rate": 2.6132039626609157e-06,
+      "loss": 0.7115,
+      "output_text_mask_loss": 0.0073,
+      "predictor_learning_rate": 0.00026132039626609155,
+      "step": 4890
+    },
+    {
+      "epoch": 0.5,
+      "gumbel_tau": 0.3155,
+      "image_masks_loss": 0.0091,
+      "llm_learning_rate": 2.6123751906762124e-06,
+      "loss": 0.7562,
+      "output_text_mask_loss": 0.0103,
+      "predictor_learning_rate": 0.00026123751906762123,
+      "step": 4891
+    },
+    {
+      "epoch": 0.5,
+      "gumbel_tau": 0.3154,
+      "image_masks_loss": 0.0086,
+      "llm_learning_rate": 2.6115464063165026e-06,
+      "loss": 0.7177,
+      "output_text_mask_loss": 0.0231,
+      "predictor_learning_rate": 0.00026115464063165025,
+      "step": 4892
+    },
+    {
+      "epoch": 0.5,
+      "gumbel_tau": 0.3153,
+      "image_masks_loss": 0.0109,
+      "llm_learning_rate": 2.6107176096730537e-06,
+      "loss": 0.8124,
+      "output_text_mask_loss": 0.0378,
+      "predictor_learning_rate": 0.0002610717609673053,
+      "step": 4893
+    },
+    {
+      "epoch": 0.5,
+      "gumbel_tau": 0.3153,
+      "image_masks_loss": 0.0155,
+      "llm_learning_rate": 2.6098888008371348e-06,
+      "loss": 0.789,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.00026098888008371345,
+      "step": 4894
+    },
+    {
+      "epoch": 0.5,
+      "gumbel_tau": 0.3152,
+      "image_masks_loss": 0.0142,
+      "llm_learning_rate": 2.6090599799000165e-06,
+      "loss": 0.7779,
+      "output_text_mask_loss": 0.0153,
+      "predictor_learning_rate": 0.00026090599799000164,
+      "step": 4895
+    },
+    {
+      "epoch": 0.5,
+      "gumbel_tau": 0.3151,
+      "image_masks_loss": 0.022,
+      "llm_learning_rate": 2.6082311469529705e-06,
+      "loss": 0.8053,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.000260823114695297,
+      "step": 4896
+    },
+    {
+      "epoch": 0.5,
+      "gumbel_tau": 0.315,
+      "image_masks_loss": 0.0075,
+      "llm_learning_rate": 2.6074023020872697e-06,
+      "loss": 0.7808,
+      "output_text_mask_loss": 0.0035,
+      "predictor_learning_rate": 0.00026074023020872695,
+      "step": 4897
+    },
+    {
+      "epoch": 0.5,
+      "gumbel_tau": 0.315,
+      "image_masks_loss": 0.0158,
+      "llm_learning_rate": 2.6065734453941883e-06,
+      "loss": 0.7875,
+      "output_text_mask_loss": 0.0078,
+      "predictor_learning_rate": 0.0002606573445394188,
+      "step": 4898
+    },
+    {
+      "epoch": 0.5,
+      "gumbel_tau": 0.3149,
+      "image_masks_loss": 0.0112,
+      "llm_learning_rate": 2.6057445769650007e-06,
+      "loss": 0.7391,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.0002605744576965001,
+      "step": 4899
+    },
+    {
+      "epoch": 0.5,
+      "gumbel_tau": 0.3148,
+      "image_masks_loss": 0.0081,
+      "llm_learning_rate": 2.6049156968909855e-06,
+      "loss": 0.6794,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.00026049156968909853,
+      "step": 4900
+    },
+    {
+      "epoch": 0.5,
+      "gumbel_tau": 0.3147,
+      "image_masks_loss": 0.0246,
+      "llm_learning_rate": 2.60408680526342e-06,
+      "loss": 0.6954,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.000260408680526342,
+      "step": 4901
+    },
+    {
+      "epoch": 0.5,
+      "gumbel_tau": 0.3147,
+      "image_masks_loss": 0.0237,
+      "llm_learning_rate": 2.603257902173584e-06,
+      "loss": 0.7416,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.00026032579021735834,
+      "step": 4902
+    },
+    {
+      "epoch": 0.5,
+      "gumbel_tau": 0.3146,
+      "image_masks_loss": 0.0141,
+      "llm_learning_rate": 2.602428987712757e-06,
+      "loss": 0.7152,
+      "output_text_mask_loss": 0.0546,
+      "predictor_learning_rate": 0.0002602428987712757,
+      "step": 4903
+    },
+    {
+      "epoch": 0.5,
+      "gumbel_tau": 0.3145,
+      "image_masks_loss": 0.0197,
+      "llm_learning_rate": 2.6016000619722226e-06,
+      "loss": 0.7094,
+      "output_text_mask_loss": 0.0381,
+      "predictor_learning_rate": 0.00026016000619722226,
+      "step": 4904
+    },
+    {
+      "epoch": 0.5,
+      "gumbel_tau": 0.3144,
+      "image_masks_loss": 0.0372,
+      "llm_learning_rate": 2.600771125043262e-06,
+      "loss": 0.7701,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0002600771125043262,
+      "step": 4905
+    },
+    {
+      "epoch": 0.5,
+      "gumbel_tau": 0.3144,
+      "image_masks_loss": 0.0022,
+      "llm_learning_rate": 2.5999421770171616e-06,
+      "loss": 0.7911,
+      "output_text_mask_loss": 0.0656,
+      "predictor_learning_rate": 0.0002599942177017162,
+      "step": 4906
+    },
+    {
+      "epoch": 0.5,
+      "gumbel_tau": 0.3143,
+      "image_masks_loss": 0.0184,
+      "llm_learning_rate": 2.599113217985206e-06,
+      "loss": 0.7498,
+      "output_text_mask_loss": 0.0078,
+      "predictor_learning_rate": 0.00025991132179852054,
+      "step": 4907
+    },
+    {
+      "epoch": 0.5,
+      "gumbel_tau": 0.3142,
+      "image_masks_loss": 0.0187,
+      "llm_learning_rate": 2.598284248038682e-06,
+      "loss": 0.7716,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.0002598284248038682,
+      "step": 4908
+    },
+    {
+      "epoch": 0.5,
+      "gumbel_tau": 0.3141,
+      "image_masks_loss": 0.0125,
+      "llm_learning_rate": 2.597455267268877e-06,
+      "loss": 0.7291,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00025974552672688767,
+      "step": 4909
+    },
+    {
+      "epoch": 0.5,
+      "gumbel_tau": 0.3141,
+      "image_masks_loss": 0.006,
+      "llm_learning_rate": 2.596626275767082e-06,
+      "loss": 0.804,
+      "output_text_mask_loss": 0.0916,
+      "predictor_learning_rate": 0.0002596626275767082,
+      "step": 4910
+    },
+    {
+      "epoch": 0.5,
+      "gumbel_tau": 0.314,
+      "image_masks_loss": 0.0148,
+      "llm_learning_rate": 2.595797273624585e-06,
+      "loss": 0.7351,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 0.0002595797273624585,
+      "step": 4911
+    },
+    {
+      "epoch": 0.5,
+      "gumbel_tau": 0.3139,
+      "image_masks_loss": 0.0172,
+      "llm_learning_rate": 2.59496826093268e-06,
+      "loss": 0.7648,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00025949682609326797,
+      "step": 4912
+    },
+    {
+      "epoch": 0.5,
+      "gumbel_tau": 0.3138,
+      "image_masks_loss": 0.0196,
+      "llm_learning_rate": 2.5941392377826587e-06,
+      "loss": 0.7436,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.00025941392377826586,
+      "step": 4913
+    },
+    {
+      "epoch": 0.5,
+      "gumbel_tau": 0.3138,
+      "image_masks_loss": 0.0122,
+      "llm_learning_rate": 2.5933102042658137e-06,
+      "loss": 0.6514,
+      "output_text_mask_loss": 0.0067,
+      "predictor_learning_rate": 0.00025933102042658134,
+      "step": 4914
+    },
+    {
+      "epoch": 0.5,
+      "gumbel_tau": 0.3137,
+      "image_masks_loss": 0.0239,
+      "llm_learning_rate": 2.592481160473442e-06,
+      "loss": 0.71,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.00025924811604734416,
+      "step": 4915
+    },
+    {
+      "epoch": 0.5,
+      "gumbel_tau": 0.3136,
+      "image_masks_loss": 0.0188,
+      "llm_learning_rate": 2.591652106496838e-06,
+      "loss": 0.7151,
+      "output_text_mask_loss": 0.0101,
+      "predictor_learning_rate": 0.0002591652106496838,
+      "step": 4916
+    },
+    {
+      "epoch": 0.5,
+      "gumbel_tau": 0.3136,
+      "image_masks_loss": 0.005,
+      "llm_learning_rate": 2.5908230424273008e-06,
+      "loss": 0.7168,
+      "output_text_mask_loss": 0.0429,
+      "predictor_learning_rate": 0.0002590823042427301,
+      "step": 4917
+    },
+    {
+      "epoch": 0.5,
+      "gumbel_tau": 0.3135,
+      "image_masks_loss": 0.0153,
+      "llm_learning_rate": 2.5899939683561277e-06,
+      "loss": 0.6987,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00025899939683561276,
+      "step": 4918
+    },
+    {
+      "epoch": 0.5,
+      "gumbel_tau": 0.3134,
+      "image_masks_loss": 0.007,
+      "llm_learning_rate": 2.589164884374618e-06,
+      "loss": 0.7439,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0002589164884374618,
+      "step": 4919
+    },
+    {
+      "epoch": 0.5,
+      "gumbel_tau": 0.3133,
+      "image_masks_loss": 0.015,
+      "llm_learning_rate": 2.5883357905740735e-06,
+      "loss": 0.6533,
+      "output_text_mask_loss": 0.037,
+      "predictor_learning_rate": 0.00025883357905740736,
+      "step": 4920
+    },
+    {
+      "epoch": 0.5,
+      "gumbel_tau": 0.3133,
+      "image_masks_loss": 0.0263,
+      "llm_learning_rate": 2.5875066870457936e-06,
+      "loss": 0.7747,
+      "output_text_mask_loss": 0.0055,
+      "predictor_learning_rate": 0.00025875066870457937,
+      "step": 4921
+    },
+    {
+      "epoch": 0.5,
+      "gumbel_tau": 0.3132,
+      "image_masks_loss": 0.0043,
+      "llm_learning_rate": 2.5866775738810833e-06,
+      "loss": 0.7169,
+      "output_text_mask_loss": 0.0284,
+      "predictor_learning_rate": 0.0002586677573881083,
+      "step": 4922
+    },
+    {
+      "epoch": 0.5,
+      "gumbel_tau": 0.3131,
+      "image_masks_loss": 0.0108,
+      "llm_learning_rate": 2.5858484511712455e-06,
+      "loss": 0.7268,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.0002585848451171245,
+      "step": 4923
+    },
+    {
+      "epoch": 0.5,
+      "gumbel_tau": 0.313,
+      "image_masks_loss": 0.0112,
+      "llm_learning_rate": 2.585019319007585e-06,
+      "loss": 0.687,
+      "output_text_mask_loss": 0.033,
+      "predictor_learning_rate": 0.0002585019319007585,
+      "step": 4924
+    },
+    {
+      "epoch": 0.5,
+      "gumbel_tau": 0.313,
+      "image_masks_loss": 0.0238,
+      "llm_learning_rate": 2.584190177481408e-06,
+      "loss": 0.7754,
+      "output_text_mask_loss": 0.0193,
+      "predictor_learning_rate": 0.00025841901774814076,
+      "step": 4925
+    },
+    {
+      "epoch": 0.5,
+      "gumbel_tau": 0.3129,
+      "image_masks_loss": 0.0107,
+      "llm_learning_rate": 2.58336102668402e-06,
+      "loss": 0.711,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.000258336102668402,
+      "step": 4926
+    },
+    {
+      "epoch": 0.5,
+      "gumbel_tau": 0.3128,
+      "image_masks_loss": 0.0059,
+      "llm_learning_rate": 2.5825318667067313e-06,
+      "loss": 0.8388,
+      "output_text_mask_loss": 0.0656,
+      "predictor_learning_rate": 0.0002582531866706731,
+      "step": 4927
+    },
+    {
+      "epoch": 0.5,
+      "gumbel_tau": 0.3127,
+      "image_masks_loss": 0.0122,
+      "llm_learning_rate": 2.5817026976408487e-06,
+      "loss": 0.8337,
+      "output_text_mask_loss": 0.0193,
+      "predictor_learning_rate": 0.0002581702697640849,
+      "step": 4928
+    },
+    {
+      "epoch": 0.51,
+      "gumbel_tau": 0.3127,
+      "image_masks_loss": 0.0119,
+      "llm_learning_rate": 2.5808735195776836e-06,
+      "loss": 0.6455,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.0002580873519577683,
+      "step": 4929
+    },
+    {
+      "epoch": 0.51,
+      "gumbel_tau": 0.3126,
+      "image_masks_loss": 0.0177,
+      "llm_learning_rate": 2.5800443326085467e-06,
+      "loss": 0.7325,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00025800443326085465,
+      "step": 4930
+    },
+    {
+      "epoch": 0.51,
+      "gumbel_tau": 0.3125,
+      "image_masks_loss": 0.036,
+      "llm_learning_rate": 2.579215136824749e-06,
+      "loss": 0.6947,
+      "output_text_mask_loss": 0.0233,
+      "predictor_learning_rate": 0.00025792151368247485,
+      "step": 4931
+    },
+    {
+      "epoch": 0.51,
+      "gumbel_tau": 0.3124,
+      "image_masks_loss": 0.028,
+      "llm_learning_rate": 2.5783859323176043e-06,
+      "loss": 0.6942,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 0.0002578385932317604,
+      "step": 4932
+    },
+    {
+      "epoch": 0.51,
+      "gumbel_tau": 0.3124,
+      "image_masks_loss": 0.0082,
+      "llm_learning_rate": 2.5775567191784257e-06,
+      "loss": 0.7219,
+      "output_text_mask_loss": 0.0025,
+      "predictor_learning_rate": 0.0002577556719178426,
+      "step": 4933
+    },
+    {
+      "epoch": 0.51,
+      "gumbel_tau": 0.3123,
+      "image_masks_loss": 0.0181,
+      "llm_learning_rate": 2.5767274974985292e-06,
+      "loss": 0.7001,
+      "output_text_mask_loss": 0.0055,
+      "predictor_learning_rate": 0.0002576727497498529,
+      "step": 4934
+    },
+    {
+      "epoch": 0.51,
+      "gumbel_tau": 0.3122,
+      "image_masks_loss": 0.0353,
+      "llm_learning_rate": 2.5758982673692284e-06,
+      "loss": 0.7457,
+      "output_text_mask_loss": 0.0156,
+      "predictor_learning_rate": 0.0002575898267369228,
+      "step": 4935
+    },
+    {
+      "epoch": 0.51,
+      "gumbel_tau": 0.3122,
+      "image_masks_loss": 0.0259,
+      "llm_learning_rate": 2.5750690288818426e-06,
+      "loss": 0.7377,
+      "output_text_mask_loss": 0.0038,
+      "predictor_learning_rate": 0.00025750690288818426,
+      "step": 4936
+    },
+    {
+      "epoch": 0.51,
+      "gumbel_tau": 0.3121,
+      "image_masks_loss": 0.0118,
+      "llm_learning_rate": 2.5742397821276876e-06,
+      "loss": 0.6746,
+      "output_text_mask_loss": 0.0119,
+      "predictor_learning_rate": 0.00025742397821276877,
+      "step": 4937
+    },
+    {
+      "epoch": 0.51,
+      "gumbel_tau": 0.312,
+      "image_masks_loss": 0.0148,
+      "llm_learning_rate": 2.5734105271980824e-06,
+      "loss": 0.7984,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 0.00025734105271980824,
+      "step": 4938
+    },
+    {
+      "epoch": 0.51,
+      "gumbel_tau": 0.3119,
+      "image_masks_loss": 0.0162,
+      "llm_learning_rate": 2.5725812641843474e-06,
+      "loss": 0.8297,
+      "output_text_mask_loss": 0.0292,
+      "predictor_learning_rate": 0.0002572581264184347,
+      "step": 4939
+    },
+    {
+      "epoch": 0.51,
+      "gumbel_tau": 0.3119,
+      "image_masks_loss": 0.0121,
+      "llm_learning_rate": 2.571751993177801e-06,
+      "loss": 0.7804,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.0002571751993177801,
+      "step": 4940
+    },
+    {
+      "epoch": 0.51,
+      "gumbel_tau": 0.3118,
+      "image_masks_loss": 0.0048,
+      "llm_learning_rate": 2.570922714269765e-06,
+      "loss": 0.8067,
+      "output_text_mask_loss": 0.0124,
+      "predictor_learning_rate": 0.0002570922714269765,
+      "step": 4941
+    },
+    {
+      "epoch": 0.51,
+      "gumbel_tau": 0.3117,
+      "image_masks_loss": 0.0123,
+      "llm_learning_rate": 2.570093427551562e-06,
+      "loss": 0.7877,
+      "output_text_mask_loss": 0.0248,
+      "predictor_learning_rate": 0.00025700934275515617,
+      "step": 4942
+    },
+    {
+      "epoch": 0.51,
+      "gumbel_tau": 0.3116,
+      "image_masks_loss": 0.0043,
+      "llm_learning_rate": 2.569264133114515e-06,
+      "loss": 0.7094,
+      "output_text_mask_loss": 0.0065,
+      "predictor_learning_rate": 0.0002569264133114515,
+      "step": 4943
+    },
+    {
+      "epoch": 0.51,
+      "gumbel_tau": 0.3116,
+      "image_masks_loss": 0.0435,
+      "llm_learning_rate": 2.5684348310499467e-06,
+      "loss": 0.639,
+      "output_text_mask_loss": 0.0158,
+      "predictor_learning_rate": 0.0002568434831049947,
+      "step": 4944
+    },
+    {
+      "epoch": 0.51,
+      "gumbel_tau": 0.3115,
+      "image_masks_loss": 0.0094,
+      "llm_learning_rate": 2.567605521449183e-06,
+      "loss": 0.7839,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.0002567605521449183,
+      "step": 4945
+    },
+    {
+      "epoch": 0.51,
+      "gumbel_tau": 0.3114,
+      "image_masks_loss": 0.0205,
+      "llm_learning_rate": 2.5667762044035487e-06,
+      "loss": 0.7521,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.0002566776204403549,
+      "step": 4946
+    },
+    {
+      "epoch": 0.51,
+      "gumbel_tau": 0.3113,
+      "image_masks_loss": 0.0036,
+      "llm_learning_rate": 2.5659468800043704e-06,
+      "loss": 0.744,
+      "output_text_mask_loss": 0.0294,
+      "predictor_learning_rate": 0.000256594688000437,
+      "step": 4947
+    },
+    {
+      "epoch": 0.51,
+      "gumbel_tau": 0.3113,
+      "image_masks_loss": 0.0094,
+      "llm_learning_rate": 2.565117548342974e-06,
+      "loss": 0.7364,
+      "output_text_mask_loss": 0.0041,
+      "predictor_learning_rate": 0.0002565117548342974,
+      "step": 4948
+    },
+    {
+      "epoch": 0.51,
+      "gumbel_tau": 0.3112,
+      "image_masks_loss": 0.013,
+      "llm_learning_rate": 2.5642882095106885e-06,
+      "loss": 0.6679,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.00025642882095106885,
+      "step": 4949
+    },
+    {
+      "epoch": 0.51,
+      "gumbel_tau": 0.3111,
+      "image_masks_loss": 0.0154,
+      "llm_learning_rate": 2.563458863598842e-06,
+      "loss": 0.7544,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0002563458863598842,
+      "step": 4950
+    },
+    {
+      "epoch": 0.51,
+      "gumbel_tau": 0.311,
+      "image_masks_loss": 0.0125,
+      "llm_learning_rate": 2.562629510698764e-06,
+      "loss": 0.7118,
+      "output_text_mask_loss": 0.0206,
+      "predictor_learning_rate": 0.0002562629510698764,
+      "step": 4951
+    },
+    {
+      "epoch": 0.51,
+      "gumbel_tau": 0.311,
+      "image_masks_loss": 0.007,
+      "llm_learning_rate": 2.561800150901785e-06,
+      "loss": 0.7595,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00025618001509017846,
+      "step": 4952
+    },
+    {
+      "epoch": 0.51,
+      "gumbel_tau": 0.3109,
+      "image_masks_loss": 0.007,
+      "llm_learning_rate": 2.560970784299236e-06,
+      "loss": 0.7052,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0002560970784299236,
+      "step": 4953
+    },
+    {
+      "epoch": 0.51,
+      "gumbel_tau": 0.3108,
+      "image_masks_loss": 0.0234,
+      "llm_learning_rate": 2.560141410982447e-06,
+      "loss": 0.7533,
+      "output_text_mask_loss": 0.1221,
+      "predictor_learning_rate": 0.00025601414109824473,
+      "step": 4954
+    },
+    {
+      "epoch": 0.51,
+      "gumbel_tau": 0.3108,
+      "image_masks_loss": 0.0106,
+      "llm_learning_rate": 2.5593120310427535e-06,
+      "loss": 0.6757,
+      "output_text_mask_loss": 0.0094,
+      "predictor_learning_rate": 0.0002559312031042753,
+      "step": 4955
+    },
+    {
+      "epoch": 0.51,
+      "gumbel_tau": 0.3107,
+      "image_masks_loss": 0.0092,
+      "llm_learning_rate": 2.5584826445714867e-06,
+      "loss": 0.7367,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 0.00025584826445714865,
+      "step": 4956
+    },
+    {
+      "epoch": 0.51,
+      "gumbel_tau": 0.3106,
+      "image_masks_loss": 0.0072,
+      "llm_learning_rate": 2.5576532516599806e-06,
+      "loss": 0.7445,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00025576532516599805,
+      "step": 4957
+    },
+    {
+      "epoch": 0.51,
+      "gumbel_tau": 0.3105,
+      "image_masks_loss": 0.0298,
+      "llm_learning_rate": 2.55682385239957e-06,
+      "loss": 0.7315,
+      "output_text_mask_loss": 0.0622,
+      "predictor_learning_rate": 0.000255682385239957,
+      "step": 4958
+    },
+    {
+      "epoch": 0.51,
+      "gumbel_tau": 0.3105,
+      "image_masks_loss": 0.0125,
+      "llm_learning_rate": 2.55599444688159e-06,
+      "loss": 0.7787,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.000255599444688159,
+      "step": 4959
+    },
+    {
+      "epoch": 0.51,
+      "gumbel_tau": 0.3104,
+      "image_masks_loss": 0.0064,
+      "llm_learning_rate": 2.5551650351973765e-06,
+      "loss": 0.7277,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00025551650351973763,
+      "step": 4960
+    },
+    {
+      "epoch": 0.51,
+      "gumbel_tau": 0.3103,
+      "image_masks_loss": 0.0205,
+      "llm_learning_rate": 2.5543356174382673e-06,
+      "loss": 0.7611,
+      "output_text_mask_loss": 0.0023,
+      "predictor_learning_rate": 0.0002554335617438267,
+      "step": 4961
+    },
+    {
+      "epoch": 0.51,
+      "gumbel_tau": 0.3102,
+      "image_masks_loss": 0.0435,
+      "llm_learning_rate": 2.5535061936955984e-06,
+      "loss": 0.7344,
+      "output_text_mask_loss": 0.0021,
+      "predictor_learning_rate": 0.0002553506193695598,
+      "step": 4962
+    },
+    {
+      "epoch": 0.51,
+      "gumbel_tau": 0.3102,
+      "image_masks_loss": 0.0118,
+      "llm_learning_rate": 2.5526767640607087e-06,
+      "loss": 0.6118,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00025526767640607085,
+      "step": 4963
+    },
+    {
+      "epoch": 0.51,
+      "gumbel_tau": 0.3101,
+      "image_masks_loss": 0.0288,
+      "llm_learning_rate": 2.551847328624936e-06,
+      "loss": 0.6552,
+      "output_text_mask_loss": 0.0162,
+      "predictor_learning_rate": 0.00025518473286249355,
+      "step": 4964
+    },
+    {
+      "epoch": 0.51,
+      "gumbel_tau": 0.31,
+      "image_masks_loss": 0.0135,
+      "llm_learning_rate": 2.5510178874796197e-06,
+      "loss": 0.7655,
+      "output_text_mask_loss": 0.0086,
+      "predictor_learning_rate": 0.00025510178874796195,
+      "step": 4965
+    },
+    {
+      "epoch": 0.51,
+      "gumbel_tau": 0.3099,
+      "image_masks_loss": 0.0256,
+      "llm_learning_rate": 2.550188440716101e-06,
+      "loss": 0.6958,
+      "output_text_mask_loss": 0.0071,
+      "predictor_learning_rate": 0.00025501884407161005,
+      "step": 4966
+    },
+    {
+      "epoch": 0.51,
+      "gumbel_tau": 0.3099,
+      "image_masks_loss": 0.0097,
+      "llm_learning_rate": 2.5493589884257187e-06,
+      "loss": 0.7468,
+      "output_text_mask_loss": 0.0055,
+      "predictor_learning_rate": 0.00025493589884257185,
+      "step": 4967
+    },
+    {
+      "epoch": 0.51,
+      "gumbel_tau": 0.3098,
+      "image_masks_loss": 0.0027,
+      "llm_learning_rate": 2.548529530699815e-06,
+      "loss": 0.6701,
+      "output_text_mask_loss": 0.001,
+      "predictor_learning_rate": 0.00025485295306998146,
+      "step": 4968
+    },
+    {
+      "epoch": 0.51,
+      "gumbel_tau": 0.3097,
+      "image_masks_loss": 0.0093,
+      "llm_learning_rate": 2.547700067629731e-06,
+      "loss": 0.757,
+      "output_text_mask_loss": 0.021,
+      "predictor_learning_rate": 0.0002547700067629731,
+      "step": 4969
+    },
+    {
+      "epoch": 0.51,
+      "gumbel_tau": 0.3097,
+      "image_masks_loss": 0.0095,
+      "llm_learning_rate": 2.5468705993068088e-06,
+      "loss": 0.7519,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0002546870599306809,
+      "step": 4970
+    },
+    {
+      "epoch": 0.51,
+      "gumbel_tau": 0.3096,
+      "image_masks_loss": 0.0105,
+      "llm_learning_rate": 2.546041125822392e-06,
+      "loss": 0.817,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.0002546041125822392,
+      "step": 4971
+    },
+    {
+      "epoch": 0.51,
+      "gumbel_tau": 0.3095,
+      "image_masks_loss": 0.0211,
+      "llm_learning_rate": 2.5452116472678244e-06,
+      "loss": 0.7388,
+      "output_text_mask_loss": 0.006,
+      "predictor_learning_rate": 0.00025452116472678246,
+      "step": 4972
+    },
+    {
+      "epoch": 0.51,
+      "gumbel_tau": 0.3094,
+      "image_masks_loss": 0.0095,
+      "llm_learning_rate": 2.5443821637344495e-06,
+      "loss": 0.7406,
+      "output_text_mask_loss": 0.0193,
+      "predictor_learning_rate": 0.00025443821637344494,
+      "step": 4973
+    },
+    {
+      "epoch": 0.51,
+      "gumbel_tau": 0.3094,
+      "image_masks_loss": 0.0089,
+      "llm_learning_rate": 2.5435526753136115e-06,
+      "loss": 0.8067,
+      "output_text_mask_loss": 0.0124,
+      "predictor_learning_rate": 0.0002543552675313611,
+      "step": 4974
+    },
+    {
+      "epoch": 0.51,
+      "gumbel_tau": 0.3093,
+      "image_masks_loss": 0.0233,
+      "llm_learning_rate": 2.5427231820966565e-06,
+      "loss": 0.8265,
+      "output_text_mask_loss": 0.0668,
+      "predictor_learning_rate": 0.00025427231820966564,
+      "step": 4975
+    },
+    {
+      "epoch": 0.51,
+      "gumbel_tau": 0.3092,
+      "image_masks_loss": 0.0221,
+      "llm_learning_rate": 2.5418936841749285e-06,
+      "loss": 0.8004,
+      "output_text_mask_loss": 0.0051,
+      "predictor_learning_rate": 0.00025418936841749284,
+      "step": 4976
+    },
+    {
+      "epoch": 0.51,
+      "gumbel_tau": 0.3091,
+      "image_masks_loss": 0.0064,
+      "llm_learning_rate": 2.541064181639776e-06,
+      "loss": 0.7984,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.00025410641816397755,
+      "step": 4977
+    },
+    {
+      "epoch": 0.51,
+      "gumbel_tau": 0.3091,
+      "image_masks_loss": 0.0109,
+      "llm_learning_rate": 2.5402346745825436e-06,
+      "loss": 0.7225,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.0002540234674582543,
+      "step": 4978
+    },
+    {
+      "epoch": 0.51,
+      "gumbel_tau": 0.309,
+      "image_masks_loss": 0.0158,
+      "llm_learning_rate": 2.5394051630945793e-06,
+      "loss": 0.7148,
+      "output_text_mask_loss": 0.0324,
+      "predictor_learning_rate": 0.00025394051630945794,
+      "step": 4979
+    },
+    {
+      "epoch": 0.51,
+      "gumbel_tau": 0.3089,
+      "image_masks_loss": 0.0144,
+      "llm_learning_rate": 2.53857564726723e-06,
+      "loss": 0.7311,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.000253857564726723,
+      "step": 4980
+    },
+    {
+      "epoch": 0.51,
+      "gumbel_tau": 0.3089,
+      "image_masks_loss": 0.0133,
+      "llm_learning_rate": 2.5377461271918453e-06,
+      "loss": 0.583,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0002537746127191845,
+      "step": 4981
+    },
+    {
+      "epoch": 0.51,
+      "gumbel_tau": 0.3088,
+      "image_masks_loss": 0.0261,
+      "llm_learning_rate": 2.536916602959773e-06,
+      "loss": 0.7446,
+      "output_text_mask_loss": 0.0058,
+      "predictor_learning_rate": 0.0002536916602959773,
+      "step": 4982
+    },
+    {
+      "epoch": 0.51,
+      "gumbel_tau": 0.3087,
+      "image_masks_loss": 0.0171,
+      "llm_learning_rate": 2.536087074662361e-06,
+      "loss": 0.6845,
+      "output_text_mask_loss": 0.0345,
+      "predictor_learning_rate": 0.0002536087074662361,
+      "step": 4983
+    },
+    {
+      "epoch": 0.51,
+      "gumbel_tau": 0.3086,
+      "image_masks_loss": 0.013,
+      "llm_learning_rate": 2.5352575423909616e-06,
+      "loss": 0.7343,
+      "output_text_mask_loss": 0.0132,
+      "predictor_learning_rate": 0.00025352575423909614,
+      "step": 4984
+    },
+    {
+      "epoch": 0.51,
+      "gumbel_tau": 0.3086,
+      "image_masks_loss": 0.0127,
+      "llm_learning_rate": 2.534428006236923e-06,
+      "loss": 0.7936,
+      "output_text_mask_loss": 0.0086,
+      "predictor_learning_rate": 0.0002534428006236923,
+      "step": 4985
+    },
+    {
+      "epoch": 0.51,
+      "gumbel_tau": 0.3085,
+      "image_masks_loss": 0.0168,
+      "llm_learning_rate": 2.533598466291595e-06,
+      "loss": 0.724,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.00025335984662915945,
+      "step": 4986
+    },
+    {
+      "epoch": 0.51,
+      "gumbel_tau": 0.3084,
+      "image_masks_loss": 0.0211,
+      "llm_learning_rate": 2.5327689226463287e-06,
+      "loss": 0.7325,
+      "output_text_mask_loss": 0.0526,
+      "predictor_learning_rate": 0.00025327689226463287,
+      "step": 4987
+    },
+    {
+      "epoch": 0.51,
+      "gumbel_tau": 0.3083,
+      "image_masks_loss": 0.0128,
+      "llm_learning_rate": 2.5319393753924766e-06,
+      "loss": 0.7319,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00025319393753924765,
+      "step": 4988
+    },
+    {
+      "epoch": 0.51,
+      "gumbel_tau": 0.3083,
+      "image_masks_loss": 0.0082,
+      "llm_learning_rate": 2.5311098246213883e-06,
+      "loss": 0.7474,
+      "output_text_mask_loss": 0.0081,
+      "predictor_learning_rate": 0.0002531109824621388,
+      "step": 4989
+    },
+    {
+      "epoch": 0.51,
+      "gumbel_tau": 0.3082,
+      "image_masks_loss": 0.0098,
+      "llm_learning_rate": 2.5302802704244172e-06,
+      "loss": 0.7594,
+      "output_text_mask_loss": 0.0303,
+      "predictor_learning_rate": 0.0002530280270424417,
+      "step": 4990
+    },
+    {
+      "epoch": 0.51,
+      "gumbel_tau": 0.3081,
+      "image_masks_loss": 0.0135,
+      "llm_learning_rate": 2.529450712892916e-06,
+      "loss": 0.8525,
+      "output_text_mask_loss": 0.0391,
+      "predictor_learning_rate": 0.00025294507128929155,
+      "step": 4991
+    },
+    {
+      "epoch": 0.51,
+      "gumbel_tau": 0.3081,
+      "image_masks_loss": 0.01,
+      "llm_learning_rate": 2.528621152118236e-06,
+      "loss": 0.7209,
+      "output_text_mask_loss": 0.0534,
+      "predictor_learning_rate": 0.00025286211521182355,
+      "step": 4992
+    },
+    {
+      "epoch": 0.51,
+      "gumbel_tau": 0.308,
+      "image_masks_loss": 0.0163,
+      "llm_learning_rate": 2.5277915881917303e-06,
+      "loss": 0.7524,
+      "output_text_mask_loss": 0.0086,
+      "predictor_learning_rate": 0.000252779158819173,
+      "step": 4993
+    },
+    {
+      "epoch": 0.51,
+      "gumbel_tau": 0.3079,
+      "image_masks_loss": 0.0217,
+      "llm_learning_rate": 2.526962021204754e-06,
+      "loss": 0.6946,
+      "output_text_mask_loss": 0.0043,
+      "predictor_learning_rate": 0.00025269620212047537,
+      "step": 4994
+    },
+    {
+      "epoch": 0.51,
+      "gumbel_tau": 0.3078,
+      "image_masks_loss": 0.0082,
+      "llm_learning_rate": 2.52613245124866e-06,
+      "loss": 0.7289,
+      "output_text_mask_loss": 0.0208,
+      "predictor_learning_rate": 0.000252613245124866,
+      "step": 4995
+    },
+    {
+      "epoch": 0.51,
+      "gumbel_tau": 0.3078,
+      "image_masks_loss": 0.0347,
+      "llm_learning_rate": 2.525302878414802e-06,
+      "loss": 0.6649,
+      "output_text_mask_loss": 0.0065,
+      "predictor_learning_rate": 0.0002525302878414802,
+      "step": 4996
+    },
+    {
+      "epoch": 0.51,
+      "gumbel_tau": 0.3077,
+      "image_masks_loss": 0.0175,
+      "llm_learning_rate": 2.5244733027945343e-06,
+      "loss": 0.7689,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.00025244733027945346,
+      "step": 4997
+    },
+    {
+      "epoch": 0.51,
+      "gumbel_tau": 0.3076,
+      "image_masks_loss": 0.0061,
+      "llm_learning_rate": 2.5236437244792128e-06,
+      "loss": 0.749,
+      "output_text_mask_loss": 0.0038,
+      "predictor_learning_rate": 0.0002523643724479213,
+      "step": 4998
+    },
+    {
+      "epoch": 0.51,
+      "gumbel_tau": 0.3075,
+      "image_masks_loss": 0.01,
+      "llm_learning_rate": 2.5228141435601907e-06,
+      "loss": 0.7926,
+      "output_text_mask_loss": 0.0511,
+      "predictor_learning_rate": 0.00025228141435601904,
+      "step": 4999
+    },
+    {
+      "epoch": 0.51,
+      "gumbel_tau": 0.3075,
+      "image_masks_loss": 0.011,
+      "llm_learning_rate": 2.5219845601288245e-06,
+      "loss": 0.7409,
+      "output_text_mask_loss": 0.0107,
+      "predictor_learning_rate": 0.00025219845601288243,
+      "step": 5000
+    },
+    {
+      "epoch": 0.51,
+      "gumbel_tau": 0.3074,
+      "image_masks_loss": 0.01,
+      "llm_learning_rate": 2.5211549742764703e-06,
+      "loss": 0.7996,
+      "output_text_mask_loss": 0.046,
+      "predictor_learning_rate": 0.000252115497427647,
+      "step": 5001
+    },
+    {
+      "epoch": 0.51,
+      "gumbel_tau": 0.3073,
+      "image_masks_loss": 0.0142,
+      "llm_learning_rate": 2.5203253860944825e-06,
+      "loss": 0.7127,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.00025203253860944824,
+      "step": 5002
+    },
+    {
+      "epoch": 0.51,
+      "gumbel_tau": 0.3073,
+      "image_masks_loss": 0.0095,
+      "llm_learning_rate": 2.5194957956742178e-06,
+      "loss": 0.6833,
+      "output_text_mask_loss": 0.0032,
+      "predictor_learning_rate": 0.00025194957956742173,
+      "step": 5003
+    },
+    {
+      "epoch": 0.51,
+      "gumbel_tau": 0.3072,
+      "image_masks_loss": 0.0128,
+      "llm_learning_rate": 2.5186662031070337e-06,
+      "loss": 0.7433,
+      "output_text_mask_loss": 0.0107,
+      "predictor_learning_rate": 0.00025186662031070336,
+      "step": 5004
+    },
+    {
+      "epoch": 0.51,
+      "gumbel_tau": 0.3071,
+      "image_masks_loss": 0.0101,
+      "llm_learning_rate": 2.5178366084842852e-06,
+      "loss": 0.7311,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.0002517836608484285,
+      "step": 5005
+    },
+    {
+      "epoch": 0.51,
+      "gumbel_tau": 0.307,
+      "image_masks_loss": 0.0063,
+      "llm_learning_rate": 2.517007011897329e-06,
+      "loss": 0.7255,
+      "output_text_mask_loss": 0.0062,
+      "predictor_learning_rate": 0.0002517007011897329,
+      "step": 5006
+    },
+    {
+      "epoch": 0.51,
+      "gumbel_tau": 0.307,
+      "image_masks_loss": 0.0068,
+      "llm_learning_rate": 2.5161774134375233e-06,
+      "loss": 0.7495,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.0002516177413437523,
+      "step": 5007
+    },
+    {
+      "epoch": 0.51,
+      "gumbel_tau": 0.3069,
+      "image_masks_loss": 0.0082,
+      "llm_learning_rate": 2.515347813196225e-06,
+      "loss": 0.7543,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0002515347813196225,
+      "step": 5008
+    },
+    {
+      "epoch": 0.51,
+      "gumbel_tau": 0.3068,
+      "image_masks_loss": 0.0246,
+      "llm_learning_rate": 2.514518211264791e-06,
+      "loss": 0.8367,
+      "output_text_mask_loss": 0.0071,
+      "predictor_learning_rate": 0.00025145182112647907,
+      "step": 5009
+    },
+    {
+      "epoch": 0.51,
+      "gumbel_tau": 0.3067,
+      "image_masks_loss": 0.0257,
+      "llm_learning_rate": 2.513688607734579e-06,
+      "loss": 0.7917,
+      "output_text_mask_loss": 0.0012,
+      "predictor_learning_rate": 0.0002513688607734579,
+      "step": 5010
+    },
+    {
+      "epoch": 0.51,
+      "gumbel_tau": 0.3067,
+      "image_masks_loss": 0.0183,
+      "llm_learning_rate": 2.5128590026969474e-06,
+      "loss": 0.7675,
+      "output_text_mask_loss": 0.0032,
+      "predictor_learning_rate": 0.00025128590026969475,
+      "step": 5011
+    },
+    {
+      "epoch": 0.51,
+      "gumbel_tau": 0.3066,
+      "image_masks_loss": 0.0154,
+      "llm_learning_rate": 2.512029396243254e-06,
+      "loss": 0.7254,
+      "output_text_mask_loss": 0.0081,
+      "predictor_learning_rate": 0.0002512029396243254,
+      "step": 5012
+    },
+    {
+      "epoch": 0.51,
+      "gumbel_tau": 0.3065,
+      "image_masks_loss": 0.0089,
+      "llm_learning_rate": 2.5111997884648566e-06,
+      "loss": 0.6877,
+      "output_text_mask_loss": 0.0012,
+      "predictor_learning_rate": 0.00025111997884648564,
+      "step": 5013
+    },
+    {
+      "epoch": 0.51,
+      "gumbel_tau": 0.3065,
+      "image_masks_loss": 0.009,
+      "llm_learning_rate": 2.510370179453113e-06,
+      "loss": 0.687,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.0002510370179453113,
+      "step": 5014
+    },
+    {
+      "epoch": 0.51,
+      "gumbel_tau": 0.3064,
+      "image_masks_loss": 0.0076,
+      "llm_learning_rate": 2.5095405692993825e-06,
+      "loss": 0.7121,
+      "output_text_mask_loss": 0.0181,
+      "predictor_learning_rate": 0.0002509540569299382,
+      "step": 5015
+    },
+    {
+      "epoch": 0.51,
+      "gumbel_tau": 0.3063,
+      "image_masks_loss": 0.0226,
+      "llm_learning_rate": 2.5087109580950235e-06,
+      "loss": 0.7598,
+      "output_text_mask_loss": 0.0139,
+      "predictor_learning_rate": 0.00025087109580950237,
+      "step": 5016
+    },
+    {
+      "epoch": 0.51,
+      "gumbel_tau": 0.3062,
+      "image_masks_loss": 0.0149,
+      "llm_learning_rate": 2.5078813459313938e-06,
+      "loss": 0.7906,
+      "output_text_mask_loss": 0.0774,
+      "predictor_learning_rate": 0.00025078813459313937,
+      "step": 5017
+    },
+    {
+      "epoch": 0.51,
+      "gumbel_tau": 0.3062,
+      "image_masks_loss": 0.0254,
+      "llm_learning_rate": 2.5070517328998534e-06,
+      "loss": 0.6504,
+      "output_text_mask_loss": 0.074,
+      "predictor_learning_rate": 0.0002507051732899853,
+      "step": 5018
+    },
+    {
+      "epoch": 0.51,
+      "gumbel_tau": 0.3061,
+      "image_masks_loss": 0.0079,
+      "llm_learning_rate": 2.5062221190917596e-06,
+      "loss": 0.716,
+      "output_text_mask_loss": 0.006,
+      "predictor_learning_rate": 0.00025062221190917596,
+      "step": 5019
+    },
+    {
+      "epoch": 0.51,
+      "gumbel_tau": 0.306,
+      "image_masks_loss": 0.0048,
+      "llm_learning_rate": 2.505392504598473e-06,
+      "loss": 0.6383,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 0.0002505392504598473,
+      "step": 5020
+    },
+    {
+      "epoch": 0.51,
+      "gumbel_tau": 0.306,
+      "image_masks_loss": 0.0117,
+      "llm_learning_rate": 2.5045628895113518e-06,
+      "loss": 0.7364,
+      "output_text_mask_loss": 0.0075,
+      "predictor_learning_rate": 0.00025045628895113513,
+      "step": 5021
+    },
+    {
+      "epoch": 0.51,
+      "gumbel_tau": 0.3059,
+      "image_masks_loss": 0.0201,
+      "llm_learning_rate": 2.5037332739217545e-06,
+      "loss": 0.8152,
+      "output_text_mask_loss": 0.0105,
+      "predictor_learning_rate": 0.00025037332739217546,
+      "step": 5022
+    },
+    {
+      "epoch": 0.51,
+      "gumbel_tau": 0.3058,
+      "image_masks_loss": 0.0169,
+      "llm_learning_rate": 2.5029036579210414e-06,
+      "loss": 0.6681,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00025029036579210414,
+      "step": 5023
+    },
+    {
+      "epoch": 0.51,
+      "gumbel_tau": 0.3057,
+      "image_masks_loss": 0.017,
+      "llm_learning_rate": 2.502074041600571e-06,
+      "loss": 0.8089,
+      "output_text_mask_loss": 0.0012,
+      "predictor_learning_rate": 0.00025020740416005705,
+      "step": 5024
+    },
+    {
+      "epoch": 0.51,
+      "gumbel_tau": 0.3057,
+      "image_masks_loss": 0.0079,
+      "llm_learning_rate": 2.5012444250517016e-06,
+      "loss": 0.7163,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 0.00025012444250517015,
+      "step": 5025
+    },
+    {
+      "epoch": 0.51,
+      "gumbel_tau": 0.3056,
+      "image_masks_loss": 0.0203,
+      "llm_learning_rate": 2.5004148083657943e-06,
+      "loss": 0.7464,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00025004148083657943,
+      "step": 5026
+    },
+    {
+      "epoch": 0.52,
+      "gumbel_tau": 0.3055,
+      "image_masks_loss": 0.0034,
+      "llm_learning_rate": 2.4995851916342074e-06,
+      "loss": 0.708,
+      "output_text_mask_loss": 0.0081,
+      "predictor_learning_rate": 0.0002499585191634207,
+      "step": 5027
+    },
+    {
+      "epoch": 0.52,
+      "gumbel_tau": 0.3054,
+      "image_masks_loss": 0.0074,
+      "llm_learning_rate": 2.4987555749482992e-06,
+      "loss": 0.7244,
+      "output_text_mask_loss": 0.037,
+      "predictor_learning_rate": 0.0002498755574948299,
+      "step": 5028
+    },
+    {
+      "epoch": 0.52,
+      "gumbel_tau": 0.3054,
+      "image_masks_loss": 0.0098,
+      "llm_learning_rate": 2.4979259583994304e-06,
+      "loss": 0.7479,
+      "output_text_mask_loss": 0.0034,
+      "predictor_learning_rate": 0.00024979259583994307,
+      "step": 5029
+    },
+    {
+      "epoch": 0.52,
+      "gumbel_tau": 0.3053,
+      "image_masks_loss": 0.0212,
+      "llm_learning_rate": 2.49709634207896e-06,
+      "loss": 0.7351,
+      "output_text_mask_loss": 0.0024,
+      "predictor_learning_rate": 0.000249709634207896,
+      "step": 5030
+    },
+    {
+      "epoch": 0.52,
+      "gumbel_tau": 0.3052,
+      "image_masks_loss": 0.0129,
+      "llm_learning_rate": 2.4962667260782463e-06,
+      "loss": 0.7338,
+      "output_text_mask_loss": 0.005,
+      "predictor_learning_rate": 0.0002496266726078246,
+      "step": 5031
+    },
+    {
+      "epoch": 0.52,
+      "gumbel_tau": 0.3052,
+      "image_masks_loss": 0.0217,
+      "llm_learning_rate": 2.495437110488649e-06,
+      "loss": 0.763,
+      "output_text_mask_loss": 0.017,
+      "predictor_learning_rate": 0.0002495437110488649,
+      "step": 5032
+    },
+    {
+      "epoch": 0.52,
+      "gumbel_tau": 0.3051,
+      "image_masks_loss": 0.0123,
+      "llm_learning_rate": 2.4946074954015275e-06,
+      "loss": 0.7536,
+      "output_text_mask_loss": 0.0039,
+      "predictor_learning_rate": 0.0002494607495401527,
+      "step": 5033
+    },
+    {
+      "epoch": 0.52,
+      "gumbel_tau": 0.305,
+      "image_masks_loss": 0.0148,
+      "llm_learning_rate": 2.4937778809082408e-06,
+      "loss": 0.8001,
+      "output_text_mask_loss": 0.0132,
+      "predictor_learning_rate": 0.00024937778809082405,
+      "step": 5034
+    },
+    {
+      "epoch": 0.52,
+      "gumbel_tau": 0.3049,
+      "image_masks_loss": 0.01,
+      "llm_learning_rate": 2.4929482671001474e-06,
+      "loss": 0.696,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00024929482671001474,
+      "step": 5035
+    },
+    {
+      "epoch": 0.52,
+      "gumbel_tau": 0.3049,
+      "image_masks_loss": 0.0114,
+      "llm_learning_rate": 2.4921186540686066e-06,
+      "loss": 0.6863,
+      "output_text_mask_loss": 0.0374,
+      "predictor_learning_rate": 0.00024921186540686064,
+      "step": 5036
+    },
+    {
+      "epoch": 0.52,
+      "gumbel_tau": 0.3048,
+      "image_masks_loss": 0.0121,
+      "llm_learning_rate": 2.4912890419049773e-06,
+      "loss": 0.7967,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.0002491289041904977,
+      "step": 5037
+    },
+    {
+      "epoch": 0.52,
+      "gumbel_tau": 0.3047,
+      "image_masks_loss": 0.006,
+      "llm_learning_rate": 2.4904594307006184e-06,
+      "loss": 0.7398,
+      "output_text_mask_loss": 0.0088,
+      "predictor_learning_rate": 0.0002490459430700618,
+      "step": 5038
+    },
+    {
+      "epoch": 0.52,
+      "gumbel_tau": 0.3047,
+      "image_masks_loss": 0.0234,
+      "llm_learning_rate": 2.4896298205468873e-06,
+      "loss": 0.739,
+      "output_text_mask_loss": 0.0204,
+      "predictor_learning_rate": 0.0002489629820546887,
+      "step": 5039
+    },
+    {
+      "epoch": 0.52,
+      "gumbel_tau": 0.3046,
+      "image_masks_loss": 0.0115,
+      "llm_learning_rate": 2.4888002115351442e-06,
+      "loss": 0.8094,
+      "output_text_mask_loss": 0.0014,
+      "predictor_learning_rate": 0.0002488800211535144,
+      "step": 5040
+    },
+    {
+      "epoch": 0.52,
+      "gumbel_tau": 0.3045,
+      "image_masks_loss": 0.0178,
+      "llm_learning_rate": 2.4879706037567468e-06,
+      "loss": 0.7316,
+      "output_text_mask_loss": 0.0269,
+      "predictor_learning_rate": 0.00024879706037567466,
+      "step": 5041
+    },
+    {
+      "epoch": 0.52,
+      "gumbel_tau": 0.3044,
+      "image_masks_loss": 0.0125,
+      "llm_learning_rate": 2.4871409973030525e-06,
+      "loss": 0.7237,
+      "output_text_mask_loss": 0.0065,
+      "predictor_learning_rate": 0.0002487140997303052,
+      "step": 5042
+    },
+    {
+      "epoch": 0.52,
+      "gumbel_tau": 0.3044,
+      "image_masks_loss": 0.0176,
+      "llm_learning_rate": 2.4863113922654212e-06,
+      "loss": 0.6689,
+      "output_text_mask_loss": 0.0233,
+      "predictor_learning_rate": 0.0002486311392265421,
+      "step": 5043
+    },
+    {
+      "epoch": 0.52,
+      "gumbel_tau": 0.3043,
+      "image_masks_loss": 0.0186,
+      "llm_learning_rate": 2.4854817887352095e-06,
+      "loss": 0.7483,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.00024854817887352094,
+      "step": 5044
+    },
+    {
+      "epoch": 0.52,
+      "gumbel_tau": 0.3042,
+      "image_masks_loss": 0.0065,
+      "llm_learning_rate": 2.4846521868037755e-06,
+      "loss": 0.8172,
+      "output_text_mask_loss": 0.024,
+      "predictor_learning_rate": 0.00024846521868037753,
+      "step": 5045
+    },
+    {
+      "epoch": 0.52,
+      "gumbel_tau": 0.3042,
+      "image_masks_loss": 0.0161,
+      "llm_learning_rate": 2.483822586562478e-06,
+      "loss": 0.8194,
+      "output_text_mask_loss": 0.006,
+      "predictor_learning_rate": 0.0002483822586562478,
+      "step": 5046
+    },
+    {
+      "epoch": 0.52,
+      "gumbel_tau": 0.3041,
+      "image_masks_loss": 0.0109,
+      "llm_learning_rate": 2.482992988102672e-06,
+      "loss": 0.7349,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00024829929881026717,
+      "step": 5047
+    },
+    {
+      "epoch": 0.52,
+      "gumbel_tau": 0.304,
+      "image_masks_loss": 0.0248,
+      "llm_learning_rate": 2.482163391515716e-06,
+      "loss": 0.8564,
+      "output_text_mask_loss": 0.0006,
+      "predictor_learning_rate": 0.0002482163391515716,
+      "step": 5048
+    },
+    {
+      "epoch": 0.52,
+      "gumbel_tau": 0.3039,
+      "image_masks_loss": 0.0043,
+      "llm_learning_rate": 2.481333796892967e-06,
+      "loss": 0.7279,
+      "output_text_mask_loss": 0.0034,
+      "predictor_learning_rate": 0.0002481333796892967,
+      "step": 5049
+    },
+    {
+      "epoch": 0.52,
+      "gumbel_tau": 0.3039,
+      "image_masks_loss": 0.0085,
+      "llm_learning_rate": 2.480504204325783e-06,
+      "loss": 0.755,
+      "output_text_mask_loss": 0.0145,
+      "predictor_learning_rate": 0.0002480504204325783,
+      "step": 5050
+    },
+    {
+      "epoch": 0.52,
+      "gumbel_tau": 0.3038,
+      "image_masks_loss": 0.0256,
+      "llm_learning_rate": 2.479674613905518e-06,
+      "loss": 0.8003,
+      "output_text_mask_loss": 0.0065,
+      "predictor_learning_rate": 0.0002479674613905518,
+      "step": 5051
+    },
+    {
+      "epoch": 0.52,
+      "gumbel_tau": 0.3037,
+      "image_masks_loss": 0.0118,
+      "llm_learning_rate": 2.4788450257235305e-06,
+      "loss": 0.8115,
+      "output_text_mask_loss": 0.017,
+      "predictor_learning_rate": 0.00024788450257235306,
+      "step": 5052
+    },
+    {
+      "epoch": 0.52,
+      "gumbel_tau": 0.3037,
+      "image_masks_loss": 0.0224,
+      "llm_learning_rate": 2.478015439871176e-06,
+      "loss": 0.7535,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.0002478015439871176,
+      "step": 5053
+    },
+    {
+      "epoch": 0.52,
+      "gumbel_tau": 0.3036,
+      "image_masks_loss": 0.0148,
+      "llm_learning_rate": 2.47718585643981e-06,
+      "loss": 0.8668,
+      "output_text_mask_loss": 0.0507,
+      "predictor_learning_rate": 0.000247718585643981,
+      "step": 5054
+    },
+    {
+      "epoch": 0.52,
+      "gumbel_tau": 0.3035,
+      "image_masks_loss": 0.0203,
+      "llm_learning_rate": 2.476356275520788e-06,
+      "loss": 0.8573,
+      "output_text_mask_loss": 0.0204,
+      "predictor_learning_rate": 0.0002476356275520788,
+      "step": 5055
+    },
+    {
+      "epoch": 0.52,
+      "gumbel_tau": 0.3034,
+      "image_masks_loss": 0.0074,
+      "llm_learning_rate": 2.475526697205466e-06,
+      "loss": 0.7188,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.0002475526697205466,
+      "step": 5056
+    },
+    {
+      "epoch": 0.52,
+      "gumbel_tau": 0.3034,
+      "image_masks_loss": 0.0101,
+      "llm_learning_rate": 2.474697121585199e-06,
+      "loss": 0.7164,
+      "output_text_mask_loss": 0.0034,
+      "predictor_learning_rate": 0.00024746971215851987,
+      "step": 5057
+    },
+    {
+      "epoch": 0.52,
+      "gumbel_tau": 0.3033,
+      "image_masks_loss": 0.0078,
+      "llm_learning_rate": 2.4738675487513404e-06,
+      "loss": 0.7388,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 0.000247386754875134,
+      "step": 5058
+    },
+    {
+      "epoch": 0.52,
+      "gumbel_tau": 0.3032,
+      "image_masks_loss": 0.0125,
+      "llm_learning_rate": 2.4730379787952464e-06,
+      "loss": 0.6948,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.00024730379787952464,
+      "step": 5059
+    },
+    {
+      "epoch": 0.52,
+      "gumbel_tau": 0.3032,
+      "image_masks_loss": 0.0166,
+      "llm_learning_rate": 2.4722084118082697e-06,
+      "loss": 0.6639,
+      "output_text_mask_loss": 0.0185,
+      "predictor_learning_rate": 0.00024722084118082696,
+      "step": 5060
+    },
+    {
+      "epoch": 0.52,
+      "gumbel_tau": 0.3031,
+      "image_masks_loss": 0.0089,
+      "llm_learning_rate": 2.471378847881765e-06,
+      "loss": 0.7964,
+      "output_text_mask_loss": 0.0307,
+      "predictor_learning_rate": 0.00024713788478817646,
+      "step": 5061
+    },
+    {
+      "epoch": 0.52,
+      "gumbel_tau": 0.303,
+      "image_masks_loss": 0.0098,
+      "llm_learning_rate": 2.4705492871070845e-06,
+      "loss": 0.8133,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.00024705492871070846,
+      "step": 5062
+    },
+    {
+      "epoch": 0.52,
+      "gumbel_tau": 0.3029,
+      "image_masks_loss": 0.0086,
+      "llm_learning_rate": 2.4697197295755827e-06,
+      "loss": 0.7337,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.0002469719729575583,
+      "step": 5063
+    },
+    {
+      "epoch": 0.52,
+      "gumbel_tau": 0.3029,
+      "image_masks_loss": 0.0068,
+      "llm_learning_rate": 2.4688901753786125e-06,
+      "loss": 0.6794,
+      "output_text_mask_loss": 0.0107,
+      "predictor_learning_rate": 0.0002468890175378613,
+      "step": 5064
+    },
+    {
+      "epoch": 0.52,
+      "gumbel_tau": 0.3028,
+      "image_masks_loss": 0.0166,
+      "llm_learning_rate": 2.468060624607525e-06,
+      "loss": 0.6984,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00024680606246075247,
+      "step": 5065
+    },
+    {
+      "epoch": 0.52,
+      "gumbel_tau": 0.3027,
+      "image_masks_loss": 0.0083,
+      "llm_learning_rate": 2.4672310773536725e-06,
+      "loss": 0.7672,
+      "output_text_mask_loss": 0.0168,
+      "predictor_learning_rate": 0.00024672310773536725,
+      "step": 5066
+    },
+    {
+      "epoch": 0.52,
+      "gumbel_tau": 0.3027,
+      "image_masks_loss": 0.0123,
+      "llm_learning_rate": 2.4664015337084064e-06,
+      "loss": 0.7445,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.0002466401533708406,
+      "step": 5067
+    },
+    {
+      "epoch": 0.52,
+      "gumbel_tau": 0.3026,
+      "image_masks_loss": 0.0081,
+      "llm_learning_rate": 2.4655719937630785e-06,
+      "loss": 0.7629,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.00024655719937630784,
+      "step": 5068
+    },
+    {
+      "epoch": 0.52,
+      "gumbel_tau": 0.3025,
+      "image_masks_loss": 0.0175,
+      "llm_learning_rate": 2.464742457609039e-06,
+      "loss": 0.7942,
+      "output_text_mask_loss": 0.0185,
+      "predictor_learning_rate": 0.0002464742457609039,
+      "step": 5069
+    },
+    {
+      "epoch": 0.52,
+      "gumbel_tau": 0.3024,
+      "image_masks_loss": 0.0151,
+      "llm_learning_rate": 2.463912925337639e-06,
+      "loss": 0.7257,
+      "output_text_mask_loss": 0.0034,
+      "predictor_learning_rate": 0.0002463912925337639,
+      "step": 5070
+    },
+    {
+      "epoch": 0.52,
+      "gumbel_tau": 0.3024,
+      "image_masks_loss": 0.0099,
+      "llm_learning_rate": 2.463083397040228e-06,
+      "loss": 0.826,
+      "output_text_mask_loss": 0.001,
+      "predictor_learning_rate": 0.00024630833970402277,
+      "step": 5071
+    },
+    {
+      "epoch": 0.52,
+      "gumbel_tau": 0.3023,
+      "image_masks_loss": 0.0138,
+      "llm_learning_rate": 2.4622538728081555e-06,
+      "loss": 0.7585,
+      "output_text_mask_loss": 0.0058,
+      "predictor_learning_rate": 0.0002462253872808155,
+      "step": 5072
+    },
+    {
+      "epoch": 0.52,
+      "gumbel_tau": 0.3022,
+      "image_masks_loss": 0.0028,
+      "llm_learning_rate": 2.4614243527327704e-06,
+      "loss": 0.694,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00024614243527327703,
+      "step": 5073
+    },
+    {
+      "epoch": 0.52,
+      "gumbel_tau": 0.3022,
+      "image_masks_loss": 0.0121,
+      "llm_learning_rate": 2.4605948369054216e-06,
+      "loss": 0.7518,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0002460594836905421,
+      "step": 5074
+    },
+    {
+      "epoch": 0.52,
+      "gumbel_tau": 0.3021,
+      "image_masks_loss": 0.0226,
+      "llm_learning_rate": 2.4597653254174573e-06,
+      "loss": 0.7481,
+      "output_text_mask_loss": 0.0058,
+      "predictor_learning_rate": 0.0002459765325417457,
+      "step": 5075
+    },
+    {
+      "epoch": 0.52,
+      "gumbel_tau": 0.302,
+      "image_masks_loss": 0.0175,
+      "llm_learning_rate": 2.458935818360225e-06,
+      "loss": 0.7827,
+      "output_text_mask_loss": 0.0137,
+      "predictor_learning_rate": 0.00024589358183602246,
+      "step": 5076
+    },
+    {
+      "epoch": 0.52,
+      "gumbel_tau": 0.3019,
+      "image_masks_loss": 0.0119,
+      "llm_learning_rate": 2.458106315825072e-06,
+      "loss": 0.7668,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.00024581063158250717,
+      "step": 5077
+    },
+    {
+      "epoch": 0.52,
+      "gumbel_tau": 0.3019,
+      "image_masks_loss": 0.0158,
+      "llm_learning_rate": 2.457276817903344e-06,
+      "loss": 0.7854,
+      "output_text_mask_loss": 0.0077,
+      "predictor_learning_rate": 0.00024572768179033437,
+      "step": 5078
+    },
+    {
+      "epoch": 0.52,
+      "gumbel_tau": 0.3018,
+      "image_masks_loss": 0.0176,
+      "llm_learning_rate": 2.456447324686389e-06,
+      "loss": 0.6904,
+      "output_text_mask_loss": 0.0313,
+      "predictor_learning_rate": 0.00024564473246863885,
+      "step": 5079
+    },
+    {
+      "epoch": 0.52,
+      "gumbel_tau": 0.3017,
+      "image_masks_loss": 0.0114,
+      "llm_learning_rate": 2.455617836265551e-06,
+      "loss": 0.878,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.0002455617836265551,
+      "step": 5080
+    },
+    {
+      "epoch": 0.52,
+      "gumbel_tau": 0.3017,
+      "image_masks_loss": 0.0081,
+      "llm_learning_rate": 2.454788352732175e-06,
+      "loss": 0.6808,
+      "output_text_mask_loss": 0.0025,
+      "predictor_learning_rate": 0.0002454788352732175,
+      "step": 5081
+    },
+    {
+      "epoch": 0.52,
+      "gumbel_tau": 0.3016,
+      "image_masks_loss": 0.0142,
+      "llm_learning_rate": 2.453958874177608e-06,
+      "loss": 0.7737,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.0002453958874177608,
+      "step": 5082
+    },
+    {
+      "epoch": 0.52,
+      "gumbel_tau": 0.3015,
+      "image_masks_loss": 0.0288,
+      "llm_learning_rate": 2.453129400693192e-06,
+      "loss": 0.7662,
+      "output_text_mask_loss": 0.001,
+      "predictor_learning_rate": 0.0002453129400693192,
+      "step": 5083
+    },
+    {
+      "epoch": 0.52,
+      "gumbel_tau": 0.3014,
+      "image_masks_loss": 0.0104,
+      "llm_learning_rate": 2.4522999323702703e-06,
+      "loss": 0.7069,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00024522999323702704,
+      "step": 5084
+    },
+    {
+      "epoch": 0.52,
+      "gumbel_tau": 0.3014,
+      "image_masks_loss": 0.0229,
+      "llm_learning_rate": 2.4514704693001864e-06,
+      "loss": 0.7937,
+      "output_text_mask_loss": 0.0301,
+      "predictor_learning_rate": 0.0002451470469300186,
+      "step": 5085
+    },
+    {
+      "epoch": 0.52,
+      "gumbel_tau": 0.3013,
+      "image_masks_loss": 0.0233,
+      "llm_learning_rate": 2.4506410115742825e-06,
+      "loss": 0.7998,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.00024506410115742827,
+      "step": 5086
+    },
+    {
+      "epoch": 0.52,
+      "gumbel_tau": 0.3012,
+      "image_masks_loss": 0.0134,
+      "llm_learning_rate": 2.4498115592839e-06,
+      "loss": 0.6955,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.00024498115592839,
+      "step": 5087
+    },
+    {
+      "epoch": 0.52,
+      "gumbel_tau": 0.3012,
+      "image_masks_loss": 0.0137,
+      "llm_learning_rate": 2.4489821125203807e-06,
+      "loss": 0.8025,
+      "output_text_mask_loss": 0.0345,
+      "predictor_learning_rate": 0.00024489821125203806,
+      "step": 5088
+    },
+    {
+      "epoch": 0.52,
+      "gumbel_tau": 0.3011,
+      "image_masks_loss": 0.0104,
+      "llm_learning_rate": 2.448152671375065e-06,
+      "loss": 0.7044,
+      "output_text_mask_loss": 0.0076,
+      "predictor_learning_rate": 0.0002448152671375065,
+      "step": 5089
+    },
+    {
+      "epoch": 0.52,
+      "gumbel_tau": 0.301,
+      "image_masks_loss": 0.0041,
+      "llm_learning_rate": 2.447323235939292e-06,
+      "loss": 0.7636,
+      "output_text_mask_loss": 0.053,
+      "predictor_learning_rate": 0.0002447323235939292,
+      "step": 5090
+    },
+    {
+      "epoch": 0.52,
+      "gumbel_tau": 0.3009,
+      "image_masks_loss": 0.0135,
+      "llm_learning_rate": 2.446493806304402e-06,
+      "loss": 0.735,
+      "output_text_mask_loss": 0.0176,
+      "predictor_learning_rate": 0.0002446493806304402,
+      "step": 5091
+    },
+    {
+      "epoch": 0.52,
+      "gumbel_tau": 0.3009,
+      "image_masks_loss": 0.006,
+      "llm_learning_rate": 2.445664382561733e-06,
+      "loss": 0.7112,
+      "output_text_mask_loss": 0.0244,
+      "predictor_learning_rate": 0.0002445664382561733,
+      "step": 5092
+    },
+    {
+      "epoch": 0.52,
+      "gumbel_tau": 0.3008,
+      "image_masks_loss": 0.0073,
+      "llm_learning_rate": 2.444834964802624e-06,
+      "loss": 0.7663,
+      "output_text_mask_loss": 0.005,
+      "predictor_learning_rate": 0.0002444834964802624,
+      "step": 5093
+    },
+    {
+      "epoch": 0.52,
+      "gumbel_tau": 0.3007,
+      "image_masks_loss": 0.0108,
+      "llm_learning_rate": 2.4440055531184105e-06,
+      "loss": 0.7675,
+      "output_text_mask_loss": 0.0227,
+      "predictor_learning_rate": 0.00024440055531184103,
+      "step": 5094
+    },
+    {
+      "epoch": 0.52,
+      "gumbel_tau": 0.3007,
+      "image_masks_loss": 0.0134,
+      "llm_learning_rate": 2.4431761476004305e-06,
+      "loss": 0.7634,
+      "output_text_mask_loss": 0.0334,
+      "predictor_learning_rate": 0.00024431761476004303,
+      "step": 5095
+    },
+    {
+      "epoch": 0.52,
+      "gumbel_tau": 0.3006,
+      "image_masks_loss": 0.0214,
+      "llm_learning_rate": 2.4423467483400203e-06,
+      "loss": 0.7669,
+      "output_text_mask_loss": 0.0039,
+      "predictor_learning_rate": 0.000244234674834002,
+      "step": 5096
+    },
+    {
+      "epoch": 0.52,
+      "gumbel_tau": 0.3005,
+      "image_masks_loss": 0.0145,
+      "llm_learning_rate": 2.4415173554285138e-06,
+      "loss": 0.7768,
+      "output_text_mask_loss": 0.0423,
+      "predictor_learning_rate": 0.00024415173554285136,
+      "step": 5097
+    },
+    {
+      "epoch": 0.52,
+      "gumbel_tau": 0.3004,
+      "image_masks_loss": 0.0034,
+      "llm_learning_rate": 2.440687968957247e-06,
+      "loss": 0.7647,
+      "output_text_mask_loss": 0.0374,
+      "predictor_learning_rate": 0.00024406879689572468,
+      "step": 5098
+    },
+    {
+      "epoch": 0.52,
+      "gumbel_tau": 0.3004,
+      "image_masks_loss": 0.0265,
+      "llm_learning_rate": 2.439858589017553e-06,
+      "loss": 0.6709,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.00024398585890175525,
+      "step": 5099
+    },
+    {
+      "epoch": 0.52,
+      "gumbel_tau": 0.3003,
+      "image_masks_loss": 0.0075,
+      "llm_learning_rate": 2.439029215700765e-06,
+      "loss": 0.7658,
+      "output_text_mask_loss": 0.0193,
+      "predictor_learning_rate": 0.0002439029215700765,
+      "step": 5100
+    },
+    {
+      "epoch": 0.52,
+      "gumbel_tau": 0.3002,
+      "image_masks_loss": 0.0107,
+      "llm_learning_rate": 2.4381998490982164e-06,
+      "loss": 0.6709,
+      "output_text_mask_loss": 0.0292,
+      "predictor_learning_rate": 0.00024381998490982163,
+      "step": 5101
+    },
+    {
+      "epoch": 0.52,
+      "gumbel_tau": 0.3002,
+      "image_masks_loss": 0.0109,
+      "llm_learning_rate": 2.4373704893012373e-06,
+      "loss": 0.6957,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0002437370489301237,
+      "step": 5102
+    },
+    {
+      "epoch": 0.52,
+      "gumbel_tau": 0.3001,
+      "image_masks_loss": 0.0055,
+      "llm_learning_rate": 2.436541136401159e-06,
+      "loss": 0.8017,
+      "output_text_mask_loss": 0.0134,
+      "predictor_learning_rate": 0.0002436541136401159,
+      "step": 5103
+    },
+    {
+      "epoch": 0.52,
+      "gumbel_tau": 0.3,
+      "image_masks_loss": 0.0064,
+      "llm_learning_rate": 2.4357117904893127e-06,
+      "loss": 0.7339,
+      "output_text_mask_loss": 0.0324,
+      "predictor_learning_rate": 0.00024357117904893124,
+      "step": 5104
+    },
+    {
+      "epoch": 0.52,
+      "gumbel_tau": 0.3,
+      "image_masks_loss": 0.0113,
+      "llm_learning_rate": 2.434882451657027e-06,
+      "loss": 0.7099,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0002434882451657027,
+      "step": 5105
+    },
+    {
+      "epoch": 0.52,
+      "gumbel_tau": 0.2999,
+      "image_masks_loss": 0.0093,
+      "llm_learning_rate": 2.4340531199956305e-06,
+      "loss": 0.7471,
+      "output_text_mask_loss": 0.0603,
+      "predictor_learning_rate": 0.00024340531199956304,
+      "step": 5106
+    },
+    {
+      "epoch": 0.52,
+      "gumbel_tau": 0.2998,
+      "image_masks_loss": 0.0141,
+      "llm_learning_rate": 2.433223795596452e-06,
+      "loss": 0.844,
+      "output_text_mask_loss": 0.0568,
+      "predictor_learning_rate": 0.00024332237955964516,
+      "step": 5107
+    },
+    {
+      "epoch": 0.52,
+      "gumbel_tau": 0.2997,
+      "image_masks_loss": 0.0044,
+      "llm_learning_rate": 2.4323944785508173e-06,
+      "loss": 0.763,
+      "output_text_mask_loss": 0.0042,
+      "predictor_learning_rate": 0.00024323944785508173,
+      "step": 5108
+    },
+    {
+      "epoch": 0.52,
+      "gumbel_tau": 0.2997,
+      "image_masks_loss": 0.0051,
+      "llm_learning_rate": 2.4315651689500538e-06,
+      "loss": 0.7287,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 0.00024315651689500538,
+      "step": 5109
+    },
+    {
+      "epoch": 0.52,
+      "gumbel_tau": 0.2996,
+      "image_masks_loss": 0.0418,
+      "llm_learning_rate": 2.4307358668854854e-06,
+      "loss": 0.7357,
+      "output_text_mask_loss": 0.0055,
+      "predictor_learning_rate": 0.00024307358668854853,
+      "step": 5110
+    },
+    {
+      "epoch": 0.52,
+      "gumbel_tau": 0.2995,
+      "image_masks_loss": 0.0069,
+      "llm_learning_rate": 2.4299065724484382e-06,
+      "loss": 0.8582,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.0002429906572448438,
+      "step": 5111
+    },
+    {
+      "epoch": 0.52,
+      "gumbel_tau": 0.2995,
+      "image_masks_loss": 0.0082,
+      "llm_learning_rate": 2.4290772857302354e-06,
+      "loss": 0.762,
+      "output_text_mask_loss": 0.0263,
+      "predictor_learning_rate": 0.00024290772857302355,
+      "step": 5112
+    },
+    {
+      "epoch": 0.52,
+      "gumbel_tau": 0.2994,
+      "image_masks_loss": 0.0076,
+      "llm_learning_rate": 2.4282480068222e-06,
+      "loss": 0.7713,
+      "output_text_mask_loss": 0.0212,
+      "predictor_learning_rate": 0.00024282480068221995,
+      "step": 5113
+    },
+    {
+      "epoch": 0.52,
+      "gumbel_tau": 0.2993,
+      "image_masks_loss": 0.0114,
+      "llm_learning_rate": 2.4274187358156535e-06,
+      "loss": 0.7252,
+      "output_text_mask_loss": 0.005,
+      "predictor_learning_rate": 0.00024274187358156533,
+      "step": 5114
+    },
+    {
+      "epoch": 0.52,
+      "gumbel_tau": 0.2992,
+      "image_masks_loss": 0.0159,
+      "llm_learning_rate": 2.426589472801918e-06,
+      "loss": 0.8142,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.00024265894728019177,
+      "step": 5115
+    },
+    {
+      "epoch": 0.52,
+      "gumbel_tau": 0.2992,
+      "image_masks_loss": 0.0159,
+      "llm_learning_rate": 2.4257602178723128e-06,
+      "loss": 0.735,
+      "output_text_mask_loss": 0.0004,
+      "predictor_learning_rate": 0.00024257602178723127,
+      "step": 5116
+    },
+    {
+      "epoch": 0.52,
+      "gumbel_tau": 0.2991,
+      "image_masks_loss": 0.0117,
+      "llm_learning_rate": 2.4249309711181578e-06,
+      "loss": 0.8303,
+      "output_text_mask_loss": 0.016,
+      "predictor_learning_rate": 0.00024249309711181575,
+      "step": 5117
+    },
+    {
+      "epoch": 0.52,
+      "gumbel_tau": 0.299,
+      "image_masks_loss": 0.0122,
+      "llm_learning_rate": 2.4241017326307716e-06,
+      "loss": 0.8375,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 0.00024241017326307712,
+      "step": 5118
+    },
+    {
+      "epoch": 0.52,
+      "gumbel_tau": 0.299,
+      "image_masks_loss": 0.019,
+      "llm_learning_rate": 2.4232725025014716e-06,
+      "loss": 0.7781,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.00024232725025014716,
+      "step": 5119
+    },
+    {
+      "epoch": 0.52,
+      "gumbel_tau": 0.2989,
+      "image_masks_loss": 0.0153,
+      "llm_learning_rate": 2.422443280821575e-06,
+      "loss": 0.6354,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.00024224432808215752,
+      "step": 5120
+    },
+    {
+      "epoch": 0.52,
+      "gumbel_tau": 0.2988,
+      "image_masks_loss": 0.0113,
+      "llm_learning_rate": 2.421614067682397e-06,
+      "loss": 0.6048,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00024216140676823968,
+      "step": 5121
+    },
+    {
+      "epoch": 0.52,
+      "gumbel_tau": 0.2987,
+      "image_masks_loss": 0.0029,
+      "llm_learning_rate": 2.4207848631752516e-06,
+      "loss": 0.8468,
+      "output_text_mask_loss": 0.0055,
+      "predictor_learning_rate": 0.00024207848631752516,
+      "step": 5122
+    },
+    {
+      "epoch": 0.52,
+      "gumbel_tau": 0.2987,
+      "image_masks_loss": 0.0259,
+      "llm_learning_rate": 2.4199556673914545e-06,
+      "loss": 0.7375,
+      "output_text_mask_loss": 0.0063,
+      "predictor_learning_rate": 0.00024199556673914544,
+      "step": 5123
+    },
+    {
+      "epoch": 0.53,
+      "gumbel_tau": 0.2986,
+      "image_masks_loss": 0.0181,
+      "llm_learning_rate": 2.4191264804223173e-06,
+      "loss": 0.7495,
+      "output_text_mask_loss": 0.0086,
+      "predictor_learning_rate": 0.0002419126480422317,
+      "step": 5124
+    },
+    {
+      "epoch": 0.53,
+      "gumbel_tau": 0.2985,
+      "image_masks_loss": 0.0112,
+      "llm_learning_rate": 2.418297302359152e-06,
+      "loss": 0.7334,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0002418297302359152,
+      "step": 5125
+    },
+    {
+      "epoch": 0.53,
+      "gumbel_tau": 0.2985,
+      "image_masks_loss": 0.0151,
+      "llm_learning_rate": 2.4174681332932696e-06,
+      "loss": 0.772,
+      "output_text_mask_loss": 0.0086,
+      "predictor_learning_rate": 0.00024174681332932694,
+      "step": 5126
+    },
+    {
+      "epoch": 0.53,
+      "gumbel_tau": 0.2984,
+      "image_masks_loss": 0.0024,
+      "llm_learning_rate": 2.4166389733159802e-06,
+      "loss": 0.797,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.00024166389733159803,
+      "step": 5127
+    },
+    {
+      "epoch": 0.53,
+      "gumbel_tau": 0.2983,
+      "image_masks_loss": 0.0071,
+      "llm_learning_rate": 2.4158098225185935e-06,
+      "loss": 0.7136,
+      "output_text_mask_loss": 0.017,
+      "predictor_learning_rate": 0.0002415809822518593,
+      "step": 5128
+    },
+    {
+      "epoch": 0.53,
+      "gumbel_tau": 0.2983,
+      "image_masks_loss": 0.0147,
+      "llm_learning_rate": 2.4149806809924155e-06,
+      "loss": 0.6813,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.0002414980680992415,
+      "step": 5129
+    },
+    {
+      "epoch": 0.53,
+      "gumbel_tau": 0.2982,
+      "image_masks_loss": 0.0078,
+      "llm_learning_rate": 2.414151548828755e-06,
+      "loss": 0.8023,
+      "output_text_mask_loss": 0.041,
+      "predictor_learning_rate": 0.00024141515488287546,
+      "step": 5130
+    },
+    {
+      "epoch": 0.53,
+      "gumbel_tau": 0.2981,
+      "image_masks_loss": 0.0193,
+      "llm_learning_rate": 2.413322426118917e-06,
+      "loss": 0.7765,
+      "output_text_mask_loss": 0.0141,
+      "predictor_learning_rate": 0.00024133224261189168,
+      "step": 5131
+    },
+    {
+      "epoch": 0.53,
+      "gumbel_tau": 0.298,
+      "image_masks_loss": 0.0109,
+      "llm_learning_rate": 2.412493312954207e-06,
+      "loss": 0.6628,
+      "output_text_mask_loss": 0.0055,
+      "predictor_learning_rate": 0.00024124933129542067,
+      "step": 5132
+    },
+    {
+      "epoch": 0.53,
+      "gumbel_tau": 0.298,
+      "image_masks_loss": 0.0088,
+      "llm_learning_rate": 2.411664209425927e-06,
+      "loss": 0.7327,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0002411664209425927,
+      "step": 5133
+    },
+    {
+      "epoch": 0.53,
+      "gumbel_tau": 0.2979,
+      "image_masks_loss": 0.0068,
+      "llm_learning_rate": 2.4108351156253825e-06,
+      "loss": 0.7944,
+      "output_text_mask_loss": 0.0298,
+      "predictor_learning_rate": 0.00024108351156253822,
+      "step": 5134
+    },
+    {
+      "epoch": 0.53,
+      "gumbel_tau": 0.2978,
+      "image_masks_loss": 0.0158,
+      "llm_learning_rate": 2.4100060316438727e-06,
+      "loss": 0.7947,
+      "output_text_mask_loss": 0.0193,
+      "predictor_learning_rate": 0.00024100060316438725,
+      "step": 5135
+    },
+    {
+      "epoch": 0.53,
+      "gumbel_tau": 0.2978,
+      "image_masks_loss": 0.0062,
+      "llm_learning_rate": 2.409176957572699e-06,
+      "loss": 0.7448,
+      "output_text_mask_loss": 0.0231,
+      "predictor_learning_rate": 0.0002409176957572699,
+      "step": 5136
+    },
+    {
+      "epoch": 0.53,
+      "gumbel_tau": 0.2977,
+      "image_masks_loss": 0.0084,
+      "llm_learning_rate": 2.4083478935031616e-06,
+      "loss": 0.8553,
+      "output_text_mask_loss": 0.0032,
+      "predictor_learning_rate": 0.00024083478935031617,
+      "step": 5137
+    },
+    {
+      "epoch": 0.53,
+      "gumbel_tau": 0.2976,
+      "image_masks_loss": 0.0092,
+      "llm_learning_rate": 2.4075188395265594e-06,
+      "loss": 0.6131,
+      "output_text_mask_loss": 0.0004,
+      "predictor_learning_rate": 0.0002407518839526559,
+      "step": 5138
+    },
+    {
+      "epoch": 0.53,
+      "gumbel_tau": 0.2976,
+      "image_masks_loss": 0.0091,
+      "llm_learning_rate": 2.406689795734187e-06,
+      "loss": 0.7849,
+      "output_text_mask_loss": 0.0462,
+      "predictor_learning_rate": 0.0002406689795734187,
+      "step": 5139
+    },
+    {
+      "epoch": 0.53,
+      "gumbel_tau": 0.2975,
+      "image_masks_loss": 0.0092,
+      "llm_learning_rate": 2.405860762217343e-06,
+      "loss": 0.6992,
+      "output_text_mask_loss": 0.0078,
+      "predictor_learning_rate": 0.0002405860762217343,
+      "step": 5140
+    },
+    {
+      "epoch": 0.53,
+      "gumbel_tau": 0.2974,
+      "image_masks_loss": 0.0107,
+      "llm_learning_rate": 2.4050317390673212e-06,
+      "loss": 0.7973,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.0002405031739067321,
+      "step": 5141
+    },
+    {
+      "epoch": 0.53,
+      "gumbel_tau": 0.2973,
+      "image_masks_loss": 0.015,
+      "llm_learning_rate": 2.4042027263754153e-06,
+      "loss": 0.7872,
+      "output_text_mask_loss": 0.0105,
+      "predictor_learning_rate": 0.0002404202726375415,
+      "step": 5142
+    },
+    {
+      "epoch": 0.53,
+      "gumbel_tau": 0.2973,
+      "image_masks_loss": 0.0192,
+      "llm_learning_rate": 2.403373724232919e-06,
+      "loss": 0.7196,
+      "output_text_mask_loss": 0.0549,
+      "predictor_learning_rate": 0.0002403373724232919,
+      "step": 5143
+    },
+    {
+      "epoch": 0.53,
+      "gumbel_tau": 0.2972,
+      "image_masks_loss": 0.0074,
+      "llm_learning_rate": 2.4025447327311236e-06,
+      "loss": 0.819,
+      "output_text_mask_loss": 0.0187,
+      "predictor_learning_rate": 0.00024025447327311234,
+      "step": 5144
+    },
+    {
+      "epoch": 0.53,
+      "gumbel_tau": 0.2971,
+      "image_masks_loss": 0.01,
+      "llm_learning_rate": 2.4017157519613187e-06,
+      "loss": 0.7472,
+      "output_text_mask_loss": 0.0454,
+      "predictor_learning_rate": 0.00024017157519613187,
+      "step": 5145
+    },
+    {
+      "epoch": 0.53,
+      "gumbel_tau": 0.2971,
+      "image_masks_loss": 0.0068,
+      "llm_learning_rate": 2.4008867820147946e-06,
+      "loss": 0.6858,
+      "output_text_mask_loss": 0.0322,
+      "predictor_learning_rate": 0.00024008867820147944,
+      "step": 5146
+    },
+    {
+      "epoch": 0.53,
+      "gumbel_tau": 0.297,
+      "image_masks_loss": 0.0236,
+      "llm_learning_rate": 2.400057822982839e-06,
+      "loss": 0.7793,
+      "output_text_mask_loss": 0.0039,
+      "predictor_learning_rate": 0.00024000578229828388,
+      "step": 5147
+    },
+    {
+      "epoch": 0.53,
+      "gumbel_tau": 0.2969,
+      "image_masks_loss": 0.0037,
+      "llm_learning_rate": 2.3992288749567383e-06,
+      "loss": 0.7381,
+      "output_text_mask_loss": 0.0322,
+      "predictor_learning_rate": 0.00023992288749567382,
+      "step": 5148
+    },
+    {
+      "epoch": 0.53,
+      "gumbel_tau": 0.2969,
+      "image_masks_loss": 0.0054,
+      "llm_learning_rate": 2.398399938027778e-06,
+      "loss": 0.7856,
+      "output_text_mask_loss": 0.0368,
+      "predictor_learning_rate": 0.0002398399938027778,
+      "step": 5149
+    },
+    {
+      "epoch": 0.53,
+      "gumbel_tau": 0.2968,
+      "image_masks_loss": 0.0059,
+      "llm_learning_rate": 2.3975710122872433e-06,
+      "loss": 0.7709,
+      "output_text_mask_loss": 0.0235,
+      "predictor_learning_rate": 0.00023975710122872433,
+      "step": 5150
+    },
+    {
+      "epoch": 0.53,
+      "gumbel_tau": 0.2967,
+      "image_masks_loss": 0.025,
+      "llm_learning_rate": 2.396742097826417e-06,
+      "loss": 0.7983,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00023967420978264167,
+      "step": 5151
+    },
+    {
+      "epoch": 0.53,
+      "gumbel_tau": 0.2966,
+      "image_masks_loss": 0.0086,
+      "llm_learning_rate": 2.3959131947365803e-06,
+      "loss": 0.6955,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.000239591319473658,
+      "step": 5152
+    },
+    {
+      "epoch": 0.53,
+      "gumbel_tau": 0.2966,
+      "image_masks_loss": 0.0126,
+      "llm_learning_rate": 2.3950843031090145e-06,
+      "loss": 0.7265,
+      "output_text_mask_loss": 0.0039,
+      "predictor_learning_rate": 0.00023950843031090145,
+      "step": 5153
+    },
+    {
+      "epoch": 0.53,
+      "gumbel_tau": 0.2965,
+      "image_masks_loss": 0.0162,
+      "llm_learning_rate": 2.3942554230349997e-06,
+      "loss": 0.6867,
+      "output_text_mask_loss": 0.0124,
+      "predictor_learning_rate": 0.00023942554230349994,
+      "step": 5154
+    },
+    {
+      "epoch": 0.53,
+      "gumbel_tau": 0.2964,
+      "image_masks_loss": 0.0232,
+      "llm_learning_rate": 2.3934265546058126e-06,
+      "loss": 0.7807,
+      "output_text_mask_loss": 0.0275,
+      "predictor_learning_rate": 0.00023934265546058126,
+      "step": 5155
+    },
+    {
+      "epoch": 0.53,
+      "gumbel_tau": 0.2964,
+      "image_masks_loss": 0.0061,
+      "llm_learning_rate": 2.3925976979127307e-06,
+      "loss": 0.6828,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.00023925976979127304,
+      "step": 5156
+    },
+    {
+      "epoch": 0.53,
+      "gumbel_tau": 0.2963,
+      "image_masks_loss": 0.011,
+      "llm_learning_rate": 2.3917688530470308e-06,
+      "loss": 0.6175,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 0.00023917688530470304,
+      "step": 5157
+    },
+    {
+      "epoch": 0.53,
+      "gumbel_tau": 0.2962,
+      "image_masks_loss": 0.0109,
+      "llm_learning_rate": 2.3909400200999844e-06,
+      "loss": 0.653,
+      "output_text_mask_loss": 0.079,
+      "predictor_learning_rate": 0.00023909400200999842,
+      "step": 5158
+    },
+    {
+      "epoch": 0.53,
+      "gumbel_tau": 0.2962,
+      "image_masks_loss": 0.0126,
+      "llm_learning_rate": 2.390111199162866e-06,
+      "loss": 0.6978,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0002390111199162866,
+      "step": 5159
+    },
+    {
+      "epoch": 0.53,
+      "gumbel_tau": 0.2961,
+      "image_masks_loss": 0.0147,
+      "llm_learning_rate": 2.3892823903269475e-06,
+      "loss": 0.7809,
+      "output_text_mask_loss": 0.0145,
+      "predictor_learning_rate": 0.00023892823903269475,
+      "step": 5160
+    },
+    {
+      "epoch": 0.53,
+      "gumbel_tau": 0.296,
+      "image_masks_loss": 0.0299,
+      "llm_learning_rate": 2.3884535936834986e-06,
+      "loss": 0.8643,
+      "output_text_mask_loss": 0.0041,
+      "predictor_learning_rate": 0.00023884535936834984,
+      "step": 5161
+    },
+    {
+      "epoch": 0.53,
+      "gumbel_tau": 0.2959,
+      "image_masks_loss": 0.0148,
+      "llm_learning_rate": 2.3876248093237884e-06,
+      "loss": 0.7303,
+      "output_text_mask_loss": 0.0081,
+      "predictor_learning_rate": 0.00023876248093237886,
+      "step": 5162
+    },
+    {
+      "epoch": 0.53,
+      "gumbel_tau": 0.2959,
+      "image_masks_loss": 0.0277,
+      "llm_learning_rate": 2.3867960373390847e-06,
+      "loss": 0.7159,
+      "output_text_mask_loss": 0.0006,
+      "predictor_learning_rate": 0.00023867960373390846,
+      "step": 5163
+    },
+    {
+      "epoch": 0.53,
+      "gumbel_tau": 0.2958,
+      "image_masks_loss": 0.0159,
+      "llm_learning_rate": 2.3859672778206537e-06,
+      "loss": 0.7343,
+      "output_text_mask_loss": 0.0061,
+      "predictor_learning_rate": 0.00023859672778206534,
+      "step": 5164
+    },
+    {
+      "epoch": 0.53,
+      "gumbel_tau": 0.2957,
+      "image_masks_loss": 0.0076,
+      "llm_learning_rate": 2.3851385308597592e-06,
+      "loss": 0.8476,
+      "output_text_mask_loss": 0.0248,
+      "predictor_learning_rate": 0.00023851385308597592,
+      "step": 5165
+    },
+    {
+      "epoch": 0.53,
+      "gumbel_tau": 0.2957,
+      "image_masks_loss": 0.0059,
+      "llm_learning_rate": 2.3843097965476662e-06,
+      "loss": 0.7301,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.00023843097965476663,
+      "step": 5166
+    },
+    {
+      "epoch": 0.53,
+      "gumbel_tau": 0.2956,
+      "image_masks_loss": 0.0042,
+      "llm_learning_rate": 2.383481074975636e-06,
+      "loss": 0.7541,
+      "output_text_mask_loss": 0.0012,
+      "predictor_learning_rate": 0.0002383481074975636,
+      "step": 5167
+    },
+    {
+      "epoch": 0.53,
+      "gumbel_tau": 0.2955,
+      "image_masks_loss": 0.0127,
+      "llm_learning_rate": 2.3826523662349295e-06,
+      "loss": 0.8536,
+      "output_text_mask_loss": 0.0065,
+      "predictor_learning_rate": 0.00023826523662349292,
+      "step": 5168
+    },
+    {
+      "epoch": 0.53,
+      "gumbel_tau": 0.2955,
+      "image_masks_loss": 0.0141,
+      "llm_learning_rate": 2.3818236704168053e-06,
+      "loss": 0.6857,
+      "output_text_mask_loss": 0.0389,
+      "predictor_learning_rate": 0.00023818236704168053,
+      "step": 5169
+    },
+    {
+      "epoch": 0.53,
+      "gumbel_tau": 0.2954,
+      "image_masks_loss": 0.0154,
+      "llm_learning_rate": 2.3809949876125224e-06,
+      "loss": 0.7395,
+      "output_text_mask_loss": 0.0039,
+      "predictor_learning_rate": 0.00023809949876125221,
+      "step": 5170
+    },
+    {
+      "epoch": 0.53,
+      "gumbel_tau": 0.2953,
+      "image_masks_loss": 0.0078,
+      "llm_learning_rate": 2.3801663179133368e-06,
+      "loss": 0.6935,
+      "output_text_mask_loss": 0.0058,
+      "predictor_learning_rate": 0.00023801663179133365,
+      "step": 5171
+    },
+    {
+      "epoch": 0.53,
+      "gumbel_tau": 0.2952,
+      "image_masks_loss": 0.0202,
+      "llm_learning_rate": 2.3793376614105018e-06,
+      "loss": 0.7475,
+      "output_text_mask_loss": 0.0229,
+      "predictor_learning_rate": 0.00023793376614105016,
+      "step": 5172
+    },
+    {
+      "epoch": 0.53,
+      "gumbel_tau": 0.2952,
+      "image_masks_loss": 0.0216,
+      "llm_learning_rate": 2.3785090181952734e-06,
+      "loss": 0.7801,
+      "output_text_mask_loss": 0.001,
+      "predictor_learning_rate": 0.00023785090181952732,
+      "step": 5173
+    },
+    {
+      "epoch": 0.53,
+      "gumbel_tau": 0.2951,
+      "image_masks_loss": 0.0196,
+      "llm_learning_rate": 2.3776803883589024e-06,
+      "loss": 0.8155,
+      "output_text_mask_loss": 0.0444,
+      "predictor_learning_rate": 0.00023776803883589022,
+      "step": 5174
+    },
+    {
+      "epoch": 0.53,
+      "gumbel_tau": 0.295,
+      "image_masks_loss": 0.0058,
+      "llm_learning_rate": 2.37685177199264e-06,
+      "loss": 0.7991,
+      "output_text_mask_loss": 0.0908,
+      "predictor_learning_rate": 0.000237685177199264,
+      "step": 5175
+    },
+    {
+      "epoch": 0.53,
+      "gumbel_tau": 0.295,
+      "image_masks_loss": 0.0076,
+      "llm_learning_rate": 2.376023169187734e-06,
+      "loss": 0.7042,
+      "output_text_mask_loss": 0.0198,
+      "predictor_learning_rate": 0.0002376023169187734,
+      "step": 5176
+    },
+    {
+      "epoch": 0.53,
+      "gumbel_tau": 0.2949,
+      "image_masks_loss": 0.0359,
+      "llm_learning_rate": 2.3751945800354326e-06,
+      "loss": 0.7027,
+      "output_text_mask_loss": 0.0488,
+      "predictor_learning_rate": 0.00023751945800354324,
+      "step": 5177
+    },
+    {
+      "epoch": 0.53,
+      "gumbel_tau": 0.2948,
+      "image_masks_loss": 0.0077,
+      "llm_learning_rate": 2.3743660046269827e-06,
+      "loss": 0.7931,
+      "output_text_mask_loss": 0.0039,
+      "predictor_learning_rate": 0.00023743660046269825,
+      "step": 5178
+    },
+    {
+      "epoch": 0.53,
+      "gumbel_tau": 0.2948,
+      "image_masks_loss": 0.0145,
+      "llm_learning_rate": 2.373537443053628e-06,
+      "loss": 0.7375,
+      "output_text_mask_loss": 0.0023,
+      "predictor_learning_rate": 0.00023735374430536277,
+      "step": 5179
+    },
+    {
+      "epoch": 0.53,
+      "gumbel_tau": 0.2947,
+      "image_masks_loss": 0.0185,
+      "llm_learning_rate": 2.3727088954066117e-06,
+      "loss": 0.7326,
+      "output_text_mask_loss": 0.0113,
+      "predictor_learning_rate": 0.00023727088954066114,
+      "step": 5180
+    },
+    {
+      "epoch": 0.53,
+      "gumbel_tau": 0.2946,
+      "image_masks_loss": 0.0053,
+      "llm_learning_rate": 2.3718803617771747e-06,
+      "loss": 0.6842,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00023718803617771744,
+      "step": 5181
+    },
+    {
+      "epoch": 0.53,
+      "gumbel_tau": 0.2946,
+      "image_masks_loss": 0.0212,
+      "llm_learning_rate": 2.3710518422565583e-06,
+      "loss": 0.7699,
+      "output_text_mask_loss": 0.0294,
+      "predictor_learning_rate": 0.0002371051842256558,
+      "step": 5182
+    },
+    {
+      "epoch": 0.53,
+      "gumbel_tau": 0.2945,
+      "image_masks_loss": 0.0254,
+      "llm_learning_rate": 2.3702233369360003e-06,
+      "loss": 0.7078,
+      "output_text_mask_loss": 0.0019,
+      "predictor_learning_rate": 0.00023702233369360002,
+      "step": 5183
+    },
+    {
+      "epoch": 0.53,
+      "gumbel_tau": 0.2944,
+      "image_masks_loss": 0.0171,
+      "llm_learning_rate": 2.369394845906737e-06,
+      "loss": 0.8319,
+      "output_text_mask_loss": 0.0296,
+      "predictor_learning_rate": 0.00023693948459067366,
+      "step": 5184
+    },
+    {
+      "epoch": 0.53,
+      "gumbel_tau": 0.2943,
+      "image_masks_loss": 0.0072,
+      "llm_learning_rate": 2.3685663692600043e-06,
+      "loss": 0.7316,
+      "output_text_mask_loss": 0.017,
+      "predictor_learning_rate": 0.00023685663692600044,
+      "step": 5185
+    },
+    {
+      "epoch": 0.53,
+      "gumbel_tau": 0.2943,
+      "image_masks_loss": 0.0094,
+      "llm_learning_rate": 2.367737907087036e-06,
+      "loss": 0.7914,
+      "output_text_mask_loss": 0.0081,
+      "predictor_learning_rate": 0.0002367737907087036,
+      "step": 5186
+    },
+    {
+      "epoch": 0.53,
+      "gumbel_tau": 0.2942,
+      "image_masks_loss": 0.02,
+      "llm_learning_rate": 2.3669094594790647e-06,
+      "loss": 0.6931,
+      "output_text_mask_loss": 0.0141,
+      "predictor_learning_rate": 0.00023669094594790644,
+      "step": 5187
+    },
+    {
+      "epoch": 0.53,
+      "gumbel_tau": 0.2941,
+      "image_masks_loss": 0.0063,
+      "llm_learning_rate": 2.3660810265273188e-06,
+      "loss": 0.6602,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00023660810265273185,
+      "step": 5188
+    },
+    {
+      "epoch": 0.53,
+      "gumbel_tau": 0.2941,
+      "image_masks_loss": 0.0254,
+      "llm_learning_rate": 2.3652526083230292e-06,
+      "loss": 0.7614,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.0002365252608323029,
+      "step": 5189
+    },
+    {
+      "epoch": 0.53,
+      "gumbel_tau": 0.294,
+      "image_masks_loss": 0.0201,
+      "llm_learning_rate": 2.3644242049574224e-06,
+      "loss": 0.7708,
+      "output_text_mask_loss": 0.0185,
+      "predictor_learning_rate": 0.00023644242049574224,
+      "step": 5190
+    },
+    {
+      "epoch": 0.53,
+      "gumbel_tau": 0.2939,
+      "image_masks_loss": 0.0063,
+      "llm_learning_rate": 2.363595816521724e-06,
+      "loss": 0.7712,
+      "output_text_mask_loss": 0.013,
+      "predictor_learning_rate": 0.0002363595816521724,
+      "step": 5191
+    },
+    {
+      "epoch": 0.53,
+      "gumbel_tau": 0.2939,
+      "image_masks_loss": 0.0131,
+      "llm_learning_rate": 2.3627674431071584e-06,
+      "loss": 0.831,
+      "output_text_mask_loss": 0.0238,
+      "predictor_learning_rate": 0.00023627674431071584,
+      "step": 5192
+    },
+    {
+      "epoch": 0.53,
+      "gumbel_tau": 0.2938,
+      "image_masks_loss": 0.0072,
+      "llm_learning_rate": 2.361939084804948e-06,
+      "loss": 0.7397,
+      "output_text_mask_loss": 0.009,
+      "predictor_learning_rate": 0.00023619390848049476,
+      "step": 5193
+    },
+    {
+      "epoch": 0.53,
+      "gumbel_tau": 0.2937,
+      "image_masks_loss": 0.0113,
+      "llm_learning_rate": 2.3611107417063133e-06,
+      "loss": 0.6743,
+      "output_text_mask_loss": 0.0092,
+      "predictor_learning_rate": 0.00023611107417063128,
+      "step": 5194
+    },
+    {
+      "epoch": 0.53,
+      "gumbel_tau": 0.2936,
+      "image_masks_loss": 0.0197,
+      "llm_learning_rate": 2.3602824139024722e-06,
+      "loss": 0.7704,
+      "output_text_mask_loss": 0.0322,
+      "predictor_learning_rate": 0.00023602824139024723,
+      "step": 5195
+    },
+    {
+      "epoch": 0.53,
+      "gumbel_tau": 0.2936,
+      "image_masks_loss": 0.0113,
+      "llm_learning_rate": 2.3594541014846437e-06,
+      "loss": 0.7415,
+      "output_text_mask_loss": 0.0071,
+      "predictor_learning_rate": 0.00023594541014846436,
+      "step": 5196
+    },
+    {
+      "epoch": 0.53,
+      "gumbel_tau": 0.2935,
+      "image_masks_loss": 0.0217,
+      "llm_learning_rate": 2.358625804544042e-06,
+      "loss": 0.7843,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 0.00023586258045440415,
+      "step": 5197
+    },
+    {
+      "epoch": 0.53,
+      "gumbel_tau": 0.2934,
+      "image_masks_loss": 0.013,
+      "llm_learning_rate": 2.3577975231718813e-06,
+      "loss": 0.7669,
+      "output_text_mask_loss": 0.0012,
+      "predictor_learning_rate": 0.00023577975231718813,
+      "step": 5198
+    },
+    {
+      "epoch": 0.53,
+      "gumbel_tau": 0.2934,
+      "image_masks_loss": 0.0108,
+      "llm_learning_rate": 2.3569692574593746e-06,
+      "loss": 0.7626,
+      "output_text_mask_loss": 0.0015,
+      "predictor_learning_rate": 0.00023569692574593746,
+      "step": 5199
+    },
+    {
+      "epoch": 0.53,
+      "gumbel_tau": 0.2933,
+      "image_masks_loss": 0.0201,
+      "llm_learning_rate": 2.3561410074977314e-06,
+      "loss": 0.692,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00023561410074977312,
+      "step": 5200
+    },
+    {
+      "epoch": 0.53,
+      "gumbel_tau": 0.2932,
+      "image_masks_loss": 0.0043,
+      "llm_learning_rate": 2.355312773378161e-06,
+      "loss": 0.7951,
+      "output_text_mask_loss": 0.0038,
+      "predictor_learning_rate": 0.0002355312773378161,
+      "step": 5201
+    },
+    {
+      "epoch": 0.53,
+      "gumbel_tau": 0.2932,
+      "image_masks_loss": 0.0085,
+      "llm_learning_rate": 2.354484555191871e-06,
+      "loss": 0.7685,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.00023544845551918707,
+      "step": 5202
+    },
+    {
+      "epoch": 0.53,
+      "gumbel_tau": 0.2931,
+      "image_masks_loss": 0.0081,
+      "llm_learning_rate": 2.3536563530300654e-06,
+      "loss": 0.7805,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.00023536563530300654,
+      "step": 5203
+    },
+    {
+      "epoch": 0.53,
+      "gumbel_tau": 0.293,
+      "image_masks_loss": 0.0134,
+      "llm_learning_rate": 2.3528281669839474e-06,
+      "loss": 0.8129,
+      "output_text_mask_loss": 0.0694,
+      "predictor_learning_rate": 0.00023528281669839475,
+      "step": 5204
+    },
+    {
+      "epoch": 0.53,
+      "gumbel_tau": 0.293,
+      "image_masks_loss": 0.0098,
+      "llm_learning_rate": 2.3519999971447206e-06,
+      "loss": 0.6991,
+      "output_text_mask_loss": 0.0061,
+      "predictor_learning_rate": 0.00023519999971447202,
+      "step": 5205
+    },
+    {
+      "epoch": 0.53,
+      "gumbel_tau": 0.2929,
+      "image_masks_loss": 0.0114,
+      "llm_learning_rate": 2.3511718436035835e-06,
+      "loss": 0.7416,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00023511718436035835,
+      "step": 5206
+    },
+    {
+      "epoch": 0.53,
+      "gumbel_tau": 0.2928,
+      "image_masks_loss": 0.0103,
+      "llm_learning_rate": 2.350343706451734e-06,
+      "loss": 0.7635,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.0002350343706451734,
+      "step": 5207
+    },
+    {
+      "epoch": 0.53,
+      "gumbel_tau": 0.2927,
+      "image_masks_loss": 0.0083,
+      "llm_learning_rate": 2.3495155857803696e-06,
+      "loss": 0.8015,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.00023495155857803694,
+      "step": 5208
+    },
+    {
+      "epoch": 0.53,
+      "gumbel_tau": 0.2927,
+      "image_masks_loss": 0.012,
+      "llm_learning_rate": 2.3486874816806836e-06,
+      "loss": 0.673,
+      "output_text_mask_loss": 0.0126,
+      "predictor_learning_rate": 0.00023486874816806835,
+      "step": 5209
+    },
+    {
+      "epoch": 0.53,
+      "gumbel_tau": 0.2926,
+      "image_masks_loss": 0.0088,
+      "llm_learning_rate": 2.347859394243869e-06,
+      "loss": 0.8112,
+      "output_text_mask_loss": 0.0024,
+      "predictor_learning_rate": 0.0002347859394243869,
+      "step": 5210
+    },
+    {
+      "epoch": 0.53,
+      "gumbel_tau": 0.2925,
+      "image_masks_loss": 0.0093,
+      "llm_learning_rate": 2.3470313235611173e-06,
+      "loss": 0.7778,
+      "output_text_mask_loss": 0.0393,
+      "predictor_learning_rate": 0.0002347031323561117,
+      "step": 5211
+    },
+    {
+      "epoch": 0.53,
+      "gumbel_tau": 0.2925,
+      "image_masks_loss": 0.0095,
+      "llm_learning_rate": 2.3462032697236177e-06,
+      "loss": 0.8077,
+      "output_text_mask_loss": 0.0488,
+      "predictor_learning_rate": 0.00023462032697236176,
+      "step": 5212
+    },
+    {
+      "epoch": 0.53,
+      "gumbel_tau": 0.2924,
+      "image_masks_loss": 0.0084,
+      "llm_learning_rate": 2.3453752328225554e-06,
+      "loss": 0.7229,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00023453752328225552,
+      "step": 5213
+    },
+    {
+      "epoch": 0.53,
+      "gumbel_tau": 0.2923,
+      "image_masks_loss": 0.0083,
+      "llm_learning_rate": 2.3445472129491173e-06,
+      "loss": 0.7613,
+      "output_text_mask_loss": 0.0156,
+      "predictor_learning_rate": 0.00023445472129491173,
+      "step": 5214
+    },
+    {
+      "epoch": 0.53,
+      "gumbel_tau": 0.2923,
+      "image_masks_loss": 0.0103,
+      "llm_learning_rate": 2.343719210194486e-06,
+      "loss": 0.6246,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 0.0002343719210194486,
+      "step": 5215
+    },
+    {
+      "epoch": 0.53,
+      "gumbel_tau": 0.2922,
+      "image_masks_loss": 0.0104,
+      "llm_learning_rate": 2.3428912246498432e-06,
+      "loss": 0.8208,
+      "output_text_mask_loss": 0.0198,
+      "predictor_learning_rate": 0.00023428912246498433,
+      "step": 5216
+    },
+    {
+      "epoch": 0.53,
+      "gumbel_tau": 0.2921,
+      "image_masks_loss": 0.0214,
+      "llm_learning_rate": 2.3420632564063692e-06,
+      "loss": 0.7741,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00023420632564063692,
+      "step": 5217
+    },
+    {
+      "epoch": 0.53,
+      "gumbel_tau": 0.2921,
+      "image_masks_loss": 0.0196,
+      "llm_learning_rate": 2.3412353055552413e-06,
+      "loss": 0.8012,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.00023412353055552414,
+      "step": 5218
+    },
+    {
+      "epoch": 0.53,
+      "gumbel_tau": 0.292,
+      "image_masks_loss": 0.011,
+      "llm_learning_rate": 2.3404073721876353e-06,
+      "loss": 0.7299,
+      "output_text_mask_loss": 0.0029,
+      "predictor_learning_rate": 0.00023404073721876352,
+      "step": 5219
+    },
+    {
+      "epoch": 0.53,
+      "gumbel_tau": 0.2919,
+      "image_masks_loss": 0.0209,
+      "llm_learning_rate": 2.339579456394724e-06,
+      "loss": 0.7901,
+      "output_text_mask_loss": 0.017,
+      "predictor_learning_rate": 0.0002339579456394724,
+      "step": 5220
+    },
+    {
+      "epoch": 0.53,
+      "gumbel_tau": 0.2919,
+      "image_masks_loss": 0.0191,
+      "llm_learning_rate": 2.3387515582676812e-06,
+      "loss": 0.7839,
+      "output_text_mask_loss": 0.0641,
+      "predictor_learning_rate": 0.0002338751558267681,
+      "step": 5221
+    },
+    {
+      "epoch": 0.54,
+      "gumbel_tau": 0.2918,
+      "image_masks_loss": 0.0207,
+      "llm_learning_rate": 2.3379236778976757e-06,
+      "loss": 0.6899,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 0.00023379236778976755,
+      "step": 5222
+    },
+    {
+      "epoch": 0.54,
+      "gumbel_tau": 0.2917,
+      "image_masks_loss": 0.0055,
+      "llm_learning_rate": 2.3370958153758753e-06,
+      "loss": 0.7166,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.0002337095815375875,
+      "step": 5223
+    },
+    {
+      "epoch": 0.54,
+      "gumbel_tau": 0.2916,
+      "image_masks_loss": 0.0111,
+      "llm_learning_rate": 2.336267970793447e-06,
+      "loss": 0.7207,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0002336267970793447,
+      "step": 5224
+    },
+    {
+      "epoch": 0.54,
+      "gumbel_tau": 0.2916,
+      "image_masks_loss": 0.0087,
+      "llm_learning_rate": 2.335440144241555e-06,
+      "loss": 0.6696,
+      "output_text_mask_loss": 0.0096,
+      "predictor_learning_rate": 0.00023354401442415547,
+      "step": 5225
+    },
+    {
+      "epoch": 0.54,
+      "gumbel_tau": 0.2915,
+      "image_masks_loss": 0.0074,
+      "llm_learning_rate": 2.33461233581136e-06,
+      "loss": 0.7238,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00023346123358113603,
+      "step": 5226
+    },
+    {
+      "epoch": 0.54,
+      "gumbel_tau": 0.2914,
+      "image_masks_loss": 0.0057,
+      "llm_learning_rate": 2.333784545594023e-06,
+      "loss": 0.698,
+      "output_text_mask_loss": 0.0488,
+      "predictor_learning_rate": 0.00023337845455940229,
+      "step": 5227
+    },
+    {
+      "epoch": 0.54,
+      "gumbel_tau": 0.2914,
+      "image_masks_loss": 0.0133,
+      "llm_learning_rate": 2.3329567736807025e-06,
+      "loss": 0.7895,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00023329567736807025,
+      "step": 5228
+    },
+    {
+      "epoch": 0.54,
+      "gumbel_tau": 0.2913,
+      "image_masks_loss": 0.0175,
+      "llm_learning_rate": 2.3321290201625544e-06,
+      "loss": 0.8113,
+      "output_text_mask_loss": 0.013,
+      "predictor_learning_rate": 0.00023321290201625542,
+      "step": 5229
+    },
+    {
+      "epoch": 0.54,
+      "gumbel_tau": 0.2912,
+      "image_masks_loss": 0.0089,
+      "llm_learning_rate": 2.3313012851307314e-06,
+      "loss": 0.6949,
+      "output_text_mask_loss": 0.0286,
+      "predictor_learning_rate": 0.00023313012851307313,
+      "step": 5230
+    },
+    {
+      "epoch": 0.54,
+      "gumbel_tau": 0.2912,
+      "image_masks_loss": 0.0027,
+      "llm_learning_rate": 2.3304735686763886e-06,
+      "loss": 0.7563,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00023304735686763885,
+      "step": 5231
+    },
+    {
+      "epoch": 0.54,
+      "gumbel_tau": 0.2911,
+      "image_masks_loss": 0.0117,
+      "llm_learning_rate": 2.3296458708906728e-06,
+      "loss": 0.7678,
+      "output_text_mask_loss": 0.0189,
+      "predictor_learning_rate": 0.00023296458708906725,
+      "step": 5232
+    },
+    {
+      "epoch": 0.54,
+      "gumbel_tau": 0.291,
+      "image_masks_loss": 0.0052,
+      "llm_learning_rate": 2.328818191864733e-06,
+      "loss": 0.6869,
+      "output_text_mask_loss": 0.0576,
+      "predictor_learning_rate": 0.00023288181918647332,
+      "step": 5233
+    },
+    {
+      "epoch": 0.54,
+      "gumbel_tau": 0.291,
+      "image_masks_loss": 0.0069,
+      "llm_learning_rate": 2.3279905316897157e-06,
+      "loss": 0.7837,
+      "output_text_mask_loss": 0.0235,
+      "predictor_learning_rate": 0.00023279905316897154,
+      "step": 5234
+    },
+    {
+      "epoch": 0.54,
+      "gumbel_tau": 0.2909,
+      "image_masks_loss": 0.0134,
+      "llm_learning_rate": 2.327162890456764e-06,
+      "loss": 0.7174,
+      "output_text_mask_loss": 0.021,
+      "predictor_learning_rate": 0.0002327162890456764,
+      "step": 5235
+    },
+    {
+      "epoch": 0.54,
+      "gumbel_tau": 0.2908,
+      "image_masks_loss": 0.0078,
+      "llm_learning_rate": 2.326335268257019e-06,
+      "loss": 0.7039,
+      "output_text_mask_loss": 0.0221,
+      "predictor_learning_rate": 0.0002326335268257019,
+      "step": 5236
+    },
+    {
+      "epoch": 0.54,
+      "gumbel_tau": 0.2908,
+      "image_masks_loss": 0.0112,
+      "llm_learning_rate": 2.325507665181622e-06,
+      "loss": 0.7298,
+      "output_text_mask_loss": 0.0095,
+      "predictor_learning_rate": 0.0002325507665181622,
+      "step": 5237
+    },
+    {
+      "epoch": 0.54,
+      "gumbel_tau": 0.2907,
+      "image_masks_loss": 0.0116,
+      "llm_learning_rate": 2.3246800813217095e-06,
+      "loss": 0.7632,
+      "output_text_mask_loss": 0.0231,
+      "predictor_learning_rate": 0.00023246800813217093,
+      "step": 5238
+    },
+    {
+      "epoch": 0.54,
+      "gumbel_tau": 0.2906,
+      "image_masks_loss": 0.0221,
+      "llm_learning_rate": 2.323852516768416e-06,
+      "loss": 0.7951,
+      "output_text_mask_loss": 0.0179,
+      "predictor_learning_rate": 0.0002323852516768416,
+      "step": 5239
+    },
+    {
+      "epoch": 0.54,
+      "gumbel_tau": 0.2905,
+      "image_masks_loss": 0.0212,
+      "llm_learning_rate": 2.323024971612877e-06,
+      "loss": 0.7158,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00023230249716128768,
+      "step": 5240
+    },
+    {
+      "epoch": 0.54,
+      "gumbel_tau": 0.2905,
+      "image_masks_loss": 0.0122,
+      "llm_learning_rate": 2.3221974459462217e-06,
+      "loss": 0.7529,
+      "output_text_mask_loss": 0.0055,
+      "predictor_learning_rate": 0.00023221974459462216,
+      "step": 5241
+    },
+    {
+      "epoch": 0.54,
+      "gumbel_tau": 0.2904,
+      "image_masks_loss": 0.0057,
+      "llm_learning_rate": 2.32136993985958e-06,
+      "loss": 0.7118,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.00023213699398595797,
+      "step": 5242
+    },
+    {
+      "epoch": 0.54,
+      "gumbel_tau": 0.2903,
+      "image_masks_loss": 0.0132,
+      "llm_learning_rate": 2.3205424534440776e-06,
+      "loss": 0.6464,
+      "output_text_mask_loss": 0.0078,
+      "predictor_learning_rate": 0.00023205424534440775,
+      "step": 5243
+    },
+    {
+      "epoch": 0.54,
+      "gumbel_tau": 0.2903,
+      "image_masks_loss": 0.0096,
+      "llm_learning_rate": 2.3197149867908407e-06,
+      "loss": 0.7513,
+      "output_text_mask_loss": 0.0124,
+      "predictor_learning_rate": 0.00023197149867908406,
+      "step": 5244
+    },
+    {
+      "epoch": 0.54,
+      "gumbel_tau": 0.2902,
+      "image_masks_loss": 0.006,
+      "llm_learning_rate": 2.318887539990991e-06,
+      "loss": 0.6429,
+      "output_text_mask_loss": 0.0124,
+      "predictor_learning_rate": 0.00023188875399909908,
+      "step": 5245
+    },
+    {
+      "epoch": 0.54,
+      "gumbel_tau": 0.2901,
+      "image_masks_loss": 0.0108,
+      "llm_learning_rate": 2.3180601131356484e-06,
+      "loss": 0.702,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00023180601131356481,
+      "step": 5246
+    },
+    {
+      "epoch": 0.54,
+      "gumbel_tau": 0.2901,
+      "image_masks_loss": 0.0143,
+      "llm_learning_rate": 2.3172327063159317e-06,
+      "loss": 0.738,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.00023172327063159315,
+      "step": 5247
+    },
+    {
+      "epoch": 0.54,
+      "gumbel_tau": 0.29,
+      "image_masks_loss": 0.0123,
+      "llm_learning_rate": 2.3164053196229566e-06,
+      "loss": 0.7781,
+      "output_text_mask_loss": 0.0019,
+      "predictor_learning_rate": 0.00023164053196229564,
+      "step": 5248
+    },
+    {
+      "epoch": 0.54,
+      "gumbel_tau": 0.2899,
+      "image_masks_loss": 0.018,
+      "llm_learning_rate": 2.315577953147837e-06,
+      "loss": 0.7377,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.0002315577953147837,
+      "step": 5249
+    },
+    {
+      "epoch": 0.54,
+      "gumbel_tau": 0.2899,
+      "image_masks_loss": 0.0065,
+      "llm_learning_rate": 2.3147506069816837e-06,
+      "loss": 0.7641,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00023147506069816834,
+      "step": 5250
+    },
+    {
+      "epoch": 0.54,
+      "gumbel_tau": 0.2898,
+      "image_masks_loss": 0.0191,
+      "llm_learning_rate": 2.313923281215606e-06,
+      "loss": 0.7838,
+      "output_text_mask_loss": 0.0052,
+      "predictor_learning_rate": 0.00023139232812156057,
+      "step": 5251
+    },
+    {
+      "epoch": 0.54,
+      "gumbel_tau": 0.2897,
+      "image_masks_loss": 0.0046,
+      "llm_learning_rate": 2.313095975940711e-06,
+      "loss": 0.7558,
+      "output_text_mask_loss": 0.0231,
+      "predictor_learning_rate": 0.00023130959759407106,
+      "step": 5252
+    },
+    {
+      "epoch": 0.54,
+      "gumbel_tau": 0.2897,
+      "image_masks_loss": 0.0127,
+      "llm_learning_rate": 2.312268691248103e-06,
+      "loss": 0.6639,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.0002312268691248103,
+      "step": 5253
+    },
+    {
+      "epoch": 0.54,
+      "gumbel_tau": 0.2896,
+      "image_masks_loss": 0.0187,
+      "llm_learning_rate": 2.3114414272288857e-06,
+      "loss": 0.7846,
+      "output_text_mask_loss": 0.0248,
+      "predictor_learning_rate": 0.00023114414272288857,
+      "step": 5254
+    },
+    {
+      "epoch": 0.54,
+      "gumbel_tau": 0.2895,
+      "image_masks_loss": 0.0074,
+      "llm_learning_rate": 2.3106141839741576e-06,
+      "loss": 0.7132,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.00023106141839741574,
+      "step": 5255
+    },
+    {
+      "epoch": 0.54,
+      "gumbel_tau": 0.2895,
+      "image_masks_loss": 0.0193,
+      "llm_learning_rate": 2.309786961575018e-06,
+      "loss": 0.7043,
+      "output_text_mask_loss": 0.0037,
+      "predictor_learning_rate": 0.00023097869615750177,
+      "step": 5256
+    },
+    {
+      "epoch": 0.54,
+      "gumbel_tau": 0.2894,
+      "image_masks_loss": 0.0228,
+      "llm_learning_rate": 2.308959760122562e-06,
+      "loss": 0.7754,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.00023089597601225615,
+      "step": 5257
+    },
+    {
+      "epoch": 0.54,
+      "gumbel_tau": 0.2893,
+      "image_masks_loss": 0.019,
+      "llm_learning_rate": 2.3081325797078825e-06,
+      "loss": 0.694,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00023081325797078822,
+      "step": 5258
+    },
+    {
+      "epoch": 0.54,
+      "gumbel_tau": 0.2892,
+      "image_masks_loss": 0.0113,
+      "llm_learning_rate": 2.3073054204220703e-06,
+      "loss": 0.7197,
+      "output_text_mask_loss": 0.037,
+      "predictor_learning_rate": 0.000230730542042207,
+      "step": 5259
+    },
+    {
+      "epoch": 0.54,
+      "gumbel_tau": 0.2892,
+      "image_masks_loss": 0.0161,
+      "llm_learning_rate": 2.3064782823562147e-06,
+      "loss": 0.6201,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00023064782823562145,
+      "step": 5260
+    },
+    {
+      "epoch": 0.54,
+      "gumbel_tau": 0.2891,
+      "image_masks_loss": 0.0275,
+      "llm_learning_rate": 2.305651165601402e-06,
+      "loss": 0.7144,
+      "output_text_mask_loss": 0.0075,
+      "predictor_learning_rate": 0.0002305651165601402,
+      "step": 5261
+    },
+    {
+      "epoch": 0.54,
+      "gumbel_tau": 0.289,
+      "image_masks_loss": 0.0067,
+      "llm_learning_rate": 2.3048240702487158e-06,
+      "loss": 0.7314,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00023048240702487155,
+      "step": 5262
+    },
+    {
+      "epoch": 0.54,
+      "gumbel_tau": 0.289,
+      "image_masks_loss": 0.0059,
+      "llm_learning_rate": 2.303996996389237e-06,
+      "loss": 0.7711,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.00023039969963892367,
+      "step": 5263
+    },
+    {
+      "epoch": 0.54,
+      "gumbel_tau": 0.2889,
+      "image_masks_loss": 0.0081,
+      "llm_learning_rate": 2.303169944114046e-06,
+      "loss": 0.7431,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0002303169944114046,
+      "step": 5264
+    },
+    {
+      "epoch": 0.54,
+      "gumbel_tau": 0.2888,
+      "image_masks_loss": 0.0225,
+      "llm_learning_rate": 2.3023429135142194e-06,
+      "loss": 0.735,
+      "output_text_mask_loss": 0.0004,
+      "predictor_learning_rate": 0.00023023429135142194,
+      "step": 5265
+    },
+    {
+      "epoch": 0.54,
+      "gumbel_tau": 0.2888,
+      "image_masks_loss": 0.0236,
+      "llm_learning_rate": 2.3015159046808307e-06,
+      "loss": 0.7272,
+      "output_text_mask_loss": 0.0052,
+      "predictor_learning_rate": 0.00023015159046808306,
+      "step": 5266
+    },
+    {
+      "epoch": 0.54,
+      "gumbel_tau": 0.2887,
+      "image_masks_loss": 0.0222,
+      "llm_learning_rate": 2.3006889177049535e-06,
+      "loss": 0.7285,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 0.00023006889177049533,
+      "step": 5267
+    },
+    {
+      "epoch": 0.54,
+      "gumbel_tau": 0.2886,
+      "image_masks_loss": 0.0176,
+      "llm_learning_rate": 2.299861952677657e-06,
+      "loss": 0.8703,
+      "output_text_mask_loss": 0.0004,
+      "predictor_learning_rate": 0.00022998619526776572,
+      "step": 5268
+    },
+    {
+      "epoch": 0.54,
+      "gumbel_tau": 0.2886,
+      "image_masks_loss": 0.0029,
+      "llm_learning_rate": 2.2990350096900076e-06,
+      "loss": 0.7668,
+      "output_text_mask_loss": 0.0014,
+      "predictor_learning_rate": 0.00022990350096900076,
+      "step": 5269
+    },
+    {
+      "epoch": 0.54,
+      "gumbel_tau": 0.2885,
+      "image_masks_loss": 0.0092,
+      "llm_learning_rate": 2.2982080888330704e-06,
+      "loss": 0.7666,
+      "output_text_mask_loss": 0.0496,
+      "predictor_learning_rate": 0.00022982080888330703,
+      "step": 5270
+    },
+    {
+      "epoch": 0.54,
+      "gumbel_tau": 0.2884,
+      "image_masks_loss": 0.0077,
+      "llm_learning_rate": 2.2973811901979084e-06,
+      "loss": 0.7256,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.0002297381190197908,
+      "step": 5271
+    },
+    {
+      "epoch": 0.54,
+      "gumbel_tau": 0.2884,
+      "image_masks_loss": 0.0108,
+      "llm_learning_rate": 2.29655431387558e-06,
+      "loss": 0.7628,
+      "output_text_mask_loss": 0.0113,
+      "predictor_learning_rate": 0.000229655431387558,
+      "step": 5272
+    },
+    {
+      "epoch": 0.54,
+      "gumbel_tau": 0.2883,
+      "image_masks_loss": 0.0071,
+      "llm_learning_rate": 2.2957274599571443e-06,
+      "loss": 0.815,
+      "output_text_mask_loss": 0.0158,
+      "predictor_learning_rate": 0.00022957274599571444,
+      "step": 5273
+    },
+    {
+      "epoch": 0.54,
+      "gumbel_tau": 0.2882,
+      "image_masks_loss": 0.0098,
+      "llm_learning_rate": 2.2949006285336557e-06,
+      "loss": 0.7128,
+      "output_text_mask_loss": 0.0179,
+      "predictor_learning_rate": 0.00022949006285336558,
+      "step": 5274
+    },
+    {
+      "epoch": 0.54,
+      "gumbel_tau": 0.2882,
+      "image_masks_loss": 0.008,
+      "llm_learning_rate": 2.2940738196961664e-06,
+      "loss": 0.7485,
+      "output_text_mask_loss": 0.0713,
+      "predictor_learning_rate": 0.00022940738196961663,
+      "step": 5275
+    },
+    {
+      "epoch": 0.54,
+      "gumbel_tau": 0.2881,
+      "image_masks_loss": 0.0121,
+      "llm_learning_rate": 2.2932470335357273e-06,
+      "loss": 0.7021,
+      "output_text_mask_loss": 0.0889,
+      "predictor_learning_rate": 0.00022932470335357272,
+      "step": 5276
+    },
+    {
+      "epoch": 0.54,
+      "gumbel_tau": 0.288,
+      "image_masks_loss": 0.0081,
+      "llm_learning_rate": 2.2924202701433852e-06,
+      "loss": 0.6934,
+      "output_text_mask_loss": 0.0095,
+      "predictor_learning_rate": 0.00022924202701433852,
+      "step": 5277
+    },
+    {
+      "epoch": 0.54,
+      "gumbel_tau": 0.288,
+      "image_masks_loss": 0.0317,
+      "llm_learning_rate": 2.2915935296101855e-06,
+      "loss": 0.6597,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 0.00022915935296101856,
+      "step": 5278
+    },
+    {
+      "epoch": 0.54,
+      "gumbel_tau": 0.2879,
+      "image_masks_loss": 0.0147,
+      "llm_learning_rate": 2.29076681202717e-06,
+      "loss": 0.768,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.000229076681202717,
+      "step": 5279
+    },
+    {
+      "epoch": 0.54,
+      "gumbel_tau": 0.2878,
+      "image_masks_loss": 0.0204,
+      "llm_learning_rate": 2.2899401174853796e-06,
+      "loss": 0.7684,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 0.00022899401174853796,
+      "step": 5280
+    },
+    {
+      "epoch": 0.54,
+      "gumbel_tau": 0.2878,
+      "image_masks_loss": 0.0112,
+      "llm_learning_rate": 2.2891134460758515e-06,
+      "loss": 0.7099,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00022891134460758513,
+      "step": 5281
+    },
+    {
+      "epoch": 0.54,
+      "gumbel_tau": 0.2877,
+      "image_masks_loss": 0.0064,
+      "llm_learning_rate": 2.2882867978896197e-06,
+      "loss": 0.774,
+      "output_text_mask_loss": 0.0446,
+      "predictor_learning_rate": 0.00022882867978896195,
+      "step": 5282
+    },
+    {
+      "epoch": 0.54,
+      "gumbel_tau": 0.2876,
+      "image_masks_loss": 0.0051,
+      "llm_learning_rate": 2.287460173017718e-06,
+      "loss": 0.7045,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0002287460173017718,
+      "step": 5283
+    },
+    {
+      "epoch": 0.54,
+      "gumbel_tau": 0.2875,
+      "image_masks_loss": 0.0117,
+      "llm_learning_rate": 2.2866335715511755e-06,
+      "loss": 0.7254,
+      "output_text_mask_loss": 0.0246,
+      "predictor_learning_rate": 0.00022866335715511754,
+      "step": 5284
+    },
+    {
+      "epoch": 0.54,
+      "gumbel_tau": 0.2875,
+      "image_masks_loss": 0.0324,
+      "llm_learning_rate": 2.2858069935810193e-06,
+      "loss": 0.763,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 0.0002285806993581019,
+      "step": 5285
+    },
+    {
+      "epoch": 0.54,
+      "gumbel_tau": 0.2874,
+      "image_masks_loss": 0.0111,
+      "llm_learning_rate": 2.2849804391982748e-06,
+      "loss": 0.6995,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00022849804391982746,
+      "step": 5286
+    },
+    {
+      "epoch": 0.54,
+      "gumbel_tau": 0.2873,
+      "image_masks_loss": 0.0246,
+      "llm_learning_rate": 2.2841539084939623e-06,
+      "loss": 0.7191,
+      "output_text_mask_loss": 0.0185,
+      "predictor_learning_rate": 0.0002284153908493962,
+      "step": 5287
+    },
+    {
+      "epoch": 0.54,
+      "gumbel_tau": 0.2873,
+      "image_masks_loss": 0.0102,
+      "llm_learning_rate": 2.283327401559102e-06,
+      "loss": 0.7497,
+      "output_text_mask_loss": 0.0103,
+      "predictor_learning_rate": 0.00022833274015591015,
+      "step": 5288
+    },
+    {
+      "epoch": 0.54,
+      "gumbel_tau": 0.2872,
+      "image_masks_loss": 0.0136,
+      "llm_learning_rate": 2.282500918484711e-06,
+      "loss": 0.9232,
+      "output_text_mask_loss": 0.0607,
+      "predictor_learning_rate": 0.0002282500918484711,
+      "step": 5289
+    },
+    {
+      "epoch": 0.54,
+      "gumbel_tau": 0.2871,
+      "image_masks_loss": 0.0123,
+      "llm_learning_rate": 2.2816744593618037e-06,
+      "loss": 0.7146,
+      "output_text_mask_loss": 0.0188,
+      "predictor_learning_rate": 0.00022816744593618036,
+      "step": 5290
+    },
+    {
+      "epoch": 0.54,
+      "gumbel_tau": 0.2871,
+      "image_masks_loss": 0.0125,
+      "llm_learning_rate": 2.2808480242813907e-06,
+      "loss": 0.731,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.00022808480242813902,
+      "step": 5291
+    },
+    {
+      "epoch": 0.54,
+      "gumbel_tau": 0.287,
+      "image_masks_loss": 0.0181,
+      "llm_learning_rate": 2.2800216133344816e-06,
+      "loss": 0.7522,
+      "output_text_mask_loss": 0.0019,
+      "predictor_learning_rate": 0.00022800216133344814,
+      "step": 5292
+    },
+    {
+      "epoch": 0.54,
+      "gumbel_tau": 0.2869,
+      "image_masks_loss": 0.01,
+      "llm_learning_rate": 2.2791952266120825e-06,
+      "loss": 0.7433,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00022791952266120824,
+      "step": 5293
+    },
+    {
+      "epoch": 0.54,
+      "gumbel_tau": 0.2869,
+      "image_masks_loss": 0.0237,
+      "llm_learning_rate": 2.278368864205197e-06,
+      "loss": 0.8275,
+      "output_text_mask_loss": 0.0032,
+      "predictor_learning_rate": 0.0002278368864205197,
+      "step": 5294
+    },
+    {
+      "epoch": 0.54,
+      "gumbel_tau": 0.2868,
+      "image_masks_loss": 0.0126,
+      "llm_learning_rate": 2.2775425262048254e-06,
+      "loss": 0.6567,
+      "output_text_mask_loss": 0.0189,
+      "predictor_learning_rate": 0.00022775425262048252,
+      "step": 5295
+    },
+    {
+      "epoch": 0.54,
+      "gumbel_tau": 0.2867,
+      "image_masks_loss": 0.0068,
+      "llm_learning_rate": 2.276716212701967e-06,
+      "loss": 0.7098,
+      "output_text_mask_loss": 0.0383,
+      "predictor_learning_rate": 0.00022767162127019667,
+      "step": 5296
+    },
+    {
+      "epoch": 0.54,
+      "gumbel_tau": 0.2867,
+      "image_masks_loss": 0.0095,
+      "llm_learning_rate": 2.275889923787616e-06,
+      "loss": 0.7806,
+      "output_text_mask_loss": 0.103,
+      "predictor_learning_rate": 0.0002275889923787616,
+      "step": 5297
+    },
+    {
+      "epoch": 0.54,
+      "gumbel_tau": 0.2866,
+      "image_masks_loss": 0.0178,
+      "llm_learning_rate": 2.2750636595527657e-06,
+      "loss": 0.7309,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00022750636595527654,
+      "step": 5298
+    },
+    {
+      "epoch": 0.54,
+      "gumbel_tau": 0.2865,
+      "image_masks_loss": 0.0071,
+      "llm_learning_rate": 2.2742374200884065e-06,
+      "loss": 0.6684,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00022742374200884064,
+      "step": 5299
+    },
+    {
+      "epoch": 0.54,
+      "gumbel_tau": 0.2865,
+      "image_masks_loss": 0.0081,
+      "llm_learning_rate": 2.2734112054855256e-06,
+      "loss": 0.8501,
+      "output_text_mask_loss": 0.0778,
+      "predictor_learning_rate": 0.00022734112054855255,
+      "step": 5300
+    },
+    {
+      "epoch": 0.54,
+      "gumbel_tau": 0.2864,
+      "image_masks_loss": 0.0096,
+      "llm_learning_rate": 2.2725850158351076e-06,
+      "loss": 0.78,
+      "output_text_mask_loss": 0.0141,
+      "predictor_learning_rate": 0.00022725850158351074,
+      "step": 5301
+    },
+    {
+      "epoch": 0.54,
+      "gumbel_tau": 0.2863,
+      "image_masks_loss": 0.0061,
+      "llm_learning_rate": 2.2717588512281334e-06,
+      "loss": 0.7518,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.00022717588512281335,
+      "step": 5302
+    },
+    {
+      "epoch": 0.54,
+      "gumbel_tau": 0.2863,
+      "image_masks_loss": 0.0075,
+      "llm_learning_rate": 2.2709327117555838e-06,
+      "loss": 0.7654,
+      "output_text_mask_loss": 0.017,
+      "predictor_learning_rate": 0.00022709327117555837,
+      "step": 5303
+    },
+    {
+      "epoch": 0.54,
+      "gumbel_tau": 0.2862,
+      "image_masks_loss": 0.008,
+      "llm_learning_rate": 2.2701065975084343e-06,
+      "loss": 0.6867,
+      "output_text_mask_loss": 0.0364,
+      "predictor_learning_rate": 0.0002270106597508434,
+      "step": 5304
+    },
+    {
+      "epoch": 0.54,
+      "gumbel_tau": 0.2861,
+      "image_masks_loss": 0.0054,
+      "llm_learning_rate": 2.2692805085776583e-06,
+      "loss": 0.7053,
+      "output_text_mask_loss": 0.0153,
+      "predictor_learning_rate": 0.0002269280508577658,
+      "step": 5305
+    },
+    {
+      "epoch": 0.54,
+      "gumbel_tau": 0.2861,
+      "image_masks_loss": 0.0193,
+      "llm_learning_rate": 2.268454445054226e-06,
+      "loss": 0.7946,
+      "output_text_mask_loss": 0.0235,
+      "predictor_learning_rate": 0.0002268454445054226,
+      "step": 5306
+    },
+    {
+      "epoch": 0.54,
+      "gumbel_tau": 0.286,
+      "image_masks_loss": 0.0143,
+      "llm_learning_rate": 2.267628407029106e-06,
+      "loss": 0.763,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 0.0002267628407029106,
+      "step": 5307
+    },
+    {
+      "epoch": 0.54,
+      "gumbel_tau": 0.2859,
+      "image_masks_loss": 0.0148,
+      "llm_learning_rate": 2.2668023945932636e-06,
+      "loss": 0.7418,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00022668023945932636,
+      "step": 5308
+    },
+    {
+      "epoch": 0.54,
+      "gumbel_tau": 0.2859,
+      "image_masks_loss": 0.0066,
+      "llm_learning_rate": 2.265976407837661e-06,
+      "loss": 0.6771,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.0002265976407837661,
+      "step": 5309
+    },
+    {
+      "epoch": 0.54,
+      "gumbel_tau": 0.2858,
+      "image_masks_loss": 0.0235,
+      "llm_learning_rate": 2.2651504468532577e-06,
+      "loss": 0.8043,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.00022651504468532577,
+      "step": 5310
+    },
+    {
+      "epoch": 0.54,
+      "gumbel_tau": 0.2857,
+      "image_masks_loss": 0.0086,
+      "llm_learning_rate": 2.26432451173101e-06,
+      "loss": 0.7822,
+      "output_text_mask_loss": 0.0275,
+      "predictor_learning_rate": 0.00022643245117310098,
+      "step": 5311
+    },
+    {
+      "epoch": 0.54,
+      "gumbel_tau": 0.2857,
+      "image_masks_loss": 0.0152,
+      "llm_learning_rate": 2.2634986025618723e-06,
+      "loss": 0.6914,
+      "output_text_mask_loss": 0.0006,
+      "predictor_learning_rate": 0.0002263498602561872,
+      "step": 5312
+    },
+    {
+      "epoch": 0.54,
+      "gumbel_tau": 0.2856,
+      "image_masks_loss": 0.0046,
+      "llm_learning_rate": 2.262672719436795e-06,
+      "loss": 0.7369,
+      "output_text_mask_loss": 0.0137,
+      "predictor_learning_rate": 0.00022626727194367948,
+      "step": 5313
+    },
+    {
+      "epoch": 0.54,
+      "gumbel_tau": 0.2855,
+      "image_masks_loss": 0.0113,
+      "llm_learning_rate": 2.2618468624467265e-06,
+      "loss": 0.692,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00022618468624467262,
+      "step": 5314
+    },
+    {
+      "epoch": 0.54,
+      "gumbel_tau": 0.2855,
+      "image_masks_loss": 0.0257,
+      "llm_learning_rate": 2.261021031682612e-06,
+      "loss": 0.8171,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.00022610210316826118,
+      "step": 5315
+    },
+    {
+      "epoch": 0.54,
+      "gumbel_tau": 0.2854,
+      "image_masks_loss": 0.0101,
+      "llm_learning_rate": 2.260195227235394e-06,
+      "loss": 0.6832,
+      "output_text_mask_loss": 0.02,
+      "predictor_learning_rate": 0.0002260195227235394,
+      "step": 5316
+    },
+    {
+      "epoch": 0.54,
+      "gumbel_tau": 0.2853,
+      "image_masks_loss": 0.016,
+      "llm_learning_rate": 2.259369449196012e-06,
+      "loss": 0.754,
+      "output_text_mask_loss": 0.0137,
+      "predictor_learning_rate": 0.00022593694491960116,
+      "step": 5317
+    },
+    {
+      "epoch": 0.54,
+      "gumbel_tau": 0.2852,
+      "image_masks_loss": 0.0106,
+      "llm_learning_rate": 2.258543697655401e-06,
+      "loss": 0.8098,
+      "output_text_mask_loss": 0.0056,
+      "predictor_learning_rate": 0.0002258543697655401,
+      "step": 5318
+    },
+    {
+      "epoch": 0.54,
+      "gumbel_tau": 0.2852,
+      "image_masks_loss": 0.006,
+      "llm_learning_rate": 2.257717972704497e-06,
+      "loss": 0.8473,
+      "output_text_mask_loss": 0.0238,
+      "predictor_learning_rate": 0.00022577179727044968,
+      "step": 5319
+    },
+    {
+      "epoch": 0.55,
+      "gumbel_tau": 0.2851,
+      "image_masks_loss": 0.0049,
+      "llm_learning_rate": 2.2568922744342294e-06,
+      "loss": 0.7898,
+      "output_text_mask_loss": 0.103,
+      "predictor_learning_rate": 0.0002256892274434229,
+      "step": 5320
+    },
+    {
+      "epoch": 0.55,
+      "gumbel_tau": 0.285,
+      "image_masks_loss": 0.0052,
+      "llm_learning_rate": 2.256066602935525e-06,
+      "loss": 0.7313,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0002256066602935525,
+      "step": 5321
+    },
+    {
+      "epoch": 0.55,
+      "gumbel_tau": 0.285,
+      "image_masks_loss": 0.012,
+      "llm_learning_rate": 2.255240958299311e-06,
+      "loss": 0.7379,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.00022552409582993107,
+      "step": 5322
+    },
+    {
+      "epoch": 0.55,
+      "gumbel_tau": 0.2849,
+      "image_masks_loss": 0.0036,
+      "llm_learning_rate": 2.254415340616508e-06,
+      "loss": 0.8289,
+      "output_text_mask_loss": 0.0404,
+      "predictor_learning_rate": 0.00022544153406165078,
+      "step": 5323
+    },
+    {
+      "epoch": 0.55,
+      "gumbel_tau": 0.2848,
+      "image_masks_loss": 0.002,
+      "llm_learning_rate": 2.253589749978034e-06,
+      "loss": 0.7923,
+      "output_text_mask_loss": 0.0284,
+      "predictor_learning_rate": 0.00022535897499780338,
+      "step": 5324
+    },
+    {
+      "epoch": 0.55,
+      "gumbel_tau": 0.2848,
+      "image_masks_loss": 0.0134,
+      "llm_learning_rate": 2.2527641864748057e-06,
+      "loss": 0.7795,
+      "output_text_mask_loss": 0.0088,
+      "predictor_learning_rate": 0.00022527641864748056,
+      "step": 5325
+    },
+    {
+      "epoch": 0.55,
+      "gumbel_tau": 0.2847,
+      "image_masks_loss": 0.0207,
+      "llm_learning_rate": 2.251938650197736e-06,
+      "loss": 0.6981,
+      "output_text_mask_loss": 0.0257,
+      "predictor_learning_rate": 0.00022519386501977358,
+      "step": 5326
+    },
+    {
+      "epoch": 0.55,
+      "gumbel_tau": 0.2846,
+      "image_masks_loss": 0.0109,
+      "llm_learning_rate": 2.251113141237734e-06,
+      "loss": 0.6749,
+      "output_text_mask_loss": 0.0374,
+      "predictor_learning_rate": 0.00022511131412377336,
+      "step": 5327
+    },
+    {
+      "epoch": 0.55,
+      "gumbel_tau": 0.2846,
+      "image_masks_loss": 0.0093,
+      "llm_learning_rate": 2.2502876596857076e-06,
+      "loss": 0.8262,
+      "output_text_mask_loss": 0.0309,
+      "predictor_learning_rate": 0.00022502876596857074,
+      "step": 5328
+    },
+    {
+      "epoch": 0.55,
+      "gumbel_tau": 0.2845,
+      "image_masks_loss": 0.0159,
+      "llm_learning_rate": 2.2494622056325606e-06,
+      "loss": 0.748,
+      "output_text_mask_loss": 0.0115,
+      "predictor_learning_rate": 0.00022494622056325607,
+      "step": 5329
+    },
+    {
+      "epoch": 0.55,
+      "gumbel_tau": 0.2844,
+      "image_masks_loss": 0.0103,
+      "llm_learning_rate": 2.248636779169193e-06,
+      "loss": 0.7385,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.0002248636779169193,
+      "step": 5330
+    },
+    {
+      "epoch": 0.55,
+      "gumbel_tau": 0.2844,
+      "image_masks_loss": 0.0103,
+      "llm_learning_rate": 2.2478113803865038e-06,
+      "loss": 0.7502,
+      "output_text_mask_loss": 0.0111,
+      "predictor_learning_rate": 0.00022478113803865035,
+      "step": 5331
+    },
+    {
+      "epoch": 0.55,
+      "gumbel_tau": 0.2843,
+      "image_masks_loss": 0.0043,
+      "llm_learning_rate": 2.246986009375387e-06,
+      "loss": 0.7073,
+      "output_text_mask_loss": 0.0137,
+      "predictor_learning_rate": 0.0002246986009375387,
+      "step": 5332
+    },
+    {
+      "epoch": 0.55,
+      "gumbel_tau": 0.2842,
+      "image_masks_loss": 0.0072,
+      "llm_learning_rate": 2.2461606662267342e-06,
+      "loss": 0.6555,
+      "output_text_mask_loss": 0.0042,
+      "predictor_learning_rate": 0.0002246160666226734,
+      "step": 5333
+    },
+    {
+      "epoch": 0.55,
+      "gumbel_tau": 0.2842,
+      "image_masks_loss": 0.0089,
+      "llm_learning_rate": 2.245335351031434e-06,
+      "loss": 0.7471,
+      "output_text_mask_loss": 0.0565,
+      "predictor_learning_rate": 0.0002245335351031434,
+      "step": 5334
+    },
+    {
+      "epoch": 0.55,
+      "gumbel_tau": 0.2841,
+      "image_masks_loss": 0.0029,
+      "llm_learning_rate": 2.2445100638803726e-06,
+      "loss": 0.6905,
+      "output_text_mask_loss": 0.0039,
+      "predictor_learning_rate": 0.00022445100638803726,
+      "step": 5335
+    },
+    {
+      "epoch": 0.55,
+      "gumbel_tau": 0.284,
+      "image_masks_loss": 0.0126,
+      "llm_learning_rate": 2.243684804864432e-06,
+      "loss": 0.747,
+      "output_text_mask_loss": 0.0067,
+      "predictor_learning_rate": 0.0002243684804864432,
+      "step": 5336
+    },
+    {
+      "epoch": 0.55,
+      "gumbel_tau": 0.284,
+      "image_masks_loss": 0.0121,
+      "llm_learning_rate": 2.2428595740744917e-06,
+      "loss": 0.7004,
+      "output_text_mask_loss": 0.0019,
+      "predictor_learning_rate": 0.00022428595740744915,
+      "step": 5337
+    },
+    {
+      "epoch": 0.55,
+      "gumbel_tau": 0.2839,
+      "image_masks_loss": 0.0135,
+      "llm_learning_rate": 2.2420343716014274e-06,
+      "loss": 0.67,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00022420343716014274,
+      "step": 5338
+    },
+    {
+      "epoch": 0.55,
+      "gumbel_tau": 0.2838,
+      "image_masks_loss": 0.0072,
+      "llm_learning_rate": 2.2412091975361135e-06,
+      "loss": 0.7153,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.00022412091975361131,
+      "step": 5339
+    },
+    {
+      "epoch": 0.55,
+      "gumbel_tau": 0.2838,
+      "image_masks_loss": 0.0206,
+      "llm_learning_rate": 2.2403840519694187e-06,
+      "loss": 0.7784,
+      "output_text_mask_loss": 0.0156,
+      "predictor_learning_rate": 0.00022403840519694186,
+      "step": 5340
+    },
+    {
+      "epoch": 0.55,
+      "gumbel_tau": 0.2837,
+      "image_masks_loss": 0.006,
+      "llm_learning_rate": 2.23955893499221e-06,
+      "loss": 0.8438,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.000223955893499221,
+      "step": 5341
+    },
+    {
+      "epoch": 0.55,
+      "gumbel_tau": 0.2836,
+      "image_masks_loss": 0.0106,
+      "llm_learning_rate": 2.238733846695353e-06,
+      "loss": 0.7635,
+      "output_text_mask_loss": 0.0324,
+      "predictor_learning_rate": 0.0002238733846695353,
+      "step": 5342
+    },
+    {
+      "epoch": 0.55,
+      "gumbel_tau": 0.2836,
+      "image_masks_loss": 0.0044,
+      "llm_learning_rate": 2.237908787169705e-06,
+      "loss": 0.7586,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.00022379087871697052,
+      "step": 5343
+    },
+    {
+      "epoch": 0.55,
+      "gumbel_tau": 0.2835,
+      "image_masks_loss": 0.0277,
+      "llm_learning_rate": 2.2370837565061264e-06,
+      "loss": 0.817,
+      "output_text_mask_loss": 0.0368,
+      "predictor_learning_rate": 0.00022370837565061262,
+      "step": 5344
+    },
+    {
+      "epoch": 0.55,
+      "gumbel_tau": 0.2834,
+      "image_masks_loss": 0.0067,
+      "llm_learning_rate": 2.2362587547954696e-06,
+      "loss": 0.7355,
+      "output_text_mask_loss": 0.0006,
+      "predictor_learning_rate": 0.00022362587547954693,
+      "step": 5345
+    },
+    {
+      "epoch": 0.55,
+      "gumbel_tau": 0.2834,
+      "image_masks_loss": 0.0082,
+      "llm_learning_rate": 2.2354337821285857e-06,
+      "loss": 0.6989,
+      "output_text_mask_loss": 0.0084,
+      "predictor_learning_rate": 0.00022354337821285856,
+      "step": 5346
+    },
+    {
+      "epoch": 0.55,
+      "gumbel_tau": 0.2833,
+      "image_masks_loss": 0.0146,
+      "llm_learning_rate": 2.2346088385963237e-06,
+      "loss": 0.7339,
+      "output_text_mask_loss": 0.0355,
+      "predictor_learning_rate": 0.00022346088385963235,
+      "step": 5347
+    },
+    {
+      "epoch": 0.55,
+      "gumbel_tau": 0.2832,
+      "image_masks_loss": 0.0097,
+      "llm_learning_rate": 2.2337839242895275e-06,
+      "loss": 0.7288,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0002233783924289527,
+      "step": 5348
+    },
+    {
+      "epoch": 0.55,
+      "gumbel_tau": 0.2832,
+      "image_masks_loss": 0.0053,
+      "llm_learning_rate": 2.2329590392990384e-06,
+      "loss": 0.7328,
+      "output_text_mask_loss": 0.0038,
+      "predictor_learning_rate": 0.0002232959039299038,
+      "step": 5349
+    },
+    {
+      "epoch": 0.55,
+      "gumbel_tau": 0.2831,
+      "image_masks_loss": 0.0045,
+      "llm_learning_rate": 2.232134183715694e-06,
+      "loss": 0.7402,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.0002232134183715694,
+      "step": 5350
+    },
+    {
+      "epoch": 0.55,
+      "gumbel_tau": 0.283,
+      "image_masks_loss": 0.0214,
+      "llm_learning_rate": 2.2313093576303307e-06,
+      "loss": 0.6917,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.00022313093576303307,
+      "step": 5351
+    },
+    {
+      "epoch": 0.55,
+      "gumbel_tau": 0.283,
+      "image_masks_loss": 0.004,
+      "llm_learning_rate": 2.2304845611337794e-06,
+      "loss": 0.7081,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.00022304845611337794,
+      "step": 5352
+    },
+    {
+      "epoch": 0.55,
+      "gumbel_tau": 0.2829,
+      "image_masks_loss": 0.0125,
+      "llm_learning_rate": 2.229659794316868e-06,
+      "loss": 0.7684,
+      "output_text_mask_loss": 0.0101,
+      "predictor_learning_rate": 0.00022296597943168679,
+      "step": 5353
+    },
+    {
+      "epoch": 0.55,
+      "gumbel_tau": 0.2828,
+      "image_masks_loss": 0.0099,
+      "llm_learning_rate": 2.2288350572704226e-06,
+      "loss": 0.7685,
+      "output_text_mask_loss": 0.0103,
+      "predictor_learning_rate": 0.00022288350572704225,
+      "step": 5354
+    },
+    {
+      "epoch": 0.55,
+      "gumbel_tau": 0.2828,
+      "image_masks_loss": 0.013,
+      "llm_learning_rate": 2.228010350085265e-06,
+      "loss": 0.8046,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0002228010350085265,
+      "step": 5355
+    },
+    {
+      "epoch": 0.55,
+      "gumbel_tau": 0.2827,
+      "image_masks_loss": 0.0128,
+      "llm_learning_rate": 2.2271856728522136e-06,
+      "loss": 0.6805,
+      "output_text_mask_loss": 0.0177,
+      "predictor_learning_rate": 0.00022271856728522133,
+      "step": 5356
+    },
+    {
+      "epoch": 0.55,
+      "gumbel_tau": 0.2826,
+      "image_masks_loss": 0.0045,
+      "llm_learning_rate": 2.2263610256620824e-06,
+      "loss": 0.7507,
+      "output_text_mask_loss": 0.0135,
+      "predictor_learning_rate": 0.00022263610256620823,
+      "step": 5357
+    },
+    {
+      "epoch": 0.55,
+      "gumbel_tau": 0.2826,
+      "image_masks_loss": 0.0199,
+      "llm_learning_rate": 2.2255364086056854e-06,
+      "loss": 0.6922,
+      "output_text_mask_loss": 0.0193,
+      "predictor_learning_rate": 0.00022255364086056853,
+      "step": 5358
+    },
+    {
+      "epoch": 0.55,
+      "gumbel_tau": 0.2825,
+      "image_masks_loss": 0.0112,
+      "llm_learning_rate": 2.2247118217738307e-06,
+      "loss": 0.701,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00022247118217738303,
+      "step": 5359
+    },
+    {
+      "epoch": 0.55,
+      "gumbel_tau": 0.2824,
+      "image_masks_loss": 0.0093,
+      "llm_learning_rate": 2.2238872652573234e-06,
+      "loss": 0.7499,
+      "output_text_mask_loss": 0.0174,
+      "predictor_learning_rate": 0.00022238872652573233,
+      "step": 5360
+    },
+    {
+      "epoch": 0.55,
+      "gumbel_tau": 0.2824,
+      "image_masks_loss": 0.0031,
+      "llm_learning_rate": 2.2230627391469653e-06,
+      "loss": 0.8393,
+      "output_text_mask_loss": 0.0322,
+      "predictor_learning_rate": 0.0002223062739146965,
+      "step": 5361
+    },
+    {
+      "epoch": 0.55,
+      "gumbel_tau": 0.2823,
+      "image_masks_loss": 0.0267,
+      "llm_learning_rate": 2.2222382435335545e-06,
+      "loss": 0.7607,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00022222382435335545,
+      "step": 5362
+    },
+    {
+      "epoch": 0.55,
+      "gumbel_tau": 0.2822,
+      "image_masks_loss": 0.0277,
+      "llm_learning_rate": 2.221413778507888e-06,
+      "loss": 0.6357,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.00022214137785078877,
+      "step": 5363
+    },
+    {
+      "epoch": 0.55,
+      "gumbel_tau": 0.2822,
+      "image_masks_loss": 0.009,
+      "llm_learning_rate": 2.220589344160756e-06,
+      "loss": 0.7664,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.0002220589344160756,
+      "step": 5364
+    },
+    {
+      "epoch": 0.55,
+      "gumbel_tau": 0.2821,
+      "image_masks_loss": 0.0022,
+      "llm_learning_rate": 2.2197649405829486e-06,
+      "loss": 0.7488,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00022197649405829485,
+      "step": 5365
+    },
+    {
+      "epoch": 0.55,
+      "gumbel_tau": 0.282,
+      "image_masks_loss": 0.0122,
+      "llm_learning_rate": 2.218940567865249e-06,
+      "loss": 0.7108,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00022189405678652488,
+      "step": 5366
+    },
+    {
+      "epoch": 0.55,
+      "gumbel_tau": 0.282,
+      "image_masks_loss": 0.0113,
+      "llm_learning_rate": 2.2181162260984407e-06,
+      "loss": 0.7588,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00022181162260984406,
+      "step": 5367
+    },
+    {
+      "epoch": 0.55,
+      "gumbel_tau": 0.2819,
+      "image_masks_loss": 0.006,
+      "llm_learning_rate": 2.2172919153733017e-06,
+      "loss": 0.7031,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.00022172919153733015,
+      "step": 5368
+    },
+    {
+      "epoch": 0.55,
+      "gumbel_tau": 0.2818,
+      "image_masks_loss": 0.0085,
+      "llm_learning_rate": 2.216467635780606e-06,
+      "loss": 0.7014,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 0.00022164676357806057,
+      "step": 5369
+    },
+    {
+      "epoch": 0.55,
+      "gumbel_tau": 0.2818,
+      "image_masks_loss": 0.0161,
+      "llm_learning_rate": 2.2156433874111265e-06,
+      "loss": 0.826,
+      "output_text_mask_loss": 0.0055,
+      "predictor_learning_rate": 0.00022156433874111263,
+      "step": 5370
+    },
+    {
+      "epoch": 0.55,
+      "gumbel_tau": 0.2817,
+      "image_masks_loss": 0.0105,
+      "llm_learning_rate": 2.21481917035563e-06,
+      "loss": 0.7554,
+      "output_text_mask_loss": 0.0132,
+      "predictor_learning_rate": 0.000221481917035563,
+      "step": 5371
+    },
+    {
+      "epoch": 0.55,
+      "gumbel_tau": 0.2816,
+      "image_masks_loss": 0.0093,
+      "llm_learning_rate": 2.213994984704882e-06,
+      "loss": 0.7789,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.00022139949847048822,
+      "step": 5372
+    },
+    {
+      "epoch": 0.55,
+      "gumbel_tau": 0.2816,
+      "image_masks_loss": 0.0137,
+      "llm_learning_rate": 2.213170830549643e-06,
+      "loss": 0.7448,
+      "output_text_mask_loss": 0.0105,
+      "predictor_learning_rate": 0.0002213170830549643,
+      "step": 5373
+    },
+    {
+      "epoch": 0.55,
+      "gumbel_tau": 0.2815,
+      "image_masks_loss": 0.0097,
+      "llm_learning_rate": 2.2123467079806716e-06,
+      "loss": 0.7576,
+      "output_text_mask_loss": 0.0193,
+      "predictor_learning_rate": 0.00022123467079806713,
+      "step": 5374
+    },
+    {
+      "epoch": 0.55,
+      "gumbel_tau": 0.2814,
+      "image_masks_loss": 0.0172,
+      "llm_learning_rate": 2.2115226170887213e-06,
+      "loss": 0.7131,
+      "output_text_mask_loss": 0.0172,
+      "predictor_learning_rate": 0.00022115226170887212,
+      "step": 5375
+    },
+    {
+      "epoch": 0.55,
+      "gumbel_tau": 0.2814,
+      "image_masks_loss": 0.0026,
+      "llm_learning_rate": 2.2106985579645423e-06,
+      "loss": 0.6949,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.00022106985579645422,
+      "step": 5376
+    },
+    {
+      "epoch": 0.55,
+      "gumbel_tau": 0.2813,
+      "image_masks_loss": 0.0134,
+      "llm_learning_rate": 2.2098745306988834e-06,
+      "loss": 0.7037,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.00022098745306988832,
+      "step": 5377
+    },
+    {
+      "epoch": 0.55,
+      "gumbel_tau": 0.2812,
+      "image_masks_loss": 0.0078,
+      "llm_learning_rate": 2.209050535382487e-06,
+      "loss": 0.7269,
+      "output_text_mask_loss": 0.0252,
+      "predictor_learning_rate": 0.00022090505353824868,
+      "step": 5378
+    },
+    {
+      "epoch": 0.55,
+      "gumbel_tau": 0.2812,
+      "image_masks_loss": 0.0113,
+      "llm_learning_rate": 2.2082265721060944e-06,
+      "loss": 0.8059,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00022082265721060944,
+      "step": 5379
+    },
+    {
+      "epoch": 0.55,
+      "gumbel_tau": 0.2811,
+      "image_masks_loss": 0.0033,
+      "llm_learning_rate": 2.207402640960441e-06,
+      "loss": 0.7401,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.00022074026409604408,
+      "step": 5380
+    },
+    {
+      "epoch": 0.55,
+      "gumbel_tau": 0.281,
+      "image_masks_loss": 0.0139,
+      "llm_learning_rate": 2.2065787420362606e-06,
+      "loss": 0.7211,
+      "output_text_mask_loss": 0.0202,
+      "predictor_learning_rate": 0.00022065787420362605,
+      "step": 5381
+    },
+    {
+      "epoch": 0.55,
+      "gumbel_tau": 0.281,
+      "image_masks_loss": 0.0145,
+      "llm_learning_rate": 2.205754875424282e-06,
+      "loss": 0.6489,
+      "output_text_mask_loss": 0.0004,
+      "predictor_learning_rate": 0.0002205754875424282,
+      "step": 5382
+    },
+    {
+      "epoch": 0.55,
+      "gumbel_tau": 0.2809,
+      "image_masks_loss": 0.0119,
+      "llm_learning_rate": 2.204931041215233e-06,
+      "loss": 0.6727,
+      "output_text_mask_loss": 0.0238,
+      "predictor_learning_rate": 0.00022049310412152327,
+      "step": 5383
+    },
+    {
+      "epoch": 0.55,
+      "gumbel_tau": 0.2808,
+      "image_masks_loss": 0.0034,
+      "llm_learning_rate": 2.204107239499834e-06,
+      "loss": 0.7291,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.00022041072394998342,
+      "step": 5384
+    },
+    {
+      "epoch": 0.55,
+      "gumbel_tau": 0.2808,
+      "image_masks_loss": 0.0083,
+      "llm_learning_rate": 2.203283470368805e-06,
+      "loss": 0.7134,
+      "output_text_mask_loss": 0.0378,
+      "predictor_learning_rate": 0.00022032834703688048,
+      "step": 5385
+    },
+    {
+      "epoch": 0.55,
+      "gumbel_tau": 0.2807,
+      "image_masks_loss": 0.0101,
+      "llm_learning_rate": 2.202459733912862e-06,
+      "loss": 0.6835,
+      "output_text_mask_loss": 0.0462,
+      "predictor_learning_rate": 0.00022024597339128616,
+      "step": 5386
+    },
+    {
+      "epoch": 0.55,
+      "gumbel_tau": 0.2806,
+      "image_masks_loss": 0.0078,
+      "llm_learning_rate": 2.2016360302227153e-06,
+      "loss": 0.8069,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.00022016360302227152,
+      "step": 5387
+    },
+    {
+      "epoch": 0.55,
+      "gumbel_tau": 0.2806,
+      "image_masks_loss": 0.0071,
+      "llm_learning_rate": 2.2008123593890734e-06,
+      "loss": 0.7313,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.00022008123593890734,
+      "step": 5388
+    },
+    {
+      "epoch": 0.55,
+      "gumbel_tau": 0.2805,
+      "image_masks_loss": 0.0064,
+      "llm_learning_rate": 2.1999887215026405e-06,
+      "loss": 0.7756,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00021999887215026404,
+      "step": 5389
+    },
+    {
+      "epoch": 0.55,
+      "gumbel_tau": 0.2804,
+      "image_masks_loss": 0.0109,
+      "llm_learning_rate": 2.199165116654119e-06,
+      "loss": 0.7355,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00021991651166541185,
+      "step": 5390
+    },
+    {
+      "epoch": 0.55,
+      "gumbel_tau": 0.2804,
+      "image_masks_loss": 0.0107,
+      "llm_learning_rate": 2.198341544934204e-06,
+      "loss": 0.8317,
+      "output_text_mask_loss": 0.0116,
+      "predictor_learning_rate": 0.00021983415449342039,
+      "step": 5391
+    },
+    {
+      "epoch": 0.55,
+      "gumbel_tau": 0.2803,
+      "image_masks_loss": 0.0079,
+      "llm_learning_rate": 2.19751800643359e-06,
+      "loss": 0.7839,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00021975180064335897,
+      "step": 5392
+    },
+    {
+      "epoch": 0.55,
+      "gumbel_tau": 0.2802,
+      "image_masks_loss": 0.01,
+      "llm_learning_rate": 2.196694501242967e-06,
+      "loss": 0.8349,
+      "output_text_mask_loss": 0.0185,
+      "predictor_learning_rate": 0.0002196694501242967,
+      "step": 5393
+    },
+    {
+      "epoch": 0.55,
+      "gumbel_tau": 0.2802,
+      "image_masks_loss": 0.0162,
+      "llm_learning_rate": 2.195871029453022e-06,
+      "loss": 0.5792,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00021958710294530214,
+      "step": 5394
+    },
+    {
+      "epoch": 0.55,
+      "gumbel_tau": 0.2801,
+      "image_masks_loss": 0.0064,
+      "llm_learning_rate": 2.1950475911544354e-06,
+      "loss": 0.7323,
+      "output_text_mask_loss": 0.0732,
+      "predictor_learning_rate": 0.00021950475911544355,
+      "step": 5395
+    },
+    {
+      "epoch": 0.55,
+      "gumbel_tau": 0.28,
+      "image_masks_loss": 0.0133,
+      "llm_learning_rate": 2.1942241864378876e-06,
+      "loss": 0.7373,
+      "output_text_mask_loss": 0.0076,
+      "predictor_learning_rate": 0.00021942241864378876,
+      "step": 5396
+    },
+    {
+      "epoch": 0.55,
+      "gumbel_tau": 0.28,
+      "image_masks_loss": 0.0041,
+      "llm_learning_rate": 2.193400815394054e-06,
+      "loss": 0.7196,
+      "output_text_mask_loss": 0.0158,
+      "predictor_learning_rate": 0.00021934008153940538,
+      "step": 5397
+    },
+    {
+      "epoch": 0.55,
+      "gumbel_tau": 0.2799,
+      "image_masks_loss": 0.0048,
+      "llm_learning_rate": 2.192577478113606e-06,
+      "loss": 0.8039,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 0.00021925774781136057,
+      "step": 5398
+    },
+    {
+      "epoch": 0.55,
+      "gumbel_tau": 0.2799,
+      "image_masks_loss": 0.0134,
+      "llm_learning_rate": 2.19175417468721e-06,
+      "loss": 0.763,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.000219175417468721,
+      "step": 5399
+    },
+    {
+      "epoch": 0.55,
+      "gumbel_tau": 0.2798,
+      "image_masks_loss": 0.0338,
+      "llm_learning_rate": 2.190930905205531e-06,
+      "loss": 0.8064,
+      "output_text_mask_loss": 0.0006,
+      "predictor_learning_rate": 0.0002190930905205531,
+      "step": 5400
+    },
+    {
+      "epoch": 0.55,
+      "gumbel_tau": 0.2797,
+      "image_masks_loss": 0.0079,
+      "llm_learning_rate": 2.190107669759229e-06,
+      "loss": 0.7387,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.00021901076697592288,
+      "step": 5401
+    },
+    {
+      "epoch": 0.55,
+      "gumbel_tau": 0.2797,
+      "image_masks_loss": 0.0203,
+      "llm_learning_rate": 2.1892844684389608e-06,
+      "loss": 0.815,
+      "output_text_mask_loss": 0.0081,
+      "predictor_learning_rate": 0.00021892844684389607,
+      "step": 5402
+    },
+    {
+      "epoch": 0.55,
+      "gumbel_tau": 0.2796,
+      "image_masks_loss": 0.0076,
+      "llm_learning_rate": 2.1884613013353793e-06,
+      "loss": 0.6369,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0002188461301335379,
+      "step": 5403
+    },
+    {
+      "epoch": 0.55,
+      "gumbel_tau": 0.2795,
+      "image_masks_loss": 0.0088,
+      "llm_learning_rate": 2.187638168539133e-06,
+      "loss": 0.7783,
+      "output_text_mask_loss": 0.0404,
+      "predictor_learning_rate": 0.00021876381685391328,
+      "step": 5404
+    },
+    {
+      "epoch": 0.55,
+      "gumbel_tau": 0.2795,
+      "image_masks_loss": 0.0091,
+      "llm_learning_rate": 2.186815070140867e-06,
+      "loss": 0.7526,
+      "output_text_mask_loss": 0.0004,
+      "predictor_learning_rate": 0.00021868150701408667,
+      "step": 5405
+    },
+    {
+      "epoch": 0.55,
+      "gumbel_tau": 0.2794,
+      "image_masks_loss": 0.0093,
+      "llm_learning_rate": 2.1859920062312234e-06,
+      "loss": 0.7494,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.00021859920062312232,
+      "step": 5406
+    },
+    {
+      "epoch": 0.55,
+      "gumbel_tau": 0.2793,
+      "image_masks_loss": 0.0074,
+      "llm_learning_rate": 2.1851689769008394e-06,
+      "loss": 0.7733,
+      "output_text_mask_loss": 0.0012,
+      "predictor_learning_rate": 0.00021851689769008392,
+      "step": 5407
+    },
+    {
+      "epoch": 0.55,
+      "gumbel_tau": 0.2793,
+      "image_masks_loss": 0.0193,
+      "llm_learning_rate": 2.1843459822403487e-06,
+      "loss": 0.7144,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00021843459822403486,
+      "step": 5408
+    },
+    {
+      "epoch": 0.55,
+      "gumbel_tau": 0.2792,
+      "image_masks_loss": 0.0065,
+      "llm_learning_rate": 2.183523022340382e-06,
+      "loss": 0.7921,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.00021835230223403817,
+      "step": 5409
+    },
+    {
+      "epoch": 0.55,
+      "gumbel_tau": 0.2791,
+      "image_masks_loss": 0.0077,
+      "llm_learning_rate": 2.182700097291565e-06,
+      "loss": 0.7421,
+      "output_text_mask_loss": 0.0141,
+      "predictor_learning_rate": 0.00021827000972915648,
+      "step": 5410
+    },
+    {
+      "epoch": 0.55,
+      "gumbel_tau": 0.2791,
+      "image_masks_loss": 0.012,
+      "llm_learning_rate": 2.18187720718452e-06,
+      "loss": 0.8243,
+      "output_text_mask_loss": 0.0067,
+      "predictor_learning_rate": 0.00021818772071845196,
+      "step": 5411
+    },
+    {
+      "epoch": 0.55,
+      "gumbel_tau": 0.279,
+      "image_masks_loss": 0.0074,
+      "llm_learning_rate": 2.181054352109865e-06,
+      "loss": 0.7481,
+      "output_text_mask_loss": 0.0032,
+      "predictor_learning_rate": 0.00021810543521098646,
+      "step": 5412
+    },
+    {
+      "epoch": 0.55,
+      "gumbel_tau": 0.2789,
+      "image_masks_loss": 0.0111,
+      "llm_learning_rate": 2.1802315321582155e-06,
+      "loss": 0.7717,
+      "output_text_mask_loss": 0.0156,
+      "predictor_learning_rate": 0.00021802315321582156,
+      "step": 5413
+    },
+    {
+      "epoch": 0.55,
+      "gumbel_tau": 0.2789,
+      "image_masks_loss": 0.0123,
+      "llm_learning_rate": 2.179408747420182e-06,
+      "loss": 0.7399,
+      "output_text_mask_loss": 0.0146,
+      "predictor_learning_rate": 0.0002179408747420182,
+      "step": 5414
+    },
+    {
+      "epoch": 0.55,
+      "gumbel_tau": 0.2788,
+      "image_masks_loss": 0.036,
+      "llm_learning_rate": 2.1785859979863713e-06,
+      "loss": 0.7698,
+      "output_text_mask_loss": 0.004,
+      "predictor_learning_rate": 0.00021785859979863708,
+      "step": 5415
+    },
+    {
+      "epoch": 0.55,
+      "gumbel_tau": 0.2787,
+      "image_masks_loss": 0.0062,
+      "llm_learning_rate": 2.1777632839473873e-06,
+      "loss": 0.6164,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00021777632839473874,
+      "step": 5416
+    },
+    {
+      "epoch": 0.56,
+      "gumbel_tau": 0.2787,
+      "image_masks_loss": 0.0515,
+      "llm_learning_rate": 2.176940605393827e-06,
+      "loss": 0.6601,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00021769406053938272,
+      "step": 5417
+    },
+    {
+      "epoch": 0.56,
+      "gumbel_tau": 0.2786,
+      "image_masks_loss": 0.0319,
+      "llm_learning_rate": 2.1761179624162874e-06,
+      "loss": 0.8302,
+      "output_text_mask_loss": 0.0221,
+      "predictor_learning_rate": 0.00021761179624162874,
+      "step": 5418
+    },
+    {
+      "epoch": 0.56,
+      "gumbel_tau": 0.2785,
+      "image_masks_loss": 0.0102,
+      "llm_learning_rate": 2.175295355105359e-06,
+      "loss": 0.7011,
+      "output_text_mask_loss": 0.006,
+      "predictor_learning_rate": 0.0002175295355105359,
+      "step": 5419
+    },
+    {
+      "epoch": 0.56,
+      "gumbel_tau": 0.2785,
+      "image_masks_loss": 0.0056,
+      "llm_learning_rate": 2.1744727835516295e-06,
+      "loss": 0.754,
+      "output_text_mask_loss": 0.016,
+      "predictor_learning_rate": 0.00021744727835516296,
+      "step": 5420
+    },
+    {
+      "epoch": 0.56,
+      "gumbel_tau": 0.2784,
+      "image_masks_loss": 0.0187,
+      "llm_learning_rate": 2.1736502478456816e-06,
+      "loss": 0.8109,
+      "output_text_mask_loss": 0.0023,
+      "predictor_learning_rate": 0.00021736502478456816,
+      "step": 5421
+    },
+    {
+      "epoch": 0.56,
+      "gumbel_tau": 0.2783,
+      "image_masks_loss": 0.0131,
+      "llm_learning_rate": 2.172827748078096e-06,
+      "loss": 0.6576,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 0.00021728277480780958,
+      "step": 5422
+    },
+    {
+      "epoch": 0.56,
+      "gumbel_tau": 0.2783,
+      "image_masks_loss": 0.0132,
+      "llm_learning_rate": 2.1720052843394474e-06,
+      "loss": 0.7352,
+      "output_text_mask_loss": 0.046,
+      "predictor_learning_rate": 0.00021720052843394474,
+      "step": 5423
+    },
+    {
+      "epoch": 0.56,
+      "gumbel_tau": 0.2782,
+      "image_masks_loss": 0.0176,
+      "llm_learning_rate": 2.1711828567203068e-06,
+      "loss": 0.6749,
+      "output_text_mask_loss": 0.0026,
+      "predictor_learning_rate": 0.00021711828567203069,
+      "step": 5424
+    },
+    {
+      "epoch": 0.56,
+      "gumbel_tau": 0.2781,
+      "image_masks_loss": 0.015,
+      "llm_learning_rate": 2.170360465311243e-06,
+      "loss": 0.7376,
+      "output_text_mask_loss": 0.0109,
+      "predictor_learning_rate": 0.0002170360465311243,
+      "step": 5425
+    },
+    {
+      "epoch": 0.56,
+      "gumbel_tau": 0.2781,
+      "image_masks_loss": 0.0193,
+      "llm_learning_rate": 2.169538110202819e-06,
+      "loss": 0.842,
+      "output_text_mask_loss": 0.0088,
+      "predictor_learning_rate": 0.0002169538110202819,
+      "step": 5426
+    },
+    {
+      "epoch": 0.56,
+      "gumbel_tau": 0.278,
+      "image_masks_loss": 0.0061,
+      "llm_learning_rate": 2.1687157914855946e-06,
+      "loss": 0.7216,
+      "output_text_mask_loss": 0.0671,
+      "predictor_learning_rate": 0.00021687157914855945,
+      "step": 5427
+    },
+    {
+      "epoch": 0.56,
+      "gumbel_tau": 0.2779,
+      "image_masks_loss": 0.0109,
+      "llm_learning_rate": 2.1678935092501243e-06,
+      "loss": 0.7973,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.0002167893509250124,
+      "step": 5428
+    },
+    {
+      "epoch": 0.56,
+      "gumbel_tau": 0.2779,
+      "image_masks_loss": 0.0105,
+      "llm_learning_rate": 2.1670712635869617e-06,
+      "loss": 0.6381,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.00021670712635869614,
+      "step": 5429
+    },
+    {
+      "epoch": 0.56,
+      "gumbel_tau": 0.2778,
+      "image_masks_loss": 0.0143,
+      "llm_learning_rate": 2.166249054586652e-06,
+      "loss": 0.7844,
+      "output_text_mask_loss": 0.0553,
+      "predictor_learning_rate": 0.0002166249054586652,
+      "step": 5430
+    },
+    {
+      "epoch": 0.56,
+      "gumbel_tau": 0.2777,
+      "image_masks_loss": 0.0047,
+      "llm_learning_rate": 2.16542688233974e-06,
+      "loss": 0.725,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.000216542688233974,
+      "step": 5431
+    },
+    {
+      "epoch": 0.56,
+      "gumbel_tau": 0.2777,
+      "image_masks_loss": 0.0108,
+      "llm_learning_rate": 2.1646047469367656e-06,
+      "loss": 0.7506,
+      "output_text_mask_loss": 0.0115,
+      "predictor_learning_rate": 0.00021646047469367656,
+      "step": 5432
+    },
+    {
+      "epoch": 0.56,
+      "gumbel_tau": 0.2776,
+      "image_masks_loss": 0.0097,
+      "llm_learning_rate": 2.163782648468263e-06,
+      "loss": 0.7414,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.00021637826484682634,
+      "step": 5433
+    },
+    {
+      "epoch": 0.56,
+      "gumbel_tau": 0.2775,
+      "image_masks_loss": 0.0196,
+      "llm_learning_rate": 2.162960587024765e-06,
+      "loss": 0.7376,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00021629605870247646,
+      "step": 5434
+    },
+    {
+      "epoch": 0.56,
+      "gumbel_tau": 0.2775,
+      "image_masks_loss": 0.009,
+      "llm_learning_rate": 2.1621385626967976e-06,
+      "loss": 0.8098,
+      "output_text_mask_loss": 0.0012,
+      "predictor_learning_rate": 0.00021621385626967974,
+      "step": 5435
+    },
+    {
+      "epoch": 0.56,
+      "gumbel_tau": 0.2774,
+      "image_masks_loss": 0.0158,
+      "llm_learning_rate": 2.161316575574884e-06,
+      "loss": 0.757,
+      "output_text_mask_loss": 0.0103,
+      "predictor_learning_rate": 0.00021613165755748838,
+      "step": 5436
+    },
+    {
+      "epoch": 0.56,
+      "gumbel_tau": 0.2774,
+      "image_masks_loss": 0.0216,
+      "llm_learning_rate": 2.160494625749542e-06,
+      "loss": 0.8009,
+      "output_text_mask_loss": 0.0238,
+      "predictor_learning_rate": 0.00021604946257495423,
+      "step": 5437
+    },
+    {
+      "epoch": 0.56,
+      "gumbel_tau": 0.2773,
+      "image_masks_loss": 0.0076,
+      "llm_learning_rate": 2.15967271331129e-06,
+      "loss": 0.7557,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00021596727133112897,
+      "step": 5438
+    },
+    {
+      "epoch": 0.56,
+      "gumbel_tau": 0.2772,
+      "image_masks_loss": 0.0087,
+      "llm_learning_rate": 2.1588508383506364e-06,
+      "loss": 0.7761,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.00021588508383506362,
+      "step": 5439
+    },
+    {
+      "epoch": 0.56,
+      "gumbel_tau": 0.2772,
+      "image_masks_loss": 0.0254,
+      "llm_learning_rate": 2.1580290009580875e-06,
+      "loss": 0.7862,
+      "output_text_mask_loss": 0.0242,
+      "predictor_learning_rate": 0.00021580290009580872,
+      "step": 5440
+    },
+    {
+      "epoch": 0.56,
+      "gumbel_tau": 0.2771,
+      "image_masks_loss": 0.005,
+      "llm_learning_rate": 2.1572072012241474e-06,
+      "loss": 0.7478,
+      "output_text_mask_loss": 0.0193,
+      "predictor_learning_rate": 0.00021572072012241472,
+      "step": 5441
+    },
+    {
+      "epoch": 0.56,
+      "gumbel_tau": 0.277,
+      "image_masks_loss": 0.0186,
+      "llm_learning_rate": 2.156385439239314e-06,
+      "loss": 0.7686,
+      "output_text_mask_loss": 0.0116,
+      "predictor_learning_rate": 0.00021563854392393138,
+      "step": 5442
+    },
+    {
+      "epoch": 0.56,
+      "gumbel_tau": 0.277,
+      "image_masks_loss": 0.022,
+      "llm_learning_rate": 2.155563715094081e-06,
+      "loss": 0.7992,
+      "output_text_mask_loss": 0.0046,
+      "predictor_learning_rate": 0.0002155563715094081,
+      "step": 5443
+    },
+    {
+      "epoch": 0.56,
+      "gumbel_tau": 0.2769,
+      "image_masks_loss": 0.0149,
+      "llm_learning_rate": 2.154742028878939e-06,
+      "loss": 0.8007,
+      "output_text_mask_loss": 0.0603,
+      "predictor_learning_rate": 0.00021547420288789388,
+      "step": 5444
+    },
+    {
+      "epoch": 0.56,
+      "gumbel_tau": 0.2768,
+      "image_masks_loss": 0.0166,
+      "llm_learning_rate": 2.153920380684374e-06,
+      "loss": 0.7484,
+      "output_text_mask_loss": 0.0109,
+      "predictor_learning_rate": 0.0002153920380684374,
+      "step": 5445
+    },
+    {
+      "epoch": 0.56,
+      "gumbel_tau": 0.2768,
+      "image_masks_loss": 0.0118,
+      "llm_learning_rate": 2.1530987706008676e-06,
+      "loss": 0.6438,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.00021530987706008673,
+      "step": 5446
+    },
+    {
+      "epoch": 0.56,
+      "gumbel_tau": 0.2767,
+      "image_masks_loss": 0.0378,
+      "llm_learning_rate": 2.1522771987188966e-06,
+      "loss": 0.604,
+      "output_text_mask_loss": 0.0231,
+      "predictor_learning_rate": 0.00021522771987188963,
+      "step": 5447
+    },
+    {
+      "epoch": 0.56,
+      "gumbel_tau": 0.2766,
+      "image_masks_loss": 0.0214,
+      "llm_learning_rate": 2.1514556651289354e-06,
+      "loss": 0.7725,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.0002151455665128935,
+      "step": 5448
+    },
+    {
+      "epoch": 0.56,
+      "gumbel_tau": 0.2766,
+      "image_masks_loss": 0.0062,
+      "llm_learning_rate": 2.150634169921453e-06,
+      "loss": 0.6574,
+      "output_text_mask_loss": 0.013,
+      "predictor_learning_rate": 0.00021506341699214528,
+      "step": 5449
+    },
+    {
+      "epoch": 0.56,
+      "gumbel_tau": 0.2765,
+      "image_masks_loss": 0.0159,
+      "llm_learning_rate": 2.149812713186914e-06,
+      "loss": 0.7322,
+      "output_text_mask_loss": 0.0477,
+      "predictor_learning_rate": 0.00021498127131869138,
+      "step": 5450
+    },
+    {
+      "epoch": 0.56,
+      "gumbel_tau": 0.2764,
+      "image_masks_loss": 0.0091,
+      "llm_learning_rate": 2.148991295015778e-06,
+      "loss": 0.789,
+      "output_text_mask_loss": 0.0143,
+      "predictor_learning_rate": 0.0002148991295015778,
+      "step": 5451
+    },
+    {
+      "epoch": 0.56,
+      "gumbel_tau": 0.2764,
+      "image_masks_loss": 0.0216,
+      "llm_learning_rate": 2.1481699154985037e-06,
+      "loss": 0.7561,
+      "output_text_mask_loss": 0.0072,
+      "predictor_learning_rate": 0.00021481699154985036,
+      "step": 5452
+    },
+    {
+      "epoch": 0.56,
+      "gumbel_tau": 0.2763,
+      "image_masks_loss": 0.0157,
+      "llm_learning_rate": 2.147348574725542e-06,
+      "loss": 0.6721,
+      "output_text_mask_loss": 0.0927,
+      "predictor_learning_rate": 0.0002147348574725542,
+      "step": 5453
+    },
+    {
+      "epoch": 0.56,
+      "gumbel_tau": 0.2762,
+      "image_masks_loss": 0.0065,
+      "llm_learning_rate": 2.1465272727873404e-06,
+      "loss": 0.7346,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.00021465272727873403,
+      "step": 5454
+    },
+    {
+      "epoch": 0.56,
+      "gumbel_tau": 0.2762,
+      "image_masks_loss": 0.005,
+      "llm_learning_rate": 2.145706009774343e-06,
+      "loss": 0.7094,
+      "output_text_mask_loss": 0.0081,
+      "predictor_learning_rate": 0.00021457060097743427,
+      "step": 5455
+    },
+    {
+      "epoch": 0.56,
+      "gumbel_tau": 0.2761,
+      "image_masks_loss": 0.0093,
+      "llm_learning_rate": 2.144884785776988e-06,
+      "loss": 0.7235,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.00021448847857769883,
+      "step": 5456
+    },
+    {
+      "epoch": 0.56,
+      "gumbel_tau": 0.276,
+      "image_masks_loss": 0.0116,
+      "llm_learning_rate": 2.1440636008857123e-06,
+      "loss": 0.8198,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.00021440636008857123,
+      "step": 5457
+    },
+    {
+      "epoch": 0.56,
+      "gumbel_tau": 0.276,
+      "image_masks_loss": 0.0122,
+      "llm_learning_rate": 2.1432424551909454e-06,
+      "loss": 0.7757,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.00021432424551909453,
+      "step": 5458
+    },
+    {
+      "epoch": 0.56,
+      "gumbel_tau": 0.2759,
+      "image_masks_loss": 0.0188,
+      "llm_learning_rate": 2.1424213487831145e-06,
+      "loss": 0.7853,
+      "output_text_mask_loss": 0.0113,
+      "predictor_learning_rate": 0.00021424213487831144,
+      "step": 5459
+    },
+    {
+      "epoch": 0.56,
+      "gumbel_tau": 0.2759,
+      "image_masks_loss": 0.006,
+      "llm_learning_rate": 2.14160028175264e-06,
+      "loss": 0.7616,
+      "output_text_mask_loss": 0.0231,
+      "predictor_learning_rate": 0.000214160028175264,
+      "step": 5460
+    },
+    {
+      "epoch": 0.56,
+      "gumbel_tau": 0.2758,
+      "image_masks_loss": 0.0259,
+      "llm_learning_rate": 2.140779254189942e-06,
+      "loss": 0.7518,
+      "output_text_mask_loss": 0.0332,
+      "predictor_learning_rate": 0.0002140779254189942,
+      "step": 5461
+    },
+    {
+      "epoch": 0.56,
+      "gumbel_tau": 0.2757,
+      "image_masks_loss": 0.0136,
+      "llm_learning_rate": 2.139958266185432e-06,
+      "loss": 0.7981,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0002139958266185432,
+      "step": 5462
+    },
+    {
+      "epoch": 0.56,
+      "gumbel_tau": 0.2757,
+      "image_masks_loss": 0.023,
+      "llm_learning_rate": 2.1391373178295194e-06,
+      "loss": 0.7246,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.00021391373178295193,
+      "step": 5463
+    },
+    {
+      "epoch": 0.56,
+      "gumbel_tau": 0.2756,
+      "image_masks_loss": 0.0225,
+      "llm_learning_rate": 2.1383164092126096e-06,
+      "loss": 0.7967,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.00021383164092126096,
+      "step": 5464
+    },
+    {
+      "epoch": 0.56,
+      "gumbel_tau": 0.2755,
+      "image_masks_loss": 0.0046,
+      "llm_learning_rate": 2.1374955404251027e-06,
+      "loss": 0.857,
+      "output_text_mask_loss": 0.0054,
+      "predictor_learning_rate": 0.00021374955404251025,
+      "step": 5465
+    },
+    {
+      "epoch": 0.56,
+      "gumbel_tau": 0.2755,
+      "image_masks_loss": 0.0172,
+      "llm_learning_rate": 2.136674711557394e-06,
+      "loss": 0.7025,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0002136674711557394,
+      "step": 5466
+    },
+    {
+      "epoch": 0.56,
+      "gumbel_tau": 0.2754,
+      "image_masks_loss": 0.0183,
+      "llm_learning_rate": 2.135853922699875e-06,
+      "loss": 0.7417,
+      "output_text_mask_loss": 0.0042,
+      "predictor_learning_rate": 0.00021358539226998748,
+      "step": 5467
+    },
+    {
+      "epoch": 0.56,
+      "gumbel_tau": 0.2753,
+      "image_masks_loss": 0.0216,
+      "llm_learning_rate": 2.1350331739429334e-06,
+      "loss": 0.7016,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 0.00021350331739429336,
+      "step": 5468
+    },
+    {
+      "epoch": 0.56,
+      "gumbel_tau": 0.2753,
+      "image_masks_loss": 0.0121,
+      "llm_learning_rate": 2.134212465376952e-06,
+      "loss": 0.741,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0002134212465376952,
+      "step": 5469
+    },
+    {
+      "epoch": 0.56,
+      "gumbel_tau": 0.2752,
+      "image_masks_loss": 0.0263,
+      "llm_learning_rate": 2.1333917970923084e-06,
+      "loss": 0.7165,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.0002133391797092308,
+      "step": 5470
+    },
+    {
+      "epoch": 0.56,
+      "gumbel_tau": 0.2751,
+      "image_masks_loss": 0.0087,
+      "llm_learning_rate": 2.1325711691793766e-06,
+      "loss": 0.7212,
+      "output_text_mask_loss": 0.0062,
+      "predictor_learning_rate": 0.00021325711691793763,
+      "step": 5471
+    },
+    {
+      "epoch": 0.56,
+      "gumbel_tau": 0.2751,
+      "image_masks_loss": 0.0112,
+      "llm_learning_rate": 2.1317505817285265e-06,
+      "loss": 0.8562,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.00021317505817285265,
+      "step": 5472
+    },
+    {
+      "epoch": 0.56,
+      "gumbel_tau": 0.275,
+      "image_masks_loss": 0.0112,
+      "llm_learning_rate": 2.1309300348301225e-06,
+      "loss": 0.6663,
+      "output_text_mask_loss": 0.0137,
+      "predictor_learning_rate": 0.00021309300348301225,
+      "step": 5473
+    },
+    {
+      "epoch": 0.56,
+      "gumbel_tau": 0.2749,
+      "image_masks_loss": 0.0079,
+      "llm_learning_rate": 2.130109528574525e-06,
+      "loss": 0.762,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 0.0002130109528574525,
+      "step": 5474
+    },
+    {
+      "epoch": 0.56,
+      "gumbel_tau": 0.2749,
+      "image_masks_loss": 0.014,
+      "llm_learning_rate": 2.1292890630520903e-06,
+      "loss": 0.645,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.000212928906305209,
+      "step": 5475
+    },
+    {
+      "epoch": 0.56,
+      "gumbel_tau": 0.2748,
+      "image_masks_loss": 0.0118,
+      "llm_learning_rate": 2.1284686383531687e-06,
+      "loss": 0.6905,
+      "output_text_mask_loss": 0.0177,
+      "predictor_learning_rate": 0.00021284686383531688,
+      "step": 5476
+    },
+    {
+      "epoch": 0.56,
+      "gumbel_tau": 0.2747,
+      "image_masks_loss": 0.0179,
+      "llm_learning_rate": 2.1276482545681093e-06,
+      "loss": 0.741,
+      "output_text_mask_loss": 0.0052,
+      "predictor_learning_rate": 0.00021276482545681092,
+      "step": 5477
+    },
+    {
+      "epoch": 0.56,
+      "gumbel_tau": 0.2747,
+      "image_masks_loss": 0.0282,
+      "llm_learning_rate": 2.126827911787253e-06,
+      "loss": 0.7599,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0002126827911787253,
+      "step": 5478
+    },
+    {
+      "epoch": 0.56,
+      "gumbel_tau": 0.2746,
+      "image_masks_loss": 0.0051,
+      "llm_learning_rate": 2.126007610100938e-06,
+      "loss": 0.779,
+      "output_text_mask_loss": 0.0116,
+      "predictor_learning_rate": 0.0002126007610100938,
+      "step": 5479
+    },
+    {
+      "epoch": 0.56,
+      "gumbel_tau": 0.2746,
+      "image_masks_loss": 0.0093,
+      "llm_learning_rate": 2.1251873495994983e-06,
+      "loss": 0.7622,
+      "output_text_mask_loss": 0.0595,
+      "predictor_learning_rate": 0.00021251873495994983,
+      "step": 5480
+    },
+    {
+      "epoch": 0.56,
+      "gumbel_tau": 0.2745,
+      "image_masks_loss": 0.0103,
+      "llm_learning_rate": 2.1243671303732628e-06,
+      "loss": 0.6995,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0002124367130373263,
+      "step": 5481
+    },
+    {
+      "epoch": 0.56,
+      "gumbel_tau": 0.2744,
+      "image_masks_loss": 0.0077,
+      "llm_learning_rate": 2.123546952512556e-06,
+      "loss": 0.749,
+      "output_text_mask_loss": 0.0058,
+      "predictor_learning_rate": 0.00021235469525125557,
+      "step": 5482
+    },
+    {
+      "epoch": 0.56,
+      "gumbel_tau": 0.2744,
+      "image_masks_loss": 0.0244,
+      "llm_learning_rate": 2.122726816107696e-06,
+      "loss": 0.6501,
+      "output_text_mask_loss": 0.0402,
+      "predictor_learning_rate": 0.00021227268161076957,
+      "step": 5483
+    },
+    {
+      "epoch": 0.56,
+      "gumbel_tau": 0.2743,
+      "image_masks_loss": 0.0058,
+      "llm_learning_rate": 2.121906721249e-06,
+      "loss": 0.637,
+      "output_text_mask_loss": 0.0039,
+      "predictor_learning_rate": 0.0002121906721249,
+      "step": 5484
+    },
+    {
+      "epoch": 0.56,
+      "gumbel_tau": 0.2742,
+      "image_masks_loss": 0.0135,
+      "llm_learning_rate": 2.1210866680267783e-06,
+      "loss": 0.7473,
+      "output_text_mask_loss": 0.0233,
+      "predictor_learning_rate": 0.00021210866680267782,
+      "step": 5485
+    },
+    {
+      "epoch": 0.56,
+      "gumbel_tau": 0.2742,
+      "image_masks_loss": 0.0162,
+      "llm_learning_rate": 2.120266656531336e-06,
+      "loss": 0.6665,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.00021202666565313359,
+      "step": 5486
+    },
+    {
+      "epoch": 0.56,
+      "gumbel_tau": 0.2741,
+      "image_masks_loss": 0.028,
+      "llm_learning_rate": 2.119446686852976e-06,
+      "loss": 0.7596,
+      "output_text_mask_loss": 0.021,
+      "predictor_learning_rate": 0.00021194466868529758,
+      "step": 5487
+    },
+    {
+      "epoch": 0.56,
+      "gumbel_tau": 0.274,
+      "image_masks_loss": 0.0152,
+      "llm_learning_rate": 2.118626759081994e-06,
+      "loss": 0.7961,
+      "output_text_mask_loss": 0.0012,
+      "predictor_learning_rate": 0.00021186267590819942,
+      "step": 5488
+    },
+    {
+      "epoch": 0.56,
+      "gumbel_tau": 0.274,
+      "image_masks_loss": 0.01,
+      "llm_learning_rate": 2.1178068733086835e-06,
+      "loss": 0.7454,
+      "output_text_mask_loss": 0.0326,
+      "predictor_learning_rate": 0.00021178068733086834,
+      "step": 5489
+    },
+    {
+      "epoch": 0.56,
+      "gumbel_tau": 0.2739,
+      "image_masks_loss": 0.0172,
+      "llm_learning_rate": 2.116987029623331e-06,
+      "loss": 0.7145,
+      "output_text_mask_loss": 0.0012,
+      "predictor_learning_rate": 0.0002116987029623331,
+      "step": 5490
+    },
+    {
+      "epoch": 0.56,
+      "gumbel_tau": 0.2738,
+      "image_masks_loss": 0.0109,
+      "llm_learning_rate": 2.1161672281162204e-06,
+      "loss": 0.7395,
+      "output_text_mask_loss": 0.0019,
+      "predictor_learning_rate": 0.00021161672281162203,
+      "step": 5491
+    },
+    {
+      "epoch": 0.56,
+      "gumbel_tau": 0.2738,
+      "image_masks_loss": 0.0083,
+      "llm_learning_rate": 2.1153474688776288e-06,
+      "loss": 0.6253,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.00021153474688776286,
+      "step": 5492
+    },
+    {
+      "epoch": 0.56,
+      "gumbel_tau": 0.2737,
+      "image_masks_loss": 0.0022,
+      "llm_learning_rate": 2.114527751997831e-06,
+      "loss": 0.7198,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0002114527751997831,
+      "step": 5493
+    },
+    {
+      "epoch": 0.56,
+      "gumbel_tau": 0.2736,
+      "image_masks_loss": 0.0128,
+      "llm_learning_rate": 2.113708077567096e-06,
+      "loss": 0.7454,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 0.0002113708077567096,
+      "step": 5494
+    },
+    {
+      "epoch": 0.56,
+      "gumbel_tau": 0.2736,
+      "image_masks_loss": 0.0049,
+      "llm_learning_rate": 2.1128884456756875e-06,
+      "loss": 0.7221,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.00021128884456756875,
+      "step": 5495
+    },
+    {
+      "epoch": 0.56,
+      "gumbel_tau": 0.2735,
+      "image_masks_loss": 0.0196,
+      "llm_learning_rate": 2.1120688564138656e-06,
+      "loss": 0.8145,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.00021120688564138653,
+      "step": 5496
+    },
+    {
+      "epoch": 0.56,
+      "gumbel_tau": 0.2735,
+      "image_masks_loss": 0.0133,
+      "llm_learning_rate": 2.1112493098718856e-06,
+      "loss": 0.718,
+      "output_text_mask_loss": 0.0238,
+      "predictor_learning_rate": 0.00021112493098718855,
+      "step": 5497
+    },
+    {
+      "epoch": 0.56,
+      "gumbel_tau": 0.2734,
+      "image_masks_loss": 0.0171,
+      "llm_learning_rate": 2.1104298061399977e-06,
+      "loss": 0.7876,
+      "output_text_mask_loss": 0.0156,
+      "predictor_learning_rate": 0.00021104298061399974,
+      "step": 5498
+    },
+    {
+      "epoch": 0.56,
+      "gumbel_tau": 0.2733,
+      "image_masks_loss": 0.0072,
+      "llm_learning_rate": 2.1096103453084464e-06,
+      "loss": 0.8133,
+      "output_text_mask_loss": 0.0208,
+      "predictor_learning_rate": 0.00021096103453084462,
+      "step": 5499
+    },
+    {
+      "epoch": 0.56,
+      "gumbel_tau": 0.2733,
+      "image_masks_loss": 0.0117,
+      "llm_learning_rate": 2.108790927467474e-06,
+      "loss": 0.8562,
+      "output_text_mask_loss": 0.0282,
+      "predictor_learning_rate": 0.00021087909274674737,
+      "step": 5500
+    },
+    {
+      "epoch": 0.56,
+      "gumbel_tau": 0.2732,
+      "image_masks_loss": 0.0194,
+      "llm_learning_rate": 2.1079715527073163e-06,
+      "loss": 0.6845,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0002107971552707316,
+      "step": 5501
+    },
+    {
+      "epoch": 0.56,
+      "gumbel_tau": 0.2731,
+      "image_masks_loss": 0.0063,
+      "llm_learning_rate": 2.1071522211182044e-06,
+      "loss": 0.7716,
+      "output_text_mask_loss": 0.0309,
+      "predictor_learning_rate": 0.00021071522211182044,
+      "step": 5502
+    },
+    {
+      "epoch": 0.56,
+      "gumbel_tau": 0.2731,
+      "image_masks_loss": 0.01,
+      "llm_learning_rate": 2.1063329327903644e-06,
+      "loss": 0.7701,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 0.00021063329327903642,
+      "step": 5503
+    },
+    {
+      "epoch": 0.56,
+      "gumbel_tau": 0.273,
+      "image_masks_loss": 0.0175,
+      "llm_learning_rate": 2.1055136878140188e-06,
+      "loss": 0.6503,
+      "output_text_mask_loss": 0.0587,
+      "predictor_learning_rate": 0.00021055136878140188,
+      "step": 5504
+    },
+    {
+      "epoch": 0.56,
+      "gumbel_tau": 0.2729,
+      "image_masks_loss": 0.0095,
+      "llm_learning_rate": 2.104694486279385e-06,
+      "loss": 0.6632,
+      "output_text_mask_loss": 0.0105,
+      "predictor_learning_rate": 0.00021046944862793848,
+      "step": 5505
+    },
+    {
+      "epoch": 0.56,
+      "gumbel_tau": 0.2729,
+      "image_masks_loss": 0.01,
+      "llm_learning_rate": 2.1038753282766746e-06,
+      "loss": 0.7784,
+      "output_text_mask_loss": 0.0103,
+      "predictor_learning_rate": 0.00021038753282766743,
+      "step": 5506
+    },
+    {
+      "epoch": 0.56,
+      "gumbel_tau": 0.2728,
+      "image_masks_loss": 0.0043,
+      "llm_learning_rate": 2.1030562138960953e-06,
+      "loss": 0.6774,
+      "output_text_mask_loss": 0.0158,
+      "predictor_learning_rate": 0.0002103056213896095,
+      "step": 5507
+    },
+    {
+      "epoch": 0.56,
+      "gumbel_tau": 0.2727,
+      "image_masks_loss": 0.0033,
+      "llm_learning_rate": 2.1022371432278503e-06,
+      "loss": 0.6806,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.000210223714322785,
+      "step": 5508
+    },
+    {
+      "epoch": 0.56,
+      "gumbel_tau": 0.2727,
+      "image_masks_loss": 0.011,
+      "llm_learning_rate": 2.1014181163621375e-06,
+      "loss": 0.7741,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.00021014181163621373,
+      "step": 5509
+    },
+    {
+      "epoch": 0.56,
+      "gumbel_tau": 0.2726,
+      "image_masks_loss": 0.0136,
+      "llm_learning_rate": 2.100599133389149e-06,
+      "loss": 0.703,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00021005991333891488,
+      "step": 5510
+    },
+    {
+      "epoch": 0.56,
+      "gumbel_tau": 0.2726,
+      "image_masks_loss": 0.0078,
+      "llm_learning_rate": 2.0997801943990734e-06,
+      "loss": 0.6519,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00020997801943990733,
+      "step": 5511
+    },
+    {
+      "epoch": 0.56,
+      "gumbel_tau": 0.2725,
+      "image_masks_loss": 0.028,
+      "llm_learning_rate": 2.098961299482094e-06,
+      "loss": 0.6916,
+      "output_text_mask_loss": 0.0166,
+      "predictor_learning_rate": 0.0002098961299482094,
+      "step": 5512
+    },
+    {
+      "epoch": 0.56,
+      "gumbel_tau": 0.2724,
+      "image_masks_loss": 0.0126,
+      "llm_learning_rate": 2.09814244872839e-06,
+      "loss": 0.7694,
+      "output_text_mask_loss": 0.0086,
+      "predictor_learning_rate": 0.00020981424487283897,
+      "step": 5513
+    },
+    {
+      "epoch": 0.56,
+      "gumbel_tau": 0.2724,
+      "image_masks_loss": 0.0121,
+      "llm_learning_rate": 2.0973236422281346e-06,
+      "loss": 0.6969,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00020973236422281344,
+      "step": 5514
+    },
+    {
+      "epoch": 0.57,
+      "gumbel_tau": 0.2723,
+      "image_masks_loss": 0.0081,
+      "llm_learning_rate": 2.096504880071496e-06,
+      "loss": 0.7457,
+      "output_text_mask_loss": 0.0229,
+      "predictor_learning_rate": 0.00020965048800714958,
+      "step": 5515
+    },
+    {
+      "epoch": 0.57,
+      "gumbel_tau": 0.2722,
+      "image_masks_loss": 0.02,
+      "llm_learning_rate": 2.0956861623486393e-06,
+      "loss": 0.7357,
+      "output_text_mask_loss": 0.0107,
+      "predictor_learning_rate": 0.00020956861623486393,
+      "step": 5516
+    },
+    {
+      "epoch": 0.57,
+      "gumbel_tau": 0.2722,
+      "image_masks_loss": 0.0139,
+      "llm_learning_rate": 2.0948674891497234e-06,
+      "loss": 0.6781,
+      "output_text_mask_loss": 0.0216,
+      "predictor_learning_rate": 0.00020948674891497232,
+      "step": 5517
+    },
+    {
+      "epoch": 0.57,
+      "gumbel_tau": 0.2721,
+      "image_masks_loss": 0.0135,
+      "llm_learning_rate": 2.0940488605649015e-06,
+      "loss": 0.7357,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 0.00020940488605649014,
+      "step": 5518
+    },
+    {
+      "epoch": 0.57,
+      "gumbel_tau": 0.272,
+      "image_masks_loss": 0.0068,
+      "llm_learning_rate": 2.093230276684323e-06,
+      "loss": 0.7865,
+      "output_text_mask_loss": 0.0265,
+      "predictor_learning_rate": 0.00020932302766843227,
+      "step": 5519
+    },
+    {
+      "epoch": 0.57,
+      "gumbel_tau": 0.272,
+      "image_masks_loss": 0.0142,
+      "llm_learning_rate": 2.092411737598133e-06,
+      "loss": 0.6335,
+      "output_text_mask_loss": 0.0029,
+      "predictor_learning_rate": 0.00020924117375981327,
+      "step": 5520
+    },
+    {
+      "epoch": 0.57,
+      "gumbel_tau": 0.2719,
+      "image_masks_loss": 0.0054,
+      "llm_learning_rate": 2.0915932433964705e-06,
+      "loss": 0.776,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.000209159324339647,
+      "step": 5521
+    },
+    {
+      "epoch": 0.57,
+      "gumbel_tau": 0.2718,
+      "image_masks_loss": 0.0181,
+      "llm_learning_rate": 2.0907747941694686e-06,
+      "loss": 0.7144,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 0.00020907747941694685,
+      "step": 5522
+    },
+    {
+      "epoch": 0.57,
+      "gumbel_tau": 0.2718,
+      "image_masks_loss": 0.0073,
+      "llm_learning_rate": 2.089956390007259e-06,
+      "loss": 0.7689,
+      "output_text_mask_loss": 0.001,
+      "predictor_learning_rate": 0.00020899563900072588,
+      "step": 5523
+    },
+    {
+      "epoch": 0.57,
+      "gumbel_tau": 0.2717,
+      "image_masks_loss": 0.012,
+      "llm_learning_rate": 2.089138030999965e-06,
+      "loss": 0.8364,
+      "output_text_mask_loss": 0.0187,
+      "predictor_learning_rate": 0.00020891380309999647,
+      "step": 5524
+    },
+    {
+      "epoch": 0.57,
+      "gumbel_tau": 0.2717,
+      "image_masks_loss": 0.018,
+      "llm_learning_rate": 2.0883197172377056e-06,
+      "loss": 0.733,
+      "output_text_mask_loss": 0.0202,
+      "predictor_learning_rate": 0.00020883197172377055,
+      "step": 5525
+    },
+    {
+      "epoch": 0.57,
+      "gumbel_tau": 0.2716,
+      "image_masks_loss": 0.009,
+      "llm_learning_rate": 2.0875014488105962e-06,
+      "loss": 0.7116,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.0002087501448810596,
+      "step": 5526
+    },
+    {
+      "epoch": 0.57,
+      "gumbel_tau": 0.2715,
+      "image_masks_loss": 0.0138,
+      "llm_learning_rate": 2.0866832258087474e-06,
+      "loss": 0.6787,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.0002086683225808747,
+      "step": 5527
+    },
+    {
+      "epoch": 0.57,
+      "gumbel_tau": 0.2715,
+      "image_masks_loss": 0.0103,
+      "llm_learning_rate": 2.0858650483222605e-06,
+      "loss": 0.7924,
+      "output_text_mask_loss": 0.0019,
+      "predictor_learning_rate": 0.00020858650483222602,
+      "step": 5528
+    },
+    {
+      "epoch": 0.57,
+      "gumbel_tau": 0.2714,
+      "image_masks_loss": 0.0211,
+      "llm_learning_rate": 2.0850469164412376e-06,
+      "loss": 0.8467,
+      "output_text_mask_loss": 0.066,
+      "predictor_learning_rate": 0.00020850469164412376,
+      "step": 5529
+    },
+    {
+      "epoch": 0.57,
+      "gumbel_tau": 0.2713,
+      "image_masks_loss": 0.0121,
+      "llm_learning_rate": 2.0842288302557727e-06,
+      "loss": 0.8269,
+      "output_text_mask_loss": 0.0052,
+      "predictor_learning_rate": 0.00020842288302557724,
+      "step": 5530
+    },
+    {
+      "epoch": 0.57,
+      "gumbel_tau": 0.2713,
+      "image_masks_loss": 0.0069,
+      "llm_learning_rate": 2.0834107898559542e-06,
+      "loss": 0.7592,
+      "output_text_mask_loss": 0.0042,
+      "predictor_learning_rate": 0.00020834107898559542,
+      "step": 5531
+    },
+    {
+      "epoch": 0.57,
+      "gumbel_tau": 0.2712,
+      "image_masks_loss": 0.0092,
+      "llm_learning_rate": 2.0825927953318683e-06,
+      "loss": 0.7733,
+      "output_text_mask_loss": 0.0078,
+      "predictor_learning_rate": 0.00020825927953318683,
+      "step": 5532
+    },
+    {
+      "epoch": 0.57,
+      "gumbel_tau": 0.2711,
+      "image_masks_loss": 0.0056,
+      "llm_learning_rate": 2.0817748467735936e-06,
+      "loss": 0.7196,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.00020817748467735935,
+      "step": 5533
+    },
+    {
+      "epoch": 0.57,
+      "gumbel_tau": 0.2711,
+      "image_masks_loss": 0.01,
+      "llm_learning_rate": 2.080956944271204e-06,
+      "loss": 0.713,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0002080956944271204,
+      "step": 5534
+    },
+    {
+      "epoch": 0.57,
+      "gumbel_tau": 0.271,
+      "image_masks_loss": 0.0161,
+      "llm_learning_rate": 2.0801390879147687e-06,
+      "loss": 0.7209,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.00020801390879147686,
+      "step": 5535
+    },
+    {
+      "epoch": 0.57,
+      "gumbel_tau": 0.271,
+      "image_masks_loss": 0.0105,
+      "llm_learning_rate": 2.0793212777943527e-06,
+      "loss": 0.7938,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00020793212777943525,
+      "step": 5536
+    },
+    {
+      "epoch": 0.57,
+      "gumbel_tau": 0.2709,
+      "image_masks_loss": 0.0127,
+      "llm_learning_rate": 2.0785035140000147e-06,
+      "loss": 0.8135,
+      "output_text_mask_loss": 0.0103,
+      "predictor_learning_rate": 0.00020785035140000147,
+      "step": 5537
+    },
+    {
+      "epoch": 0.57,
+      "gumbel_tau": 0.2708,
+      "image_masks_loss": 0.0149,
+      "llm_learning_rate": 2.0776857966218083e-06,
+      "loss": 0.7542,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.00020776857966218082,
+      "step": 5538
+    },
+    {
+      "epoch": 0.57,
+      "gumbel_tau": 0.2708,
+      "image_masks_loss": 0.016,
+      "llm_learning_rate": 2.076868125749783e-06,
+      "loss": 0.8106,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 0.00020768681257497828,
+      "step": 5539
+    },
+    {
+      "epoch": 0.57,
+      "gumbel_tau": 0.2707,
+      "image_masks_loss": 0.0095,
+      "llm_learning_rate": 2.076050501473982e-06,
+      "loss": 0.7848,
+      "output_text_mask_loss": 0.001,
+      "predictor_learning_rate": 0.0002076050501473982,
+      "step": 5540
+    },
+    {
+      "epoch": 0.57,
+      "gumbel_tau": 0.2706,
+      "image_masks_loss": 0.0189,
+      "llm_learning_rate": 2.0752329238844444e-06,
+      "loss": 0.8134,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.00020752329238844441,
+      "step": 5541
+    },
+    {
+      "epoch": 0.57,
+      "gumbel_tau": 0.2706,
+      "image_masks_loss": 0.0079,
+      "llm_learning_rate": 2.0744153930712023e-06,
+      "loss": 0.7511,
+      "output_text_mask_loss": 0.1289,
+      "predictor_learning_rate": 0.00020744153930712022,
+      "step": 5542
+    },
+    {
+      "epoch": 0.57,
+      "gumbel_tau": 0.2705,
+      "image_masks_loss": 0.0093,
+      "llm_learning_rate": 2.0735979091242863e-06,
+      "loss": 0.6962,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.0002073597909124286,
+      "step": 5543
+    },
+    {
+      "epoch": 0.57,
+      "gumbel_tau": 0.2704,
+      "image_masks_loss": 0.0077,
+      "llm_learning_rate": 2.072780472133718e-06,
+      "loss": 0.7495,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00020727804721337175,
+      "step": 5544
+    },
+    {
+      "epoch": 0.57,
+      "gumbel_tau": 0.2704,
+      "image_masks_loss": 0.0086,
+      "llm_learning_rate": 2.0719630821895147e-06,
+      "loss": 0.6863,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00020719630821895146,
+      "step": 5545
+    },
+    {
+      "epoch": 0.57,
+      "gumbel_tau": 0.2703,
+      "image_masks_loss": 0.0107,
+      "llm_learning_rate": 2.0711457393816925e-06,
+      "loss": 0.7941,
+      "output_text_mask_loss": 0.002,
+      "predictor_learning_rate": 0.00020711457393816923,
+      "step": 5546
+    },
+    {
+      "epoch": 0.57,
+      "gumbel_tau": 0.2702,
+      "image_masks_loss": 0.0084,
+      "llm_learning_rate": 2.0703284438002545e-06,
+      "loss": 0.6579,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.00020703284438002546,
+      "step": 5547
+    },
+    {
+      "epoch": 0.57,
+      "gumbel_tau": 0.2702,
+      "image_masks_loss": 0.0194,
+      "llm_learning_rate": 2.0695111955352066e-06,
+      "loss": 0.7847,
+      "output_text_mask_loss": 0.0292,
+      "predictor_learning_rate": 0.00020695111955352065,
+      "step": 5548
+    },
+    {
+      "epoch": 0.57,
+      "gumbel_tau": 0.2701,
+      "image_masks_loss": 0.0144,
+      "llm_learning_rate": 2.0686939946765445e-06,
+      "loss": 0.7527,
+      "output_text_mask_loss": 0.0052,
+      "predictor_learning_rate": 0.00020686939946765442,
+      "step": 5549
+    },
+    {
+      "epoch": 0.57,
+      "gumbel_tau": 0.2701,
+      "image_masks_loss": 0.0075,
+      "llm_learning_rate": 2.06787684131426e-06,
+      "loss": 0.7653,
+      "output_text_mask_loss": 0.0027,
+      "predictor_learning_rate": 0.00020678768413142599,
+      "step": 5550
+    },
+    {
+      "epoch": 0.57,
+      "gumbel_tau": 0.27,
+      "image_masks_loss": 0.0219,
+      "llm_learning_rate": 2.0670597355383397e-06,
+      "loss": 0.7902,
+      "output_text_mask_loss": 0.0035,
+      "predictor_learning_rate": 0.00020670597355383397,
+      "step": 5551
+    },
+    {
+      "epoch": 0.57,
+      "gumbel_tau": 0.2699,
+      "image_masks_loss": 0.0116,
+      "llm_learning_rate": 2.066242677438767e-06,
+      "loss": 0.7887,
+      "output_text_mask_loss": 0.008,
+      "predictor_learning_rate": 0.00020662426774387666,
+      "step": 5552
+    },
+    {
+      "epoch": 0.57,
+      "gumbel_tau": 0.2699,
+      "image_masks_loss": 0.0114,
+      "llm_learning_rate": 2.065425667105516e-06,
+      "loss": 0.8002,
+      "output_text_mask_loss": 0.0782,
+      "predictor_learning_rate": 0.00020654256671055158,
+      "step": 5553
+    },
+    {
+      "epoch": 0.57,
+      "gumbel_tau": 0.2698,
+      "image_masks_loss": 0.0083,
+      "llm_learning_rate": 2.064608704628558e-06,
+      "loss": 0.7859,
+      "output_text_mask_loss": 0.0126,
+      "predictor_learning_rate": 0.0002064608704628558,
+      "step": 5554
+    },
+    {
+      "epoch": 0.57,
+      "gumbel_tau": 0.2697,
+      "image_masks_loss": 0.0147,
+      "llm_learning_rate": 2.0637917900978604e-06,
+      "loss": 0.7543,
+      "output_text_mask_loss": 0.0362,
+      "predictor_learning_rate": 0.000206379179009786,
+      "step": 5555
+    },
+    {
+      "epoch": 0.57,
+      "gumbel_tau": 0.2697,
+      "image_masks_loss": 0.0168,
+      "llm_learning_rate": 2.062974923603382e-06,
+      "loss": 0.7603,
+      "output_text_mask_loss": 0.0334,
+      "predictor_learning_rate": 0.0002062974923603382,
+      "step": 5556
+    },
+    {
+      "epoch": 0.57,
+      "gumbel_tau": 0.2696,
+      "image_masks_loss": 0.012,
+      "llm_learning_rate": 2.062158105235078e-06,
+      "loss": 0.7621,
+      "output_text_mask_loss": 0.0317,
+      "predictor_learning_rate": 0.0002062158105235078,
+      "step": 5557
+    },
+    {
+      "epoch": 0.57,
+      "gumbel_tau": 0.2695,
+      "image_masks_loss": 0.0041,
+      "llm_learning_rate": 2.0613413350828993e-06,
+      "loss": 0.7105,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00020613413350828992,
+      "step": 5558
+    },
+    {
+      "epoch": 0.57,
+      "gumbel_tau": 0.2695,
+      "image_masks_loss": 0.0057,
+      "llm_learning_rate": 2.0605246132367896e-06,
+      "loss": 0.8088,
+      "output_text_mask_loss": 0.0113,
+      "predictor_learning_rate": 0.00020605246132367894,
+      "step": 5559
+    },
+    {
+      "epoch": 0.57,
+      "gumbel_tau": 0.2694,
+      "image_masks_loss": 0.0123,
+      "llm_learning_rate": 2.0597079397866887e-06,
+      "loss": 0.8042,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00020597079397866887,
+      "step": 5560
+    },
+    {
+      "epoch": 0.57,
+      "gumbel_tau": 0.2694,
+      "image_masks_loss": 0.0069,
+      "llm_learning_rate": 2.0588913148225295e-06,
+      "loss": 0.7358,
+      "output_text_mask_loss": 0.0401,
+      "predictor_learning_rate": 0.00020588913148225294,
+      "step": 5561
+    },
+    {
+      "epoch": 0.57,
+      "gumbel_tau": 0.2693,
+      "image_masks_loss": 0.0123,
+      "llm_learning_rate": 2.058074738434242e-06,
+      "loss": 0.7566,
+      "output_text_mask_loss": 0.0038,
+      "predictor_learning_rate": 0.00020580747384342418,
+      "step": 5562
+    },
+    {
+      "epoch": 0.57,
+      "gumbel_tau": 0.2692,
+      "image_masks_loss": 0.0084,
+      "llm_learning_rate": 2.0572582107117482e-06,
+      "loss": 0.7953,
+      "output_text_mask_loss": 0.0004,
+      "predictor_learning_rate": 0.00020572582107117484,
+      "step": 5563
+    },
+    {
+      "epoch": 0.57,
+      "gumbel_tau": 0.2692,
+      "image_masks_loss": 0.0073,
+      "llm_learning_rate": 2.0564417317449674e-06,
+      "loss": 0.7706,
+      "output_text_mask_loss": 0.0172,
+      "predictor_learning_rate": 0.00020564417317449674,
+      "step": 5564
+    },
+    {
+      "epoch": 0.57,
+      "gumbel_tau": 0.2691,
+      "image_masks_loss": 0.0048,
+      "llm_learning_rate": 2.05562530162381e-06,
+      "loss": 0.7589,
+      "output_text_mask_loss": 0.0271,
+      "predictor_learning_rate": 0.00020556253016238102,
+      "step": 5565
+    },
+    {
+      "epoch": 0.57,
+      "gumbel_tau": 0.269,
+      "image_masks_loss": 0.0035,
+      "llm_learning_rate": 2.054808920438185e-06,
+      "loss": 0.718,
+      "output_text_mask_loss": 0.0324,
+      "predictor_learning_rate": 0.00020548089204381847,
+      "step": 5566
+    },
+    {
+      "epoch": 0.57,
+      "gumbel_tau": 0.269,
+      "image_masks_loss": 0.0201,
+      "llm_learning_rate": 2.053992588277992e-06,
+      "loss": 0.7549,
+      "output_text_mask_loss": 0.001,
+      "predictor_learning_rate": 0.0002053992588277992,
+      "step": 5567
+    },
+    {
+      "epoch": 0.57,
+      "gumbel_tau": 0.2689,
+      "image_masks_loss": 0.0072,
+      "llm_learning_rate": 2.0531763052331297e-06,
+      "loss": 0.6466,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00020531763052331295,
+      "step": 5568
+    },
+    {
+      "epoch": 0.57,
+      "gumbel_tau": 0.2688,
+      "image_masks_loss": 0.0118,
+      "llm_learning_rate": 2.0523600713934876e-06,
+      "loss": 0.7168,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00020523600713934876,
+      "step": 5569
+    },
+    {
+      "epoch": 0.57,
+      "gumbel_tau": 0.2688,
+      "image_masks_loss": 0.0078,
+      "llm_learning_rate": 2.0515438868489514e-06,
+      "loss": 0.7467,
+      "output_text_mask_loss": 0.0992,
+      "predictor_learning_rate": 0.00020515438868489512,
+      "step": 5570
+    },
+    {
+      "epoch": 0.57,
+      "gumbel_tau": 0.2687,
+      "image_masks_loss": 0.0104,
+      "llm_learning_rate": 2.0507277516894015e-06,
+      "loss": 0.7779,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 0.00020507277516894012,
+      "step": 5571
+    },
+    {
+      "epoch": 0.57,
+      "gumbel_tau": 0.2687,
+      "image_masks_loss": 0.0064,
+      "llm_learning_rate": 2.0499116660047123e-06,
+      "loss": 0.7658,
+      "output_text_mask_loss": 0.0124,
+      "predictor_learning_rate": 0.0002049911666004712,
+      "step": 5572
+    },
+    {
+      "epoch": 0.57,
+      "gumbel_tau": 0.2686,
+      "image_masks_loss": 0.0073,
+      "llm_learning_rate": 2.0490956298847527e-06,
+      "loss": 0.8216,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.00020490956298847526,
+      "step": 5573
+    },
+    {
+      "epoch": 0.57,
+      "gumbel_tau": 0.2685,
+      "image_masks_loss": 0.0112,
+      "llm_learning_rate": 2.0482796434193867e-06,
+      "loss": 0.6127,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.00020482796434193864,
+      "step": 5574
+    },
+    {
+      "epoch": 0.57,
+      "gumbel_tau": 0.2685,
+      "image_masks_loss": 0.0073,
+      "llm_learning_rate": 2.0474637066984724e-06,
+      "loss": 0.7082,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00020474637066984725,
+      "step": 5575
+    },
+    {
+      "epoch": 0.57,
+      "gumbel_tau": 0.2684,
+      "image_masks_loss": 0.0158,
+      "llm_learning_rate": 2.0466478198118632e-06,
+      "loss": 0.7979,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.0002046647819811863,
+      "step": 5576
+    },
+    {
+      "epoch": 0.57,
+      "gumbel_tau": 0.2683,
+      "image_masks_loss": 0.012,
+      "llm_learning_rate": 2.0458319828494056e-06,
+      "loss": 0.685,
+      "output_text_mask_loss": 0.0162,
+      "predictor_learning_rate": 0.00020458319828494054,
+      "step": 5577
+    },
+    {
+      "epoch": 0.57,
+      "gumbel_tau": 0.2683,
+      "image_masks_loss": 0.0075,
+      "llm_learning_rate": 2.045016195900942e-06,
+      "loss": 0.7944,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00020450161959009416,
+      "step": 5578
+    },
+    {
+      "epoch": 0.57,
+      "gumbel_tau": 0.2682,
+      "image_masks_loss": 0.0057,
+      "llm_learning_rate": 2.044200459056308e-06,
+      "loss": 0.7486,
+      "output_text_mask_loss": 0.0198,
+      "predictor_learning_rate": 0.0002044200459056308,
+      "step": 5579
+    },
+    {
+      "epoch": 0.57,
+      "gumbel_tau": 0.2682,
+      "image_masks_loss": 0.0083,
+      "llm_learning_rate": 2.043384772405335e-06,
+      "loss": 0.6546,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00020433847724053348,
+      "step": 5580
+    },
+    {
+      "epoch": 0.57,
+      "gumbel_tau": 0.2681,
+      "image_masks_loss": 0.0085,
+      "llm_learning_rate": 2.042569136037847e-06,
+      "loss": 0.7211,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0002042569136037847,
+      "step": 5581
+    },
+    {
+      "epoch": 0.57,
+      "gumbel_tau": 0.268,
+      "image_masks_loss": 0.0067,
+      "llm_learning_rate": 2.0417535500436658e-06,
+      "loss": 0.7307,
+      "output_text_mask_loss": 0.0032,
+      "predictor_learning_rate": 0.00020417535500436655,
+      "step": 5582
+    },
+    {
+      "epoch": 0.57,
+      "gumbel_tau": 0.268,
+      "image_masks_loss": 0.0064,
+      "llm_learning_rate": 2.0409380145126043e-06,
+      "loss": 0.8097,
+      "output_text_mask_loss": 0.0202,
+      "predictor_learning_rate": 0.00020409380145126042,
+      "step": 5583
+    },
+    {
+      "epoch": 0.57,
+      "gumbel_tau": 0.2679,
+      "image_masks_loss": 0.0113,
+      "llm_learning_rate": 2.040122529534471e-06,
+      "loss": 0.7493,
+      "output_text_mask_loss": 0.0183,
+      "predictor_learning_rate": 0.00020401225295344709,
+      "step": 5584
+    },
+    {
+      "epoch": 0.57,
+      "gumbel_tau": 0.2678,
+      "image_masks_loss": 0.0067,
+      "llm_learning_rate": 2.039307095199069e-06,
+      "loss": 0.7744,
+      "output_text_mask_loss": 0.0036,
+      "predictor_learning_rate": 0.0002039307095199069,
+      "step": 5585
+    },
+    {
+      "epoch": 0.57,
+      "gumbel_tau": 0.2678,
+      "image_masks_loss": 0.0159,
+      "llm_learning_rate": 2.0384917115961953e-06,
+      "loss": 0.7209,
+      "output_text_mask_loss": 0.0116,
+      "predictor_learning_rate": 0.00020384917115961952,
+      "step": 5586
+    },
+    {
+      "epoch": 0.57,
+      "gumbel_tau": 0.2677,
+      "image_masks_loss": 0.0042,
+      "llm_learning_rate": 2.037676378815643e-06,
+      "loss": 0.6725,
+      "output_text_mask_loss": 0.0324,
+      "predictor_learning_rate": 0.00020376763788156426,
+      "step": 5587
+    },
+    {
+      "epoch": 0.57,
+      "gumbel_tau": 0.2676,
+      "image_masks_loss": 0.0096,
+      "llm_learning_rate": 2.0368610969471974e-06,
+      "loss": 0.7567,
+      "output_text_mask_loss": 0.0116,
+      "predictor_learning_rate": 0.0002036861096947197,
+      "step": 5588
+    },
+    {
+      "epoch": 0.57,
+      "gumbel_tau": 0.2676,
+      "image_masks_loss": 0.0116,
+      "llm_learning_rate": 2.0360458660806395e-06,
+      "loss": 0.7471,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0002036045866080639,
+      "step": 5589
+    },
+    {
+      "epoch": 0.57,
+      "gumbel_tau": 0.2675,
+      "image_masks_loss": 0.0082,
+      "llm_learning_rate": 2.0352306863057435e-06,
+      "loss": 0.7268,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.0002035230686305743,
+      "step": 5590
+    },
+    {
+      "epoch": 0.57,
+      "gumbel_tau": 0.2675,
+      "image_masks_loss": 0.0044,
+      "llm_learning_rate": 2.03441555771228e-06,
+      "loss": 0.7874,
+      "output_text_mask_loss": 0.017,
+      "predictor_learning_rate": 0.000203441555771228,
+      "step": 5591
+    },
+    {
+      "epoch": 0.57,
+      "gumbel_tau": 0.2674,
+      "image_masks_loss": 0.0133,
+      "llm_learning_rate": 2.033600480390013e-06,
+      "loss": 0.7849,
+      "output_text_mask_loss": 0.0101,
+      "predictor_learning_rate": 0.00020336004803900128,
+      "step": 5592
+    },
+    {
+      "epoch": 0.57,
+      "gumbel_tau": 0.2673,
+      "image_masks_loss": 0.0042,
+      "llm_learning_rate": 2.032785454428699e-06,
+      "loss": 0.7588,
+      "output_text_mask_loss": 0.0172,
+      "predictor_learning_rate": 0.00020327854544286987,
+      "step": 5593
+    },
+    {
+      "epoch": 0.57,
+      "gumbel_tau": 0.2673,
+      "image_masks_loss": 0.0108,
+      "llm_learning_rate": 2.0319704799180918e-06,
+      "loss": 0.7858,
+      "output_text_mask_loss": 0.0114,
+      "predictor_learning_rate": 0.00020319704799180916,
+      "step": 5594
+    },
+    {
+      "epoch": 0.57,
+      "gumbel_tau": 0.2672,
+      "image_masks_loss": 0.0099,
+      "llm_learning_rate": 2.031155556947938e-06,
+      "loss": 0.729,
+      "output_text_mask_loss": 0.1038,
+      "predictor_learning_rate": 0.00020311555569479378,
+      "step": 5595
+    },
+    {
+      "epoch": 0.57,
+      "gumbel_tau": 0.2671,
+      "image_masks_loss": 0.0149,
+      "llm_learning_rate": 2.0303406856079786e-06,
+      "loss": 0.7526,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 0.00020303406856079788,
+      "step": 5596
+    },
+    {
+      "epoch": 0.57,
+      "gumbel_tau": 0.2671,
+      "image_masks_loss": 0.0097,
+      "llm_learning_rate": 2.0295258659879487e-06,
+      "loss": 0.6658,
+      "output_text_mask_loss": 0.001,
+      "predictor_learning_rate": 0.00020295258659879485,
+      "step": 5597
+    },
+    {
+      "epoch": 0.57,
+      "gumbel_tau": 0.267,
+      "image_masks_loss": 0.0106,
+      "llm_learning_rate": 2.0287110981775787e-06,
+      "loss": 0.812,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00020287110981775786,
+      "step": 5598
+    },
+    {
+      "epoch": 0.57,
+      "gumbel_tau": 0.267,
+      "image_masks_loss": 0.0209,
+      "llm_learning_rate": 2.0278963822665925e-06,
+      "loss": 0.7441,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00020278963822665922,
+      "step": 5599
+    },
+    {
+      "epoch": 0.57,
+      "gumbel_tau": 0.2669,
+      "image_masks_loss": 0.007,
+      "llm_learning_rate": 2.0270817183447073e-06,
+      "loss": 0.7755,
+      "output_text_mask_loss": 0.0214,
+      "predictor_learning_rate": 0.0002027081718344707,
+      "step": 5600
+    },
+    {
+      "epoch": 0.57,
+      "gumbel_tau": 0.2668,
+      "image_masks_loss": 0.0082,
+      "llm_learning_rate": 2.0262671065016386e-06,
+      "loss": 0.7051,
+      "output_text_mask_loss": 0.0072,
+      "predictor_learning_rate": 0.00020262671065016387,
+      "step": 5601
+    },
+    {
+      "epoch": 0.57,
+      "gumbel_tau": 0.2668,
+      "image_masks_loss": 0.0129,
+      "llm_learning_rate": 2.0254525468270895e-06,
+      "loss": 0.7515,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.00020254525468270897,
+      "step": 5602
+    },
+    {
+      "epoch": 0.57,
+      "gumbel_tau": 0.2667,
+      "image_masks_loss": 0.0133,
+      "llm_learning_rate": 2.0246380394107636e-06,
+      "loss": 0.744,
+      "output_text_mask_loss": 0.0217,
+      "predictor_learning_rate": 0.00020246380394107637,
+      "step": 5603
+    },
+    {
+      "epoch": 0.57,
+      "gumbel_tau": 0.2666,
+      "image_masks_loss": 0.0238,
+      "llm_learning_rate": 2.023823584342356e-06,
+      "loss": 0.7192,
+      "output_text_mask_loss": 0.0231,
+      "predictor_learning_rate": 0.00020238235843423558,
+      "step": 5604
+    },
+    {
+      "epoch": 0.57,
+      "gumbel_tau": 0.2666,
+      "image_masks_loss": 0.0171,
+      "llm_learning_rate": 2.0230091817115556e-06,
+      "loss": 0.6837,
+      "output_text_mask_loss": 0.0113,
+      "predictor_learning_rate": 0.00020230091817115555,
+      "step": 5605
+    },
+    {
+      "epoch": 0.57,
+      "gumbel_tau": 0.2665,
+      "image_masks_loss": 0.0088,
+      "llm_learning_rate": 2.022194831608046e-06,
+      "loss": 0.7612,
+      "output_text_mask_loss": 0.0208,
+      "predictor_learning_rate": 0.00020221948316080456,
+      "step": 5606
+    },
+    {
+      "epoch": 0.57,
+      "gumbel_tau": 0.2664,
+      "image_masks_loss": 0.0123,
+      "llm_learning_rate": 2.021380534121506e-06,
+      "loss": 0.7956,
+      "output_text_mask_loss": 0.0075,
+      "predictor_learning_rate": 0.0002021380534121506,
+      "step": 5607
+    },
+    {
+      "epoch": 0.57,
+      "gumbel_tau": 0.2664,
+      "image_masks_loss": 0.0049,
+      "llm_learning_rate": 2.0205662893416086e-06,
+      "loss": 0.7688,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 0.00020205662893416082,
+      "step": 5608
+    },
+    {
+      "epoch": 0.57,
+      "gumbel_tau": 0.2663,
+      "image_masks_loss": 0.004,
+      "llm_learning_rate": 2.0197520973580183e-06,
+      "loss": 0.806,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.0002019752097358018,
+      "step": 5609
+    },
+    {
+      "epoch": 0.57,
+      "gumbel_tau": 0.2663,
+      "image_masks_loss": 0.0041,
+      "llm_learning_rate": 2.0189379582603967e-06,
+      "loss": 0.728,
+      "output_text_mask_loss": 0.0036,
+      "predictor_learning_rate": 0.00020189379582603967,
+      "step": 5610
+    },
+    {
+      "epoch": 0.57,
+      "gumbel_tau": 0.2662,
+      "image_masks_loss": 0.0071,
+      "llm_learning_rate": 2.0181238721383992e-06,
+      "loss": 0.7435,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00020181238721383993,
+      "step": 5611
+    },
+    {
+      "epoch": 0.57,
+      "gumbel_tau": 0.2661,
+      "image_masks_loss": 0.013,
+      "llm_learning_rate": 2.0173098390816747e-06,
+      "loss": 0.8119,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.00020173098390816743,
+      "step": 5612
+    },
+    {
+      "epoch": 0.58,
+      "gumbel_tau": 0.2661,
+      "image_masks_loss": 0.0036,
+      "llm_learning_rate": 2.0164958591798645e-06,
+      "loss": 0.6248,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00020164958591798644,
+      "step": 5613
+    },
+    {
+      "epoch": 0.58,
+      "gumbel_tau": 0.266,
+      "image_masks_loss": 0.0138,
+      "llm_learning_rate": 2.015681932522608e-06,
+      "loss": 0.753,
+      "output_text_mask_loss": 0.0443,
+      "predictor_learning_rate": 0.0002015681932522608,
+      "step": 5614
+    },
+    {
+      "epoch": 0.58,
+      "gumbel_tau": 0.2659,
+      "image_masks_loss": 0.0019,
+      "llm_learning_rate": 2.0148680591995362e-06,
+      "loss": 0.8073,
+      "output_text_mask_loss": 0.0458,
+      "predictor_learning_rate": 0.00020148680591995362,
+      "step": 5615
+    },
+    {
+      "epoch": 0.58,
+      "gumbel_tau": 0.2659,
+      "image_masks_loss": 0.0076,
+      "llm_learning_rate": 2.0140542393002736e-06,
+      "loss": 0.7698,
+      "output_text_mask_loss": 0.021,
+      "predictor_learning_rate": 0.00020140542393002733,
+      "step": 5616
+    },
+    {
+      "epoch": 0.58,
+      "gumbel_tau": 0.2658,
+      "image_masks_loss": 0.0111,
+      "llm_learning_rate": 2.013240472914441e-06,
+      "loss": 0.6803,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.0002013240472914441,
+      "step": 5617
+    },
+    {
+      "epoch": 0.58,
+      "gumbel_tau": 0.2658,
+      "image_masks_loss": 0.0108,
+      "llm_learning_rate": 2.012426760131652e-06,
+      "loss": 0.7152,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.00020124267601316517,
+      "step": 5618
+    },
+    {
+      "epoch": 0.58,
+      "gumbel_tau": 0.2657,
+      "image_masks_loss": 0.015,
+      "llm_learning_rate": 2.011613101041514e-06,
+      "loss": 0.7357,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.00020116131010415136,
+      "step": 5619
+    },
+    {
+      "epoch": 0.58,
+      "gumbel_tau": 0.2656,
+      "image_masks_loss": 0.003,
+      "llm_learning_rate": 2.0107994957336295e-06,
+      "loss": 0.6621,
+      "output_text_mask_loss": 0.0336,
+      "predictor_learning_rate": 0.00020107994957336295,
+      "step": 5620
+    },
+    {
+      "epoch": 0.58,
+      "gumbel_tau": 0.2656,
+      "image_masks_loss": 0.0054,
+      "llm_learning_rate": 2.0099859442975937e-06,
+      "loss": 0.6905,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.00020099859442975938,
+      "step": 5621
+    },
+    {
+      "epoch": 0.58,
+      "gumbel_tau": 0.2655,
+      "image_masks_loss": 0.0196,
+      "llm_learning_rate": 2.009172446822997e-06,
+      "loss": 0.7959,
+      "output_text_mask_loss": 0.0088,
+      "predictor_learning_rate": 0.00020091724468229968,
+      "step": 5622
+    },
+    {
+      "epoch": 0.58,
+      "gumbel_tau": 0.2654,
+      "image_masks_loss": 0.0031,
+      "llm_learning_rate": 2.008359003399424e-06,
+      "loss": 0.7296,
+      "output_text_mask_loss": 0.0023,
+      "predictor_learning_rate": 0.0002008359003399424,
+      "step": 5623
+    },
+    {
+      "epoch": 0.58,
+      "gumbel_tau": 0.2654,
+      "image_masks_loss": 0.0115,
+      "llm_learning_rate": 2.007545614116453e-06,
+      "loss": 0.6737,
+      "output_text_mask_loss": 0.0244,
+      "predictor_learning_rate": 0.00020075456141164526,
+      "step": 5624
+    },
+    {
+      "epoch": 0.58,
+      "gumbel_tau": 0.2653,
+      "image_masks_loss": 0.0075,
+      "llm_learning_rate": 2.0067322790636546e-06,
+      "loss": 0.7175,
+      "output_text_mask_loss": 0.0042,
+      "predictor_learning_rate": 0.00020067322790636547,
+      "step": 5625
+    },
+    {
+      "epoch": 0.58,
+      "gumbel_tau": 0.2653,
+      "image_masks_loss": 0.008,
+      "llm_learning_rate": 2.0059189983305976e-06,
+      "loss": 0.688,
+      "output_text_mask_loss": 0.0086,
+      "predictor_learning_rate": 0.00020059189983305975,
+      "step": 5626
+    },
+    {
+      "epoch": 0.58,
+      "gumbel_tau": 0.2652,
+      "image_masks_loss": 0.0059,
+      "llm_learning_rate": 2.0051057720068405e-06,
+      "loss": 0.8137,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.00020051057720068405,
+      "step": 5627
+    },
+    {
+      "epoch": 0.58,
+      "gumbel_tau": 0.2651,
+      "image_masks_loss": 0.0078,
+      "llm_learning_rate": 2.0042926001819383e-06,
+      "loss": 0.6796,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.00020042926001819384,
+      "step": 5628
+    },
+    {
+      "epoch": 0.58,
+      "gumbel_tau": 0.2651,
+      "image_masks_loss": 0.0072,
+      "llm_learning_rate": 2.003479482945439e-06,
+      "loss": 0.7462,
+      "output_text_mask_loss": 0.0134,
+      "predictor_learning_rate": 0.00020034794829454388,
+      "step": 5629
+    },
+    {
+      "epoch": 0.58,
+      "gumbel_tau": 0.265,
+      "image_masks_loss": 0.0035,
+      "llm_learning_rate": 2.002666420386885e-06,
+      "loss": 0.7261,
+      "output_text_mask_loss": 0.0082,
+      "predictor_learning_rate": 0.0002002666420386885,
+      "step": 5630
+    },
+    {
+      "epoch": 0.58,
+      "gumbel_tau": 0.2649,
+      "image_masks_loss": 0.0024,
+      "llm_learning_rate": 2.001853412595813e-06,
+      "loss": 0.736,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.00020018534125958127,
+      "step": 5631
+    },
+    {
+      "epoch": 0.58,
+      "gumbel_tau": 0.2649,
+      "image_masks_loss": 0.0082,
+      "llm_learning_rate": 2.0010404596617516e-06,
+      "loss": 0.7761,
+      "output_text_mask_loss": 0.0137,
+      "predictor_learning_rate": 0.00020010404596617514,
+      "step": 5632
+    },
+    {
+      "epoch": 0.58,
+      "gumbel_tau": 0.2648,
+      "image_masks_loss": 0.0061,
+      "llm_learning_rate": 2.000227561674227e-06,
+      "loss": 0.7837,
+      "output_text_mask_loss": 0.0101,
+      "predictor_learning_rate": 0.0002000227561674227,
+      "step": 5633
+    },
+    {
+      "epoch": 0.58,
+      "gumbel_tau": 0.2648,
+      "image_masks_loss": 0.0148,
+      "llm_learning_rate": 1.999414718722757e-06,
+      "loss": 0.8049,
+      "output_text_mask_loss": 0.0343,
+      "predictor_learning_rate": 0.00019994147187227569,
+      "step": 5634
+    },
+    {
+      "epoch": 0.58,
+      "gumbel_tau": 0.2647,
+      "image_masks_loss": 0.0051,
+      "llm_learning_rate": 1.9986019308968527e-06,
+      "loss": 0.6771,
+      "output_text_mask_loss": 0.0275,
+      "predictor_learning_rate": 0.00019986019308968525,
+      "step": 5635
+    },
+    {
+      "epoch": 0.58,
+      "gumbel_tau": 0.2646,
+      "image_masks_loss": 0.0059,
+      "llm_learning_rate": 1.99778919828602e-06,
+      "loss": 0.7178,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 0.000199778919828602,
+      "step": 5636
+    },
+    {
+      "epoch": 0.58,
+      "gumbel_tau": 0.2646,
+      "image_masks_loss": 0.0117,
+      "llm_learning_rate": 1.99697652097976e-06,
+      "loss": 0.6994,
+      "output_text_mask_loss": 0.0439,
+      "predictor_learning_rate": 0.000199697652097976,
+      "step": 5637
+    },
+    {
+      "epoch": 0.58,
+      "gumbel_tau": 0.2645,
+      "image_masks_loss": 0.0062,
+      "llm_learning_rate": 1.9961638990675666e-06,
+      "loss": 0.7511,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.00019961638990675666,
+      "step": 5638
+    },
+    {
+      "epoch": 0.58,
+      "gumbel_tau": 0.2644,
+      "image_masks_loss": 0.0248,
+      "llm_learning_rate": 1.9953513326389263e-06,
+      "loss": 0.7358,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.0001995351332638926,
+      "step": 5639
+    },
+    {
+      "epoch": 0.58,
+      "gumbel_tau": 0.2644,
+      "image_masks_loss": 0.0099,
+      "llm_learning_rate": 1.9945388217833215e-06,
+      "loss": 0.8129,
+      "output_text_mask_loss": 0.0124,
+      "predictor_learning_rate": 0.00019945388217833214,
+      "step": 5640
+    },
+    {
+      "epoch": 0.58,
+      "gumbel_tau": 0.2643,
+      "image_masks_loss": 0.0058,
+      "llm_learning_rate": 1.9937263665902267e-06,
+      "loss": 0.6911,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.00019937263665902266,
+      "step": 5641
+    },
+    {
+      "epoch": 0.58,
+      "gumbel_tau": 0.2643,
+      "image_masks_loss": 0.0184,
+      "llm_learning_rate": 1.9929139671491125e-06,
+      "loss": 0.754,
+      "output_text_mask_loss": 0.0347,
+      "predictor_learning_rate": 0.00019929139671491126,
+      "step": 5642
+    },
+    {
+      "epoch": 0.58,
+      "gumbel_tau": 0.2642,
+      "image_masks_loss": 0.0068,
+      "llm_learning_rate": 1.9921016235494423e-06,
+      "loss": 0.7553,
+      "output_text_mask_loss": 0.0465,
+      "predictor_learning_rate": 0.00019921016235494423,
+      "step": 5643
+    },
+    {
+      "epoch": 0.58,
+      "gumbel_tau": 0.2641,
+      "image_masks_loss": 0.0167,
+      "llm_learning_rate": 1.991289335880672e-06,
+      "loss": 0.7491,
+      "output_text_mask_loss": 0.0496,
+      "predictor_learning_rate": 0.0001991289335880672,
+      "step": 5644
+    },
+    {
+      "epoch": 0.58,
+      "gumbel_tau": 0.2641,
+      "image_masks_loss": 0.0122,
+      "llm_learning_rate": 1.990477104232253e-06,
+      "loss": 0.7698,
+      "output_text_mask_loss": 0.001,
+      "predictor_learning_rate": 0.0001990477104232253,
+      "step": 5645
+    },
+    {
+      "epoch": 0.58,
+      "gumbel_tau": 0.264,
+      "image_masks_loss": 0.0078,
+      "llm_learning_rate": 1.9896649286936306e-06,
+      "loss": 0.6782,
+      "output_text_mask_loss": 0.001,
+      "predictor_learning_rate": 0.00019896649286936304,
+      "step": 5646
+    },
+    {
+      "epoch": 0.58,
+      "gumbel_tau": 0.2639,
+      "image_masks_loss": 0.0055,
+      "llm_learning_rate": 1.988852809354243e-06,
+      "loss": 0.7461,
+      "output_text_mask_loss": 0.0322,
+      "predictor_learning_rate": 0.0001988852809354243,
+      "step": 5647
+    },
+    {
+      "epoch": 0.58,
+      "gumbel_tau": 0.2639,
+      "image_masks_loss": 0.0171,
+      "llm_learning_rate": 1.988040746303522e-06,
+      "loss": 0.8162,
+      "output_text_mask_loss": 0.0145,
+      "predictor_learning_rate": 0.00019880407463035218,
+      "step": 5648
+    },
+    {
+      "epoch": 0.58,
+      "gumbel_tau": 0.2638,
+      "image_masks_loss": 0.0054,
+      "llm_learning_rate": 1.9872287396308947e-06,
+      "loss": 0.7299,
+      "output_text_mask_loss": 0.0467,
+      "predictor_learning_rate": 0.00019872287396308948,
+      "step": 5649
+    },
+    {
+      "epoch": 0.58,
+      "gumbel_tau": 0.2638,
+      "image_masks_loss": 0.0186,
+      "llm_learning_rate": 1.9864167894257814e-06,
+      "loss": 0.71,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00019864167894257813,
+      "step": 5650
+    },
+    {
+      "epoch": 0.58,
+      "gumbel_tau": 0.2637,
+      "image_masks_loss": 0.015,
+      "llm_learning_rate": 1.9856048957775946e-06,
+      "loss": 0.6416,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.00019856048957775943,
+      "step": 5651
+    },
+    {
+      "epoch": 0.58,
+      "gumbel_tau": 0.2636,
+      "image_masks_loss": 0.0175,
+      "llm_learning_rate": 1.984793058775742e-06,
+      "loss": 0.7492,
+      "output_text_mask_loss": 0.0139,
+      "predictor_learning_rate": 0.0001984793058775742,
+      "step": 5652
+    },
+    {
+      "epoch": 0.58,
+      "gumbel_tau": 0.2636,
+      "image_masks_loss": 0.0154,
+      "llm_learning_rate": 1.983981278509626e-06,
+      "loss": 0.8077,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.00019839812785096258,
+      "step": 5653
+    },
+    {
+      "epoch": 0.58,
+      "gumbel_tau": 0.2635,
+      "image_masks_loss": 0.0062,
+      "llm_learning_rate": 1.9831695550686407e-06,
+      "loss": 0.7333,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00019831695550686406,
+      "step": 5654
+    },
+    {
+      "epoch": 0.58,
+      "gumbel_tau": 0.2634,
+      "image_masks_loss": 0.0135,
+      "llm_learning_rate": 1.982357888542175e-06,
+      "loss": 0.7181,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.0001982357888542175,
+      "step": 5655
+    },
+    {
+      "epoch": 0.58,
+      "gumbel_tau": 0.2634,
+      "image_masks_loss": 0.0069,
+      "llm_learning_rate": 1.981546279019612e-06,
+      "loss": 0.8189,
+      "output_text_mask_loss": 0.0094,
+      "predictor_learning_rate": 0.0001981546279019612,
+      "step": 5656
+    },
+    {
+      "epoch": 0.58,
+      "gumbel_tau": 0.2633,
+      "image_masks_loss": 0.0095,
+      "llm_learning_rate": 1.9807347265903285e-06,
+      "loss": 0.6839,
+      "output_text_mask_loss": 0.0189,
+      "predictor_learning_rate": 0.00019807347265903284,
+      "step": 5657
+    },
+    {
+      "epoch": 0.58,
+      "gumbel_tau": 0.2633,
+      "image_masks_loss": 0.007,
+      "llm_learning_rate": 1.9799232313436925e-06,
+      "loss": 0.7637,
+      "output_text_mask_loss": 0.0139,
+      "predictor_learning_rate": 0.00019799232313436925,
+      "step": 5658
+    },
+    {
+      "epoch": 0.58,
+      "gumbel_tau": 0.2632,
+      "image_masks_loss": 0.0103,
+      "llm_learning_rate": 1.979111793369069e-06,
+      "loss": 0.6561,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0001979111793369069,
+      "step": 5659
+    },
+    {
+      "epoch": 0.58,
+      "gumbel_tau": 0.2631,
+      "image_masks_loss": 0.0196,
+      "llm_learning_rate": 1.9783004127558147e-06,
+      "loss": 0.6959,
+      "output_text_mask_loss": 0.0055,
+      "predictor_learning_rate": 0.00019783004127558147,
+      "step": 5660
+    },
+    {
+      "epoch": 0.58,
+      "gumbel_tau": 0.2631,
+      "image_masks_loss": 0.0322,
+      "llm_learning_rate": 1.9774890895932808e-06,
+      "loss": 0.7961,
+      "output_text_mask_loss": 0.0313,
+      "predictor_learning_rate": 0.00019774890895932806,
+      "step": 5661
+    },
+    {
+      "epoch": 0.58,
+      "gumbel_tau": 0.263,
+      "image_masks_loss": 0.0085,
+      "llm_learning_rate": 1.9766778239708126e-06,
+      "loss": 0.6848,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.00019766778239708124,
+      "step": 5662
+    },
+    {
+      "epoch": 0.58,
+      "gumbel_tau": 0.263,
+      "image_masks_loss": 0.0418,
+      "llm_learning_rate": 1.975866615977748e-06,
+      "loss": 0.7178,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.0001975866615977748,
+      "step": 5663
+    },
+    {
+      "epoch": 0.58,
+      "gumbel_tau": 0.2629,
+      "image_masks_loss": 0.0096,
+      "llm_learning_rate": 1.9750554657034186e-06,
+      "loss": 0.7026,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.00019750554657034186,
+      "step": 5664
+    },
+    {
+      "epoch": 0.58,
+      "gumbel_tau": 0.2628,
+      "image_masks_loss": 0.0061,
+      "llm_learning_rate": 1.9742443732371514e-06,
+      "loss": 0.8043,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0001974244373237151,
+      "step": 5665
+    },
+    {
+      "epoch": 0.58,
+      "gumbel_tau": 0.2628,
+      "image_masks_loss": 0.0067,
+      "llm_learning_rate": 1.9734333386682646e-06,
+      "loss": 0.7741,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.00019734333386682646,
+      "step": 5666
+    },
+    {
+      "epoch": 0.58,
+      "gumbel_tau": 0.2627,
+      "image_masks_loss": 0.0057,
+      "llm_learning_rate": 1.9726223620860712e-06,
+      "loss": 0.7591,
+      "output_text_mask_loss": 0.0231,
+      "predictor_learning_rate": 0.0001972622362086071,
+      "step": 5667
+    },
+    {
+      "epoch": 0.58,
+      "gumbel_tau": 0.2626,
+      "image_masks_loss": 0.0169,
+      "llm_learning_rate": 1.971811443579878e-06,
+      "loss": 0.6661,
+      "output_text_mask_loss": 0.0023,
+      "predictor_learning_rate": 0.0001971811443579878,
+      "step": 5668
+    },
+    {
+      "epoch": 0.58,
+      "gumbel_tau": 0.2626,
+      "image_masks_loss": 0.0072,
+      "llm_learning_rate": 1.9710005832389852e-06,
+      "loss": 0.7325,
+      "output_text_mask_loss": 0.0305,
+      "predictor_learning_rate": 0.00019710005832389852,
+      "step": 5669
+    },
+    {
+      "epoch": 0.58,
+      "gumbel_tau": 0.2625,
+      "image_masks_loss": 0.0149,
+      "llm_learning_rate": 1.970189781152687e-06,
+      "loss": 0.7232,
+      "output_text_mask_loss": 0.0534,
+      "predictor_learning_rate": 0.00019701897811526866,
+      "step": 5670
+    },
+    {
+      "epoch": 0.58,
+      "gumbel_tau": 0.2625,
+      "image_masks_loss": 0.0045,
+      "llm_learning_rate": 1.969379037410269e-06,
+      "loss": 0.6732,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00019693790374102687,
+      "step": 5671
+    },
+    {
+      "epoch": 0.58,
+      "gumbel_tau": 0.2624,
+      "image_masks_loss": 0.0023,
+      "llm_learning_rate": 1.9685683521010142e-06,
+      "loss": 0.843,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 0.0001968568352101014,
+      "step": 5672
+    },
+    {
+      "epoch": 0.58,
+      "gumbel_tau": 0.2623,
+      "image_masks_loss": 0.0102,
+      "llm_learning_rate": 1.967757725314196e-06,
+      "loss": 0.8339,
+      "output_text_mask_loss": 0.0568,
+      "predictor_learning_rate": 0.00019677577253141958,
+      "step": 5673
+    },
+    {
+      "epoch": 0.58,
+      "gumbel_tau": 0.2623,
+      "image_masks_loss": 0.0106,
+      "llm_learning_rate": 1.966947157139083e-06,
+      "loss": 0.7666,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.0001966947157139083,
+      "step": 5674
+    },
+    {
+      "epoch": 0.58,
+      "gumbel_tau": 0.2622,
+      "image_masks_loss": 0.0073,
+      "llm_learning_rate": 1.9661366476649367e-06,
+      "loss": 0.6432,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00019661366476649367,
+      "step": 5675
+    },
+    {
+      "epoch": 0.58,
+      "gumbel_tau": 0.2621,
+      "image_masks_loss": 0.0138,
+      "llm_learning_rate": 1.9653261969810113e-06,
+      "loss": 0.7665,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.00019653261969810112,
+      "step": 5676
+    },
+    {
+      "epoch": 0.58,
+      "gumbel_tau": 0.2621,
+      "image_masks_loss": 0.0067,
+      "llm_learning_rate": 1.9645158051765556e-06,
+      "loss": 0.6213,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 0.00019645158051765556,
+      "step": 5677
+    },
+    {
+      "epoch": 0.58,
+      "gumbel_tau": 0.262,
+      "image_masks_loss": 0.0142,
+      "llm_learning_rate": 1.963705472340813e-06,
+      "loss": 0.7474,
+      "output_text_mask_loss": 0.0039,
+      "predictor_learning_rate": 0.00019637054723408126,
+      "step": 5678
+    },
+    {
+      "epoch": 0.58,
+      "gumbel_tau": 0.262,
+      "image_masks_loss": 0.0192,
+      "llm_learning_rate": 1.9628951985630177e-06,
+      "loss": 0.7016,
+      "output_text_mask_loss": 0.0103,
+      "predictor_learning_rate": 0.00019628951985630174,
+      "step": 5679
+    },
+    {
+      "epoch": 0.58,
+      "gumbel_tau": 0.2619,
+      "image_masks_loss": 0.0077,
+      "llm_learning_rate": 1.962084983932399e-06,
+      "loss": 0.718,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0001962084983932399,
+      "step": 5680
+    },
+    {
+      "epoch": 0.58,
+      "gumbel_tau": 0.2618,
+      "image_masks_loss": 0.0195,
+      "llm_learning_rate": 1.961274828538181e-06,
+      "loss": 0.7623,
+      "output_text_mask_loss": 0.0198,
+      "predictor_learning_rate": 0.0001961274828538181,
+      "step": 5681
+    },
+    {
+      "epoch": 0.58,
+      "gumbel_tau": 0.2618,
+      "image_masks_loss": 0.0051,
+      "llm_learning_rate": 1.9604647324695786e-06,
+      "loss": 0.7141,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.00019604647324695786,
+      "step": 5682
+    },
+    {
+      "epoch": 0.58,
+      "gumbel_tau": 0.2617,
+      "image_masks_loss": 0.0098,
+      "llm_learning_rate": 1.9596546958158017e-06,
+      "loss": 0.7557,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00019596546958158015,
+      "step": 5683
+    },
+    {
+      "epoch": 0.58,
+      "gumbel_tau": 0.2617,
+      "image_masks_loss": 0.008,
+      "llm_learning_rate": 1.958844718666053e-06,
+      "loss": 0.7427,
+      "output_text_mask_loss": 0.0248,
+      "predictor_learning_rate": 0.00019588447186660526,
+      "step": 5684
+    },
+    {
+      "epoch": 0.58,
+      "gumbel_tau": 0.2616,
+      "image_masks_loss": 0.0232,
+      "llm_learning_rate": 1.9580348011095295e-06,
+      "loss": 0.8044,
+      "output_text_mask_loss": 0.0378,
+      "predictor_learning_rate": 0.00019580348011095293,
+      "step": 5685
+    },
+    {
+      "epoch": 0.58,
+      "gumbel_tau": 0.2615,
+      "image_masks_loss": 0.032,
+      "llm_learning_rate": 1.9572249432354206e-06,
+      "loss": 0.7383,
+      "output_text_mask_loss": 0.0387,
+      "predictor_learning_rate": 0.00019572249432354207,
+      "step": 5686
+    },
+    {
+      "epoch": 0.58,
+      "gumbel_tau": 0.2615,
+      "image_masks_loss": 0.0143,
+      "llm_learning_rate": 1.95641514513291e-06,
+      "loss": 0.7312,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00019564151451329095,
+      "step": 5687
+    },
+    {
+      "epoch": 0.58,
+      "gumbel_tau": 0.2614,
+      "image_masks_loss": 0.0055,
+      "llm_learning_rate": 1.9556054068911747e-06,
+      "loss": 0.7608,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00019556054068911743,
+      "step": 5688
+    },
+    {
+      "epoch": 0.58,
+      "gumbel_tau": 0.2613,
+      "image_masks_loss": 0.0767,
+      "llm_learning_rate": 1.954795728599384e-06,
+      "loss": 0.7754,
+      "output_text_mask_loss": 0.0202,
+      "predictor_learning_rate": 0.0001954795728599384,
+      "step": 5689
+    },
+    {
+      "epoch": 0.58,
+      "gumbel_tau": 0.2613,
+      "image_masks_loss": 0.0534,
+      "llm_learning_rate": 1.953986110346703e-06,
+      "loss": 0.6774,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 0.00019539861103467026,
+      "step": 5690
+    },
+    {
+      "epoch": 0.58,
+      "gumbel_tau": 0.2612,
+      "image_masks_loss": 0.0265,
+      "llm_learning_rate": 1.9531765522222863e-06,
+      "loss": 0.6504,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00019531765522222863,
+      "step": 5691
+    },
+    {
+      "epoch": 0.58,
+      "gumbel_tau": 0.2612,
+      "image_masks_loss": 0.0184,
+      "llm_learning_rate": 1.9523670543152867e-06,
+      "loss": 0.6799,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00019523670543152865,
+      "step": 5692
+    },
+    {
+      "epoch": 0.58,
+      "gumbel_tau": 0.2611,
+      "image_masks_loss": 0.033,
+      "llm_learning_rate": 1.951557616714847e-06,
+      "loss": 0.7005,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0001951557616714847,
+      "step": 5693
+    },
+    {
+      "epoch": 0.58,
+      "gumbel_tau": 0.261,
+      "image_masks_loss": 0.0051,
+      "llm_learning_rate": 1.9507482395101044e-06,
+      "loss": 0.7698,
+      "output_text_mask_loss": 0.0233,
+      "predictor_learning_rate": 0.00019507482395101044,
+      "step": 5694
+    },
+    {
+      "epoch": 0.58,
+      "gumbel_tau": 0.261,
+      "image_masks_loss": 0.0595,
+      "llm_learning_rate": 1.9499389227901887e-06,
+      "loss": 0.7298,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.00019499389227901887,
+      "step": 5695
+    },
+    {
+      "epoch": 0.58,
+      "gumbel_tau": 0.2609,
+      "image_masks_loss": 0.0549,
+      "llm_learning_rate": 1.9491296666442236e-06,
+      "loss": 0.7629,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00019491296666442236,
+      "step": 5696
+    },
+    {
+      "epoch": 0.58,
+      "gumbel_tau": 0.2609,
+      "image_masks_loss": 0.0202,
+      "llm_learning_rate": 1.948320471161327e-06,
+      "loss": 0.7803,
+      "output_text_mask_loss": 0.0202,
+      "predictor_learning_rate": 0.0001948320471161327,
+      "step": 5697
+    },
+    {
+      "epoch": 0.58,
+      "gumbel_tau": 0.2608,
+      "image_masks_loss": 0.0271,
+      "llm_learning_rate": 1.947511336430609e-06,
+      "loss": 0.7725,
+      "output_text_mask_loss": 0.0124,
+      "predictor_learning_rate": 0.0001947511336430609,
+      "step": 5698
+    },
+    {
+      "epoch": 0.58,
+      "gumbel_tau": 0.2607,
+      "image_masks_loss": 0.0301,
+      "llm_learning_rate": 1.9467022625411735e-06,
+      "loss": 0.8192,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.0001946702262541173,
+      "step": 5699
+    },
+    {
+      "epoch": 0.58,
+      "gumbel_tau": 0.2607,
+      "image_masks_loss": 0.0131,
+      "llm_learning_rate": 1.9458932495821165e-06,
+      "loss": 0.6976,
+      "output_text_mask_loss": 0.0267,
+      "predictor_learning_rate": 0.00019458932495821163,
+      "step": 5700
+    },
+    {
+      "epoch": 0.58,
+      "gumbel_tau": 0.2606,
+      "image_masks_loss": 0.0137,
+      "llm_learning_rate": 1.9450842976425295e-06,
+      "loss": 0.8007,
+      "output_text_mask_loss": 0.0254,
+      "predictor_learning_rate": 0.00019450842976425291,
+      "step": 5701
+    },
+    {
+      "epoch": 0.58,
+      "gumbel_tau": 0.2605,
+      "image_masks_loss": 0.0149,
+      "llm_learning_rate": 1.9442754068114954e-06,
+      "loss": 0.7769,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.00019442754068114954,
+      "step": 5702
+    },
+    {
+      "epoch": 0.58,
+      "gumbel_tau": 0.2605,
+      "image_masks_loss": 0.0126,
+      "llm_learning_rate": 1.9434665771780915e-06,
+      "loss": 0.72,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00019434665771780912,
+      "step": 5703
+    },
+    {
+      "epoch": 0.58,
+      "gumbel_tau": 0.2604,
+      "image_masks_loss": 0.0332,
+      "llm_learning_rate": 1.9426578088313877e-06,
+      "loss": 0.7941,
+      "output_text_mask_loss": 0.0062,
+      "predictor_learning_rate": 0.00019426578088313875,
+      "step": 5704
+    },
+    {
+      "epoch": 0.58,
+      "gumbel_tau": 0.2604,
+      "image_masks_loss": 0.0181,
+      "llm_learning_rate": 1.9418491018604477e-06,
+      "loss": 0.7701,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00019418491018604475,
+      "step": 5705
+    },
+    {
+      "epoch": 0.58,
+      "gumbel_tau": 0.2603,
+      "image_masks_loss": 0.0121,
+      "llm_learning_rate": 1.9410404563543278e-06,
+      "loss": 0.7387,
+      "output_text_mask_loss": 0.0092,
+      "predictor_learning_rate": 0.00019410404563543278,
+      "step": 5706
+    },
+    {
+      "epoch": 0.58,
+      "gumbel_tau": 0.2602,
+      "image_masks_loss": 0.0271,
+      "llm_learning_rate": 1.9402318724020773e-06,
+      "loss": 0.7217,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.00019402318724020772,
+      "step": 5707
+    },
+    {
+      "epoch": 0.58,
+      "gumbel_tau": 0.2602,
+      "image_masks_loss": 0.0096,
+      "llm_learning_rate": 1.9394233500927406e-06,
+      "loss": 0.5781,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00019394233500927405,
+      "step": 5708
+    },
+    {
+      "epoch": 0.58,
+      "gumbel_tau": 0.2601,
+      "image_masks_loss": 0.0237,
+      "llm_learning_rate": 1.9386148895153527e-06,
+      "loss": 0.7364,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00019386148895153527,
+      "step": 5709
+    },
+    {
+      "epoch": 0.59,
+      "gumbel_tau": 0.2601,
+      "image_masks_loss": 0.0406,
+      "llm_learning_rate": 1.9378064907589436e-06,
+      "loss": 0.6749,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 0.00019378064907589436,
+      "step": 5710
+    },
+    {
+      "epoch": 0.59,
+      "gumbel_tau": 0.26,
+      "image_masks_loss": 0.0162,
+      "llm_learning_rate": 1.9369981539125364e-06,
+      "loss": 0.7732,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.00019369981539125365,
+      "step": 5711
+    },
+    {
+      "epoch": 0.59,
+      "gumbel_tau": 0.2599,
+      "image_masks_loss": 0.034,
+      "llm_learning_rate": 1.936189879065147e-06,
+      "loss": 0.7819,
+      "output_text_mask_loss": 0.0427,
+      "predictor_learning_rate": 0.0001936189879065147,
+      "step": 5712
+    },
+    {
+      "epoch": 0.59,
+      "gumbel_tau": 0.2599,
+      "image_masks_loss": 0.0079,
+      "llm_learning_rate": 1.9353816663057835e-06,
+      "loss": 0.7109,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00019353816663057833,
+      "step": 5713
+    },
+    {
+      "epoch": 0.59,
+      "gumbel_tau": 0.2598,
+      "image_masks_loss": 0.0118,
+      "llm_learning_rate": 1.934573515723448e-06,
+      "loss": 0.7786,
+      "output_text_mask_loss": 0.0174,
+      "predictor_learning_rate": 0.00019345735157234482,
+      "step": 5714
+    },
+    {
+      "epoch": 0.59,
+      "gumbel_tau": 0.2597,
+      "image_masks_loss": 0.0303,
+      "llm_learning_rate": 1.933765427407137e-06,
+      "loss": 0.8105,
+      "output_text_mask_loss": 0.017,
+      "predictor_learning_rate": 0.0001933765427407137,
+      "step": 5715
+    },
+    {
+      "epoch": 0.59,
+      "gumbel_tau": 0.2597,
+      "image_masks_loss": 0.0133,
+      "llm_learning_rate": 1.932957401445838e-06,
+      "loss": 0.7993,
+      "output_text_mask_loss": 0.0353,
+      "predictor_learning_rate": 0.0001932957401445838,
+      "step": 5716
+    },
+    {
+      "epoch": 0.59,
+      "gumbel_tau": 0.2596,
+      "image_masks_loss": 0.0254,
+      "llm_learning_rate": 1.932149437928533e-06,
+      "loss": 0.8023,
+      "output_text_mask_loss": 0.0372,
+      "predictor_learning_rate": 0.0001932149437928533,
+      "step": 5717
+    },
+    {
+      "epoch": 0.59,
+      "gumbel_tau": 0.2596,
+      "image_masks_loss": 0.0151,
+      "llm_learning_rate": 1.931341536944197e-06,
+      "loss": 0.7198,
+      "output_text_mask_loss": 0.0257,
+      "predictor_learning_rate": 0.00019313415369441966,
+      "step": 5718
+    },
+    {
+      "epoch": 0.59,
+      "gumbel_tau": 0.2595,
+      "image_masks_loss": 0.0188,
+      "llm_learning_rate": 1.930533698581797e-06,
+      "loss": 0.6785,
+      "output_text_mask_loss": 0.0115,
+      "predictor_learning_rate": 0.00019305336985817967,
+      "step": 5719
+    },
+    {
+      "epoch": 0.59,
+      "gumbel_tau": 0.2594,
+      "image_masks_loss": 0.0156,
+      "llm_learning_rate": 1.929725922930295e-06,
+      "loss": 0.7729,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0001929725922930295,
+      "step": 5720
+    },
+    {
+      "epoch": 0.59,
+      "gumbel_tau": 0.2594,
+      "image_masks_loss": 0.0164,
+      "llm_learning_rate": 1.9289182100786445e-06,
+      "loss": 0.6356,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00019289182100786444,
+      "step": 5721
+    },
+    {
+      "epoch": 0.59,
+      "gumbel_tau": 0.2593,
+      "image_masks_loss": 0.0044,
+      "llm_learning_rate": 1.928110560115793e-06,
+      "loss": 0.7413,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0001928110560115793,
+      "step": 5722
+    },
+    {
+      "epoch": 0.59,
+      "gumbel_tau": 0.2593,
+      "image_masks_loss": 0.0117,
+      "llm_learning_rate": 1.92730297313068e-06,
+      "loss": 0.7039,
+      "output_text_mask_loss": 0.0023,
+      "predictor_learning_rate": 0.000192730297313068,
+      "step": 5723
+    },
+    {
+      "epoch": 0.59,
+      "gumbel_tau": 0.2592,
+      "image_masks_loss": 0.0058,
+      "llm_learning_rate": 1.9264954492122394e-06,
+      "loss": 0.7632,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.00019264954492122395,
+      "step": 5724
+    },
+    {
+      "epoch": 0.59,
+      "gumbel_tau": 0.2591,
+      "image_masks_loss": 0.0104,
+      "llm_learning_rate": 1.9256879884493977e-06,
+      "loss": 0.7264,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00019256879884493975,
+      "step": 5725
+    },
+    {
+      "epoch": 0.59,
+      "gumbel_tau": 0.2591,
+      "image_masks_loss": 0.0147,
+      "llm_learning_rate": 1.924880590931074e-06,
+      "loss": 0.7778,
+      "output_text_mask_loss": 0.025,
+      "predictor_learning_rate": 0.00019248805909310738,
+      "step": 5726
+    },
+    {
+      "epoch": 0.59,
+      "gumbel_tau": 0.259,
+      "image_masks_loss": 0.0227,
+      "llm_learning_rate": 1.9240732567461796e-06,
+      "loss": 0.6905,
+      "output_text_mask_loss": 0.0035,
+      "predictor_learning_rate": 0.00019240732567461796,
+      "step": 5727
+    },
+    {
+      "epoch": 0.59,
+      "gumbel_tau": 0.259,
+      "image_masks_loss": 0.0078,
+      "llm_learning_rate": 1.923265985983622e-06,
+      "loss": 0.6314,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0001923265985983622,
+      "step": 5728
+    },
+    {
+      "epoch": 0.59,
+      "gumbel_tau": 0.2589,
+      "image_masks_loss": 0.0057,
+      "llm_learning_rate": 1.9224587787322985e-06,
+      "loss": 0.6954,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00019224587787322983,
+      "step": 5729
+    },
+    {
+      "epoch": 0.59,
+      "gumbel_tau": 0.2588,
+      "image_masks_loss": 0.0139,
+      "llm_learning_rate": 1.9216516350811e-06,
+      "loss": 0.7694,
+      "output_text_mask_loss": 0.0278,
+      "predictor_learning_rate": 0.00019216516350811,
+      "step": 5730
+    },
+    {
+      "epoch": 0.59,
+      "gumbel_tau": 0.2588,
+      "image_masks_loss": 0.0074,
+      "llm_learning_rate": 1.9208445551189135e-06,
+      "loss": 0.8044,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00019208445551189136,
+      "step": 5731
+    },
+    {
+      "epoch": 0.59,
+      "gumbel_tau": 0.2587,
+      "image_masks_loss": 0.0238,
+      "llm_learning_rate": 1.920037538934613e-06,
+      "loss": 0.7215,
+      "output_text_mask_loss": 0.0248,
+      "predictor_learning_rate": 0.00019200375389346127,
+      "step": 5732
+    },
+    {
+      "epoch": 0.59,
+      "gumbel_tau": 0.2586,
+      "image_masks_loss": 0.0106,
+      "llm_learning_rate": 1.919230586617071e-06,
+      "loss": 0.7144,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0001919230586617071,
+      "step": 5733
+    },
+    {
+      "epoch": 0.59,
+      "gumbel_tau": 0.2586,
+      "image_masks_loss": 0.0192,
+      "llm_learning_rate": 1.9184236982551504e-06,
+      "loss": 0.7305,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.00019184236982551503,
+      "step": 5734
+    },
+    {
+      "epoch": 0.59,
+      "gumbel_tau": 0.2585,
+      "image_masks_loss": 0.0222,
+      "llm_learning_rate": 1.9176168739377076e-06,
+      "loss": 0.7435,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00019176168739377077,
+      "step": 5735
+    },
+    {
+      "epoch": 0.59,
+      "gumbel_tau": 0.2585,
+      "image_masks_loss": 0.0216,
+      "llm_learning_rate": 1.9168101137535914e-06,
+      "loss": 0.8323,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 0.00019168101137535912,
+      "step": 5736
+    },
+    {
+      "epoch": 0.59,
+      "gumbel_tau": 0.2584,
+      "image_masks_loss": 0.0129,
+      "llm_learning_rate": 1.9160034177916446e-06,
+      "loss": 0.7754,
+      "output_text_mask_loss": 0.0244,
+      "predictor_learning_rate": 0.00019160034177916445,
+      "step": 5737
+    },
+    {
+      "epoch": 0.59,
+      "gumbel_tau": 0.2583,
+      "image_masks_loss": 0.0053,
+      "llm_learning_rate": 1.9151967861407026e-06,
+      "loss": 0.772,
+      "output_text_mask_loss": 0.0324,
+      "predictor_learning_rate": 0.00019151967861407026,
+      "step": 5738
+    },
+    {
+      "epoch": 0.59,
+      "gumbel_tau": 0.2583,
+      "image_masks_loss": 0.0041,
+      "llm_learning_rate": 1.9143902188895927e-06,
+      "loss": 0.744,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 0.00019143902188895927,
+      "step": 5739
+    },
+    {
+      "epoch": 0.59,
+      "gumbel_tau": 0.2582,
+      "image_masks_loss": 0.0194,
+      "llm_learning_rate": 1.9135837161271367e-06,
+      "loss": 0.5655,
+      "output_text_mask_loss": 0.0368,
+      "predictor_learning_rate": 0.00019135837161271367,
+      "step": 5740
+    },
+    {
+      "epoch": 0.59,
+      "gumbel_tau": 0.2582,
+      "image_masks_loss": 0.0155,
+      "llm_learning_rate": 1.912777277942148e-06,
+      "loss": 0.7279,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.0001912777277942148,
+      "step": 5741
+    },
+    {
+      "epoch": 0.59,
+      "gumbel_tau": 0.2581,
+      "image_masks_loss": 0.0374,
+      "llm_learning_rate": 1.9119709044234334e-06,
+      "loss": 0.6726,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00019119709044234336,
+      "step": 5742
+    },
+    {
+      "epoch": 0.59,
+      "gumbel_tau": 0.258,
+      "image_masks_loss": 0.0081,
+      "llm_learning_rate": 1.911164595659792e-06,
+      "loss": 0.766,
+      "output_text_mask_loss": 0.0248,
+      "predictor_learning_rate": 0.0001911164595659792,
+      "step": 5743
+    },
+    {
+      "epoch": 0.59,
+      "gumbel_tau": 0.258,
+      "image_masks_loss": 0.0097,
+      "llm_learning_rate": 1.9103583517400177e-06,
+      "loss": 0.7319,
+      "output_text_mask_loss": 0.0549,
+      "predictor_learning_rate": 0.00019103583517400175,
+      "step": 5744
+    },
+    {
+      "epoch": 0.59,
+      "gumbel_tau": 0.2579,
+      "image_masks_loss": 0.0101,
+      "llm_learning_rate": 1.909552172752895e-06,
+      "loss": 0.7382,
+      "output_text_mask_loss": 0.036,
+      "predictor_learning_rate": 0.00019095521727528948,
+      "step": 5745
+    },
+    {
+      "epoch": 0.59,
+      "gumbel_tau": 0.2579,
+      "image_masks_loss": 0.0182,
+      "llm_learning_rate": 1.9087460587872015e-06,
+      "loss": 0.7268,
+      "output_text_mask_loss": 0.0313,
+      "predictor_learning_rate": 0.00019087460587872014,
+      "step": 5746
+    },
+    {
+      "epoch": 0.59,
+      "gumbel_tau": 0.2578,
+      "image_masks_loss": 0.0106,
+      "llm_learning_rate": 1.9079400099317097e-06,
+      "loss": 0.7625,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.00019079400099317095,
+      "step": 5747
+    },
+    {
+      "epoch": 0.59,
+      "gumbel_tau": 0.2577,
+      "image_masks_loss": 0.0025,
+      "llm_learning_rate": 1.9071340262751825e-06,
+      "loss": 0.7693,
+      "output_text_mask_loss": 0.0217,
+      "predictor_learning_rate": 0.00019071340262751823,
+      "step": 5748
+    },
+    {
+      "epoch": 0.59,
+      "gumbel_tau": 0.2577,
+      "image_masks_loss": 0.0129,
+      "llm_learning_rate": 1.9063281079063772e-06,
+      "loss": 0.837,
+      "output_text_mask_loss": 0.0137,
+      "predictor_learning_rate": 0.0001906328107906377,
+      "step": 5749
+    },
+    {
+      "epoch": 0.59,
+      "gumbel_tau": 0.2576,
+      "image_masks_loss": 0.007,
+      "llm_learning_rate": 1.9055222549140435e-06,
+      "loss": 0.7538,
+      "output_text_mask_loss": 0.0004,
+      "predictor_learning_rate": 0.00019055222549140433,
+      "step": 5750
+    },
+    {
+      "epoch": 0.59,
+      "gumbel_tau": 0.2575,
+      "image_masks_loss": 0.0065,
+      "llm_learning_rate": 1.9047164673869226e-06,
+      "loss": 0.7065,
+      "output_text_mask_loss": 0.0282,
+      "predictor_learning_rate": 0.00019047164673869225,
+      "step": 5751
+    },
+    {
+      "epoch": 0.59,
+      "gumbel_tau": 0.2575,
+      "image_masks_loss": 0.0214,
+      "llm_learning_rate": 1.9039107454137496e-06,
+      "loss": 0.7589,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.00019039107454137493,
+      "step": 5752
+    },
+    {
+      "epoch": 0.59,
+      "gumbel_tau": 0.2574,
+      "image_masks_loss": 0.0132,
+      "llm_learning_rate": 1.9031050890832536e-06,
+      "loss": 0.7667,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.00019031050890832535,
+      "step": 5753
+    },
+    {
+      "epoch": 0.59,
+      "gumbel_tau": 0.2574,
+      "image_masks_loss": 0.0216,
+      "llm_learning_rate": 1.902299498484155e-06,
+      "loss": 0.7644,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00019022994984841548,
+      "step": 5754
+    },
+    {
+      "epoch": 0.59,
+      "gumbel_tau": 0.2573,
+      "image_masks_loss": 0.0121,
+      "llm_learning_rate": 1.9014939737051663e-06,
+      "loss": 0.7563,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00019014939737051662,
+      "step": 5755
+    },
+    {
+      "epoch": 0.59,
+      "gumbel_tau": 0.2572,
+      "image_masks_loss": 0.0129,
+      "llm_learning_rate": 1.9006885148349946e-06,
+      "loss": 0.7542,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 0.00019006885148349946,
+      "step": 5756
+    },
+    {
+      "epoch": 0.59,
+      "gumbel_tau": 0.2572,
+      "image_masks_loss": 0.0234,
+      "llm_learning_rate": 1.899883121962339e-06,
+      "loss": 0.7348,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00018998831219623388,
+      "step": 5757
+    },
+    {
+      "epoch": 0.59,
+      "gumbel_tau": 0.2571,
+      "image_masks_loss": 0.0309,
+      "llm_learning_rate": 1.8990777951758904e-06,
+      "loss": 0.7084,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.00018990777951758902,
+      "step": 5758
+    },
+    {
+      "epoch": 0.59,
+      "gumbel_tau": 0.2571,
+      "image_masks_loss": 0.0148,
+      "llm_learning_rate": 1.8982725345643332e-06,
+      "loss": 0.6935,
+      "output_text_mask_loss": 0.0099,
+      "predictor_learning_rate": 0.0001898272534564333,
+      "step": 5759
+    },
+    {
+      "epoch": 0.59,
+      "gumbel_tau": 0.257,
+      "image_masks_loss": 0.0173,
+      "llm_learning_rate": 1.897467340216345e-06,
+      "loss": 0.7397,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.0001897467340216345,
+      "step": 5760
+    },
+    {
+      "epoch": 0.59,
+      "gumbel_tau": 0.2569,
+      "image_masks_loss": 0.0112,
+      "llm_learning_rate": 1.8966622122205957e-06,
+      "loss": 0.6817,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00018966622122205956,
+      "step": 5761
+    },
+    {
+      "epoch": 0.59,
+      "gumbel_tau": 0.2569,
+      "image_masks_loss": 0.0102,
+      "llm_learning_rate": 1.8958571506657468e-06,
+      "loss": 0.7614,
+      "output_text_mask_loss": 0.119,
+      "predictor_learning_rate": 0.00018958571506657467,
+      "step": 5762
+    },
+    {
+      "epoch": 0.59,
+      "gumbel_tau": 0.2568,
+      "image_masks_loss": 0.0178,
+      "llm_learning_rate": 1.8950521556404548e-06,
+      "loss": 0.7792,
+      "output_text_mask_loss": 0.0263,
+      "predictor_learning_rate": 0.00018950521556404545,
+      "step": 5763
+    },
+    {
+      "epoch": 0.59,
+      "gumbel_tau": 0.2568,
+      "image_masks_loss": 0.014,
+      "llm_learning_rate": 1.8942472272333668e-06,
+      "loss": 0.7835,
+      "output_text_mask_loss": 0.0099,
+      "predictor_learning_rate": 0.00018942472272333666,
+      "step": 5764
+    },
+    {
+      "epoch": 0.59,
+      "gumbel_tau": 0.2567,
+      "image_masks_loss": 0.0101,
+      "llm_learning_rate": 1.8934423655331232e-06,
+      "loss": 0.704,
+      "output_text_mask_loss": 0.0012,
+      "predictor_learning_rate": 0.00018934423655331232,
+      "step": 5765
+    },
+    {
+      "epoch": 0.59,
+      "gumbel_tau": 0.2566,
+      "image_masks_loss": 0.0257,
+      "llm_learning_rate": 1.892637570628357e-06,
+      "loss": 0.7987,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.0001892637570628357,
+      "step": 5766
+    },
+    {
+      "epoch": 0.59,
+      "gumbel_tau": 0.2566,
+      "image_masks_loss": 0.0368,
+      "llm_learning_rate": 1.8918328426076948e-06,
+      "loss": 0.7764,
+      "output_text_mask_loss": 0.0116,
+      "predictor_learning_rate": 0.00018918328426076949,
+      "step": 5767
+    },
+    {
+      "epoch": 0.59,
+      "gumbel_tau": 0.2565,
+      "image_masks_loss": 0.0122,
+      "llm_learning_rate": 1.8910281815597554e-06,
+      "loss": 0.7439,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00018910281815597552,
+      "step": 5768
+    },
+    {
+      "epoch": 0.59,
+      "gumbel_tau": 0.2565,
+      "image_masks_loss": 0.0075,
+      "llm_learning_rate": 1.8902235875731485e-06,
+      "loss": 0.7368,
+      "output_text_mask_loss": 0.0542,
+      "predictor_learning_rate": 0.00018902235875731482,
+      "step": 5769
+    },
+    {
+      "epoch": 0.59,
+      "gumbel_tau": 0.2564,
+      "image_masks_loss": 0.0252,
+      "llm_learning_rate": 1.8894190607364782e-06,
+      "loss": 0.7654,
+      "output_text_mask_loss": 0.0383,
+      "predictor_learning_rate": 0.0001889419060736478,
+      "step": 5770
+    },
+    {
+      "epoch": 0.59,
+      "gumbel_tau": 0.2563,
+      "image_masks_loss": 0.0124,
+      "llm_learning_rate": 1.8886146011383405e-06,
+      "loss": 0.7775,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 0.00018886146011383404,
+      "step": 5771
+    },
+    {
+      "epoch": 0.59,
+      "gumbel_tau": 0.2563,
+      "image_masks_loss": 0.028,
+      "llm_learning_rate": 1.887810208867325e-06,
+      "loss": 0.7582,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00018878102088673248,
+      "step": 5772
+    },
+    {
+      "epoch": 0.59,
+      "gumbel_tau": 0.2562,
+      "image_masks_loss": 0.0219,
+      "llm_learning_rate": 1.8870058840120133e-06,
+      "loss": 0.73,
+      "output_text_mask_loss": 0.0286,
+      "predictor_learning_rate": 0.00018870058840120133,
+      "step": 5773
+    },
+    {
+      "epoch": 0.59,
+      "gumbel_tau": 0.2562,
+      "image_masks_loss": 0.007,
+      "llm_learning_rate": 1.886201626660979e-06,
+      "loss": 0.6256,
+      "output_text_mask_loss": NaN,
+      "predictor_learning_rate": 0.00018862016266609787,
+      "step": 5774
+    },
+    {
+      "epoch": 0.59,
+      "gumbel_tau": 0.2561,
+      "image_masks_loss": 0.0171,
+      "llm_learning_rate": 1.885397436902788e-06,
+      "loss": 0.7181,
+      "output_text_mask_loss": 0.0019,
+      "predictor_learning_rate": 0.00018853974369027878,
+      "step": 5775
+    },
+    {
+      "epoch": 0.59,
+      "gumbel_tau": 0.256,
+      "image_masks_loss": 0.0156,
+      "llm_learning_rate": 1.884593314826001e-06,
+      "loss": 0.7548,
+      "output_text_mask_loss": 0.0248,
+      "predictor_learning_rate": 0.00018845933148260008,
+      "step": 5776
+    },
+    {
+      "epoch": 0.59,
+      "gumbel_tau": 0.256,
+      "image_masks_loss": 0.0118,
+      "llm_learning_rate": 1.8837892605191688e-06,
+      "loss": 0.6882,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.00018837892605191688,
+      "step": 5777
+    },
+    {
+      "epoch": 0.59,
+      "gumbel_tau": 0.2559,
+      "image_masks_loss": 0.016,
+      "llm_learning_rate": 1.8829852740708352e-06,
+      "loss": 0.7196,
+      "output_text_mask_loss": 0.0039,
+      "predictor_learning_rate": 0.0001882985274070835,
+      "step": 5778
+    },
+    {
+      "epoch": 0.59,
+      "gumbel_tau": 0.2559,
+      "image_masks_loss": 0.0116,
+      "llm_learning_rate": 1.882181355569538e-06,
+      "loss": 0.6485,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00018821813555695377,
+      "step": 5779
+    },
+    {
+      "epoch": 0.59,
+      "gumbel_tau": 0.2558,
+      "image_masks_loss": 0.0317,
+      "llm_learning_rate": 1.8813775051038064e-06,
+      "loss": 0.6298,
+      "output_text_mask_loss": 0.0374,
+      "predictor_learning_rate": 0.00018813775051038063,
+      "step": 5780
+    },
+    {
+      "epoch": 0.59,
+      "gumbel_tau": 0.2557,
+      "image_masks_loss": 0.0112,
+      "llm_learning_rate": 1.8805737227621615e-06,
+      "loss": 0.7045,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00018805737227621614,
+      "step": 5781
+    },
+    {
+      "epoch": 0.59,
+      "gumbel_tau": 0.2557,
+      "image_masks_loss": 0.0185,
+      "llm_learning_rate": 1.8797700086331178e-06,
+      "loss": 0.7914,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.00018797700086331176,
+      "step": 5782
+    },
+    {
+      "epoch": 0.59,
+      "gumbel_tau": 0.2556,
+      "image_masks_loss": 0.0149,
+      "llm_learning_rate": 1.8789663628051824e-06,
+      "loss": 0.6189,
+      "output_text_mask_loss": 0.0054,
+      "predictor_learning_rate": 0.00018789663628051824,
+      "step": 5783
+    },
+    {
+      "epoch": 0.59,
+      "gumbel_tau": 0.2556,
+      "image_masks_loss": 0.0336,
+      "llm_learning_rate": 1.8781627853668547e-06,
+      "loss": 0.702,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00018781627853668546,
+      "step": 5784
+    },
+    {
+      "epoch": 0.59,
+      "gumbel_tau": 0.2555,
+      "image_masks_loss": 0.0082,
+      "llm_learning_rate": 1.8773592764066253e-06,
+      "loss": 0.7622,
+      "output_text_mask_loss": 0.0107,
+      "predictor_learning_rate": 0.00018773592764066252,
+      "step": 5785
+    },
+    {
+      "epoch": 0.59,
+      "gumbel_tau": 0.2554,
+      "image_masks_loss": 0.007,
+      "llm_learning_rate": 1.8765558360129802e-06,
+      "loss": 0.7877,
+      "output_text_mask_loss": 0.0014,
+      "predictor_learning_rate": 0.000187655583601298,
+      "step": 5786
+    },
+    {
+      "epoch": 0.59,
+      "gumbel_tau": 0.2554,
+      "image_masks_loss": 0.0181,
+      "llm_learning_rate": 1.8757524642743955e-06,
+      "loss": 0.7777,
+      "output_text_mask_loss": 0.046,
+      "predictor_learning_rate": 0.00018757524642743954,
+      "step": 5787
+    },
+    {
+      "epoch": 0.59,
+      "gumbel_tau": 0.2553,
+      "image_masks_loss": 0.038,
+      "llm_learning_rate": 1.8749491612793394e-06,
+      "loss": 0.8328,
+      "output_text_mask_loss": 0.0126,
+      "predictor_learning_rate": 0.00018749491612793392,
+      "step": 5788
+    },
+    {
+      "epoch": 0.59,
+      "gumbel_tau": 0.2553,
+      "image_masks_loss": 0.0168,
+      "llm_learning_rate": 1.8741459271162737e-06,
+      "loss": 0.7487,
+      "output_text_mask_loss": 0.0082,
+      "predictor_learning_rate": 0.00018741459271162737,
+      "step": 5789
+    },
+    {
+      "epoch": 0.59,
+      "gumbel_tau": 0.2552,
+      "image_masks_loss": 0.0299,
+      "llm_learning_rate": 1.8733427618736528e-06,
+      "loss": 0.7698,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00018733427618736526,
+      "step": 5790
+    },
+    {
+      "epoch": 0.59,
+      "gumbel_tau": 0.2551,
+      "image_masks_loss": 0.0174,
+      "llm_learning_rate": 1.872539665639922e-06,
+      "loss": 0.7472,
+      "output_text_mask_loss": 0.0431,
+      "predictor_learning_rate": 0.0001872539665639922,
+      "step": 5791
+    },
+    {
+      "epoch": 0.59,
+      "gumbel_tau": 0.2551,
+      "image_masks_loss": 0.0222,
+      "llm_learning_rate": 1.8717366385035218e-06,
+      "loss": 0.7121,
+      "output_text_mask_loss": 0.017,
+      "predictor_learning_rate": 0.00018717366385035216,
+      "step": 5792
+    },
+    {
+      "epoch": 0.59,
+      "gumbel_tau": 0.255,
+      "image_masks_loss": 0.0114,
+      "llm_learning_rate": 1.8709336805528819e-06,
+      "loss": 0.7357,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00018709336805528817,
+      "step": 5793
+    },
+    {
+      "epoch": 0.59,
+      "gumbel_tau": 0.2549,
+      "image_masks_loss": 0.0111,
+      "llm_learning_rate": 1.8701307918764261e-06,
+      "loss": 0.806,
+      "output_text_mask_loss": 0.0322,
+      "predictor_learning_rate": 0.0001870130791876426,
+      "step": 5794
+    },
+    {
+      "epoch": 0.59,
+      "gumbel_tau": 0.2549,
+      "image_masks_loss": 0.0138,
+      "llm_learning_rate": 1.8693279725625708e-06,
+      "loss": 0.7659,
+      "output_text_mask_loss": 0.0259,
+      "predictor_learning_rate": 0.00018693279725625706,
+      "step": 5795
+    },
+    {
+      "epoch": 0.59,
+      "gumbel_tau": 0.2548,
+      "image_masks_loss": 0.0278,
+      "llm_learning_rate": 1.868525222699724e-06,
+      "loss": 0.8339,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 0.00018685252226997238,
+      "step": 5796
+    },
+    {
+      "epoch": 0.59,
+      "gumbel_tau": 0.2548,
+      "image_masks_loss": 0.0311,
+      "llm_learning_rate": 1.8677225423762863e-06,
+      "loss": 0.7122,
+      "output_text_mask_loss": 0.001,
+      "predictor_learning_rate": 0.00018677225423762862,
+      "step": 5797
+    },
+    {
+      "epoch": 0.59,
+      "gumbel_tau": 0.2547,
+      "image_masks_loss": 0.0231,
+      "llm_learning_rate": 1.8669199316806497e-06,
+      "loss": 0.7809,
+      "output_text_mask_loss": 0.0452,
+      "predictor_learning_rate": 0.00018669199316806497,
+      "step": 5798
+    },
+    {
+      "epoch": 0.59,
+      "gumbel_tau": 0.2546,
+      "image_masks_loss": 0.0148,
+      "llm_learning_rate": 1.8661173907012012e-06,
+      "loss": 0.8006,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 0.00018661173907012012,
+      "step": 5799
+    },
+    {
+      "epoch": 0.59,
+      "gumbel_tau": 0.2546,
+      "image_masks_loss": 0.0091,
+      "llm_learning_rate": 1.8653149195263177e-06,
+      "loss": 0.759,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.00018653149195263174,
+      "step": 5800
+    },
+    {
+      "epoch": 0.59,
+      "gumbel_tau": 0.2545,
+      "image_masks_loss": 0.0257,
+      "llm_learning_rate": 1.864512518244368e-06,
+      "loss": 0.757,
+      "output_text_mask_loss": 0.0292,
+      "predictor_learning_rate": 0.00018645125182443679,
+      "step": 5801
+    },
+    {
+      "epoch": 0.59,
+      "gumbel_tau": 0.2545,
+      "image_masks_loss": 0.0261,
+      "llm_learning_rate": 1.863710186943716e-06,
+      "loss": 0.7035,
+      "output_text_mask_loss": 0.0208,
+      "predictor_learning_rate": 0.00018637101869437158,
+      "step": 5802
+    },
+    {
+      "epoch": 0.59,
+      "gumbel_tau": 0.2544,
+      "image_masks_loss": 0.0113,
+      "llm_learning_rate": 1.8629079257127156e-06,
+      "loss": 0.6981,
+      "output_text_mask_loss": 0.0147,
+      "predictor_learning_rate": 0.00018629079257127153,
+      "step": 5803
+    },
+    {
+      "epoch": 0.59,
+      "gumbel_tau": 0.2543,
+      "image_masks_loss": 0.0109,
+      "llm_learning_rate": 1.8621057346397133e-06,
+      "loss": 0.6825,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.00018621057346397132,
+      "step": 5804
+    },
+    {
+      "epoch": 0.59,
+      "gumbel_tau": 0.2543,
+      "image_masks_loss": 0.0034,
+      "llm_learning_rate": 1.861303613813049e-06,
+      "loss": 0.802,
+      "output_text_mask_loss": 0.0177,
+      "predictor_learning_rate": 0.0001861303613813049,
+      "step": 5805
+    },
+    {
+      "epoch": 0.59,
+      "gumbel_tau": 0.2542,
+      "image_masks_loss": 0.0156,
+      "llm_learning_rate": 1.8605015633210529e-06,
+      "loss": 0.7068,
+      "output_text_mask_loss": 0.0039,
+      "predictor_learning_rate": 0.0001860501563321053,
+      "step": 5806
+    },
+    {
+      "epoch": 0.59,
+      "gumbel_tau": 0.2542,
+      "image_masks_loss": 0.0098,
+      "llm_learning_rate": 1.8596995832520487e-06,
+      "loss": 0.7854,
+      "output_text_mask_loss": 0.0004,
+      "predictor_learning_rate": 0.00018596995832520486,
+      "step": 5807
+    },
+    {
+      "epoch": 0.6,
+      "gumbel_tau": 0.2541,
+      "image_masks_loss": 0.0055,
+      "llm_learning_rate": 1.8588976736943525e-06,
+      "loss": 0.7695,
+      "output_text_mask_loss": 0.0032,
+      "predictor_learning_rate": 0.00018588976736943526,
+      "step": 5808
+    },
+    {
+      "epoch": 0.6,
+      "gumbel_tau": 0.254,
+      "image_masks_loss": 0.0292,
+      "llm_learning_rate": 1.8580958347362732e-06,
+      "loss": 0.8023,
+      "output_text_mask_loss": 0.005,
+      "predictor_learning_rate": 0.00018580958347362732,
+      "step": 5809
+    },
+    {
+      "epoch": 0.6,
+      "gumbel_tau": 0.254,
+      "image_masks_loss": 0.0079,
+      "llm_learning_rate": 1.8572940664661094e-06,
+      "loss": 0.7782,
+      "output_text_mask_loss": 0.0324,
+      "predictor_learning_rate": 0.00018572940664661093,
+      "step": 5810
+    },
+    {
+      "epoch": 0.6,
+      "gumbel_tau": 0.2539,
+      "image_masks_loss": 0.0239,
+      "llm_learning_rate": 1.8564923689721554e-06,
+      "loss": 0.7502,
+      "output_text_mask_loss": 0.0072,
+      "predictor_learning_rate": 0.00018564923689721552,
+      "step": 5811
+    },
+    {
+      "epoch": 0.6,
+      "gumbel_tau": 0.2539,
+      "image_masks_loss": 0.0179,
+      "llm_learning_rate": 1.855690742342695e-06,
+      "loss": 0.6643,
+      "output_text_mask_loss": 0.0109,
+      "predictor_learning_rate": 0.0001855690742342695,
+      "step": 5812
+    },
+    {
+      "epoch": 0.6,
+      "gumbel_tau": 0.2538,
+      "image_masks_loss": 0.0092,
+      "llm_learning_rate": 1.8548891866660052e-06,
+      "loss": 0.7362,
+      "output_text_mask_loss": 0.0088,
+      "predictor_learning_rate": 0.0001854889186666005,
+      "step": 5813
+    },
+    {
+      "epoch": 0.6,
+      "gumbel_tau": 0.2537,
+      "image_masks_loss": 0.0153,
+      "llm_learning_rate": 1.8540877020303543e-06,
+      "loss": 0.7279,
+      "output_text_mask_loss": 0.02,
+      "predictor_learning_rate": 0.00018540877020303543,
+      "step": 5814
+    },
+    {
+      "epoch": 0.6,
+      "gumbel_tau": 0.2537,
+      "image_masks_loss": 0.0215,
+      "llm_learning_rate": 1.853286288524005e-06,
+      "loss": 0.7128,
+      "output_text_mask_loss": 0.0322,
+      "predictor_learning_rate": 0.00018532862885240048,
+      "step": 5815
+    },
+    {
+      "epoch": 0.6,
+      "gumbel_tau": 0.2536,
+      "image_masks_loss": 0.0257,
+      "llm_learning_rate": 1.8524849462352101e-06,
+      "loss": 0.7668,
+      "output_text_mask_loss": 0.0824,
+      "predictor_learning_rate": 0.00018524849462352101,
+      "step": 5816
+    },
+    {
+      "epoch": 0.6,
+      "gumbel_tau": 0.2536,
+      "image_masks_loss": 0.0231,
+      "llm_learning_rate": 1.851683675252215e-06,
+      "loss": 0.7863,
+      "output_text_mask_loss": 0.0137,
+      "predictor_learning_rate": 0.00018516836752522148,
+      "step": 5817
+    },
+    {
+      "epoch": 0.6,
+      "gumbel_tau": 0.2535,
+      "image_masks_loss": 0.0045,
+      "llm_learning_rate": 1.850882475663258e-06,
+      "loss": 0.8134,
+      "output_text_mask_loss": 0.0077,
+      "predictor_learning_rate": 0.00018508824756632578,
+      "step": 5818
+    },
+    {
+      "epoch": 0.6,
+      "gumbel_tau": 0.2535,
+      "image_masks_loss": 0.0081,
+      "llm_learning_rate": 1.8500813475565685e-06,
+      "loss": 0.7377,
+      "output_text_mask_loss": 0.0189,
+      "predictor_learning_rate": 0.00018500813475565684,
+      "step": 5819
+    },
+    {
+      "epoch": 0.6,
+      "gumbel_tau": 0.2534,
+      "image_masks_loss": 0.0054,
+      "llm_learning_rate": 1.8492802910203688e-06,
+      "loss": 0.6844,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00018492802910203686,
+      "step": 5820
+    },
+    {
+      "epoch": 0.6,
+      "gumbel_tau": 0.2533,
+      "image_masks_loss": 0.0103,
+      "llm_learning_rate": 1.8484793061428725e-06,
+      "loss": 0.6363,
+      "output_text_mask_loss": 0.0326,
+      "predictor_learning_rate": 0.00018484793061428724,
+      "step": 5821
+    },
+    {
+      "epoch": 0.6,
+      "gumbel_tau": 0.2533,
+      "image_masks_loss": 0.011,
+      "llm_learning_rate": 1.8476783930122862e-06,
+      "loss": 0.6788,
+      "output_text_mask_loss": 0.0219,
+      "predictor_learning_rate": 0.00018476783930122863,
+      "step": 5822
+    },
+    {
+      "epoch": 0.6,
+      "gumbel_tau": 0.2532,
+      "image_masks_loss": 0.0129,
+      "llm_learning_rate": 1.8468775517168087e-06,
+      "loss": 0.7245,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.00018468775517168086,
+      "step": 5823
+    },
+    {
+      "epoch": 0.6,
+      "gumbel_tau": 0.2532,
+      "image_masks_loss": 0.0108,
+      "llm_learning_rate": 1.8460767823446302e-06,
+      "loss": 0.7832,
+      "output_text_mask_loss": 0.0172,
+      "predictor_learning_rate": 0.000184607678234463,
+      "step": 5824
+    },
+    {
+      "epoch": 0.6,
+      "gumbel_tau": 0.2531,
+      "image_masks_loss": 0.015,
+      "llm_learning_rate": 1.8452760849839327e-06,
+      "loss": 0.7416,
+      "output_text_mask_loss": 0.0225,
+      "predictor_learning_rate": 0.00018452760849839325,
+      "step": 5825
+    },
+    {
+      "epoch": 0.6,
+      "gumbel_tau": 0.253,
+      "image_masks_loss": 0.0121,
+      "llm_learning_rate": 1.8444754597228904e-06,
+      "loss": 0.7704,
+      "output_text_mask_loss": 0.036,
+      "predictor_learning_rate": 0.00018444754597228902,
+      "step": 5826
+    },
+    {
+      "epoch": 0.6,
+      "gumbel_tau": 0.253,
+      "image_masks_loss": 0.0084,
+      "llm_learning_rate": 1.8436749066496711e-06,
+      "loss": 0.6456,
+      "output_text_mask_loss": 0.0086,
+      "predictor_learning_rate": 0.0001843674906649671,
+      "step": 5827
+    },
+    {
+      "epoch": 0.6,
+      "gumbel_tau": 0.2529,
+      "image_masks_loss": 0.0109,
+      "llm_learning_rate": 1.8428744258524335e-06,
+      "loss": 0.6865,
+      "output_text_mask_loss": 0.0023,
+      "predictor_learning_rate": 0.00018428744258524332,
+      "step": 5828
+    },
+    {
+      "epoch": 0.6,
+      "gumbel_tau": 0.2529,
+      "image_masks_loss": 0.0081,
+      "llm_learning_rate": 1.8420740174193273e-06,
+      "loss": 0.589,
+      "output_text_mask_loss": 0.0534,
+      "predictor_learning_rate": 0.0001842074017419327,
+      "step": 5829
+    },
+    {
+      "epoch": 0.6,
+      "gumbel_tau": 0.2528,
+      "image_masks_loss": 0.0081,
+      "llm_learning_rate": 1.8412736814384952e-06,
+      "loss": 0.7919,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.0001841273681438495,
+      "step": 5830
+    },
+    {
+      "epoch": 0.6,
+      "gumbel_tau": 0.2527,
+      "image_masks_loss": 0.0191,
+      "llm_learning_rate": 1.8404734179980731e-06,
+      "loss": 0.7935,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.00018404734179980732,
+      "step": 5831
+    },
+    {
+      "epoch": 0.6,
+      "gumbel_tau": 0.2527,
+      "image_masks_loss": 0.0138,
+      "llm_learning_rate": 1.8396732271861875e-06,
+      "loss": 0.7071,
+      "output_text_mask_loss": 0.0286,
+      "predictor_learning_rate": 0.00018396732271861875,
+      "step": 5832
+    },
+    {
+      "epoch": 0.6,
+      "gumbel_tau": 0.2526,
+      "image_masks_loss": 0.0048,
+      "llm_learning_rate": 1.8388731090909561e-06,
+      "loss": 0.6721,
+      "output_text_mask_loss": 0.0275,
+      "predictor_learning_rate": 0.0001838873109090956,
+      "step": 5833
+    },
+    {
+      "epoch": 0.6,
+      "gumbel_tau": 0.2526,
+      "image_masks_loss": 0.0153,
+      "llm_learning_rate": 1.8380730638004913e-06,
+      "loss": 0.7717,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00018380730638004912,
+      "step": 5834
+    },
+    {
+      "epoch": 0.6,
+      "gumbel_tau": 0.2525,
+      "image_masks_loss": 0.0041,
+      "llm_learning_rate": 1.8372730914028953e-06,
+      "loss": 0.8206,
+      "output_text_mask_loss": 0.0099,
+      "predictor_learning_rate": 0.00018372730914028952,
+      "step": 5835
+    },
+    {
+      "epoch": 0.6,
+      "gumbel_tau": 0.2524,
+      "image_masks_loss": 0.0065,
+      "llm_learning_rate": 1.8364731919862623e-06,
+      "loss": 0.8513,
+      "output_text_mask_loss": 0.0116,
+      "predictor_learning_rate": 0.00018364731919862622,
+      "step": 5836
+    },
+    {
+      "epoch": 0.6,
+      "gumbel_tau": 0.2524,
+      "image_masks_loss": 0.0224,
+      "llm_learning_rate": 1.8356733656386788e-06,
+      "loss": 0.737,
+      "output_text_mask_loss": 0.0237,
+      "predictor_learning_rate": 0.00018356733656386788,
+      "step": 5837
+    },
+    {
+      "epoch": 0.6,
+      "gumbel_tau": 0.2523,
+      "image_masks_loss": 0.0044,
+      "llm_learning_rate": 1.8348736124482246e-06,
+      "loss": 0.7284,
+      "output_text_mask_loss": 0.0012,
+      "predictor_learning_rate": 0.00018348736124482247,
+      "step": 5838
+    },
+    {
+      "epoch": 0.6,
+      "gumbel_tau": 0.2523,
+      "image_masks_loss": 0.0142,
+      "llm_learning_rate": 1.8340739325029699e-06,
+      "loss": 0.7292,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00018340739325029698,
+      "step": 5839
+    },
+    {
+      "epoch": 0.6,
+      "gumbel_tau": 0.2522,
+      "image_masks_loss": 0.0115,
+      "llm_learning_rate": 1.8332743258909764e-06,
+      "loss": 0.7353,
+      "output_text_mask_loss": 0.0058,
+      "predictor_learning_rate": 0.00018332743258909763,
+      "step": 5840
+    },
+    {
+      "epoch": 0.6,
+      "gumbel_tau": 0.2521,
+      "image_masks_loss": 0.0047,
+      "llm_learning_rate": 1.8324747927002996e-06,
+      "loss": 0.7923,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.00018324747927002993,
+      "step": 5841
+    },
+    {
+      "epoch": 0.6,
+      "gumbel_tau": 0.2521,
+      "image_masks_loss": 0.0028,
+      "llm_learning_rate": 1.831675333018986e-06,
+      "loss": 0.762,
+      "output_text_mask_loss": 0.0185,
+      "predictor_learning_rate": 0.00018316753330189858,
+      "step": 5842
+    },
+    {
+      "epoch": 0.6,
+      "gumbel_tau": 0.252,
+      "image_masks_loss": 0.0122,
+      "llm_learning_rate": 1.8308759469350728e-06,
+      "loss": 0.6914,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 0.00018308759469350727,
+      "step": 5843
+    },
+    {
+      "epoch": 0.6,
+      "gumbel_tau": 0.252,
+      "image_masks_loss": 0.0102,
+      "llm_learning_rate": 1.8300766345365911e-06,
+      "loss": 0.6929,
+      "output_text_mask_loss": 0.0177,
+      "predictor_learning_rate": 0.0001830076634536591,
+      "step": 5844
+    },
+    {
+      "epoch": 0.6,
+      "gumbel_tau": 0.2519,
+      "image_masks_loss": 0.0063,
+      "llm_learning_rate": 1.8292773959115623e-06,
+      "loss": 0.6645,
+      "output_text_mask_loss": 0.0058,
+      "predictor_learning_rate": 0.00018292773959115621,
+      "step": 5845
+    },
+    {
+      "epoch": 0.6,
+      "gumbel_tau": 0.2518,
+      "image_masks_loss": 0.0098,
+      "llm_learning_rate": 1.8284782311480004e-06,
+      "loss": 0.7301,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.00018284782311480004,
+      "step": 5846
+    },
+    {
+      "epoch": 0.6,
+      "gumbel_tau": 0.2518,
+      "image_masks_loss": 0.008,
+      "llm_learning_rate": 1.8276791403339116e-06,
+      "loss": 0.7238,
+      "output_text_mask_loss": 0.0263,
+      "predictor_learning_rate": 0.00018276791403339115,
+      "step": 5847
+    },
+    {
+      "epoch": 0.6,
+      "gumbel_tau": 0.2517,
+      "image_masks_loss": 0.0209,
+      "llm_learning_rate": 1.8268801235572935e-06,
+      "loss": 0.6474,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.00018268801235572934,
+      "step": 5848
+    },
+    {
+      "epoch": 0.6,
+      "gumbel_tau": 0.2517,
+      "image_masks_loss": 0.0174,
+      "llm_learning_rate": 1.8260811809061352e-06,
+      "loss": 0.7075,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.0001826081180906135,
+      "step": 5849
+    },
+    {
+      "epoch": 0.6,
+      "gumbel_tau": 0.2516,
+      "image_masks_loss": 0.008,
+      "llm_learning_rate": 1.8252823124684185e-06,
+      "loss": 0.7752,
+      "output_text_mask_loss": 0.0368,
+      "predictor_learning_rate": 0.00018252823124684185,
+      "step": 5850
+    },
+    {
+      "epoch": 0.6,
+      "gumbel_tau": 0.2515,
+      "image_masks_loss": 0.009,
+      "llm_learning_rate": 1.824483518332117e-06,
+      "loss": 0.6436,
+      "output_text_mask_loss": 0.0137,
+      "predictor_learning_rate": 0.00018244835183321167,
+      "step": 5851
+    },
+    {
+      "epoch": 0.6,
+      "gumbel_tau": 0.2515,
+      "image_masks_loss": 0.0142,
+      "llm_learning_rate": 1.8236847985851947e-06,
+      "loss": 0.7148,
+      "output_text_mask_loss": 0.0039,
+      "predictor_learning_rate": 0.00018236847985851946,
+      "step": 5852
+    },
+    {
+      "epoch": 0.6,
+      "gumbel_tau": 0.2514,
+      "image_masks_loss": 0.0173,
+      "llm_learning_rate": 1.822886153315608e-06,
+      "loss": 0.7061,
+      "output_text_mask_loss": 0.0081,
+      "predictor_learning_rate": 0.0001822886153315608,
+      "step": 5853
+    },
+    {
+      "epoch": 0.6,
+      "gumbel_tau": 0.2514,
+      "image_masks_loss": 0.0192,
+      "llm_learning_rate": 1.8220875826113072e-06,
+      "loss": 0.7636,
+      "output_text_mask_loss": 0.0648,
+      "predictor_learning_rate": 0.0001822087582611307,
+      "step": 5854
+    },
+    {
+      "epoch": 0.6,
+      "gumbel_tau": 0.2513,
+      "image_masks_loss": 0.0112,
+      "llm_learning_rate": 1.8212890865602317e-06,
+      "loss": 0.6866,
+      "output_text_mask_loss": 0.0029,
+      "predictor_learning_rate": 0.00018212890865602316,
+      "step": 5855
+    },
+    {
+      "epoch": 0.6,
+      "gumbel_tau": 0.2512,
+      "image_masks_loss": 0.0072,
+      "llm_learning_rate": 1.8204906652503136e-06,
+      "loss": 0.7182,
+      "output_text_mask_loss": 0.0168,
+      "predictor_learning_rate": 0.00018204906652503132,
+      "step": 5856
+    },
+    {
+      "epoch": 0.6,
+      "gumbel_tau": 0.2512,
+      "image_masks_loss": 0.0145,
+      "llm_learning_rate": 1.8196923187694771e-06,
+      "loss": 0.7032,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.0001819692318769477,
+      "step": 5857
+    },
+    {
+      "epoch": 0.6,
+      "gumbel_tau": 0.2511,
+      "image_masks_loss": 0.0207,
+      "llm_learning_rate": 1.818894047205638e-06,
+      "loss": 0.7613,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.00018188940472056378,
+      "step": 5858
+    },
+    {
+      "epoch": 0.6,
+      "gumbel_tau": 0.2511,
+      "image_masks_loss": 0.0181,
+      "llm_learning_rate": 1.8180958506467033e-06,
+      "loss": 0.8664,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 0.00018180958506467033,
+      "step": 5859
+    },
+    {
+      "epoch": 0.6,
+      "gumbel_tau": 0.251,
+      "image_masks_loss": 0.0122,
+      "llm_learning_rate": 1.817297729180572e-06,
+      "loss": 0.7772,
+      "output_text_mask_loss": 0.0099,
+      "predictor_learning_rate": 0.0001817297729180572,
+      "step": 5860
+    },
+    {
+      "epoch": 0.6,
+      "gumbel_tau": 0.251,
+      "image_masks_loss": 0.0257,
+      "llm_learning_rate": 1.816499682895137e-06,
+      "loss": 0.7443,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.00018164996828951368,
+      "step": 5861
+    },
+    {
+      "epoch": 0.6,
+      "gumbel_tau": 0.2509,
+      "image_masks_loss": 0.0062,
+      "llm_learning_rate": 1.815701711878278e-06,
+      "loss": 0.658,
+      "output_text_mask_loss": 0.0086,
+      "predictor_learning_rate": 0.0001815701711878278,
+      "step": 5862
+    },
+    {
+      "epoch": 0.6,
+      "gumbel_tau": 0.2508,
+      "image_masks_loss": 0.0039,
+      "llm_learning_rate": 1.8149038162178711e-06,
+      "loss": 0.6964,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.0001814903816217871,
+      "step": 5863
+    },
+    {
+      "epoch": 0.6,
+      "gumbel_tau": 0.2508,
+      "image_masks_loss": 0.0086,
+      "llm_learning_rate": 1.8141059960017818e-06,
+      "loss": 0.7249,
+      "output_text_mask_loss": 0.0107,
+      "predictor_learning_rate": 0.00018141059960017817,
+      "step": 5864
+    },
+    {
+      "epoch": 0.6,
+      "gumbel_tau": 0.2507,
+      "image_masks_loss": 0.0463,
+      "llm_learning_rate": 1.8133082513178674e-06,
+      "loss": 0.8429,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.00018133082513178674,
+      "step": 5865
+    },
+    {
+      "epoch": 0.6,
+      "gumbel_tau": 0.2507,
+      "image_masks_loss": 0.0142,
+      "llm_learning_rate": 1.8125105822539785e-06,
+      "loss": 0.7739,
+      "output_text_mask_loss": 0.032,
+      "predictor_learning_rate": 0.00018125105822539784,
+      "step": 5866
+    },
+    {
+      "epoch": 0.6,
+      "gumbel_tau": 0.2506,
+      "image_masks_loss": 0.0167,
+      "llm_learning_rate": 1.8117129888979556e-06,
+      "loss": 0.7084,
+      "output_text_mask_loss": 0.0063,
+      "predictor_learning_rate": 0.00018117129888979555,
+      "step": 5867
+    },
+    {
+      "epoch": 0.6,
+      "gumbel_tau": 0.2505,
+      "image_masks_loss": 0.0145,
+      "llm_learning_rate": 1.8109154713376318e-06,
+      "loss": 0.7329,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.00018109154713376316,
+      "step": 5868
+    },
+    {
+      "epoch": 0.6,
+      "gumbel_tau": 0.2505,
+      "image_masks_loss": 0.0059,
+      "llm_learning_rate": 1.8101180296608302e-06,
+      "loss": 0.7871,
+      "output_text_mask_loss": 0.1389,
+      "predictor_learning_rate": 0.00018101180296608302,
+      "step": 5869
+    },
+    {
+      "epoch": 0.6,
+      "gumbel_tau": 0.2504,
+      "image_masks_loss": 0.0147,
+      "llm_learning_rate": 1.8093206639553688e-06,
+      "loss": 0.7669,
+      "output_text_mask_loss": 0.0277,
+      "predictor_learning_rate": 0.00018093206639553685,
+      "step": 5870
+    },
+    {
+      "epoch": 0.6,
+      "gumbel_tau": 0.2504,
+      "image_masks_loss": 0.0215,
+      "llm_learning_rate": 1.8085233743090542e-06,
+      "loss": 0.7279,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0001808523374309054,
+      "step": 5871
+    },
+    {
+      "epoch": 0.6,
+      "gumbel_tau": 0.2503,
+      "image_masks_loss": 0.014,
+      "llm_learning_rate": 1.807726160809685e-06,
+      "loss": 0.6467,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.0001807726160809685,
+      "step": 5872
+    },
+    {
+      "epoch": 0.6,
+      "gumbel_tau": 0.2502,
+      "image_masks_loss": 0.0149,
+      "llm_learning_rate": 1.8069290235450538e-06,
+      "loss": 0.6616,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.00018069290235450536,
+      "step": 5873
+    },
+    {
+      "epoch": 0.6,
+      "gumbel_tau": 0.2502,
+      "image_masks_loss": 0.006,
+      "llm_learning_rate": 1.8061319626029416e-06,
+      "loss": 0.7207,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 0.00018061319626029416,
+      "step": 5874
+    },
+    {
+      "epoch": 0.6,
+      "gumbel_tau": 0.2501,
+      "image_masks_loss": 0.0237,
+      "llm_learning_rate": 1.8053349780711238e-06,
+      "loss": 0.7791,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.00018053349780711237,
+      "step": 5875
+    },
+    {
+      "epoch": 0.6,
+      "gumbel_tau": 0.2501,
+      "image_masks_loss": 0.0133,
+      "llm_learning_rate": 1.8045380700373648e-06,
+      "loss": 0.8347,
+      "output_text_mask_loss": 0.0149,
+      "predictor_learning_rate": 0.00018045380700373648,
+      "step": 5876
+    },
+    {
+      "epoch": 0.6,
+      "gumbel_tau": 0.25,
+      "image_masks_loss": 0.0143,
+      "llm_learning_rate": 1.8037412385894238e-06,
+      "loss": 0.7854,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00018037412385894235,
+      "step": 5877
+    },
+    {
+      "epoch": 0.6,
+      "gumbel_tau": 0.2499,
+      "image_masks_loss": 0.0159,
+      "llm_learning_rate": 1.8029444838150477e-06,
+      "loss": 0.7678,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.00018029444838150478,
+      "step": 5878
+    },
+    {
+      "epoch": 0.6,
+      "gumbel_tau": 0.2499,
+      "image_masks_loss": 0.0197,
+      "llm_learning_rate": 1.8021478058019784e-06,
+      "loss": 0.7278,
+      "output_text_mask_loss": 0.0023,
+      "predictor_learning_rate": 0.00018021478058019783,
+      "step": 5879
+    },
+    {
+      "epoch": 0.6,
+      "gumbel_tau": 0.2498,
+      "image_masks_loss": 0.0265,
+      "llm_learning_rate": 1.8013512046379467e-06,
+      "loss": 0.7181,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00018013512046379464,
+      "step": 5880
+    },
+    {
+      "epoch": 0.6,
+      "gumbel_tau": 0.2498,
+      "image_masks_loss": 0.0303,
+      "llm_learning_rate": 1.8005546804106758e-06,
+      "loss": 0.805,
+      "output_text_mask_loss": 0.0103,
+      "predictor_learning_rate": 0.00018005546804106756,
+      "step": 5881
+    },
+    {
+      "epoch": 0.6,
+      "gumbel_tau": 0.2497,
+      "image_masks_loss": 0.0187,
+      "llm_learning_rate": 1.7997582332078823e-06,
+      "loss": 0.8182,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.00017997582332078822,
+      "step": 5882
+    },
+    {
+      "epoch": 0.6,
+      "gumbel_tau": 0.2497,
+      "image_masks_loss": 0.0139,
+      "llm_learning_rate": 1.7989618631172714e-06,
+      "loss": 0.7178,
+      "output_text_mask_loss": 0.0305,
+      "predictor_learning_rate": 0.00017989618631172714,
+      "step": 5883
+    },
+    {
+      "epoch": 0.6,
+      "gumbel_tau": 0.2496,
+      "image_masks_loss": 0.0138,
+      "llm_learning_rate": 1.7981655702265422e-06,
+      "loss": 0.709,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 0.0001798165570226542,
+      "step": 5884
+    },
+    {
+      "epoch": 0.6,
+      "gumbel_tau": 0.2495,
+      "image_masks_loss": 0.0134,
+      "llm_learning_rate": 1.7973693546233827e-06,
+      "loss": 0.625,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.00017973693546233826,
+      "step": 5885
+    },
+    {
+      "epoch": 0.6,
+      "gumbel_tau": 0.2495,
+      "image_masks_loss": 0.016,
+      "llm_learning_rate": 1.7965732163954756e-06,
+      "loss": 0.7805,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.00017965732163954755,
+      "step": 5886
+    },
+    {
+      "epoch": 0.6,
+      "gumbel_tau": 0.2494,
+      "image_masks_loss": 0.0077,
+      "llm_learning_rate": 1.7957771556304926e-06,
+      "loss": 0.8065,
+      "output_text_mask_loss": 0.0024,
+      "predictor_learning_rate": 0.00017957771556304924,
+      "step": 5887
+    },
+    {
+      "epoch": 0.6,
+      "gumbel_tau": 0.2494,
+      "image_masks_loss": 0.0204,
+      "llm_learning_rate": 1.7949811724160976e-06,
+      "loss": 0.7189,
+      "output_text_mask_loss": 0.0141,
+      "predictor_learning_rate": 0.00017949811724160975,
+      "step": 5888
+    },
+    {
+      "epoch": 0.6,
+      "gumbel_tau": 0.2493,
+      "image_masks_loss": 0.0119,
+      "llm_learning_rate": 1.7941852668399465e-06,
+      "loss": 0.6855,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.00017941852668399463,
+      "step": 5889
+    },
+    {
+      "epoch": 0.6,
+      "gumbel_tau": 0.2492,
+      "image_masks_loss": 0.0086,
+      "llm_learning_rate": 1.793389438989686e-06,
+      "loss": 0.7318,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.00017933894389896859,
+      "step": 5890
+    },
+    {
+      "epoch": 0.6,
+      "gumbel_tau": 0.2492,
+      "image_masks_loss": 0.0124,
+      "llm_learning_rate": 1.7925936889529544e-06,
+      "loss": 0.6769,
+      "output_text_mask_loss": 0.0063,
+      "predictor_learning_rate": 0.00017925936889529544,
+      "step": 5891
+    },
+    {
+      "epoch": 0.6,
+      "gumbel_tau": 0.2491,
+      "image_masks_loss": 0.0156,
+      "llm_learning_rate": 1.791798016817381e-06,
+      "loss": 0.7024,
+      "output_text_mask_loss": 0.0378,
+      "predictor_learning_rate": 0.0001791798016817381,
+      "step": 5892
+    },
+    {
+      "epoch": 0.6,
+      "gumbel_tau": 0.2491,
+      "image_masks_loss": 0.0147,
+      "llm_learning_rate": 1.791002422670588e-06,
+      "loss": 0.8008,
+      "output_text_mask_loss": 0.0381,
+      "predictor_learning_rate": 0.00017910024226705878,
+      "step": 5893
+    },
+    {
+      "epoch": 0.6,
+      "gumbel_tau": 0.249,
+      "image_masks_loss": 0.019,
+      "llm_learning_rate": 1.7902069066001875e-06,
+      "loss": 0.7521,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 0.00017902069066001874,
+      "step": 5894
+    },
+    {
+      "epoch": 0.6,
+      "gumbel_tau": 0.2489,
+      "image_masks_loss": 0.0068,
+      "llm_learning_rate": 1.789411468693783e-06,
+      "loss": 0.7198,
+      "output_text_mask_loss": 0.0021,
+      "predictor_learning_rate": 0.0001789411468693783,
+      "step": 5895
+    },
+    {
+      "epoch": 0.6,
+      "gumbel_tau": 0.2489,
+      "image_masks_loss": 0.0069,
+      "llm_learning_rate": 1.7886161090389705e-06,
+      "loss": 0.7842,
+      "output_text_mask_loss": 0.0463,
+      "predictor_learning_rate": 0.00017886161090389704,
+      "step": 5896
+    },
+    {
+      "epoch": 0.6,
+      "gumbel_tau": 0.2488,
+      "image_masks_loss": 0.0075,
+      "llm_learning_rate": 1.7878208277233368e-06,
+      "loss": 0.7493,
+      "output_text_mask_loss": 0.0034,
+      "predictor_learning_rate": 0.0001787820827723337,
+      "step": 5897
+    },
+    {
+      "epoch": 0.6,
+      "gumbel_tau": 0.2488,
+      "image_masks_loss": 0.0143,
+      "llm_learning_rate": 1.7870256248344606e-06,
+      "loss": 0.7551,
+      "output_text_mask_loss": 0.0368,
+      "predictor_learning_rate": 0.00017870256248344604,
+      "step": 5898
+    },
+    {
+      "epoch": 0.6,
+      "gumbel_tau": 0.2487,
+      "image_masks_loss": 0.0081,
+      "llm_learning_rate": 1.7862305004599097e-06,
+      "loss": 0.7775,
+      "output_text_mask_loss": 0.0175,
+      "predictor_learning_rate": 0.00017862305004599097,
+      "step": 5899
+    },
+    {
+      "epoch": 0.6,
+      "gumbel_tau": 0.2487,
+      "image_masks_loss": 0.0096,
+      "llm_learning_rate": 1.7854354546872465e-06,
+      "loss": 0.6612,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00017854354546872463,
+      "step": 5900
+    },
+    {
+      "epoch": 0.6,
+      "gumbel_tau": 0.2486,
+      "image_masks_loss": 0.0178,
+      "llm_learning_rate": 1.7846404876040218e-06,
+      "loss": 0.7536,
+      "output_text_mask_loss": 0.0088,
+      "predictor_learning_rate": 0.0001784640487604022,
+      "step": 5901
+    },
+    {
+      "epoch": 0.6,
+      "gumbel_tau": 0.2485,
+      "image_masks_loss": 0.0206,
+      "llm_learning_rate": 1.7838455992977805e-06,
+      "loss": 0.7542,
+      "output_text_mask_loss": 0.0109,
+      "predictor_learning_rate": 0.00017838455992977804,
+      "step": 5902
+    },
+    {
+      "epoch": 0.6,
+      "gumbel_tau": 0.2485,
+      "image_masks_loss": 0.0134,
+      "llm_learning_rate": 1.7830507898560573e-06,
+      "loss": 0.7501,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00017830507898560573,
+      "step": 5903
+    },
+    {
+      "epoch": 0.6,
+      "gumbel_tau": 0.2484,
+      "image_masks_loss": 0.0146,
+      "llm_learning_rate": 1.7822560593663771e-06,
+      "loss": 0.761,
+      "output_text_mask_loss": 0.0139,
+      "predictor_learning_rate": 0.0001782256059366377,
+      "step": 5904
+    },
+    {
+      "epoch": 0.61,
+      "gumbel_tau": 0.2484,
+      "image_masks_loss": 0.0166,
+      "llm_learning_rate": 1.781461407916259e-06,
+      "loss": 0.8013,
+      "output_text_mask_loss": 0.0012,
+      "predictor_learning_rate": 0.0001781461407916259,
+      "step": 5905
+    },
+    {
+      "epoch": 0.61,
+      "gumbel_tau": 0.2483,
+      "image_masks_loss": 0.0023,
+      "llm_learning_rate": 1.7806668355932113e-06,
+      "loss": 0.716,
+      "output_text_mask_loss": 0.0261,
+      "predictor_learning_rate": 0.00017806668355932112,
+      "step": 5906
+    },
+    {
+      "epoch": 0.61,
+      "gumbel_tau": 0.2482,
+      "image_masks_loss": 0.0166,
+      "llm_learning_rate": 1.7798723424847338e-06,
+      "loss": 0.741,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.00017798723424847338,
+      "step": 5907
+    },
+    {
+      "epoch": 0.61,
+      "gumbel_tau": 0.2482,
+      "image_masks_loss": 0.0059,
+      "llm_learning_rate": 1.7790779286783176e-06,
+      "loss": 0.7459,
+      "output_text_mask_loss": 0.0023,
+      "predictor_learning_rate": 0.00017790779286783174,
+      "step": 5908
+    },
+    {
+      "epoch": 0.61,
+      "gumbel_tau": 0.2481,
+      "image_masks_loss": 0.0175,
+      "llm_learning_rate": 1.778283594261446e-06,
+      "loss": 0.7368,
+      "output_text_mask_loss": 0.0062,
+      "predictor_learning_rate": 0.00017782835942614457,
+      "step": 5909
+    },
+    {
+      "epoch": 0.61,
+      "gumbel_tau": 0.2481,
+      "image_masks_loss": 0.0117,
+      "llm_learning_rate": 1.7774893393215925e-06,
+      "loss": 0.7211,
+      "output_text_mask_loss": 0.0381,
+      "predictor_learning_rate": 0.00017774893393215924,
+      "step": 5910
+    },
+    {
+      "epoch": 0.61,
+      "gumbel_tau": 0.248,
+      "image_masks_loss": 0.0056,
+      "llm_learning_rate": 1.7766951639462222e-06,
+      "loss": 0.7674,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0001776695163946222,
+      "step": 5911
+    },
+    {
+      "epoch": 0.61,
+      "gumbel_tau": 0.248,
+      "image_masks_loss": 0.0132,
+      "llm_learning_rate": 1.7759010682227918e-06,
+      "loss": 0.7407,
+      "output_text_mask_loss": 0.0401,
+      "predictor_learning_rate": 0.00017759010682227917,
+      "step": 5912
+    },
+    {
+      "epoch": 0.61,
+      "gumbel_tau": 0.2479,
+      "image_masks_loss": 0.0159,
+      "llm_learning_rate": 1.7751070522387484e-06,
+      "loss": 0.7333,
+      "output_text_mask_loss": 0.0175,
+      "predictor_learning_rate": 0.00017751070522387482,
+      "step": 5913
+    },
+    {
+      "epoch": 0.61,
+      "gumbel_tau": 0.2478,
+      "image_masks_loss": 0.0104,
+      "llm_learning_rate": 1.7743131160815314e-06,
+      "loss": 0.7489,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.00017743131160815313,
+      "step": 5914
+    },
+    {
+      "epoch": 0.61,
+      "gumbel_tau": 0.2478,
+      "image_masks_loss": 0.0128,
+      "llm_learning_rate": 1.7735192598385697e-06,
+      "loss": 0.7697,
+      "output_text_mask_loss": 0.0004,
+      "predictor_learning_rate": 0.00017735192598385696,
+      "step": 5915
+    },
+    {
+      "epoch": 0.61,
+      "gumbel_tau": 0.2477,
+      "image_masks_loss": 0.0073,
+      "llm_learning_rate": 1.7727254835972868e-06,
+      "loss": 0.7509,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.00017727254835972866,
+      "step": 5916
+    },
+    {
+      "epoch": 0.61,
+      "gumbel_tau": 0.2477,
+      "image_masks_loss": 0.0041,
+      "llm_learning_rate": 1.7719317874450922e-06,
+      "loss": 0.6563,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0001771931787445092,
+      "step": 5917
+    },
+    {
+      "epoch": 0.61,
+      "gumbel_tau": 0.2476,
+      "image_masks_loss": 0.0081,
+      "llm_learning_rate": 1.7711381714693915e-06,
+      "loss": 0.6959,
+      "output_text_mask_loss": 0.0172,
+      "predictor_learning_rate": 0.00017711381714693913,
+      "step": 5918
+    },
+    {
+      "epoch": 0.61,
+      "gumbel_tau": 0.2475,
+      "image_masks_loss": 0.0084,
+      "llm_learning_rate": 1.7703446357575787e-06,
+      "loss": 0.7124,
+      "output_text_mask_loss": 0.0319,
+      "predictor_learning_rate": 0.00017703446357575787,
+      "step": 5919
+    },
+    {
+      "epoch": 0.61,
+      "gumbel_tau": 0.2475,
+      "image_masks_loss": 0.0035,
+      "llm_learning_rate": 1.7695511803970394e-06,
+      "loss": 0.7224,
+      "output_text_mask_loss": 0.0086,
+      "predictor_learning_rate": 0.00017695511803970392,
+      "step": 5920
+    },
+    {
+      "epoch": 0.61,
+      "gumbel_tau": 0.2474,
+      "image_masks_loss": 0.0129,
+      "llm_learning_rate": 1.7687578054751514e-06,
+      "loss": 0.7693,
+      "output_text_mask_loss": 0.0242,
+      "predictor_learning_rate": 0.00017687578054751513,
+      "step": 5921
+    },
+    {
+      "epoch": 0.61,
+      "gumbel_tau": 0.2474,
+      "image_masks_loss": 0.0062,
+      "llm_learning_rate": 1.7679645110792829e-06,
+      "loss": 0.6551,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.00017679645110792828,
+      "step": 5922
+    },
+    {
+      "epoch": 0.61,
+      "gumbel_tau": 0.2473,
+      "image_masks_loss": 0.0222,
+      "llm_learning_rate": 1.767171297296793e-06,
+      "loss": 0.8347,
+      "output_text_mask_loss": 0.0381,
+      "predictor_learning_rate": 0.0001767171297296793,
+      "step": 5923
+    },
+    {
+      "epoch": 0.61,
+      "gumbel_tau": 0.2472,
+      "image_masks_loss": 0.0136,
+      "llm_learning_rate": 1.7663781642150313e-06,
+      "loss": 0.7238,
+      "output_text_mask_loss": 0.0185,
+      "predictor_learning_rate": 0.00017663781642150313,
+      "step": 5924
+    },
+    {
+      "epoch": 0.61,
+      "gumbel_tau": 0.2472,
+      "image_masks_loss": 0.0154,
+      "llm_learning_rate": 1.765585111921341e-06,
+      "loss": 0.829,
+      "output_text_mask_loss": 0.0058,
+      "predictor_learning_rate": 0.00017655851119213409,
+      "step": 5925
+    },
+    {
+      "epoch": 0.61,
+      "gumbel_tau": 0.2471,
+      "image_masks_loss": 0.0085,
+      "llm_learning_rate": 1.764792140503054e-06,
+      "loss": 0.7671,
+      "output_text_mask_loss": 0.0062,
+      "predictor_learning_rate": 0.0001764792140503054,
+      "step": 5926
+    },
+    {
+      "epoch": 0.61,
+      "gumbel_tau": 0.2471,
+      "image_masks_loss": 0.006,
+      "llm_learning_rate": 1.7639992500474938e-06,
+      "loss": 0.8072,
+      "output_text_mask_loss": 0.0035,
+      "predictor_learning_rate": 0.00017639992500474935,
+      "step": 5927
+    },
+    {
+      "epoch": 0.61,
+      "gumbel_tau": 0.247,
+      "image_masks_loss": 0.0072,
+      "llm_learning_rate": 1.7632064406419758e-06,
+      "loss": 0.7703,
+      "output_text_mask_loss": 0.0078,
+      "predictor_learning_rate": 0.00017632064406419756,
+      "step": 5928
+    },
+    {
+      "epoch": 0.61,
+      "gumbel_tau": 0.247,
+      "image_masks_loss": 0.0167,
+      "llm_learning_rate": 1.7624137123738056e-06,
+      "loss": 0.7598,
+      "output_text_mask_loss": 0.0132,
+      "predictor_learning_rate": 0.00017624137123738054,
+      "step": 5929
+    },
+    {
+      "epoch": 0.61,
+      "gumbel_tau": 0.2469,
+      "image_masks_loss": 0.0193,
+      "llm_learning_rate": 1.7616210653302807e-06,
+      "loss": 0.7006,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.00017616210653302806,
+      "step": 5930
+    },
+    {
+      "epoch": 0.61,
+      "gumbel_tau": 0.2468,
+      "image_masks_loss": 0.0211,
+      "llm_learning_rate": 1.760828499598688e-06,
+      "loss": 0.7092,
+      "output_text_mask_loss": 0.0488,
+      "predictor_learning_rate": 0.00017608284995986878,
+      "step": 5931
+    },
+    {
+      "epoch": 0.61,
+      "gumbel_tau": 0.2468,
+      "image_masks_loss": 0.009,
+      "llm_learning_rate": 1.7600360152663076e-06,
+      "loss": 0.6528,
+      "output_text_mask_loss": 0.0362,
+      "predictor_learning_rate": 0.00017600360152663074,
+      "step": 5932
+    },
+    {
+      "epoch": 0.61,
+      "gumbel_tau": 0.2467,
+      "image_masks_loss": 0.0189,
+      "llm_learning_rate": 1.7592436124204098e-06,
+      "loss": 0.7549,
+      "output_text_mask_loss": 0.0143,
+      "predictor_learning_rate": 0.00017592436124204096,
+      "step": 5933
+    },
+    {
+      "epoch": 0.61,
+      "gumbel_tau": 0.2467,
+      "image_masks_loss": 0.0052,
+      "llm_learning_rate": 1.758451291148255e-06,
+      "loss": 0.7708,
+      "output_text_mask_loss": 0.033,
+      "predictor_learning_rate": 0.0001758451291148255,
+      "step": 5934
+    },
+    {
+      "epoch": 0.61,
+      "gumbel_tau": 0.2466,
+      "image_masks_loss": 0.0152,
+      "llm_learning_rate": 1.7576590515370967e-06,
+      "loss": 0.7115,
+      "output_text_mask_loss": 0.0147,
+      "predictor_learning_rate": 0.00017576590515370967,
+      "step": 5935
+    },
+    {
+      "epoch": 0.61,
+      "gumbel_tau": 0.2466,
+      "image_masks_loss": 0.0227,
+      "llm_learning_rate": 1.7568668936741761e-06,
+      "loss": 0.85,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0001756866893674176,
+      "step": 5936
+    },
+    {
+      "epoch": 0.61,
+      "gumbel_tau": 0.2465,
+      "image_masks_loss": 0.0033,
+      "llm_learning_rate": 1.7560748176467287e-06,
+      "loss": 0.7055,
+      "output_text_mask_loss": 0.0587,
+      "predictor_learning_rate": 0.00017560748176467287,
+      "step": 5937
+    },
+    {
+      "epoch": 0.61,
+      "gumbel_tau": 0.2464,
+      "image_masks_loss": 0.0038,
+      "llm_learning_rate": 1.7552828235419798e-06,
+      "loss": 0.76,
+      "output_text_mask_loss": 0.0336,
+      "predictor_learning_rate": 0.00017552828235419796,
+      "step": 5938
+    },
+    {
+      "epoch": 0.61,
+      "gumbel_tau": 0.2464,
+      "image_masks_loss": 0.0078,
+      "llm_learning_rate": 1.7544909114471447e-06,
+      "loss": 0.7139,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00017544909114471446,
+      "step": 5939
+    },
+    {
+      "epoch": 0.61,
+      "gumbel_tau": 0.2463,
+      "image_masks_loss": 0.0216,
+      "llm_learning_rate": 1.7536990814494306e-06,
+      "loss": 0.6959,
+      "output_text_mask_loss": 0.0387,
+      "predictor_learning_rate": 0.00017536990814494304,
+      "step": 5940
+    },
+    {
+      "epoch": 0.61,
+      "gumbel_tau": 0.2463,
+      "image_masks_loss": 0.0068,
+      "llm_learning_rate": 1.7529073336360364e-06,
+      "loss": 0.7473,
+      "output_text_mask_loss": 0.0035,
+      "predictor_learning_rate": 0.00017529073336360363,
+      "step": 5941
+    },
+    {
+      "epoch": 0.61,
+      "gumbel_tau": 0.2462,
+      "image_masks_loss": 0.0176,
+      "llm_learning_rate": 1.7521156680941509e-06,
+      "loss": 0.6557,
+      "output_text_mask_loss": 0.0058,
+      "predictor_learning_rate": 0.0001752115668094151,
+      "step": 5942
+    },
+    {
+      "epoch": 0.61,
+      "gumbel_tau": 0.2461,
+      "image_masks_loss": 0.0083,
+      "llm_learning_rate": 1.7513240849109534e-06,
+      "loss": 0.7584,
+      "output_text_mask_loss": 0.0492,
+      "predictor_learning_rate": 0.00017513240849109532,
+      "step": 5943
+    },
+    {
+      "epoch": 0.61,
+      "gumbel_tau": 0.2461,
+      "image_masks_loss": 0.0103,
+      "llm_learning_rate": 1.7505325841736158e-06,
+      "loss": 0.7555,
+      "output_text_mask_loss": 0.0429,
+      "predictor_learning_rate": 0.00017505325841736156,
+      "step": 5944
+    },
+    {
+      "epoch": 0.61,
+      "gumbel_tau": 0.246,
+      "image_masks_loss": 0.009,
+      "llm_learning_rate": 1.7497411659692993e-06,
+      "loss": 0.817,
+      "output_text_mask_loss": 0.0035,
+      "predictor_learning_rate": 0.00017497411659692991,
+      "step": 5945
+    },
+    {
+      "epoch": 0.61,
+      "gumbel_tau": 0.246,
+      "image_masks_loss": 0.0334,
+      "llm_learning_rate": 1.7489498303851565e-06,
+      "loss": 0.7731,
+      "output_text_mask_loss": 0.0137,
+      "predictor_learning_rate": 0.00017489498303851564,
+      "step": 5946
+    },
+    {
+      "epoch": 0.61,
+      "gumbel_tau": 0.2459,
+      "image_masks_loss": 0.0165,
+      "llm_learning_rate": 1.748158577508331e-06,
+      "loss": 0.7007,
+      "output_text_mask_loss": 0.0185,
+      "predictor_learning_rate": 0.00017481585775083307,
+      "step": 5947
+    },
+    {
+      "epoch": 0.61,
+      "gumbel_tau": 0.2459,
+      "image_masks_loss": 0.0146,
+      "llm_learning_rate": 1.7473674074259583e-06,
+      "loss": 0.8311,
+      "output_text_mask_loss": 0.0081,
+      "predictor_learning_rate": 0.00017473674074259583,
+      "step": 5948
+    },
+    {
+      "epoch": 0.61,
+      "gumbel_tau": 0.2458,
+      "image_masks_loss": 0.0357,
+      "llm_learning_rate": 1.746576320225163e-06,
+      "loss": 0.8125,
+      "output_text_mask_loss": 0.0071,
+      "predictor_learning_rate": 0.00017465763202251626,
+      "step": 5949
+    },
+    {
+      "epoch": 0.61,
+      "gumbel_tau": 0.2457,
+      "image_masks_loss": 0.0196,
+      "llm_learning_rate": 1.7457853159930613e-06,
+      "loss": 0.7635,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.00017457853159930613,
+      "step": 5950
+    },
+    {
+      "epoch": 0.61,
+      "gumbel_tau": 0.2457,
+      "image_masks_loss": 0.0157,
+      "llm_learning_rate": 1.7449943948167602e-06,
+      "loss": 0.6888,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.000174499439481676,
+      "step": 5951
+    },
+    {
+      "epoch": 0.61,
+      "gumbel_tau": 0.2456,
+      "image_masks_loss": 0.0252,
+      "llm_learning_rate": 1.7442035567833585e-06,
+      "loss": 0.6748,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00017442035567833582,
+      "step": 5952
+    },
+    {
+      "epoch": 0.61,
+      "gumbel_tau": 0.2456,
+      "image_masks_loss": 0.0054,
+      "llm_learning_rate": 1.7434128019799452e-06,
+      "loss": 0.6787,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0001743412801979945,
+      "step": 5953
+    },
+    {
+      "epoch": 0.61,
+      "gumbel_tau": 0.2455,
+      "image_masks_loss": 0.0184,
+      "llm_learning_rate": 1.7426221304935986e-06,
+      "loss": 0.6631,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.00017426221304935984,
+      "step": 5954
+    },
+    {
+      "epoch": 0.61,
+      "gumbel_tau": 0.2454,
+      "image_masks_loss": 0.012,
+      "llm_learning_rate": 1.7418315424113902e-06,
+      "loss": 0.7032,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.000174183154241139,
+      "step": 5955
+    },
+    {
+      "epoch": 0.61,
+      "gumbel_tau": 0.2454,
+      "image_masks_loss": 0.0228,
+      "llm_learning_rate": 1.7410410378203807e-06,
+      "loss": 0.8018,
+      "output_text_mask_loss": 0.0018,
+      "predictor_learning_rate": 0.00017410410378203807,
+      "step": 5956
+    },
+    {
+      "epoch": 0.61,
+      "gumbel_tau": 0.2453,
+      "image_masks_loss": 0.0311,
+      "llm_learning_rate": 1.740250616807623e-06,
+      "loss": 0.8002,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.0001740250616807623,
+      "step": 5957
+    },
+    {
+      "epoch": 0.61,
+      "gumbel_tau": 0.2453,
+      "image_masks_loss": 0.0219,
+      "llm_learning_rate": 1.7394602794601598e-06,
+      "loss": 0.7472,
+      "output_text_mask_loss": 0.0095,
+      "predictor_learning_rate": 0.00017394602794601596,
+      "step": 5958
+    },
+    {
+      "epoch": 0.61,
+      "gumbel_tau": 0.2452,
+      "image_masks_loss": 0.0149,
+      "llm_learning_rate": 1.7386700258650247e-06,
+      "loss": 0.7789,
+      "output_text_mask_loss": 0.012,
+      "predictor_learning_rate": 0.00017386700258650246,
+      "step": 5959
+    },
+    {
+      "epoch": 0.61,
+      "gumbel_tau": 0.2452,
+      "image_masks_loss": 0.0203,
+      "llm_learning_rate": 1.7378798561092413e-06,
+      "loss": 0.79,
+      "output_text_mask_loss": 0.0055,
+      "predictor_learning_rate": 0.00017378798561092413,
+      "step": 5960
+    },
+    {
+      "epoch": 0.61,
+      "gumbel_tau": 0.2451,
+      "image_masks_loss": 0.0187,
+      "llm_learning_rate": 1.7370897702798265e-06,
+      "loss": 0.8109,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.00017370897702798265,
+      "step": 5961
+    },
+    {
+      "epoch": 0.61,
+      "gumbel_tau": 0.245,
+      "image_masks_loss": 0.0055,
+      "llm_learning_rate": 1.7362997684637853e-06,
+      "loss": 0.6993,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0001736299768463785,
+      "step": 5962
+    },
+    {
+      "epoch": 0.61,
+      "gumbel_tau": 0.245,
+      "image_masks_loss": 0.0107,
+      "llm_learning_rate": 1.7355098507481144e-06,
+      "loss": 0.7605,
+      "output_text_mask_loss": 0.0061,
+      "predictor_learning_rate": 0.00017355098507481143,
+      "step": 5963
+    },
+    {
+      "epoch": 0.61,
+      "gumbel_tau": 0.2449,
+      "image_masks_loss": 0.0077,
+      "llm_learning_rate": 1.734720017219802e-06,
+      "loss": 0.7616,
+      "output_text_mask_loss": 0.0046,
+      "predictor_learning_rate": 0.00017347200172198018,
+      "step": 5964
+    },
+    {
+      "epoch": 0.61,
+      "gumbel_tau": 0.2449,
+      "image_masks_loss": 0.016,
+      "llm_learning_rate": 1.7339302679658257e-06,
+      "loss": 0.7207,
+      "output_text_mask_loss": 0.0324,
+      "predictor_learning_rate": 0.00017339302679658257,
+      "step": 5965
+    },
+    {
+      "epoch": 0.61,
+      "gumbel_tau": 0.2448,
+      "image_masks_loss": 0.0059,
+      "llm_learning_rate": 1.7331406030731546e-06,
+      "loss": 0.84,
+      "output_text_mask_loss": 0.001,
+      "predictor_learning_rate": 0.00017331406030731544,
+      "step": 5966
+    },
+    {
+      "epoch": 0.61,
+      "gumbel_tau": 0.2448,
+      "image_masks_loss": 0.0056,
+      "llm_learning_rate": 1.7323510226287482e-06,
+      "loss": 0.7478,
+      "output_text_mask_loss": 0.0124,
+      "predictor_learning_rate": 0.0001732351022628748,
+      "step": 5967
+    },
+    {
+      "epoch": 0.61,
+      "gumbel_tau": 0.2447,
+      "image_masks_loss": 0.0061,
+      "llm_learning_rate": 1.7315615267195573e-06,
+      "loss": 0.6786,
+      "output_text_mask_loss": 0.0012,
+      "predictor_learning_rate": 0.00017315615267195572,
+      "step": 5968
+    },
+    {
+      "epoch": 0.61,
+      "gumbel_tau": 0.2446,
+      "image_masks_loss": 0.0075,
+      "llm_learning_rate": 1.730772115432523e-06,
+      "loss": 0.6498,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00017307721154325229,
+      "step": 5969
+    },
+    {
+      "epoch": 0.61,
+      "gumbel_tau": 0.2446,
+      "image_masks_loss": 0.0117,
+      "llm_learning_rate": 1.7299827888545762e-06,
+      "loss": 0.769,
+      "output_text_mask_loss": 0.0012,
+      "predictor_learning_rate": 0.00017299827888545761,
+      "step": 5970
+    },
+    {
+      "epoch": 0.61,
+      "gumbel_tau": 0.2445,
+      "image_masks_loss": 0.0072,
+      "llm_learning_rate": 1.7291935470726402e-06,
+      "loss": 0.786,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.000172919354707264,
+      "step": 5971
+    },
+    {
+      "epoch": 0.61,
+      "gumbel_tau": 0.2445,
+      "image_masks_loss": 0.0147,
+      "llm_learning_rate": 1.728404390173628e-06,
+      "loss": 0.7021,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.0001728404390173628,
+      "step": 5972
+    },
+    {
+      "epoch": 0.61,
+      "gumbel_tau": 0.2444,
+      "image_masks_loss": 0.0114,
+      "llm_learning_rate": 1.727615318244443e-06,
+      "loss": 0.8076,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.0001727615318244443,
+      "step": 5973
+    },
+    {
+      "epoch": 0.61,
+      "gumbel_tau": 0.2443,
+      "image_masks_loss": 0.0167,
+      "llm_learning_rate": 1.7268263313719794e-06,
+      "loss": 0.7625,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00017268263313719793,
+      "step": 5974
+    },
+    {
+      "epoch": 0.61,
+      "gumbel_tau": 0.2443,
+      "image_masks_loss": 0.0059,
+      "llm_learning_rate": 1.726037429643122e-06,
+      "loss": 0.8157,
+      "output_text_mask_loss": 0.0071,
+      "predictor_learning_rate": 0.00017260374296431218,
+      "step": 5975
+    },
+    {
+      "epoch": 0.61,
+      "gumbel_tau": 0.2442,
+      "image_masks_loss": 0.0175,
+      "llm_learning_rate": 1.7252486131447466e-06,
+      "loss": 0.7318,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00017252486131447465,
+      "step": 5976
+    },
+    {
+      "epoch": 0.61,
+      "gumbel_tau": 0.2442,
+      "image_masks_loss": 0.0048,
+      "llm_learning_rate": 1.7244598819637204e-06,
+      "loss": 0.672,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00017244598819637203,
+      "step": 5977
+    },
+    {
+      "epoch": 0.61,
+      "gumbel_tau": 0.2441,
+      "image_masks_loss": 0.0157,
+      "llm_learning_rate": 1.723671236186899e-06,
+      "loss": 0.7741,
+      "output_text_mask_loss": 0.0032,
+      "predictor_learning_rate": 0.0001723671236186899,
+      "step": 5978
+    },
+    {
+      "epoch": 0.61,
+      "gumbel_tau": 0.2441,
+      "image_masks_loss": 0.0059,
+      "llm_learning_rate": 1.7228826759011298e-06,
+      "loss": 0.8447,
+      "output_text_mask_loss": 0.0698,
+      "predictor_learning_rate": 0.00017228826759011295,
+      "step": 5979
+    },
+    {
+      "epoch": 0.61,
+      "gumbel_tau": 0.244,
+      "image_masks_loss": 0.0074,
+      "llm_learning_rate": 1.7220942011932519e-06,
+      "loss": 0.72,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.00017220942011932517,
+      "step": 5980
+    },
+    {
+      "epoch": 0.61,
+      "gumbel_tau": 0.2439,
+      "image_masks_loss": 0.0111,
+      "llm_learning_rate": 1.7213058121500934e-06,
+      "loss": 0.6916,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.00017213058121500933,
+      "step": 5981
+    },
+    {
+      "epoch": 0.61,
+      "gumbel_tau": 0.2439,
+      "image_masks_loss": 0.0017,
+      "llm_learning_rate": 1.720517508858473e-06,
+      "loss": 0.841,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.00017205175088584728,
+      "step": 5982
+    },
+    {
+      "epoch": 0.61,
+      "gumbel_tau": 0.2438,
+      "image_masks_loss": 0.0128,
+      "llm_learning_rate": 1.7197292914052005e-06,
+      "loss": 0.7622,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.00017197292914052003,
+      "step": 5983
+    },
+    {
+      "epoch": 0.61,
+      "gumbel_tau": 0.2438,
+      "image_masks_loss": 0.009,
+      "llm_learning_rate": 1.7189411598770773e-06,
+      "loss": 0.7596,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0001718941159877077,
+      "step": 5984
+    },
+    {
+      "epoch": 0.61,
+      "gumbel_tau": 0.2437,
+      "image_masks_loss": 0.0261,
+      "llm_learning_rate": 1.7181531143608927e-06,
+      "loss": 0.6535,
+      "output_text_mask_loss": 0.0679,
+      "predictor_learning_rate": 0.00017181531143608926,
+      "step": 5985
+    },
+    {
+      "epoch": 0.61,
+      "gumbel_tau": 0.2437,
+      "image_masks_loss": 0.0176,
+      "llm_learning_rate": 1.7173651549434283e-06,
+      "loss": 0.7412,
+      "output_text_mask_loss": 0.0046,
+      "predictor_learning_rate": 0.00017173651549434282,
+      "step": 5986
+    },
+    {
+      "epoch": 0.61,
+      "gumbel_tau": 0.2436,
+      "image_masks_loss": 0.0088,
+      "llm_learning_rate": 1.7165772817114567e-06,
+      "loss": 0.6915,
+      "output_text_mask_loss": 0.0408,
+      "predictor_learning_rate": 0.00017165772817114567,
+      "step": 5987
+    },
+    {
+      "epoch": 0.61,
+      "gumbel_tau": 0.2435,
+      "image_masks_loss": 0.0126,
+      "llm_learning_rate": 1.71578949475174e-06,
+      "loss": 0.7709,
+      "output_text_mask_loss": 0.0052,
+      "predictor_learning_rate": 0.000171578949475174,
+      "step": 5988
+    },
+    {
+      "epoch": 0.61,
+      "gumbel_tau": 0.2435,
+      "image_masks_loss": 0.0154,
+      "llm_learning_rate": 1.715001794151031e-06,
+      "loss": 0.775,
+      "output_text_mask_loss": 0.0095,
+      "predictor_learning_rate": 0.0001715001794151031,
+      "step": 5989
+    },
+    {
+      "epoch": 0.61,
+      "gumbel_tau": 0.2434,
+      "image_masks_loss": 0.0077,
+      "llm_learning_rate": 1.7142141799960732e-06,
+      "loss": 0.8693,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.0001714214179996073,
+      "step": 5990
+    },
+    {
+      "epoch": 0.61,
+      "gumbel_tau": 0.2434,
+      "image_masks_loss": 0.0187,
+      "llm_learning_rate": 1.7134266523735999e-06,
+      "loss": 0.7034,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 0.00017134266523735996,
+      "step": 5991
+    },
+    {
+      "epoch": 0.61,
+      "gumbel_tau": 0.2433,
+      "image_masks_loss": 0.011,
+      "llm_learning_rate": 1.712639211370335e-06,
+      "loss": 0.7486,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 0.0001712639211370335,
+      "step": 5992
+    },
+    {
+      "epoch": 0.61,
+      "gumbel_tau": 0.2433,
+      "image_masks_loss": 0.02,
+      "llm_learning_rate": 1.7118518570729944e-06,
+      "loss": 0.7776,
+      "output_text_mask_loss": 0.0238,
+      "predictor_learning_rate": 0.00017118518570729944,
+      "step": 5993
+    },
+    {
+      "epoch": 0.61,
+      "gumbel_tau": 0.2432,
+      "image_masks_loss": 0.017,
+      "llm_learning_rate": 1.711064589568283e-06,
+      "loss": 0.7524,
+      "output_text_mask_loss": 0.0124,
+      "predictor_learning_rate": 0.0001711064589568283,
+      "step": 5994
+    },
+    {
+      "epoch": 0.61,
+      "gumbel_tau": 0.2431,
+      "image_masks_loss": 0.0061,
+      "llm_learning_rate": 1.7102774089428959e-06,
+      "loss": 0.7132,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00017102774089428958,
+      "step": 5995
+    },
+    {
+      "epoch": 0.61,
+      "gumbel_tau": 0.2431,
+      "image_masks_loss": 0.0248,
+      "llm_learning_rate": 1.7094903152835195e-06,
+      "loss": 0.7666,
+      "output_text_mask_loss": 0.012,
+      "predictor_learning_rate": 0.00017094903152835194,
+      "step": 5996
+    },
+    {
+      "epoch": 0.61,
+      "gumbel_tau": 0.243,
+      "image_masks_loss": 0.02,
+      "llm_learning_rate": 1.7087033086768305e-06,
+      "loss": 0.7715,
+      "output_text_mask_loss": 0.001,
+      "predictor_learning_rate": 0.00017087033086768305,
+      "step": 5997
+    },
+    {
+      "epoch": 0.61,
+      "gumbel_tau": 0.243,
+      "image_masks_loss": 0.0184,
+      "llm_learning_rate": 1.7079163892094963e-06,
+      "loss": 0.6956,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00017079163892094962,
+      "step": 5998
+    },
+    {
+      "epoch": 0.61,
+      "gumbel_tau": 0.2429,
+      "image_masks_loss": 0.0079,
+      "llm_learning_rate": 1.7071295569681728e-06,
+      "loss": 0.8151,
+      "output_text_mask_loss": 0.0233,
+      "predictor_learning_rate": 0.00017071295569681729,
+      "step": 5999
+    },
+    {
+      "epoch": 0.61,
+      "gumbel_tau": 0.2429,
+      "image_masks_loss": 0.0116,
+      "llm_learning_rate": 1.7063428120395087e-06,
+      "loss": 0.8142,
+      "output_text_mask_loss": 0.0216,
+      "predictor_learning_rate": 0.00017063428120395085,
+      "step": 6000
+    },
+    {
+      "epoch": 0.61,
+      "gumbel_tau": 0.2428,
+      "image_masks_loss": 0.0031,
+      "llm_learning_rate": 1.7055561545101423e-06,
+      "loss": 0.8089,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.00017055561545101423,
+      "step": 6001
+    },
+    {
+      "epoch": 0.61,
+      "gumbel_tau": 0.2427,
+      "image_masks_loss": 0.0114,
+      "llm_learning_rate": 1.7047695844667011e-06,
+      "loss": 0.757,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.0001704769584466701,
+      "step": 6002
+    },
+    {
+      "epoch": 0.62,
+      "gumbel_tau": 0.2427,
+      "image_masks_loss": 0.0048,
+      "llm_learning_rate": 1.7039831019958053e-06,
+      "loss": 0.6956,
+      "output_text_mask_loss": 0.0078,
+      "predictor_learning_rate": 0.00017039831019958053,
+      "step": 6003
+    },
+    {
+      "epoch": 0.62,
+      "gumbel_tau": 0.2426,
+      "image_masks_loss": 0.0233,
+      "llm_learning_rate": 1.703196707184063e-06,
+      "loss": 0.8231,
+      "output_text_mask_loss": 0.0309,
+      "predictor_learning_rate": 0.0001703196707184063,
+      "step": 6004
+    },
+    {
+      "epoch": 0.62,
+      "gumbel_tau": 0.2426,
+      "image_masks_loss": 0.0189,
+      "llm_learning_rate": 1.702410400118074e-06,
+      "loss": 0.8058,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.0001702410400118074,
+      "step": 6005
+    },
+    {
+      "epoch": 0.62,
+      "gumbel_tau": 0.2425,
+      "image_masks_loss": 0.0152,
+      "llm_learning_rate": 1.7016241808844284e-06,
+      "loss": 0.6151,
+      "output_text_mask_loss": 0.0125,
+      "predictor_learning_rate": 0.00017016241808844281,
+      "step": 6006
+    },
+    {
+      "epoch": 0.62,
+      "gumbel_tau": 0.2425,
+      "image_masks_loss": 0.0072,
+      "llm_learning_rate": 1.7008380495697063e-06,
+      "loss": 0.7242,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.00017008380495697062,
+      "step": 6007
+    },
+    {
+      "epoch": 0.62,
+      "gumbel_tau": 0.2424,
+      "image_masks_loss": 0.0037,
+      "llm_learning_rate": 1.7000520062604785e-06,
+      "loss": 0.6729,
+      "output_text_mask_loss": 0.0107,
+      "predictor_learning_rate": 0.00017000520062604785,
+      "step": 6008
+    },
+    {
+      "epoch": 0.62,
+      "gumbel_tau": 0.2423,
+      "image_masks_loss": 0.0171,
+      "llm_learning_rate": 1.6992660510433062e-06,
+      "loss": 0.7857,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.00016992660510433061,
+      "step": 6009
+    },
+    {
+      "epoch": 0.62,
+      "gumbel_tau": 0.2423,
+      "image_masks_loss": 0.0159,
+      "llm_learning_rate": 1.6984801840047394e-06,
+      "loss": 0.728,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.0001698480184004739,
+      "step": 6010
+    },
+    {
+      "epoch": 0.62,
+      "gumbel_tau": 0.2422,
+      "image_masks_loss": 0.0227,
+      "llm_learning_rate": 1.6976944052313194e-06,
+      "loss": 0.6747,
+      "output_text_mask_loss": 0.0309,
+      "predictor_learning_rate": 0.00016976944052313192,
+      "step": 6011
+    },
+    {
+      "epoch": 0.62,
+      "gumbel_tau": 0.2422,
+      "image_masks_loss": 0.0121,
+      "llm_learning_rate": 1.6969087148095787e-06,
+      "loss": 0.7336,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00016969087148095786,
+      "step": 6012
+    },
+    {
+      "epoch": 0.62,
+      "gumbel_tau": 0.2421,
+      "image_masks_loss": 0.0074,
+      "llm_learning_rate": 1.6961231128260398e-06,
+      "loss": 0.7817,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.00016961231128260397,
+      "step": 6013
+    },
+    {
+      "epoch": 0.62,
+      "gumbel_tau": 0.2421,
+      "image_masks_loss": 0.0203,
+      "llm_learning_rate": 1.6953375993672138e-06,
+      "loss": 0.7109,
+      "output_text_mask_loss": 0.0099,
+      "predictor_learning_rate": 0.00016953375993672138,
+      "step": 6014
+    },
+    {
+      "epoch": 0.62,
+      "gumbel_tau": 0.242,
+      "image_masks_loss": 0.0109,
+      "llm_learning_rate": 1.6945521745196035e-06,
+      "loss": 0.6712,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.00016945521745196032,
+      "step": 6015
+    },
+    {
+      "epoch": 0.62,
+      "gumbel_tau": 0.2419,
+      "image_masks_loss": 0.0213,
+      "llm_learning_rate": 1.6937668383697021e-06,
+      "loss": 0.6846,
+      "output_text_mask_loss": 0.0641,
+      "predictor_learning_rate": 0.0001693766838369702,
+      "step": 6016
+    },
+    {
+      "epoch": 0.62,
+      "gumbel_tau": 0.2419,
+      "image_masks_loss": 0.0096,
+      "llm_learning_rate": 1.692981591003992e-06,
+      "loss": 0.7526,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.0001692981591003992,
+      "step": 6017
+    },
+    {
+      "epoch": 0.62,
+      "gumbel_tau": 0.2418,
+      "image_masks_loss": 0.0164,
+      "llm_learning_rate": 1.6921964325089468e-06,
+      "loss": 0.7061,
+      "output_text_mask_loss": 0.0713,
+      "predictor_learning_rate": 0.00016921964325089468,
+      "step": 6018
+    },
+    {
+      "epoch": 0.62,
+      "gumbel_tau": 0.2418,
+      "image_masks_loss": 0.0219,
+      "llm_learning_rate": 1.6914113629710298e-06,
+      "loss": 0.7582,
+      "output_text_mask_loss": 0.0073,
+      "predictor_learning_rate": 0.00016914113629710298,
+      "step": 6019
+    },
+    {
+      "epoch": 0.62,
+      "gumbel_tau": 0.2417,
+      "image_masks_loss": 0.0042,
+      "llm_learning_rate": 1.690626382476695e-06,
+      "loss": 0.8017,
+      "output_text_mask_loss": 0.0259,
+      "predictor_learning_rate": 0.0001690626382476695,
+      "step": 6020
+    },
+    {
+      "epoch": 0.62,
+      "gumbel_tau": 0.2417,
+      "image_masks_loss": 0.0141,
+      "llm_learning_rate": 1.6898414911123853e-06,
+      "loss": 0.7668,
+      "output_text_mask_loss": 0.0479,
+      "predictor_learning_rate": 0.00016898414911123853,
+      "step": 6021
+    },
+    {
+      "epoch": 0.62,
+      "gumbel_tau": 0.2416,
+      "image_masks_loss": 0.0143,
+      "llm_learning_rate": 1.689056688964535e-06,
+      "loss": 0.7282,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00016890566889645348,
+      "step": 6022
+    },
+    {
+      "epoch": 0.62,
+      "gumbel_tau": 0.2415,
+      "image_masks_loss": 0.0186,
+      "llm_learning_rate": 1.6882719761195686e-06,
+      "loss": 0.7358,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.00016882719761195686,
+      "step": 6023
+    },
+    {
+      "epoch": 0.62,
+      "gumbel_tau": 0.2415,
+      "image_masks_loss": 0.0068,
+      "llm_learning_rate": 1.6874873526639004e-06,
+      "loss": 0.7288,
+      "output_text_mask_loss": 0.0286,
+      "predictor_learning_rate": 0.00016874873526639004,
+      "step": 6024
+    },
+    {
+      "epoch": 0.62,
+      "gumbel_tau": 0.2414,
+      "image_masks_loss": 0.0152,
+      "llm_learning_rate": 1.686702818683934e-06,
+      "loss": 0.7272,
+      "output_text_mask_loss": 0.0217,
+      "predictor_learning_rate": 0.0001686702818683934,
+      "step": 6025
+    },
+    {
+      "epoch": 0.62,
+      "gumbel_tau": 0.2414,
+      "image_masks_loss": 0.0292,
+      "llm_learning_rate": 1.6859183742660657e-06,
+      "loss": 0.8256,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00016859183742660655,
+      "step": 6026
+    },
+    {
+      "epoch": 0.62,
+      "gumbel_tau": 0.2413,
+      "image_masks_loss": 0.034,
+      "llm_learning_rate": 1.6851340194966798e-06,
+      "loss": 0.7194,
+      "output_text_mask_loss": 0.0012,
+      "predictor_learning_rate": 0.00016851340194966796,
+      "step": 6027
+    },
+    {
+      "epoch": 0.62,
+      "gumbel_tau": 0.2413,
+      "image_masks_loss": 0.009,
+      "llm_learning_rate": 1.6843497544621498e-06,
+      "loss": 0.7961,
+      "output_text_mask_loss": 0.021,
+      "predictor_learning_rate": 0.00016843497544621498,
+      "step": 6028
+    },
+    {
+      "epoch": 0.62,
+      "gumbel_tau": 0.2412,
+      "image_masks_loss": 0.0126,
+      "llm_learning_rate": 1.683565579248842e-06,
+      "loss": 0.8229,
+      "output_text_mask_loss": 0.0198,
+      "predictor_learning_rate": 0.0001683565579248842,
+      "step": 6029
+    },
+    {
+      "epoch": 0.62,
+      "gumbel_tau": 0.2411,
+      "image_masks_loss": 0.033,
+      "llm_learning_rate": 1.682781493943111e-06,
+      "loss": 0.7323,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00016827814939431108,
+      "step": 6030
+    },
+    {
+      "epoch": 0.62,
+      "gumbel_tau": 0.2411,
+      "image_masks_loss": 0.0261,
+      "llm_learning_rate": 1.6819974986313015e-06,
+      "loss": 0.7771,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 0.00016819974986313015,
+      "step": 6031
+    },
+    {
+      "epoch": 0.62,
+      "gumbel_tau": 0.241,
+      "image_masks_loss": 0.0178,
+      "llm_learning_rate": 1.6812135933997504e-06,
+      "loss": 0.7599,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 0.000168121359339975,
+      "step": 6032
+    },
+    {
+      "epoch": 0.62,
+      "gumbel_tau": 0.241,
+      "image_masks_loss": 0.0103,
+      "llm_learning_rate": 1.6804297783347818e-06,
+      "loss": 0.8165,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.00016804297783347817,
+      "step": 6033
+    },
+    {
+      "epoch": 0.62,
+      "gumbel_tau": 0.2409,
+      "image_masks_loss": 0.0051,
+      "llm_learning_rate": 1.679646053522711e-06,
+      "loss": 0.6936,
+      "output_text_mask_loss": 0.013,
+      "predictor_learning_rate": 0.0001679646053522711,
+      "step": 6034
+    },
+    {
+      "epoch": 0.62,
+      "gumbel_tau": 0.2409,
+      "image_masks_loss": 0.0096,
+      "llm_learning_rate": 1.6788624190498448e-06,
+      "loss": 0.7267,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00016788624190498447,
+      "step": 6035
+    },
+    {
+      "epoch": 0.62,
+      "gumbel_tau": 0.2408,
+      "image_masks_loss": 0.0132,
+      "llm_learning_rate": 1.678078875002478e-06,
+      "loss": 0.7199,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.0001678078875002478,
+      "step": 6036
+    },
+    {
+      "epoch": 0.62,
+      "gumbel_tau": 0.2407,
+      "image_masks_loss": 0.0079,
+      "llm_learning_rate": 1.6772954214668962e-06,
+      "loss": 0.7195,
+      "output_text_mask_loss": 0.0095,
+      "predictor_learning_rate": 0.0001677295421466896,
+      "step": 6037
+    },
+    {
+      "epoch": 0.62,
+      "gumbel_tau": 0.2407,
+      "image_masks_loss": 0.011,
+      "llm_learning_rate": 1.6765120585293745e-06,
+      "loss": 0.7642,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.00016765120585293745,
+      "step": 6038
+    },
+    {
+      "epoch": 0.62,
+      "gumbel_tau": 0.2406,
+      "image_masks_loss": 0.0173,
+      "llm_learning_rate": 1.6757287862761799e-06,
+      "loss": 0.8807,
+      "output_text_mask_loss": 0.0052,
+      "predictor_learning_rate": 0.00016757287862761798,
+      "step": 6039
+    },
+    {
+      "epoch": 0.62,
+      "gumbel_tau": 0.2406,
+      "image_masks_loss": 0.0118,
+      "llm_learning_rate": 1.6749456047935668e-06,
+      "loss": 0.6842,
+      "output_text_mask_loss": 0.001,
+      "predictor_learning_rate": 0.00016749456047935668,
+      "step": 6040
+    },
+    {
+      "epoch": 0.62,
+      "gumbel_tau": 0.2405,
+      "image_masks_loss": 0.0064,
+      "llm_learning_rate": 1.6741625141677815e-06,
+      "loss": 0.7115,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00016741625141677814,
+      "step": 6041
+    },
+    {
+      "epoch": 0.62,
+      "gumbel_tau": 0.2405,
+      "image_masks_loss": 0.0119,
+      "llm_learning_rate": 1.6733795144850601e-06,
+      "loss": 0.7197,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.000167337951448506,
+      "step": 6042
+    },
+    {
+      "epoch": 0.62,
+      "gumbel_tau": 0.2404,
+      "image_masks_loss": 0.0136,
+      "llm_learning_rate": 1.6725966058316275e-06,
+      "loss": 0.8428,
+      "output_text_mask_loss": 0.0187,
+      "predictor_learning_rate": 0.00016725966058316275,
+      "step": 6043
+    },
+    {
+      "epoch": 0.62,
+      "gumbel_tau": 0.2403,
+      "image_masks_loss": 0.0046,
+      "llm_learning_rate": 1.6718137882936996e-06,
+      "loss": 0.6374,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.00016718137882936994,
+      "step": 6044
+    },
+    {
+      "epoch": 0.62,
+      "gumbel_tau": 0.2403,
+      "image_masks_loss": 0.0248,
+      "llm_learning_rate": 1.6710310619574821e-06,
+      "loss": 0.7486,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.0001671031061957482,
+      "step": 6045
+    },
+    {
+      "epoch": 0.62,
+      "gumbel_tau": 0.2402,
+      "image_masks_loss": 0.0232,
+      "llm_learning_rate": 1.6702484269091715e-06,
+      "loss": 0.792,
+      "output_text_mask_loss": 0.0259,
+      "predictor_learning_rate": 0.00016702484269091715,
+      "step": 6046
+    },
+    {
+      "epoch": 0.62,
+      "gumbel_tau": 0.2402,
+      "image_masks_loss": 0.0064,
+      "llm_learning_rate": 1.669465883234951e-06,
+      "loss": 0.814,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.00016694658832349507,
+      "step": 6047
+    },
+    {
+      "epoch": 0.62,
+      "gumbel_tau": 0.2401,
+      "image_masks_loss": 0.0064,
+      "llm_learning_rate": 1.6686834310209977e-06,
+      "loss": 0.7562,
+      "output_text_mask_loss": 0.0166,
+      "predictor_learning_rate": 0.00016686834310209974,
+      "step": 6048
+    },
+    {
+      "epoch": 0.62,
+      "gumbel_tau": 0.2401,
+      "image_masks_loss": 0.0084,
+      "llm_learning_rate": 1.667901070353477e-06,
+      "loss": 0.7272,
+      "output_text_mask_loss": 0.0051,
+      "predictor_learning_rate": 0.0001667901070353477,
+      "step": 6049
+    },
+    {
+      "epoch": 0.62,
+      "gumbel_tau": 0.24,
+      "image_masks_loss": 0.0353,
+      "llm_learning_rate": 1.667118801318543e-06,
+      "loss": 0.6937,
+      "output_text_mask_loss": 0.0322,
+      "predictor_learning_rate": 0.00016671188013185428,
+      "step": 6050
+    },
+    {
+      "epoch": 0.62,
+      "gumbel_tau": 0.24,
+      "image_masks_loss": 0.0078,
+      "llm_learning_rate": 1.6663366240023424e-06,
+      "loss": 0.7602,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.00016663366240023423,
+      "step": 6051
+    },
+    {
+      "epoch": 0.62,
+      "gumbel_tau": 0.2399,
+      "image_masks_loss": 0.0125,
+      "llm_learning_rate": 1.66555453849101e-06,
+      "loss": 0.6862,
+      "output_text_mask_loss": 0.0055,
+      "predictor_learning_rate": 0.00016655545384910097,
+      "step": 6052
+    },
+    {
+      "epoch": 0.62,
+      "gumbel_tau": 0.2398,
+      "image_masks_loss": 0.0232,
+      "llm_learning_rate": 1.6647725448706697e-06,
+      "loss": 0.8485,
+      "output_text_mask_loss": 0.0244,
+      "predictor_learning_rate": 0.00016647725448706697,
+      "step": 6053
+    },
+    {
+      "epoch": 0.62,
+      "gumbel_tau": 0.2398,
+      "image_masks_loss": 0.0136,
+      "llm_learning_rate": 1.6639906432274372e-06,
+      "loss": 0.7338,
+      "output_text_mask_loss": 0.0557,
+      "predictor_learning_rate": 0.0001663990643227437,
+      "step": 6054
+    },
+    {
+      "epoch": 0.62,
+      "gumbel_tau": 0.2397,
+      "image_masks_loss": 0.0081,
+      "llm_learning_rate": 1.6632088336474178e-06,
+      "loss": 0.7914,
+      "output_text_mask_loss": 0.0717,
+      "predictor_learning_rate": 0.00016632088336474178,
+      "step": 6055
+    },
+    {
+      "epoch": 0.62,
+      "gumbel_tau": 0.2397,
+      "image_masks_loss": 0.0047,
+      "llm_learning_rate": 1.6624271162167054e-06,
+      "loss": 0.5777,
+      "output_text_mask_loss": 0.0233,
+      "predictor_learning_rate": 0.00016624271162167054,
+      "step": 6056
+    },
+    {
+      "epoch": 0.62,
+      "gumbel_tau": 0.2396,
+      "image_masks_loss": 0.0082,
+      "llm_learning_rate": 1.6616454910213838e-06,
+      "loss": 0.78,
+      "output_text_mask_loss": 0.002,
+      "predictor_learning_rate": 0.00016616454910213838,
+      "step": 6057
+    },
+    {
+      "epoch": 0.62,
+      "gumbel_tau": 0.2396,
+      "image_masks_loss": 0.028,
+      "llm_learning_rate": 1.660863958147529e-06,
+      "loss": 0.6794,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00016608639581475288,
+      "step": 6058
+    },
+    {
+      "epoch": 0.62,
+      "gumbel_tau": 0.2395,
+      "image_masks_loss": 0.0147,
+      "llm_learning_rate": 1.660082517681204e-06,
+      "loss": 0.7261,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.0001660082517681204,
+      "step": 6059
+    },
+    {
+      "epoch": 0.62,
+      "gumbel_tau": 0.2394,
+      "image_masks_loss": 0.0201,
+      "llm_learning_rate": 1.6593011697084633e-06,
+      "loss": 0.7874,
+      "output_text_mask_loss": 0.0355,
+      "predictor_learning_rate": 0.00016593011697084632,
+      "step": 6060
+    },
+    {
+      "epoch": 0.62,
+      "gumbel_tau": 0.2394,
+      "image_masks_loss": 0.0214,
+      "llm_learning_rate": 1.6585199143153494e-06,
+      "loss": 0.7501,
+      "output_text_mask_loss": 0.0006,
+      "predictor_learning_rate": 0.0001658519914315349,
+      "step": 6061
+    },
+    {
+      "epoch": 0.62,
+      "gumbel_tau": 0.2393,
+      "image_masks_loss": 0.0041,
+      "llm_learning_rate": 1.6577387515878975e-06,
+      "loss": 0.7122,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.00016577387515878976,
+      "step": 6062
+    },
+    {
+      "epoch": 0.62,
+      "gumbel_tau": 0.2393,
+      "image_masks_loss": 0.0168,
+      "llm_learning_rate": 1.6569576816121305e-06,
+      "loss": 0.7151,
+      "output_text_mask_loss": 0.0231,
+      "predictor_learning_rate": 0.00016569576816121304,
+      "step": 6063
+    },
+    {
+      "epoch": 0.62,
+      "gumbel_tau": 0.2392,
+      "image_masks_loss": 0.0148,
+      "llm_learning_rate": 1.6561767044740618e-06,
+      "loss": 0.6979,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.00016561767044740615,
+      "step": 6064
+    },
+    {
+      "epoch": 0.62,
+      "gumbel_tau": 0.2392,
+      "image_masks_loss": 0.0232,
+      "llm_learning_rate": 1.6553958202596937e-06,
+      "loss": 0.7701,
+      "output_text_mask_loss": 0.0275,
+      "predictor_learning_rate": 0.00016553958202596936,
+      "step": 6065
+    },
+    {
+      "epoch": 0.62,
+      "gumbel_tau": 0.2391,
+      "image_masks_loss": 0.0115,
+      "llm_learning_rate": 1.6546150290550184e-06,
+      "loss": 0.7506,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.00016546150290550182,
+      "step": 6066
+    },
+    {
+      "epoch": 0.62,
+      "gumbel_tau": 0.239,
+      "image_masks_loss": 0.0129,
+      "llm_learning_rate": 1.6538343309460195e-06,
+      "loss": 0.8026,
+      "output_text_mask_loss": 0.0298,
+      "predictor_learning_rate": 0.00016538343309460194,
+      "step": 6067
+    },
+    {
+      "epoch": 0.62,
+      "gumbel_tau": 0.239,
+      "image_masks_loss": 0.0058,
+      "llm_learning_rate": 1.653053726018669e-06,
+      "loss": 0.7215,
+      "output_text_mask_loss": 0.0081,
+      "predictor_learning_rate": 0.0001653053726018669,
+      "step": 6068
+    },
+    {
+      "epoch": 0.62,
+      "gumbel_tau": 0.2389,
+      "image_masks_loss": 0.0094,
+      "llm_learning_rate": 1.6522732143589288e-06,
+      "loss": 0.7652,
+      "output_text_mask_loss": 0.0086,
+      "predictor_learning_rate": 0.00016522732143589287,
+      "step": 6069
+    },
+    {
+      "epoch": 0.62,
+      "gumbel_tau": 0.2389,
+      "image_masks_loss": 0.0137,
+      "llm_learning_rate": 1.65149279605275e-06,
+      "loss": 0.7409,
+      "output_text_mask_loss": 0.0429,
+      "predictor_learning_rate": 0.00016514927960527498,
+      "step": 6070
+    },
+    {
+      "epoch": 0.62,
+      "gumbel_tau": 0.2388,
+      "image_masks_loss": 0.0141,
+      "llm_learning_rate": 1.6507124711860745e-06,
+      "loss": 0.8215,
+      "output_text_mask_loss": 0.0032,
+      "predictor_learning_rate": 0.00016507124711860743,
+      "step": 6071
+    },
+    {
+      "epoch": 0.62,
+      "gumbel_tau": 0.2388,
+      "image_masks_loss": 0.0254,
+      "llm_learning_rate": 1.6499322398448337e-06,
+      "loss": 0.759,
+      "output_text_mask_loss": 0.0081,
+      "predictor_learning_rate": 0.00016499322398448336,
+      "step": 6072
+    },
+    {
+      "epoch": 0.62,
+      "gumbel_tau": 0.2387,
+      "image_masks_loss": 0.0108,
+      "llm_learning_rate": 1.6491521021149476e-06,
+      "loss": 0.7336,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.00016491521021149474,
+      "step": 6073
+    },
+    {
+      "epoch": 0.62,
+      "gumbel_tau": 0.2387,
+      "image_masks_loss": 0.0082,
+      "llm_learning_rate": 1.6483720580823276e-06,
+      "loss": 0.8095,
+      "output_text_mask_loss": 0.0078,
+      "predictor_learning_rate": 0.00016483720580823274,
+      "step": 6074
+    },
+    {
+      "epoch": 0.62,
+      "gumbel_tau": 0.2386,
+      "image_masks_loss": 0.0105,
+      "llm_learning_rate": 1.647592107832874e-06,
+      "loss": 0.766,
+      "output_text_mask_loss": 0.0099,
+      "predictor_learning_rate": 0.00016475921078328736,
+      "step": 6075
+    },
+    {
+      "epoch": 0.62,
+      "gumbel_tau": 0.2385,
+      "image_masks_loss": 0.0122,
+      "llm_learning_rate": 1.6468122514524754e-06,
+      "loss": 0.7142,
+      "output_text_mask_loss": 0.0399,
+      "predictor_learning_rate": 0.00016468122514524752,
+      "step": 6076
+    },
+    {
+      "epoch": 0.62,
+      "gumbel_tau": 0.2385,
+      "image_masks_loss": 0.0168,
+      "llm_learning_rate": 1.646032489027012e-06,
+      "loss": 0.6089,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0001646032489027012,
+      "step": 6077
+    },
+    {
+      "epoch": 0.62,
+      "gumbel_tau": 0.2384,
+      "image_masks_loss": 0.0182,
+      "llm_learning_rate": 1.6452528206423535e-06,
+      "loss": 0.7899,
+      "output_text_mask_loss": 0.0656,
+      "predictor_learning_rate": 0.00016452528206423533,
+      "step": 6078
+    },
+    {
+      "epoch": 0.62,
+      "gumbel_tau": 0.2384,
+      "image_masks_loss": 0.0047,
+      "llm_learning_rate": 1.6444732463843582e-06,
+      "loss": 0.7019,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.0001644473246384358,
+      "step": 6079
+    },
+    {
+      "epoch": 0.62,
+      "gumbel_tau": 0.2383,
+      "image_masks_loss": 0.0165,
+      "llm_learning_rate": 1.643693766338874e-06,
+      "loss": 0.7131,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0001643693766338874,
+      "step": 6080
+    },
+    {
+      "epoch": 0.62,
+      "gumbel_tau": 0.2383,
+      "image_masks_loss": 0.0192,
+      "llm_learning_rate": 1.6429143805917404e-06,
+      "loss": 0.7089,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.00016429143805917403,
+      "step": 6081
+    },
+    {
+      "epoch": 0.62,
+      "gumbel_tau": 0.2382,
+      "image_masks_loss": 0.0204,
+      "llm_learning_rate": 1.642135089228784e-06,
+      "loss": 0.7364,
+      "output_text_mask_loss": 0.0196,
+      "predictor_learning_rate": 0.0001642135089228784,
+      "step": 6082
+    },
+    {
+      "epoch": 0.62,
+      "gumbel_tau": 0.2381,
+      "image_masks_loss": 0.0032,
+      "llm_learning_rate": 1.6413558923358233e-06,
+      "loss": 0.7309,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00016413558923358232,
+      "step": 6083
+    },
+    {
+      "epoch": 0.62,
+      "gumbel_tau": 0.2381,
+      "image_masks_loss": 0.0298,
+      "llm_learning_rate": 1.6405767899986635e-06,
+      "loss": 0.7578,
+      "output_text_mask_loss": 0.0029,
+      "predictor_learning_rate": 0.00016405767899986635,
+      "step": 6084
+    },
+    {
+      "epoch": 0.62,
+      "gumbel_tau": 0.238,
+      "image_masks_loss": 0.0114,
+      "llm_learning_rate": 1.639797782303102e-06,
+      "loss": 0.7808,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 0.0001639797782303102,
+      "step": 6085
+    },
+    {
+      "epoch": 0.62,
+      "gumbel_tau": 0.238,
+      "image_masks_loss": 0.0265,
+      "llm_learning_rate": 1.639018869334924e-06,
+      "loss": 0.7972,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.0001639018869334924,
+      "step": 6086
+    },
+    {
+      "epoch": 0.62,
+      "gumbel_tau": 0.2379,
+      "image_masks_loss": 0.0082,
+      "llm_learning_rate": 1.6382400511799068e-06,
+      "loss": 0.7457,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.00016382400511799066,
+      "step": 6087
+    },
+    {
+      "epoch": 0.62,
+      "gumbel_tau": 0.2379,
+      "image_masks_loss": 0.0269,
+      "llm_learning_rate": 1.6374613279238144e-06,
+      "loss": 0.8003,
+      "output_text_mask_loss": 0.0065,
+      "predictor_learning_rate": 0.00016374613279238143,
+      "step": 6088
+    },
+    {
+      "epoch": 0.62,
+      "gumbel_tau": 0.2378,
+      "image_masks_loss": 0.0187,
+      "llm_learning_rate": 1.636682699652401e-06,
+      "loss": 0.7547,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.00016366826996524008,
+      "step": 6089
+    },
+    {
+      "epoch": 0.62,
+      "gumbel_tau": 0.2378,
+      "image_masks_loss": 0.013,
+      "llm_learning_rate": 1.6359041664514124e-06,
+      "loss": 0.6397,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00016359041664514124,
+      "step": 6090
+    },
+    {
+      "epoch": 0.62,
+      "gumbel_tau": 0.2377,
+      "image_masks_loss": 0.0099,
+      "llm_learning_rate": 1.635125728406582e-06,
+      "loss": 0.8164,
+      "output_text_mask_loss": 0.0401,
+      "predictor_learning_rate": 0.00016351257284065817,
+      "step": 6091
+    },
+    {
+      "epoch": 0.62,
+      "gumbel_tau": 0.2376,
+      "image_masks_loss": 0.014,
+      "llm_learning_rate": 1.6343473856036323e-06,
+      "loss": 0.7378,
+      "output_text_mask_loss": 0.0378,
+      "predictor_learning_rate": 0.00016343473856036324,
+      "step": 6092
+    },
+    {
+      "epoch": 0.62,
+      "gumbel_tau": 0.2376,
+      "image_masks_loss": 0.0187,
+      "llm_learning_rate": 1.6335691381282763e-06,
+      "loss": 0.7616,
+      "output_text_mask_loss": 0.024,
+      "predictor_learning_rate": 0.00016335691381282763,
+      "step": 6093
+    },
+    {
+      "epoch": 0.62,
+      "gumbel_tau": 0.2375,
+      "image_masks_loss": 0.0216,
+      "llm_learning_rate": 1.6327909860662178e-06,
+      "loss": 0.7439,
+      "output_text_mask_loss": 0.001,
+      "predictor_learning_rate": 0.00016327909860662176,
+      "step": 6094
+    },
+    {
+      "epoch": 0.62,
+      "gumbel_tau": 0.2375,
+      "image_masks_loss": 0.022,
+      "llm_learning_rate": 1.6320129295031467e-06,
+      "loss": 0.8252,
+      "output_text_mask_loss": 0.0181,
+      "predictor_learning_rate": 0.00016320129295031467,
+      "step": 6095
+    },
+    {
+      "epoch": 0.62,
+      "gumbel_tau": 0.2374,
+      "image_masks_loss": 0.0099,
+      "llm_learning_rate": 1.6312349685247453e-06,
+      "loss": 0.775,
+      "output_text_mask_loss": 0.0336,
+      "predictor_learning_rate": 0.00016312349685247452,
+      "step": 6096
+    },
+    {
+      "epoch": 0.62,
+      "gumbel_tau": 0.2374,
+      "image_masks_loss": 0.0185,
+      "llm_learning_rate": 1.6304571032166844e-06,
+      "loss": 0.7637,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00016304571032166844,
+      "step": 6097
+    },
+    {
+      "epoch": 0.62,
+      "gumbel_tau": 0.2373,
+      "image_masks_loss": 0.0134,
+      "llm_learning_rate": 1.6296793336646243e-06,
+      "loss": 0.6387,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00016296793336646243,
+      "step": 6098
+    },
+    {
+      "epoch": 0.62,
+      "gumbel_tau": 0.2372,
+      "image_masks_loss": 0.0045,
+      "llm_learning_rate": 1.6289016599542146e-06,
+      "loss": 0.7342,
+      "output_text_mask_loss": 0.0141,
+      "predictor_learning_rate": 0.00016289016599542144,
+      "step": 6099
+    },
+    {
+      "epoch": 0.62,
+      "gumbel_tau": 0.2372,
+      "image_masks_loss": 0.002,
+      "llm_learning_rate": 1.6281240821710937e-06,
+      "loss": 0.6888,
+      "output_text_mask_loss": 0.0177,
+      "predictor_learning_rate": 0.00016281240821710937,
+      "step": 6100
+    },
+    {
+      "epoch": 0.63,
+      "gumbel_tau": 0.2371,
+      "image_masks_loss": 0.0234,
+      "llm_learning_rate": 1.6273466004008915e-06,
+      "loss": 0.6463,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.00016273466004008914,
+      "step": 6101
+    },
+    {
+      "epoch": 0.63,
+      "gumbel_tau": 0.2371,
+      "image_masks_loss": 0.0128,
+      "llm_learning_rate": 1.6265692147292259e-06,
+      "loss": 0.7407,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.00016265692147292258,
+      "step": 6102
+    },
+    {
+      "epoch": 0.63,
+      "gumbel_tau": 0.237,
+      "image_masks_loss": 0.0134,
+      "llm_learning_rate": 1.6257919252417032e-06,
+      "loss": 0.705,
+      "output_text_mask_loss": 0.0021,
+      "predictor_learning_rate": 0.00016257919252417032,
+      "step": 6103
+    },
+    {
+      "epoch": 0.63,
+      "gumbel_tau": 0.237,
+      "image_masks_loss": 0.01,
+      "llm_learning_rate": 1.6250147320239212e-06,
+      "loss": 0.7283,
+      "output_text_mask_loss": 0.0029,
+      "predictor_learning_rate": 0.0001625014732023921,
+      "step": 6104
+    },
+    {
+      "epoch": 0.63,
+      "gumbel_tau": 0.2369,
+      "image_masks_loss": 0.0178,
+      "llm_learning_rate": 1.6242376351614653e-06,
+      "loss": 0.7532,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.00016242376351614652,
+      "step": 6105
+    },
+    {
+      "epoch": 0.63,
+      "gumbel_tau": 0.2369,
+      "image_masks_loss": 0.0217,
+      "llm_learning_rate": 1.623460634739912e-06,
+      "loss": 0.7809,
+      "output_text_mask_loss": 0.0242,
+      "predictor_learning_rate": 0.0001623460634739912,
+      "step": 6106
+    },
+    {
+      "epoch": 0.63,
+      "gumbel_tau": 0.2368,
+      "image_masks_loss": 0.0062,
+      "llm_learning_rate": 1.6226837308448262e-06,
+      "loss": 0.6372,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.0001622683730844826,
+      "step": 6107
+    },
+    {
+      "epoch": 0.63,
+      "gumbel_tau": 0.2367,
+      "image_masks_loss": 0.0151,
+      "llm_learning_rate": 1.621906923561762e-06,
+      "loss": 0.7457,
+      "output_text_mask_loss": 0.0328,
+      "predictor_learning_rate": 0.0001621906923561762,
+      "step": 6108
+    },
+    {
+      "epoch": 0.63,
+      "gumbel_tau": 0.2367,
+      "image_masks_loss": 0.009,
+      "llm_learning_rate": 1.621130212976263e-06,
+      "loss": 0.6843,
+      "output_text_mask_loss": 0.046,
+      "predictor_learning_rate": 0.00016211302129762629,
+      "step": 6109
+    },
+    {
+      "epoch": 0.63,
+      "gumbel_tau": 0.2366,
+      "image_masks_loss": 0.0101,
+      "llm_learning_rate": 1.6203535991738627e-06,
+      "loss": 0.7164,
+      "output_text_mask_loss": 0.0175,
+      "predictor_learning_rate": 0.00016203535991738628,
+      "step": 6110
+    },
+    {
+      "epoch": 0.63,
+      "gumbel_tau": 0.2366,
+      "image_masks_loss": 0.0046,
+      "llm_learning_rate": 1.6195770822400837e-06,
+      "loss": 0.7016,
+      "output_text_mask_loss": 0.0546,
+      "predictor_learning_rate": 0.00016195770822400834,
+      "step": 6111
+    },
+    {
+      "epoch": 0.63,
+      "gumbel_tau": 0.2365,
+      "image_masks_loss": 0.0042,
+      "llm_learning_rate": 1.618800662260437e-06,
+      "loss": 0.7627,
+      "output_text_mask_loss": 0.0717,
+      "predictor_learning_rate": 0.00016188006622604368,
+      "step": 6112
+    },
+    {
+      "epoch": 0.63,
+      "gumbel_tau": 0.2365,
+      "image_masks_loss": 0.0109,
+      "llm_learning_rate": 1.6180243393204242e-06,
+      "loss": 0.6283,
+      "output_text_mask_loss": 0.0313,
+      "predictor_learning_rate": 0.0001618024339320424,
+      "step": 6113
+    },
+    {
+      "epoch": 0.63,
+      "gumbel_tau": 0.2364,
+      "image_masks_loss": 0.0085,
+      "llm_learning_rate": 1.6172481135055363e-06,
+      "loss": 0.7112,
+      "output_text_mask_loss": 0.0052,
+      "predictor_learning_rate": 0.0001617248113505536,
+      "step": 6114
+    },
+    {
+      "epoch": 0.63,
+      "gumbel_tau": 0.2364,
+      "image_masks_loss": 0.0095,
+      "llm_learning_rate": 1.6164719849012517e-06,
+      "loss": 0.7982,
+      "output_text_mask_loss": 0.0429,
+      "predictor_learning_rate": 0.00016164719849012517,
+      "step": 6115
+    },
+    {
+      "epoch": 0.63,
+      "gumbel_tau": 0.2363,
+      "image_masks_loss": 0.016,
+      "llm_learning_rate": 1.61569595359304e-06,
+      "loss": 0.7613,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.00016156959535930398,
+      "step": 6116
+    },
+    {
+      "epoch": 0.63,
+      "gumbel_tau": 0.2362,
+      "image_masks_loss": 0.0109,
+      "llm_learning_rate": 1.61492001966636e-06,
+      "loss": 0.7051,
+      "output_text_mask_loss": 0.0839,
+      "predictor_learning_rate": 0.000161492001966636,
+      "step": 6117
+    },
+    {
+      "epoch": 0.63,
+      "gumbel_tau": 0.2362,
+      "image_masks_loss": 0.0127,
+      "llm_learning_rate": 1.6141441832066589e-06,
+      "loss": 0.7325,
+      "output_text_mask_loss": 0.0071,
+      "predictor_learning_rate": 0.0001614144183206659,
+      "step": 6118
+    },
+    {
+      "epoch": 0.63,
+      "gumbel_tau": 0.2361,
+      "image_masks_loss": 0.003,
+      "llm_learning_rate": 1.6133684442993734e-06,
+      "loss": 0.784,
+      "output_text_mask_loss": 0.0071,
+      "predictor_learning_rate": 0.0001613368444299373,
+      "step": 6119
+    },
+    {
+      "epoch": 0.63,
+      "gumbel_tau": 0.2361,
+      "image_masks_loss": 0.0091,
+      "llm_learning_rate": 1.6125928030299308e-06,
+      "loss": 0.7434,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00016125928030299308,
+      "step": 6120
+    },
+    {
+      "epoch": 0.63,
+      "gumbel_tau": 0.236,
+      "image_masks_loss": 0.0028,
+      "llm_learning_rate": 1.611817259483744e-06,
+      "loss": 0.7388,
+      "output_text_mask_loss": 0.0086,
+      "predictor_learning_rate": 0.00016118172594837437,
+      "step": 6121
+    },
+    {
+      "epoch": 0.63,
+      "gumbel_tau": 0.236,
+      "image_masks_loss": 0.0263,
+      "llm_learning_rate": 1.6110418137462192e-06,
+      "loss": 0.7801,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.0001611041813746219,
+      "step": 6122
+    },
+    {
+      "epoch": 0.63,
+      "gumbel_tau": 0.2359,
+      "image_masks_loss": 0.0298,
+      "llm_learning_rate": 1.61026646590275e-06,
+      "loss": 0.7308,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 0.000161026646590275,
+      "step": 6123
+    },
+    {
+      "epoch": 0.63,
+      "gumbel_tau": 0.2359,
+      "image_masks_loss": 0.0081,
+      "llm_learning_rate": 1.6094912160387197e-06,
+      "loss": 0.7844,
+      "output_text_mask_loss": 0.0641,
+      "predictor_learning_rate": 0.00016094912160387196,
+      "step": 6124
+    },
+    {
+      "epoch": 0.63,
+      "gumbel_tau": 0.2358,
+      "image_masks_loss": 0.0048,
+      "llm_learning_rate": 1.6087160642394994e-06,
+      "loss": 0.6783,
+      "output_text_mask_loss": 0.0004,
+      "predictor_learning_rate": 0.00016087160642394994,
+      "step": 6125
+    },
+    {
+      "epoch": 0.63,
+      "gumbel_tau": 0.2357,
+      "image_masks_loss": 0.0203,
+      "llm_learning_rate": 1.6079410105904518e-06,
+      "loss": 0.8296,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.00016079410105904518,
+      "step": 6126
+    },
+    {
+      "epoch": 0.63,
+      "gumbel_tau": 0.2357,
+      "image_masks_loss": 0.0077,
+      "llm_learning_rate": 1.6071660551769274e-06,
+      "loss": 0.7212,
+      "output_text_mask_loss": 0.0147,
+      "predictor_learning_rate": 0.00016071660551769275,
+      "step": 6127
+    },
+    {
+      "epoch": 0.63,
+      "gumbel_tau": 0.2356,
+      "image_masks_loss": 0.0176,
+      "llm_learning_rate": 1.606391198084265e-06,
+      "loss": 0.6787,
+      "output_text_mask_loss": 0.0307,
+      "predictor_learning_rate": 0.0001606391198084265,
+      "step": 6128
+    },
+    {
+      "epoch": 0.63,
+      "gumbel_tau": 0.2356,
+      "image_masks_loss": 0.0189,
+      "llm_learning_rate": 1.6056164393977952e-06,
+      "loss": 0.7415,
+      "output_text_mask_loss": 0.0021,
+      "predictor_learning_rate": 0.00016056164393977952,
+      "step": 6129
+    },
+    {
+      "epoch": 0.63,
+      "gumbel_tau": 0.2355,
+      "image_masks_loss": 0.0036,
+      "llm_learning_rate": 1.6048417792028349e-06,
+      "loss": 0.8451,
+      "output_text_mask_loss": 0.037,
+      "predictor_learning_rate": 0.00016048417792028348,
+      "step": 6130
+    },
+    {
+      "epoch": 0.63,
+      "gumbel_tau": 0.2355,
+      "image_masks_loss": 0.0138,
+      "llm_learning_rate": 1.6040672175846916e-06,
+      "loss": 0.7814,
+      "output_text_mask_loss": 0.0353,
+      "predictor_learning_rate": 0.00016040672175846916,
+      "step": 6131
+    },
+    {
+      "epoch": 0.63,
+      "gumbel_tau": 0.2354,
+      "image_masks_loss": 0.0065,
+      "llm_learning_rate": 1.6032927546286616e-06,
+      "loss": 0.7082,
+      "output_text_mask_loss": 0.0511,
+      "predictor_learning_rate": 0.00016032927546286616,
+      "step": 6132
+    },
+    {
+      "epoch": 0.63,
+      "gumbel_tau": 0.2354,
+      "image_masks_loss": 0.0061,
+      "llm_learning_rate": 1.6025183904200315e-06,
+      "loss": 0.6801,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00016025183904200315,
+      "step": 6133
+    },
+    {
+      "epoch": 0.63,
+      "gumbel_tau": 0.2353,
+      "image_masks_loss": 0.0097,
+      "llm_learning_rate": 1.601744125044075e-06,
+      "loss": 0.8298,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00016017441250440752,
+      "step": 6134
+    },
+    {
+      "epoch": 0.63,
+      "gumbel_tau": 0.2352,
+      "image_masks_loss": 0.0042,
+      "llm_learning_rate": 1.6009699585860558e-06,
+      "loss": 0.7919,
+      "output_text_mask_loss": 0.0633,
+      "predictor_learning_rate": 0.00016009699585860557,
+      "step": 6135
+    },
+    {
+      "epoch": 0.63,
+      "gumbel_tau": 0.2352,
+      "image_masks_loss": 0.0107,
+      "llm_learning_rate": 1.6001958911312278e-06,
+      "loss": 0.711,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.00016001958911312277,
+      "step": 6136
+    },
+    {
+      "epoch": 0.63,
+      "gumbel_tau": 0.2351,
+      "image_masks_loss": 0.0103,
+      "llm_learning_rate": 1.5994219227648328e-06,
+      "loss": 0.7643,
+      "output_text_mask_loss": 0.0231,
+      "predictor_learning_rate": 0.00015994219227648328,
+      "step": 6137
+    },
+    {
+      "epoch": 0.63,
+      "gumbel_tau": 0.2351,
+      "image_masks_loss": 0.0058,
+      "llm_learning_rate": 1.598648053572101e-06,
+      "loss": 0.679,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.0001598648053572101,
+      "step": 6138
+    },
+    {
+      "epoch": 0.63,
+      "gumbel_tau": 0.235,
+      "image_masks_loss": 0.0025,
+      "llm_learning_rate": 1.5978742836382543e-06,
+      "loss": 0.809,
+      "output_text_mask_loss": 0.1175,
+      "predictor_learning_rate": 0.00015978742836382542,
+      "step": 6139
+    },
+    {
+      "epoch": 0.63,
+      "gumbel_tau": 0.235,
+      "image_masks_loss": 0.0076,
+      "llm_learning_rate": 1.5971006130485e-06,
+      "loss": 0.7176,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00015971006130485,
+      "step": 6140
+    },
+    {
+      "epoch": 0.63,
+      "gumbel_tau": 0.2349,
+      "image_masks_loss": 0.0159,
+      "llm_learning_rate": 1.5963270418880372e-06,
+      "loss": 0.657,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00015963270418880372,
+      "step": 6141
+    },
+    {
+      "epoch": 0.63,
+      "gumbel_tau": 0.2349,
+      "image_masks_loss": 0.0063,
+      "llm_learning_rate": 1.5955535702420536e-06,
+      "loss": 0.7947,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.00015955535702420534,
+      "step": 6142
+    },
+    {
+      "epoch": 0.63,
+      "gumbel_tau": 0.2348,
+      "image_masks_loss": 0.0035,
+      "llm_learning_rate": 1.5947801981957255e-06,
+      "loss": 0.7275,
+      "output_text_mask_loss": 0.0584,
+      "predictor_learning_rate": 0.00015947801981957256,
+      "step": 6143
+    },
+    {
+      "epoch": 0.63,
+      "gumbel_tau": 0.2347,
+      "image_masks_loss": 0.0244,
+      "llm_learning_rate": 1.5940069258342177e-06,
+      "loss": 0.7411,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.00015940069258342176,
+      "step": 6144
+    },
+    {
+      "epoch": 0.63,
+      "gumbel_tau": 0.2347,
+      "image_masks_loss": 0.0204,
+      "llm_learning_rate": 1.593233753242686e-06,
+      "loss": 0.716,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.00015932337532426858,
+      "step": 6145
+    },
+    {
+      "epoch": 0.63,
+      "gumbel_tau": 0.2346,
+      "image_masks_loss": 0.02,
+      "llm_learning_rate": 1.5924606805062727e-06,
+      "loss": 0.8063,
+      "output_text_mask_loss": 0.0347,
+      "predictor_learning_rate": 0.00015924606805062727,
+      "step": 6146
+    },
+    {
+      "epoch": 0.63,
+      "gumbel_tau": 0.2346,
+      "image_masks_loss": 0.0067,
+      "llm_learning_rate": 1.5916877077101113e-06,
+      "loss": 0.7333,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.0001591687707710111,
+      "step": 6147
+    },
+    {
+      "epoch": 0.63,
+      "gumbel_tau": 0.2345,
+      "image_masks_loss": 0.0124,
+      "llm_learning_rate": 1.5909148349393215e-06,
+      "loss": 0.8012,
+      "output_text_mask_loss": 0.0208,
+      "predictor_learning_rate": 0.00015909148349393216,
+      "step": 6148
+    },
+    {
+      "epoch": 0.63,
+      "gumbel_tau": 0.2345,
+      "image_masks_loss": 0.028,
+      "llm_learning_rate": 1.5901420622790162e-06,
+      "loss": 0.6986,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.0001590142062279016,
+      "step": 6149
+    },
+    {
+      "epoch": 0.63,
+      "gumbel_tau": 0.2344,
+      "image_masks_loss": 0.0111,
+      "llm_learning_rate": 1.5893693898142932e-06,
+      "loss": 0.6906,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0001589369389814293,
+      "step": 6150
+    },
+    {
+      "epoch": 0.63,
+      "gumbel_tau": 0.2344,
+      "image_masks_loss": 0.0097,
+      "llm_learning_rate": 1.5885968176302411e-06,
+      "loss": 0.6931,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0001588596817630241,
+      "step": 6151
+    },
+    {
+      "epoch": 0.63,
+      "gumbel_tau": 0.2343,
+      "image_masks_loss": 0.0035,
+      "llm_learning_rate": 1.587824345811938e-06,
+      "loss": 0.7632,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00015878243458119378,
+      "step": 6152
+    },
+    {
+      "epoch": 0.63,
+      "gumbel_tau": 0.2342,
+      "image_masks_loss": 0.0196,
+      "llm_learning_rate": 1.5870519744444495e-06,
+      "loss": 0.7295,
+      "output_text_mask_loss": 0.0062,
+      "predictor_learning_rate": 0.00015870519744444494,
+      "step": 6153
+    },
+    {
+      "epoch": 0.63,
+      "gumbel_tau": 0.2342,
+      "image_masks_loss": 0.0178,
+      "llm_learning_rate": 1.586279703612831e-06,
+      "loss": 0.6888,
+      "output_text_mask_loss": 0.0137,
+      "predictor_learning_rate": 0.0001586279703612831,
+      "step": 6154
+    },
+    {
+      "epoch": 0.63,
+      "gumbel_tau": 0.2341,
+      "image_masks_loss": 0.0112,
+      "llm_learning_rate": 1.5855075334021266e-06,
+      "loss": 0.746,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.00015855075334021264,
+      "step": 6155
+    },
+    {
+      "epoch": 0.63,
+      "gumbel_tau": 0.2341,
+      "image_masks_loss": 0.008,
+      "llm_learning_rate": 1.58473546389737e-06,
+      "loss": 0.813,
+      "output_text_mask_loss": 0.0309,
+      "predictor_learning_rate": 0.000158473546389737,
+      "step": 6156
+    },
+    {
+      "epoch": 0.63,
+      "gumbel_tau": 0.234,
+      "image_masks_loss": 0.0049,
+      "llm_learning_rate": 1.583963495183583e-06,
+      "loss": 0.7064,
+      "output_text_mask_loss": 0.0233,
+      "predictor_learning_rate": 0.00015839634951835829,
+      "step": 6157
+    },
+    {
+      "epoch": 0.63,
+      "gumbel_tau": 0.234,
+      "image_masks_loss": 0.0088,
+      "llm_learning_rate": 1.5831916273457762e-06,
+      "loss": 0.7534,
+      "output_text_mask_loss": 0.0607,
+      "predictor_learning_rate": 0.0001583191627345776,
+      "step": 6158
+    },
+    {
+      "epoch": 0.63,
+      "gumbel_tau": 0.2339,
+      "image_masks_loss": 0.006,
+      "llm_learning_rate": 1.5824198604689495e-06,
+      "loss": 0.7486,
+      "output_text_mask_loss": 0.0292,
+      "predictor_learning_rate": 0.00015824198604689492,
+      "step": 6159
+    },
+    {
+      "epoch": 0.63,
+      "gumbel_tau": 0.2339,
+      "image_masks_loss": 0.0206,
+      "llm_learning_rate": 1.5816481946380908e-06,
+      "loss": 0.844,
+      "output_text_mask_loss": 0.0124,
+      "predictor_learning_rate": 0.00015816481946380906,
+      "step": 6160
+    },
+    {
+      "epoch": 0.63,
+      "gumbel_tau": 0.2338,
+      "image_masks_loss": 0.0037,
+      "llm_learning_rate": 1.580876629938179e-06,
+      "loss": 0.6809,
+      "output_text_mask_loss": 0.0041,
+      "predictor_learning_rate": 0.0001580876629938179,
+      "step": 6161
+    },
+    {
+      "epoch": 0.63,
+      "gumbel_tau": 0.2337,
+      "image_masks_loss": 0.0086,
+      "llm_learning_rate": 1.58010516645418e-06,
+      "loss": 0.7847,
+      "output_text_mask_loss": 0.1106,
+      "predictor_learning_rate": 0.000158010516645418,
+      "step": 6162
+    },
+    {
+      "epoch": 0.63,
+      "gumbel_tau": 0.2337,
+      "image_masks_loss": 0.0088,
+      "llm_learning_rate": 1.5793338042710494e-06,
+      "loss": 0.7449,
+      "output_text_mask_loss": 0.0141,
+      "predictor_learning_rate": 0.00015793338042710493,
+      "step": 6163
+    },
+    {
+      "epoch": 0.63,
+      "gumbel_tau": 0.2336,
+      "image_masks_loss": 0.0047,
+      "llm_learning_rate": 1.57856254347373e-06,
+      "loss": 0.6319,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 0.000157856254347373,
+      "step": 6164
+    },
+    {
+      "epoch": 0.63,
+      "gumbel_tau": 0.2336,
+      "image_masks_loss": 0.0022,
+      "llm_learning_rate": 1.5777913841471566e-06,
+      "loss": 0.7359,
+      "output_text_mask_loss": 0.0018,
+      "predictor_learning_rate": 0.00015777913841471565,
+      "step": 6165
+    },
+    {
+      "epoch": 0.63,
+      "gumbel_tau": 0.2335,
+      "image_masks_loss": 0.0202,
+      "llm_learning_rate": 1.5770203263762499e-06,
+      "loss": 0.7452,
+      "output_text_mask_loss": 0.0345,
+      "predictor_learning_rate": 0.00015770203263762499,
+      "step": 6166
+    },
+    {
+      "epoch": 0.63,
+      "gumbel_tau": 0.2335,
+      "image_masks_loss": 0.0081,
+      "llm_learning_rate": 1.5762493702459203e-06,
+      "loss": 0.7293,
+      "output_text_mask_loss": 0.0256,
+      "predictor_learning_rate": 0.000157624937024592,
+      "step": 6167
+    },
+    {
+      "epoch": 0.63,
+      "gumbel_tau": 0.2334,
+      "image_masks_loss": 0.0067,
+      "llm_learning_rate": 1.575478515841068e-06,
+      "loss": 0.7581,
+      "output_text_mask_loss": 0.0216,
+      "predictor_learning_rate": 0.0001575478515841068,
+      "step": 6168
+    },
+    {
+      "epoch": 0.63,
+      "gumbel_tau": 0.2334,
+      "image_masks_loss": 0.0113,
+      "llm_learning_rate": 1.5747077632465806e-06,
+      "loss": 0.8151,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00015747077632465806,
+      "step": 6169
+    },
+    {
+      "epoch": 0.63,
+      "gumbel_tau": 0.2333,
+      "image_masks_loss": 0.0094,
+      "llm_learning_rate": 1.5739371125473358e-06,
+      "loss": 0.8298,
+      "output_text_mask_loss": 0.0029,
+      "predictor_learning_rate": 0.00015739371125473355,
+      "step": 6170
+    },
+    {
+      "epoch": 0.63,
+      "gumbel_tau": 0.2333,
+      "image_masks_loss": 0.0134,
+      "llm_learning_rate": 1.5731665638281985e-06,
+      "loss": 0.7877,
+      "output_text_mask_loss": 0.0252,
+      "predictor_learning_rate": 0.00015731665638281982,
+      "step": 6171
+    },
+    {
+      "epoch": 0.63,
+      "gumbel_tau": 0.2332,
+      "image_masks_loss": 0.0104,
+      "llm_learning_rate": 1.5723961171740237e-06,
+      "loss": 0.7609,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00015723961171740236,
+      "step": 6172
+    },
+    {
+      "epoch": 0.63,
+      "gumbel_tau": 0.2331,
+      "image_masks_loss": 0.0147,
+      "llm_learning_rate": 1.571625772669655e-06,
+      "loss": 0.6766,
+      "output_text_mask_loss": 0.0023,
+      "predictor_learning_rate": 0.00015716257726696548,
+      "step": 6173
+    },
+    {
+      "epoch": 0.63,
+      "gumbel_tau": 0.2331,
+      "image_masks_loss": 0.0225,
+      "llm_learning_rate": 1.5708555303999231e-06,
+      "loss": 0.665,
+      "output_text_mask_loss": 0.0012,
+      "predictor_learning_rate": 0.0001570855530399923,
+      "step": 6174
+    },
+    {
+      "epoch": 0.63,
+      "gumbel_tau": 0.233,
+      "image_masks_loss": 0.0113,
+      "llm_learning_rate": 1.5700853904496506e-06,
+      "loss": 0.7727,
+      "output_text_mask_loss": 0.0082,
+      "predictor_learning_rate": 0.00015700853904496503,
+      "step": 6175
+    },
+    {
+      "epoch": 0.63,
+      "gumbel_tau": 0.233,
+      "image_masks_loss": 0.0099,
+      "llm_learning_rate": 1.5693153529036467e-06,
+      "loss": 0.6724,
+      "output_text_mask_loss": 0.0107,
+      "predictor_learning_rate": 0.00015693153529036465,
+      "step": 6176
+    },
+    {
+      "epoch": 0.63,
+      "gumbel_tau": 0.2329,
+      "image_masks_loss": 0.0104,
+      "llm_learning_rate": 1.5685454178467085e-06,
+      "loss": 0.7626,
+      "output_text_mask_loss": 0.0095,
+      "predictor_learning_rate": 0.00015685454178467084,
+      "step": 6177
+    },
+    {
+      "epoch": 0.63,
+      "gumbel_tau": 0.2329,
+      "image_masks_loss": 0.0063,
+      "llm_learning_rate": 1.567775585363624e-06,
+      "loss": 0.779,
+      "output_text_mask_loss": 0.0458,
+      "predictor_learning_rate": 0.00015677755853636238,
+      "step": 6178
+    },
+    {
+      "epoch": 0.63,
+      "gumbel_tau": 0.2328,
+      "image_masks_loss": 0.0073,
+      "llm_learning_rate": 1.5670058555391682e-06,
+      "loss": 0.7949,
+      "output_text_mask_loss": 0.0603,
+      "predictor_learning_rate": 0.0001567005855539168,
+      "step": 6179
+    },
+    {
+      "epoch": 0.63,
+      "gumbel_tau": 0.2328,
+      "image_masks_loss": 0.0137,
+      "llm_learning_rate": 1.5662362284581051e-06,
+      "loss": 0.7882,
+      "output_text_mask_loss": 0.0034,
+      "predictor_learning_rate": 0.00015662362284581048,
+      "step": 6180
+    },
+    {
+      "epoch": 0.63,
+      "gumbel_tau": 0.2327,
+      "image_masks_loss": 0.0155,
+      "llm_learning_rate": 1.565466704205189e-06,
+      "loss": 0.6871,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.0001565466704205189,
+      "step": 6181
+    },
+    {
+      "epoch": 0.63,
+      "gumbel_tau": 0.2326,
+      "image_masks_loss": 0.0082,
+      "llm_learning_rate": 1.5646972828651613e-06,
+      "loss": 0.7526,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00015646972828651612,
+      "step": 6182
+    },
+    {
+      "epoch": 0.63,
+      "gumbel_tau": 0.2326,
+      "image_masks_loss": 0.0078,
+      "llm_learning_rate": 1.5639279645227517e-06,
+      "loss": 0.6608,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.00015639279645227516,
+      "step": 6183
+    },
+    {
+      "epoch": 0.63,
+      "gumbel_tau": 0.2325,
+      "image_masks_loss": 0.0133,
+      "llm_learning_rate": 1.5631587492626796e-06,
+      "loss": 0.7357,
+      "output_text_mask_loss": 0.0219,
+      "predictor_learning_rate": 0.00015631587492626796,
+      "step": 6184
+    },
+    {
+      "epoch": 0.63,
+      "gumbel_tau": 0.2325,
+      "image_masks_loss": 0.0087,
+      "llm_learning_rate": 1.562389637169653e-06,
+      "loss": 0.7438,
+      "output_text_mask_loss": 0.0064,
+      "predictor_learning_rate": 0.00015623896371696528,
+      "step": 6185
+    },
+    {
+      "epoch": 0.63,
+      "gumbel_tau": 0.2324,
+      "image_masks_loss": 0.0083,
+      "llm_learning_rate": 1.561620628328368e-06,
+      "loss": 0.7336,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.00015616206283283678,
+      "step": 6186
+    },
+    {
+      "epoch": 0.63,
+      "gumbel_tau": 0.2324,
+      "image_masks_loss": 0.0101,
+      "llm_learning_rate": 1.5608517228235093e-06,
+      "loss": 0.7371,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0001560851722823509,
+      "step": 6187
+    },
+    {
+      "epoch": 0.63,
+      "gumbel_tau": 0.2323,
+      "image_masks_loss": 0.0106,
+      "llm_learning_rate": 1.560082920739751e-06,
+      "loss": 0.8015,
+      "output_text_mask_loss": 0.0332,
+      "predictor_learning_rate": 0.00015600829207397507,
+      "step": 6188
+    },
+    {
+      "epoch": 0.63,
+      "gumbel_tau": 0.2323,
+      "image_masks_loss": 0.0292,
+      "llm_learning_rate": 1.5593142221617553e-06,
+      "loss": 0.8048,
+      "output_text_mask_loss": 0.0004,
+      "predictor_learning_rate": 0.0001559314222161755,
+      "step": 6189
+    },
+    {
+      "epoch": 0.63,
+      "gumbel_tau": 0.2322,
+      "image_masks_loss": 0.0105,
+      "llm_learning_rate": 1.5585456271741722e-06,
+      "loss": 0.6454,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0001558545627174172,
+      "step": 6190
+    },
+    {
+      "epoch": 0.63,
+      "gumbel_tau": 0.2322,
+      "image_masks_loss": 0.0094,
+      "llm_learning_rate": 1.5577771358616416e-06,
+      "loss": 0.7368,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.00015577771358616415,
+      "step": 6191
+    },
+    {
+      "epoch": 0.63,
+      "gumbel_tau": 0.2321,
+      "image_masks_loss": 0.0065,
+      "llm_learning_rate": 1.557008748308792e-06,
+      "loss": 0.7325,
+      "output_text_mask_loss": 0.0132,
+      "predictor_learning_rate": 0.0001557008748308792,
+      "step": 6192
+    },
+    {
+      "epoch": 0.63,
+      "gumbel_tau": 0.232,
+      "image_masks_loss": 0.0143,
+      "llm_learning_rate": 1.5562404646002394e-06,
+      "loss": 0.7257,
+      "output_text_mask_loss": 0.0042,
+      "predictor_learning_rate": 0.00015562404646002394,
+      "step": 6193
+    },
+    {
+      "epoch": 0.63,
+      "gumbel_tau": 0.232,
+      "image_masks_loss": 0.0168,
+      "llm_learning_rate": 1.5554722848205894e-06,
+      "loss": 0.6917,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.00015554722848205892,
+      "step": 6194
+    },
+    {
+      "epoch": 0.63,
+      "gumbel_tau": 0.2319,
+      "image_masks_loss": 0.0172,
+      "llm_learning_rate": 1.5547042090544348e-06,
+      "loss": 0.7173,
+      "output_text_mask_loss": 0.0072,
+      "predictor_learning_rate": 0.00015547042090544347,
+      "step": 6195
+    },
+    {
+      "epoch": 0.63,
+      "gumbel_tau": 0.2319,
+      "image_masks_loss": 0.0051,
+      "llm_learning_rate": 1.5539362373863579e-06,
+      "loss": 0.7764,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.00015539362373863578,
+      "step": 6196
+    },
+    {
+      "epoch": 0.63,
+      "gumbel_tau": 0.2318,
+      "image_masks_loss": 0.0145,
+      "llm_learning_rate": 1.55316836990093e-06,
+      "loss": 0.7945,
+      "output_text_mask_loss": 0.006,
+      "predictor_learning_rate": 0.000155316836990093,
+      "step": 6197
+    },
+    {
+      "epoch": 0.64,
+      "gumbel_tau": 0.2318,
+      "image_masks_loss": 0.0066,
+      "llm_learning_rate": 1.5524006066827105e-06,
+      "loss": 0.7138,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.00015524006066827104,
+      "step": 6198
+    },
+    {
+      "epoch": 0.64,
+      "gumbel_tau": 0.2317,
+      "image_masks_loss": 0.012,
+      "llm_learning_rate": 1.5516329478162468e-06,
+      "loss": 0.7378,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 0.00015516329478162467,
+      "step": 6199
+    },
+    {
+      "epoch": 0.64,
+      "gumbel_tau": 0.2317,
+      "image_masks_loss": 0.013,
+      "llm_learning_rate": 1.5508653933860748e-06,
+      "loss": 0.8142,
+      "output_text_mask_loss": 0.0076,
+      "predictor_learning_rate": 0.00015508653933860747,
+      "step": 6200
+    },
+    {
+      "epoch": 0.64,
+      "gumbel_tau": 0.2316,
+      "image_masks_loss": 0.0134,
+      "llm_learning_rate": 1.5500979434767201e-06,
+      "loss": 0.7344,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.000155009794347672,
+      "step": 6201
+    },
+    {
+      "epoch": 0.64,
+      "gumbel_tau": 0.2316,
+      "image_masks_loss": 0.0084,
+      "llm_learning_rate": 1.5493305981726958e-06,
+      "loss": 0.74,
+      "output_text_mask_loss": 0.0058,
+      "predictor_learning_rate": 0.00015493305981726958,
+      "step": 6202
+    },
+    {
+      "epoch": 0.64,
+      "gumbel_tau": 0.2315,
+      "image_masks_loss": 0.0108,
+      "llm_learning_rate": 1.5485633575585035e-06,
+      "loss": 0.7424,
+      "output_text_mask_loss": 0.0058,
+      "predictor_learning_rate": 0.00015485633575585035,
+      "step": 6203
+    },
+    {
+      "epoch": 0.64,
+      "gumbel_tau": 0.2314,
+      "image_masks_loss": 0.0104,
+      "llm_learning_rate": 1.5477962217186335e-06,
+      "loss": 0.7933,
+      "output_text_mask_loss": 0.0071,
+      "predictor_learning_rate": 0.00015477962217186334,
+      "step": 6204
+    },
+    {
+      "epoch": 0.64,
+      "gumbel_tau": 0.2314,
+      "image_masks_loss": 0.0126,
+      "llm_learning_rate": 1.5470291907375645e-06,
+      "loss": 0.7551,
+      "output_text_mask_loss": 0.0058,
+      "predictor_learning_rate": 0.00015470291907375644,
+      "step": 6205
+    },
+    {
+      "epoch": 0.64,
+      "gumbel_tau": 0.2313,
+      "image_masks_loss": 0.0107,
+      "llm_learning_rate": 1.5462622646997638e-06,
+      "loss": 0.7564,
+      "output_text_mask_loss": 0.0141,
+      "predictor_learning_rate": 0.00015462622646997636,
+      "step": 6206
+    },
+    {
+      "epoch": 0.64,
+      "gumbel_tau": 0.2313,
+      "image_masks_loss": 0.006,
+      "llm_learning_rate": 1.5454954436896864e-06,
+      "loss": 0.7982,
+      "output_text_mask_loss": 0.0038,
+      "predictor_learning_rate": 0.00015454954436896863,
+      "step": 6207
+    },
+    {
+      "epoch": 0.64,
+      "gumbel_tau": 0.2312,
+      "image_masks_loss": 0.0089,
+      "llm_learning_rate": 1.5447287277917775e-06,
+      "loss": 0.7864,
+      "output_text_mask_loss": 0.0269,
+      "predictor_learning_rate": 0.00015447287277917774,
+      "step": 6208
+    },
+    {
+      "epoch": 0.64,
+      "gumbel_tau": 0.2312,
+      "image_masks_loss": 0.0159,
+      "llm_learning_rate": 1.5439621170904687e-06,
+      "loss": 0.7573,
+      "output_text_mask_loss": 0.0025,
+      "predictor_learning_rate": 0.00015439621170904686,
+      "step": 6209
+    },
+    {
+      "epoch": 0.64,
+      "gumbel_tau": 0.2311,
+      "image_masks_loss": 0.0093,
+      "llm_learning_rate": 1.5431956116701805e-06,
+      "loss": 0.7464,
+      "output_text_mask_loss": 0.0549,
+      "predictor_learning_rate": 0.00015431956116701803,
+      "step": 6210
+    },
+    {
+      "epoch": 0.64,
+      "gumbel_tau": 0.2311,
+      "image_masks_loss": 0.0079,
+      "llm_learning_rate": 1.5424292116153233e-06,
+      "loss": 0.7596,
+      "output_text_mask_loss": 0.0071,
+      "predictor_learning_rate": 0.00015424292116153232,
+      "step": 6211
+    },
+    {
+      "epoch": 0.64,
+      "gumbel_tau": 0.231,
+      "image_masks_loss": 0.0098,
+      "llm_learning_rate": 1.5416629170102943e-06,
+      "loss": 0.8197,
+      "output_text_mask_loss": 0.0275,
+      "predictor_learning_rate": 0.00015416629170102942,
+      "step": 6212
+    },
+    {
+      "epoch": 0.64,
+      "gumbel_tau": 0.231,
+      "image_masks_loss": 0.0085,
+      "llm_learning_rate": 1.5408967279394798e-06,
+      "loss": 0.8285,
+      "output_text_mask_loss": 0.0311,
+      "predictor_learning_rate": 0.000154089672793948,
+      "step": 6213
+    },
+    {
+      "epoch": 0.64,
+      "gumbel_tau": 0.2309,
+      "image_masks_loss": 0.0106,
+      "llm_learning_rate": 1.5401306444872534e-06,
+      "loss": 0.7199,
+      "output_text_mask_loss": 0.0496,
+      "predictor_learning_rate": 0.00015401306444872535,
+      "step": 6214
+    },
+    {
+      "epoch": 0.64,
+      "gumbel_tau": 0.2308,
+      "image_masks_loss": 0.0074,
+      "llm_learning_rate": 1.5393646667379785e-06,
+      "loss": 0.7928,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.00015393646667379786,
+      "step": 6215
+    },
+    {
+      "epoch": 0.64,
+      "gumbel_tau": 0.2308,
+      "image_masks_loss": 0.0079,
+      "llm_learning_rate": 1.538598794776006e-06,
+      "loss": 0.753,
+      "output_text_mask_loss": 0.0015,
+      "predictor_learning_rate": 0.00015385987947760061,
+      "step": 6216
+    },
+    {
+      "epoch": 0.64,
+      "gumbel_tau": 0.2307,
+      "image_masks_loss": 0.0175,
+      "llm_learning_rate": 1.537833028685676e-06,
+      "loss": 0.7347,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.0001537833028685676,
+      "step": 6217
+    },
+    {
+      "epoch": 0.64,
+      "gumbel_tau": 0.2307,
+      "image_masks_loss": 0.01,
+      "llm_learning_rate": 1.537067368551316e-06,
+      "loss": 0.7294,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.00015370673685513158,
+      "step": 6218
+    },
+    {
+      "epoch": 0.64,
+      "gumbel_tau": 0.2306,
+      "image_masks_loss": 0.0176,
+      "llm_learning_rate": 1.5363018144572418e-06,
+      "loss": 0.7829,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 0.00015363018144572417,
+      "step": 6219
+    },
+    {
+      "epoch": 0.64,
+      "gumbel_tau": 0.2306,
+      "image_masks_loss": 0.0124,
+      "llm_learning_rate": 1.5355363664877588e-06,
+      "loss": 0.7559,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00015355363664877586,
+      "step": 6220
+    },
+    {
+      "epoch": 0.64,
+      "gumbel_tau": 0.2305,
+      "image_masks_loss": 0.0087,
+      "llm_learning_rate": 1.5347710247271592e-06,
+      "loss": 0.7837,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.0001534771024727159,
+      "step": 6221
+    },
+    {
+      "epoch": 0.64,
+      "gumbel_tau": 0.2305,
+      "image_masks_loss": 0.0039,
+      "llm_learning_rate": 1.5340057892597246e-06,
+      "loss": 0.6879,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.00015340057892597244,
+      "step": 6222
+    },
+    {
+      "epoch": 0.64,
+      "gumbel_tau": 0.2304,
+      "image_masks_loss": 0.0011,
+      "llm_learning_rate": 1.5332406601697237e-06,
+      "loss": 0.771,
+      "output_text_mask_loss": 0.0081,
+      "predictor_learning_rate": 0.00015332406601697236,
+      "step": 6223
+    },
+    {
+      "epoch": 0.64,
+      "gumbel_tau": 0.2304,
+      "image_masks_loss": 0.0113,
+      "llm_learning_rate": 1.5324756375414145e-06,
+      "loss": 0.7389,
+      "output_text_mask_loss": 0.0042,
+      "predictor_learning_rate": 0.00015324756375414145,
+      "step": 6224
+    },
+    {
+      "epoch": 0.64,
+      "gumbel_tau": 0.2303,
+      "image_masks_loss": 0.0095,
+      "llm_learning_rate": 1.5317107214590437e-06,
+      "loss": 0.7718,
+      "output_text_mask_loss": 0.0139,
+      "predictor_learning_rate": 0.00015317107214590436,
+      "step": 6225
+    },
+    {
+      "epoch": 0.64,
+      "gumbel_tau": 0.2302,
+      "image_masks_loss": 0.0025,
+      "llm_learning_rate": 1.5309459120068448e-06,
+      "loss": 0.6998,
+      "output_text_mask_loss": 0.0111,
+      "predictor_learning_rate": 0.00015309459120068449,
+      "step": 6226
+    },
+    {
+      "epoch": 0.64,
+      "gumbel_tau": 0.2302,
+      "image_masks_loss": 0.0028,
+      "llm_learning_rate": 1.5301812092690406e-06,
+      "loss": 0.7733,
+      "output_text_mask_loss": 0.0039,
+      "predictor_learning_rate": 0.00015301812092690406,
+      "step": 6227
+    },
+    {
+      "epoch": 0.64,
+      "gumbel_tau": 0.2301,
+      "image_masks_loss": 0.0175,
+      "llm_learning_rate": 1.5294166133298422e-06,
+      "loss": 0.6844,
+      "output_text_mask_loss": 0.0519,
+      "predictor_learning_rate": 0.0001529416613329842,
+      "step": 6228
+    },
+    {
+      "epoch": 0.64,
+      "gumbel_tau": 0.2301,
+      "image_masks_loss": 0.0162,
+      "llm_learning_rate": 1.5286521242734481e-06,
+      "loss": 0.7501,
+      "output_text_mask_loss": 0.0025,
+      "predictor_learning_rate": 0.0001528652124273448,
+      "step": 6229
+    },
+    {
+      "epoch": 0.64,
+      "gumbel_tau": 0.23,
+      "image_masks_loss": 0.0161,
+      "llm_learning_rate": 1.5278877421840448e-06,
+      "loss": 0.8111,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00015278877421840448,
+      "step": 6230
+    },
+    {
+      "epoch": 0.64,
+      "gumbel_tau": 0.23,
+      "image_masks_loss": 0.0082,
+      "llm_learning_rate": 1.5271234671458106e-06,
+      "loss": 0.736,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00015271234671458106,
+      "step": 6231
+    },
+    {
+      "epoch": 0.64,
+      "gumbel_tau": 0.2299,
+      "image_masks_loss": 0.013,
+      "llm_learning_rate": 1.5263592992429055e-06,
+      "loss": 0.7626,
+      "output_text_mask_loss": 0.0077,
+      "predictor_learning_rate": 0.00015263592992429054,
+      "step": 6232
+    },
+    {
+      "epoch": 0.64,
+      "gumbel_tau": 0.2299,
+      "image_masks_loss": 0.0104,
+      "llm_learning_rate": 1.5255952385594836e-06,
+      "loss": 0.75,
+      "output_text_mask_loss": 0.0042,
+      "predictor_learning_rate": 0.00015255952385594834,
+      "step": 6233
+    },
+    {
+      "epoch": 0.64,
+      "gumbel_tau": 0.2298,
+      "image_masks_loss": 0.0074,
+      "llm_learning_rate": 1.5248312851796844e-06,
+      "loss": 0.7454,
+      "output_text_mask_loss": 0.0362,
+      "predictor_learning_rate": 0.00015248312851796843,
+      "step": 6234
+    },
+    {
+      "epoch": 0.64,
+      "gumbel_tau": 0.2298,
+      "image_masks_loss": 0.0069,
+      "llm_learning_rate": 1.5240674391876357e-06,
+      "loss": 0.752,
+      "output_text_mask_loss": 0.0177,
+      "predictor_learning_rate": 0.00015240674391876356,
+      "step": 6235
+    },
+    {
+      "epoch": 0.64,
+      "gumbel_tau": 0.2297,
+      "image_masks_loss": 0.0138,
+      "llm_learning_rate": 1.5233037006674547e-06,
+      "loss": 0.7497,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.00015233037006674548,
+      "step": 6236
+    },
+    {
+      "epoch": 0.64,
+      "gumbel_tau": 0.2296,
+      "image_masks_loss": 0.0113,
+      "llm_learning_rate": 1.5225400697032462e-06,
+      "loss": 0.5912,
+      "output_text_mask_loss": 0.0033,
+      "predictor_learning_rate": 0.0001522540069703246,
+      "step": 6237
+    },
+    {
+      "epoch": 0.64,
+      "gumbel_tau": 0.2296,
+      "image_masks_loss": 0.0065,
+      "llm_learning_rate": 1.5217765463791017e-06,
+      "loss": 0.8032,
+      "output_text_mask_loss": 0.0572,
+      "predictor_learning_rate": 0.00015217765463791018,
+      "step": 6238
+    },
+    {
+      "epoch": 0.64,
+      "gumbel_tau": 0.2295,
+      "image_masks_loss": 0.0038,
+      "llm_learning_rate": 1.521013130779103e-06,
+      "loss": 0.6764,
+      "output_text_mask_loss": 0.0277,
+      "predictor_learning_rate": 0.00015210131307791028,
+      "step": 6239
+    },
+    {
+      "epoch": 0.64,
+      "gumbel_tau": 0.2295,
+      "image_masks_loss": 0.0063,
+      "llm_learning_rate": 1.5202498229873191e-06,
+      "loss": 0.7325,
+      "output_text_mask_loss": 0.0038,
+      "predictor_learning_rate": 0.00015202498229873192,
+      "step": 6240
+    },
+    {
+      "epoch": 0.64,
+      "gumbel_tau": 0.2294,
+      "image_masks_loss": 0.0111,
+      "llm_learning_rate": 1.519486623087807e-06,
+      "loss": 0.7418,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00015194866230878068,
+      "step": 6241
+    },
+    {
+      "epoch": 0.64,
+      "gumbel_tau": 0.2294,
+      "image_masks_loss": 0.0149,
+      "llm_learning_rate": 1.5187235311646114e-06,
+      "loss": 0.805,
+      "output_text_mask_loss": 0.0275,
+      "predictor_learning_rate": 0.00015187235311646115,
+      "step": 6242
+    },
+    {
+      "epoch": 0.64,
+      "gumbel_tau": 0.2293,
+      "image_masks_loss": 0.0144,
+      "llm_learning_rate": 1.517960547301767e-06,
+      "loss": 0.764,
+      "output_text_mask_loss": 0.0099,
+      "predictor_learning_rate": 0.00015179605473017668,
+      "step": 6243
+    },
+    {
+      "epoch": 0.64,
+      "gumbel_tau": 0.2293,
+      "image_masks_loss": 0.0083,
+      "llm_learning_rate": 1.5171976715832946e-06,
+      "loss": 0.7805,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.00015171976715832943,
+      "step": 6244
+    },
+    {
+      "epoch": 0.64,
+      "gumbel_tau": 0.2292,
+      "image_masks_loss": 0.0112,
+      "llm_learning_rate": 1.5164349040932037e-06,
+      "loss": 0.6916,
+      "output_text_mask_loss": 0.0119,
+      "predictor_learning_rate": 0.00015164349040932036,
+      "step": 6245
+    },
+    {
+      "epoch": 0.64,
+      "gumbel_tau": 0.2292,
+      "image_masks_loss": 0.0109,
+      "llm_learning_rate": 1.5156722449154908e-06,
+      "loss": 0.7125,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 0.0001515672244915491,
+      "step": 6246
+    },
+    {
+      "epoch": 0.64,
+      "gumbel_tau": 0.2291,
+      "image_masks_loss": 0.0037,
+      "llm_learning_rate": 1.514909694134144e-06,
+      "loss": 0.7313,
+      "output_text_mask_loss": 0.0198,
+      "predictor_learning_rate": 0.0001514909694134144,
+      "step": 6247
+    },
+    {
+      "epoch": 0.64,
+      "gumbel_tau": 0.2291,
+      "image_masks_loss": 0.0186,
+      "llm_learning_rate": 1.5141472518331357e-06,
+      "loss": 0.8661,
+      "output_text_mask_loss": 0.0185,
+      "predictor_learning_rate": 0.00015141472518331357,
+      "step": 6248
+    },
+    {
+      "epoch": 0.64,
+      "gumbel_tau": 0.229,
+      "image_masks_loss": 0.0082,
+      "llm_learning_rate": 1.5133849180964274e-06,
+      "loss": 0.7314,
+      "output_text_mask_loss": 0.005,
+      "predictor_learning_rate": 0.00015133849180964273,
+      "step": 6249
+    },
+    {
+      "epoch": 0.64,
+      "gumbel_tau": 0.2289,
+      "image_masks_loss": 0.0077,
+      "llm_learning_rate": 1.5126226930079712e-06,
+      "loss": 0.7633,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00015126226930079712,
+      "step": 6250
+    },
+    {
+      "epoch": 0.64,
+      "gumbel_tau": 0.2289,
+      "image_masks_loss": 0.0138,
+      "llm_learning_rate": 1.5118605766517014e-06,
+      "loss": 0.7735,
+      "output_text_mask_loss": 0.0134,
+      "predictor_learning_rate": 0.00015118605766517012,
+      "step": 6251
+    },
+    {
+      "epoch": 0.64,
+      "gumbel_tau": 0.2288,
+      "image_masks_loss": 0.0041,
+      "llm_learning_rate": 1.5110985691115468e-06,
+      "loss": 0.772,
+      "output_text_mask_loss": 0.0374,
+      "predictor_learning_rate": 0.00015110985691115466,
+      "step": 6252
+    },
+    {
+      "epoch": 0.64,
+      "gumbel_tau": 0.2288,
+      "image_masks_loss": 0.0139,
+      "llm_learning_rate": 1.5103366704714204e-06,
+      "loss": 0.7989,
+      "output_text_mask_loss": 0.0156,
+      "predictor_learning_rate": 0.00015103366704714204,
+      "step": 6253
+    },
+    {
+      "epoch": 0.64,
+      "gumbel_tau": 0.2287,
+      "image_masks_loss": 0.0124,
+      "llm_learning_rate": 1.5095748808152242e-06,
+      "loss": 0.7517,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 0.0001509574880815224,
+      "step": 6254
+    },
+    {
+      "epoch": 0.64,
+      "gumbel_tau": 0.2287,
+      "image_masks_loss": 0.0059,
+      "llm_learning_rate": 1.5088132002268476e-06,
+      "loss": 0.7284,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00015088132002268477,
+      "step": 6255
+    },
+    {
+      "epoch": 0.64,
+      "gumbel_tau": 0.2286,
+      "image_masks_loss": 0.0164,
+      "llm_learning_rate": 1.5080516287901698e-06,
+      "loss": 0.6577,
+      "output_text_mask_loss": 0.0086,
+      "predictor_learning_rate": 0.00015080516287901697,
+      "step": 6256
+    },
+    {
+      "epoch": 0.64,
+      "gumbel_tau": 0.2286,
+      "image_masks_loss": 0.0044,
+      "llm_learning_rate": 1.507290166589056e-06,
+      "loss": 0.6912,
+      "output_text_mask_loss": 0.0019,
+      "predictor_learning_rate": 0.0001507290166589056,
+      "step": 6257
+    },
+    {
+      "epoch": 0.64,
+      "gumbel_tau": 0.2285,
+      "image_masks_loss": 0.0109,
+      "llm_learning_rate": 1.5065288137073597e-06,
+      "loss": 0.7168,
+      "output_text_mask_loss": 0.0101,
+      "predictor_learning_rate": 0.00015065288137073595,
+      "step": 6258
+    },
+    {
+      "epoch": 0.64,
+      "gumbel_tau": 0.2285,
+      "image_masks_loss": 0.0064,
+      "llm_learning_rate": 1.5057675702289244e-06,
+      "loss": 0.7251,
+      "output_text_mask_loss": 0.1465,
+      "predictor_learning_rate": 0.00015057675702289243,
+      "step": 6259
+    },
+    {
+      "epoch": 0.64,
+      "gumbel_tau": 0.2284,
+      "image_masks_loss": 0.0113,
+      "llm_learning_rate": 1.505006436237578e-06,
+      "loss": 0.6911,
+      "output_text_mask_loss": 0.0263,
+      "predictor_learning_rate": 0.00015050064362375778,
+      "step": 6260
+    },
+    {
+      "epoch": 0.64,
+      "gumbel_tau": 0.2284,
+      "image_masks_loss": 0.0095,
+      "llm_learning_rate": 1.5042454118171395e-06,
+      "loss": 0.7434,
+      "output_text_mask_loss": 0.0168,
+      "predictor_learning_rate": 0.00015042454118171394,
+      "step": 6261
+    },
+    {
+      "epoch": 0.64,
+      "gumbel_tau": 0.2283,
+      "image_masks_loss": 0.0053,
+      "llm_learning_rate": 1.5034844970514135e-06,
+      "loss": 0.7419,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00015034844970514133,
+      "step": 6262
+    },
+    {
+      "epoch": 0.64,
+      "gumbel_tau": 0.2282,
+      "image_masks_loss": 0.009,
+      "llm_learning_rate": 1.502723692024195e-06,
+      "loss": 0.6661,
+      "output_text_mask_loss": 0.0343,
+      "predictor_learning_rate": 0.0001502723692024195,
+      "step": 6263
+    },
+    {
+      "epoch": 0.64,
+      "gumbel_tau": 0.2282,
+      "image_masks_loss": 0.0217,
+      "llm_learning_rate": 1.501962996819265e-06,
+      "loss": 0.8018,
+      "output_text_mask_loss": 0.0423,
+      "predictor_learning_rate": 0.0001501962996819265,
+      "step": 6264
+    },
+    {
+      "epoch": 0.64,
+      "gumbel_tau": 0.2281,
+      "image_masks_loss": 0.0148,
+      "llm_learning_rate": 1.5012024115203921e-06,
+      "loss": 0.6969,
+      "output_text_mask_loss": 0.0572,
+      "predictor_learning_rate": 0.0001501202411520392,
+      "step": 6265
+    },
+    {
+      "epoch": 0.64,
+      "gumbel_tau": 0.2281,
+      "image_masks_loss": 0.0152,
+      "llm_learning_rate": 1.5004419362113349e-06,
+      "loss": 0.7485,
+      "output_text_mask_loss": 0.0135,
+      "predictor_learning_rate": 0.00015004419362113348,
+      "step": 6266
+    },
+    {
+      "epoch": 0.64,
+      "gumbel_tau": 0.228,
+      "image_masks_loss": 0.0164,
+      "llm_learning_rate": 1.4996815709758384e-06,
+      "loss": 0.8665,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.00014996815709758383,
+      "step": 6267
+    },
+    {
+      "epoch": 0.64,
+      "gumbel_tau": 0.228,
+      "image_masks_loss": 0.0136,
+      "llm_learning_rate": 1.4989213158976356e-06,
+      "loss": 0.6856,
+      "output_text_mask_loss": 0.0231,
+      "predictor_learning_rate": 0.00014989213158976355,
+      "step": 6268
+    },
+    {
+      "epoch": 0.64,
+      "gumbel_tau": 0.2279,
+      "image_masks_loss": 0.0218,
+      "llm_learning_rate": 1.4981611710604472e-06,
+      "loss": 0.7525,
+      "output_text_mask_loss": 0.0404,
+      "predictor_learning_rate": 0.0001498161171060447,
+      "step": 6269
+    },
+    {
+      "epoch": 0.64,
+      "gumbel_tau": 0.2279,
+      "image_masks_loss": 0.0172,
+      "llm_learning_rate": 1.4974011365479818e-06,
+      "loss": 0.7623,
+      "output_text_mask_loss": 0.0355,
+      "predictor_learning_rate": 0.00014974011365479818,
+      "step": 6270
+    },
+    {
+      "epoch": 0.64,
+      "gumbel_tau": 0.2278,
+      "image_masks_loss": 0.0161,
+      "llm_learning_rate": 1.4966412124439359e-06,
+      "loss": 0.7279,
+      "output_text_mask_loss": 0.0015,
+      "predictor_learning_rate": 0.00014966412124439358,
+      "step": 6271
+    },
+    {
+      "epoch": 0.64,
+      "gumbel_tau": 0.2278,
+      "image_masks_loss": 0.006,
+      "llm_learning_rate": 1.4958813988319954e-06,
+      "loss": 0.7827,
+      "output_text_mask_loss": 0.0446,
+      "predictor_learning_rate": 0.00014958813988319954,
+      "step": 6272
+    },
+    {
+      "epoch": 0.64,
+      "gumbel_tau": 0.2277,
+      "image_masks_loss": 0.0103,
+      "llm_learning_rate": 1.4951216957958314e-06,
+      "loss": 0.771,
+      "output_text_mask_loss": 0.0399,
+      "predictor_learning_rate": 0.00014951216957958312,
+      "step": 6273
+    },
+    {
+      "epoch": 0.64,
+      "gumbel_tau": 0.2277,
+      "image_masks_loss": 0.0057,
+      "llm_learning_rate": 1.4943621034191042e-06,
+      "loss": 0.7137,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00014943621034191042,
+      "step": 6274
+    },
+    {
+      "epoch": 0.64,
+      "gumbel_tau": 0.2276,
+      "image_masks_loss": 0.0058,
+      "llm_learning_rate": 1.4936026217854626e-06,
+      "loss": 0.7312,
+      "output_text_mask_loss": 0.0162,
+      "predictor_learning_rate": 0.00014936026217854625,
+      "step": 6275
+    },
+    {
+      "epoch": 0.64,
+      "gumbel_tau": 0.2275,
+      "image_masks_loss": 0.0031,
+      "llm_learning_rate": 1.4928432509785423e-06,
+      "loss": 0.7895,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.00014928432509785422,
+      "step": 6276
+    },
+    {
+      "epoch": 0.64,
+      "gumbel_tau": 0.2275,
+      "image_masks_loss": 0.0156,
+      "llm_learning_rate": 1.4920839910819662e-06,
+      "loss": 0.8529,
+      "output_text_mask_loss": 0.0313,
+      "predictor_learning_rate": 0.00014920839910819662,
+      "step": 6277
+    },
+    {
+      "epoch": 0.64,
+      "gumbel_tau": 0.2274,
+      "image_masks_loss": 0.0041,
+      "llm_learning_rate": 1.4913248421793458e-06,
+      "loss": 0.7409,
+      "output_text_mask_loss": 0.0086,
+      "predictor_learning_rate": 0.00014913248421793458,
+      "step": 6278
+    },
+    {
+      "epoch": 0.64,
+      "gumbel_tau": 0.2274,
+      "image_masks_loss": 0.0052,
+      "llm_learning_rate": 1.4905658043542809e-06,
+      "loss": 0.8394,
+      "output_text_mask_loss": 0.0282,
+      "predictor_learning_rate": 0.00014905658043542808,
+      "step": 6279
+    },
+    {
+      "epoch": 0.64,
+      "gumbel_tau": 0.2273,
+      "image_masks_loss": 0.0054,
+      "llm_learning_rate": 1.4898068776903579e-06,
+      "loss": 0.7196,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.0001489806877690358,
+      "step": 6280
+    },
+    {
+      "epoch": 0.64,
+      "gumbel_tau": 0.2273,
+      "image_masks_loss": 0.0305,
+      "llm_learning_rate": 1.4890480622711515e-06,
+      "loss": 0.7558,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.00014890480622711515,
+      "step": 6281
+    },
+    {
+      "epoch": 0.64,
+      "gumbel_tau": 0.2272,
+      "image_masks_loss": 0.0216,
+      "llm_learning_rate": 1.4882893581802244e-06,
+      "loss": 0.7697,
+      "output_text_mask_loss": 0.0238,
+      "predictor_learning_rate": 0.00014882893581802243,
+      "step": 6282
+    },
+    {
+      "epoch": 0.64,
+      "gumbel_tau": 0.2272,
+      "image_masks_loss": 0.0063,
+      "llm_learning_rate": 1.487530765501127e-06,
+      "loss": 0.7188,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0001487530765501127,
+      "step": 6283
+    },
+    {
+      "epoch": 0.64,
+      "gumbel_tau": 0.2271,
+      "image_masks_loss": 0.0143,
+      "llm_learning_rate": 1.4867722843173965e-06,
+      "loss": 0.7988,
+      "output_text_mask_loss": 0.0137,
+      "predictor_learning_rate": 0.00014867722843173964,
+      "step": 6284
+    },
+    {
+      "epoch": 0.64,
+      "gumbel_tau": 0.2271,
+      "image_masks_loss": 0.0117,
+      "llm_learning_rate": 1.4860139147125592e-06,
+      "loss": 0.7252,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0001486013914712559,
+      "step": 6285
+    },
+    {
+      "epoch": 0.64,
+      "gumbel_tau": 0.227,
+      "image_masks_loss": 0.0101,
+      "llm_learning_rate": 1.485255656770128e-06,
+      "loss": 0.7955,
+      "output_text_mask_loss": 0.0034,
+      "predictor_learning_rate": 0.0001485255656770128,
+      "step": 6286
+    },
+    {
+      "epoch": 0.64,
+      "gumbel_tau": 0.227,
+      "image_masks_loss": 0.0068,
+      "llm_learning_rate": 1.4844975105736048e-06,
+      "loss": 0.762,
+      "output_text_mask_loss": 0.036,
+      "predictor_learning_rate": 0.00014844975105736047,
+      "step": 6287
+    },
+    {
+      "epoch": 0.64,
+      "gumbel_tau": 0.2269,
+      "image_masks_loss": 0.0093,
+      "llm_learning_rate": 1.4837394762064771e-06,
+      "loss": 0.7509,
+      "output_text_mask_loss": 0.0111,
+      "predictor_learning_rate": 0.0001483739476206477,
+      "step": 6288
+    },
+    {
+      "epoch": 0.64,
+      "gumbel_tau": 0.2268,
+      "image_masks_loss": 0.0201,
+      "llm_learning_rate": 1.482981553752222e-06,
+      "loss": 0.7711,
+      "output_text_mask_loss": 0.0217,
+      "predictor_learning_rate": 0.00014829815537522218,
+      "step": 6289
+    },
+    {
+      "epoch": 0.64,
+      "gumbel_tau": 0.2268,
+      "image_masks_loss": 0.0048,
+      "llm_learning_rate": 1.4822237432943029e-06,
+      "loss": 0.775,
+      "output_text_mask_loss": 0.0041,
+      "predictor_learning_rate": 0.0001482223743294303,
+      "step": 6290
+    },
+    {
+      "epoch": 0.64,
+      "gumbel_tau": 0.2267,
+      "image_masks_loss": 0.0195,
+      "llm_learning_rate": 1.4814660449161728e-06,
+      "loss": 0.7573,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00014814660449161726,
+      "step": 6291
+    },
+    {
+      "epoch": 0.64,
+      "gumbel_tau": 0.2267,
+      "image_masks_loss": 0.014,
+      "llm_learning_rate": 1.4807084587012705e-06,
+      "loss": 0.7798,
+      "output_text_mask_loss": 0.0111,
+      "predictor_learning_rate": 0.00014807084587012704,
+      "step": 6292
+    },
+    {
+      "epoch": 0.64,
+      "gumbel_tau": 0.2266,
+      "image_masks_loss": 0.0038,
+      "llm_learning_rate": 1.4799509847330228e-06,
+      "loss": 0.7301,
+      "output_text_mask_loss": 0.0032,
+      "predictor_learning_rate": 0.0001479950984733023,
+      "step": 6293
+    },
+    {
+      "epoch": 0.64,
+      "gumbel_tau": 0.2266,
+      "image_masks_loss": 0.0034,
+      "llm_learning_rate": 1.479193623094845e-06,
+      "loss": 0.7124,
+      "output_text_mask_loss": 0.024,
+      "predictor_learning_rate": 0.0001479193623094845,
+      "step": 6294
+    },
+    {
+      "epoch": 0.64,
+      "gumbel_tau": 0.2265,
+      "image_masks_loss": 0.0152,
+      "llm_learning_rate": 1.478436373870139e-06,
+      "loss": 0.7304,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.0001478436373870139,
+      "step": 6295
+    },
+    {
+      "epoch": 0.65,
+      "gumbel_tau": 0.2265,
+      "image_masks_loss": 0.0205,
+      "llm_learning_rate": 1.4776792371422954e-06,
+      "loss": 0.7566,
+      "output_text_mask_loss": 0.0039,
+      "predictor_learning_rate": 0.00014776792371422955,
+      "step": 6296
+    },
+    {
+      "epoch": 0.65,
+      "gumbel_tau": 0.2264,
+      "image_masks_loss": 0.0141,
+      "llm_learning_rate": 1.4769222129946906e-06,
+      "loss": 0.7318,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.00014769222129946906,
+      "step": 6297
+    },
+    {
+      "epoch": 0.65,
+      "gumbel_tau": 0.2264,
+      "image_masks_loss": 0.0077,
+      "llm_learning_rate": 1.4761653015106907e-06,
+      "loss": 0.8585,
+      "output_text_mask_loss": 0.0145,
+      "predictor_learning_rate": 0.00014761653015106909,
+      "step": 6298
+    },
+    {
+      "epoch": 0.65,
+      "gumbel_tau": 0.2263,
+      "image_masks_loss": 0.0087,
+      "llm_learning_rate": 1.475408502773649e-06,
+      "loss": 0.7314,
+      "output_text_mask_loss": 0.0549,
+      "predictor_learning_rate": 0.0001475408502773649,
+      "step": 6299
+    },
+    {
+      "epoch": 0.65,
+      "gumbel_tau": 0.2263,
+      "image_masks_loss": 0.0121,
+      "llm_learning_rate": 1.4746518168669047e-06,
+      "loss": 0.753,
+      "output_text_mask_loss": 0.0086,
+      "predictor_learning_rate": 0.00014746518168669046,
+      "step": 6300
+    },
+    {
+      "epoch": 0.65,
+      "gumbel_tau": 0.2262,
+      "image_masks_loss": 0.0118,
+      "llm_learning_rate": 1.4738952438737858e-06,
+      "loss": 0.819,
+      "output_text_mask_loss": 0.0322,
+      "predictor_learning_rate": 0.00014738952438737857,
+      "step": 6301
+    },
+    {
+      "epoch": 0.65,
+      "gumbel_tau": 0.2262,
+      "image_masks_loss": 0.0064,
+      "llm_learning_rate": 1.4731387838776088e-06,
+      "loss": 0.7432,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.00014731387838776087,
+      "step": 6302
+    },
+    {
+      "epoch": 0.65,
+      "gumbel_tau": 0.2261,
+      "image_masks_loss": 0.0161,
+      "llm_learning_rate": 1.472382436961676e-06,
+      "loss": 0.7456,
+      "output_text_mask_loss": 0.0062,
+      "predictor_learning_rate": 0.00014723824369616758,
+      "step": 6303
+    },
+    {
+      "epoch": 0.65,
+      "gumbel_tau": 0.226,
+      "image_masks_loss": 0.0181,
+      "llm_learning_rate": 1.4716262032092776e-06,
+      "loss": 0.7189,
+      "output_text_mask_loss": 0.0174,
+      "predictor_learning_rate": 0.00014716262032092776,
+      "step": 6304
+    },
+    {
+      "epoch": 0.65,
+      "gumbel_tau": 0.226,
+      "image_masks_loss": 0.0051,
+      "llm_learning_rate": 1.4708700827036943e-06,
+      "loss": 0.7054,
+      "output_text_mask_loss": 0.0136,
+      "predictor_learning_rate": 0.00014708700827036943,
+      "step": 6305
+    },
+    {
+      "epoch": 0.65,
+      "gumbel_tau": 0.2259,
+      "image_masks_loss": 0.01,
+      "llm_learning_rate": 1.4701140755281878e-06,
+      "loss": 0.642,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.00014701140755281878,
+      "step": 6306
+    },
+    {
+      "epoch": 0.65,
+      "gumbel_tau": 0.2259,
+      "image_masks_loss": 0.0103,
+      "llm_learning_rate": 1.4693581817660141e-06,
+      "loss": 0.6712,
+      "output_text_mask_loss": 0.0561,
+      "predictor_learning_rate": 0.0001469358181766014,
+      "step": 6307
+    },
+    {
+      "epoch": 0.65,
+      "gumbel_tau": 0.2258,
+      "image_masks_loss": 0.008,
+      "llm_learning_rate": 1.468602401500413e-06,
+      "loss": 0.7843,
+      "output_text_mask_loss": 0.0256,
+      "predictor_learning_rate": 0.00014686024015004128,
+      "step": 6308
+    },
+    {
+      "epoch": 0.65,
+      "gumbel_tau": 0.2258,
+      "image_masks_loss": 0.0082,
+      "llm_learning_rate": 1.4678467348146124e-06,
+      "loss": 0.8105,
+      "output_text_mask_loss": 0.0454,
+      "predictor_learning_rate": 0.00014678467348146123,
+      "step": 6309
+    },
+    {
+      "epoch": 0.65,
+      "gumbel_tau": 0.2257,
+      "image_masks_loss": 0.0057,
+      "llm_learning_rate": 1.4670911817918282e-06,
+      "loss": 0.793,
+      "output_text_mask_loss": 0.0075,
+      "predictor_learning_rate": 0.0001467091181791828,
+      "step": 6310
+    },
+    {
+      "epoch": 0.65,
+      "gumbel_tau": 0.2257,
+      "image_masks_loss": 0.0046,
+      "llm_learning_rate": 1.4663357425152644e-06,
+      "loss": 0.7137,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.00014663357425152642,
+      "step": 6311
+    },
+    {
+      "epoch": 0.65,
+      "gumbel_tau": 0.2256,
+      "image_masks_loss": 0.0068,
+      "llm_learning_rate": 1.465580417068111e-06,
+      "loss": 0.7046,
+      "output_text_mask_loss": 0.0179,
+      "predictor_learning_rate": 0.00014655804170681108,
+      "step": 6312
+    },
+    {
+      "epoch": 0.65,
+      "gumbel_tau": 0.2256,
+      "image_masks_loss": 0.0053,
+      "llm_learning_rate": 1.4648252055335454e-06,
+      "loss": 0.729,
+      "output_text_mask_loss": 0.0107,
+      "predictor_learning_rate": 0.00014648252055335454,
+      "step": 6313
+    },
+    {
+      "epoch": 0.65,
+      "gumbel_tau": 0.2255,
+      "image_masks_loss": 0.0072,
+      "llm_learning_rate": 1.4640701079947347e-06,
+      "loss": 0.6905,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.00014640701079947347,
+      "step": 6314
+    },
+    {
+      "epoch": 0.65,
+      "gumbel_tau": 0.2255,
+      "image_masks_loss": 0.0099,
+      "llm_learning_rate": 1.4633151245348308e-06,
+      "loss": 0.6818,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.00014633151245348308,
+      "step": 6315
+    },
+    {
+      "epoch": 0.65,
+      "gumbel_tau": 0.2254,
+      "image_masks_loss": 0.0075,
+      "llm_learning_rate": 1.462560255236975e-06,
+      "loss": 0.7147,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00014625602552369748,
+      "step": 6316
+    },
+    {
+      "epoch": 0.65,
+      "gumbel_tau": 0.2254,
+      "image_masks_loss": 0.0072,
+      "llm_learning_rate": 1.461805500184294e-06,
+      "loss": 0.7297,
+      "output_text_mask_loss": 0.0138,
+      "predictor_learning_rate": 0.0001461805500184294,
+      "step": 6317
+    },
+    {
+      "epoch": 0.65,
+      "gumbel_tau": 0.2253,
+      "image_masks_loss": 0.0147,
+      "llm_learning_rate": 1.4610508594599044e-06,
+      "loss": 0.8082,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00014610508594599042,
+      "step": 6318
+    },
+    {
+      "epoch": 0.65,
+      "gumbel_tau": 0.2252,
+      "image_masks_loss": 0.0088,
+      "llm_learning_rate": 1.4602963331469083e-06,
+      "loss": 0.6881,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.00014602963331469083,
+      "step": 6319
+    },
+    {
+      "epoch": 0.65,
+      "gumbel_tau": 0.2252,
+      "image_masks_loss": 0.0067,
+      "llm_learning_rate": 1.4595419213283952e-06,
+      "loss": 0.7288,
+      "output_text_mask_loss": 0.0237,
+      "predictor_learning_rate": 0.0001459541921328395,
+      "step": 6320
+    },
+    {
+      "epoch": 0.65,
+      "gumbel_tau": 0.2251,
+      "image_masks_loss": 0.0124,
+      "llm_learning_rate": 1.4587876240874438e-06,
+      "loss": 0.7741,
+      "output_text_mask_loss": 0.0425,
+      "predictor_learning_rate": 0.00014587876240874437,
+      "step": 6321
+    },
+    {
+      "epoch": 0.65,
+      "gumbel_tau": 0.2251,
+      "image_masks_loss": 0.0097,
+      "llm_learning_rate": 1.4580334415071185e-06,
+      "loss": 0.7435,
+      "output_text_mask_loss": 0.0128,
+      "predictor_learning_rate": 0.00014580334415071185,
+      "step": 6322
+    },
+    {
+      "epoch": 0.65,
+      "gumbel_tau": 0.225,
+      "image_masks_loss": 0.0103,
+      "llm_learning_rate": 1.4572793736704713e-06,
+      "loss": 0.7545,
+      "output_text_mask_loss": 0.0618,
+      "predictor_learning_rate": 0.0001457279373670471,
+      "step": 6323
+    },
+    {
+      "epoch": 0.65,
+      "gumbel_tau": 0.225,
+      "image_masks_loss": 0.0072,
+      "llm_learning_rate": 1.4565254206605422e-06,
+      "loss": 0.7445,
+      "output_text_mask_loss": 0.0108,
+      "predictor_learning_rate": 0.00014565254206605423,
+      "step": 6324
+    },
+    {
+      "epoch": 0.65,
+      "gumbel_tau": 0.2249,
+      "image_masks_loss": 0.0058,
+      "llm_learning_rate": 1.455771582560358e-06,
+      "loss": 0.7143,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.0001455771582560358,
+      "step": 6325
+    },
+    {
+      "epoch": 0.65,
+      "gumbel_tau": 0.2249,
+      "image_masks_loss": 0.0068,
+      "llm_learning_rate": 1.455017859452933e-06,
+      "loss": 0.7022,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.00014550178594529328,
+      "step": 6326
+    },
+    {
+      "epoch": 0.65,
+      "gumbel_tau": 0.2248,
+      "image_masks_loss": 0.0109,
+      "llm_learning_rate": 1.454264251421269e-06,
+      "loss": 0.6557,
+      "output_text_mask_loss": 0.0275,
+      "predictor_learning_rate": 0.0001454264251421269,
+      "step": 6327
+    },
+    {
+      "epoch": 0.65,
+      "gumbel_tau": 0.2248,
+      "image_masks_loss": 0.0197,
+      "llm_learning_rate": 1.453510758548355e-06,
+      "loss": 0.7487,
+      "output_text_mask_loss": 0.0012,
+      "predictor_learning_rate": 0.00014535107585483549,
+      "step": 6328
+    },
+    {
+      "epoch": 0.65,
+      "gumbel_tau": 0.2247,
+      "image_masks_loss": 0.0041,
+      "llm_learning_rate": 1.4527573809171671e-06,
+      "loss": 0.74,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 0.0001452757380917167,
+      "step": 6329
+    },
+    {
+      "epoch": 0.65,
+      "gumbel_tau": 0.2247,
+      "image_masks_loss": 0.0073,
+      "llm_learning_rate": 1.4520041186106692e-06,
+      "loss": 0.6595,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 0.0001452004118610669,
+      "step": 6330
+    },
+    {
+      "epoch": 0.65,
+      "gumbel_tau": 0.2246,
+      "image_masks_loss": 0.0098,
+      "llm_learning_rate": 1.4512509717118108e-06,
+      "loss": 0.5991,
+      "output_text_mask_loss": 0.0046,
+      "predictor_learning_rate": 0.00014512509717118109,
+      "step": 6331
+    },
+    {
+      "epoch": 0.65,
+      "gumbel_tau": 0.2246,
+      "image_masks_loss": 0.0073,
+      "llm_learning_rate": 1.4504979403035326e-06,
+      "loss": 0.7794,
+      "output_text_mask_loss": 0.0284,
+      "predictor_learning_rate": 0.00014504979403035325,
+      "step": 6332
+    },
+    {
+      "epoch": 0.65,
+      "gumbel_tau": 0.2245,
+      "image_masks_loss": 0.0169,
+      "llm_learning_rate": 1.449745024468759e-06,
+      "loss": 0.7558,
+      "output_text_mask_loss": 0.0134,
+      "predictor_learning_rate": 0.00014497450244687588,
+      "step": 6333
+    },
+    {
+      "epoch": 0.65,
+      "gumbel_tau": 0.2245,
+      "image_masks_loss": 0.0054,
+      "llm_learning_rate": 1.4489922242904025e-06,
+      "loss": 0.8142,
+      "output_text_mask_loss": 0.0221,
+      "predictor_learning_rate": 0.00014489922242904026,
+      "step": 6334
+    },
+    {
+      "epoch": 0.65,
+      "gumbel_tau": 0.2244,
+      "image_masks_loss": 0.0055,
+      "llm_learning_rate": 1.4482395398513633e-06,
+      "loss": 0.809,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0001448239539851363,
+      "step": 6335
+    },
+    {
+      "epoch": 0.65,
+      "gumbel_tau": 0.2243,
+      "image_masks_loss": 0.0054,
+      "llm_learning_rate": 1.4474869712345286e-06,
+      "loss": 0.746,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.00014474869712345284,
+      "step": 6336
+    },
+    {
+      "epoch": 0.65,
+      "gumbel_tau": 0.2243,
+      "image_masks_loss": 0.0111,
+      "llm_learning_rate": 1.4467345185227732e-06,
+      "loss": 0.6909,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0001446734518522773,
+      "step": 6337
+    },
+    {
+      "epoch": 0.65,
+      "gumbel_tau": 0.2242,
+      "image_masks_loss": 0.0128,
+      "llm_learning_rate": 1.4459821817989575e-06,
+      "loss": 0.7637,
+      "output_text_mask_loss": 0.0111,
+      "predictor_learning_rate": 0.00014459821817989574,
+      "step": 6338
+    },
+    {
+      "epoch": 0.65,
+      "gumbel_tau": 0.2242,
+      "image_masks_loss": 0.0078,
+      "llm_learning_rate": 1.4452299611459329e-06,
+      "loss": 0.6818,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.00014452299611459327,
+      "step": 6339
+    },
+    {
+      "epoch": 0.65,
+      "gumbel_tau": 0.2241,
+      "image_masks_loss": 0.004,
+      "llm_learning_rate": 1.444477856646534e-06,
+      "loss": 0.6277,
+      "output_text_mask_loss": 0.0015,
+      "predictor_learning_rate": 0.0001444477856646534,
+      "step": 6340
+    },
+    {
+      "epoch": 0.65,
+      "gumbel_tau": 0.2241,
+      "image_masks_loss": 0.0121,
+      "llm_learning_rate": 1.4437258683835847e-06,
+      "loss": 0.6241,
+      "output_text_mask_loss": 0.0488,
+      "predictor_learning_rate": 0.00014437258683835846,
+      "step": 6341
+    },
+    {
+      "epoch": 0.65,
+      "gumbel_tau": 0.224,
+      "image_masks_loss": 0.0067,
+      "llm_learning_rate": 1.442973996439897e-06,
+      "loss": 0.7158,
+      "output_text_mask_loss": 0.0034,
+      "predictor_learning_rate": 0.00014429739964398969,
+      "step": 6342
+    },
+    {
+      "epoch": 0.65,
+      "gumbel_tau": 0.224,
+      "image_masks_loss": 0.013,
+      "llm_learning_rate": 1.4422222408982659e-06,
+      "loss": 0.6676,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.00014422222408982659,
+      "step": 6343
+    },
+    {
+      "epoch": 0.65,
+      "gumbel_tau": 0.2239,
+      "image_masks_loss": 0.0109,
+      "llm_learning_rate": 1.4414706018414765e-06,
+      "loss": 0.6623,
+      "output_text_mask_loss": 0.0254,
+      "predictor_learning_rate": 0.00014414706018414764,
+      "step": 6344
+    },
+    {
+      "epoch": 0.65,
+      "gumbel_tau": 0.2239,
+      "image_masks_loss": 0.0103,
+      "llm_learning_rate": 1.4407190793523035e-06,
+      "loss": 0.7899,
+      "output_text_mask_loss": 0.0429,
+      "predictor_learning_rate": 0.00014407190793523035,
+      "step": 6345
+    },
+    {
+      "epoch": 0.65,
+      "gumbel_tau": 0.2238,
+      "image_masks_loss": 0.0116,
+      "llm_learning_rate": 1.439967673513505e-06,
+      "loss": 0.8244,
+      "output_text_mask_loss": 0.0153,
+      "predictor_learning_rate": 0.00014399676735135048,
+      "step": 6346
+    },
+    {
+      "epoch": 0.65,
+      "gumbel_tau": 0.2238,
+      "image_masks_loss": 0.0084,
+      "llm_learning_rate": 1.4392163844078272e-06,
+      "loss": 0.7671,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.0001439216384407827,
+      "step": 6347
+    },
+    {
+      "epoch": 0.65,
+      "gumbel_tau": 0.2237,
+      "image_masks_loss": 0.0042,
+      "llm_learning_rate": 1.4384652121180043e-06,
+      "loss": 0.6782,
+      "output_text_mask_loss": 0.0313,
+      "predictor_learning_rate": 0.0001438465212118004,
+      "step": 6348
+    },
+    {
+      "epoch": 0.65,
+      "gumbel_tau": 0.2237,
+      "image_masks_loss": 0.0086,
+      "llm_learning_rate": 1.4377141567267567e-06,
+      "loss": 0.7246,
+      "output_text_mask_loss": 0.0126,
+      "predictor_learning_rate": 0.00014377141567267564,
+      "step": 6349
+    },
+    {
+      "epoch": 0.65,
+      "gumbel_tau": 0.2236,
+      "image_masks_loss": 0.0144,
+      "llm_learning_rate": 1.436963218316792e-06,
+      "loss": 0.6947,
+      "output_text_mask_loss": 0.0103,
+      "predictor_learning_rate": 0.0001436963218316792,
+      "step": 6350
+    },
+    {
+      "epoch": 0.65,
+      "gumbel_tau": 0.2236,
+      "image_masks_loss": 0.0098,
+      "llm_learning_rate": 1.4362123969708043e-06,
+      "loss": 0.6906,
+      "output_text_mask_loss": 0.0076,
+      "predictor_learning_rate": 0.00014362123969708043,
+      "step": 6351
+    },
+    {
+      "epoch": 0.65,
+      "gumbel_tau": 0.2235,
+      "image_masks_loss": 0.0175,
+      "llm_learning_rate": 1.4354616927714785e-06,
+      "loss": 0.7711,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.00014354616927714784,
+      "step": 6352
+    },
+    {
+      "epoch": 0.65,
+      "gumbel_tau": 0.2234,
+      "image_masks_loss": 0.0137,
+      "llm_learning_rate": 1.4347111058014818e-06,
+      "loss": 0.7331,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00014347111058014817,
+      "step": 6353
+    },
+    {
+      "epoch": 0.65,
+      "gumbel_tau": 0.2234,
+      "image_masks_loss": 0.0034,
+      "llm_learning_rate": 1.4339606361434714e-06,
+      "loss": 0.6787,
+      "output_text_mask_loss": 0.0212,
+      "predictor_learning_rate": 0.00014339606361434711,
+      "step": 6354
+    },
+    {
+      "epoch": 0.65,
+      "gumbel_tau": 0.2233,
+      "image_masks_loss": 0.0165,
+      "llm_learning_rate": 1.43321028388009e-06,
+      "loss": 0.7355,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.00014332102838800898,
+      "step": 6355
+    },
+    {
+      "epoch": 0.65,
+      "gumbel_tau": 0.2233,
+      "image_masks_loss": 0.0183,
+      "llm_learning_rate": 1.4324600490939682e-06,
+      "loss": 0.751,
+      "output_text_mask_loss": 0.0146,
+      "predictor_learning_rate": 0.0001432460049093968,
+      "step": 6356
+    },
+    {
+      "epoch": 0.65,
+      "gumbel_tau": 0.2232,
+      "image_masks_loss": 0.0053,
+      "llm_learning_rate": 1.4317099318677242e-06,
+      "loss": 0.769,
+      "output_text_mask_loss": 0.0014,
+      "predictor_learning_rate": 0.0001431709931867724,
+      "step": 6357
+    },
+    {
+      "epoch": 0.65,
+      "gumbel_tau": 0.2232,
+      "image_masks_loss": 0.0129,
+      "llm_learning_rate": 1.4309599322839605e-06,
+      "loss": 0.854,
+      "output_text_mask_loss": 0.0023,
+      "predictor_learning_rate": 0.00014309599322839602,
+      "step": 6358
+    },
+    {
+      "epoch": 0.65,
+      "gumbel_tau": 0.2231,
+      "image_masks_loss": 0.0131,
+      "llm_learning_rate": 1.4302100504252714e-06,
+      "loss": 0.6654,
+      "output_text_mask_loss": 0.0097,
+      "predictor_learning_rate": 0.00014302100504252712,
+      "step": 6359
+    },
+    {
+      "epoch": 0.65,
+      "gumbel_tau": 0.2231,
+      "image_masks_loss": 0.0112,
+      "llm_learning_rate": 1.4294602863742344e-06,
+      "loss": 0.7208,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.00014294602863742343,
+      "step": 6360
+    },
+    {
+      "epoch": 0.65,
+      "gumbel_tau": 0.223,
+      "image_masks_loss": 0.0235,
+      "llm_learning_rate": 1.428710640213416e-06,
+      "loss": 0.779,
+      "output_text_mask_loss": 0.0113,
+      "predictor_learning_rate": 0.0001428710640213416,
+      "step": 6361
+    },
+    {
+      "epoch": 0.65,
+      "gumbel_tau": 0.223,
+      "image_masks_loss": 0.0048,
+      "llm_learning_rate": 1.4279611120253673e-06,
+      "loss": 0.7153,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00014279611120253673,
+      "step": 6362
+    },
+    {
+      "epoch": 0.65,
+      "gumbel_tau": 0.2229,
+      "image_masks_loss": 0.0142,
+      "llm_learning_rate": 1.427211701892628e-06,
+      "loss": 0.7197,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 0.00014272117018926278,
+      "step": 6363
+    },
+    {
+      "epoch": 0.65,
+      "gumbel_tau": 0.2229,
+      "image_masks_loss": 0.0077,
+      "llm_learning_rate": 1.4264624098977265e-06,
+      "loss": 0.6566,
+      "output_text_mask_loss": 0.0004,
+      "predictor_learning_rate": 0.00014264624098977264,
+      "step": 6364
+    },
+    {
+      "epoch": 0.65,
+      "gumbel_tau": 0.2228,
+      "image_masks_loss": 0.0062,
+      "llm_learning_rate": 1.4257132361231754e-06,
+      "loss": 0.7794,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.00014257132361231753,
+      "step": 6365
+    },
+    {
+      "epoch": 0.65,
+      "gumbel_tau": 0.2228,
+      "image_masks_loss": 0.003,
+      "llm_learning_rate": 1.4249641806514757e-06,
+      "loss": 0.7761,
+      "output_text_mask_loss": 0.0071,
+      "predictor_learning_rate": 0.00014249641806514756,
+      "step": 6366
+    },
+    {
+      "epoch": 0.65,
+      "gumbel_tau": 0.2227,
+      "image_masks_loss": 0.0063,
+      "llm_learning_rate": 1.4242152435651152e-06,
+      "loss": 0.7509,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.00014242152435651152,
+      "step": 6367
+    },
+    {
+      "epoch": 0.65,
+      "gumbel_tau": 0.2227,
+      "image_masks_loss": 0.0163,
+      "llm_learning_rate": 1.4234664249465682e-06,
+      "loss": 0.6837,
+      "output_text_mask_loss": 0.0014,
+      "predictor_learning_rate": 0.00014234664249465682,
+      "step": 6368
+    },
+    {
+      "epoch": 0.65,
+      "gumbel_tau": 0.2226,
+      "image_masks_loss": 0.0089,
+      "llm_learning_rate": 1.4227177248782964e-06,
+      "loss": 0.7569,
+      "output_text_mask_loss": 0.0023,
+      "predictor_learning_rate": 0.00014227177248782964,
+      "step": 6369
+    },
+    {
+      "epoch": 0.65,
+      "gumbel_tau": 0.2226,
+      "image_masks_loss": 0.0187,
+      "llm_learning_rate": 1.4219691434427472e-06,
+      "loss": 0.7128,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.0001421969143442747,
+      "step": 6370
+    },
+    {
+      "epoch": 0.65,
+      "gumbel_tau": 0.2225,
+      "image_masks_loss": 0.0234,
+      "llm_learning_rate": 1.4212206807223583e-06,
+      "loss": 0.6735,
+      "output_text_mask_loss": 0.0488,
+      "predictor_learning_rate": 0.00014212206807223582,
+      "step": 6371
+    },
+    {
+      "epoch": 0.65,
+      "gumbel_tau": 0.2225,
+      "image_masks_loss": 0.0012,
+      "llm_learning_rate": 1.4204723367995504e-06,
+      "loss": 0.7021,
+      "output_text_mask_loss": 0.0463,
+      "predictor_learning_rate": 0.00014204723367995504,
+      "step": 6372
+    },
+    {
+      "epoch": 0.65,
+      "gumbel_tau": 0.2224,
+      "image_masks_loss": 0.0257,
+      "llm_learning_rate": 1.419724111756734e-06,
+      "loss": 0.7701,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.0001419724111756734,
+      "step": 6373
+    },
+    {
+      "epoch": 0.65,
+      "gumbel_tau": 0.2223,
+      "image_masks_loss": 0.0134,
+      "llm_learning_rate": 1.4189760056763041e-06,
+      "loss": 0.7257,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 0.0001418976005676304,
+      "step": 6374
+    },
+    {
+      "epoch": 0.65,
+      "gumbel_tau": 0.2223,
+      "image_masks_loss": 0.0081,
+      "llm_learning_rate": 1.418228018640645e-06,
+      "loss": 0.7379,
+      "output_text_mask_loss": 0.0095,
+      "predictor_learning_rate": 0.00014182280186406449,
+      "step": 6375
+    },
+    {
+      "epoch": 0.65,
+      "gumbel_tau": 0.2222,
+      "image_masks_loss": 0.0078,
+      "llm_learning_rate": 1.4174801507321253e-06,
+      "loss": 0.7774,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.00014174801507321253,
+      "step": 6376
+    },
+    {
+      "epoch": 0.65,
+      "gumbel_tau": 0.2222,
+      "image_masks_loss": 0.0092,
+      "llm_learning_rate": 1.416732402033102e-06,
+      "loss": 0.8038,
+      "output_text_mask_loss": 0.0052,
+      "predictor_learning_rate": 0.0001416732402033102,
+      "step": 6377
+    },
+    {
+      "epoch": 0.65,
+      "gumbel_tau": 0.2221,
+      "image_masks_loss": 0.0236,
+      "llm_learning_rate": 1.4159847726259204e-06,
+      "loss": 0.749,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00014159847726259205,
+      "step": 6378
+    },
+    {
+      "epoch": 0.65,
+      "gumbel_tau": 0.2221,
+      "image_masks_loss": 0.0061,
+      "llm_learning_rate": 1.415237262592911e-06,
+      "loss": 0.7776,
+      "output_text_mask_loss": 0.0282,
+      "predictor_learning_rate": 0.0001415237262592911,
+      "step": 6379
+    },
+    {
+      "epoch": 0.65,
+      "gumbel_tau": 0.222,
+      "image_masks_loss": 0.0068,
+      "llm_learning_rate": 1.41448987201639e-06,
+      "loss": 0.8086,
+      "output_text_mask_loss": 0.0041,
+      "predictor_learning_rate": 0.00014144898720163898,
+      "step": 6380
+    },
+    {
+      "epoch": 0.65,
+      "gumbel_tau": 0.222,
+      "image_masks_loss": 0.0116,
+      "llm_learning_rate": 1.413742600978662e-06,
+      "loss": 0.7956,
+      "output_text_mask_loss": 0.0034,
+      "predictor_learning_rate": 0.0001413742600978662,
+      "step": 6381
+    },
+    {
+      "epoch": 0.65,
+      "gumbel_tau": 0.2219,
+      "image_masks_loss": 0.0093,
+      "llm_learning_rate": 1.412995449562018e-06,
+      "loss": 0.7954,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.00014129954495620178,
+      "step": 6382
+    },
+    {
+      "epoch": 0.65,
+      "gumbel_tau": 0.2219,
+      "image_masks_loss": 0.0146,
+      "llm_learning_rate": 1.4122484178487355e-06,
+      "loss": 0.6904,
+      "output_text_mask_loss": 0.0179,
+      "predictor_learning_rate": 0.00014122484178487354,
+      "step": 6383
+    },
+    {
+      "epoch": 0.65,
+      "gumbel_tau": 0.2218,
+      "image_masks_loss": 0.01,
+      "llm_learning_rate": 1.4115015059210812e-06,
+      "loss": 0.8266,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 0.0001411501505921081,
+      "step": 6384
+    },
+    {
+      "epoch": 0.65,
+      "gumbel_tau": 0.2218,
+      "image_masks_loss": 0.0066,
+      "llm_learning_rate": 1.4107547138613053e-06,
+      "loss": 0.7272,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.00014107547138613051,
+      "step": 6385
+    },
+    {
+      "epoch": 0.65,
+      "gumbel_tau": 0.2217,
+      "image_masks_loss": 0.0111,
+      "llm_learning_rate": 1.4100080417516466e-06,
+      "loss": 0.6831,
+      "output_text_mask_loss": 0.0294,
+      "predictor_learning_rate": 0.00014100080417516465,
+      "step": 6386
+    },
+    {
+      "epoch": 0.65,
+      "gumbel_tau": 0.2217,
+      "image_masks_loss": 0.0101,
+      "llm_learning_rate": 1.4092614896743303e-06,
+      "loss": 0.6896,
+      "output_text_mask_loss": 0.0416,
+      "predictor_learning_rate": 0.00014092614896743302,
+      "step": 6387
+    },
+    {
+      "epoch": 0.65,
+      "gumbel_tau": 0.2216,
+      "image_masks_loss": 0.012,
+      "llm_learning_rate": 1.4085150577115683e-06,
+      "loss": 0.8368,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 0.00014085150577115683,
+      "step": 6388
+    },
+    {
+      "epoch": 0.65,
+      "gumbel_tau": 0.2216,
+      "image_masks_loss": 0.0155,
+      "llm_learning_rate": 1.4077687459455593e-06,
+      "loss": 0.7279,
+      "output_text_mask_loss": 0.0275,
+      "predictor_learning_rate": 0.00014077687459455594,
+      "step": 6389
+    },
+    {
+      "epoch": 0.65,
+      "gumbel_tau": 0.2215,
+      "image_masks_loss": 0.009,
+      "llm_learning_rate": 1.407022554458488e-06,
+      "loss": 0.6688,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.0001407022554458488,
+      "step": 6390
+    },
+    {
+      "epoch": 0.65,
+      "gumbel_tau": 0.2215,
+      "image_masks_loss": 0.0057,
+      "llm_learning_rate": 1.406276483332528e-06,
+      "loss": 0.7043,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.0001406276483332528,
+      "step": 6391
+    },
+    {
+      "epoch": 0.65,
+      "gumbel_tau": 0.2214,
+      "image_masks_loss": 0.0183,
+      "llm_learning_rate": 1.4055305326498385e-06,
+      "loss": 0.8212,
+      "output_text_mask_loss": 0.0078,
+      "predictor_learning_rate": 0.00014055305326498386,
+      "step": 6392
+    },
+    {
+      "epoch": 0.66,
+      "gumbel_tau": 0.2214,
+      "image_masks_loss": 0.0103,
+      "llm_learning_rate": 1.4047847024925643e-06,
+      "loss": 0.6078,
+      "output_text_mask_loss": 0.0248,
+      "predictor_learning_rate": 0.00014047847024925642,
+      "step": 6393
+    },
+    {
+      "epoch": 0.66,
+      "gumbel_tau": 0.2213,
+      "image_masks_loss": 0.0131,
+      "llm_learning_rate": 1.404038992942838e-06,
+      "loss": 0.8128,
+      "output_text_mask_loss": 0.024,
+      "predictor_learning_rate": 0.0001404038992942838,
+      "step": 6394
+    },
+    {
+      "epoch": 0.66,
+      "gumbel_tau": 0.2212,
+      "image_masks_loss": 0.0205,
+      "llm_learning_rate": 1.4032934040827792e-06,
+      "loss": 0.7585,
+      "output_text_mask_loss": 0.0204,
+      "predictor_learning_rate": 0.00014032934040827792,
+      "step": 6395
+    },
+    {
+      "epoch": 0.66,
+      "gumbel_tau": 0.2212,
+      "image_masks_loss": 0.013,
+      "llm_learning_rate": 1.4025479359944932e-06,
+      "loss": 0.6367,
+      "output_text_mask_loss": 0.0219,
+      "predictor_learning_rate": 0.00014025479359944931,
+      "step": 6396
+    },
+    {
+      "epoch": 0.66,
+      "gumbel_tau": 0.2211,
+      "image_masks_loss": 0.0189,
+      "llm_learning_rate": 1.4018025887600722e-06,
+      "loss": 0.764,
+      "output_text_mask_loss": 0.0118,
+      "predictor_learning_rate": 0.00014018025887600722,
+      "step": 6397
+    },
+    {
+      "epoch": 0.66,
+      "gumbel_tau": 0.2211,
+      "image_masks_loss": 0.0223,
+      "llm_learning_rate": 1.4010573624615986e-06,
+      "loss": 0.7151,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00014010573624615986,
+      "step": 6398
+    },
+    {
+      "epoch": 0.66,
+      "gumbel_tau": 0.221,
+      "image_masks_loss": 0.0093,
+      "llm_learning_rate": 1.4003122571811345e-06,
+      "loss": 0.7609,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 0.00014003122571811343,
+      "step": 6399
+    },
+    {
+      "epoch": 0.66,
+      "gumbel_tau": 0.221,
+      "image_masks_loss": 0.0083,
+      "llm_learning_rate": 1.3995672730007343e-06,
+      "loss": 0.7781,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0001399567273000734,
+      "step": 6400
+    },
+    {
+      "epoch": 0.66,
+      "gumbel_tau": 0.2209,
+      "image_masks_loss": 0.0177,
+      "llm_learning_rate": 1.3988224100024372e-06,
+      "loss": 0.7003,
+      "output_text_mask_loss": 0.0565,
+      "predictor_learning_rate": 0.00013988224100024371,
+      "step": 6401
+    },
+    {
+      "epoch": 0.66,
+      "gumbel_tau": 0.2209,
+      "image_masks_loss": 0.0104,
+      "llm_learning_rate": 1.3980776682682681e-06,
+      "loss": 0.7989,
+      "output_text_mask_loss": 0.017,
+      "predictor_learning_rate": 0.0001398077668268268,
+      "step": 6402
+    },
+    {
+      "epoch": 0.66,
+      "gumbel_tau": 0.2208,
+      "image_masks_loss": 0.0182,
+      "llm_learning_rate": 1.3973330478802416e-06,
+      "loss": 0.7182,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.00013973330478802415,
+      "step": 6403
+    },
+    {
+      "epoch": 0.66,
+      "gumbel_tau": 0.2208,
+      "image_masks_loss": 0.0205,
+      "llm_learning_rate": 1.3965885489203556e-06,
+      "loss": 0.8022,
+      "output_text_mask_loss": 0.001,
+      "predictor_learning_rate": 0.00013965885489203555,
+      "step": 6404
+    },
+    {
+      "epoch": 0.66,
+      "gumbel_tau": 0.2207,
+      "image_masks_loss": 0.0133,
+      "llm_learning_rate": 1.3958441714705968e-06,
+      "loss": 0.7205,
+      "output_text_mask_loss": 0.0488,
+      "predictor_learning_rate": 0.00013958441714705965,
+      "step": 6405
+    },
+    {
+      "epoch": 0.66,
+      "gumbel_tau": 0.2207,
+      "image_masks_loss": 0.0104,
+      "llm_learning_rate": 1.3950999156129372e-06,
+      "loss": 0.7228,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00013950999156129372,
+      "step": 6406
+    },
+    {
+      "epoch": 0.66,
+      "gumbel_tau": 0.2206,
+      "image_masks_loss": 0.0119,
+      "llm_learning_rate": 1.3943557814293357e-06,
+      "loss": 0.7386,
+      "output_text_mask_loss": 0.0137,
+      "predictor_learning_rate": 0.00013943557814293357,
+      "step": 6407
+    },
+    {
+      "epoch": 0.66,
+      "gumbel_tau": 0.2206,
+      "image_masks_loss": 0.0049,
+      "llm_learning_rate": 1.3936117690017386e-06,
+      "loss": 0.7741,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 0.00013936117690017385,
+      "step": 6408
+    },
+    {
+      "epoch": 0.66,
+      "gumbel_tau": 0.2205,
+      "image_masks_loss": 0.009,
+      "llm_learning_rate": 1.3928678784120766e-06,
+      "loss": 0.7605,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00013928678784120766,
+      "step": 6409
+    },
+    {
+      "epoch": 0.66,
+      "gumbel_tau": 0.2205,
+      "image_masks_loss": 0.0118,
+      "llm_learning_rate": 1.3921241097422711e-06,
+      "loss": 0.7118,
+      "output_text_mask_loss": 0.0179,
+      "predictor_learning_rate": 0.0001392124109742271,
+      "step": 6410
+    },
+    {
+      "epoch": 0.66,
+      "gumbel_tau": 0.2204,
+      "image_masks_loss": 0.0138,
+      "llm_learning_rate": 1.3913804630742262e-06,
+      "loss": 0.7768,
+      "output_text_mask_loss": 0.0584,
+      "predictor_learning_rate": 0.0001391380463074226,
+      "step": 6411
+    },
+    {
+      "epoch": 0.66,
+      "gumbel_tau": 0.2204,
+      "image_masks_loss": 0.0114,
+      "llm_learning_rate": 1.3906369384898339e-06,
+      "loss": 0.7372,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.00013906369384898338,
+      "step": 6412
+    },
+    {
+      "epoch": 0.66,
+      "gumbel_tau": 0.2203,
+      "image_masks_loss": 0.0151,
+      "llm_learning_rate": 1.3898935360709732e-06,
+      "loss": 0.6929,
+      "output_text_mask_loss": 0.0202,
+      "predictor_learning_rate": 0.0001389893536070973,
+      "step": 6413
+    },
+    {
+      "epoch": 0.66,
+      "gumbel_tau": 0.2203,
+      "image_masks_loss": 0.0149,
+      "llm_learning_rate": 1.3891502558995088e-06,
+      "loss": 0.7796,
+      "output_text_mask_loss": 0.0404,
+      "predictor_learning_rate": 0.00013891502558995087,
+      "step": 6414
+    },
+    {
+      "epoch": 0.66,
+      "gumbel_tau": 0.2202,
+      "image_masks_loss": 0.0145,
+      "llm_learning_rate": 1.3884070980572927e-06,
+      "loss": 0.7953,
+      "output_text_mask_loss": 0.0076,
+      "predictor_learning_rate": 0.00013884070980572927,
+      "step": 6415
+    },
+    {
+      "epoch": 0.66,
+      "gumbel_tau": 0.2202,
+      "image_masks_loss": 0.0063,
+      "llm_learning_rate": 1.387664062626163e-06,
+      "loss": 0.7179,
+      "output_text_mask_loss": 0.0223,
+      "predictor_learning_rate": 0.0001387664062626163,
+      "step": 6416
+    },
+    {
+      "epoch": 0.66,
+      "gumbel_tau": 0.2201,
+      "image_masks_loss": 0.0154,
+      "llm_learning_rate": 1.386921149687944e-06,
+      "loss": 0.7051,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 0.00013869211496879438,
+      "step": 6417
+    },
+    {
+      "epoch": 0.66,
+      "gumbel_tau": 0.22,
+      "image_masks_loss": 0.0161,
+      "llm_learning_rate": 1.3861783593244476e-06,
+      "loss": 0.7424,
+      "output_text_mask_loss": 0.0263,
+      "predictor_learning_rate": 0.00013861783593244477,
+      "step": 6418
+    },
+    {
+      "epoch": 0.66,
+      "gumbel_tau": 0.22,
+      "image_masks_loss": 0.012,
+      "llm_learning_rate": 1.385435691617471e-06,
+      "loss": 0.7223,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.0001385435691617471,
+      "step": 6419
+    },
+    {
+      "epoch": 0.66,
+      "gumbel_tau": 0.2199,
+      "image_masks_loss": 0.0085,
+      "llm_learning_rate": 1.384693146648799e-06,
+      "loss": 0.7298,
+      "output_text_mask_loss": 0.0012,
+      "predictor_learning_rate": 0.00013846931466487988,
+      "step": 6420
+    },
+    {
+      "epoch": 0.66,
+      "gumbel_tau": 0.2199,
+      "image_masks_loss": 0.0055,
+      "llm_learning_rate": 1.3839507245002015e-06,
+      "loss": 0.723,
+      "output_text_mask_loss": 0.0078,
+      "predictor_learning_rate": 0.00013839507245002014,
+      "step": 6421
+    },
+    {
+      "epoch": 0.66,
+      "gumbel_tau": 0.2198,
+      "image_masks_loss": 0.0176,
+      "llm_learning_rate": 1.3832084252534353e-06,
+      "loss": 0.789,
+      "output_text_mask_loss": 0.0698,
+      "predictor_learning_rate": 0.0001383208425253435,
+      "step": 6422
+    },
+    {
+      "epoch": 0.66,
+      "gumbel_tau": 0.2198,
+      "image_masks_loss": 0.0143,
+      "llm_learning_rate": 1.3824662489902457e-06,
+      "loss": 0.8051,
+      "output_text_mask_loss": 0.024,
+      "predictor_learning_rate": 0.00013824662489902455,
+      "step": 6423
+    },
+    {
+      "epoch": 0.66,
+      "gumbel_tau": 0.2197,
+      "image_masks_loss": 0.0202,
+      "llm_learning_rate": 1.381724195792362e-06,
+      "loss": 0.7042,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.0001381724195792362,
+      "step": 6424
+    },
+    {
+      "epoch": 0.66,
+      "gumbel_tau": 0.2197,
+      "image_masks_loss": 0.0138,
+      "llm_learning_rate": 1.3809822657415013e-06,
+      "loss": 0.8782,
+      "output_text_mask_loss": 0.0006,
+      "predictor_learning_rate": 0.0001380982265741501,
+      "step": 6425
+    },
+    {
+      "epoch": 0.66,
+      "gumbel_tau": 0.2196,
+      "image_masks_loss": 0.0093,
+      "llm_learning_rate": 1.3802404589193653e-06,
+      "loss": 0.6902,
+      "output_text_mask_loss": 0.0385,
+      "predictor_learning_rate": 0.0001380240458919365,
+      "step": 6426
+    },
+    {
+      "epoch": 0.66,
+      "gumbel_tau": 0.2196,
+      "image_masks_loss": 0.0087,
+      "llm_learning_rate": 1.3794987754076443e-06,
+      "loss": 0.6885,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00013794987754076442,
+      "step": 6427
+    },
+    {
+      "epoch": 0.66,
+      "gumbel_tau": 0.2195,
+      "image_masks_loss": 0.0225,
+      "llm_learning_rate": 1.378757215288014e-06,
+      "loss": 0.6565,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.0001378757215288014,
+      "step": 6428
+    },
+    {
+      "epoch": 0.66,
+      "gumbel_tau": 0.2195,
+      "image_masks_loss": 0.0156,
+      "llm_learning_rate": 1.3780157786421361e-06,
+      "loss": 0.7854,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.0001378015778642136,
+      "step": 6429
+    },
+    {
+      "epoch": 0.66,
+      "gumbel_tau": 0.2194,
+      "image_masks_loss": 0.0154,
+      "llm_learning_rate": 1.37727446555166e-06,
+      "loss": 0.7616,
+      "output_text_mask_loss": 0.0034,
+      "predictor_learning_rate": 0.000137727446555166,
+      "step": 6430
+    },
+    {
+      "epoch": 0.66,
+      "gumbel_tau": 0.2194,
+      "image_masks_loss": 0.0207,
+      "llm_learning_rate": 1.3765332760982212e-06,
+      "loss": 0.7705,
+      "output_text_mask_loss": 0.001,
+      "predictor_learning_rate": 0.00013765332760982212,
+      "step": 6431
+    },
+    {
+      "epoch": 0.66,
+      "gumbel_tau": 0.2193,
+      "image_masks_loss": 0.0076,
+      "llm_learning_rate": 1.3757922103634402e-06,
+      "loss": 0.7183,
+      "output_text_mask_loss": 0.0016,
+      "predictor_learning_rate": 0.00013757922103634402,
+      "step": 6432
+    },
+    {
+      "epoch": 0.66,
+      "gumbel_tau": 0.2193,
+      "image_masks_loss": 0.0154,
+      "llm_learning_rate": 1.3750512684289246e-06,
+      "loss": 0.7129,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00013750512684289245,
+      "step": 6433
+    },
+    {
+      "epoch": 0.66,
+      "gumbel_tau": 0.2192,
+      "image_masks_loss": 0.0118,
+      "llm_learning_rate": 1.3743104503762697e-06,
+      "loss": 0.7161,
+      "output_text_mask_loss": 0.0029,
+      "predictor_learning_rate": 0.00013743104503762695,
+      "step": 6434
+    },
+    {
+      "epoch": 0.66,
+      "gumbel_tau": 0.2192,
+      "image_masks_loss": 0.0212,
+      "llm_learning_rate": 1.3735697562870553e-06,
+      "loss": 0.743,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.00013735697562870553,
+      "step": 6435
+    },
+    {
+      "epoch": 0.66,
+      "gumbel_tau": 0.2191,
+      "image_masks_loss": 0.0317,
+      "llm_learning_rate": 1.372829186242848e-06,
+      "loss": 0.7055,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.00013728291862428479,
+      "step": 6436
+    },
+    {
+      "epoch": 0.66,
+      "gumbel_tau": 0.2191,
+      "image_masks_loss": 0.0096,
+      "llm_learning_rate": 1.3720887403252018e-06,
+      "loss": 0.7321,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.00013720887403252015,
+      "step": 6437
+    },
+    {
+      "epoch": 0.66,
+      "gumbel_tau": 0.219,
+      "image_masks_loss": 0.0113,
+      "llm_learning_rate": 1.3713484186156555e-06,
+      "loss": 0.7325,
+      "output_text_mask_loss": 0.0004,
+      "predictor_learning_rate": 0.00013713484186156556,
+      "step": 6438
+    },
+    {
+      "epoch": 0.66,
+      "gumbel_tau": 0.219,
+      "image_masks_loss": 0.0099,
+      "llm_learning_rate": 1.3706082211957355e-06,
+      "loss": 0.6838,
+      "output_text_mask_loss": 0.0278,
+      "predictor_learning_rate": 0.00013706082211957355,
+      "step": 6439
+    },
+    {
+      "epoch": 0.66,
+      "gumbel_tau": 0.2189,
+      "image_masks_loss": 0.0023,
+      "llm_learning_rate": 1.369868148146954e-06,
+      "loss": 0.737,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 0.0001369868148146954,
+      "step": 6440
+    },
+    {
+      "epoch": 0.66,
+      "gumbel_tau": 0.2189,
+      "image_masks_loss": 0.0062,
+      "llm_learning_rate": 1.3691281995508076e-06,
+      "loss": 0.7395,
+      "output_text_mask_loss": 0.0082,
+      "predictor_learning_rate": 0.00013691281995508077,
+      "step": 6441
+    },
+    {
+      "epoch": 0.66,
+      "gumbel_tau": 0.2188,
+      "image_masks_loss": 0.0098,
+      "llm_learning_rate": 1.3683883754887844e-06,
+      "loss": 0.7902,
+      "output_text_mask_loss": 0.0664,
+      "predictor_learning_rate": 0.00013683883754887844,
+      "step": 6442
+    },
+    {
+      "epoch": 0.66,
+      "gumbel_tau": 0.2188,
+      "image_masks_loss": 0.0076,
+      "llm_learning_rate": 1.3676486760423538e-06,
+      "loss": 0.6959,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00013676486760423538,
+      "step": 6443
+    },
+    {
+      "epoch": 0.66,
+      "gumbel_tau": 0.2187,
+      "image_masks_loss": 0.0046,
+      "llm_learning_rate": 1.366909101292973e-06,
+      "loss": 0.91,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.0001366909101292973,
+      "step": 6444
+    },
+    {
+      "epoch": 0.66,
+      "gumbel_tau": 0.2187,
+      "image_masks_loss": 0.009,
+      "llm_learning_rate": 1.366169651322086e-06,
+      "loss": 0.756,
+      "output_text_mask_loss": 0.0231,
+      "predictor_learning_rate": 0.00013661696513220858,
+      "step": 6445
+    },
+    {
+      "epoch": 0.66,
+      "gumbel_tau": 0.2186,
+      "image_masks_loss": 0.0064,
+      "llm_learning_rate": 1.3654303262111223e-06,
+      "loss": 0.7389,
+      "output_text_mask_loss": 0.0004,
+      "predictor_learning_rate": 0.00013654303262111223,
+      "step": 6446
+    },
+    {
+      "epoch": 0.66,
+      "gumbel_tau": 0.2185,
+      "image_masks_loss": 0.0074,
+      "llm_learning_rate": 1.3646911260414985e-06,
+      "loss": 0.7955,
+      "output_text_mask_loss": 0.0162,
+      "predictor_learning_rate": 0.00013646911260414984,
+      "step": 6447
+    },
+    {
+      "epoch": 0.66,
+      "gumbel_tau": 0.2185,
+      "image_masks_loss": 0.0041,
+      "llm_learning_rate": 1.3639520508946156e-06,
+      "loss": 0.6881,
+      "output_text_mask_loss": 0.001,
+      "predictor_learning_rate": 0.00013639520508946154,
+      "step": 6448
+    },
+    {
+      "epoch": 0.66,
+      "gumbel_tau": 0.2184,
+      "image_masks_loss": 0.0134,
+      "llm_learning_rate": 1.3632131008518642e-06,
+      "loss": 0.7218,
+      "output_text_mask_loss": 0.0441,
+      "predictor_learning_rate": 0.0001363213100851864,
+      "step": 6449
+    },
+    {
+      "epoch": 0.66,
+      "gumbel_tau": 0.2184,
+      "image_masks_loss": 0.007,
+      "llm_learning_rate": 1.3624742759946183e-06,
+      "loss": 0.7375,
+      "output_text_mask_loss": 0.0174,
+      "predictor_learning_rate": 0.00013624742759946182,
+      "step": 6450
+    },
+    {
+      "epoch": 0.66,
+      "gumbel_tau": 0.2183,
+      "image_masks_loss": 0.0076,
+      "llm_learning_rate": 1.361735576404239e-06,
+      "loss": 0.6021,
+      "output_text_mask_loss": 0.0032,
+      "predictor_learning_rate": 0.0001361735576404239,
+      "step": 6451
+    },
+    {
+      "epoch": 0.66,
+      "gumbel_tau": 0.2183,
+      "image_masks_loss": 0.0108,
+      "llm_learning_rate": 1.3609970021620733e-06,
+      "loss": 0.756,
+      "output_text_mask_loss": 0.0576,
+      "predictor_learning_rate": 0.00013609970021620732,
+      "step": 6452
+    },
+    {
+      "epoch": 0.66,
+      "gumbel_tau": 0.2182,
+      "image_masks_loss": 0.0143,
+      "llm_learning_rate": 1.3602585533494545e-06,
+      "loss": 0.7036,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.00013602585533494544,
+      "step": 6453
+    },
+    {
+      "epoch": 0.66,
+      "gumbel_tau": 0.2182,
+      "image_masks_loss": 0.0184,
+      "llm_learning_rate": 1.3595202300477028e-06,
+      "loss": 0.7247,
+      "output_text_mask_loss": 0.0248,
+      "predictor_learning_rate": 0.00013595202300477026,
+      "step": 6454
+    },
+    {
+      "epoch": 0.66,
+      "gumbel_tau": 0.2181,
+      "image_masks_loss": 0.0063,
+      "llm_learning_rate": 1.3587820323381235e-06,
+      "loss": 0.7353,
+      "output_text_mask_loss": 0.0174,
+      "predictor_learning_rate": 0.00013587820323381232,
+      "step": 6455
+    },
+    {
+      "epoch": 0.66,
+      "gumbel_tau": 0.2181,
+      "image_masks_loss": 0.0158,
+      "llm_learning_rate": 1.358043960302009e-06,
+      "loss": 0.8018,
+      "output_text_mask_loss": 0.0177,
+      "predictor_learning_rate": 0.0001358043960302009,
+      "step": 6456
+    },
+    {
+      "epoch": 0.66,
+      "gumbel_tau": 0.218,
+      "image_masks_loss": 0.0082,
+      "llm_learning_rate": 1.357306014020637e-06,
+      "loss": 0.8062,
+      "output_text_mask_loss": 0.005,
+      "predictor_learning_rate": 0.0001357306014020637,
+      "step": 6457
+    },
+    {
+      "epoch": 0.66,
+      "gumbel_tau": 0.218,
+      "image_masks_loss": 0.0085,
+      "llm_learning_rate": 1.3565681935752722e-06,
+      "loss": 0.7493,
+      "output_text_mask_loss": 0.0027,
+      "predictor_learning_rate": 0.00013565681935752722,
+      "step": 6458
+    },
+    {
+      "epoch": 0.66,
+      "gumbel_tau": 0.2179,
+      "image_masks_loss": 0.009,
+      "llm_learning_rate": 1.3558304990471643e-06,
+      "loss": 0.6921,
+      "output_text_mask_loss": 0.0098,
+      "predictor_learning_rate": 0.00013558304990471642,
+      "step": 6459
+    },
+    {
+      "epoch": 0.66,
+      "gumbel_tau": 0.2179,
+      "image_masks_loss": 0.0128,
+      "llm_learning_rate": 1.3550929305175508e-06,
+      "loss": 0.7374,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.00013550929305175508,
+      "step": 6460
+    },
+    {
+      "epoch": 0.66,
+      "gumbel_tau": 0.2178,
+      "image_masks_loss": 0.0069,
+      "llm_learning_rate": 1.354355488067653e-06,
+      "loss": 0.814,
+      "output_text_mask_loss": 0.0511,
+      "predictor_learning_rate": 0.0001354355488067653,
+      "step": 6461
+    },
+    {
+      "epoch": 0.66,
+      "gumbel_tau": 0.2178,
+      "image_masks_loss": 0.0029,
+      "llm_learning_rate": 1.3536181717786816e-06,
+      "loss": 0.7529,
+      "output_text_mask_loss": 0.0534,
+      "predictor_learning_rate": 0.00013536181717786817,
+      "step": 6462
+    },
+    {
+      "epoch": 0.66,
+      "gumbel_tau": 0.2177,
+      "image_masks_loss": 0.0129,
+      "llm_learning_rate": 1.3528809817318307e-06,
+      "loss": 0.7401,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00013528809817318305,
+      "step": 6463
+    },
+    {
+      "epoch": 0.66,
+      "gumbel_tau": 0.2177,
+      "image_masks_loss": 0.0217,
+      "llm_learning_rate": 1.3521439180082812e-06,
+      "loss": 0.7192,
+      "output_text_mask_loss": 0.0195,
+      "predictor_learning_rate": 0.00013521439180082812,
+      "step": 6464
+    },
+    {
+      "epoch": 0.66,
+      "gumbel_tau": 0.2176,
+      "image_masks_loss": 0.0077,
+      "llm_learning_rate": 1.3514069806892e-06,
+      "loss": 0.8068,
+      "output_text_mask_loss": 0.0054,
+      "predictor_learning_rate": 0.00013514069806891998,
+      "step": 6465
+    },
+    {
+      "epoch": 0.66,
+      "gumbel_tau": 0.2176,
+      "image_masks_loss": 0.0231,
+      "llm_learning_rate": 1.3506701698557406e-06,
+      "loss": 0.6862,
+      "output_text_mask_loss": 0.0101,
+      "predictor_learning_rate": 0.00013506701698557405,
+      "step": 6466
+    },
+    {
+      "epoch": 0.66,
+      "gumbel_tau": 0.2175,
+      "image_masks_loss": 0.008,
+      "llm_learning_rate": 1.349933485589042e-06,
+      "loss": 0.7908,
+      "output_text_mask_loss": 0.0076,
+      "predictor_learning_rate": 0.00013499334855890418,
+      "step": 6467
+    },
+    {
+      "epoch": 0.66,
+      "gumbel_tau": 0.2175,
+      "image_masks_loss": 0.0027,
+      "llm_learning_rate": 1.349196927970229e-06,
+      "loss": 0.7219,
+      "output_text_mask_loss": 0.0267,
+      "predictor_learning_rate": 0.00013491969279702288,
+      "step": 6468
+    },
+    {
+      "epoch": 0.66,
+      "gumbel_tau": 0.2174,
+      "image_masks_loss": 0.0252,
+      "llm_learning_rate": 1.3484604970804142e-06,
+      "loss": 0.7778,
+      "output_text_mask_loss": 0.0353,
+      "predictor_learning_rate": 0.0001348460497080414,
+      "step": 6469
+    },
+    {
+      "epoch": 0.66,
+      "gumbel_tau": 0.2174,
+      "image_masks_loss": 0.0212,
+      "llm_learning_rate": 1.3477241930006946e-06,
+      "loss": 0.7423,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00013477241930006945,
+      "step": 6470
+    },
+    {
+      "epoch": 0.66,
+      "gumbel_tau": 0.2173,
+      "image_masks_loss": 0.0213,
+      "llm_learning_rate": 1.3469880158121532e-06,
+      "loss": 0.7317,
+      "output_text_mask_loss": 0.0088,
+      "predictor_learning_rate": 0.00013469880158121532,
+      "step": 6471
+    },
+    {
+      "epoch": 0.66,
+      "gumbel_tau": 0.2173,
+      "image_masks_loss": 0.0099,
+      "llm_learning_rate": 1.346251965595861e-06,
+      "loss": 0.7337,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.0001346251965595861,
+      "step": 6472
+    },
+    {
+      "epoch": 0.66,
+      "gumbel_tau": 0.2172,
+      "image_masks_loss": 0.0078,
+      "llm_learning_rate": 1.34551604243287e-06,
+      "loss": 0.8038,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.000134551604243287,
+      "step": 6473
+    },
+    {
+      "epoch": 0.66,
+      "gumbel_tau": 0.2172,
+      "image_masks_loss": 0.0123,
+      "llm_learning_rate": 1.3447802464042252e-06,
+      "loss": 0.7515,
+      "output_text_mask_loss": 0.021,
+      "predictor_learning_rate": 0.00013447802464042254,
+      "step": 6474
+    },
+    {
+      "epoch": 0.66,
+      "gumbel_tau": 0.2171,
+      "image_masks_loss": 0.0058,
+      "llm_learning_rate": 1.3440445775909526e-06,
+      "loss": 0.5924,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00013440445775909527,
+      "step": 6475
+    },
+    {
+      "epoch": 0.66,
+      "gumbel_tau": 0.2171,
+      "image_masks_loss": 0.0059,
+      "llm_learning_rate": 1.3433090360740661e-06,
+      "loss": 0.7022,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 0.00013433090360740661,
+      "step": 6476
+    },
+    {
+      "epoch": 0.66,
+      "gumbel_tau": 0.217,
+      "image_masks_loss": 0.0046,
+      "llm_learning_rate": 1.3425736219345647e-06,
+      "loss": 0.7639,
+      "output_text_mask_loss": 0.0058,
+      "predictor_learning_rate": 0.00013425736219345647,
+      "step": 6477
+    },
+    {
+      "epoch": 0.66,
+      "gumbel_tau": 0.217,
+      "image_masks_loss": 0.0071,
+      "llm_learning_rate": 1.3418383352534344e-06,
+      "loss": 0.6546,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00013418383352534343,
+      "step": 6478
+    },
+    {
+      "epoch": 0.66,
+      "gumbel_tau": 0.2169,
+      "image_masks_loss": 0.0073,
+      "llm_learning_rate": 1.3411031761116462e-06,
+      "loss": 0.7754,
+      "output_text_mask_loss": 0.0322,
+      "predictor_learning_rate": 0.0001341103176111646,
+      "step": 6479
+    },
+    {
+      "epoch": 0.66,
+      "gumbel_tau": 0.2169,
+      "image_masks_loss": 0.0167,
+      "llm_learning_rate": 1.3403681445901568e-06,
+      "loss": 0.7955,
+      "output_text_mask_loss": 0.0055,
+      "predictor_learning_rate": 0.00013403681445901567,
+      "step": 6480
+    },
+    {
+      "epoch": 0.66,
+      "gumbel_tau": 0.2168,
+      "image_masks_loss": 0.0225,
+      "llm_learning_rate": 1.339633240769911e-06,
+      "loss": 0.711,
+      "output_text_mask_loss": 0.0208,
+      "predictor_learning_rate": 0.0001339633240769911,
+      "step": 6481
+    },
+    {
+      "epoch": 0.66,
+      "gumbel_tau": 0.2168,
+      "image_masks_loss": 0.0075,
+      "llm_learning_rate": 1.3388984647318377e-06,
+      "loss": 0.6875,
+      "output_text_mask_loss": 0.0065,
+      "predictor_learning_rate": 0.00013388984647318375,
+      "step": 6482
+    },
+    {
+      "epoch": 0.66,
+      "gumbel_tau": 0.2167,
+      "image_masks_loss": 0.0091,
+      "llm_learning_rate": 1.338163816556851e-06,
+      "loss": 0.7485,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.00013381638165568512,
+      "step": 6483
+    },
+    {
+      "epoch": 0.66,
+      "gumbel_tau": 0.2166,
+      "image_masks_loss": 0.0084,
+      "llm_learning_rate": 1.3374292963258534e-06,
+      "loss": 0.7416,
+      "output_text_mask_loss": 0.0324,
+      "predictor_learning_rate": 0.00013374292963258533,
+      "step": 6484
+    },
+    {
+      "epoch": 0.66,
+      "gumbel_tau": 0.2166,
+      "image_masks_loss": 0.0113,
+      "llm_learning_rate": 1.3366949041197308e-06,
+      "loss": 0.6838,
+      "output_text_mask_loss": 0.0565,
+      "predictor_learning_rate": 0.00013366949041197307,
+      "step": 6485
+    },
+    {
+      "epoch": 0.66,
+      "gumbel_tau": 0.2165,
+      "image_masks_loss": 0.0043,
+      "llm_learning_rate": 1.335960640019357e-06,
+      "loss": 0.7828,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.00013359606400193568,
+      "step": 6486
+    },
+    {
+      "epoch": 0.66,
+      "gumbel_tau": 0.2165,
+      "image_masks_loss": 0.0073,
+      "llm_learning_rate": 1.335226504105589e-06,
+      "loss": 0.7701,
+      "output_text_mask_loss": 0.0055,
+      "predictor_learning_rate": 0.0001335226504105589,
+      "step": 6487
+    },
+    {
+      "epoch": 0.66,
+      "gumbel_tau": 0.2164,
+      "image_masks_loss": 0.0098,
+      "llm_learning_rate": 1.3344924964592737e-06,
+      "loss": 0.802,
+      "output_text_mask_loss": 0.061,
+      "predictor_learning_rate": 0.00013344924964592736,
+      "step": 6488
+    },
+    {
+      "epoch": 0.66,
+      "gumbel_tau": 0.2164,
+      "image_masks_loss": 0.0167,
+      "llm_learning_rate": 1.333758617161241e-06,
+      "loss": 0.6913,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0001333758617161241,
+      "step": 6489
+    },
+    {
+      "epoch": 0.66,
+      "gumbel_tau": 0.2163,
+      "image_masks_loss": 0.003,
+      "llm_learning_rate": 1.3330248662923068e-06,
+      "loss": 0.7528,
+      "output_text_mask_loss": 0.0713,
+      "predictor_learning_rate": 0.00013330248662923066,
+      "step": 6490
+    },
+    {
+      "epoch": 0.67,
+      "gumbel_tau": 0.2163,
+      "image_masks_loss": 0.0015,
+      "llm_learning_rate": 1.332291243933275e-06,
+      "loss": 0.6875,
+      "output_text_mask_loss": 0.0204,
+      "predictor_learning_rate": 0.0001332291243933275,
+      "step": 6491
+    },
+    {
+      "epoch": 0.67,
+      "gumbel_tau": 0.2162,
+      "image_masks_loss": 0.0105,
+      "llm_learning_rate": 1.3315577501649312e-06,
+      "loss": 0.8075,
+      "output_text_mask_loss": 0.008,
+      "predictor_learning_rate": 0.00013315577501649313,
+      "step": 6492
+    },
+    {
+      "epoch": 0.67,
+      "gumbel_tau": 0.2162,
+      "image_masks_loss": 0.013,
+      "llm_learning_rate": 1.3308243850680496e-06,
+      "loss": 0.6928,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00013308243850680495,
+      "step": 6493
+    },
+    {
+      "epoch": 0.67,
+      "gumbel_tau": 0.2161,
+      "image_masks_loss": 0.0056,
+      "llm_learning_rate": 1.330091148723392e-06,
+      "loss": 0.7303,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.0001330091148723392,
+      "step": 6494
+    },
+    {
+      "epoch": 0.67,
+      "gumbel_tau": 0.2161,
+      "image_masks_loss": 0.0292,
+      "llm_learning_rate": 1.3293580412117026e-06,
+      "loss": 0.7411,
+      "output_text_mask_loss": 0.0534,
+      "predictor_learning_rate": 0.00013293580412117027,
+      "step": 6495
+    },
+    {
+      "epoch": 0.67,
+      "gumbel_tau": 0.216,
+      "image_masks_loss": 0.0064,
+      "llm_learning_rate": 1.3286250626137136e-06,
+      "loss": 0.748,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00013286250626137137,
+      "step": 6496
+    },
+    {
+      "epoch": 0.67,
+      "gumbel_tau": 0.216,
+      "image_masks_loss": 0.0221,
+      "llm_learning_rate": 1.3278922130101418e-06,
+      "loss": 0.6656,
+      "output_text_mask_loss": 0.0748,
+      "predictor_learning_rate": 0.00013278922130101418,
+      "step": 6497
+    },
+    {
+      "epoch": 0.67,
+      "gumbel_tau": 0.2159,
+      "image_masks_loss": 0.0112,
+      "llm_learning_rate": 1.32715949248169e-06,
+      "loss": 0.6855,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.000132715949248169,
+      "step": 6498
+    },
+    {
+      "epoch": 0.67,
+      "gumbel_tau": 0.2159,
+      "image_masks_loss": 0.0213,
+      "llm_learning_rate": 1.3264269011090474e-06,
+      "loss": 0.7381,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00013264269011090474,
+      "step": 6499
+    },
+    {
+      "epoch": 0.67,
+      "gumbel_tau": 0.2158,
+      "image_masks_loss": 0.0086,
+      "llm_learning_rate": 1.3256944389728877e-06,
+      "loss": 0.7499,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00013256944389728874,
+      "step": 6500
+    },
+    {
+      "epoch": 0.67,
+      "gumbel_tau": 0.2158,
+      "image_masks_loss": 0.0045,
+      "llm_learning_rate": 1.324962106153873e-06,
+      "loss": 0.8088,
+      "output_text_mask_loss": 0.005,
+      "predictor_learning_rate": 0.0001324962106153873,
+      "step": 6501
+    },
+    {
+      "epoch": 0.67,
+      "gumbel_tau": 0.2157,
+      "image_masks_loss": 0.0051,
+      "llm_learning_rate": 1.3242299027326484e-06,
+      "loss": 0.6464,
+      "output_text_mask_loss": 0.0626,
+      "predictor_learning_rate": 0.00013242299027326482,
+      "step": 6502
+    },
+    {
+      "epoch": 0.67,
+      "gumbel_tau": 0.2157,
+      "image_masks_loss": 0.0058,
+      "llm_learning_rate": 1.3234978287898458e-06,
+      "loss": 0.7436,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.00013234978287898457,
+      "step": 6503
+    },
+    {
+      "epoch": 0.67,
+      "gumbel_tau": 0.2156,
+      "image_masks_loss": 0.0212,
+      "llm_learning_rate": 1.3227658844060828e-06,
+      "loss": 0.7862,
+      "output_text_mask_loss": 0.0309,
+      "predictor_learning_rate": 0.00013227658844060827,
+      "step": 6504
+    },
+    {
+      "epoch": 0.67,
+      "gumbel_tau": 0.2156,
+      "image_masks_loss": 0.0183,
+      "llm_learning_rate": 1.322034069661963e-06,
+      "loss": 0.7466,
+      "output_text_mask_loss": 0.0277,
+      "predictor_learning_rate": 0.0001322034069661963,
+      "step": 6505
+    },
+    {
+      "epoch": 0.67,
+      "gumbel_tau": 0.2155,
+      "image_masks_loss": 0.0159,
+      "llm_learning_rate": 1.3213023846380751e-06,
+      "loss": 0.7772,
+      "output_text_mask_loss": 0.0809,
+      "predictor_learning_rate": 0.0001321302384638075,
+      "step": 6506
+    },
+    {
+      "epoch": 0.67,
+      "gumbel_tau": 0.2155,
+      "image_masks_loss": 0.0125,
+      "llm_learning_rate": 1.320570829414993e-06,
+      "loss": 0.6848,
+      "output_text_mask_loss": 0.0058,
+      "predictor_learning_rate": 0.0001320570829414993,
+      "step": 6507
+    },
+    {
+      "epoch": 0.67,
+      "gumbel_tau": 0.2154,
+      "image_masks_loss": 0.0239,
+      "llm_learning_rate": 1.3198394040732792e-06,
+      "loss": 0.7603,
+      "output_text_mask_loss": 0.001,
+      "predictor_learning_rate": 0.00013198394040732792,
+      "step": 6508
+    },
+    {
+      "epoch": 0.67,
+      "gumbel_tau": 0.2154,
+      "image_masks_loss": 0.0139,
+      "llm_learning_rate": 1.3191081086934805e-06,
+      "loss": 0.7638,
+      "output_text_mask_loss": 0.0282,
+      "predictor_learning_rate": 0.00013191081086934804,
+      "step": 6509
+    },
+    {
+      "epoch": 0.67,
+      "gumbel_tau": 0.2153,
+      "image_masks_loss": 0.0035,
+      "llm_learning_rate": 1.3183769433561255e-06,
+      "loss": 0.7362,
+      "output_text_mask_loss": 0.0006,
+      "predictor_learning_rate": 0.00013183769433561254,
+      "step": 6510
+    },
+    {
+      "epoch": 0.67,
+      "gumbel_tau": 0.2153,
+      "image_masks_loss": 0.0095,
+      "llm_learning_rate": 1.3176459081417336e-06,
+      "loss": 0.712,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00013176459081417336,
+      "step": 6511
+    },
+    {
+      "epoch": 0.67,
+      "gumbel_tau": 0.2152,
+      "image_masks_loss": 0.0134,
+      "llm_learning_rate": 1.316915003130807e-06,
+      "loss": 0.7942,
+      "output_text_mask_loss": 0.0315,
+      "predictor_learning_rate": 0.0001316915003130807,
+      "step": 6512
+    },
+    {
+      "epoch": 0.67,
+      "gumbel_tau": 0.2152,
+      "image_masks_loss": 0.0102,
+      "llm_learning_rate": 1.3161842284038365e-06,
+      "loss": 0.6324,
+      "output_text_mask_loss": 0.0204,
+      "predictor_learning_rate": 0.00013161842284038365,
+      "step": 6513
+    },
+    {
+      "epoch": 0.67,
+      "gumbel_tau": 0.2151,
+      "image_masks_loss": 0.0137,
+      "llm_learning_rate": 1.3154535840412957e-06,
+      "loss": 0.6815,
+      "output_text_mask_loss": 0.0162,
+      "predictor_learning_rate": 0.00013154535840412956,
+      "step": 6514
+    },
+    {
+      "epoch": 0.67,
+      "gumbel_tau": 0.2151,
+      "image_masks_loss": 0.0085,
+      "llm_learning_rate": 1.3147230701236447e-06,
+      "loss": 0.7577,
+      "output_text_mask_loss": 0.0012,
+      "predictor_learning_rate": 0.00013147230701236447,
+      "step": 6515
+    },
+    {
+      "epoch": 0.67,
+      "gumbel_tau": 0.215,
+      "image_masks_loss": 0.0143,
+      "llm_learning_rate": 1.313992686731329e-06,
+      "loss": 0.7386,
+      "output_text_mask_loss": 0.0051,
+      "predictor_learning_rate": 0.0001313992686731329,
+      "step": 6516
+    },
+    {
+      "epoch": 0.67,
+      "gumbel_tau": 0.215,
+      "image_masks_loss": 0.0094,
+      "llm_learning_rate": 1.3132624339447805e-06,
+      "loss": 0.7764,
+      "output_text_mask_loss": 0.0078,
+      "predictor_learning_rate": 0.00013132624339447804,
+      "step": 6517
+    },
+    {
+      "epoch": 0.67,
+      "gumbel_tau": 0.2149,
+      "image_masks_loss": 0.0123,
+      "llm_learning_rate": 1.3125323118444162e-06,
+      "loss": 0.7221,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.0001312532311844416,
+      "step": 6518
+    },
+    {
+      "epoch": 0.67,
+      "gumbel_tau": 0.2149,
+      "image_masks_loss": 0.0137,
+      "llm_learning_rate": 1.3118023205106379e-06,
+      "loss": 0.6855,
+      "output_text_mask_loss": 0.0195,
+      "predictor_learning_rate": 0.00013118023205106378,
+      "step": 6519
+    },
+    {
+      "epoch": 0.67,
+      "gumbel_tau": 0.2148,
+      "image_masks_loss": 0.0357,
+      "llm_learning_rate": 1.3110724600238354e-06,
+      "loss": 0.7334,
+      "output_text_mask_loss": 0.0372,
+      "predictor_learning_rate": 0.00013110724600238352,
+      "step": 6520
+    },
+    {
+      "epoch": 0.67,
+      "gumbel_tau": 0.2148,
+      "image_masks_loss": 0.0083,
+      "llm_learning_rate": 1.3103427304643818e-06,
+      "loss": 0.7226,
+      "output_text_mask_loss": 0.0105,
+      "predictor_learning_rate": 0.00013103427304643817,
+      "step": 6521
+    },
+    {
+      "epoch": 0.67,
+      "gumbel_tau": 0.2147,
+      "image_masks_loss": 0.007,
+      "llm_learning_rate": 1.3096131319126365e-06,
+      "loss": 0.649,
+      "output_text_mask_loss": 0.0362,
+      "predictor_learning_rate": 0.00013096131319126364,
+      "step": 6522
+    },
+    {
+      "epoch": 0.67,
+      "gumbel_tau": 0.2147,
+      "image_masks_loss": 0.0201,
+      "llm_learning_rate": 1.308883664448945e-06,
+      "loss": 0.7306,
+      "output_text_mask_loss": 0.0193,
+      "predictor_learning_rate": 0.00013088836644489449,
+      "step": 6523
+    },
+    {
+      "epoch": 0.67,
+      "gumbel_tau": 0.2146,
+      "image_masks_loss": 0.0102,
+      "llm_learning_rate": 1.3081543281536373e-06,
+      "loss": 0.8076,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.00013081543281536374,
+      "step": 6524
+    },
+    {
+      "epoch": 0.67,
+      "gumbel_tau": 0.2146,
+      "image_masks_loss": 0.0099,
+      "llm_learning_rate": 1.3074251231070295e-06,
+      "loss": 0.7589,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.00013074251231070294,
+      "step": 6525
+    },
+    {
+      "epoch": 0.67,
+      "gumbel_tau": 0.2145,
+      "image_masks_loss": 0.0057,
+      "llm_learning_rate": 1.3066960493894226e-06,
+      "loss": 0.8041,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.00013066960493894226,
+      "step": 6526
+    },
+    {
+      "epoch": 0.67,
+      "gumbel_tau": 0.2145,
+      "image_masks_loss": 0.0111,
+      "llm_learning_rate": 1.3059671070811058e-06,
+      "loss": 0.632,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.00013059671070811058,
+      "step": 6527
+    },
+    {
+      "epoch": 0.67,
+      "gumbel_tau": 0.2144,
+      "image_masks_loss": 0.0149,
+      "llm_learning_rate": 1.305238296262352e-06,
+      "loss": 0.6378,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00013052382962623518,
+      "step": 6528
+    },
+    {
+      "epoch": 0.67,
+      "gumbel_tau": 0.2144,
+      "image_masks_loss": 0.0093,
+      "llm_learning_rate": 1.304509617013417e-06,
+      "loss": 0.7895,
+      "output_text_mask_loss": 0.0418,
+      "predictor_learning_rate": 0.0001304509617013417,
+      "step": 6529
+    },
+    {
+      "epoch": 0.67,
+      "gumbel_tau": 0.2143,
+      "image_masks_loss": 0.0074,
+      "llm_learning_rate": 1.303781069414546e-06,
+      "loss": 0.8243,
+      "output_text_mask_loss": 0.034,
+      "predictor_learning_rate": 0.00013037810694145458,
+      "step": 6530
+    },
+    {
+      "epoch": 0.67,
+      "gumbel_tau": 0.2143,
+      "image_masks_loss": 0.0243,
+      "llm_learning_rate": 1.3030526535459675e-06,
+      "loss": 0.7395,
+      "output_text_mask_loss": 0.017,
+      "predictor_learning_rate": 0.00013030526535459674,
+      "step": 6531
+    },
+    {
+      "epoch": 0.67,
+      "gumbel_tau": 0.2142,
+      "image_masks_loss": 0.0193,
+      "llm_learning_rate": 1.3023243694878967e-06,
+      "loss": 0.6592,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.00013023243694878966,
+      "step": 6532
+    },
+    {
+      "epoch": 0.67,
+      "gumbel_tau": 0.2142,
+      "image_masks_loss": 0.0151,
+      "llm_learning_rate": 1.3015962173205345e-06,
+      "loss": 0.7976,
+      "output_text_mask_loss": 0.0032,
+      "predictor_learning_rate": 0.00013015962173205345,
+      "step": 6533
+    },
+    {
+      "epoch": 0.67,
+      "gumbel_tau": 0.2141,
+      "image_masks_loss": 0.0128,
+      "llm_learning_rate": 1.3008681971240664e-06,
+      "loss": 0.7195,
+      "output_text_mask_loss": 0.0296,
+      "predictor_learning_rate": 0.00013008681971240664,
+      "step": 6534
+    },
+    {
+      "epoch": 0.67,
+      "gumbel_tau": 0.2141,
+      "image_masks_loss": 0.0046,
+      "llm_learning_rate": 1.3001403089786635e-06,
+      "loss": 0.7543,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 0.00013001403089786634,
+      "step": 6535
+    },
+    {
+      "epoch": 0.67,
+      "gumbel_tau": 0.214,
+      "image_masks_loss": 0.0033,
+      "llm_learning_rate": 1.2994125529644821e-06,
+      "loss": 0.7413,
+      "output_text_mask_loss": 0.0322,
+      "predictor_learning_rate": 0.00012994125529644822,
+      "step": 6536
+    },
+    {
+      "epoch": 0.67,
+      "gumbel_tau": 0.214,
+      "image_masks_loss": 0.013,
+      "llm_learning_rate": 1.298684929161665e-06,
+      "loss": 0.7784,
+      "output_text_mask_loss": 0.021,
+      "predictor_learning_rate": 0.00012986849291616649,
+      "step": 6537
+    },
+    {
+      "epoch": 0.67,
+      "gumbel_tau": 0.2139,
+      "image_masks_loss": 0.0076,
+      "llm_learning_rate": 1.2979574376503387e-06,
+      "loss": 0.7057,
+      "output_text_mask_loss": 0.0014,
+      "predictor_learning_rate": 0.00012979574376503387,
+      "step": 6538
+    },
+    {
+      "epoch": 0.67,
+      "gumbel_tau": 0.2139,
+      "image_masks_loss": 0.011,
+      "llm_learning_rate": 1.2972300785106161e-06,
+      "loss": 0.7454,
+      "output_text_mask_loss": 0.0141,
+      "predictor_learning_rate": 0.00012972300785106162,
+      "step": 6539
+    },
+    {
+      "epoch": 0.67,
+      "gumbel_tau": 0.2138,
+      "image_masks_loss": 0.0044,
+      "llm_learning_rate": 1.2965028518225976e-06,
+      "loss": 0.7212,
+      "output_text_mask_loss": 0.025,
+      "predictor_learning_rate": 0.00012965028518225973,
+      "step": 6540
+    },
+    {
+      "epoch": 0.67,
+      "gumbel_tau": 0.2138,
+      "image_masks_loss": 0.0026,
+      "llm_learning_rate": 1.2957757576663655e-06,
+      "loss": 0.8,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.00012957757576663654,
+      "step": 6541
+    },
+    {
+      "epoch": 0.67,
+      "gumbel_tau": 0.2137,
+      "image_masks_loss": 0.0063,
+      "llm_learning_rate": 1.2950487961219891e-06,
+      "loss": 0.695,
+      "output_text_mask_loss": 0.0444,
+      "predictor_learning_rate": 0.0001295048796121989,
+      "step": 6542
+    },
+    {
+      "epoch": 0.67,
+      "gumbel_tau": 0.2137,
+      "image_masks_loss": 0.0091,
+      "llm_learning_rate": 1.2943219672695233e-06,
+      "loss": 0.7578,
+      "output_text_mask_loss": 0.0454,
+      "predictor_learning_rate": 0.00012943219672695235,
+      "step": 6543
+    },
+    {
+      "epoch": 0.67,
+      "gumbel_tau": 0.2136,
+      "image_masks_loss": 0.0087,
+      "llm_learning_rate": 1.293595271189008e-06,
+      "loss": 0.8084,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.0001293595271189008,
+      "step": 6544
+    },
+    {
+      "epoch": 0.67,
+      "gumbel_tau": 0.2136,
+      "image_masks_loss": 0.0034,
+      "llm_learning_rate": 1.2928687079604687e-06,
+      "loss": 0.7096,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.00012928687079604686,
+      "step": 6545
+    },
+    {
+      "epoch": 0.67,
+      "gumbel_tau": 0.2135,
+      "image_masks_loss": 0.0034,
+      "llm_learning_rate": 1.2921422776639159e-06,
+      "loss": 0.785,
+      "output_text_mask_loss": 0.0443,
+      "predictor_learning_rate": 0.00012921422776639158,
+      "step": 6546
+    },
+    {
+      "epoch": 0.67,
+      "gumbel_tau": 0.2135,
+      "image_masks_loss": 0.0131,
+      "llm_learning_rate": 1.291415980379346e-06,
+      "loss": 0.7678,
+      "output_text_mask_loss": 0.0139,
+      "predictor_learning_rate": 0.0001291415980379346,
+      "step": 6547
+    },
+    {
+      "epoch": 0.67,
+      "gumbel_tau": 0.2134,
+      "image_masks_loss": 0.0129,
+      "llm_learning_rate": 1.2906898161867404e-06,
+      "loss": 0.7056,
+      "output_text_mask_loss": 0.0143,
+      "predictor_learning_rate": 0.00012906898161867402,
+      "step": 6548
+    },
+    {
+      "epoch": 0.67,
+      "gumbel_tau": 0.2134,
+      "image_masks_loss": 0.0092,
+      "llm_learning_rate": 1.289963785166066e-06,
+      "loss": 0.5749,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.00012899637851660658,
+      "step": 6549
+    },
+    {
+      "epoch": 0.67,
+      "gumbel_tau": 0.2133,
+      "image_masks_loss": 0.0052,
+      "llm_learning_rate": 1.2892378873972743e-06,
+      "loss": 0.7206,
+      "output_text_mask_loss": 0.0072,
+      "predictor_learning_rate": 0.00012892378873972742,
+      "step": 6550
+    },
+    {
+      "epoch": 0.67,
+      "gumbel_tau": 0.2133,
+      "image_masks_loss": 0.0221,
+      "llm_learning_rate": 1.288512122960304e-06,
+      "loss": 0.6654,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0001288512122960304,
+      "step": 6551
+    },
+    {
+      "epoch": 0.67,
+      "gumbel_tau": 0.2132,
+      "image_masks_loss": 0.0063,
+      "llm_learning_rate": 1.2877864919350757e-06,
+      "loss": 0.7858,
+      "output_text_mask_loss": 0.0233,
+      "predictor_learning_rate": 0.00012877864919350756,
+      "step": 6552
+    },
+    {
+      "epoch": 0.67,
+      "gumbel_tau": 0.2132,
+      "image_masks_loss": 0.0083,
+      "llm_learning_rate": 1.2870609944015003e-06,
+      "loss": 0.7631,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00012870609944015,
+      "step": 6553
+    },
+    {
+      "epoch": 0.67,
+      "gumbel_tau": 0.2131,
+      "image_masks_loss": 0.0092,
+      "llm_learning_rate": 1.2863356304394696e-06,
+      "loss": 0.7208,
+      "output_text_mask_loss": 0.0058,
+      "predictor_learning_rate": 0.00012863356304394696,
+      "step": 6554
+    },
+    {
+      "epoch": 0.67,
+      "gumbel_tau": 0.2131,
+      "image_masks_loss": 0.0157,
+      "llm_learning_rate": 1.285610400128863e-06,
+      "loss": 0.745,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0001285610400128863,
+      "step": 6555
+    },
+    {
+      "epoch": 0.67,
+      "gumbel_tau": 0.213,
+      "image_masks_loss": 0.0155,
+      "llm_learning_rate": 1.284885303549544e-06,
+      "loss": 0.7087,
+      "output_text_mask_loss": 0.069,
+      "predictor_learning_rate": 0.0001284885303549544,
+      "step": 6556
+    },
+    {
+      "epoch": 0.67,
+      "gumbel_tau": 0.213,
+      "image_masks_loss": 0.0134,
+      "llm_learning_rate": 1.284160340781362e-06,
+      "loss": 0.7685,
+      "output_text_mask_loss": 0.0052,
+      "predictor_learning_rate": 0.00012841603407813619,
+      "step": 6557
+    },
+    {
+      "epoch": 0.67,
+      "gumbel_tau": 0.2129,
+      "image_masks_loss": 0.0089,
+      "llm_learning_rate": 1.2834355119041507e-06,
+      "loss": 0.7566,
+      "output_text_mask_loss": 0.0126,
+      "predictor_learning_rate": 0.0001283435511904151,
+      "step": 6558
+    },
+    {
+      "epoch": 0.67,
+      "gumbel_tau": 0.2129,
+      "image_masks_loss": 0.0053,
+      "llm_learning_rate": 1.2827108169977304e-06,
+      "loss": 0.7551,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00012827108169977304,
+      "step": 6559
+    },
+    {
+      "epoch": 0.67,
+      "gumbel_tau": 0.2128,
+      "image_masks_loss": 0.0132,
+      "llm_learning_rate": 1.2819862561419071e-06,
+      "loss": 0.7677,
+      "output_text_mask_loss": 0.0096,
+      "predictor_learning_rate": 0.0001281986256141907,
+      "step": 6560
+    },
+    {
+      "epoch": 0.67,
+      "gumbel_tau": 0.2127,
+      "image_masks_loss": 0.0027,
+      "llm_learning_rate": 1.2812618294164703e-06,
+      "loss": 0.6709,
+      "output_text_mask_loss": 0.0127,
+      "predictor_learning_rate": 0.00012812618294164702,
+      "step": 6561
+    },
+    {
+      "epoch": 0.67,
+      "gumbel_tau": 0.2127,
+      "image_masks_loss": 0.0028,
+      "llm_learning_rate": 1.2805375369011957e-06,
+      "loss": 0.7069,
+      "output_text_mask_loss": 0.0081,
+      "predictor_learning_rate": 0.00012805375369011958,
+      "step": 6562
+    },
+    {
+      "epoch": 0.67,
+      "gumbel_tau": 0.2126,
+      "image_masks_loss": 0.009,
+      "llm_learning_rate": 1.2798133786758434e-06,
+      "loss": 0.7053,
+      "output_text_mask_loss": 0.0032,
+      "predictor_learning_rate": 0.00012798133786758435,
+      "step": 6563
+    },
+    {
+      "epoch": 0.67,
+      "gumbel_tau": 0.2126,
+      "image_masks_loss": 0.0094,
+      "llm_learning_rate": 1.2790893548201599e-06,
+      "loss": 0.754,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.000127908935482016,
+      "step": 6564
+    },
+    {
+      "epoch": 0.67,
+      "gumbel_tau": 0.2125,
+      "image_masks_loss": 0.0066,
+      "llm_learning_rate": 1.278365465413876e-06,
+      "loss": 0.6745,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0001278365465413876,
+      "step": 6565
+    },
+    {
+      "epoch": 0.67,
+      "gumbel_tau": 0.2125,
+      "image_masks_loss": 0.0158,
+      "llm_learning_rate": 1.2776417105367084e-06,
+      "loss": 0.6609,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00012776417105367084,
+      "step": 6566
+    },
+    {
+      "epoch": 0.67,
+      "gumbel_tau": 0.2124,
+      "image_masks_loss": 0.0056,
+      "llm_learning_rate": 1.276918090268358e-06,
+      "loss": 0.7758,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.0001276918090268358,
+      "step": 6567
+    },
+    {
+      "epoch": 0.67,
+      "gumbel_tau": 0.2124,
+      "image_masks_loss": 0.0141,
+      "llm_learning_rate": 1.276194604688512e-06,
+      "loss": 0.751,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0001276194604688512,
+      "step": 6568
+    },
+    {
+      "epoch": 0.67,
+      "gumbel_tau": 0.2123,
+      "image_masks_loss": 0.005,
+      "llm_learning_rate": 1.275471253876842e-06,
+      "loss": 0.727,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.0001275471253876842,
+      "step": 6569
+    },
+    {
+      "epoch": 0.67,
+      "gumbel_tau": 0.2123,
+      "image_masks_loss": 0.0055,
+      "llm_learning_rate": 1.274748037913005e-06,
+      "loss": 0.7988,
+      "output_text_mask_loss": 0.0063,
+      "predictor_learning_rate": 0.0001274748037913005,
+      "step": 6570
+    },
+    {
+      "epoch": 0.67,
+      "gumbel_tau": 0.2122,
+      "image_masks_loss": 0.0203,
+      "llm_learning_rate": 1.2740249568766424e-06,
+      "loss": 0.7243,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 0.00012740249568766424,
+      "step": 6571
+    },
+    {
+      "epoch": 0.67,
+      "gumbel_tau": 0.2122,
+      "image_masks_loss": 0.0105,
+      "llm_learning_rate": 1.2733020108473832e-06,
+      "loss": 0.7804,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.0001273302010847383,
+      "step": 6572
+    },
+    {
+      "epoch": 0.67,
+      "gumbel_tau": 0.2121,
+      "image_masks_loss": 0.0172,
+      "llm_learning_rate": 1.272579199904839e-06,
+      "loss": 0.7224,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.0001272579199904839,
+      "step": 6573
+    },
+    {
+      "epoch": 0.67,
+      "gumbel_tau": 0.2121,
+      "image_masks_loss": 0.0053,
+      "llm_learning_rate": 1.271856524128607e-06,
+      "loss": 0.7342,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.0001271856524128607,
+      "step": 6574
+    },
+    {
+      "epoch": 0.67,
+      "gumbel_tau": 0.212,
+      "image_masks_loss": 0.0072,
+      "llm_learning_rate": 1.2711339835982705e-06,
+      "loss": 0.6219,
+      "output_text_mask_loss": 0.0012,
+      "predictor_learning_rate": 0.00012711339835982705,
+      "step": 6575
+    },
+    {
+      "epoch": 0.67,
+      "gumbel_tau": 0.212,
+      "image_masks_loss": 0.0124,
+      "llm_learning_rate": 1.2704115783933968e-06,
+      "loss": 0.763,
+      "output_text_mask_loss": 0.0378,
+      "predictor_learning_rate": 0.00012704115783933968,
+      "step": 6576
+    },
+    {
+      "epoch": 0.67,
+      "gumbel_tau": 0.2119,
+      "image_masks_loss": 0.0101,
+      "llm_learning_rate": 1.2696893085935385e-06,
+      "loss": 0.5764,
+      "output_text_mask_loss": 0.0275,
+      "predictor_learning_rate": 0.00012696893085935385,
+      "step": 6577
+    },
+    {
+      "epoch": 0.67,
+      "gumbel_tau": 0.2119,
+      "image_masks_loss": 0.0026,
+      "llm_learning_rate": 1.2689671742782336e-06,
+      "loss": 0.6692,
+      "output_text_mask_loss": 0.0115,
+      "predictor_learning_rate": 0.00012689671742782335,
+      "step": 6578
+    },
+    {
+      "epoch": 0.67,
+      "gumbel_tau": 0.2118,
+      "image_masks_loss": 0.0097,
+      "llm_learning_rate": 1.2682451755270062e-06,
+      "loss": 0.7205,
+      "output_text_mask_loss": 0.002,
+      "predictor_learning_rate": 0.0001268245175527006,
+      "step": 6579
+    },
+    {
+      "epoch": 0.67,
+      "gumbel_tau": 0.2118,
+      "image_masks_loss": 0.0102,
+      "llm_learning_rate": 1.2675233124193637e-06,
+      "loss": 0.7529,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.00012675233124193635,
+      "step": 6580
+    },
+    {
+      "epoch": 0.67,
+      "gumbel_tau": 0.2117,
+      "image_masks_loss": 0.0257,
+      "llm_learning_rate": 1.2668015850347993e-06,
+      "loss": 0.9015,
+      "output_text_mask_loss": 0.0198,
+      "predictor_learning_rate": 0.00012668015850347993,
+      "step": 6581
+    },
+    {
+      "epoch": 0.67,
+      "gumbel_tau": 0.2117,
+      "image_masks_loss": 0.0117,
+      "llm_learning_rate": 1.266079993452791e-06,
+      "loss": 0.6741,
+      "output_text_mask_loss": 0.0039,
+      "predictor_learning_rate": 0.00012660799934527907,
+      "step": 6582
+    },
+    {
+      "epoch": 0.67,
+      "gumbel_tau": 0.2116,
+      "image_masks_loss": 0.0112,
+      "llm_learning_rate": 1.2653585377528033e-06,
+      "loss": 0.7418,
+      "output_text_mask_loss": 0.0039,
+      "predictor_learning_rate": 0.0001265358537752803,
+      "step": 6583
+    },
+    {
+      "epoch": 0.67,
+      "gumbel_tau": 0.2116,
+      "image_masks_loss": 0.0178,
+      "llm_learning_rate": 1.2646372180142819e-06,
+      "loss": 0.6972,
+      "output_text_mask_loss": 0.0038,
+      "predictor_learning_rate": 0.00012646372180142817,
+      "step": 6584
+    },
+    {
+      "epoch": 0.67,
+      "gumbel_tau": 0.2115,
+      "image_masks_loss": 0.0142,
+      "llm_learning_rate": 1.2639160343166622e-06,
+      "loss": 0.7601,
+      "output_text_mask_loss": 0.0229,
+      "predictor_learning_rate": 0.00012639160343166624,
+      "step": 6585
+    },
+    {
+      "epoch": 0.67,
+      "gumbel_tau": 0.2115,
+      "image_masks_loss": 0.0102,
+      "llm_learning_rate": 1.2631949867393622e-06,
+      "loss": 0.8408,
+      "output_text_mask_loss": 0.0259,
+      "predictor_learning_rate": 0.00012631949867393622,
+      "step": 6586
+    },
+    {
+      "epoch": 0.67,
+      "gumbel_tau": 0.2114,
+      "image_masks_loss": 0.021,
+      "llm_learning_rate": 1.2624740753617854e-06,
+      "loss": 0.6681,
+      "output_text_mask_loss": 0.0305,
+      "predictor_learning_rate": 0.00012624740753617854,
+      "step": 6587
+    },
+    {
+      "epoch": 0.68,
+      "gumbel_tau": 0.2114,
+      "image_masks_loss": 0.0088,
+      "llm_learning_rate": 1.2617533002633198e-06,
+      "loss": 0.6427,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.00012617533002633196,
+      "step": 6588
+    },
+    {
+      "epoch": 0.68,
+      "gumbel_tau": 0.2113,
+      "image_masks_loss": 0.0108,
+      "llm_learning_rate": 1.2610326615233388e-06,
+      "loss": 0.7484,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00012610326615233386,
+      "step": 6589
+    },
+    {
+      "epoch": 0.68,
+      "gumbel_tau": 0.2113,
+      "image_masks_loss": 0.0046,
+      "llm_learning_rate": 1.2603121592212005e-06,
+      "loss": 0.7132,
+      "output_text_mask_loss": 0.0118,
+      "predictor_learning_rate": 0.00012603121592212005,
+      "step": 6590
+    },
+    {
+      "epoch": 0.68,
+      "gumbel_tau": 0.2112,
+      "image_masks_loss": 0.0038,
+      "llm_learning_rate": 1.2595917934362478e-06,
+      "loss": 0.7908,
+      "output_text_mask_loss": 0.0622,
+      "predictor_learning_rate": 0.00012595917934362478,
+      "step": 6591
+    },
+    {
+      "epoch": 0.68,
+      "gumbel_tau": 0.2112,
+      "image_masks_loss": 0.0075,
+      "llm_learning_rate": 1.2588715642478108e-06,
+      "loss": 0.617,
+      "output_text_mask_loss": 0.0058,
+      "predictor_learning_rate": 0.00012588715642478107,
+      "step": 6592
+    },
+    {
+      "epoch": 0.68,
+      "gumbel_tau": 0.2111,
+      "image_masks_loss": 0.0075,
+      "llm_learning_rate": 1.2581514717352016e-06,
+      "loss": 0.773,
+      "output_text_mask_loss": 0.0077,
+      "predictor_learning_rate": 0.00012581514717352016,
+      "step": 6593
+    },
+    {
+      "epoch": 0.68,
+      "gumbel_tau": 0.2111,
+      "image_masks_loss": 0.012,
+      "llm_learning_rate": 1.2574315159777184e-06,
+      "loss": 0.7128,
+      "output_text_mask_loss": 0.0172,
+      "predictor_learning_rate": 0.00012574315159777184,
+      "step": 6594
+    },
+    {
+      "epoch": 0.68,
+      "gumbel_tau": 0.211,
+      "image_masks_loss": 0.0039,
+      "llm_learning_rate": 1.2567116970546444e-06,
+      "loss": 0.7692,
+      "output_text_mask_loss": 0.0332,
+      "predictor_learning_rate": 0.00012567116970546444,
+      "step": 6595
+    },
+    {
+      "epoch": 0.68,
+      "gumbel_tau": 0.211,
+      "image_masks_loss": 0.0074,
+      "llm_learning_rate": 1.2559920150452475e-06,
+      "loss": 0.7011,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 0.00012559920150452476,
+      "step": 6596
+    },
+    {
+      "epoch": 0.68,
+      "gumbel_tau": 0.211,
+      "image_masks_loss": 0.0207,
+      "llm_learning_rate": 1.255272470028781e-06,
+      "loss": 0.6725,
+      "output_text_mask_loss": 0.0137,
+      "predictor_learning_rate": 0.0001255272470028781,
+      "step": 6597
+    },
+    {
+      "epoch": 0.68,
+      "gumbel_tau": 0.2109,
+      "image_masks_loss": 0.0093,
+      "llm_learning_rate": 1.2545530620844814e-06,
+      "loss": 0.6833,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.00012545530620844815,
+      "step": 6598
+    },
+    {
+      "epoch": 0.68,
+      "gumbel_tau": 0.2109,
+      "image_masks_loss": 0.019,
+      "llm_learning_rate": 1.253833791291574e-06,
+      "loss": 0.6246,
+      "output_text_mask_loss": 0.0193,
+      "predictor_learning_rate": 0.0001253833791291574,
+      "step": 6599
+    },
+    {
+      "epoch": 0.68,
+      "gumbel_tau": 0.2108,
+      "image_masks_loss": 0.007,
+      "llm_learning_rate": 1.2531146577292651e-06,
+      "loss": 0.7028,
+      "output_text_mask_loss": 0.0051,
+      "predictor_learning_rate": 0.00012531146577292651,
+      "step": 6600
+    },
+    {
+      "epoch": 0.68,
+      "gumbel_tau": 0.2108,
+      "image_masks_loss": 0.003,
+      "llm_learning_rate": 1.2523956614767473e-06,
+      "loss": 0.8003,
+      "output_text_mask_loss": 0.0223,
+      "predictor_learning_rate": 0.00012523956614767475,
+      "step": 6601
+    },
+    {
+      "epoch": 0.68,
+      "gumbel_tau": 0.2107,
+      "image_masks_loss": 0.0048,
+      "llm_learning_rate": 1.2516768026131993e-06,
+      "loss": 0.7607,
+      "output_text_mask_loss": 0.0019,
+      "predictor_learning_rate": 0.0001251676802613199,
+      "step": 6602
+    },
+    {
+      "epoch": 0.68,
+      "gumbel_tau": 0.2107,
+      "image_masks_loss": 0.0102,
+      "llm_learning_rate": 1.2509580812177805e-06,
+      "loss": 0.7576,
+      "output_text_mask_loss": 0.0039,
+      "predictor_learning_rate": 0.00012509580812177803,
+      "step": 6603
+    },
+    {
+      "epoch": 0.68,
+      "gumbel_tau": 0.2106,
+      "image_masks_loss": 0.0157,
+      "llm_learning_rate": 1.2502394973696409e-06,
+      "loss": 0.7287,
+      "output_text_mask_loss": 0.0038,
+      "predictor_learning_rate": 0.00012502394973696408,
+      "step": 6604
+    },
+    {
+      "epoch": 0.68,
+      "gumbel_tau": 0.2106,
+      "image_masks_loss": 0.0189,
+      "llm_learning_rate": 1.2495210511479113e-06,
+      "loss": 0.7456,
+      "output_text_mask_loss": 0.0246,
+      "predictor_learning_rate": 0.00012495210511479114,
+      "step": 6605
+    },
+    {
+      "epoch": 0.68,
+      "gumbel_tau": 0.2105,
+      "image_masks_loss": 0.0286,
+      "llm_learning_rate": 1.248802742631709e-06,
+      "loss": 0.8199,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.0001248802742631709,
+      "step": 6606
+    },
+    {
+      "epoch": 0.68,
+      "gumbel_tau": 0.2105,
+      "image_masks_loss": 0.012,
+      "llm_learning_rate": 1.2480845719001358e-06,
+      "loss": 0.7359,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.00012480845719001356,
+      "step": 6607
+    },
+    {
+      "epoch": 0.68,
+      "gumbel_tau": 0.2104,
+      "image_masks_loss": 0.0101,
+      "llm_learning_rate": 1.2473665390322779e-06,
+      "loss": 0.702,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00012473665390322778,
+      "step": 6608
+    },
+    {
+      "epoch": 0.68,
+      "gumbel_tau": 0.2104,
+      "image_masks_loss": 0.0104,
+      "llm_learning_rate": 1.2466486441072068e-06,
+      "loss": 0.7576,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.00012466486441072068,
+      "step": 6609
+    },
+    {
+      "epoch": 0.68,
+      "gumbel_tau": 0.2103,
+      "image_masks_loss": 0.0078,
+      "llm_learning_rate": 1.245930887203978e-06,
+      "loss": 0.7743,
+      "output_text_mask_loss": 0.038,
+      "predictor_learning_rate": 0.0001245930887203978,
+      "step": 6610
+    },
+    {
+      "epoch": 0.68,
+      "gumbel_tau": 0.2103,
+      "image_masks_loss": 0.0226,
+      "llm_learning_rate": 1.2452132684016338e-06,
+      "loss": 0.8156,
+      "output_text_mask_loss": 0.0185,
+      "predictor_learning_rate": 0.00012452132684016338,
+      "step": 6611
+    },
+    {
+      "epoch": 0.68,
+      "gumbel_tau": 0.2102,
+      "image_masks_loss": 0.0131,
+      "llm_learning_rate": 1.2444957877791996e-06,
+      "loss": 0.7598,
+      "output_text_mask_loss": 0.0572,
+      "predictor_learning_rate": 0.00012444957877791996,
+      "step": 6612
+    },
+    {
+      "epoch": 0.68,
+      "gumbel_tau": 0.2102,
+      "image_masks_loss": 0.01,
+      "llm_learning_rate": 1.243778445415686e-06,
+      "loss": 0.7902,
+      "output_text_mask_loss": 0.0098,
+      "predictor_learning_rate": 0.00012437784454156857,
+      "step": 6613
+    },
+    {
+      "epoch": 0.68,
+      "gumbel_tau": 0.2101,
+      "image_masks_loss": 0.0298,
+      "llm_learning_rate": 1.243061241390088e-06,
+      "loss": 0.709,
+      "output_text_mask_loss": 0.0313,
+      "predictor_learning_rate": 0.0001243061241390088,
+      "step": 6614
+    },
+    {
+      "epoch": 0.68,
+      "gumbel_tau": 0.2101,
+      "image_masks_loss": 0.0183,
+      "llm_learning_rate": 1.242344175781386e-06,
+      "loss": 0.7613,
+      "output_text_mask_loss": 0.0019,
+      "predictor_learning_rate": 0.00012423441757813857,
+      "step": 6615
+    },
+    {
+      "epoch": 0.68,
+      "gumbel_tau": 0.21,
+      "image_masks_loss": 0.022,
+      "llm_learning_rate": 1.2416272486685443e-06,
+      "loss": 0.7303,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.00012416272486685443,
+      "step": 6616
+    },
+    {
+      "epoch": 0.68,
+      "gumbel_tau": 0.21,
+      "image_masks_loss": 0.012,
+      "llm_learning_rate": 1.240910460130512e-06,
+      "loss": 0.7503,
+      "output_text_mask_loss": 0.0626,
+      "predictor_learning_rate": 0.00012409104601305117,
+      "step": 6617
+    },
+    {
+      "epoch": 0.68,
+      "gumbel_tau": 0.2099,
+      "image_masks_loss": 0.0334,
+      "llm_learning_rate": 1.2401938102462255e-06,
+      "loss": 0.8186,
+      "output_text_mask_loss": 0.0706,
+      "predictor_learning_rate": 0.00012401938102462253,
+      "step": 6618
+    },
+    {
+      "epoch": 0.68,
+      "gumbel_tau": 0.2099,
+      "image_masks_loss": 0.01,
+      "llm_learning_rate": 1.2394772990946024e-06,
+      "loss": 0.7471,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 0.00012394772990946025,
+      "step": 6619
+    },
+    {
+      "epoch": 0.68,
+      "gumbel_tau": 0.2098,
+      "image_masks_loss": 0.016,
+      "llm_learning_rate": 1.238760926754548e-06,
+      "loss": 0.7934,
+      "output_text_mask_loss": 0.0347,
+      "predictor_learning_rate": 0.0001238760926754548,
+      "step": 6620
+    },
+    {
+      "epoch": 0.68,
+      "gumbel_tau": 0.2098,
+      "image_masks_loss": 0.0158,
+      "llm_learning_rate": 1.2380446933049487e-06,
+      "loss": 0.7837,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00012380446933049486,
+      "step": 6621
+    },
+    {
+      "epoch": 0.68,
+      "gumbel_tau": 0.2097,
+      "image_masks_loss": 0.0201,
+      "llm_learning_rate": 1.2373285988246787e-06,
+      "loss": 0.723,
+      "output_text_mask_loss": 0.0071,
+      "predictor_learning_rate": 0.00012373285988246784,
+      "step": 6622
+    },
+    {
+      "epoch": 0.68,
+      "gumbel_tau": 0.2097,
+      "image_masks_loss": 0.0229,
+      "llm_learning_rate": 1.2366126433925949e-06,
+      "loss": 0.7266,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.00012366126433925947,
+      "step": 6623
+    },
+    {
+      "epoch": 0.68,
+      "gumbel_tau": 0.2096,
+      "image_masks_loss": 0.01,
+      "llm_learning_rate": 1.2358968270875418e-06,
+      "loss": 0.7685,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00012358968270875418,
+      "step": 6624
+    },
+    {
+      "epoch": 0.68,
+      "gumbel_tau": 0.2096,
+      "image_masks_loss": 0.0338,
+      "llm_learning_rate": 1.235181149988346e-06,
+      "loss": 0.7588,
+      "output_text_mask_loss": 0.0061,
+      "predictor_learning_rate": 0.0001235181149988346,
+      "step": 6625
+    },
+    {
+      "epoch": 0.68,
+      "gumbel_tau": 0.2095,
+      "image_masks_loss": 0.0254,
+      "llm_learning_rate": 1.2344656121738188e-06,
+      "loss": 0.7047,
+      "output_text_mask_loss": 0.0035,
+      "predictor_learning_rate": 0.00012344656121738188,
+      "step": 6626
+    },
+    {
+      "epoch": 0.68,
+      "gumbel_tau": 0.2095,
+      "image_masks_loss": 0.01,
+      "llm_learning_rate": 1.2337502137227574e-06,
+      "loss": 0.6881,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.00012337502137227575,
+      "step": 6627
+    },
+    {
+      "epoch": 0.68,
+      "gumbel_tau": 0.2094,
+      "image_masks_loss": 0.0209,
+      "llm_learning_rate": 1.2330349547139434e-06,
+      "loss": 0.6544,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00012330349547139434,
+      "step": 6628
+    },
+    {
+      "epoch": 0.68,
+      "gumbel_tau": 0.2094,
+      "image_masks_loss": 0.0286,
+      "llm_learning_rate": 1.2323198352261418e-06,
+      "loss": 0.6896,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.00012323198352261417,
+      "step": 6629
+    },
+    {
+      "epoch": 0.68,
+      "gumbel_tau": 0.2093,
+      "image_masks_loss": 0.0063,
+      "llm_learning_rate": 1.2316048553381027e-06,
+      "loss": 0.7366,
+      "output_text_mask_loss": 0.0025,
+      "predictor_learning_rate": 0.00012316048553381027,
+      "step": 6630
+    },
+    {
+      "epoch": 0.68,
+      "gumbel_tau": 0.2093,
+      "image_masks_loss": 0.012,
+      "llm_learning_rate": 1.230890015128563e-06,
+      "loss": 0.7842,
+      "output_text_mask_loss": 0.0032,
+      "predictor_learning_rate": 0.0001230890015128563,
+      "step": 6631
+    },
+    {
+      "epoch": 0.68,
+      "gumbel_tau": 0.2092,
+      "image_masks_loss": 0.0271,
+      "llm_learning_rate": 1.2301753146762418e-06,
+      "loss": 0.7631,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.00012301753146762417,
+      "step": 6632
+    },
+    {
+      "epoch": 0.68,
+      "gumbel_tau": 0.2092,
+      "image_masks_loss": 0.0099,
+      "llm_learning_rate": 1.2294607540598435e-06,
+      "loss": 0.7602,
+      "output_text_mask_loss": 0.0055,
+      "predictor_learning_rate": 0.00012294607540598435,
+      "step": 6633
+    },
+    {
+      "epoch": 0.68,
+      "gumbel_tau": 0.2091,
+      "image_masks_loss": 0.0353,
+      "llm_learning_rate": 1.2287463333580567e-06,
+      "loss": 0.5617,
+      "output_text_mask_loss": 0.0248,
+      "predictor_learning_rate": 0.00012287463333580568,
+      "step": 6634
+    },
+    {
+      "epoch": 0.68,
+      "gumbel_tau": 0.2091,
+      "image_masks_loss": 0.0299,
+      "llm_learning_rate": 1.2280320526495556e-06,
+      "loss": 0.763,
+      "output_text_mask_loss": 0.009,
+      "predictor_learning_rate": 0.00012280320526495556,
+      "step": 6635
+    },
+    {
+      "epoch": 0.68,
+      "gumbel_tau": 0.209,
+      "image_masks_loss": 0.0052,
+      "llm_learning_rate": 1.2273179120129978e-06,
+      "loss": 0.7395,
+      "output_text_mask_loss": 0.0549,
+      "predictor_learning_rate": 0.00012273179120129976,
+      "step": 6636
+    },
+    {
+      "epoch": 0.68,
+      "gumbel_tau": 0.209,
+      "image_masks_loss": 0.0168,
+      "llm_learning_rate": 1.2266039115270253e-06,
+      "loss": 0.7406,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00012266039115270253,
+      "step": 6637
+    },
+    {
+      "epoch": 0.68,
+      "gumbel_tau": 0.2089,
+      "image_masks_loss": 0.0148,
+      "llm_learning_rate": 1.2258900512702677e-06,
+      "loss": 0.6576,
+      "output_text_mask_loss": 0.0056,
+      "predictor_learning_rate": 0.00012258900512702675,
+      "step": 6638
+    },
+    {
+      "epoch": 0.68,
+      "gumbel_tau": 0.2089,
+      "image_masks_loss": 0.0029,
+      "llm_learning_rate": 1.2251763313213363e-06,
+      "loss": 0.7254,
+      "output_text_mask_loss": 0.0223,
+      "predictor_learning_rate": 0.00012251763313213363,
+      "step": 6639
+    },
+    {
+      "epoch": 0.68,
+      "gumbel_tau": 0.2088,
+      "image_masks_loss": 0.0213,
+      "llm_learning_rate": 1.224462751758826e-06,
+      "loss": 0.7317,
+      "output_text_mask_loss": 0.0244,
+      "predictor_learning_rate": 0.0001224462751758826,
+      "step": 6640
+    },
+    {
+      "epoch": 0.68,
+      "gumbel_tau": 0.2088,
+      "image_masks_loss": 0.0136,
+      "llm_learning_rate": 1.2237493126613186e-06,
+      "loss": 0.7327,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00012237493126613185,
+      "step": 6641
+    },
+    {
+      "epoch": 0.68,
+      "gumbel_tau": 0.2087,
+      "image_masks_loss": 0.0097,
+      "llm_learning_rate": 1.2230360141073788e-06,
+      "loss": 0.6913,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.00012230360141073787,
+      "step": 6642
+    },
+    {
+      "epoch": 0.68,
+      "gumbel_tau": 0.2087,
+      "image_masks_loss": 0.0073,
+      "llm_learning_rate": 1.2223228561755584e-06,
+      "loss": 0.6861,
+      "output_text_mask_loss": 0.0488,
+      "predictor_learning_rate": 0.00012223228561755584,
+      "step": 6643
+    },
+    {
+      "epoch": 0.68,
+      "gumbel_tau": 0.2086,
+      "image_masks_loss": 0.0119,
+      "llm_learning_rate": 1.221609838944391e-06,
+      "loss": 0.7044,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00012216098389443908,
+      "step": 6644
+    },
+    {
+      "epoch": 0.68,
+      "gumbel_tau": 0.2086,
+      "image_masks_loss": 0.0208,
+      "llm_learning_rate": 1.2208969624923955e-06,
+      "loss": 0.7298,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00012208969624923955,
+      "step": 6645
+    },
+    {
+      "epoch": 0.68,
+      "gumbel_tau": 0.2085,
+      "image_masks_loss": 0.0126,
+      "llm_learning_rate": 1.2201842268980757e-06,
+      "loss": 0.7433,
+      "output_text_mask_loss": 0.0085,
+      "predictor_learning_rate": 0.00012201842268980756,
+      "step": 6646
+    },
+    {
+      "epoch": 0.68,
+      "gumbel_tau": 0.2085,
+      "image_masks_loss": 0.0091,
+      "llm_learning_rate": 1.2194716322399195e-06,
+      "loss": 0.8132,
+      "output_text_mask_loss": 0.053,
+      "predictor_learning_rate": 0.00012194716322399193,
+      "step": 6647
+    },
+    {
+      "epoch": 0.68,
+      "gumbel_tau": 0.2084,
+      "image_masks_loss": 0.0062,
+      "llm_learning_rate": 1.2187591785963993e-06,
+      "loss": 0.7777,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00012187591785963992,
+      "step": 6648
+    },
+    {
+      "epoch": 0.68,
+      "gumbel_tau": 0.2084,
+      "image_masks_loss": 0.0113,
+      "llm_learning_rate": 1.2180468660459712e-06,
+      "loss": 0.7872,
+      "output_text_mask_loss": 0.0012,
+      "predictor_learning_rate": 0.00012180468660459713,
+      "step": 6649
+    },
+    {
+      "epoch": 0.68,
+      "gumbel_tau": 0.2083,
+      "image_masks_loss": 0.008,
+      "llm_learning_rate": 1.217334694667079e-06,
+      "loss": 0.8098,
+      "output_text_mask_loss": 0.0107,
+      "predictor_learning_rate": 0.00012173346946670788,
+      "step": 6650
+    },
+    {
+      "epoch": 0.68,
+      "gumbel_tau": 0.2083,
+      "image_masks_loss": 0.0171,
+      "llm_learning_rate": 1.216622664538147e-06,
+      "loss": 0.8028,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 0.00012166226645381468,
+      "step": 6651
+    },
+    {
+      "epoch": 0.68,
+      "gumbel_tau": 0.2082,
+      "image_masks_loss": 0.0244,
+      "llm_learning_rate": 1.2159107757375858e-06,
+      "loss": 0.7137,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.00012159107757375857,
+      "step": 6652
+    },
+    {
+      "epoch": 0.68,
+      "gumbel_tau": 0.2082,
+      "image_masks_loss": 0.0021,
+      "llm_learning_rate": 1.21519902834379e-06,
+      "loss": 0.6507,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.000121519902834379,
+      "step": 6653
+    },
+    {
+      "epoch": 0.68,
+      "gumbel_tau": 0.2081,
+      "image_masks_loss": 0.0141,
+      "llm_learning_rate": 1.2144874224351391e-06,
+      "loss": 0.7387,
+      "output_text_mask_loss": 0.0015,
+      "predictor_learning_rate": 0.0001214487422435139,
+      "step": 6654
+    },
+    {
+      "epoch": 0.68,
+      "gumbel_tau": 0.2081,
+      "image_masks_loss": 0.0067,
+      "llm_learning_rate": 1.2137759580899966e-06,
+      "loss": 0.6612,
+      "output_text_mask_loss": 0.0252,
+      "predictor_learning_rate": 0.00012137759580899966,
+      "step": 6655
+    },
+    {
+      "epoch": 0.68,
+      "gumbel_tau": 0.208,
+      "image_masks_loss": 0.0131,
+      "llm_learning_rate": 1.2130646353867095e-06,
+      "loss": 0.5977,
+      "output_text_mask_loss": 0.0004,
+      "predictor_learning_rate": 0.00012130646353867094,
+      "step": 6656
+    },
+    {
+      "epoch": 0.68,
+      "gumbel_tau": 0.208,
+      "image_masks_loss": 0.0137,
+      "llm_learning_rate": 1.2123534544036133e-06,
+      "loss": 0.7919,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.00012123534544036133,
+      "step": 6657
+    },
+    {
+      "epoch": 0.68,
+      "gumbel_tau": 0.2079,
+      "image_masks_loss": 0.0186,
+      "llm_learning_rate": 1.211642415219022e-06,
+      "loss": 0.6599,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.0001211642415219022,
+      "step": 6658
+    },
+    {
+      "epoch": 0.68,
+      "gumbel_tau": 0.2079,
+      "image_masks_loss": 0.0048,
+      "llm_learning_rate": 1.2109315179112374e-06,
+      "loss": 0.7657,
+      "output_text_mask_loss": 0.0355,
+      "predictor_learning_rate": 0.00012109315179112374,
+      "step": 6659
+    },
+    {
+      "epoch": 0.68,
+      "gumbel_tau": 0.2078,
+      "image_masks_loss": 0.0222,
+      "llm_learning_rate": 1.2102207625585454e-06,
+      "loss": 0.7195,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00012102207625585454,
+      "step": 6660
+    },
+    {
+      "epoch": 0.68,
+      "gumbel_tau": 0.2078,
+      "image_masks_loss": 0.0225,
+      "llm_learning_rate": 1.2095101492392162e-06,
+      "loss": 0.7317,
+      "output_text_mask_loss": 0.0061,
+      "predictor_learning_rate": 0.00012095101492392161,
+      "step": 6661
+    },
+    {
+      "epoch": 0.68,
+      "gumbel_tau": 0.2077,
+      "image_masks_loss": 0.0097,
+      "llm_learning_rate": 1.2087996780315025e-06,
+      "loss": 0.6674,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00012087996780315025,
+      "step": 6662
+    },
+    {
+      "epoch": 0.68,
+      "gumbel_tau": 0.2077,
+      "image_masks_loss": 0.0122,
+      "llm_learning_rate": 1.2080893490136456e-06,
+      "loss": 0.6904,
+      "output_text_mask_loss": 0.0032,
+      "predictor_learning_rate": 0.00012080893490136455,
+      "step": 6663
+    },
+    {
+      "epoch": 0.68,
+      "gumbel_tau": 0.2076,
+      "image_masks_loss": 0.0151,
+      "llm_learning_rate": 1.207379162263867e-06,
+      "loss": 0.7066,
+      "output_text_mask_loss": 0.0198,
+      "predictor_learning_rate": 0.00012073791622638669,
+      "step": 6664
+    },
+    {
+      "epoch": 0.68,
+      "gumbel_tau": 0.2076,
+      "image_masks_loss": 0.0133,
+      "llm_learning_rate": 1.206669117860374e-06,
+      "loss": 0.7523,
+      "output_text_mask_loss": 0.0118,
+      "predictor_learning_rate": 0.00012066691178603739,
+      "step": 6665
+    },
+    {
+      "epoch": 0.68,
+      "gumbel_tau": 0.2075,
+      "image_masks_loss": 0.0234,
+      "llm_learning_rate": 1.2059592158813587e-06,
+      "loss": 0.7471,
+      "output_text_mask_loss": 0.0357,
+      "predictor_learning_rate": 0.00012059592158813588,
+      "step": 6666
+    },
+    {
+      "epoch": 0.68,
+      "gumbel_tau": 0.2075,
+      "image_masks_loss": 0.0075,
+      "llm_learning_rate": 1.2052494564049971e-06,
+      "loss": 0.7577,
+      "output_text_mask_loss": 0.0238,
+      "predictor_learning_rate": 0.0001205249456404997,
+      "step": 6667
+    },
+    {
+      "epoch": 0.68,
+      "gumbel_tau": 0.2074,
+      "image_masks_loss": 0.011,
+      "llm_learning_rate": 1.204539839509449e-06,
+      "loss": 0.7039,
+      "output_text_mask_loss": 0.0595,
+      "predictor_learning_rate": 0.0001204539839509449,
+      "step": 6668
+    },
+    {
+      "epoch": 0.68,
+      "gumbel_tau": 0.2074,
+      "image_masks_loss": 0.0142,
+      "llm_learning_rate": 1.2038303652728587e-06,
+      "loss": 0.7095,
+      "output_text_mask_loss": 0.0385,
+      "predictor_learning_rate": 0.00012038303652728585,
+      "step": 6669
+    },
+    {
+      "epoch": 0.68,
+      "gumbel_tau": 0.2073,
+      "image_masks_loss": 0.0087,
+      "llm_learning_rate": 1.2031210337733565e-06,
+      "loss": 0.7234,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.00012031210337733565,
+      "step": 6670
+    },
+    {
+      "epoch": 0.68,
+      "gumbel_tau": 0.2073,
+      "image_masks_loss": 0.0174,
+      "llm_learning_rate": 1.2024118450890545e-06,
+      "loss": 0.7289,
+      "output_text_mask_loss": 0.0519,
+      "predictor_learning_rate": 0.00012024118450890544,
+      "step": 6671
+    },
+    {
+      "epoch": 0.68,
+      "gumbel_tau": 0.2073,
+      "image_masks_loss": 0.015,
+      "llm_learning_rate": 1.2017027992980507e-06,
+      "loss": 0.7953,
+      "output_text_mask_loss": 0.0105,
+      "predictor_learning_rate": 0.00012017027992980505,
+      "step": 6672
+    },
+    {
+      "epoch": 0.68,
+      "gumbel_tau": 0.2072,
+      "image_masks_loss": 0.0126,
+      "llm_learning_rate": 1.2009938964784261e-06,
+      "loss": 0.7669,
+      "output_text_mask_loss": 0.0019,
+      "predictor_learning_rate": 0.00012009938964784261,
+      "step": 6673
+    },
+    {
+      "epoch": 0.68,
+      "gumbel_tau": 0.2072,
+      "image_masks_loss": 0.006,
+      "llm_learning_rate": 1.2002851367082475e-06,
+      "loss": 0.6914,
+      "output_text_mask_loss": 0.0185,
+      "predictor_learning_rate": 0.00012002851367082473,
+      "step": 6674
+    },
+    {
+      "epoch": 0.68,
+      "gumbel_tau": 0.2071,
+      "image_masks_loss": 0.0111,
+      "llm_learning_rate": 1.1995765200655647e-06,
+      "loss": 0.6679,
+      "output_text_mask_loss": 0.0216,
+      "predictor_learning_rate": 0.00011995765200655646,
+      "step": 6675
+    },
+    {
+      "epoch": 0.68,
+      "gumbel_tau": 0.2071,
+      "image_masks_loss": 0.0214,
+      "llm_learning_rate": 1.1988680466284116e-06,
+      "loss": 0.6328,
+      "output_text_mask_loss": 0.0109,
+      "predictor_learning_rate": 0.00011988680466284116,
+      "step": 6676
+    },
+    {
+      "epoch": 0.68,
+      "gumbel_tau": 0.207,
+      "image_masks_loss": 0.0079,
+      "llm_learning_rate": 1.1981597164748077e-06,
+      "loss": 0.7424,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.00011981597164748075,
+      "step": 6677
+    },
+    {
+      "epoch": 0.68,
+      "gumbel_tau": 0.207,
+      "image_masks_loss": 0.005,
+      "llm_learning_rate": 1.1974515296827554e-06,
+      "loss": 0.7248,
+      "output_text_mask_loss": 0.0109,
+      "predictor_learning_rate": 0.00011974515296827554,
+      "step": 6678
+    },
+    {
+      "epoch": 0.68,
+      "gumbel_tau": 0.2069,
+      "image_masks_loss": 0.0152,
+      "llm_learning_rate": 1.196743486330242e-06,
+      "loss": 0.7071,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.00011967434863302418,
+      "step": 6679
+    },
+    {
+      "epoch": 0.68,
+      "gumbel_tau": 0.2069,
+      "image_masks_loss": 0.0236,
+      "llm_learning_rate": 1.196035586495239e-06,
+      "loss": 0.7158,
+      "output_text_mask_loss": 0.0105,
+      "predictor_learning_rate": 0.00011960355864952388,
+      "step": 6680
+    },
+    {
+      "epoch": 0.68,
+      "gumbel_tau": 0.2068,
+      "image_masks_loss": 0.0085,
+      "llm_learning_rate": 1.1953278302557004e-06,
+      "loss": 0.6482,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00011953278302557002,
+      "step": 6681
+    },
+    {
+      "epoch": 0.68,
+      "gumbel_tau": 0.2068,
+      "image_masks_loss": 0.0121,
+      "llm_learning_rate": 1.1946202176895682e-06,
+      "loss": 0.6714,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00011946202176895682,
+      "step": 6682
+    },
+    {
+      "epoch": 0.68,
+      "gumbel_tau": 0.2067,
+      "image_masks_loss": 0.0027,
+      "llm_learning_rate": 1.193912748874765e-06,
+      "loss": 0.7239,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.0001193912748874765,
+      "step": 6683
+    },
+    {
+      "epoch": 0.68,
+      "gumbel_tau": 0.2067,
+      "image_masks_loss": 0.0142,
+      "llm_learning_rate": 1.1932054238891997e-06,
+      "loss": 0.6032,
+      "output_text_mask_loss": 0.0071,
+      "predictor_learning_rate": 0.00011932054238891995,
+      "step": 6684
+    },
+    {
+      "epoch": 0.68,
+      "gumbel_tau": 0.2066,
+      "image_masks_loss": 0.0075,
+      "llm_learning_rate": 1.1924982428107634e-06,
+      "loss": 0.6256,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.00011924982428107633,
+      "step": 6685
+    },
+    {
+      "epoch": 0.69,
+      "gumbel_tau": 0.2066,
+      "image_masks_loss": 0.0106,
+      "llm_learning_rate": 1.1917912057173333e-06,
+      "loss": 0.7702,
+      "output_text_mask_loss": 0.0313,
+      "predictor_learning_rate": 0.00011917912057173333,
+      "step": 6686
+    },
+    {
+      "epoch": 0.69,
+      "gumbel_tau": 0.2065,
+      "image_masks_loss": 0.0093,
+      "llm_learning_rate": 1.1910843126867694e-06,
+      "loss": 0.7443,
+      "output_text_mask_loss": 0.0406,
+      "predictor_learning_rate": 0.00011910843126867694,
+      "step": 6687
+    },
+    {
+      "epoch": 0.69,
+      "gumbel_tau": 0.2065,
+      "image_masks_loss": 0.0051,
+      "llm_learning_rate": 1.1903775637969154e-06,
+      "loss": 0.8355,
+      "output_text_mask_loss": 0.001,
+      "predictor_learning_rate": 0.00011903775637969152,
+      "step": 6688
+    },
+    {
+      "epoch": 0.69,
+      "gumbel_tau": 0.2064,
+      "image_masks_loss": 0.0154,
+      "llm_learning_rate": 1.1896709591256023e-06,
+      "loss": 0.745,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.00011896709591256022,
+      "step": 6689
+    },
+    {
+      "epoch": 0.69,
+      "gumbel_tau": 0.2064,
+      "image_masks_loss": 0.0115,
+      "llm_learning_rate": 1.1889644987506416e-06,
+      "loss": 0.7529,
+      "output_text_mask_loss": 0.0065,
+      "predictor_learning_rate": 0.00011889644987506414,
+      "step": 6690
+    },
+    {
+      "epoch": 0.69,
+      "gumbel_tau": 0.2063,
+      "image_masks_loss": 0.0036,
+      "llm_learning_rate": 1.1882581827498307e-06,
+      "loss": 0.8074,
+      "output_text_mask_loss": 0.0014,
+      "predictor_learning_rate": 0.00011882581827498307,
+      "step": 6691
+    },
+    {
+      "epoch": 0.69,
+      "gumbel_tau": 0.2063,
+      "image_masks_loss": 0.0175,
+      "llm_learning_rate": 1.1875520112009507e-06,
+      "loss": 0.7964,
+      "output_text_mask_loss": 0.0025,
+      "predictor_learning_rate": 0.00011875520112009505,
+      "step": 6692
+    },
+    {
+      "epoch": 0.69,
+      "gumbel_tau": 0.2062,
+      "image_masks_loss": 0.0103,
+      "llm_learning_rate": 1.1868459841817662e-06,
+      "loss": 0.7744,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00011868459841817662,
+      "step": 6693
+    },
+    {
+      "epoch": 0.69,
+      "gumbel_tau": 0.2062,
+      "image_masks_loss": 0.0121,
+      "llm_learning_rate": 1.186140101770027e-06,
+      "loss": 0.7281,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 0.00011861401017700269,
+      "step": 6694
+    },
+    {
+      "epoch": 0.69,
+      "gumbel_tau": 0.2061,
+      "image_masks_loss": 0.0117,
+      "llm_learning_rate": 1.1854343640434662e-06,
+      "loss": 0.7379,
+      "output_text_mask_loss": 0.0355,
+      "predictor_learning_rate": 0.00011854343640434661,
+      "step": 6695
+    },
+    {
+      "epoch": 0.69,
+      "gumbel_tau": 0.2061,
+      "image_masks_loss": 0.0074,
+      "llm_learning_rate": 1.1847287710798014e-06,
+      "loss": 0.7485,
+      "output_text_mask_loss": 0.0244,
+      "predictor_learning_rate": 0.00011847287710798013,
+      "step": 6696
+    },
+    {
+      "epoch": 0.69,
+      "gumbel_tau": 0.206,
+      "image_masks_loss": 0.016,
+      "llm_learning_rate": 1.1840233229567342e-06,
+      "loss": 0.6613,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.00011840233229567341,
+      "step": 6697
+    },
+    {
+      "epoch": 0.69,
+      "gumbel_tau": 0.206,
+      "image_masks_loss": 0.029,
+      "llm_learning_rate": 1.1833180197519495e-06,
+      "loss": 0.7014,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00011833180197519494,
+      "step": 6698
+    },
+    {
+      "epoch": 0.69,
+      "gumbel_tau": 0.2059,
+      "image_masks_loss": 0.0048,
+      "llm_learning_rate": 1.1826128615431174e-06,
+      "loss": 0.7359,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.00011826128615431172,
+      "step": 6699
+    },
+    {
+      "epoch": 0.69,
+      "gumbel_tau": 0.2059,
+      "image_masks_loss": 0.0066,
+      "llm_learning_rate": 1.1819078484078914e-06,
+      "loss": 0.8007,
+      "output_text_mask_loss": 0.0219,
+      "predictor_learning_rate": 0.00011819078484078912,
+      "step": 6700
+    },
+    {
+      "epoch": 0.69,
+      "gumbel_tau": 0.2058,
+      "image_masks_loss": 0.0141,
+      "llm_learning_rate": 1.181202980423908e-06,
+      "loss": 0.6525,
+      "output_text_mask_loss": 0.0124,
+      "predictor_learning_rate": 0.0001181202980423908,
+      "step": 6701
+    },
+    {
+      "epoch": 0.69,
+      "gumbel_tau": 0.2058,
+      "image_masks_loss": 0.0049,
+      "llm_learning_rate": 1.1804982576687913e-06,
+      "loss": 0.7063,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 0.00011804982576687912,
+      "step": 6702
+    },
+    {
+      "epoch": 0.69,
+      "gumbel_tau": 0.2057,
+      "image_masks_loss": 0.017,
+      "llm_learning_rate": 1.1797936802201452e-06,
+      "loss": 0.7614,
+      "output_text_mask_loss": 0.0431,
+      "predictor_learning_rate": 0.0001179793680220145,
+      "step": 6703
+    },
+    {
+      "epoch": 0.69,
+      "gumbel_tau": 0.2057,
+      "image_masks_loss": 0.0144,
+      "llm_learning_rate": 1.1790892481555595e-06,
+      "loss": 0.735,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00011790892481555596,
+      "step": 6704
+    },
+    {
+      "epoch": 0.69,
+      "gumbel_tau": 0.2056,
+      "image_masks_loss": 0.0224,
+      "llm_learning_rate": 1.1783849615526083e-06,
+      "loss": 0.7484,
+      "output_text_mask_loss": 0.013,
+      "predictor_learning_rate": 0.00011783849615526082,
+      "step": 6705
+    },
+    {
+      "epoch": 0.69,
+      "gumbel_tau": 0.2056,
+      "image_masks_loss": 0.0092,
+      "llm_learning_rate": 1.1776808204888488e-06,
+      "loss": 0.7131,
+      "output_text_mask_loss": 0.0077,
+      "predictor_learning_rate": 0.00011776808204888487,
+      "step": 6706
+    },
+    {
+      "epoch": 0.69,
+      "gumbel_tau": 0.2055,
+      "image_masks_loss": 0.019,
+      "llm_learning_rate": 1.1769768250418227e-06,
+      "loss": 0.7071,
+      "output_text_mask_loss": 0.0065,
+      "predictor_learning_rate": 0.00011769768250418225,
+      "step": 6707
+    },
+    {
+      "epoch": 0.69,
+      "gumbel_tau": 0.2055,
+      "image_masks_loss": 0.0052,
+      "llm_learning_rate": 1.1762729752890544e-06,
+      "loss": 0.7482,
+      "output_text_mask_loss": 0.0441,
+      "predictor_learning_rate": 0.00011762729752890545,
+      "step": 6708
+    },
+    {
+      "epoch": 0.69,
+      "gumbel_tau": 0.2054,
+      "image_masks_loss": 0.0073,
+      "llm_learning_rate": 1.1755692713080558e-06,
+      "loss": 0.7044,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00011755692713080557,
+      "step": 6709
+    },
+    {
+      "epoch": 0.69,
+      "gumbel_tau": 0.2054,
+      "image_masks_loss": 0.0126,
+      "llm_learning_rate": 1.1748657131763192e-06,
+      "loss": 0.6719,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00011748657131763191,
+      "step": 6710
+    },
+    {
+      "epoch": 0.69,
+      "gumbel_tau": 0.2054,
+      "image_masks_loss": 0.0072,
+      "llm_learning_rate": 1.1741623009713216e-06,
+      "loss": 0.7219,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.00011741623009713215,
+      "step": 6711
+    },
+    {
+      "epoch": 0.69,
+      "gumbel_tau": 0.2053,
+      "image_masks_loss": 0.0117,
+      "llm_learning_rate": 1.1734590347705246e-06,
+      "loss": 0.734,
+      "output_text_mask_loss": 0.0063,
+      "predictor_learning_rate": 0.00011734590347705246,
+      "step": 6712
+    },
+    {
+      "epoch": 0.69,
+      "gumbel_tau": 0.2053,
+      "image_masks_loss": 0.0209,
+      "llm_learning_rate": 1.1727559146513736e-06,
+      "loss": 0.6287,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00011727559146513736,
+      "step": 6713
+    },
+    {
+      "epoch": 0.69,
+      "gumbel_tau": 0.2052,
+      "image_masks_loss": 0.0057,
+      "llm_learning_rate": 1.1720529406912976e-06,
+      "loss": 0.7544,
+      "output_text_mask_loss": 0.0097,
+      "predictor_learning_rate": 0.00011720529406912975,
+      "step": 6714
+    },
+    {
+      "epoch": 0.69,
+      "gumbel_tau": 0.2052,
+      "image_masks_loss": 0.0169,
+      "llm_learning_rate": 1.1713501129677095e-06,
+      "loss": 0.7593,
+      "output_text_mask_loss": 0.0587,
+      "predictor_learning_rate": 0.00011713501129677096,
+      "step": 6715
+    },
+    {
+      "epoch": 0.69,
+      "gumbel_tau": 0.2051,
+      "image_masks_loss": 0.0128,
+      "llm_learning_rate": 1.1706474315580065e-06,
+      "loss": 0.6727,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.00011706474315580065,
+      "step": 6716
+    },
+    {
+      "epoch": 0.69,
+      "gumbel_tau": 0.2051,
+      "image_masks_loss": 0.0055,
+      "llm_learning_rate": 1.1699448965395693e-06,
+      "loss": 0.653,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00011699448965395692,
+      "step": 6717
+    },
+    {
+      "epoch": 0.69,
+      "gumbel_tau": 0.205,
+      "image_masks_loss": 0.0122,
+      "llm_learning_rate": 1.169242507989763e-06,
+      "loss": 0.7081,
+      "output_text_mask_loss": 0.001,
+      "predictor_learning_rate": 0.00011692425079897628,
+      "step": 6718
+    },
+    {
+      "epoch": 0.69,
+      "gumbel_tau": 0.205,
+      "image_masks_loss": 0.0154,
+      "llm_learning_rate": 1.1685402659859355e-06,
+      "loss": 0.7695,
+      "output_text_mask_loss": 0.058,
+      "predictor_learning_rate": 0.00011685402659859354,
+      "step": 6719
+    },
+    {
+      "epoch": 0.69,
+      "gumbel_tau": 0.2049,
+      "image_masks_loss": 0.0175,
+      "llm_learning_rate": 1.1678381706054187e-06,
+      "loss": 0.7035,
+      "output_text_mask_loss": 0.0107,
+      "predictor_learning_rate": 0.00011678381706054186,
+      "step": 6720
+    },
+    {
+      "epoch": 0.69,
+      "gumbel_tau": 0.2049,
+      "image_masks_loss": 0.0062,
+      "llm_learning_rate": 1.1671362219255309e-06,
+      "loss": 0.798,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 0.00011671362219255307,
+      "step": 6721
+    },
+    {
+      "epoch": 0.69,
+      "gumbel_tau": 0.2048,
+      "image_masks_loss": 0.0098,
+      "llm_learning_rate": 1.1664344200235711e-06,
+      "loss": 0.8615,
+      "output_text_mask_loss": 0.0307,
+      "predictor_learning_rate": 0.0001166434420023571,
+      "step": 6722
+    },
+    {
+      "epoch": 0.69,
+      "gumbel_tau": 0.2048,
+      "image_masks_loss": 0.0128,
+      "llm_learning_rate": 1.1657327649768233e-06,
+      "loss": 0.7308,
+      "output_text_mask_loss": 0.0324,
+      "predictor_learning_rate": 0.00011657327649768232,
+      "step": 6723
+    },
+    {
+      "epoch": 0.69,
+      "gumbel_tau": 0.2047,
+      "image_masks_loss": 0.0159,
+      "llm_learning_rate": 1.1650312568625554e-06,
+      "loss": 0.7396,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00011650312568625554,
+      "step": 6724
+    },
+    {
+      "epoch": 0.69,
+      "gumbel_tau": 0.2047,
+      "image_masks_loss": 0.0087,
+      "llm_learning_rate": 1.1643298957580192e-06,
+      "loss": 0.6818,
+      "output_text_mask_loss": 0.0275,
+      "predictor_learning_rate": 0.00011643298957580192,
+      "step": 6725
+    },
+    {
+      "epoch": 0.69,
+      "gumbel_tau": 0.2046,
+      "image_masks_loss": 0.008,
+      "llm_learning_rate": 1.1636286817404499e-06,
+      "loss": 0.7258,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00011636286817404498,
+      "step": 6726
+    },
+    {
+      "epoch": 0.69,
+      "gumbel_tau": 0.2046,
+      "image_masks_loss": 0.0085,
+      "llm_learning_rate": 1.1629276148870658e-06,
+      "loss": 0.7104,
+      "output_text_mask_loss": 0.0076,
+      "predictor_learning_rate": 0.00011629276148870657,
+      "step": 6727
+    },
+    {
+      "epoch": 0.69,
+      "gumbel_tau": 0.2045,
+      "image_masks_loss": 0.0099,
+      "llm_learning_rate": 1.1622266952750719e-06,
+      "loss": 0.7791,
+      "output_text_mask_loss": 0.0038,
+      "predictor_learning_rate": 0.00011622266952750718,
+      "step": 6728
+    },
+    {
+      "epoch": 0.69,
+      "gumbel_tau": 0.2045,
+      "image_masks_loss": 0.0107,
+      "llm_learning_rate": 1.1615259229816543e-06,
+      "loss": 0.7597,
+      "output_text_mask_loss": 0.028,
+      "predictor_learning_rate": 0.00011615259229816541,
+      "step": 6729
+    },
+    {
+      "epoch": 0.69,
+      "gumbel_tau": 0.2044,
+      "image_masks_loss": 0.0103,
+      "llm_learning_rate": 1.1608252980839831e-06,
+      "loss": 0.6859,
+      "output_text_mask_loss": 0.0081,
+      "predictor_learning_rate": 0.00011608252980839832,
+      "step": 6730
+    },
+    {
+      "epoch": 0.69,
+      "gumbel_tau": 0.2044,
+      "image_masks_loss": 0.0029,
+      "llm_learning_rate": 1.1601248206592141e-06,
+      "loss": 0.7152,
+      "output_text_mask_loss": 0.0073,
+      "predictor_learning_rate": 0.00011601248206592141,
+      "step": 6731
+    },
+    {
+      "epoch": 0.69,
+      "gumbel_tau": 0.2043,
+      "image_masks_loss": 0.0092,
+      "llm_learning_rate": 1.1594244907844831e-06,
+      "loss": 0.7067,
+      "output_text_mask_loss": 0.0147,
+      "predictor_learning_rate": 0.0001159424490784483,
+      "step": 6732
+    },
+    {
+      "epoch": 0.69,
+      "gumbel_tau": 0.2043,
+      "image_masks_loss": 0.0113,
+      "llm_learning_rate": 1.1587243085369127e-06,
+      "loss": 0.8397,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00011587243085369126,
+      "step": 6733
+    },
+    {
+      "epoch": 0.69,
+      "gumbel_tau": 0.2042,
+      "image_masks_loss": 0.0078,
+      "llm_learning_rate": 1.1580242739936095e-06,
+      "loss": 0.7068,
+      "output_text_mask_loss": 0.0025,
+      "predictor_learning_rate": 0.00011580242739936095,
+      "step": 6734
+    },
+    {
+      "epoch": 0.69,
+      "gumbel_tau": 0.2042,
+      "image_masks_loss": 0.0226,
+      "llm_learning_rate": 1.1573243872316628e-06,
+      "loss": 0.8322,
+      "output_text_mask_loss": 0.0075,
+      "predictor_learning_rate": 0.00011573243872316627,
+      "step": 6735
+    },
+    {
+      "epoch": 0.69,
+      "gumbel_tau": 0.2041,
+      "image_masks_loss": 0.0257,
+      "llm_learning_rate": 1.1566246483281451e-06,
+      "loss": 0.7768,
+      "output_text_mask_loss": 0.0032,
+      "predictor_learning_rate": 0.00011566246483281451,
+      "step": 6736
+    },
+    {
+      "epoch": 0.69,
+      "gumbel_tau": 0.2041,
+      "image_masks_loss": 0.0157,
+      "llm_learning_rate": 1.1559250573601133e-06,
+      "loss": 0.7628,
+      "output_text_mask_loss": 0.0292,
+      "predictor_learning_rate": 0.00011559250573601132,
+      "step": 6737
+    },
+    {
+      "epoch": 0.69,
+      "gumbel_tau": 0.204,
+      "image_masks_loss": 0.0093,
+      "llm_learning_rate": 1.155225614404608e-06,
+      "loss": 0.7464,
+      "output_text_mask_loss": 0.0195,
+      "predictor_learning_rate": 0.00011552256144046081,
+      "step": 6738
+    },
+    {
+      "epoch": 0.69,
+      "gumbel_tau": 0.204,
+      "image_masks_loss": 0.004,
+      "llm_learning_rate": 1.1545263195386537e-06,
+      "loss": 0.6791,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00011545263195386536,
+      "step": 6739
+    },
+    {
+      "epoch": 0.69,
+      "gumbel_tau": 0.204,
+      "image_masks_loss": 0.0108,
+      "llm_learning_rate": 1.1538271728392567e-06,
+      "loss": 0.7259,
+      "output_text_mask_loss": 0.0223,
+      "predictor_learning_rate": 0.00011538271728392566,
+      "step": 6740
+    },
+    {
+      "epoch": 0.69,
+      "gumbel_tau": 0.2039,
+      "image_masks_loss": 0.004,
+      "llm_learning_rate": 1.1531281743834109e-06,
+      "loss": 0.7251,
+      "output_text_mask_loss": 0.0175,
+      "predictor_learning_rate": 0.00011531281743834107,
+      "step": 6741
+    },
+    {
+      "epoch": 0.69,
+      "gumbel_tau": 0.2039,
+      "image_masks_loss": 0.009,
+      "llm_learning_rate": 1.1524293242480905e-06,
+      "loss": 0.7756,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 0.00011524293242480904,
+      "step": 6742
+    },
+    {
+      "epoch": 0.69,
+      "gumbel_tau": 0.2038,
+      "image_masks_loss": 0.0092,
+      "llm_learning_rate": 1.1517306225102548e-06,
+      "loss": 0.7642,
+      "output_text_mask_loss": 0.0181,
+      "predictor_learning_rate": 0.00011517306225102548,
+      "step": 6743
+    },
+    {
+      "epoch": 0.69,
+      "gumbel_tau": 0.2038,
+      "image_masks_loss": 0.0126,
+      "llm_learning_rate": 1.151032069246846e-06,
+      "loss": 0.769,
+      "output_text_mask_loss": 0.0051,
+      "predictor_learning_rate": 0.0001151032069246846,
+      "step": 6744
+    },
+    {
+      "epoch": 0.69,
+      "gumbel_tau": 0.2037,
+      "image_masks_loss": 0.0148,
+      "llm_learning_rate": 1.1503336645347903e-06,
+      "loss": 0.6832,
+      "output_text_mask_loss": 0.0039,
+      "predictor_learning_rate": 0.00011503336645347903,
+      "step": 6745
+    },
+    {
+      "epoch": 0.69,
+      "gumbel_tau": 0.2037,
+      "image_masks_loss": 0.0056,
+      "llm_learning_rate": 1.1496354084509977e-06,
+      "loss": 0.672,
+      "output_text_mask_loss": 0.0025,
+      "predictor_learning_rate": 0.00011496354084509977,
+      "step": 6746
+    },
+    {
+      "epoch": 0.69,
+      "gumbel_tau": 0.2036,
+      "image_masks_loss": 0.0046,
+      "llm_learning_rate": 1.1489373010723612e-06,
+      "loss": 0.6621,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.00011489373010723612,
+      "step": 6747
+    },
+    {
+      "epoch": 0.69,
+      "gumbel_tau": 0.2036,
+      "image_masks_loss": 0.0101,
+      "llm_learning_rate": 1.148239342475759e-06,
+      "loss": 0.7138,
+      "output_text_mask_loss": 0.0096,
+      "predictor_learning_rate": 0.00011482393424757589,
+      "step": 6748
+    },
+    {
+      "epoch": 0.69,
+      "gumbel_tau": 0.2035,
+      "image_masks_loss": 0.014,
+      "llm_learning_rate": 1.1475415327380515e-06,
+      "loss": 0.7186,
+      "output_text_mask_loss": 0.0086,
+      "predictor_learning_rate": 0.00011475415327380514,
+      "step": 6749
+    },
+    {
+      "epoch": 0.69,
+      "gumbel_tau": 0.2035,
+      "image_masks_loss": 0.0128,
+      "llm_learning_rate": 1.1468438719360838e-06,
+      "loss": 0.7106,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.00011468438719360838,
+      "step": 6750
+    },
+    {
+      "epoch": 0.69,
+      "gumbel_tau": 0.2034,
+      "image_masks_loss": 0.0071,
+      "llm_learning_rate": 1.1461463601466821e-06,
+      "loss": 0.6483,
+      "output_text_mask_loss": 0.0034,
+      "predictor_learning_rate": 0.0001146146360146682,
+      "step": 6751
+    },
+    {
+      "epoch": 0.69,
+      "gumbel_tau": 0.2034,
+      "image_masks_loss": 0.0121,
+      "llm_learning_rate": 1.1454489974466579e-06,
+      "loss": 0.7473,
+      "output_text_mask_loss": 0.0095,
+      "predictor_learning_rate": 0.00011454489974466578,
+      "step": 6752
+    },
+    {
+      "epoch": 0.69,
+      "gumbel_tau": 0.2033,
+      "image_masks_loss": 0.0061,
+      "llm_learning_rate": 1.144751783912808e-06,
+      "loss": 0.7679,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.0001144751783912808,
+      "step": 6753
+    },
+    {
+      "epoch": 0.69,
+      "gumbel_tau": 0.2033,
+      "image_masks_loss": 0.0197,
+      "llm_learning_rate": 1.1440547196219104e-06,
+      "loss": 0.7514,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 0.00011440547196219104,
+      "step": 6754
+    },
+    {
+      "epoch": 0.69,
+      "gumbel_tau": 0.2032,
+      "image_masks_loss": 0.0062,
+      "llm_learning_rate": 1.1433578046507273e-06,
+      "loss": 0.6714,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.00011433578046507273,
+      "step": 6755
+    },
+    {
+      "epoch": 0.69,
+      "gumbel_tau": 0.2032,
+      "image_masks_loss": 0.0048,
+      "llm_learning_rate": 1.1426610390760048e-06,
+      "loss": 0.7886,
+      "output_text_mask_loss": 0.0038,
+      "predictor_learning_rate": 0.00011426610390760048,
+      "step": 6756
+    },
+    {
+      "epoch": 0.69,
+      "gumbel_tau": 0.2031,
+      "image_masks_loss": 0.0091,
+      "llm_learning_rate": 1.1419644229744717e-06,
+      "loss": 0.6099,
+      "output_text_mask_loss": 0.017,
+      "predictor_learning_rate": 0.00011419644229744716,
+      "step": 6757
+    },
+    {
+      "epoch": 0.69,
+      "gumbel_tau": 0.2031,
+      "image_masks_loss": 0.0031,
+      "llm_learning_rate": 1.1412679564228416e-06,
+      "loss": 0.766,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00011412679564228415,
+      "step": 6758
+    },
+    {
+      "epoch": 0.69,
+      "gumbel_tau": 0.203,
+      "image_masks_loss": 0.0078,
+      "llm_learning_rate": 1.1405716394978102e-06,
+      "loss": 0.7126,
+      "output_text_mask_loss": 0.0071,
+      "predictor_learning_rate": 0.00011405716394978102,
+      "step": 6759
+    },
+    {
+      "epoch": 0.69,
+      "gumbel_tau": 0.203,
+      "image_masks_loss": 0.0144,
+      "llm_learning_rate": 1.1398754722760572e-06,
+      "loss": 0.6908,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00011398754722760571,
+      "step": 6760
+    },
+    {
+      "epoch": 0.69,
+      "gumbel_tau": 0.2029,
+      "image_masks_loss": 0.0135,
+      "llm_learning_rate": 1.1391794548342477e-06,
+      "loss": 0.7248,
+      "output_text_mask_loss": 0.0076,
+      "predictor_learning_rate": 0.00011391794548342476,
+      "step": 6761
+    },
+    {
+      "epoch": 0.69,
+      "gumbel_tau": 0.2029,
+      "image_masks_loss": 0.0076,
+      "llm_learning_rate": 1.1384835872490277e-06,
+      "loss": 0.735,
+      "output_text_mask_loss": 0.0105,
+      "predictor_learning_rate": 0.00011384835872490277,
+      "step": 6762
+    },
+    {
+      "epoch": 0.69,
+      "gumbel_tau": 0.2028,
+      "image_masks_loss": 0.0051,
+      "llm_learning_rate": 1.1377878695970277e-06,
+      "loss": 0.7681,
+      "output_text_mask_loss": 0.0137,
+      "predictor_learning_rate": 0.00011377878695970276,
+      "step": 6763
+    },
+    {
+      "epoch": 0.69,
+      "gumbel_tau": 0.2028,
+      "image_masks_loss": 0.0115,
+      "llm_learning_rate": 1.1370923019548615e-06,
+      "loss": 0.651,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 0.00011370923019548615,
+      "step": 6764
+    },
+    {
+      "epoch": 0.69,
+      "gumbel_tau": 0.2028,
+      "image_masks_loss": 0.0104,
+      "llm_learning_rate": 1.1363968843991272e-06,
+      "loss": 0.7033,
+      "output_text_mask_loss": 0.0399,
+      "predictor_learning_rate": 0.0001136396884399127,
+      "step": 6765
+    },
+    {
+      "epoch": 0.69,
+      "gumbel_tau": 0.2027,
+      "image_masks_loss": 0.0069,
+      "llm_learning_rate": 1.135701617006405e-06,
+      "loss": 0.8066,
+      "output_text_mask_loss": 0.029,
+      "predictor_learning_rate": 0.00011357016170064049,
+      "step": 6766
+    },
+    {
+      "epoch": 0.69,
+      "gumbel_tau": 0.2027,
+      "image_masks_loss": 0.0056,
+      "llm_learning_rate": 1.1350064998532587e-06,
+      "loss": 0.7305,
+      "output_text_mask_loss": 0.0082,
+      "predictor_learning_rate": 0.00011350064998532586,
+      "step": 6767
+    },
+    {
+      "epoch": 0.69,
+      "gumbel_tau": 0.2026,
+      "image_masks_loss": 0.0209,
+      "llm_learning_rate": 1.1343115330162388e-06,
+      "loss": 0.7772,
+      "output_text_mask_loss": 0.0553,
+      "predictor_learning_rate": 0.00011343115330162387,
+      "step": 6768
+    },
+    {
+      "epoch": 0.69,
+      "gumbel_tau": 0.2026,
+      "image_masks_loss": 0.0091,
+      "llm_learning_rate": 1.1336167165718742e-06,
+      "loss": 0.7349,
+      "output_text_mask_loss": 0.0202,
+      "predictor_learning_rate": 0.00011336167165718742,
+      "step": 6769
+    },
+    {
+      "epoch": 0.69,
+      "gumbel_tau": 0.2025,
+      "image_masks_loss": 0.0137,
+      "llm_learning_rate": 1.13292205059668e-06,
+      "loss": 0.6431,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.000113292205059668,
+      "step": 6770
+    },
+    {
+      "epoch": 0.69,
+      "gumbel_tau": 0.2025,
+      "image_masks_loss": 0.0016,
+      "llm_learning_rate": 1.1322275351671547e-06,
+      "loss": 0.7053,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 0.00011322275351671546,
+      "step": 6771
+    },
+    {
+      "epoch": 0.69,
+      "gumbel_tau": 0.2024,
+      "image_masks_loss": 0.0077,
+      "llm_learning_rate": 1.131533170359779e-06,
+      "loss": 0.5931,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.0001131533170359779,
+      "step": 6772
+    },
+    {
+      "epoch": 0.69,
+      "gumbel_tau": 0.2024,
+      "image_masks_loss": 0.0117,
+      "llm_learning_rate": 1.1308389562510196e-06,
+      "loss": 0.7406,
+      "output_text_mask_loss": 0.0404,
+      "predictor_learning_rate": 0.00011308389562510196,
+      "step": 6773
+    },
+    {
+      "epoch": 0.69,
+      "gumbel_tau": 0.2023,
+      "image_masks_loss": 0.0071,
+      "llm_learning_rate": 1.130144892917324e-06,
+      "loss": 0.7079,
+      "output_text_mask_loss": 0.0193,
+      "predictor_learning_rate": 0.00011301448929173239,
+      "step": 6774
+    },
+    {
+      "epoch": 0.69,
+      "gumbel_tau": 0.2023,
+      "image_masks_loss": 0.0216,
+      "llm_learning_rate": 1.129450980435124e-06,
+      "loss": 0.7838,
+      "output_text_mask_loss": 0.0305,
+      "predictor_learning_rate": 0.00011294509804351238,
+      "step": 6775
+    },
+    {
+      "epoch": 0.69,
+      "gumbel_tau": 0.2022,
+      "image_masks_loss": 0.0034,
+      "llm_learning_rate": 1.128757218880835e-06,
+      "loss": 0.6741,
+      "output_text_mask_loss": 0.0185,
+      "predictor_learning_rate": 0.00011287572188808348,
+      "step": 6776
+    },
+    {
+      "epoch": 0.69,
+      "gumbel_tau": 0.2022,
+      "image_masks_loss": 0.019,
+      "llm_learning_rate": 1.1280636083308553e-06,
+      "loss": 0.6905,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.00011280636083308551,
+      "step": 6777
+    },
+    {
+      "epoch": 0.69,
+      "gumbel_tau": 0.2021,
+      "image_masks_loss": 0.0129,
+      "llm_learning_rate": 1.1273701488615665e-06,
+      "loss": 0.734,
+      "output_text_mask_loss": 0.0526,
+      "predictor_learning_rate": 0.00011273701488615665,
+      "step": 6778
+    },
+    {
+      "epoch": 0.69,
+      "gumbel_tau": 0.2021,
+      "image_masks_loss": 0.0152,
+      "llm_learning_rate": 1.126676840549334e-06,
+      "loss": 0.6959,
+      "output_text_mask_loss": 0.0198,
+      "predictor_learning_rate": 0.00011266768405493338,
+      "step": 6779
+    },
+    {
+      "epoch": 0.69,
+      "gumbel_tau": 0.202,
+      "image_masks_loss": 0.0042,
+      "llm_learning_rate": 1.1259836834705074e-06,
+      "loss": 0.8082,
+      "output_text_mask_loss": 0.0504,
+      "predictor_learning_rate": 0.00011259836834705073,
+      "step": 6780
+    },
+    {
+      "epoch": 0.69,
+      "gumbel_tau": 0.202,
+      "image_masks_loss": 0.008,
+      "llm_learning_rate": 1.125290677701418e-06,
+      "loss": 0.6447,
+      "output_text_mask_loss": 0.0039,
+      "predictor_learning_rate": 0.0001125290677701418,
+      "step": 6781
+    },
+    {
+      "epoch": 0.69,
+      "gumbel_tau": 0.2019,
+      "image_masks_loss": 0.0098,
+      "llm_learning_rate": 1.1245978233183813e-06,
+      "loss": 0.796,
+      "output_text_mask_loss": 0.0208,
+      "predictor_learning_rate": 0.00011245978233183812,
+      "step": 6782
+    },
+    {
+      "epoch": 0.69,
+      "gumbel_tau": 0.2019,
+      "image_masks_loss": 0.0106,
+      "llm_learning_rate": 1.123905120397696e-06,
+      "loss": 0.6987,
+      "output_text_mask_loss": 0.0227,
+      "predictor_learning_rate": 0.0001123905120397696,
+      "step": 6783
+    },
+    {
+      "epoch": 0.7,
+      "gumbel_tau": 0.2018,
+      "image_masks_loss": 0.0143,
+      "llm_learning_rate": 1.123212569015644e-06,
+      "loss": 0.7851,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 0.00011232125690156439,
+      "step": 6784
+    },
+    {
+      "epoch": 0.7,
+      "gumbel_tau": 0.2018,
+      "image_masks_loss": 0.0075,
+      "llm_learning_rate": 1.1225201692484903e-06,
+      "loss": 0.732,
+      "output_text_mask_loss": 0.0206,
+      "predictor_learning_rate": 0.00011225201692484901,
+      "step": 6785
+    },
+    {
+      "epoch": 0.7,
+      "gumbel_tau": 0.2018,
+      "image_masks_loss": 0.0105,
+      "llm_learning_rate": 1.121827921172483e-06,
+      "loss": 0.7012,
+      "output_text_mask_loss": 0.0153,
+      "predictor_learning_rate": 0.0001121827921172483,
+      "step": 6786
+    },
+    {
+      "epoch": 0.7,
+      "gumbel_tau": 0.2017,
+      "image_masks_loss": 0.0212,
+      "llm_learning_rate": 1.1211358248638568e-06,
+      "loss": 0.7482,
+      "output_text_mask_loss": 0.0107,
+      "predictor_learning_rate": 0.00011211358248638568,
+      "step": 6787
+    },
+    {
+      "epoch": 0.7,
+      "gumbel_tau": 0.2017,
+      "image_masks_loss": 0.0168,
+      "llm_learning_rate": 1.1204438803988238e-06,
+      "loss": 0.6933,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00011204438803988237,
+      "step": 6788
+    },
+    {
+      "epoch": 0.7,
+      "gumbel_tau": 0.2016,
+      "image_masks_loss": 0.0098,
+      "llm_learning_rate": 1.1197520878535833e-06,
+      "loss": 0.7885,
+      "output_text_mask_loss": 0.0124,
+      "predictor_learning_rate": 0.00011197520878535833,
+      "step": 6789
+    },
+    {
+      "epoch": 0.7,
+      "gumbel_tau": 0.2016,
+      "image_masks_loss": 0.006,
+      "llm_learning_rate": 1.1190604473043175e-06,
+      "loss": 0.7294,
+      "output_text_mask_loss": 0.025,
+      "predictor_learning_rate": 0.00011190604473043176,
+      "step": 6790
+    },
+    {
+      "epoch": 0.7,
+      "gumbel_tau": 0.2015,
+      "image_masks_loss": 0.0029,
+      "llm_learning_rate": 1.1183689588271912e-06,
+      "loss": 0.7269,
+      "output_text_mask_loss": 0.0096,
+      "predictor_learning_rate": 0.00011183689588271912,
+      "step": 6791
+    },
+    {
+      "epoch": 0.7,
+      "gumbel_tau": 0.2015,
+      "image_masks_loss": 0.0071,
+      "llm_learning_rate": 1.1176776224983519e-06,
+      "loss": 0.7017,
+      "output_text_mask_loss": 0.0147,
+      "predictor_learning_rate": 0.00011176776224983518,
+      "step": 6792
+    },
+    {
+      "epoch": 0.7,
+      "gumbel_tau": 0.2014,
+      "image_masks_loss": 0.0077,
+      "llm_learning_rate": 1.1169864383939321e-06,
+      "loss": 0.7772,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0001116986438393932,
+      "step": 6793
+    },
+    {
+      "epoch": 0.7,
+      "gumbel_tau": 0.2014,
+      "image_masks_loss": 0.0072,
+      "llm_learning_rate": 1.1162954065900468e-06,
+      "loss": 0.6989,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 0.00011162954065900468,
+      "step": 6794
+    },
+    {
+      "epoch": 0.7,
+      "gumbel_tau": 0.2013,
+      "image_masks_loss": 0.004,
+      "llm_learning_rate": 1.1156045271627933e-06,
+      "loss": 0.7587,
+      "output_text_mask_loss": 0.0081,
+      "predictor_learning_rate": 0.00011156045271627932,
+      "step": 6795
+    },
+    {
+      "epoch": 0.7,
+      "gumbel_tau": 0.2013,
+      "image_masks_loss": 0.0094,
+      "llm_learning_rate": 1.1149138001882529e-06,
+      "loss": 0.7982,
+      "output_text_mask_loss": 0.0275,
+      "predictor_learning_rate": 0.00011149138001882528,
+      "step": 6796
+    },
+    {
+      "epoch": 0.7,
+      "gumbel_tau": 0.2012,
+      "image_masks_loss": 0.0334,
+      "llm_learning_rate": 1.1142232257424897e-06,
+      "loss": 0.7982,
+      "output_text_mask_loss": 0.0032,
+      "predictor_learning_rate": 0.00011142232257424897,
+      "step": 6797
+    },
+    {
+      "epoch": 0.7,
+      "gumbel_tau": 0.2012,
+      "image_masks_loss": 0.0088,
+      "llm_learning_rate": 1.113532803901552e-06,
+      "loss": 0.6457,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00011135328039015519,
+      "step": 6798
+    },
+    {
+      "epoch": 0.7,
+      "gumbel_tau": 0.2011,
+      "image_masks_loss": 0.012,
+      "llm_learning_rate": 1.112842534741469e-06,
+      "loss": 0.7282,
+      "output_text_mask_loss": 0.0054,
+      "predictor_learning_rate": 0.00011128425347414689,
+      "step": 6799
+    },
+    {
+      "epoch": 0.7,
+      "gumbel_tau": 0.2011,
+      "image_masks_loss": 0.0065,
+      "llm_learning_rate": 1.1121524183382568e-06,
+      "loss": 0.7948,
+      "output_text_mask_loss": 0.0014,
+      "predictor_learning_rate": 0.00011121524183382567,
+      "step": 6800
+    },
+    {
+      "epoch": 0.7,
+      "gumbel_tau": 0.201,
+      "image_masks_loss": 0.0126,
+      "llm_learning_rate": 1.1114624547679114e-06,
+      "loss": 0.742,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 0.00011114624547679114,
+      "step": 6801
+    },
+    {
+      "epoch": 0.7,
+      "gumbel_tau": 0.201,
+      "image_masks_loss": 0.0053,
+      "llm_learning_rate": 1.1107726441064137e-06,
+      "loss": 0.7761,
+      "output_text_mask_loss": 0.0456,
+      "predictor_learning_rate": 0.00011107726441064137,
+      "step": 6802
+    },
+    {
+      "epoch": 0.7,
+      "gumbel_tau": 0.2009,
+      "image_masks_loss": 0.0076,
+      "llm_learning_rate": 1.1100829864297266e-06,
+      "loss": 0.6315,
+      "output_text_mask_loss": 0.0012,
+      "predictor_learning_rate": 0.00011100829864297266,
+      "step": 6803
+    },
+    {
+      "epoch": 0.7,
+      "gumbel_tau": 0.2009,
+      "image_masks_loss": 0.0107,
+      "llm_learning_rate": 1.109393481813797e-06,
+      "loss": 0.7507,
+      "output_text_mask_loss": 0.0231,
+      "predictor_learning_rate": 0.00011093934818137969,
+      "step": 6804
+    },
+    {
+      "epoch": 0.7,
+      "gumbel_tau": 0.2008,
+      "image_masks_loss": 0.0067,
+      "llm_learning_rate": 1.1087041303345547e-06,
+      "loss": 0.7854,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.00011087041303345546,
+      "step": 6805
+    },
+    {
+      "epoch": 0.7,
+      "gumbel_tau": 0.2008,
+      "image_masks_loss": 0.0117,
+      "llm_learning_rate": 1.1080149320679123e-06,
+      "loss": 0.6438,
+      "output_text_mask_loss": 0.0401,
+      "predictor_learning_rate": 0.00011080149320679123,
+      "step": 6806
+    },
+    {
+      "epoch": 0.7,
+      "gumbel_tau": 0.2008,
+      "image_masks_loss": 0.0133,
+      "llm_learning_rate": 1.1073258870897664e-06,
+      "loss": 0.6301,
+      "output_text_mask_loss": 0.0062,
+      "predictor_learning_rate": 0.00011073258870897662,
+      "step": 6807
+    },
+    {
+      "epoch": 0.7,
+      "gumbel_tau": 0.2007,
+      "image_masks_loss": 0.0083,
+      "llm_learning_rate": 1.1066369954759956e-06,
+      "loss": 0.7504,
+      "output_text_mask_loss": 0.0452,
+      "predictor_learning_rate": 0.00011066369954759956,
+      "step": 6808
+    },
+    {
+      "epoch": 0.7,
+      "gumbel_tau": 0.2007,
+      "image_masks_loss": 0.0104,
+      "llm_learning_rate": 1.1059482573024627e-06,
+      "loss": 0.5836,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00011059482573024626,
+      "step": 6809
+    },
+    {
+      "epoch": 0.7,
+      "gumbel_tau": 0.2006,
+      "image_masks_loss": 0.02,
+      "llm_learning_rate": 1.105259672645013e-06,
+      "loss": 0.7619,
+      "output_text_mask_loss": 0.0061,
+      "predictor_learning_rate": 0.00011052596726450129,
+      "step": 6810
+    },
+    {
+      "epoch": 0.7,
+      "gumbel_tau": 0.2006,
+      "image_masks_loss": 0.0051,
+      "llm_learning_rate": 1.1045712415794734e-06,
+      "loss": 0.7536,
+      "output_text_mask_loss": 0.0039,
+      "predictor_learning_rate": 0.00011045712415794734,
+      "step": 6811
+    },
+    {
+      "epoch": 0.7,
+      "gumbel_tau": 0.2005,
+      "image_masks_loss": 0.0099,
+      "llm_learning_rate": 1.103882964181658e-06,
+      "loss": 0.678,
+      "output_text_mask_loss": 0.0076,
+      "predictor_learning_rate": 0.00011038829641816581,
+      "step": 6812
+    },
+    {
+      "epoch": 0.7,
+      "gumbel_tau": 0.2005,
+      "image_masks_loss": 0.0103,
+      "llm_learning_rate": 1.1031948405273604e-06,
+      "loss": 0.7463,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.00011031948405273603,
+      "step": 6813
+    },
+    {
+      "epoch": 0.7,
+      "gumbel_tau": 0.2004,
+      "image_masks_loss": 0.0095,
+      "llm_learning_rate": 1.1025068706923583e-06,
+      "loss": 0.7842,
+      "output_text_mask_loss": 0.0075,
+      "predictor_learning_rate": 0.00011025068706923583,
+      "step": 6814
+    },
+    {
+      "epoch": 0.7,
+      "gumbel_tau": 0.2004,
+      "image_masks_loss": 0.0065,
+      "llm_learning_rate": 1.1018190547524122e-06,
+      "loss": 0.7601,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.00011018190547524121,
+      "step": 6815
+    },
+    {
+      "epoch": 0.7,
+      "gumbel_tau": 0.2003,
+      "image_masks_loss": 0.0129,
+      "llm_learning_rate": 1.1011313927832663e-06,
+      "loss": 0.711,
+      "output_text_mask_loss": 0.0105,
+      "predictor_learning_rate": 0.00011011313927832662,
+      "step": 6816
+    },
+    {
+      "epoch": 0.7,
+      "gumbel_tau": 0.2003,
+      "image_masks_loss": 0.0065,
+      "llm_learning_rate": 1.1004438848606471e-06,
+      "loss": 0.7155,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00011004438848606471,
+      "step": 6817
+    },
+    {
+      "epoch": 0.7,
+      "gumbel_tau": 0.2002,
+      "image_masks_loss": 0.0086,
+      "llm_learning_rate": 1.0997565310602637e-06,
+      "loss": 0.8041,
+      "output_text_mask_loss": 0.0187,
+      "predictor_learning_rate": 0.00010997565310602636,
+      "step": 6818
+    },
+    {
+      "epoch": 0.7,
+      "gumbel_tau": 0.2002,
+      "image_masks_loss": 0.02,
+      "llm_learning_rate": 1.0990693314578114e-06,
+      "loss": 0.7802,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00010990693314578112,
+      "step": 6819
+    },
+    {
+      "epoch": 0.7,
+      "gumbel_tau": 0.2001,
+      "image_masks_loss": 0.0062,
+      "llm_learning_rate": 1.0983822861289642e-06,
+      "loss": 0.8082,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.00010983822861289641,
+      "step": 6820
+    },
+    {
+      "epoch": 0.7,
+      "gumbel_tau": 0.2001,
+      "image_masks_loss": 0.0049,
+      "llm_learning_rate": 1.097695395149382e-06,
+      "loss": 0.7406,
+      "output_text_mask_loss": 0.0071,
+      "predictor_learning_rate": 0.00010976953951493818,
+      "step": 6821
+    },
+    {
+      "epoch": 0.7,
+      "gumbel_tau": 0.2,
+      "image_masks_loss": 0.0071,
+      "llm_learning_rate": 1.0970086585947062e-06,
+      "loss": 0.7919,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.00010970086585947061,
+      "step": 6822
+    },
+    {
+      "epoch": 0.7,
+      "gumbel_tau": 0.2,
+      "image_masks_loss": 0.0079,
+      "llm_learning_rate": 1.0963220765405623e-06,
+      "loss": 0.6658,
+      "output_text_mask_loss": 0.0322,
+      "predictor_learning_rate": 0.0001096322076540562,
+      "step": 6823
+    },
+    {
+      "epoch": 0.7,
+      "gumbel_tau": 0.2,
+      "image_masks_loss": 0.0157,
+      "llm_learning_rate": 1.0956356490625574e-06,
+      "loss": 0.7493,
+      "output_text_mask_loss": 0.0071,
+      "predictor_learning_rate": 0.00010956356490625575,
+      "step": 6824
+    },
+    {
+      "epoch": 0.7,
+      "gumbel_tau": 0.1999,
+      "image_masks_loss": 0.0045,
+      "llm_learning_rate": 1.094949376236283e-06,
+      "loss": 0.7466,
+      "output_text_mask_loss": 0.0226,
+      "predictor_learning_rate": 0.00010949493762362831,
+      "step": 6825
+    },
+    {
+      "epoch": 0.7,
+      "gumbel_tau": 0.1999,
+      "image_masks_loss": 0.0087,
+      "llm_learning_rate": 1.0942632581373133e-06,
+      "loss": 0.7706,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.00010942632581373132,
+      "step": 6826
+    },
+    {
+      "epoch": 0.7,
+      "gumbel_tau": 0.1998,
+      "image_masks_loss": 0.0125,
+      "llm_learning_rate": 1.0935772948412045e-06,
+      "loss": 0.6436,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.00010935772948412045,
+      "step": 6827
+    },
+    {
+      "epoch": 0.7,
+      "gumbel_tau": 0.1998,
+      "image_masks_loss": 0.0137,
+      "llm_learning_rate": 1.0928914864234965e-06,
+      "loss": 0.7506,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.00010928914864234965,
+      "step": 6828
+    },
+    {
+      "epoch": 0.7,
+      "gumbel_tau": 0.1997,
+      "image_masks_loss": 0.0053,
+      "llm_learning_rate": 1.0922058329597128e-06,
+      "loss": 0.6151,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 0.00010922058329597126,
+      "step": 6829
+    },
+    {
+      "epoch": 0.7,
+      "gumbel_tau": 0.1997,
+      "image_masks_loss": 0.0148,
+      "llm_learning_rate": 1.0915203345253581e-06,
+      "loss": 0.7709,
+      "output_text_mask_loss": 0.017,
+      "predictor_learning_rate": 0.0001091520334525358,
+      "step": 6830
+    },
+    {
+      "epoch": 0.7,
+      "gumbel_tau": 0.1996,
+      "image_masks_loss": 0.0084,
+      "llm_learning_rate": 1.0908349911959204e-06,
+      "loss": 0.6835,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00010908349911959204,
+      "step": 6831
+    },
+    {
+      "epoch": 0.7,
+      "gumbel_tau": 0.1996,
+      "image_masks_loss": 0.0069,
+      "llm_learning_rate": 1.0901498030468736e-06,
+      "loss": 0.7476,
+      "output_text_mask_loss": 0.0134,
+      "predictor_learning_rate": 0.00010901498030468734,
+      "step": 6832
+    },
+    {
+      "epoch": 0.7,
+      "gumbel_tau": 0.1995,
+      "image_masks_loss": 0.009,
+      "llm_learning_rate": 1.0894647701536704e-06,
+      "loss": 0.7377,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00010894647701536704,
+      "step": 6833
+    },
+    {
+      "epoch": 0.7,
+      "gumbel_tau": 0.1995,
+      "image_masks_loss": 0.011,
+      "llm_learning_rate": 1.0887798925917486e-06,
+      "loss": 0.6929,
+      "output_text_mask_loss": 0.008,
+      "predictor_learning_rate": 0.00010887798925917486,
+      "step": 6834
+    },
+    {
+      "epoch": 0.7,
+      "gumbel_tau": 0.1994,
+      "image_masks_loss": 0.003,
+      "llm_learning_rate": 1.0880951704365287e-06,
+      "loss": 0.6947,
+      "output_text_mask_loss": 0.0145,
+      "predictor_learning_rate": 0.00010880951704365285,
+      "step": 6835
+    },
+    {
+      "epoch": 0.7,
+      "gumbel_tau": 0.1994,
+      "image_masks_loss": 0.0155,
+      "llm_learning_rate": 1.0874106037634132e-06,
+      "loss": 0.7301,
+      "output_text_mask_loss": 0.0885,
+      "predictor_learning_rate": 0.00010874106037634132,
+      "step": 6836
+    },
+    {
+      "epoch": 0.7,
+      "gumbel_tau": 0.1993,
+      "image_masks_loss": 0.009,
+      "llm_learning_rate": 1.0867261926477888e-06,
+      "loss": 0.7609,
+      "output_text_mask_loss": 0.0135,
+      "predictor_learning_rate": 0.00010867261926477887,
+      "step": 6837
+    },
+    {
+      "epoch": 0.7,
+      "gumbel_tau": 0.1993,
+      "image_masks_loss": 0.0144,
+      "llm_learning_rate": 1.0860419371650228e-06,
+      "loss": 0.7177,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00010860419371650226,
+      "step": 6838
+    },
+    {
+      "epoch": 0.7,
+      "gumbel_tau": 0.1992,
+      "image_masks_loss": 0.0067,
+      "llm_learning_rate": 1.0853578373904691e-06,
+      "loss": 0.7035,
+      "output_text_mask_loss": 0.0041,
+      "predictor_learning_rate": 0.0001085357837390469,
+      "step": 6839
+    },
+    {
+      "epoch": 0.7,
+      "gumbel_tau": 0.1992,
+      "image_masks_loss": 0.0077,
+      "llm_learning_rate": 1.0846738933994614e-06,
+      "loss": 0.7563,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00010846738933994612,
+      "step": 6840
+    },
+    {
+      "epoch": 0.7,
+      "gumbel_tau": 0.1992,
+      "image_masks_loss": 0.0119,
+      "llm_learning_rate": 1.083990105267317e-06,
+      "loss": 0.7663,
+      "output_text_mask_loss": 0.0175,
+      "predictor_learning_rate": 0.00010839901052673168,
+      "step": 6841
+    },
+    {
+      "epoch": 0.7,
+      "gumbel_tau": 0.1991,
+      "image_masks_loss": 0.0058,
+      "llm_learning_rate": 1.0833064730693364e-06,
+      "loss": 0.764,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.00010833064730693363,
+      "step": 6842
+    },
+    {
+      "epoch": 0.7,
+      "gumbel_tau": 0.1991,
+      "image_masks_loss": 0.0127,
+      "llm_learning_rate": 1.0826229968808022e-06,
+      "loss": 0.7183,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.00010826229968808021,
+      "step": 6843
+    },
+    {
+      "epoch": 0.7,
+      "gumbel_tau": 0.199,
+      "image_masks_loss": 0.0174,
+      "llm_learning_rate": 1.081939676776981e-06,
+      "loss": 0.7255,
+      "output_text_mask_loss": 0.0076,
+      "predictor_learning_rate": 0.0001081939676776981,
+      "step": 6844
+    },
+    {
+      "epoch": 0.7,
+      "gumbel_tau": 0.199,
+      "image_masks_loss": 0.0046,
+      "llm_learning_rate": 1.081256512833121e-06,
+      "loss": 0.6795,
+      "output_text_mask_loss": 0.0208,
+      "predictor_learning_rate": 0.00010812565128331208,
+      "step": 6845
+    },
+    {
+      "epoch": 0.7,
+      "gumbel_tau": 0.1989,
+      "image_masks_loss": 0.0234,
+      "llm_learning_rate": 1.080573505124454e-06,
+      "loss": 0.79,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 0.00010805735051244539,
+      "step": 6846
+    },
+    {
+      "epoch": 0.7,
+      "gumbel_tau": 0.1989,
+      "image_masks_loss": 0.0122,
+      "llm_learning_rate": 1.0798906537261945e-06,
+      "loss": 0.7625,
+      "output_text_mask_loss": 0.0081,
+      "predictor_learning_rate": 0.00010798906537261943,
+      "step": 6847
+    },
+    {
+      "epoch": 0.7,
+      "gumbel_tau": 0.1988,
+      "image_masks_loss": 0.013,
+      "llm_learning_rate": 1.079207958713539e-06,
+      "loss": 0.752,
+      "output_text_mask_loss": 0.0162,
+      "predictor_learning_rate": 0.0001079207958713539,
+      "step": 6848
+    },
+    {
+      "epoch": 0.7,
+      "gumbel_tau": 0.1988,
+      "image_masks_loss": 0.0063,
+      "llm_learning_rate": 1.078525420161668e-06,
+      "loss": 0.7437,
+      "output_text_mask_loss": 0.0141,
+      "predictor_learning_rate": 0.00010785254201616678,
+      "step": 6849
+    },
+    {
+      "epoch": 0.7,
+      "gumbel_tau": 0.1987,
+      "image_masks_loss": 0.0075,
+      "llm_learning_rate": 1.077843038145743e-06,
+      "loss": 0.737,
+      "output_text_mask_loss": 0.0233,
+      "predictor_learning_rate": 0.0001077843038145743,
+      "step": 6850
+    },
+    {
+      "epoch": 0.7,
+      "gumbel_tau": 0.1987,
+      "image_masks_loss": 0.0122,
+      "llm_learning_rate": 1.0771608127409113e-06,
+      "loss": 0.7995,
+      "output_text_mask_loss": 0.001,
+      "predictor_learning_rate": 0.00010771608127409113,
+      "step": 6851
+    },
+    {
+      "epoch": 0.7,
+      "gumbel_tau": 0.1986,
+      "image_masks_loss": 0.0157,
+      "llm_learning_rate": 1.0764787440223002e-06,
+      "loss": 0.6511,
+      "output_text_mask_loss": 0.0096,
+      "predictor_learning_rate": 0.00010764787440223,
+      "step": 6852
+    },
+    {
+      "epoch": 0.7,
+      "gumbel_tau": 0.1986,
+      "image_masks_loss": 0.0045,
+      "llm_learning_rate": 1.0757968320650202e-06,
+      "loss": 0.7362,
+      "output_text_mask_loss": 0.0238,
+      "predictor_learning_rate": 0.00010757968320650203,
+      "step": 6853
+    },
+    {
+      "epoch": 0.7,
+      "gumbel_tau": 0.1985,
+      "image_masks_loss": 0.0123,
+      "llm_learning_rate": 1.075115076944166e-06,
+      "loss": 0.7588,
+      "output_text_mask_loss": 0.0824,
+      "predictor_learning_rate": 0.00010751150769441659,
+      "step": 6854
+    },
+    {
+      "epoch": 0.7,
+      "gumbel_tau": 0.1985,
+      "image_masks_loss": 0.0069,
+      "llm_learning_rate": 1.0744334787348129e-06,
+      "loss": 0.7431,
+      "output_text_mask_loss": 0.0231,
+      "predictor_learning_rate": 0.00010744334787348129,
+      "step": 6855
+    },
+    {
+      "epoch": 0.7,
+      "gumbel_tau": 0.1984,
+      "image_masks_loss": 0.007,
+      "llm_learning_rate": 1.0737520375120206e-06,
+      "loss": 0.7335,
+      "output_text_mask_loss": 0.0169,
+      "predictor_learning_rate": 0.00010737520375120205,
+      "step": 6856
+    },
+    {
+      "epoch": 0.7,
+      "gumbel_tau": 0.1984,
+      "image_masks_loss": 0.0083,
+      "llm_learning_rate": 1.0730707533508297e-06,
+      "loss": 0.6879,
+      "output_text_mask_loss": 0.0511,
+      "predictor_learning_rate": 0.00010730707533508296,
+      "step": 6857
+    },
+    {
+      "epoch": 0.7,
+      "gumbel_tau": 0.1984,
+      "image_masks_loss": 0.017,
+      "llm_learning_rate": 1.0723896263262667e-06,
+      "loss": 0.7916,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00010723896263262667,
+      "step": 6858
+    },
+    {
+      "epoch": 0.7,
+      "gumbel_tau": 0.1983,
+      "image_masks_loss": 0.0105,
+      "llm_learning_rate": 1.071708656513338e-06,
+      "loss": 0.674,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.0001071708656513338,
+      "step": 6859
+    },
+    {
+      "epoch": 0.7,
+      "gumbel_tau": 0.1983,
+      "image_masks_loss": 0.0121,
+      "llm_learning_rate": 1.0710278439870336e-06,
+      "loss": 0.7686,
+      "output_text_mask_loss": 0.0322,
+      "predictor_learning_rate": 0.00010710278439870336,
+      "step": 6860
+    },
+    {
+      "epoch": 0.7,
+      "gumbel_tau": 0.1982,
+      "image_masks_loss": 0.0067,
+      "llm_learning_rate": 1.0703471888223265e-06,
+      "loss": 0.7885,
+      "output_text_mask_loss": 0.0109,
+      "predictor_learning_rate": 0.00010703471888223265,
+      "step": 6861
+    },
+    {
+      "epoch": 0.7,
+      "gumbel_tau": 0.1982,
+      "image_masks_loss": 0.0146,
+      "llm_learning_rate": 1.0696666910941709e-06,
+      "loss": 0.7577,
+      "output_text_mask_loss": 0.009,
+      "predictor_learning_rate": 0.00010696666910941708,
+      "step": 6862
+    },
+    {
+      "epoch": 0.7,
+      "gumbel_tau": 0.1981,
+      "image_masks_loss": 0.0248,
+      "llm_learning_rate": 1.068986350877504e-06,
+      "loss": 0.7377,
+      "output_text_mask_loss": 0.0062,
+      "predictor_learning_rate": 0.00010689863508775038,
+      "step": 6863
+    },
+    {
+      "epoch": 0.7,
+      "gumbel_tau": 0.1981,
+      "image_masks_loss": 0.0105,
+      "llm_learning_rate": 1.0683061682472486e-06,
+      "loss": 0.689,
+      "output_text_mask_loss": 0.005,
+      "predictor_learning_rate": 0.00010683061682472486,
+      "step": 6864
+    },
+    {
+      "epoch": 0.7,
+      "gumbel_tau": 0.198,
+      "image_masks_loss": 0.0099,
+      "llm_learning_rate": 1.0676261432783069e-06,
+      "loss": 0.6622,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 0.00010676261432783069,
+      "step": 6865
+    },
+    {
+      "epoch": 0.7,
+      "gumbel_tau": 0.198,
+      "image_masks_loss": 0.0124,
+      "llm_learning_rate": 1.066946276045565e-06,
+      "loss": 0.6808,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.00010669462760455648,
+      "step": 6866
+    },
+    {
+      "epoch": 0.7,
+      "gumbel_tau": 0.1979,
+      "image_masks_loss": 0.0091,
+      "llm_learning_rate": 1.066266566623891e-06,
+      "loss": 0.7759,
+      "output_text_mask_loss": 0.0124,
+      "predictor_learning_rate": 0.00010662665666238908,
+      "step": 6867
+    },
+    {
+      "epoch": 0.7,
+      "gumbel_tau": 0.1979,
+      "image_masks_loss": 0.0167,
+      "llm_learning_rate": 1.065587015088136e-06,
+      "loss": 0.7394,
+      "output_text_mask_loss": 0.0334,
+      "predictor_learning_rate": 0.00010655870150881358,
+      "step": 6868
+    },
+    {
+      "epoch": 0.7,
+      "gumbel_tau": 0.1978,
+      "image_masks_loss": 0.0028,
+      "llm_learning_rate": 1.064907621513134e-06,
+      "loss": 0.7846,
+      "output_text_mask_loss": 0.0116,
+      "predictor_learning_rate": 0.00010649076215131339,
+      "step": 6869
+    },
+    {
+      "epoch": 0.7,
+      "gumbel_tau": 0.1978,
+      "image_masks_loss": 0.0046,
+      "llm_learning_rate": 1.0642283859737006e-06,
+      "loss": 0.6537,
+      "output_text_mask_loss": 0.0248,
+      "predictor_learning_rate": 0.00010642283859737006,
+      "step": 6870
+    },
+    {
+      "epoch": 0.7,
+      "gumbel_tau": 0.1977,
+      "image_masks_loss": 0.0088,
+      "llm_learning_rate": 1.0635493085446363e-06,
+      "loss": 0.7151,
+      "output_text_mask_loss": 0.005,
+      "predictor_learning_rate": 0.00010635493085446362,
+      "step": 6871
+    },
+    {
+      "epoch": 0.7,
+      "gumbel_tau": 0.1977,
+      "image_masks_loss": 0.0075,
+      "llm_learning_rate": 1.0628703893007213e-06,
+      "loss": 0.8257,
+      "output_text_mask_loss": 0.0187,
+      "predictor_learning_rate": 0.00010628703893007213,
+      "step": 6872
+    },
+    {
+      "epoch": 0.7,
+      "gumbel_tau": 0.1977,
+      "image_masks_loss": 0.0137,
+      "llm_learning_rate": 1.0621916283167205e-06,
+      "loss": 0.7327,
+      "output_text_mask_loss": 0.0231,
+      "predictor_learning_rate": 0.00010621916283167205,
+      "step": 6873
+    },
+    {
+      "epoch": 0.7,
+      "gumbel_tau": 0.1976,
+      "image_masks_loss": 0.0134,
+      "llm_learning_rate": 1.0615130256673803e-06,
+      "loss": 0.7365,
+      "output_text_mask_loss": 0.0309,
+      "predictor_learning_rate": 0.00010615130256673803,
+      "step": 6874
+    },
+    {
+      "epoch": 0.7,
+      "gumbel_tau": 0.1976,
+      "image_masks_loss": 0.0096,
+      "llm_learning_rate": 1.0608345814274296e-06,
+      "loss": 0.7289,
+      "output_text_mask_loss": 0.0278,
+      "predictor_learning_rate": 0.00010608345814274295,
+      "step": 6875
+    },
+    {
+      "epoch": 0.7,
+      "gumbel_tau": 0.1975,
+      "image_masks_loss": 0.0082,
+      "llm_learning_rate": 1.0601562956715802e-06,
+      "loss": 0.794,
+      "output_text_mask_loss": 0.013,
+      "predictor_learning_rate": 0.00010601562956715802,
+      "step": 6876
+    },
+    {
+      "epoch": 0.7,
+      "gumbel_tau": 0.1975,
+      "image_masks_loss": 0.0084,
+      "llm_learning_rate": 1.0594781684745263e-06,
+      "loss": 0.7419,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00010594781684745261,
+      "step": 6877
+    },
+    {
+      "epoch": 0.7,
+      "gumbel_tau": 0.1974,
+      "image_masks_loss": 0.0076,
+      "llm_learning_rate": 1.0588001999109453e-06,
+      "loss": 0.7991,
+      "output_text_mask_loss": 0.0355,
+      "predictor_learning_rate": 0.00010588001999109454,
+      "step": 6878
+    },
+    {
+      "epoch": 0.7,
+      "gumbel_tau": 0.1974,
+      "image_masks_loss": 0.0067,
+      "llm_learning_rate": 1.0581223900554972e-06,
+      "loss": 0.7225,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 0.00010581223900554971,
+      "step": 6879
+    },
+    {
+      "epoch": 0.7,
+      "gumbel_tau": 0.1973,
+      "image_masks_loss": 0.0049,
+      "llm_learning_rate": 1.0574447389828235e-06,
+      "loss": 0.7561,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 0.00010574447389828235,
+      "step": 6880
+    },
+    {
+      "epoch": 0.71,
+      "gumbel_tau": 0.1973,
+      "image_masks_loss": 0.0066,
+      "llm_learning_rate": 1.0567672467675473e-06,
+      "loss": 0.7499,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.00010567672467675474,
+      "step": 6881
+    },
+    {
+      "epoch": 0.71,
+      "gumbel_tau": 0.1972,
+      "image_masks_loss": 0.0065,
+      "llm_learning_rate": 1.0560899134842761e-06,
+      "loss": 0.7336,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0001056089913484276,
+      "step": 6882
+    },
+    {
+      "epoch": 0.71,
+      "gumbel_tau": 0.1972,
+      "image_masks_loss": 0.0077,
+      "llm_learning_rate": 1.0554127392076008e-06,
+      "loss": 0.7113,
+      "output_text_mask_loss": 0.0275,
+      "predictor_learning_rate": 0.00010554127392076007,
+      "step": 6883
+    },
+    {
+      "epoch": 0.71,
+      "gumbel_tau": 0.1971,
+      "image_masks_loss": 0.0054,
+      "llm_learning_rate": 1.054735724012092e-06,
+      "loss": 0.7052,
+      "output_text_mask_loss": 0.0446,
+      "predictor_learning_rate": 0.00010547357240120919,
+      "step": 6884
+    },
+    {
+      "epoch": 0.71,
+      "gumbel_tau": 0.1971,
+      "image_masks_loss": 0.0081,
+      "llm_learning_rate": 1.0540588679723046e-06,
+      "loss": 0.6995,
+      "output_text_mask_loss": 0.0041,
+      "predictor_learning_rate": 0.00010540588679723045,
+      "step": 6885
+    },
+    {
+      "epoch": 0.71,
+      "gumbel_tau": 0.197,
+      "image_masks_loss": 0.0093,
+      "llm_learning_rate": 1.0533821711627753e-06,
+      "loss": 0.7792,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.00010533821711627753,
+      "step": 6886
+    },
+    {
+      "epoch": 0.71,
+      "gumbel_tau": 0.197,
+      "image_masks_loss": 0.0061,
+      "llm_learning_rate": 1.0527056336580236e-06,
+      "loss": 0.7375,
+      "output_text_mask_loss": 0.0047,
+      "predictor_learning_rate": 0.00010527056336580235,
+      "step": 6887
+    },
+    {
+      "epoch": 0.71,
+      "gumbel_tau": 0.197,
+      "image_masks_loss": 0.0087,
+      "llm_learning_rate": 1.0520292555325512e-06,
+      "loss": 0.8576,
+      "output_text_mask_loss": 0.0041,
+      "predictor_learning_rate": 0.00010520292555325512,
+      "step": 6888
+    },
+    {
+      "epoch": 0.71,
+      "gumbel_tau": 0.1969,
+      "image_masks_loss": 0.0106,
+      "llm_learning_rate": 1.0513530368608416e-06,
+      "loss": 0.6823,
+      "output_text_mask_loss": 0.0431,
+      "predictor_learning_rate": 0.00010513530368608415,
+      "step": 6889
+    },
+    {
+      "epoch": 0.71,
+      "gumbel_tau": 0.1969,
+      "image_masks_loss": 0.0097,
+      "llm_learning_rate": 1.0506769777173632e-06,
+      "loss": 0.668,
+      "output_text_mask_loss": 0.0086,
+      "predictor_learning_rate": 0.00010506769777173631,
+      "step": 6890
+    },
+    {
+      "epoch": 0.71,
+      "gumbel_tau": 0.1968,
+      "image_masks_loss": 0.0051,
+      "llm_learning_rate": 1.050001078176564e-06,
+      "loss": 0.7066,
+      "output_text_mask_loss": 0.0101,
+      "predictor_learning_rate": 0.0001050001078176564,
+      "step": 6891
+    },
+    {
+      "epoch": 0.71,
+      "gumbel_tau": 0.1968,
+      "image_masks_loss": 0.0092,
+      "llm_learning_rate": 1.0493253383128763e-06,
+      "loss": 0.7297,
+      "output_text_mask_loss": 0.0081,
+      "predictor_learning_rate": 0.00010493253383128762,
+      "step": 6892
+    },
+    {
+      "epoch": 0.71,
+      "gumbel_tau": 0.1967,
+      "image_masks_loss": 0.0244,
+      "llm_learning_rate": 1.0486497582007132e-06,
+      "loss": 0.7663,
+      "output_text_mask_loss": 0.053,
+      "predictor_learning_rate": 0.00010486497582007132,
+      "step": 6893
+    },
+    {
+      "epoch": 0.71,
+      "gumbel_tau": 0.1967,
+      "image_masks_loss": 0.014,
+      "llm_learning_rate": 1.0479743379144715e-06,
+      "loss": 0.7569,
+      "output_text_mask_loss": 0.008,
+      "predictor_learning_rate": 0.00010479743379144716,
+      "step": 6894
+    },
+    {
+      "epoch": 0.71,
+      "gumbel_tau": 0.1966,
+      "image_masks_loss": 0.0057,
+      "llm_learning_rate": 1.0472990775285302e-06,
+      "loss": 0.7755,
+      "output_text_mask_loss": 0.0633,
+      "predictor_learning_rate": 0.00010472990775285302,
+      "step": 6895
+    },
+    {
+      "epoch": 0.71,
+      "gumbel_tau": 0.1966,
+      "image_masks_loss": 0.0086,
+      "llm_learning_rate": 1.046623977117249e-06,
+      "loss": 0.7045,
+      "output_text_mask_loss": 0.0204,
+      "predictor_learning_rate": 0.00010466239771172489,
+      "step": 6896
+    },
+    {
+      "epoch": 0.71,
+      "gumbel_tau": 0.1965,
+      "image_masks_loss": 0.0111,
+      "llm_learning_rate": 1.045949036754974e-06,
+      "loss": 0.7458,
+      "output_text_mask_loss": 0.0549,
+      "predictor_learning_rate": 0.0001045949036754974,
+      "step": 6897
+    },
+    {
+      "epoch": 0.71,
+      "gumbel_tau": 0.1965,
+      "image_masks_loss": 0.0039,
+      "llm_learning_rate": 1.0452742565160304e-06,
+      "loss": 0.6596,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00010452742565160303,
+      "step": 6898
+    },
+    {
+      "epoch": 0.71,
+      "gumbel_tau": 0.1964,
+      "image_masks_loss": 0.0098,
+      "llm_learning_rate": 1.0445996364747254e-06,
+      "loss": 0.7214,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.00010445996364747254,
+      "step": 6899
+    },
+    {
+      "epoch": 0.71,
+      "gumbel_tau": 0.1964,
+      "image_masks_loss": 0.0101,
+      "llm_learning_rate": 1.0439251767053499e-06,
+      "loss": 0.7531,
+      "output_text_mask_loss": 0.0511,
+      "predictor_learning_rate": 0.00010439251767053498,
+      "step": 6900
+    },
+    {
+      "epoch": 0.71,
+      "gumbel_tau": 0.1964,
+      "image_masks_loss": 0.0036,
+      "llm_learning_rate": 1.0432508772821772e-06,
+      "loss": 0.7031,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.00010432508772821772,
+      "step": 6901
+    },
+    {
+      "epoch": 0.71,
+      "gumbel_tau": 0.1963,
+      "image_masks_loss": 0.0094,
+      "llm_learning_rate": 1.042576738279462e-06,
+      "loss": 0.7502,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.00010425767382794621,
+      "step": 6902
+    },
+    {
+      "epoch": 0.71,
+      "gumbel_tau": 0.1963,
+      "image_masks_loss": 0.0093,
+      "llm_learning_rate": 1.0419027597714434e-06,
+      "loss": 0.785,
+      "output_text_mask_loss": 0.0045,
+      "predictor_learning_rate": 0.00010419027597714434,
+      "step": 6903
+    },
+    {
+      "epoch": 0.71,
+      "gumbel_tau": 0.1962,
+      "image_masks_loss": 0.0113,
+      "llm_learning_rate": 1.0412289418323409e-06,
+      "loss": 0.7358,
+      "output_text_mask_loss": 0.0231,
+      "predictor_learning_rate": 0.00010412289418323409,
+      "step": 6904
+    },
+    {
+      "epoch": 0.71,
+      "gumbel_tau": 0.1962,
+      "image_masks_loss": 0.003,
+      "llm_learning_rate": 1.0405552845363565e-06,
+      "loss": 0.7284,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00010405552845363563,
+      "step": 6905
+    },
+    {
+      "epoch": 0.71,
+      "gumbel_tau": 0.1961,
+      "image_masks_loss": 0.0152,
+      "llm_learning_rate": 1.039881787957675e-06,
+      "loss": 0.7612,
+      "output_text_mask_loss": 0.0139,
+      "predictor_learning_rate": 0.00010398817879576749,
+      "step": 6906
+    },
+    {
+      "epoch": 0.71,
+      "gumbel_tau": 0.1961,
+      "image_masks_loss": 0.0091,
+      "llm_learning_rate": 1.0392084521704632e-06,
+      "loss": 0.788,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00010392084521704632,
+      "step": 6907
+    },
+    {
+      "epoch": 0.71,
+      "gumbel_tau": 0.196,
+      "image_masks_loss": 0.0056,
+      "llm_learning_rate": 1.0385352772488704e-06,
+      "loss": 0.7467,
+      "output_text_mask_loss": 0.0041,
+      "predictor_learning_rate": 0.00010385352772488705,
+      "step": 6908
+    },
+    {
+      "epoch": 0.71,
+      "gumbel_tau": 0.196,
+      "image_masks_loss": 0.0052,
+      "llm_learning_rate": 1.0378622632670274e-06,
+      "loss": 0.7161,
+      "output_text_mask_loss": 0.0126,
+      "predictor_learning_rate": 0.00010378622632670273,
+      "step": 6909
+    },
+    {
+      "epoch": 0.71,
+      "gumbel_tau": 0.1959,
+      "image_masks_loss": 0.0117,
+      "llm_learning_rate": 1.0371894102990496e-06,
+      "loss": 0.7245,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.00010371894102990494,
+      "step": 6910
+    },
+    {
+      "epoch": 0.71,
+      "gumbel_tau": 0.1959,
+      "image_masks_loss": 0.0147,
+      "llm_learning_rate": 1.0365167184190319e-06,
+      "loss": 0.7805,
+      "output_text_mask_loss": 0.0128,
+      "predictor_learning_rate": 0.00010365167184190318,
+      "step": 6911
+    },
+    {
+      "epoch": 0.71,
+      "gumbel_tau": 0.1958,
+      "image_masks_loss": 0.009,
+      "llm_learning_rate": 1.0358441877010534e-06,
+      "loss": 0.7577,
+      "output_text_mask_loss": 0.0126,
+      "predictor_learning_rate": 0.00010358441877010532,
+      "step": 6912
+    },
+    {
+      "epoch": 0.71,
+      "gumbel_tau": 0.1958,
+      "image_masks_loss": 0.0054,
+      "llm_learning_rate": 1.0351718182191741e-06,
+      "loss": 0.7233,
+      "output_text_mask_loss": 0.013,
+      "predictor_learning_rate": 0.00010351718182191742,
+      "step": 6913
+    },
+    {
+      "epoch": 0.71,
+      "gumbel_tau": 0.1957,
+      "image_masks_loss": 0.0072,
+      "llm_learning_rate": 1.0344996100474369e-06,
+      "loss": 0.7392,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.00010344996100474368,
+      "step": 6914
+    },
+    {
+      "epoch": 0.71,
+      "gumbel_tau": 0.1957,
+      "image_masks_loss": 0.0044,
+      "llm_learning_rate": 1.0338275632598668e-06,
+      "loss": 0.713,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.00010338275632598668,
+      "step": 6915
+    },
+    {
+      "epoch": 0.71,
+      "gumbel_tau": 0.1957,
+      "image_masks_loss": 0.0128,
+      "llm_learning_rate": 1.0331556779304705e-06,
+      "loss": 0.7386,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 0.00010331556779304705,
+      "step": 6916
+    },
+    {
+      "epoch": 0.71,
+      "gumbel_tau": 0.1956,
+      "image_masks_loss": 0.0108,
+      "llm_learning_rate": 1.0324839541332399e-06,
+      "loss": 0.7227,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 0.00010324839541332398,
+      "step": 6917
+    },
+    {
+      "epoch": 0.71,
+      "gumbel_tau": 0.1956,
+      "image_masks_loss": 0.0108,
+      "llm_learning_rate": 1.0318123919421443e-06,
+      "loss": 0.7411,
+      "output_text_mask_loss": 0.0025,
+      "predictor_learning_rate": 0.00010318123919421443,
+      "step": 6918
+    },
+    {
+      "epoch": 0.71,
+      "gumbel_tau": 0.1955,
+      "image_masks_loss": 0.0085,
+      "llm_learning_rate": 1.0311409914311382e-06,
+      "loss": 0.7516,
+      "output_text_mask_loss": 0.0467,
+      "predictor_learning_rate": 0.00010311409914311382,
+      "step": 6919
+    },
+    {
+      "epoch": 0.71,
+      "gumbel_tau": 0.1955,
+      "image_masks_loss": 0.0112,
+      "llm_learning_rate": 1.030469752674158e-06,
+      "loss": 0.6703,
+      "output_text_mask_loss": 0.0172,
+      "predictor_learning_rate": 0.0001030469752674158,
+      "step": 6920
+    },
+    {
+      "epoch": 0.71,
+      "gumbel_tau": 0.1954,
+      "image_masks_loss": 0.0055,
+      "llm_learning_rate": 1.029798675745121e-06,
+      "loss": 0.7914,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.00010297986757451209,
+      "step": 6921
+    },
+    {
+      "epoch": 0.71,
+      "gumbel_tau": 0.1954,
+      "image_masks_loss": 0.0134,
+      "llm_learning_rate": 1.0291277607179292e-06,
+      "loss": 0.7442,
+      "output_text_mask_loss": 0.0004,
+      "predictor_learning_rate": 0.00010291277607179292,
+      "step": 6922
+    },
+    {
+      "epoch": 0.71,
+      "gumbel_tau": 0.1953,
+      "image_masks_loss": 0.0063,
+      "llm_learning_rate": 1.028457007666465e-06,
+      "loss": 0.726,
+      "output_text_mask_loss": 0.0277,
+      "predictor_learning_rate": 0.00010284570076664649,
+      "step": 6923
+    },
+    {
+      "epoch": 0.71,
+      "gumbel_tau": 0.1953,
+      "image_masks_loss": 0.0064,
+      "llm_learning_rate": 1.0277864166645924e-06,
+      "loss": 0.7148,
+      "output_text_mask_loss": 0.0058,
+      "predictor_learning_rate": 0.00010277864166645922,
+      "step": 6924
+    },
+    {
+      "epoch": 0.71,
+      "gumbel_tau": 0.1952,
+      "image_masks_loss": 0.0108,
+      "llm_learning_rate": 1.0271159877861592e-06,
+      "loss": 0.6542,
+      "output_text_mask_loss": 0.0165,
+      "predictor_learning_rate": 0.00010271159877861591,
+      "step": 6925
+    },
+    {
+      "epoch": 0.71,
+      "gumbel_tau": 0.1952,
+      "image_masks_loss": 0.0098,
+      "llm_learning_rate": 1.026445721104994e-06,
+      "loss": 0.8141,
+      "output_text_mask_loss": 0.0116,
+      "predictor_learning_rate": 0.00010264457211049938,
+      "step": 6926
+    },
+    {
+      "epoch": 0.71,
+      "gumbel_tau": 0.1952,
+      "image_masks_loss": 0.0139,
+      "llm_learning_rate": 1.0257756166949082e-06,
+      "loss": 0.7549,
+      "output_text_mask_loss": 0.0141,
+      "predictor_learning_rate": 0.00010257756166949081,
+      "step": 6927
+    },
+    {
+      "epoch": 0.71,
+      "gumbel_tau": 0.1951,
+      "image_masks_loss": 0.0022,
+      "llm_learning_rate": 1.0251056746296944e-06,
+      "loss": 0.7594,
+      "output_text_mask_loss": 0.0229,
+      "predictor_learning_rate": 0.00010251056746296944,
+      "step": 6928
+    },
+    {
+      "epoch": 0.71,
+      "gumbel_tau": 0.1951,
+      "image_masks_loss": 0.0063,
+      "llm_learning_rate": 1.0244358949831296e-06,
+      "loss": 0.8474,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.00010244358949831295,
+      "step": 6929
+    },
+    {
+      "epoch": 0.71,
+      "gumbel_tau": 0.195,
+      "image_masks_loss": 0.0051,
+      "llm_learning_rate": 1.0237662778289708e-06,
+      "loss": 0.7181,
+      "output_text_mask_loss": 0.0234,
+      "predictor_learning_rate": 0.00010237662778289707,
+      "step": 6930
+    },
+    {
+      "epoch": 0.71,
+      "gumbel_tau": 0.195,
+      "image_masks_loss": 0.0072,
+      "llm_learning_rate": 1.023096823240958e-06,
+      "loss": 0.753,
+      "output_text_mask_loss": 0.0179,
+      "predictor_learning_rate": 0.0001023096823240958,
+      "step": 6931
+    },
+    {
+      "epoch": 0.71,
+      "gumbel_tau": 0.1949,
+      "image_masks_loss": 0.0057,
+      "llm_learning_rate": 1.0224275312928123e-06,
+      "loss": 0.7265,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.00010224275312928122,
+      "step": 6932
+    },
+    {
+      "epoch": 0.71,
+      "gumbel_tau": 0.1949,
+      "image_masks_loss": 0.0078,
+      "llm_learning_rate": 1.0217584020582384e-06,
+      "loss": 0.7886,
+      "output_text_mask_loss": 0.0148,
+      "predictor_learning_rate": 0.00010217584020582382,
+      "step": 6933
+    },
+    {
+      "epoch": 0.71,
+      "gumbel_tau": 0.1948,
+      "image_masks_loss": 0.0017,
+      "llm_learning_rate": 1.0210894356109214e-06,
+      "loss": 0.7735,
+      "output_text_mask_loss": 0.0854,
+      "predictor_learning_rate": 0.00010210894356109213,
+      "step": 6934
+    },
+    {
+      "epoch": 0.71,
+      "gumbel_tau": 0.1948,
+      "image_masks_loss": 0.0075,
+      "llm_learning_rate": 1.0204206320245302e-06,
+      "loss": 0.7729,
+      "output_text_mask_loss": 0.0162,
+      "predictor_learning_rate": 0.00010204206320245302,
+      "step": 6935
+    },
+    {
+      "epoch": 0.71,
+      "gumbel_tau": 0.1947,
+      "image_masks_loss": 0.0065,
+      "llm_learning_rate": 1.0197519913727146e-06,
+      "loss": 0.7176,
+      "output_text_mask_loss": 0.0012,
+      "predictor_learning_rate": 0.00010197519913727146,
+      "step": 6936
+    },
+    {
+      "epoch": 0.71,
+      "gumbel_tau": 0.1947,
+      "image_masks_loss": 0.0036,
+      "llm_learning_rate": 1.0190835137291068e-06,
+      "loss": 0.6922,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00010190835137291068,
+      "step": 6937
+    },
+    {
+      "epoch": 0.71,
+      "gumbel_tau": 0.1946,
+      "image_masks_loss": 0.0118,
+      "llm_learning_rate": 1.018415199167321e-06,
+      "loss": 0.6905,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.0001018415199167321,
+      "step": 6938
+    },
+    {
+      "epoch": 0.71,
+      "gumbel_tau": 0.1946,
+      "image_masks_loss": 0.0084,
+      "llm_learning_rate": 1.0177470477609538e-06,
+      "loss": 0.7652,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.00010177470477609537,
+      "step": 6939
+    },
+    {
+      "epoch": 0.71,
+      "gumbel_tau": 0.1946,
+      "image_masks_loss": 0.0118,
+      "llm_learning_rate": 1.0170790595835834e-06,
+      "loss": 0.7603,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.00010170790595835832,
+      "step": 6940
+    },
+    {
+      "epoch": 0.71,
+      "gumbel_tau": 0.1945,
+      "image_masks_loss": 0.005,
+      "llm_learning_rate": 1.0164112347087688e-06,
+      "loss": 0.6358,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00010164112347087687,
+      "step": 6941
+    },
+    {
+      "epoch": 0.71,
+      "gumbel_tau": 0.1945,
+      "image_masks_loss": 0.0042,
+      "llm_learning_rate": 1.0157435732100545e-06,
+      "loss": 0.7631,
+      "output_text_mask_loss": 0.0185,
+      "predictor_learning_rate": 0.00010157435732100545,
+      "step": 6942
+    },
+    {
+      "epoch": 0.71,
+      "gumbel_tau": 0.1944,
+      "image_masks_loss": 0.0047,
+      "llm_learning_rate": 1.0150760751609642e-06,
+      "loss": 0.6827,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 0.00010150760751609642,
+      "step": 6943
+    },
+    {
+      "epoch": 0.71,
+      "gumbel_tau": 0.1944,
+      "image_masks_loss": 0.0025,
+      "llm_learning_rate": 1.014408740635004e-06,
+      "loss": 0.6673,
+      "output_text_mask_loss": 0.0187,
+      "predictor_learning_rate": 0.00010144087406350039,
+      "step": 6944
+    },
+    {
+      "epoch": 0.71,
+      "gumbel_tau": 0.1943,
+      "image_masks_loss": 0.0107,
+      "llm_learning_rate": 1.0137415697056621e-06,
+      "loss": 0.7479,
+      "output_text_mask_loss": 0.0099,
+      "predictor_learning_rate": 0.0001013741569705662,
+      "step": 6945
+    },
+    {
+      "epoch": 0.71,
+      "gumbel_tau": 0.1943,
+      "image_masks_loss": 0.0073,
+      "llm_learning_rate": 1.0130745624464092e-06,
+      "loss": 0.6658,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 0.00010130745624464091,
+      "step": 6946
+    },
+    {
+      "epoch": 0.71,
+      "gumbel_tau": 0.1942,
+      "image_masks_loss": 0.01,
+      "llm_learning_rate": 1.0124077189306974e-06,
+      "loss": 0.7734,
+      "output_text_mask_loss": 0.0034,
+      "predictor_learning_rate": 0.00010124077189306974,
+      "step": 6947
+    },
+    {
+      "epoch": 0.71,
+      "gumbel_tau": 0.1942,
+      "image_masks_loss": 0.0087,
+      "llm_learning_rate": 1.01174103923196e-06,
+      "loss": 0.7153,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00010117410392319601,
+      "step": 6948
+    },
+    {
+      "epoch": 0.71,
+      "gumbel_tau": 0.1941,
+      "image_masks_loss": 0.0132,
+      "llm_learning_rate": 1.0110745234236157e-06,
+      "loss": 0.7407,
+      "output_text_mask_loss": 0.0172,
+      "predictor_learning_rate": 0.00010110745234236157,
+      "step": 6949
+    },
+    {
+      "epoch": 0.71,
+      "gumbel_tau": 0.1941,
+      "image_masks_loss": 0.0077,
+      "llm_learning_rate": 1.0104081715790611e-06,
+      "loss": 0.6562,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.00010104081715790611,
+      "step": 6950
+    },
+    {
+      "epoch": 0.71,
+      "gumbel_tau": 0.194,
+      "image_masks_loss": 0.017,
+      "llm_learning_rate": 1.0097419837716766e-06,
+      "loss": 0.7135,
+      "output_text_mask_loss": 0.0256,
+      "predictor_learning_rate": 0.00010097419837716765,
+      "step": 6951
+    },
+    {
+      "epoch": 0.71,
+      "gumbel_tau": 0.194,
+      "image_masks_loss": 0.0096,
+      "llm_learning_rate": 1.0090759600748242e-06,
+      "loss": 0.79,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.00010090759600748242,
+      "step": 6952
+    },
+    {
+      "epoch": 0.71,
+      "gumbel_tau": 0.194,
+      "image_masks_loss": 0.0037,
+      "llm_learning_rate": 1.008410100561848e-06,
+      "loss": 0.6284,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 0.00010084101005618479,
+      "step": 6953
+    },
+    {
+      "epoch": 0.71,
+      "gumbel_tau": 0.1939,
+      "image_masks_loss": 0.0056,
+      "llm_learning_rate": 1.007744405306074e-06,
+      "loss": 0.7447,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 0.00010077444053060739,
+      "step": 6954
+    },
+    {
+      "epoch": 0.71,
+      "gumbel_tau": 0.1939,
+      "image_masks_loss": 0.0044,
+      "llm_learning_rate": 1.00707887438081e-06,
+      "loss": 0.6786,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.00010070788743808098,
+      "step": 6955
+    },
+    {
+      "epoch": 0.71,
+      "gumbel_tau": 0.1938,
+      "image_masks_loss": 0.02,
+      "llm_learning_rate": 1.0064135078593454e-06,
+      "loss": 0.6645,
+      "output_text_mask_loss": 0.0035,
+      "predictor_learning_rate": 0.00010064135078593453,
+      "step": 6956
+    },
+    {
+      "epoch": 0.71,
+      "gumbel_tau": 0.1938,
+      "image_masks_loss": 0.0028,
+      "llm_learning_rate": 1.0057483058149526e-06,
+      "loss": 0.7082,
+      "output_text_mask_loss": 0.0242,
+      "predictor_learning_rate": 0.00010057483058149524,
+      "step": 6957
+    },
+    {
+      "epoch": 0.71,
+      "gumbel_tau": 0.1937,
+      "image_masks_loss": 0.0086,
+      "llm_learning_rate": 1.0050832683208842e-06,
+      "loss": 0.8014,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.0001005083268320884,
+      "step": 6958
+    },
+    {
+      "epoch": 0.71,
+      "gumbel_tau": 0.1937,
+      "image_masks_loss": 0.0114,
+      "llm_learning_rate": 1.0044183954503763e-06,
+      "loss": 0.7156,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 0.00010044183954503764,
+      "step": 6959
+    },
+    {
+      "epoch": 0.71,
+      "gumbel_tau": 0.1936,
+      "image_masks_loss": 0.018,
+      "llm_learning_rate": 1.003753687276645e-06,
+      "loss": 0.8016,
+      "output_text_mask_loss": 0.0124,
+      "predictor_learning_rate": 0.0001003753687276645,
+      "step": 6960
+    },
+    {
+      "epoch": 0.71,
+      "gumbel_tau": 0.1936,
+      "image_masks_loss": 0.0085,
+      "llm_learning_rate": 1.0030891438728915e-06,
+      "loss": 0.7363,
+      "output_text_mask_loss": 0.0058,
+      "predictor_learning_rate": 0.00010030891438728914,
+      "step": 6961
+    },
+    {
+      "epoch": 0.71,
+      "gumbel_tau": 0.1935,
+      "image_masks_loss": 0.0081,
+      "llm_learning_rate": 1.0024247653122956e-06,
+      "loss": 0.7339,
+      "output_text_mask_loss": 0.0263,
+      "predictor_learning_rate": 0.00010024247653122956,
+      "step": 6962
+    },
+    {
+      "epoch": 0.71,
+      "gumbel_tau": 0.1935,
+      "image_masks_loss": 0.0159,
+      "llm_learning_rate": 1.0017605516680204e-06,
+      "loss": 0.7393,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 0.00010017605516680203,
+      "step": 6963
+    },
+    {
+      "epoch": 0.71,
+      "gumbel_tau": 0.1935,
+      "image_masks_loss": 0.0042,
+      "llm_learning_rate": 1.0010965030132103e-06,
+      "loss": 0.6882,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 0.00010010965030132103,
+      "step": 6964
+    },
+    {
+      "epoch": 0.71,
+      "gumbel_tau": 0.1934,
+      "image_masks_loss": 0.0049,
+      "llm_learning_rate": 1.0004326194209921e-06,
+      "loss": 0.7515,
+      "output_text_mask_loss": 0.0058,
+      "predictor_learning_rate": 0.0001000432619420992,
+      "step": 6965
+    },
+    {
+      "epoch": 0.71,
+      "gumbel_tau": 0.1934,
+      "image_masks_loss": 0.012,
+      "llm_learning_rate": 9.997689009644736e-07,
+      "loss": 0.7332,
+      "output_text_mask_loss": 0.0526,
+      "predictor_learning_rate": 9.997689009644736e-05,
+      "step": 6966
+    },
+    {
+      "epoch": 0.71,
+      "gumbel_tau": 0.1933,
+      "image_masks_loss": 0.0092,
+      "llm_learning_rate": 9.99105347716745e-07,
+      "loss": 0.7558,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 9.99105347716745e-05,
+      "step": 6967
+    },
+    {
+      "epoch": 0.71,
+      "gumbel_tau": 0.1933,
+      "image_masks_loss": 0.0124,
+      "llm_learning_rate": 9.984419597508794e-07,
+      "loss": 0.7037,
+      "output_text_mask_loss": 0.0086,
+      "predictor_learning_rate": 9.984419597508793e-05,
+      "step": 6968
+    },
+    {
+      "epoch": 0.71,
+      "gumbel_tau": 0.1932,
+      "image_masks_loss": 0.0051,
+      "llm_learning_rate": 9.977787371399295e-07,
+      "loss": 0.7677,
+      "output_text_mask_loss": 0.0319,
+      "predictor_learning_rate": 9.977787371399294e-05,
+      "step": 6969
+    },
+    {
+      "epoch": 0.71,
+      "gumbel_tau": 0.1932,
+      "image_masks_loss": 0.0075,
+      "llm_learning_rate": 9.971156799569312e-07,
+      "loss": 0.7601,
+      "output_text_mask_loss": 0.0192,
+      "predictor_learning_rate": 9.971156799569312e-05,
+      "step": 6970
+    },
+    {
+      "epoch": 0.71,
+      "gumbel_tau": 0.1931,
+      "image_masks_loss": 0.0149,
+      "llm_learning_rate": 9.96452788274902e-07,
+      "loss": 0.8193,
+      "output_text_mask_loss": 0.042,
+      "predictor_learning_rate": 9.964527882749019e-05,
+      "step": 6971
+    },
+    {
+      "epoch": 0.71,
+      "gumbel_tau": 0.1931,
+      "image_masks_loss": 0.006,
+      "llm_learning_rate": 9.957900621668416e-07,
+      "loss": 0.6926,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 9.957900621668414e-05,
+      "step": 6972
+    },
+    {
+      "epoch": 0.71,
+      "gumbel_tau": 0.193,
+      "image_masks_loss": 0.0216,
+      "llm_learning_rate": 9.951275017057279e-07,
+      "loss": 0.8123,
+      "output_text_mask_loss": 0.0143,
+      "predictor_learning_rate": 9.951275017057279e-05,
+      "step": 6973
+    },
+    {
+      "epoch": 0.71,
+      "gumbel_tau": 0.193,
+      "image_masks_loss": 0.0023,
+      "llm_learning_rate": 9.94465106964527e-07,
+      "loss": 0.7393,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 9.944651069645269e-05,
+      "step": 6974
+    },
+    {
+      "epoch": 0.71,
+      "gumbel_tau": 0.193,
+      "image_masks_loss": 0.0079,
+      "llm_learning_rate": 9.93802878016182e-07,
+      "loss": 0.8219,
+      "output_text_mask_loss": 0.0259,
+      "predictor_learning_rate": 9.938028780161818e-05,
+      "step": 6975
+    },
+    {
+      "epoch": 0.71,
+      "gumbel_tau": 0.1929,
+      "image_masks_loss": 0.0036,
+      "llm_learning_rate": 9.931408149336187e-07,
+      "loss": 0.7164,
+      "output_text_mask_loss": 0.0221,
+      "predictor_learning_rate": 9.931408149336185e-05,
+      "step": 6976
+    },
+    {
+      "epoch": 0.71,
+      "gumbel_tau": 0.1929,
+      "image_masks_loss": 0.0034,
+      "llm_learning_rate": 9.924789177897455e-07,
+      "loss": 0.7142,
+      "output_text_mask_loss": 0.0324,
+      "predictor_learning_rate": 9.924789177897453e-05,
+      "step": 6977
+    },
+    {
+      "epoch": 0.71,
+      "gumbel_tau": 0.1928,
+      "image_masks_loss": 0.0108,
+      "llm_learning_rate": 9.918171866574514e-07,
+      "loss": 0.7573,
+      "output_text_mask_loss": 0.0618,
+      "predictor_learning_rate": 9.918171866574513e-05,
+      "step": 6978
+    },
+    {
+      "epoch": 0.72,
+      "gumbel_tau": 0.1928,
+      "image_masks_loss": 0.0052,
+      "llm_learning_rate": 9.911556216096083e-07,
+      "loss": 0.5911,
+      "output_text_mask_loss": 0.0137,
+      "predictor_learning_rate": 9.911556216096082e-05,
+      "step": 6979
+    },
+    {
+      "epoch": 0.72,
+      "gumbel_tau": 0.1927,
+      "image_masks_loss": 0.0116,
+      "llm_learning_rate": 9.904942227190682e-07,
+      "loss": 0.755,
+      "output_text_mask_loss": 0.0212,
+      "predictor_learning_rate": 9.90494222719068e-05,
+      "step": 6980
+    },
+    {
+      "epoch": 0.72,
+      "gumbel_tau": 0.1927,
+      "image_masks_loss": 0.0097,
+      "llm_learning_rate": 9.898329900586673e-07,
+      "loss": 0.7942,
+      "output_text_mask_loss": 0.0055,
+      "predictor_learning_rate": 9.898329900586672e-05,
+      "step": 6981
+    },
+    {
+      "epoch": 0.72,
+      "gumbel_tau": 0.1926,
+      "image_masks_loss": 0.0188,
+      "llm_learning_rate": 9.891719237012212e-07,
+      "loss": 0.7375,
+      "output_text_mask_loss": 0.0231,
+      "predictor_learning_rate": 9.891719237012211e-05,
+      "step": 6982
+    },
+    {
+      "epoch": 0.72,
+      "gumbel_tau": 0.1926,
+      "image_masks_loss": 0.0061,
+      "llm_learning_rate": 9.885110237195282e-07,
+      "loss": 0.6863,
+      "output_text_mask_loss": 0.021,
+      "predictor_learning_rate": 9.885110237195283e-05,
+      "step": 6983
+    },
+    {
+      "epoch": 0.72,
+      "gumbel_tau": 0.1925,
+      "image_masks_loss": 0.0128,
+      "llm_learning_rate": 9.878502901863682e-07,
+      "loss": 0.7248,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 9.878502901863681e-05,
+      "step": 6984
+    },
+    {
+      "epoch": 0.72,
+      "gumbel_tau": 0.1925,
+      "image_masks_loss": 0.009,
+      "llm_learning_rate": 9.871897231745023e-07,
+      "loss": 0.6487,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 9.871897231745023e-05,
+      "step": 6985
+    },
+    {
+      "epoch": 0.72,
+      "gumbel_tau": 0.1925,
+      "image_masks_loss": 0.0147,
+      "llm_learning_rate": 9.865293227566737e-07,
+      "loss": 0.7551,
+      "output_text_mask_loss": 0.0041,
+      "predictor_learning_rate": 9.865293227566735e-05,
+      "step": 6986
+    },
+    {
+      "epoch": 0.72,
+      "gumbel_tau": 0.1924,
+      "image_masks_loss": 0.0033,
+      "llm_learning_rate": 9.85869089005606e-07,
+      "loss": 0.7185,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 9.85869089005606e-05,
+      "step": 6987
+    },
+    {
+      "epoch": 0.72,
+      "gumbel_tau": 0.1924,
+      "image_masks_loss": 0.0154,
+      "llm_learning_rate": 9.852090219940081e-07,
+      "loss": 0.7397,
+      "output_text_mask_loss": 0.0169,
+      "predictor_learning_rate": 9.85209021994008e-05,
+      "step": 6988
+    },
+    {
+      "epoch": 0.72,
+      "gumbel_tau": 0.1923,
+      "image_masks_loss": 0.0095,
+      "llm_learning_rate": 9.845491217945665e-07,
+      "loss": 0.7913,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 9.845491217945665e-05,
+      "step": 6989
+    },
+    {
+      "epoch": 0.72,
+      "gumbel_tau": 0.1923,
+      "image_masks_loss": 0.0154,
+      "llm_learning_rate": 9.838893884799514e-07,
+      "loss": 0.7206,
+      "output_text_mask_loss": 0.0073,
+      "predictor_learning_rate": 9.838893884799514e-05,
+      "step": 6990
+    },
+    {
+      "epoch": 0.72,
+      "gumbel_tau": 0.1922,
+      "image_masks_loss": 0.0079,
+      "llm_learning_rate": 9.832298221228148e-07,
+      "loss": 0.7474,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 9.832298221228148e-05,
+      "step": 6991
+    },
+    {
+      "epoch": 0.72,
+      "gumbel_tau": 0.1922,
+      "image_masks_loss": 0.0127,
+      "llm_learning_rate": 9.825704227957874e-07,
+      "loss": 0.7076,
+      "output_text_mask_loss": 0.0175,
+      "predictor_learning_rate": 9.825704227957874e-05,
+      "step": 6992
+    },
+    {
+      "epoch": 0.72,
+      "gumbel_tau": 0.1921,
+      "image_masks_loss": 0.0069,
+      "llm_learning_rate": 9.819111905714845e-07,
+      "loss": 0.7661,
+      "output_text_mask_loss": 0.0511,
+      "predictor_learning_rate": 9.819111905714845e-05,
+      "step": 6993
+    },
+    {
+      "epoch": 0.72,
+      "gumbel_tau": 0.1921,
+      "image_masks_loss": 0.0422,
+      "llm_learning_rate": 9.812521255225032e-07,
+      "loss": 0.7964,
+      "output_text_mask_loss": 0.01,
+      "predictor_learning_rate": 9.812521255225032e-05,
+      "step": 6994
+    },
+    {
+      "epoch": 0.72,
+      "gumbel_tau": 0.192,
+      "image_masks_loss": 0.0046,
+      "llm_learning_rate": 9.805932277214212e-07,
+      "loss": 0.7522,
+      "output_text_mask_loss": 0.0139,
+      "predictor_learning_rate": 9.805932277214211e-05,
+      "step": 6995
+    },
+    {
+      "epoch": 0.72,
+      "gumbel_tau": 0.192,
+      "image_masks_loss": 0.0129,
+      "llm_learning_rate": 9.799344972407967e-07,
+      "loss": 0.7787,
+      "output_text_mask_loss": 0.013,
+      "predictor_learning_rate": 9.799344972407967e-05,
+      "step": 6996
+    },
+    {
+      "epoch": 0.72,
+      "gumbel_tau": 0.192,
+      "image_masks_loss": 0.0074,
+      "llm_learning_rate": 9.792759341531715e-07,
+      "loss": 0.7039,
+      "output_text_mask_loss": 0.0147,
+      "predictor_learning_rate": 9.792759341531716e-05,
+      "step": 6997
+    },
+    {
+      "epoch": 0.72,
+      "gumbel_tau": 0.1919,
+      "image_masks_loss": 0.013,
+      "llm_learning_rate": 9.786175385310678e-07,
+      "loss": 0.6416,
+      "output_text_mask_loss": 0.0168,
+      "predictor_learning_rate": 9.786175385310677e-05,
+      "step": 6998
+    },
+    {
+      "epoch": 0.72,
+      "gumbel_tau": 0.1919,
+      "image_masks_loss": 0.0119,
+      "llm_learning_rate": 9.779593104469894e-07,
+      "loss": 0.7443,
+      "output_text_mask_loss": 0.0202,
+      "predictor_learning_rate": 9.779593104469895e-05,
+      "step": 6999
+    },
+    {
+      "epoch": 0.72,
+      "gumbel_tau": 0.1918,
+      "image_masks_loss": 0.009,
+      "llm_learning_rate": 9.77301249973421e-07,
+      "loss": 0.6644,
+      "output_text_mask_loss": 0.0116,
+      "predictor_learning_rate": 9.77301249973421e-05,
+      "step": 7000
+    },
+    {
+      "epoch": 0.72,
+      "gumbel_tau": 0.1918,
+      "image_masks_loss": 0.0079,
+      "llm_learning_rate": 9.766433571828318e-07,
+      "loss": 0.6986,
+      "output_text_mask_loss": 0.0041,
+      "predictor_learning_rate": 9.766433571828318e-05,
+      "step": 7001
+    },
+    {
+      "epoch": 0.72,
+      "gumbel_tau": 0.1917,
+      "image_masks_loss": 0.0171,
+      "llm_learning_rate": 9.75985632147669e-07,
+      "loss": 0.7078,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 9.75985632147669e-05,
+      "step": 7002
+    },
+    {
+      "epoch": 0.72,
+      "gumbel_tau": 0.1917,
+      "image_masks_loss": 0.0058,
+      "llm_learning_rate": 9.753280749403628e-07,
+      "loss": 0.7401,
+      "output_text_mask_loss": 0.0465,
+      "predictor_learning_rate": 9.753280749403629e-05,
+      "step": 7003
+    },
+    {
+      "epoch": 0.72,
+      "gumbel_tau": 0.1916,
+      "image_masks_loss": 0.0109,
+      "llm_learning_rate": 9.746706856333255e-07,
+      "loss": 0.7292,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 9.746706856333254e-05,
+      "step": 7004
+    },
+    {
+      "epoch": 0.72,
+      "gumbel_tau": 0.1916,
+      "image_masks_loss": 0.0077,
+      "llm_learning_rate": 9.740134642989495e-07,
+      "loss": 0.7086,
+      "output_text_mask_loss": 0.0374,
+      "predictor_learning_rate": 9.740134642989495e-05,
+      "step": 7005
+    },
+    {
+      "epoch": 0.72,
+      "gumbel_tau": 0.1915,
+      "image_masks_loss": 0.0104,
+      "llm_learning_rate": 9.7335641100961e-07,
+      "loss": 0.7932,
+      "output_text_mask_loss": 0.0019,
+      "predictor_learning_rate": 9.733564110096099e-05,
+      "step": 7006
+    },
+    {
+      "epoch": 0.72,
+      "gumbel_tau": 0.1915,
+      "image_masks_loss": 0.0087,
+      "llm_learning_rate": 9.726995258376621e-07,
+      "loss": 0.7538,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 9.72699525837662e-05,
+      "step": 7007
+    },
+    {
+      "epoch": 0.72,
+      "gumbel_tau": 0.1915,
+      "image_masks_loss": 0.013,
+      "llm_learning_rate": 9.72042808855445e-07,
+      "loss": 0.7783,
+      "output_text_mask_loss": 0.0004,
+      "predictor_learning_rate": 9.72042808855445e-05,
+      "step": 7008
+    },
+    {
+      "epoch": 0.72,
+      "gumbel_tau": 0.1914,
+      "image_masks_loss": 0.0309,
+      "llm_learning_rate": 9.713862601352784e-07,
+      "loss": 0.766,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 9.713862601352783e-05,
+      "step": 7009
+    },
+    {
+      "epoch": 0.72,
+      "gumbel_tau": 0.1914,
+      "image_masks_loss": 0.0218,
+      "llm_learning_rate": 9.707298797494604e-07,
+      "loss": 0.701,
+      "output_text_mask_loss": 0.0101,
+      "predictor_learning_rate": 9.707298797494604e-05,
+      "step": 7010
+    },
+    {
+      "epoch": 0.72,
+      "gumbel_tau": 0.1913,
+      "image_masks_loss": 0.0125,
+      "llm_learning_rate": 9.70073667770275e-07,
+      "loss": 0.7261,
+      "output_text_mask_loss": 0.0057,
+      "predictor_learning_rate": 9.700736677702748e-05,
+      "step": 7011
+    },
+    {
+      "epoch": 0.72,
+      "gumbel_tau": 0.1913,
+      "image_masks_loss": 0.0082,
+      "llm_learning_rate": 9.694176242699835e-07,
+      "loss": 0.764,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 9.694176242699835e-05,
+      "step": 7012
+    },
+    {
+      "epoch": 0.72,
+      "gumbel_tau": 0.1912,
+      "image_masks_loss": 0.0081,
+      "llm_learning_rate": 9.687617493208336e-07,
+      "loss": 0.6802,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 9.687617493208337e-05,
+      "step": 7013
+    },
+    {
+      "epoch": 0.72,
+      "gumbel_tau": 0.1912,
+      "image_masks_loss": 0.0081,
+      "llm_learning_rate": 9.681060429950507e-07,
+      "loss": 0.8074,
+      "output_text_mask_loss": 0.0076,
+      "predictor_learning_rate": 9.681060429950505e-05,
+      "step": 7014
+    },
+    {
+      "epoch": 0.72,
+      "gumbel_tau": 0.1911,
+      "image_masks_loss": 0.0173,
+      "llm_learning_rate": 9.674505053648425e-07,
+      "loss": 0.6906,
+      "output_text_mask_loss": 0.1312,
+      "predictor_learning_rate": 9.674505053648424e-05,
+      "step": 7015
+    },
+    {
+      "epoch": 0.72,
+      "gumbel_tau": 0.1911,
+      "image_masks_loss": 0.0165,
+      "llm_learning_rate": 9.66795136502398e-07,
+      "loss": 0.8012,
+      "output_text_mask_loss": 0.0439,
+      "predictor_learning_rate": 9.66795136502398e-05,
+      "step": 7016
+    },
+    {
+      "epoch": 0.72,
+      "gumbel_tau": 0.1911,
+      "image_masks_loss": 0.0072,
+      "llm_learning_rate": 9.661399364798884e-07,
+      "loss": 0.7307,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 9.661399364798885e-05,
+      "step": 7017
+    },
+    {
+      "epoch": 0.72,
+      "gumbel_tau": 0.191,
+      "image_masks_loss": 0.011,
+      "llm_learning_rate": 9.654849053694653e-07,
+      "loss": 0.764,
+      "output_text_mask_loss": 0.0103,
+      "predictor_learning_rate": 9.654849053694654e-05,
+      "step": 7018
+    },
+    {
+      "epoch": 0.72,
+      "gumbel_tau": 0.191,
+      "image_masks_loss": 0.0077,
+      "llm_learning_rate": 9.64830043243262e-07,
+      "loss": 0.6198,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 9.64830043243262e-05,
+      "step": 7019
+    },
+    {
+      "epoch": 0.72,
+      "gumbel_tau": 0.1909,
+      "image_masks_loss": 0.0091,
+      "llm_learning_rate": 9.641753501733942e-07,
+      "loss": 0.7341,
+      "output_text_mask_loss": 0.0275,
+      "predictor_learning_rate": 9.641753501733941e-05,
+      "step": 7020
+    },
+    {
+      "epoch": 0.72,
+      "gumbel_tau": 0.1909,
+      "image_masks_loss": 0.0061,
+      "llm_learning_rate": 9.635208262319578e-07,
+      "loss": 0.6797,
+      "output_text_mask_loss": 0.0542,
+      "predictor_learning_rate": 9.635208262319578e-05,
+      "step": 7021
+    },
+    {
+      "epoch": 0.72,
+      "gumbel_tau": 0.1908,
+      "image_masks_loss": 0.0044,
+      "llm_learning_rate": 9.628664714910307e-07,
+      "loss": 0.8387,
+      "output_text_mask_loss": 0.0261,
+      "predictor_learning_rate": 9.628664714910307e-05,
+      "step": 7022
+    },
+    {
+      "epoch": 0.72,
+      "gumbel_tau": 0.1908,
+      "image_masks_loss": 0.0201,
+      "llm_learning_rate": 9.622122860226714e-07,
+      "loss": 0.8133,
+      "output_text_mask_loss": 0.0141,
+      "predictor_learning_rate": 9.622122860226715e-05,
+      "step": 7023
+    },
+    {
+      "epoch": 0.72,
+      "gumbel_tau": 0.1907,
+      "image_masks_loss": 0.0046,
+      "llm_learning_rate": 9.615582698989204e-07,
+      "loss": 0.8304,
+      "output_text_mask_loss": 0.0223,
+      "predictor_learning_rate": 9.615582698989203e-05,
+      "step": 7024
+    },
+    {
+      "epoch": 0.72,
+      "gumbel_tau": 0.1907,
+      "image_masks_loss": 0.0115,
+      "llm_learning_rate": 9.609044231917996e-07,
+      "loss": 0.7609,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 9.609044231917996e-05,
+      "step": 7025
+    },
+    {
+      "epoch": 0.72,
+      "gumbel_tau": 0.1906,
+      "image_masks_loss": 0.0047,
+      "llm_learning_rate": 9.60250745973311e-07,
+      "loss": 0.7853,
+      "output_text_mask_loss": 0.0374,
+      "predictor_learning_rate": 9.60250745973311e-05,
+      "step": 7026
+    },
+    {
+      "epoch": 0.72,
+      "gumbel_tau": 0.1906,
+      "image_masks_loss": 0.0075,
+      "llm_learning_rate": 9.595972383154406e-07,
+      "loss": 0.6943,
+      "output_text_mask_loss": 0.0317,
+      "predictor_learning_rate": 9.595972383154405e-05,
+      "step": 7027
+    },
+    {
+      "epoch": 0.72,
+      "gumbel_tau": 0.1906,
+      "image_masks_loss": 0.0026,
+      "llm_learning_rate": 9.589439002901546e-07,
+      "loss": 0.7074,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 9.589439002901546e-05,
+      "step": 7028
+    },
+    {
+      "epoch": 0.72,
+      "gumbel_tau": 0.1905,
+      "image_masks_loss": 0.0132,
+      "llm_learning_rate": 9.582907319693976e-07,
+      "loss": 0.6668,
+      "output_text_mask_loss": 0.0071,
+      "predictor_learning_rate": 9.582907319693975e-05,
+      "step": 7029
+    },
+    {
+      "epoch": 0.72,
+      "gumbel_tau": 0.1905,
+      "image_masks_loss": 0.0148,
+      "llm_learning_rate": 9.576377334250997e-07,
+      "loss": 0.6979,
+      "output_text_mask_loss": 0.0025,
+      "predictor_learning_rate": 9.576377334250997e-05,
+      "step": 7030
+    },
+    {
+      "epoch": 0.72,
+      "gumbel_tau": 0.1904,
+      "image_masks_loss": 0.0053,
+      "llm_learning_rate": 9.569849047291698e-07,
+      "loss": 0.6934,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 9.569849047291698e-05,
+      "step": 7031
+    },
+    {
+      "epoch": 0.72,
+      "gumbel_tau": 0.1904,
+      "image_masks_loss": 0.0109,
+      "llm_learning_rate": 9.563322459534982e-07,
+      "loss": 0.8026,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 9.563322459534982e-05,
+      "step": 7032
+    },
+    {
+      "epoch": 0.72,
+      "gumbel_tau": 0.1903,
+      "image_masks_loss": 0.0215,
+      "llm_learning_rate": 9.556797571699592e-07,
+      "loss": 0.7599,
+      "output_text_mask_loss": 0.074,
+      "predictor_learning_rate": 9.55679757169959e-05,
+      "step": 7033
+    },
+    {
+      "epoch": 0.72,
+      "gumbel_tau": 0.1903,
+      "image_masks_loss": 0.0226,
+      "llm_learning_rate": 9.550274384504049e-07,
+      "loss": 0.7273,
+      "output_text_mask_loss": 0.0402,
+      "predictor_learning_rate": 9.550274384504049e-05,
+      "step": 7034
+    },
+    {
+      "epoch": 0.72,
+      "gumbel_tau": 0.1902,
+      "image_masks_loss": 0.0148,
+      "llm_learning_rate": 9.543752898666705e-07,
+      "loss": 0.7284,
+      "output_text_mask_loss": 0.0065,
+      "predictor_learning_rate": 9.543752898666705e-05,
+      "step": 7035
+    },
+    {
+      "epoch": 0.72,
+      "gumbel_tau": 0.1902,
+      "image_masks_loss": 0.0051,
+      "llm_learning_rate": 9.537233114905719e-07,
+      "loss": 0.7631,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 9.537233114905719e-05,
+      "step": 7036
+    },
+    {
+      "epoch": 0.72,
+      "gumbel_tau": 0.1902,
+      "image_masks_loss": 0.0106,
+      "llm_learning_rate": 9.530715033939064e-07,
+      "loss": 0.7529,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 9.530715033939064e-05,
+      "step": 7037
+    },
+    {
+      "epoch": 0.72,
+      "gumbel_tau": 0.1901,
+      "image_masks_loss": 0.0118,
+      "llm_learning_rate": 9.524198656484526e-07,
+      "loss": 0.7835,
+      "output_text_mask_loss": 0.0139,
+      "predictor_learning_rate": 9.524198656484526e-05,
+      "step": 7038
+    },
+    {
+      "epoch": 0.72,
+      "gumbel_tau": 0.1901,
+      "image_masks_loss": 0.0068,
+      "llm_learning_rate": 9.517683983259696e-07,
+      "loss": 0.7145,
+      "output_text_mask_loss": 0.0324,
+      "predictor_learning_rate": 9.517683983259695e-05,
+      "step": 7039
+    },
+    {
+      "epoch": 0.72,
+      "gumbel_tau": 0.19,
+      "image_masks_loss": 0.0105,
+      "llm_learning_rate": 9.511171014981996e-07,
+      "loss": 0.7536,
+      "output_text_mask_loss": 0.0116,
+      "predictor_learning_rate": 9.511171014981995e-05,
+      "step": 7040
+    },
+    {
+      "epoch": 0.72,
+      "gumbel_tau": 0.19,
+      "image_masks_loss": 0.0091,
+      "llm_learning_rate": 9.504659752368647e-07,
+      "loss": 0.8121,
+      "output_text_mask_loss": 0.0141,
+      "predictor_learning_rate": 9.504659752368647e-05,
+      "step": 7041
+    },
+    {
+      "epoch": 0.72,
+      "gumbel_tau": 0.1899,
+      "image_masks_loss": 0.0081,
+      "llm_learning_rate": 9.498150196136679e-07,
+      "loss": 0.6759,
+      "output_text_mask_loss": 0.0055,
+      "predictor_learning_rate": 9.498150196136679e-05,
+      "step": 7042
+    },
+    {
+      "epoch": 0.72,
+      "gumbel_tau": 0.1899,
+      "image_masks_loss": 0.0123,
+      "llm_learning_rate": 9.491642347002941e-07,
+      "loss": 0.7387,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 9.491642347002941e-05,
+      "step": 7043
+    },
+    {
+      "epoch": 0.72,
+      "gumbel_tau": 0.1898,
+      "image_masks_loss": 0.018,
+      "llm_learning_rate": 9.485136205684093e-07,
+      "loss": 0.7987,
+      "output_text_mask_loss": 0.0021,
+      "predictor_learning_rate": 9.485136205684092e-05,
+      "step": 7044
+    },
+    {
+      "epoch": 0.72,
+      "gumbel_tau": 0.1898,
+      "image_masks_loss": 0.0301,
+      "llm_learning_rate": 9.478631772896602e-07,
+      "loss": 0.7732,
+      "output_text_mask_loss": 0.0103,
+      "predictor_learning_rate": 9.4786317728966e-05,
+      "step": 7045
+    },
+    {
+      "epoch": 0.72,
+      "gumbel_tau": 0.1897,
+      "image_masks_loss": 0.0237,
+      "llm_learning_rate": 9.472129049356754e-07,
+      "loss": 0.717,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 9.472129049356753e-05,
+      "step": 7046
+    },
+    {
+      "epoch": 0.72,
+      "gumbel_tau": 0.1897,
+      "image_masks_loss": 0.012,
+      "llm_learning_rate": 9.465628035780641e-07,
+      "loss": 0.7493,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 9.465628035780641e-05,
+      "step": 7047
+    },
+    {
+      "epoch": 0.72,
+      "gumbel_tau": 0.1897,
+      "image_masks_loss": 0.0157,
+      "llm_learning_rate": 9.459128732884171e-07,
+      "loss": 0.7113,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 9.45912873288417e-05,
+      "step": 7048
+    },
+    {
+      "epoch": 0.72,
+      "gumbel_tau": 0.1896,
+      "image_masks_loss": 0.0171,
+      "llm_learning_rate": 9.452631141383059e-07,
+      "loss": 0.7804,
+      "output_text_mask_loss": 0.0134,
+      "predictor_learning_rate": 9.452631141383059e-05,
+      "step": 7049
+    },
+    {
+      "epoch": 0.72,
+      "gumbel_tau": 0.1896,
+      "image_masks_loss": 0.0024,
+      "llm_learning_rate": 9.44613526199284e-07,
+      "loss": 0.7225,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 9.446135261992839e-05,
+      "step": 7050
+    },
+    {
+      "epoch": 0.72,
+      "gumbel_tau": 0.1895,
+      "image_masks_loss": 0.0093,
+      "llm_learning_rate": 9.439641095428837e-07,
+      "loss": 0.7421,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 9.439641095428836e-05,
+      "step": 7051
+    },
+    {
+      "epoch": 0.72,
+      "gumbel_tau": 0.1895,
+      "image_masks_loss": 0.0194,
+      "llm_learning_rate": 9.433148642406229e-07,
+      "loss": 0.7927,
+      "output_text_mask_loss": 0.0112,
+      "predictor_learning_rate": 9.433148642406229e-05,
+      "step": 7052
+    },
+    {
+      "epoch": 0.72,
+      "gumbel_tau": 0.1894,
+      "image_masks_loss": 0.0058,
+      "llm_learning_rate": 9.426657903639963e-07,
+      "loss": 0.6609,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 9.426657903639962e-05,
+      "step": 7053
+    },
+    {
+      "epoch": 0.72,
+      "gumbel_tau": 0.1894,
+      "image_masks_loss": 0.0115,
+      "llm_learning_rate": 9.420168879844818e-07,
+      "loss": 0.7014,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 9.420168879844817e-05,
+      "step": 7054
+    },
+    {
+      "epoch": 0.72,
+      "gumbel_tau": 0.1893,
+      "image_masks_loss": 0.0045,
+      "llm_learning_rate": 9.413681571735381e-07,
+      "loss": 0.6341,
+      "output_text_mask_loss": 0.0063,
+      "predictor_learning_rate": 9.41368157173538e-05,
+      "step": 7055
+    },
+    {
+      "epoch": 0.72,
+      "gumbel_tau": 0.1893,
+      "image_masks_loss": 0.0091,
+      "llm_learning_rate": 9.407195980026043e-07,
+      "loss": 0.7257,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 9.407195980026043e-05,
+      "step": 7056
+    },
+    {
+      "epoch": 0.72,
+      "gumbel_tau": 0.1893,
+      "image_masks_loss": 0.025,
+      "llm_learning_rate": 9.400712105431017e-07,
+      "loss": 0.7367,
+      "output_text_mask_loss": 0.0117,
+      "predictor_learning_rate": 9.400712105431017e-05,
+      "step": 7057
+    },
+    {
+      "epoch": 0.72,
+      "gumbel_tau": 0.1892,
+      "image_masks_loss": 0.0083,
+      "llm_learning_rate": 9.39422994866431e-07,
+      "loss": 0.6779,
+      "output_text_mask_loss": 0.0483,
+      "predictor_learning_rate": 9.39422994866431e-05,
+      "step": 7058
+    },
+    {
+      "epoch": 0.72,
+      "gumbel_tau": 0.1892,
+      "image_masks_loss": 0.0144,
+      "llm_learning_rate": 9.387749510439772e-07,
+      "loss": 0.6785,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 9.387749510439772e-05,
+      "step": 7059
+    },
+    {
+      "epoch": 0.72,
+      "gumbel_tau": 0.1891,
+      "image_masks_loss": 0.0146,
+      "llm_learning_rate": 9.381270791471034e-07,
+      "loss": 0.8093,
+      "output_text_mask_loss": 0.0193,
+      "predictor_learning_rate": 9.381270791471033e-05,
+      "step": 7060
+    },
+    {
+      "epoch": 0.72,
+      "gumbel_tau": 0.1891,
+      "image_masks_loss": 0.0051,
+      "llm_learning_rate": 9.374793792471542e-07,
+      "loss": 0.7338,
+      "output_text_mask_loss": 0.0071,
+      "predictor_learning_rate": 9.374793792471541e-05,
+      "step": 7061
+    },
+    {
+      "epoch": 0.72,
+      "gumbel_tau": 0.189,
+      "image_masks_loss": 0.0098,
+      "llm_learning_rate": 9.368318514154564e-07,
+      "loss": 0.6589,
+      "output_text_mask_loss": 0.0058,
+      "predictor_learning_rate": 9.368318514154564e-05,
+      "step": 7062
+    },
+    {
+      "epoch": 0.72,
+      "gumbel_tau": 0.189,
+      "image_masks_loss": 0.0144,
+      "llm_learning_rate": 9.361844957233168e-07,
+      "loss": 0.7646,
+      "output_text_mask_loss": 0.0082,
+      "predictor_learning_rate": 9.361844957233167e-05,
+      "step": 7063
+    },
+    {
+      "epoch": 0.72,
+      "gumbel_tau": 0.1889,
+      "image_masks_loss": 0.0043,
+      "llm_learning_rate": 9.355373122420239e-07,
+      "loss": 0.6697,
+      "output_text_mask_loss": 0.0187,
+      "predictor_learning_rate": 9.355373122420238e-05,
+      "step": 7064
+    },
+    {
+      "epoch": 0.72,
+      "gumbel_tau": 0.1889,
+      "image_masks_loss": 0.0013,
+      "llm_learning_rate": 9.34890301042847e-07,
+      "loss": 0.7755,
+      "output_text_mask_loss": 0.0252,
+      "predictor_learning_rate": 9.348903010428469e-05,
+      "step": 7065
+    },
+    {
+      "epoch": 0.72,
+      "gumbel_tau": 0.1889,
+      "image_masks_loss": 0.0113,
+      "llm_learning_rate": 9.342434621970362e-07,
+      "loss": 0.7592,
+      "output_text_mask_loss": 0.0393,
+      "predictor_learning_rate": 9.342434621970361e-05,
+      "step": 7066
+    },
+    {
+      "epoch": 0.72,
+      "gumbel_tau": 0.1888,
+      "image_masks_loss": 0.0111,
+      "llm_learning_rate": 9.335967957758226e-07,
+      "loss": 0.7738,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 9.335967957758226e-05,
+      "step": 7067
+    },
+    {
+      "epoch": 0.72,
+      "gumbel_tau": 0.1888,
+      "image_masks_loss": 0.0115,
+      "llm_learning_rate": 9.329503018504193e-07,
+      "loss": 0.7539,
+      "output_text_mask_loss": 0.0231,
+      "predictor_learning_rate": 9.329503018504193e-05,
+      "step": 7068
+    },
+    {
+      "epoch": 0.72,
+      "gumbel_tau": 0.1887,
+      "image_masks_loss": 0.0132,
+      "llm_learning_rate": 9.323039804920192e-07,
+      "loss": 0.819,
+      "output_text_mask_loss": 0.0055,
+      "predictor_learning_rate": 9.323039804920192e-05,
+      "step": 7069
+    },
+    {
+      "epoch": 0.72,
+      "gumbel_tau": 0.1887,
+      "image_masks_loss": 0.0157,
+      "llm_learning_rate": 9.316578317717964e-07,
+      "loss": 0.6384,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 9.316578317717964e-05,
+      "step": 7070
+    },
+    {
+      "epoch": 0.72,
+      "gumbel_tau": 0.1886,
+      "image_masks_loss": 0.0137,
+      "llm_learning_rate": 9.310118557609058e-07,
+      "loss": 0.7969,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 9.310118557609057e-05,
+      "step": 7071
+    },
+    {
+      "epoch": 0.72,
+      "gumbel_tau": 0.1886,
+      "image_masks_loss": 0.0147,
+      "llm_learning_rate": 9.303660525304853e-07,
+      "loss": 0.7637,
+      "output_text_mask_loss": 0.0158,
+      "predictor_learning_rate": 9.303660525304852e-05,
+      "step": 7072
+    },
+    {
+      "epoch": 0.72,
+      "gumbel_tau": 0.1885,
+      "image_masks_loss": 0.0153,
+      "llm_learning_rate": 9.297204221516512e-07,
+      "loss": 0.6957,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 9.297204221516511e-05,
+      "step": 7073
+    },
+    {
+      "epoch": 0.72,
+      "gumbel_tau": 0.1885,
+      "image_masks_loss": 0.0121,
+      "llm_learning_rate": 9.290749646955022e-07,
+      "loss": 0.8129,
+      "output_text_mask_loss": 0.0004,
+      "predictor_learning_rate": 9.290749646955021e-05,
+      "step": 7074
+    },
+    {
+      "epoch": 0.72,
+      "gumbel_tau": 0.1885,
+      "image_masks_loss": 0.0087,
+      "llm_learning_rate": 9.284296802331169e-07,
+      "loss": 0.7536,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 9.284296802331169e-05,
+      "step": 7075
+    },
+    {
+      "epoch": 0.72,
+      "gumbel_tau": 0.1884,
+      "image_masks_loss": 0.0101,
+      "llm_learning_rate": 9.277845688355558e-07,
+      "loss": 0.7831,
+      "output_text_mask_loss": 0.0081,
+      "predictor_learning_rate": 9.277845688355558e-05,
+      "step": 7076
+    },
+    {
+      "epoch": 0.73,
+      "gumbel_tau": 0.1884,
+      "image_masks_loss": 0.0094,
+      "llm_learning_rate": 9.2713963057386e-07,
+      "loss": 0.7664,
+      "output_text_mask_loss": 0.0058,
+      "predictor_learning_rate": 9.2713963057386e-05,
+      "step": 7077
+    },
+    {
+      "epoch": 0.73,
+      "gumbel_tau": 0.1883,
+      "image_masks_loss": 0.0136,
+      "llm_learning_rate": 9.264948655190506e-07,
+      "loss": 0.7177,
+      "output_text_mask_loss": 0.005,
+      "predictor_learning_rate": 9.264948655190505e-05,
+      "step": 7078
+    },
+    {
+      "epoch": 0.73,
+      "gumbel_tau": 0.1883,
+      "image_masks_loss": 0.0032,
+      "llm_learning_rate": 9.258502737421324e-07,
+      "loss": 0.7602,
+      "output_text_mask_loss": 0.0004,
+      "predictor_learning_rate": 9.258502737421323e-05,
+      "step": 7079
+    },
+    {
+      "epoch": 0.73,
+      "gumbel_tau": 0.1882,
+      "image_masks_loss": 0.0052,
+      "llm_learning_rate": 9.25205855314088e-07,
+      "loss": 0.6802,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 9.25205855314088e-05,
+      "step": 7080
+    },
+    {
+      "epoch": 0.73,
+      "gumbel_tau": 0.1882,
+      "image_masks_loss": 0.0042,
+      "llm_learning_rate": 9.24561610305883e-07,
+      "loss": 0.754,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 9.245616103058829e-05,
+      "step": 7081
+    },
+    {
+      "epoch": 0.73,
+      "gumbel_tau": 0.1881,
+      "image_masks_loss": 0.0106,
+      "llm_learning_rate": 9.239175387884625e-07,
+      "loss": 0.6833,
+      "output_text_mask_loss": 0.0038,
+      "predictor_learning_rate": 9.239175387884623e-05,
+      "step": 7082
+    },
+    {
+      "epoch": 0.73,
+      "gumbel_tau": 0.1881,
+      "image_masks_loss": 0.0063,
+      "llm_learning_rate": 9.23273640832753e-07,
+      "loss": 0.717,
+      "output_text_mask_loss": 0.0223,
+      "predictor_learning_rate": 9.23273640832753e-05,
+      "step": 7083
+    },
+    {
+      "epoch": 0.73,
+      "gumbel_tau": 0.1881,
+      "image_masks_loss": 0.0058,
+      "llm_learning_rate": 9.226299165096622e-07,
+      "loss": 0.7438,
+      "output_text_mask_loss": 0.0095,
+      "predictor_learning_rate": 9.226299165096622e-05,
+      "step": 7084
+    },
+    {
+      "epoch": 0.73,
+      "gumbel_tau": 0.188,
+      "image_masks_loss": 0.0069,
+      "llm_learning_rate": 9.219863658900785e-07,
+      "loss": 0.7374,
+      "output_text_mask_loss": 0.0441,
+      "predictor_learning_rate": 9.219863658900784e-05,
+      "step": 7085
+    },
+    {
+      "epoch": 0.73,
+      "gumbel_tau": 0.188,
+      "image_masks_loss": 0.0146,
+      "llm_learning_rate": 9.213429890448713e-07,
+      "loss": 0.6844,
+      "output_text_mask_loss": 0.0042,
+      "predictor_learning_rate": 9.213429890448714e-05,
+      "step": 7086
+    },
+    {
+      "epoch": 0.73,
+      "gumbel_tau": 0.1879,
+      "image_masks_loss": 0.0068,
+      "llm_learning_rate": 9.206997860448902e-07,
+      "loss": 0.7921,
+      "output_text_mask_loss": 0.0126,
+      "predictor_learning_rate": 9.206997860448902e-05,
+      "step": 7087
+    },
+    {
+      "epoch": 0.73,
+      "gumbel_tau": 0.1879,
+      "image_masks_loss": 0.0066,
+      "llm_learning_rate": 9.200567569609666e-07,
+      "loss": 0.6872,
+      "output_text_mask_loss": 0.0012,
+      "predictor_learning_rate": 9.200567569609666e-05,
+      "step": 7088
+    },
+    {
+      "epoch": 0.73,
+      "gumbel_tau": 0.1878,
+      "image_masks_loss": 0.0124,
+      "llm_learning_rate": 9.194139018639119e-07,
+      "loss": 0.8089,
+      "output_text_mask_loss": 0.0083,
+      "predictor_learning_rate": 9.194139018639119e-05,
+      "step": 7089
+    },
+    {
+      "epoch": 0.73,
+      "gumbel_tau": 0.1878,
+      "image_masks_loss": 0.0132,
+      "llm_learning_rate": 9.187712208245184e-07,
+      "loss": 0.6589,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 9.187712208245183e-05,
+      "step": 7090
+    },
+    {
+      "epoch": 0.73,
+      "gumbel_tau": 0.1877,
+      "image_masks_loss": 0.0123,
+      "llm_learning_rate": 9.181287139135605e-07,
+      "loss": 0.6646,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 9.181287139135605e-05,
+      "step": 7091
+    },
+    {
+      "epoch": 0.73,
+      "gumbel_tau": 0.1877,
+      "image_masks_loss": 0.0037,
+      "llm_learning_rate": 9.174863812017925e-07,
+      "loss": 0.7327,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 9.174863812017925e-05,
+      "step": 7092
+    },
+    {
+      "epoch": 0.73,
+      "gumbel_tau": 0.1877,
+      "image_masks_loss": 0.0098,
+      "llm_learning_rate": 9.168442227599491e-07,
+      "loss": 0.6774,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 9.16844222759949e-05,
+      "step": 7093
+    },
+    {
+      "epoch": 0.73,
+      "gumbel_tau": 0.1876,
+      "image_masks_loss": 0.012,
+      "llm_learning_rate": 9.162022386587461e-07,
+      "loss": 0.6254,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 9.16202238658746e-05,
+      "step": 7094
+    },
+    {
+      "epoch": 0.73,
+      "gumbel_tau": 0.1876,
+      "image_masks_loss": 0.0125,
+      "llm_learning_rate": 9.155604289688804e-07,
+      "loss": 0.7973,
+      "output_text_mask_loss": 0.0107,
+      "predictor_learning_rate": 9.155604289688804e-05,
+      "step": 7095
+    },
+    {
+      "epoch": 0.73,
+      "gumbel_tau": 0.1875,
+      "image_masks_loss": 0.0064,
+      "llm_learning_rate": 9.149187937610294e-07,
+      "loss": 0.6636,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 9.149187937610292e-05,
+      "step": 7096
+    },
+    {
+      "epoch": 0.73,
+      "gumbel_tau": 0.1875,
+      "image_masks_loss": 0.0045,
+      "llm_learning_rate": 9.142773331058508e-07,
+      "loss": 0.7446,
+      "output_text_mask_loss": 0.0055,
+      "predictor_learning_rate": 9.142773331058507e-05,
+      "step": 7097
+    },
+    {
+      "epoch": 0.73,
+      "gumbel_tau": 0.1874,
+      "image_masks_loss": 0.0036,
+      "llm_learning_rate": 9.13636047073985e-07,
+      "loss": 0.6442,
+      "output_text_mask_loss": 0.0124,
+      "predictor_learning_rate": 9.13636047073985e-05,
+      "step": 7098
+    },
+    {
+      "epoch": 0.73,
+      "gumbel_tau": 0.1874,
+      "image_masks_loss": 0.0037,
+      "llm_learning_rate": 9.129949357360512e-07,
+      "loss": 0.7841,
+      "output_text_mask_loss": 0.0151,
+      "predictor_learning_rate": 9.129949357360512e-05,
+      "step": 7099
+    },
+    {
+      "epoch": 0.73,
+      "gumbel_tau": 0.1873,
+      "image_masks_loss": 0.0046,
+      "llm_learning_rate": 9.123539991626501e-07,
+      "loss": 0.6937,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 9.123539991626501e-05,
+      "step": 7100
+    },
+    {
+      "epoch": 0.73,
+      "gumbel_tau": 0.1873,
+      "image_masks_loss": 0.0124,
+      "llm_learning_rate": 9.117132374243628e-07,
+      "loss": 0.7478,
+      "output_text_mask_loss": 0.0019,
+      "predictor_learning_rate": 9.117132374243629e-05,
+      "step": 7101
+    },
+    {
+      "epoch": 0.73,
+      "gumbel_tau": 0.1873,
+      "image_masks_loss": 0.0227,
+      "llm_learning_rate": 9.110726505917525e-07,
+      "loss": 0.6624,
+      "output_text_mask_loss": 0.0137,
+      "predictor_learning_rate": 9.110726505917524e-05,
+      "step": 7102
+    },
+    {
+      "epoch": 0.73,
+      "gumbel_tau": 0.1872,
+      "image_masks_loss": 0.0018,
+      "llm_learning_rate": 9.104322387353592e-07,
+      "loss": 0.7878,
+      "output_text_mask_loss": 0.0034,
+      "predictor_learning_rate": 9.104322387353592e-05,
+      "step": 7103
+    },
+    {
+      "epoch": 0.73,
+      "gumbel_tau": 0.1872,
+      "image_masks_loss": 0.007,
+      "llm_learning_rate": 9.097920019257094e-07,
+      "loss": 0.7857,
+      "output_text_mask_loss": 0.0116,
+      "predictor_learning_rate": 9.097920019257095e-05,
+      "step": 7104
+    },
+    {
+      "epoch": 0.73,
+      "gumbel_tau": 0.1871,
+      "image_masks_loss": 0.0082,
+      "llm_learning_rate": 9.091519402333066e-07,
+      "loss": 0.7529,
+      "output_text_mask_loss": 0.0029,
+      "predictor_learning_rate": 9.091519402333065e-05,
+      "step": 7105
+    },
+    {
+      "epoch": 0.73,
+      "gumbel_tau": 0.1871,
+      "image_masks_loss": 0.0133,
+      "llm_learning_rate": 9.085120537286355e-07,
+      "loss": 0.7489,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 9.085120537286353e-05,
+      "step": 7106
+    },
+    {
+      "epoch": 0.73,
+      "gumbel_tau": 0.187,
+      "image_masks_loss": 0.0034,
+      "llm_learning_rate": 9.07872342482162e-07,
+      "loss": 0.7196,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 9.07872342482162e-05,
+      "step": 7107
+    },
+    {
+      "epoch": 0.73,
+      "gumbel_tau": 0.187,
+      "image_masks_loss": 0.0072,
+      "llm_learning_rate": 9.072328065643324e-07,
+      "loss": 0.7299,
+      "output_text_mask_loss": 0.0162,
+      "predictor_learning_rate": 9.072328065643324e-05,
+      "step": 7108
+    },
+    {
+      "epoch": 0.73,
+      "gumbel_tau": 0.1869,
+      "image_masks_loss": 0.0117,
+      "llm_learning_rate": 9.065934460455736e-07,
+      "loss": 0.6617,
+      "output_text_mask_loss": 0.0071,
+      "predictor_learning_rate": 9.065934460455735e-05,
+      "step": 7109
+    },
+    {
+      "epoch": 0.73,
+      "gumbel_tau": 0.1869,
+      "image_masks_loss": 0.0102,
+      "llm_learning_rate": 9.059542609962932e-07,
+      "loss": 0.7606,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 9.059542609962932e-05,
+      "step": 7110
+    },
+    {
+      "epoch": 0.73,
+      "gumbel_tau": 0.1869,
+      "image_masks_loss": 0.0037,
+      "llm_learning_rate": 9.05315251486881e-07,
+      "loss": 0.7176,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 9.053152514868809e-05,
+      "step": 7111
+    },
+    {
+      "epoch": 0.73,
+      "gumbel_tau": 0.1868,
+      "image_masks_loss": 0.0068,
+      "llm_learning_rate": 9.046764175877051e-07,
+      "loss": 0.7508,
+      "output_text_mask_loss": 0.0185,
+      "predictor_learning_rate": 9.046764175877051e-05,
+      "step": 7112
+    },
+    {
+      "epoch": 0.73,
+      "gumbel_tau": 0.1868,
+      "image_masks_loss": 0.0073,
+      "llm_learning_rate": 9.040377593691158e-07,
+      "loss": 0.686,
+      "output_text_mask_loss": 0.0286,
+      "predictor_learning_rate": 9.040377593691157e-05,
+      "step": 7113
+    },
+    {
+      "epoch": 0.73,
+      "gumbel_tau": 0.1867,
+      "image_masks_loss": 0.0135,
+      "llm_learning_rate": 9.033992769014433e-07,
+      "loss": 0.7169,
+      "output_text_mask_loss": 0.0076,
+      "predictor_learning_rate": 9.033992769014433e-05,
+      "step": 7114
+    },
+    {
+      "epoch": 0.73,
+      "gumbel_tau": 0.1867,
+      "image_masks_loss": 0.0098,
+      "llm_learning_rate": 9.027609702549986e-07,
+      "loss": 0.6944,
+      "output_text_mask_loss": 0.0429,
+      "predictor_learning_rate": 9.027609702549986e-05,
+      "step": 7115
+    },
+    {
+      "epoch": 0.73,
+      "gumbel_tau": 0.1866,
+      "image_masks_loss": 0.0042,
+      "llm_learning_rate": 9.021228395000736e-07,
+      "loss": 0.699,
+      "output_text_mask_loss": 0.0202,
+      "predictor_learning_rate": 9.021228395000736e-05,
+      "step": 7116
+    },
+    {
+      "epoch": 0.73,
+      "gumbel_tau": 0.1866,
+      "image_masks_loss": 0.0094,
+      "llm_learning_rate": 9.014848847069396e-07,
+      "loss": 0.7253,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 9.014848847069395e-05,
+      "step": 7117
+    },
+    {
+      "epoch": 0.73,
+      "gumbel_tau": 0.1866,
+      "image_masks_loss": 0.0106,
+      "llm_learning_rate": 9.008471059458518e-07,
+      "loss": 0.6958,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 9.008471059458517e-05,
+      "step": 7118
+    },
+    {
+      "epoch": 0.73,
+      "gumbel_tau": 0.1865,
+      "image_masks_loss": 0.0065,
+      "llm_learning_rate": 9.002095032870426e-07,
+      "loss": 0.7207,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 9.002095032870425e-05,
+      "step": 7119
+    },
+    {
+      "epoch": 0.73,
+      "gumbel_tau": 0.1865,
+      "image_masks_loss": 0.0078,
+      "llm_learning_rate": 8.995720768007271e-07,
+      "loss": 0.6605,
+      "output_text_mask_loss": 0.0275,
+      "predictor_learning_rate": 8.99572076800727e-05,
+      "step": 7120
+    },
+    {
+      "epoch": 0.73,
+      "gumbel_tau": 0.1864,
+      "image_masks_loss": 0.0082,
+      "llm_learning_rate": 8.989348265570985e-07,
+      "loss": 0.6775,
+      "output_text_mask_loss": 0.0305,
+      "predictor_learning_rate": 8.989348265570985e-05,
+      "step": 7121
+    },
+    {
+      "epoch": 0.73,
+      "gumbel_tau": 0.1864,
+      "image_masks_loss": 0.0139,
+      "llm_learning_rate": 8.98297752626332e-07,
+      "loss": 0.7908,
+      "output_text_mask_loss": 0.0065,
+      "predictor_learning_rate": 8.98297752626332e-05,
+      "step": 7122
+    },
+    {
+      "epoch": 0.73,
+      "gumbel_tau": 0.1863,
+      "image_masks_loss": 0.0073,
+      "llm_learning_rate": 8.976608550785856e-07,
+      "loss": 0.748,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 8.976608550785856e-05,
+      "step": 7123
+    },
+    {
+      "epoch": 0.73,
+      "gumbel_tau": 0.1863,
+      "image_masks_loss": 0.0037,
+      "llm_learning_rate": 8.970241339839949e-07,
+      "loss": 0.6674,
+      "output_text_mask_loss": 0.0085,
+      "predictor_learning_rate": 8.970241339839949e-05,
+      "step": 7124
+    },
+    {
+      "epoch": 0.73,
+      "gumbel_tau": 0.1862,
+      "image_masks_loss": 0.0166,
+      "llm_learning_rate": 8.963875894126772e-07,
+      "loss": 0.7252,
+      "output_text_mask_loss": 0.0254,
+      "predictor_learning_rate": 8.963875894126771e-05,
+      "step": 7125
+    },
+    {
+      "epoch": 0.73,
+      "gumbel_tau": 0.1862,
+      "image_masks_loss": 0.0188,
+      "llm_learning_rate": 8.9575122143473e-07,
+      "loss": 0.6501,
+      "output_text_mask_loss": 0.037,
+      "predictor_learning_rate": 8.9575122143473e-05,
+      "step": 7126
+    },
+    {
+      "epoch": 0.73,
+      "gumbel_tau": 0.1862,
+      "image_masks_loss": 0.0116,
+      "llm_learning_rate": 8.951150301202316e-07,
+      "loss": 0.7228,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 8.951150301202315e-05,
+      "step": 7127
+    },
+    {
+      "epoch": 0.73,
+      "gumbel_tau": 0.1861,
+      "image_masks_loss": 0.0082,
+      "llm_learning_rate": 8.944790155392407e-07,
+      "loss": 0.724,
+      "output_text_mask_loss": 0.0322,
+      "predictor_learning_rate": 8.944790155392407e-05,
+      "step": 7128
+    },
+    {
+      "epoch": 0.73,
+      "gumbel_tau": 0.1861,
+      "image_masks_loss": 0.0156,
+      "llm_learning_rate": 8.938431777617962e-07,
+      "loss": 0.7021,
+      "output_text_mask_loss": 0.0014,
+      "predictor_learning_rate": 8.938431777617961e-05,
+      "step": 7129
+    },
+    {
+      "epoch": 0.73,
+      "gumbel_tau": 0.186,
+      "image_masks_loss": 0.0048,
+      "llm_learning_rate": 8.93207516857919e-07,
+      "loss": 0.7788,
+      "output_text_mask_loss": 0.037,
+      "predictor_learning_rate": 8.93207516857919e-05,
+      "step": 7130
+    },
+    {
+      "epoch": 0.73,
+      "gumbel_tau": 0.186,
+      "image_masks_loss": 0.015,
+      "llm_learning_rate": 8.925720328976092e-07,
+      "loss": 0.7876,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 8.925720328976092e-05,
+      "step": 7131
+    },
+    {
+      "epoch": 0.73,
+      "gumbel_tau": 0.1859,
+      "image_masks_loss": 0.012,
+      "llm_learning_rate": 8.919367259508476e-07,
+      "loss": 0.7547,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 8.919367259508477e-05,
+      "step": 7132
+    },
+    {
+      "epoch": 0.73,
+      "gumbel_tau": 0.1859,
+      "image_masks_loss": 0.0033,
+      "llm_learning_rate": 8.913015960875956e-07,
+      "loss": 0.7694,
+      "output_text_mask_loss": 0.0309,
+      "predictor_learning_rate": 8.913015960875955e-05,
+      "step": 7133
+    },
+    {
+      "epoch": 0.73,
+      "gumbel_tau": 0.1858,
+      "image_masks_loss": 0.0115,
+      "llm_learning_rate": 8.906666433777949e-07,
+      "loss": 0.7211,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 8.90666643377795e-05,
+      "step": 7134
+    },
+    {
+      "epoch": 0.73,
+      "gumbel_tau": 0.1858,
+      "image_masks_loss": 0.0057,
+      "llm_learning_rate": 8.900318678913683e-07,
+      "loss": 0.7466,
+      "output_text_mask_loss": 0.0176,
+      "predictor_learning_rate": 8.900318678913683e-05,
+      "step": 7135
+    },
+    {
+      "epoch": 0.73,
+      "gumbel_tau": 0.1858,
+      "image_masks_loss": 0.005,
+      "llm_learning_rate": 8.893972696982176e-07,
+      "loss": 0.6513,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 8.893972696982175e-05,
+      "step": 7136
+    },
+    {
+      "epoch": 0.73,
+      "gumbel_tau": 0.1857,
+      "image_masks_loss": 0.0111,
+      "llm_learning_rate": 8.887628488682276e-07,
+      "loss": 0.6757,
+      "output_text_mask_loss": 0.0595,
+      "predictor_learning_rate": 8.887628488682276e-05,
+      "step": 7137
+    },
+    {
+      "epoch": 0.73,
+      "gumbel_tau": 0.1857,
+      "image_masks_loss": 0.0223,
+      "llm_learning_rate": 8.881286054712618e-07,
+      "loss": 0.6132,
+      "output_text_mask_loss": 0.0019,
+      "predictor_learning_rate": 8.881286054712617e-05,
+      "step": 7138
+    },
+    {
+      "epoch": 0.73,
+      "gumbel_tau": 0.1856,
+      "image_masks_loss": 0.0029,
+      "llm_learning_rate": 8.874945395771654e-07,
+      "loss": 0.7762,
+      "output_text_mask_loss": 0.0202,
+      "predictor_learning_rate": 8.874945395771652e-05,
+      "step": 7139
+    },
+    {
+      "epoch": 0.73,
+      "gumbel_tau": 0.1856,
+      "image_masks_loss": 0.0122,
+      "llm_learning_rate": 8.868606512557607e-07,
+      "loss": 0.7628,
+      "output_text_mask_loss": 0.0012,
+      "predictor_learning_rate": 8.868606512557606e-05,
+      "step": 7140
+    },
+    {
+      "epoch": 0.73,
+      "gumbel_tau": 0.1855,
+      "image_masks_loss": 0.0034,
+      "llm_learning_rate": 8.862269405768542e-07,
+      "loss": 0.7223,
+      "output_text_mask_loss": 0.0039,
+      "predictor_learning_rate": 8.862269405768542e-05,
+      "step": 7141
+    },
+    {
+      "epoch": 0.73,
+      "gumbel_tau": 0.1855,
+      "image_masks_loss": 0.0069,
+      "llm_learning_rate": 8.85593407610231e-07,
+      "loss": 0.6918,
+      "output_text_mask_loss": 0.0111,
+      "predictor_learning_rate": 8.85593407610231e-05,
+      "step": 7142
+    },
+    {
+      "epoch": 0.73,
+      "gumbel_tau": 0.1855,
+      "image_masks_loss": 0.0103,
+      "llm_learning_rate": 8.849600524256585e-07,
+      "loss": 0.7373,
+      "output_text_mask_loss": 0.0038,
+      "predictor_learning_rate": 8.849600524256585e-05,
+      "step": 7143
+    },
+    {
+      "epoch": 0.73,
+      "gumbel_tau": 0.1854,
+      "image_masks_loss": 0.0044,
+      "llm_learning_rate": 8.843268750928822e-07,
+      "loss": 0.769,
+      "output_text_mask_loss": 0.0111,
+      "predictor_learning_rate": 8.843268750928821e-05,
+      "step": 7144
+    },
+    {
+      "epoch": 0.73,
+      "gumbel_tau": 0.1854,
+      "image_masks_loss": 0.0122,
+      "llm_learning_rate": 8.836938756816294e-07,
+      "loss": 0.7324,
+      "output_text_mask_loss": 0.0004,
+      "predictor_learning_rate": 8.836938756816293e-05,
+      "step": 7145
+    },
+    {
+      "epoch": 0.73,
+      "gumbel_tau": 0.1853,
+      "image_masks_loss": 0.0165,
+      "llm_learning_rate": 8.830610542616069e-07,
+      "loss": 0.7425,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 8.830610542616069e-05,
+      "step": 7146
+    },
+    {
+      "epoch": 0.73,
+      "gumbel_tau": 0.1853,
+      "image_masks_loss": 0.0084,
+      "llm_learning_rate": 8.824284109025031e-07,
+      "loss": 0.724,
+      "output_text_mask_loss": 0.0208,
+      "predictor_learning_rate": 8.82428410902503e-05,
+      "step": 7147
+    },
+    {
+      "epoch": 0.73,
+      "gumbel_tau": 0.1852,
+      "image_masks_loss": 0.0102,
+      "llm_learning_rate": 8.817959456739853e-07,
+      "loss": 0.7662,
+      "output_text_mask_loss": 0.0223,
+      "predictor_learning_rate": 8.817959456739851e-05,
+      "step": 7148
+    },
+    {
+      "epoch": 0.73,
+      "gumbel_tau": 0.1852,
+      "image_masks_loss": 0.0059,
+      "llm_learning_rate": 8.811636586457017e-07,
+      "loss": 0.7382,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 8.811636586457017e-05,
+      "step": 7149
+    },
+    {
+      "epoch": 0.73,
+      "gumbel_tau": 0.1851,
+      "image_masks_loss": 0.0044,
+      "llm_learning_rate": 8.805315498872826e-07,
+      "loss": 0.776,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 8.805315498872826e-05,
+      "step": 7150
+    },
+    {
+      "epoch": 0.73,
+      "gumbel_tau": 0.1851,
+      "image_masks_loss": 0.0118,
+      "llm_learning_rate": 8.798996194683365e-07,
+      "loss": 0.7119,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 8.798996194683364e-05,
+      "step": 7151
+    },
+    {
+      "epoch": 0.73,
+      "gumbel_tau": 0.1851,
+      "image_masks_loss": 0.0041,
+      "llm_learning_rate": 8.792678674584532e-07,
+      "loss": 0.6182,
+      "output_text_mask_loss": 0.0099,
+      "predictor_learning_rate": 8.792678674584531e-05,
+      "step": 7152
+    },
+    {
+      "epoch": 0.73,
+      "gumbel_tau": 0.185,
+      "image_masks_loss": 0.0102,
+      "llm_learning_rate": 8.78636293927202e-07,
+      "loss": 0.7322,
+      "output_text_mask_loss": 0.0198,
+      "predictor_learning_rate": 8.786362939272019e-05,
+      "step": 7153
+    },
+    {
+      "epoch": 0.73,
+      "gumbel_tau": 0.185,
+      "image_masks_loss": 0.0122,
+      "llm_learning_rate": 8.780048989441336e-07,
+      "loss": 0.7237,
+      "output_text_mask_loss": 0.006,
+      "predictor_learning_rate": 8.780048989441336e-05,
+      "step": 7154
+    },
+    {
+      "epoch": 0.73,
+      "gumbel_tau": 0.1849,
+      "image_masks_loss": 0.0062,
+      "llm_learning_rate": 8.773736825787788e-07,
+      "loss": 0.6822,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 8.773736825787786e-05,
+      "step": 7155
+    },
+    {
+      "epoch": 0.73,
+      "gumbel_tau": 0.1849,
+      "image_masks_loss": 0.0028,
+      "llm_learning_rate": 8.767426449006472e-07,
+      "loss": 0.758,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 8.767426449006471e-05,
+      "step": 7156
+    },
+    {
+      "epoch": 0.73,
+      "gumbel_tau": 0.1848,
+      "image_masks_loss": 0.0144,
+      "llm_learning_rate": 8.761117859792323e-07,
+      "loss": 0.6689,
+      "output_text_mask_loss": 0.0227,
+      "predictor_learning_rate": 8.761117859792322e-05,
+      "step": 7157
+    },
+    {
+      "epoch": 0.73,
+      "gumbel_tau": 0.1848,
+      "image_masks_loss": 0.0156,
+      "llm_learning_rate": 8.754811058840054e-07,
+      "loss": 0.7256,
+      "output_text_mask_loss": 0.0124,
+      "predictor_learning_rate": 8.754811058840053e-05,
+      "step": 7158
+    },
+    {
+      "epoch": 0.73,
+      "gumbel_tau": 0.1848,
+      "image_masks_loss": 0.0104,
+      "llm_learning_rate": 8.748506046844165e-07,
+      "loss": 0.8287,
+      "output_text_mask_loss": 0.0056,
+      "predictor_learning_rate": 8.748506046844165e-05,
+      "step": 7159
+    },
+    {
+      "epoch": 0.73,
+      "gumbel_tau": 0.1847,
+      "image_masks_loss": 0.0095,
+      "llm_learning_rate": 8.742202824498991e-07,
+      "loss": 0.6592,
+      "output_text_mask_loss": 0.0067,
+      "predictor_learning_rate": 8.74220282449899e-05,
+      "step": 7160
+    },
+    {
+      "epoch": 0.73,
+      "gumbel_tau": 0.1847,
+      "image_masks_loss": 0.0072,
+      "llm_learning_rate": 8.735901392498647e-07,
+      "loss": 0.7013,
+      "output_text_mask_loss": 0.0214,
+      "predictor_learning_rate": 8.735901392498646e-05,
+      "step": 7161
+    },
+    {
+      "epoch": 0.73,
+      "gumbel_tau": 0.1846,
+      "image_masks_loss": 0.0093,
+      "llm_learning_rate": 8.729601751537076e-07,
+      "loss": 0.7549,
+      "output_text_mask_loss": 0.0117,
+      "predictor_learning_rate": 8.729601751537076e-05,
+      "step": 7162
+    },
+    {
+      "epoch": 0.73,
+      "gumbel_tau": 0.1846,
+      "image_masks_loss": 0.0096,
+      "llm_learning_rate": 8.723303902308003e-07,
+      "loss": 0.7305,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 8.723303902308002e-05,
+      "step": 7163
+    },
+    {
+      "epoch": 0.73,
+      "gumbel_tau": 0.1845,
+      "image_masks_loss": 0.0072,
+      "llm_learning_rate": 8.717007845504958e-07,
+      "loss": 0.7032,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 8.717007845504957e-05,
+      "step": 7164
+    },
+    {
+      "epoch": 0.73,
+      "gumbel_tau": 0.1845,
+      "image_masks_loss": 0.0067,
+      "llm_learning_rate": 8.710713581821275e-07,
+      "loss": 0.6476,
+      "output_text_mask_loss": 0.0004,
+      "predictor_learning_rate": 8.710713581821275e-05,
+      "step": 7165
+    },
+    {
+      "epoch": 0.73,
+      "gumbel_tau": 0.1845,
+      "image_masks_loss": 0.0085,
+      "llm_learning_rate": 8.704421111950098e-07,
+      "loss": 0.6597,
+      "output_text_mask_loss": 0.017,
+      "predictor_learning_rate": 8.704421111950098e-05,
+      "step": 7166
+    },
+    {
+      "epoch": 0.73,
+      "gumbel_tau": 0.1844,
+      "image_masks_loss": 0.0087,
+      "llm_learning_rate": 8.698130436584365e-07,
+      "loss": 0.8258,
+      "output_text_mask_loss": 0.0034,
+      "predictor_learning_rate": 8.698130436584364e-05,
+      "step": 7167
+    },
+    {
+      "epoch": 0.73,
+      "gumbel_tau": 0.1844,
+      "image_masks_loss": 0.0175,
+      "llm_learning_rate": 8.69184155641681e-07,
+      "loss": 0.713,
+      "output_text_mask_loss": 0.0168,
+      "predictor_learning_rate": 8.69184155641681e-05,
+      "step": 7168
+    },
+    {
+      "epoch": 0.73,
+      "gumbel_tau": 0.1843,
+      "image_masks_loss": 0.0237,
+      "llm_learning_rate": 8.685554472139998e-07,
+      "loss": 0.6712,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 8.685554472139998e-05,
+      "step": 7169
+    },
+    {
+      "epoch": 0.73,
+      "gumbel_tau": 0.1843,
+      "image_masks_loss": 0.0056,
+      "llm_learning_rate": 8.679269184446266e-07,
+      "loss": 0.7068,
+      "output_text_mask_loss": 0.0109,
+      "predictor_learning_rate": 8.679269184446265e-05,
+      "step": 7170
+    },
+    {
+      "epoch": 0.73,
+      "gumbel_tau": 0.1842,
+      "image_masks_loss": 0.0047,
+      "llm_learning_rate": 8.672985694027764e-07,
+      "loss": 0.7168,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 8.672985694027762e-05,
+      "step": 7171
+    },
+    {
+      "epoch": 0.73,
+      "gumbel_tau": 0.1842,
+      "image_masks_loss": 0.0058,
+      "llm_learning_rate": 8.666704001576443e-07,
+      "loss": 0.6524,
+      "output_text_mask_loss": 0.0526,
+      "predictor_learning_rate": 8.666704001576443e-05,
+      "step": 7172
+    },
+    {
+      "epoch": 0.73,
+      "gumbel_tau": 0.1841,
+      "image_masks_loss": 0.0067,
+      "llm_learning_rate": 8.66042410778406e-07,
+      "loss": 0.777,
+      "output_text_mask_loss": 0.042,
+      "predictor_learning_rate": 8.660424107784059e-05,
+      "step": 7173
+    },
+    {
+      "epoch": 0.74,
+      "gumbel_tau": 0.1841,
+      "image_masks_loss": 0.0109,
+      "llm_learning_rate": 8.654146013342169e-07,
+      "loss": 0.7208,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 8.654146013342168e-05,
+      "step": 7174
+    },
+    {
+      "epoch": 0.74,
+      "gumbel_tau": 0.1841,
+      "image_masks_loss": 0.01,
+      "llm_learning_rate": 8.64786971894212e-07,
+      "loss": 0.698,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 8.647869718942119e-05,
+      "step": 7175
+    },
+    {
+      "epoch": 0.74,
+      "gumbel_tau": 0.184,
+      "image_masks_loss": 0.0083,
+      "llm_learning_rate": 8.641595225275096e-07,
+      "loss": 0.714,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 8.641595225275095e-05,
+      "step": 7176
+    },
+    {
+      "epoch": 0.74,
+      "gumbel_tau": 0.184,
+      "image_masks_loss": 0.0066,
+      "llm_learning_rate": 8.635322533032034e-07,
+      "loss": 0.7809,
+      "output_text_mask_loss": 0.0173,
+      "predictor_learning_rate": 8.635322533032034e-05,
+      "step": 7177
+    },
+    {
+      "epoch": 0.74,
+      "gumbel_tau": 0.1839,
+      "image_masks_loss": 0.0059,
+      "llm_learning_rate": 8.629051642903708e-07,
+      "loss": 0.7473,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 8.629051642903707e-05,
+      "step": 7178
+    },
+    {
+      "epoch": 0.74,
+      "gumbel_tau": 0.1839,
+      "image_masks_loss": 0.0086,
+      "llm_learning_rate": 8.622782555580677e-07,
+      "loss": 0.6804,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 8.622782555580677e-05,
+      "step": 7179
+    },
+    {
+      "epoch": 0.74,
+      "gumbel_tau": 0.1838,
+      "image_masks_loss": 0.0048,
+      "llm_learning_rate": 8.616515271753311e-07,
+      "loss": 0.653,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 8.616515271753312e-05,
+      "step": 7180
+    },
+    {
+      "epoch": 0.74,
+      "gumbel_tau": 0.1838,
+      "image_masks_loss": 0.0045,
+      "llm_learning_rate": 8.610249792111769e-07,
+      "loss": 0.7475,
+      "output_text_mask_loss": 0.0071,
+      "predictor_learning_rate": 8.610249792111769e-05,
+      "step": 7181
+    },
+    {
+      "epoch": 0.74,
+      "gumbel_tau": 0.1838,
+      "image_masks_loss": 0.0066,
+      "llm_learning_rate": 8.603986117346036e-07,
+      "loss": 0.6671,
+      "output_text_mask_loss": 0.024,
+      "predictor_learning_rate": 8.603986117346035e-05,
+      "step": 7182
+    },
+    {
+      "epoch": 0.74,
+      "gumbel_tau": 0.1837,
+      "image_masks_loss": 0.0142,
+      "llm_learning_rate": 8.59772424814587e-07,
+      "loss": 0.6959,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 8.597724248145869e-05,
+      "step": 7183
+    },
+    {
+      "epoch": 0.74,
+      "gumbel_tau": 0.1837,
+      "image_masks_loss": 0.0066,
+      "llm_learning_rate": 8.591464185200843e-07,
+      "loss": 0.7114,
+      "output_text_mask_loss": 0.0023,
+      "predictor_learning_rate": 8.591464185200842e-05,
+      "step": 7184
+    },
+    {
+      "epoch": 0.74,
+      "gumbel_tau": 0.1836,
+      "image_masks_loss": 0.0079,
+      "llm_learning_rate": 8.585205929200333e-07,
+      "loss": 0.71,
+      "output_text_mask_loss": 0.0056,
+      "predictor_learning_rate": 8.585205929200332e-05,
+      "step": 7185
+    },
+    {
+      "epoch": 0.74,
+      "gumbel_tau": 0.1836,
+      "image_masks_loss": 0.0122,
+      "llm_learning_rate": 8.578949480833503e-07,
+      "loss": 0.8019,
+      "output_text_mask_loss": 0.0025,
+      "predictor_learning_rate": 8.578949480833503e-05,
+      "step": 7186
+    },
+    {
+      "epoch": 0.74,
+      "gumbel_tau": 0.1835,
+      "image_masks_loss": 0.0102,
+      "llm_learning_rate": 8.572694840789336e-07,
+      "loss": 0.5707,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 8.572694840789336e-05,
+      "step": 7187
+    },
+    {
+      "epoch": 0.74,
+      "gumbel_tau": 0.1835,
+      "image_masks_loss": 0.0207,
+      "llm_learning_rate": 8.566442009756596e-07,
+      "loss": 0.7448,
+      "output_text_mask_loss": 0.0065,
+      "predictor_learning_rate": 8.566442009756594e-05,
+      "step": 7188
+    },
+    {
+      "epoch": 0.74,
+      "gumbel_tau": 0.1835,
+      "image_masks_loss": 0.0046,
+      "llm_learning_rate": 8.560190988423873e-07,
+      "loss": 0.7528,
+      "output_text_mask_loss": 0.0038,
+      "predictor_learning_rate": 8.560190988423871e-05,
+      "step": 7189
+    },
+    {
+      "epoch": 0.74,
+      "gumbel_tau": 0.1834,
+      "image_masks_loss": 0.0122,
+      "llm_learning_rate": 8.553941777479538e-07,
+      "loss": 0.7567,
+      "output_text_mask_loss": 0.0162,
+      "predictor_learning_rate": 8.553941777479537e-05,
+      "step": 7190
+    },
+    {
+      "epoch": 0.74,
+      "gumbel_tau": 0.1834,
+      "image_masks_loss": 0.0052,
+      "llm_learning_rate": 8.547694377611765e-07,
+      "loss": 0.6871,
+      "output_text_mask_loss": 0.0013,
+      "predictor_learning_rate": 8.547694377611765e-05,
+      "step": 7191
+    },
+    {
+      "epoch": 0.74,
+      "gumbel_tau": 0.1833,
+      "image_masks_loss": 0.0112,
+      "llm_learning_rate": 8.541448789508533e-07,
+      "loss": 0.7277,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 8.541448789508533e-05,
+      "step": 7192
+    },
+    {
+      "epoch": 0.74,
+      "gumbel_tau": 0.1833,
+      "image_masks_loss": 0.0226,
+      "llm_learning_rate": 8.535205013857623e-07,
+      "loss": 0.7768,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 8.535205013857623e-05,
+      "step": 7193
+    },
+    {
+      "epoch": 0.74,
+      "gumbel_tau": 0.1832,
+      "image_masks_loss": 0.0045,
+      "llm_learning_rate": 8.528963051346608e-07,
+      "loss": 0.706,
+      "output_text_mask_loss": 0.0082,
+      "predictor_learning_rate": 8.528963051346608e-05,
+      "step": 7194
+    },
+    {
+      "epoch": 0.74,
+      "gumbel_tau": 0.1832,
+      "image_masks_loss": 0.0064,
+      "llm_learning_rate": 8.52272290266287e-07,
+      "loss": 0.781,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 8.522722902662869e-05,
+      "step": 7195
+    },
+    {
+      "epoch": 0.74,
+      "gumbel_tau": 0.1832,
+      "image_masks_loss": 0.0183,
+      "llm_learning_rate": 8.516484568493588e-07,
+      "loss": 0.7364,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 8.516484568493587e-05,
+      "step": 7196
+    },
+    {
+      "epoch": 0.74,
+      "gumbel_tau": 0.1831,
+      "image_masks_loss": 0.0061,
+      "llm_learning_rate": 8.510248049525741e-07,
+      "loss": 0.6551,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 8.51024804952574e-05,
+      "step": 7197
+    },
+    {
+      "epoch": 0.74,
+      "gumbel_tau": 0.1831,
+      "image_masks_loss": 0.0077,
+      "llm_learning_rate": 8.504013346446108e-07,
+      "loss": 0.7663,
+      "output_text_mask_loss": 0.0185,
+      "predictor_learning_rate": 8.504013346446107e-05,
+      "step": 7198
+    },
+    {
+      "epoch": 0.74,
+      "gumbel_tau": 0.183,
+      "image_masks_loss": 0.0035,
+      "llm_learning_rate": 8.497780459941269e-07,
+      "loss": 0.6565,
+      "output_text_mask_loss": 0.0124,
+      "predictor_learning_rate": 8.497780459941267e-05,
+      "step": 7199
+    },
+    {
+      "epoch": 0.74,
+      "gumbel_tau": 0.183,
+      "image_masks_loss": 0.0033,
+      "llm_learning_rate": 8.491549390697596e-07,
+      "loss": 0.7645,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 8.491549390697596e-05,
+      "step": 7200
+    },
+    {
+      "epoch": 0.74,
+      "gumbel_tau": 0.1829,
+      "image_masks_loss": 0.0013,
+      "llm_learning_rate": 8.485320139401282e-07,
+      "loss": 0.7839,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 8.485320139401282e-05,
+      "step": 7201
+    },
+    {
+      "epoch": 0.74,
+      "gumbel_tau": 0.1829,
+      "image_masks_loss": 0.0084,
+      "llm_learning_rate": 8.479092706738299e-07,
+      "loss": 0.6431,
+      "output_text_mask_loss": 0.0086,
+      "predictor_learning_rate": 8.479092706738298e-05,
+      "step": 7202
+    },
+    {
+      "epoch": 0.74,
+      "gumbel_tau": 0.1828,
+      "image_masks_loss": 0.0076,
+      "llm_learning_rate": 8.472867093394427e-07,
+      "loss": 0.7122,
+      "output_text_mask_loss": 0.0072,
+      "predictor_learning_rate": 8.472867093394426e-05,
+      "step": 7203
+    },
+    {
+      "epoch": 0.74,
+      "gumbel_tau": 0.1828,
+      "image_masks_loss": 0.0038,
+      "llm_learning_rate": 8.466643300055241e-07,
+      "loss": 0.7633,
+      "output_text_mask_loss": 0.0198,
+      "predictor_learning_rate": 8.466643300055241e-05,
+      "step": 7204
+    },
+    {
+      "epoch": 0.74,
+      "gumbel_tau": 0.1828,
+      "image_masks_loss": 0.011,
+      "llm_learning_rate": 8.460421327406123e-07,
+      "loss": 0.6587,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 8.460421327406123e-05,
+      "step": 7205
+    },
+    {
+      "epoch": 0.74,
+      "gumbel_tau": 0.1827,
+      "image_masks_loss": 0.0046,
+      "llm_learning_rate": 8.454201176132251e-07,
+      "loss": 0.6575,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 8.454201176132251e-05,
+      "step": 7206
+    },
+    {
+      "epoch": 0.74,
+      "gumbel_tau": 0.1827,
+      "image_masks_loss": 0.0093,
+      "llm_learning_rate": 8.44798284691859e-07,
+      "loss": 0.6964,
+      "output_text_mask_loss": 0.0038,
+      "predictor_learning_rate": 8.447982846918589e-05,
+      "step": 7207
+    },
+    {
+      "epoch": 0.74,
+      "gumbel_tau": 0.1826,
+      "image_masks_loss": 0.0061,
+      "llm_learning_rate": 8.441766340449934e-07,
+      "loss": 0.7149,
+      "output_text_mask_loss": 0.0019,
+      "predictor_learning_rate": 8.441766340449933e-05,
+      "step": 7208
+    },
+    {
+      "epoch": 0.74,
+      "gumbel_tau": 0.1826,
+      "image_masks_loss": 0.0176,
+      "llm_learning_rate": 8.435551657410849e-07,
+      "loss": 0.7383,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 8.435551657410848e-05,
+      "step": 7209
+    },
+    {
+      "epoch": 0.74,
+      "gumbel_tau": 0.1825,
+      "image_masks_loss": 0.0128,
+      "llm_learning_rate": 8.429338798485715e-07,
+      "loss": 0.6925,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 8.429338798485713e-05,
+      "step": 7210
+    },
+    {
+      "epoch": 0.74,
+      "gumbel_tau": 0.1825,
+      "image_masks_loss": 0.015,
+      "llm_learning_rate": 8.423127764358699e-07,
+      "loss": 0.7489,
+      "output_text_mask_loss": 0.016,
+      "predictor_learning_rate": 8.423127764358698e-05,
+      "step": 7211
+    },
+    {
+      "epoch": 0.74,
+      "gumbel_tau": 0.1825,
+      "image_masks_loss": 0.0095,
+      "llm_learning_rate": 8.416918555713776e-07,
+      "loss": 0.6815,
+      "output_text_mask_loss": 0.0252,
+      "predictor_learning_rate": 8.416918555713776e-05,
+      "step": 7212
+    },
+    {
+      "epoch": 0.74,
+      "gumbel_tau": 0.1824,
+      "image_masks_loss": 0.0043,
+      "llm_learning_rate": 8.410711173234723e-07,
+      "loss": 0.6574,
+      "output_text_mask_loss": 0.0014,
+      "predictor_learning_rate": 8.410711173234722e-05,
+      "step": 7213
+    },
+    {
+      "epoch": 0.74,
+      "gumbel_tau": 0.1824,
+      "image_masks_loss": 0.008,
+      "llm_learning_rate": 8.404505617605105e-07,
+      "loss": 0.7399,
+      "output_text_mask_loss": 0.0225,
+      "predictor_learning_rate": 8.404505617605105e-05,
+      "step": 7214
+    },
+    {
+      "epoch": 0.74,
+      "gumbel_tau": 0.1823,
+      "image_masks_loss": 0.0102,
+      "llm_learning_rate": 8.398301889508292e-07,
+      "loss": 0.7375,
+      "output_text_mask_loss": 0.0038,
+      "predictor_learning_rate": 8.398301889508293e-05,
+      "step": 7215
+    },
+    {
+      "epoch": 0.74,
+      "gumbel_tau": 0.1823,
+      "image_masks_loss": 0.009,
+      "llm_learning_rate": 8.392099989627456e-07,
+      "loss": 0.6928,
+      "output_text_mask_loss": 0.0086,
+      "predictor_learning_rate": 8.392099989627455e-05,
+      "step": 7216
+    },
+    {
+      "epoch": 0.74,
+      "gumbel_tau": 0.1822,
+      "image_masks_loss": 0.0085,
+      "llm_learning_rate": 8.385899918645563e-07,
+      "loss": 0.823,
+      "output_text_mask_loss": 0.001,
+      "predictor_learning_rate": 8.385899918645562e-05,
+      "step": 7217
+    },
+    {
+      "epoch": 0.74,
+      "gumbel_tau": 0.1822,
+      "image_masks_loss": 0.0053,
+      "llm_learning_rate": 8.379701677245375e-07,
+      "loss": 0.7676,
+      "output_text_mask_loss": 0.0015,
+      "predictor_learning_rate": 8.379701677245374e-05,
+      "step": 7218
+    },
+    {
+      "epoch": 0.74,
+      "gumbel_tau": 0.1822,
+      "image_masks_loss": 0.0085,
+      "llm_learning_rate": 8.373505266109458e-07,
+      "loss": 0.7728,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 8.373505266109457e-05,
+      "step": 7219
+    },
+    {
+      "epoch": 0.74,
+      "gumbel_tau": 0.1821,
+      "image_masks_loss": 0.0127,
+      "llm_learning_rate": 8.367310685920169e-07,
+      "loss": 0.7441,
+      "output_text_mask_loss": 0.0065,
+      "predictor_learning_rate": 8.367310685920169e-05,
+      "step": 7220
+    },
+    {
+      "epoch": 0.74,
+      "gumbel_tau": 0.1821,
+      "image_masks_loss": 0.0051,
+      "llm_learning_rate": 8.361117937359681e-07,
+      "loss": 0.7014,
+      "output_text_mask_loss": 0.0137,
+      "predictor_learning_rate": 8.36111793735968e-05,
+      "step": 7221
+    },
+    {
+      "epoch": 0.74,
+      "gumbel_tau": 0.182,
+      "image_masks_loss": 0.0056,
+      "llm_learning_rate": 8.354927021109949e-07,
+      "loss": 0.681,
+      "output_text_mask_loss": 0.021,
+      "predictor_learning_rate": 8.354927021109948e-05,
+      "step": 7222
+    },
+    {
+      "epoch": 0.74,
+      "gumbel_tau": 0.182,
+      "image_masks_loss": 0.0049,
+      "llm_learning_rate": 8.34873793785273e-07,
+      "loss": 0.6351,
+      "output_text_mask_loss": 0.0058,
+      "predictor_learning_rate": 8.34873793785273e-05,
+      "step": 7223
+    },
+    {
+      "epoch": 0.74,
+      "gumbel_tau": 0.1819,
+      "image_masks_loss": 0.0119,
+      "llm_learning_rate": 8.342550688269577e-07,
+      "loss": 0.6871,
+      "output_text_mask_loss": 0.0006,
+      "predictor_learning_rate": 8.342550688269576e-05,
+      "step": 7224
+    },
+    {
+      "epoch": 0.74,
+      "gumbel_tau": 0.1819,
+      "image_masks_loss": 0.0119,
+      "llm_learning_rate": 8.336365273041843e-07,
+      "loss": 0.7531,
+      "output_text_mask_loss": 0.0778,
+      "predictor_learning_rate": 8.336365273041843e-05,
+      "step": 7225
+    },
+    {
+      "epoch": 0.74,
+      "gumbel_tau": 0.1819,
+      "image_masks_loss": 0.007,
+      "llm_learning_rate": 8.330181692850686e-07,
+      "loss": 0.6295,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 8.330181692850685e-05,
+      "step": 7226
+    },
+    {
+      "epoch": 0.74,
+      "gumbel_tau": 0.1818,
+      "image_masks_loss": 0.0077,
+      "llm_learning_rate": 8.323999948377043e-07,
+      "loss": 0.6576,
+      "output_text_mask_loss": 0.0035,
+      "predictor_learning_rate": 8.323999948377043e-05,
+      "step": 7227
+    },
+    {
+      "epoch": 0.74,
+      "gumbel_tau": 0.1818,
+      "image_masks_loss": 0.0079,
+      "llm_learning_rate": 8.317820040301675e-07,
+      "loss": 0.793,
+      "output_text_mask_loss": 0.0238,
+      "predictor_learning_rate": 8.317820040301674e-05,
+      "step": 7228
+    },
+    {
+      "epoch": 0.74,
+      "gumbel_tau": 0.1817,
+      "image_masks_loss": 0.0185,
+      "llm_learning_rate": 8.311641969305126e-07,
+      "loss": 0.7746,
+      "output_text_mask_loss": 0.0058,
+      "predictor_learning_rate": 8.311641969305125e-05,
+      "step": 7229
+    },
+    {
+      "epoch": 0.74,
+      "gumbel_tau": 0.1817,
+      "image_masks_loss": 0.0014,
+      "llm_learning_rate": 8.305465736067733e-07,
+      "loss": 0.7609,
+      "output_text_mask_loss": 0.0238,
+      "predictor_learning_rate": 8.305465736067732e-05,
+      "step": 7230
+    },
+    {
+      "epoch": 0.74,
+      "gumbel_tau": 0.1816,
+      "image_masks_loss": 0.0055,
+      "llm_learning_rate": 8.299291341269639e-07,
+      "loss": 0.7083,
+      "output_text_mask_loss": 0.0098,
+      "predictor_learning_rate": 8.299291341269639e-05,
+      "step": 7231
+    },
+    {
+      "epoch": 0.74,
+      "gumbel_tau": 0.1816,
+      "image_masks_loss": 0.004,
+      "llm_learning_rate": 8.29311878559079e-07,
+      "loss": 0.714,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 8.29311878559079e-05,
+      "step": 7232
+    },
+    {
+      "epoch": 0.74,
+      "gumbel_tau": 0.1816,
+      "image_masks_loss": 0.0024,
+      "llm_learning_rate": 8.286948069710896e-07,
+      "loss": 0.7792,
+      "output_text_mask_loss": 0.0223,
+      "predictor_learning_rate": 8.286948069710894e-05,
+      "step": 7233
+    },
+    {
+      "epoch": 0.74,
+      "gumbel_tau": 0.1815,
+      "image_masks_loss": 0.0132,
+      "llm_learning_rate": 8.280779194309518e-07,
+      "loss": 0.7118,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 8.280779194309518e-05,
+      "step": 7234
+    },
+    {
+      "epoch": 0.74,
+      "gumbel_tau": 0.1815,
+      "image_masks_loss": 0.0095,
+      "llm_learning_rate": 8.274612160065973e-07,
+      "loss": 0.8014,
+      "output_text_mask_loss": 0.0189,
+      "predictor_learning_rate": 8.274612160065973e-05,
+      "step": 7235
+    },
+    {
+      "epoch": 0.74,
+      "gumbel_tau": 0.1814,
+      "image_masks_loss": 0.0112,
+      "llm_learning_rate": 8.268446967659394e-07,
+      "loss": 0.6858,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 8.268446967659393e-05,
+      "step": 7236
+    },
+    {
+      "epoch": 0.74,
+      "gumbel_tau": 0.1814,
+      "image_masks_loss": 0.0022,
+      "llm_learning_rate": 8.262283617768701e-07,
+      "loss": 0.7249,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 8.262283617768701e-05,
+      "step": 7237
+    },
+    {
+      "epoch": 0.74,
+      "gumbel_tau": 0.1813,
+      "image_masks_loss": 0.0037,
+      "llm_learning_rate": 8.25612211107262e-07,
+      "loss": 0.757,
+      "output_text_mask_loss": 0.0309,
+      "predictor_learning_rate": 8.25612211107262e-05,
+      "step": 7238
+    },
+    {
+      "epoch": 0.74,
+      "gumbel_tau": 0.1813,
+      "image_masks_loss": 0.0053,
+      "llm_learning_rate": 8.249962448249665e-07,
+      "loss": 0.7449,
+      "output_text_mask_loss": 0.0082,
+      "predictor_learning_rate": 8.249962448249665e-05,
+      "step": 7239
+    },
+    {
+      "epoch": 0.74,
+      "gumbel_tau": 0.1813,
+      "image_masks_loss": 0.0027,
+      "llm_learning_rate": 8.243804629978147e-07,
+      "loss": 0.7197,
+      "output_text_mask_loss": 0.0098,
+      "predictor_learning_rate": 8.243804629978147e-05,
+      "step": 7240
+    },
+    {
+      "epoch": 0.74,
+      "gumbel_tau": 0.1812,
+      "image_masks_loss": 0.0058,
+      "llm_learning_rate": 8.237648656936195e-07,
+      "loss": 0.7324,
+      "output_text_mask_loss": 0.0033,
+      "predictor_learning_rate": 8.237648656936195e-05,
+      "step": 7241
+    },
+    {
+      "epoch": 0.74,
+      "gumbel_tau": 0.1812,
+      "image_masks_loss": 0.0072,
+      "llm_learning_rate": 8.231494529801709e-07,
+      "loss": 0.6232,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 8.231494529801708e-05,
+      "step": 7242
+    },
+    {
+      "epoch": 0.74,
+      "gumbel_tau": 0.1811,
+      "image_masks_loss": 0.0028,
+      "llm_learning_rate": 8.225342249252396e-07,
+      "loss": 0.8079,
+      "output_text_mask_loss": 0.0061,
+      "predictor_learning_rate": 8.225342249252396e-05,
+      "step": 7243
+    },
+    {
+      "epoch": 0.74,
+      "gumbel_tau": 0.1811,
+      "image_masks_loss": 0.0041,
+      "llm_learning_rate": 8.21919181596576e-07,
+      "loss": 0.6535,
+      "output_text_mask_loss": 0.0153,
+      "predictor_learning_rate": 8.21919181596576e-05,
+      "step": 7244
+    },
+    {
+      "epoch": 0.74,
+      "gumbel_tau": 0.181,
+      "image_masks_loss": 0.006,
+      "llm_learning_rate": 8.213043230619097e-07,
+      "loss": 0.7272,
+      "output_text_mask_loss": 0.0015,
+      "predictor_learning_rate": 8.213043230619097e-05,
+      "step": 7245
+    },
+    {
+      "epoch": 0.74,
+      "gumbel_tau": 0.181,
+      "image_masks_loss": 0.0036,
+      "llm_learning_rate": 8.206896493889507e-07,
+      "loss": 0.7793,
+      "output_text_mask_loss": 0.0484,
+      "predictor_learning_rate": 8.206896493889506e-05,
+      "step": 7246
+    },
+    {
+      "epoch": 0.74,
+      "gumbel_tau": 0.181,
+      "image_masks_loss": 0.0086,
+      "llm_learning_rate": 8.200751606453872e-07,
+      "loss": 0.7052,
+      "output_text_mask_loss": 0.0101,
+      "predictor_learning_rate": 8.200751606453871e-05,
+      "step": 7247
+    },
+    {
+      "epoch": 0.74,
+      "gumbel_tau": 0.1809,
+      "image_masks_loss": 0.0102,
+      "llm_learning_rate": 8.194608568988893e-07,
+      "loss": 0.8532,
+      "output_text_mask_loss": 0.0024,
+      "predictor_learning_rate": 8.194608568988893e-05,
+      "step": 7248
+    },
+    {
+      "epoch": 0.74,
+      "gumbel_tau": 0.1809,
+      "image_masks_loss": 0.0183,
+      "llm_learning_rate": 8.188467382171056e-07,
+      "loss": 0.7666,
+      "output_text_mask_loss": 0.0076,
+      "predictor_learning_rate": 8.188467382171055e-05,
+      "step": 7249
+    },
+    {
+      "epoch": 0.74,
+      "gumbel_tau": 0.1808,
+      "image_masks_loss": 0.0033,
+      "llm_learning_rate": 8.182328046676643e-07,
+      "loss": 0.7195,
+      "output_text_mask_loss": 0.069,
+      "predictor_learning_rate": 8.182328046676644e-05,
+      "step": 7250
+    },
+    {
+      "epoch": 0.74,
+      "gumbel_tau": 0.1808,
+      "image_masks_loss": 0.0044,
+      "llm_learning_rate": 8.176190563181718e-07,
+      "loss": 0.742,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 8.176190563181718e-05,
+      "step": 7251
+    },
+    {
+      "epoch": 0.74,
+      "gumbel_tau": 0.1807,
+      "image_masks_loss": 0.0059,
+      "llm_learning_rate": 8.17005493236215e-07,
+      "loss": 0.6942,
+      "output_text_mask_loss": 0.0018,
+      "predictor_learning_rate": 8.170054932362151e-05,
+      "step": 7252
+    },
+    {
+      "epoch": 0.74,
+      "gumbel_tau": 0.1807,
+      "image_masks_loss": 0.0095,
+      "llm_learning_rate": 8.163921154893633e-07,
+      "loss": 0.6986,
+      "output_text_mask_loss": 0.0801,
+      "predictor_learning_rate": 8.163921154893634e-05,
+      "step": 7253
+    },
+    {
+      "epoch": 0.74,
+      "gumbel_tau": 0.1807,
+      "image_masks_loss": 0.0029,
+      "llm_learning_rate": 8.157789231451618e-07,
+      "loss": 0.7355,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 8.157789231451618e-05,
+      "step": 7254
+    },
+    {
+      "epoch": 0.74,
+      "gumbel_tau": 0.1806,
+      "image_masks_loss": 0.0057,
+      "llm_learning_rate": 8.151659162711367e-07,
+      "loss": 0.6758,
+      "output_text_mask_loss": 0.02,
+      "predictor_learning_rate": 8.151659162711367e-05,
+      "step": 7255
+    },
+    {
+      "epoch": 0.74,
+      "gumbel_tau": 0.1806,
+      "image_masks_loss": 0.0014,
+      "llm_learning_rate": 8.145530949347938e-07,
+      "loss": 0.7178,
+      "output_text_mask_loss": 0.0231,
+      "predictor_learning_rate": 8.145530949347937e-05,
+      "step": 7256
+    },
+    {
+      "epoch": 0.74,
+      "gumbel_tau": 0.1805,
+      "image_masks_loss": 0.0073,
+      "llm_learning_rate": 8.139404592036182e-07,
+      "loss": 0.6965,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 8.139404592036182e-05,
+      "step": 7257
+    },
+    {
+      "epoch": 0.74,
+      "gumbel_tau": 0.1805,
+      "image_masks_loss": 0.0053,
+      "llm_learning_rate": 8.133280091450746e-07,
+      "loss": 0.7932,
+      "output_text_mask_loss": 0.0086,
+      "predictor_learning_rate": 8.133280091450746e-05,
+      "step": 7258
+    },
+    {
+      "epoch": 0.74,
+      "gumbel_tau": 0.1804,
+      "image_masks_loss": 0.0106,
+      "llm_learning_rate": 8.127157448266069e-07,
+      "loss": 0.6299,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 8.127157448266068e-05,
+      "step": 7259
+    },
+    {
+      "epoch": 0.74,
+      "gumbel_tau": 0.1804,
+      "image_masks_loss": 0.0102,
+      "llm_learning_rate": 8.121036663156404e-07,
+      "loss": 0.7506,
+      "output_text_mask_loss": 0.0099,
+      "predictor_learning_rate": 8.121036663156403e-05,
+      "step": 7260
+    },
+    {
+      "epoch": 0.74,
+      "gumbel_tau": 0.1804,
+      "image_masks_loss": 0.0074,
+      "llm_learning_rate": 8.11491773679578e-07,
+      "loss": 0.7113,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 8.11491773679578e-05,
+      "step": 7261
+    },
+    {
+      "epoch": 0.74,
+      "gumbel_tau": 0.1803,
+      "image_masks_loss": 0.0045,
+      "llm_learning_rate": 8.108800669858025e-07,
+      "loss": 0.733,
+      "output_text_mask_loss": 0.0124,
+      "predictor_learning_rate": 8.108800669858024e-05,
+      "step": 7262
+    },
+    {
+      "epoch": 0.74,
+      "gumbel_tau": 0.1803,
+      "image_masks_loss": 0.004,
+      "llm_learning_rate": 8.10268546301676e-07,
+      "loss": 0.7508,
+      "output_text_mask_loss": 0.0041,
+      "predictor_learning_rate": 8.10268546301676e-05,
+      "step": 7263
+    },
+    {
+      "epoch": 0.74,
+      "gumbel_tau": 0.1802,
+      "image_masks_loss": 0.0075,
+      "llm_learning_rate": 8.096572116945411e-07,
+      "loss": 0.7829,
+      "output_text_mask_loss": 0.036,
+      "predictor_learning_rate": 8.09657211694541e-05,
+      "step": 7264
+    },
+    {
+      "epoch": 0.74,
+      "gumbel_tau": 0.1802,
+      "image_masks_loss": 0.0058,
+      "llm_learning_rate": 8.090460632317193e-07,
+      "loss": 0.7555,
+      "output_text_mask_loss": 0.0147,
+      "predictor_learning_rate": 8.090460632317192e-05,
+      "step": 7265
+    },
+    {
+      "epoch": 0.74,
+      "gumbel_tau": 0.1802,
+      "image_masks_loss": 0.0057,
+      "llm_learning_rate": 8.084351009805105e-07,
+      "loss": 0.6638,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 8.084351009805105e-05,
+      "step": 7266
+    },
+    {
+      "epoch": 0.74,
+      "gumbel_tau": 0.1801,
+      "image_masks_loss": 0.0076,
+      "llm_learning_rate": 8.078243250081971e-07,
+      "loss": 0.8201,
+      "output_text_mask_loss": 0.0618,
+      "predictor_learning_rate": 8.07824325008197e-05,
+      "step": 7267
+    },
+    {
+      "epoch": 0.74,
+      "gumbel_tau": 0.1801,
+      "image_masks_loss": 0.0068,
+      "llm_learning_rate": 8.072137353820384e-07,
+      "loss": 0.7737,
+      "output_text_mask_loss": 0.0561,
+      "predictor_learning_rate": 8.072137353820383e-05,
+      "step": 7268
+    },
+    {
+      "epoch": 0.74,
+      "gumbel_tau": 0.18,
+      "image_masks_loss": 0.0074,
+      "llm_learning_rate": 8.066033321692743e-07,
+      "loss": 0.8115,
+      "output_text_mask_loss": 0.0484,
+      "predictor_learning_rate": 8.066033321692743e-05,
+      "step": 7269
+    },
+    {
+      "epoch": 0.74,
+      "gumbel_tau": 0.18,
+      "image_masks_loss": 0.0043,
+      "llm_learning_rate": 8.059931154371223e-07,
+      "loss": 0.7215,
+      "output_text_mask_loss": 0.0763,
+      "predictor_learning_rate": 8.059931154371222e-05,
+      "step": 7270
+    },
+    {
+      "epoch": 0.74,
+      "gumbel_tau": 0.1799,
+      "image_masks_loss": 0.0086,
+      "llm_learning_rate": 8.053830852527816e-07,
+      "loss": 0.7788,
+      "output_text_mask_loss": 0.0231,
+      "predictor_learning_rate": 8.053830852527815e-05,
+      "step": 7271
+    },
+    {
+      "epoch": 0.75,
+      "gumbel_tau": 0.1799,
+      "image_masks_loss": 0.0048,
+      "llm_learning_rate": 8.047732416834294e-07,
+      "loss": 0.731,
+      "output_text_mask_loss": 0.0877,
+      "predictor_learning_rate": 8.047732416834293e-05,
+      "step": 7272
+    },
+    {
+      "epoch": 0.75,
+      "gumbel_tau": 0.1799,
+      "image_masks_loss": 0.0162,
+      "llm_learning_rate": 8.041635847962248e-07,
+      "loss": 0.7671,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 8.041635847962248e-05,
+      "step": 7273
+    },
+    {
+      "epoch": 0.75,
+      "gumbel_tau": 0.1798,
+      "image_masks_loss": 0.0084,
+      "llm_learning_rate": 8.035541146583034e-07,
+      "loss": 0.8318,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 8.035541146583033e-05,
+      "step": 7274
+    },
+    {
+      "epoch": 0.75,
+      "gumbel_tau": 0.1798,
+      "image_masks_loss": 0.0057,
+      "llm_learning_rate": 8.029448313367815e-07,
+      "loss": 0.7949,
+      "output_text_mask_loss": 0.0284,
+      "predictor_learning_rate": 8.029448313367815e-05,
+      "step": 7275
+    },
+    {
+      "epoch": 0.75,
+      "gumbel_tau": 0.1797,
+      "image_masks_loss": 0.0065,
+      "llm_learning_rate": 8.023357348987549e-07,
+      "loss": 0.7167,
+      "output_text_mask_loss": 0.0105,
+      "predictor_learning_rate": 8.023357348987548e-05,
+      "step": 7276
+    },
+    {
+      "epoch": 0.75,
+      "gumbel_tau": 0.1797,
+      "image_masks_loss": 0.0079,
+      "llm_learning_rate": 8.017268254112986e-07,
+      "loss": 0.748,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 8.017268254112986e-05,
+      "step": 7277
+    },
+    {
+      "epoch": 0.75,
+      "gumbel_tau": 0.1796,
+      "image_masks_loss": 0.007,
+      "llm_learning_rate": 8.011181029414672e-07,
+      "loss": 0.6517,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 8.011181029414671e-05,
+      "step": 7278
+    },
+    {
+      "epoch": 0.75,
+      "gumbel_tau": 0.1796,
+      "image_masks_loss": 0.0081,
+      "llm_learning_rate": 8.005095675562932e-07,
+      "loss": 0.7277,
+      "output_text_mask_loss": 0.0259,
+      "predictor_learning_rate": 8.005095675562932e-05,
+      "step": 7279
+    },
+    {
+      "epoch": 0.75,
+      "gumbel_tau": 0.1796,
+      "image_masks_loss": 0.0068,
+      "llm_learning_rate": 7.999012193227923e-07,
+      "loss": 0.6994,
+      "output_text_mask_loss": 0.0484,
+      "predictor_learning_rate": 7.999012193227923e-05,
+      "step": 7280
+    },
+    {
+      "epoch": 0.75,
+      "gumbel_tau": 0.1795,
+      "image_masks_loss": 0.0129,
+      "llm_learning_rate": 7.992930583079558e-07,
+      "loss": 0.6137,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 7.992930583079558e-05,
+      "step": 7281
+    },
+    {
+      "epoch": 0.75,
+      "gumbel_tau": 0.1795,
+      "image_masks_loss": 0.0046,
+      "llm_learning_rate": 7.986850845787558e-07,
+      "loss": 0.754,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 7.986850845787558e-05,
+      "step": 7282
+    },
+    {
+      "epoch": 0.75,
+      "gumbel_tau": 0.1794,
+      "image_masks_loss": 0.0079,
+      "llm_learning_rate": 7.980772982021439e-07,
+      "loss": 0.7184,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 7.98077298202144e-05,
+      "step": 7283
+    },
+    {
+      "epoch": 0.75,
+      "gumbel_tau": 0.1794,
+      "image_masks_loss": 0.0114,
+      "llm_learning_rate": 7.974696992450509e-07,
+      "loss": 0.7639,
+      "output_text_mask_loss": 0.0618,
+      "predictor_learning_rate": 7.974696992450508e-05,
+      "step": 7284
+    },
+    {
+      "epoch": 0.75,
+      "gumbel_tau": 0.1793,
+      "image_masks_loss": 0.0037,
+      "llm_learning_rate": 7.968622877743868e-07,
+      "loss": 0.7725,
+      "output_text_mask_loss": 0.0107,
+      "predictor_learning_rate": 7.968622877743867e-05,
+      "step": 7285
+    },
+    {
+      "epoch": 0.75,
+      "gumbel_tau": 0.1793,
+      "image_masks_loss": 0.0088,
+      "llm_learning_rate": 7.962550638570402e-07,
+      "loss": 0.6264,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 7.962550638570401e-05,
+      "step": 7286
+    },
+    {
+      "epoch": 0.75,
+      "gumbel_tau": 0.1793,
+      "image_masks_loss": 0.0044,
+      "llm_learning_rate": 7.956480275598826e-07,
+      "loss": 0.737,
+      "output_text_mask_loss": 0.0231,
+      "predictor_learning_rate": 7.956480275598827e-05,
+      "step": 7287
+    },
+    {
+      "epoch": 0.75,
+      "gumbel_tau": 0.1792,
+      "image_masks_loss": 0.0022,
+      "llm_learning_rate": 7.950411789497597e-07,
+      "loss": 0.7704,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 7.950411789497597e-05,
+      "step": 7288
+    },
+    {
+      "epoch": 0.75,
+      "gumbel_tau": 0.1792,
+      "image_masks_loss": 0.0157,
+      "llm_learning_rate": 7.944345180935001e-07,
+      "loss": 0.7084,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 7.944345180935e-05,
+      "step": 7289
+    },
+    {
+      "epoch": 0.75,
+      "gumbel_tau": 0.1791,
+      "image_masks_loss": 0.0127,
+      "llm_learning_rate": 7.9382804505791e-07,
+      "loss": 0.7592,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 7.9382804505791e-05,
+      "step": 7290
+    },
+    {
+      "epoch": 0.75,
+      "gumbel_tau": 0.1791,
+      "image_masks_loss": 0.0102,
+      "llm_learning_rate": 7.932217599097753e-07,
+      "loss": 0.6777,
+      "output_text_mask_loss": 0.0023,
+      "predictor_learning_rate": 7.932217599097751e-05,
+      "step": 7291
+    },
+    {
+      "epoch": 0.75,
+      "gumbel_tau": 0.179,
+      "image_masks_loss": 0.0071,
+      "llm_learning_rate": 7.926156627158629e-07,
+      "loss": 0.6792,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 7.926156627158629e-05,
+      "step": 7292
+    },
+    {
+      "epoch": 0.75,
+      "gumbel_tau": 0.179,
+      "image_masks_loss": 0.0035,
+      "llm_learning_rate": 7.920097535429167e-07,
+      "loss": 0.7296,
+      "output_text_mask_loss": 0.0324,
+      "predictor_learning_rate": 7.920097535429166e-05,
+      "step": 7293
+    },
+    {
+      "epoch": 0.75,
+      "gumbel_tau": 0.179,
+      "image_masks_loss": 0.0074,
+      "llm_learning_rate": 7.914040324576608e-07,
+      "loss": 0.685,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 7.914040324576607e-05,
+      "step": 7294
+    },
+    {
+      "epoch": 0.75,
+      "gumbel_tau": 0.1789,
+      "image_masks_loss": 0.0042,
+      "llm_learning_rate": 7.907984995267987e-07,
+      "loss": 0.7057,
+      "output_text_mask_loss": 0.0078,
+      "predictor_learning_rate": 7.907984995267986e-05,
+      "step": 7295
+    },
+    {
+      "epoch": 0.75,
+      "gumbel_tau": 0.1789,
+      "image_masks_loss": 0.0036,
+      "llm_learning_rate": 7.901931548170127e-07,
+      "loss": 0.7503,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 7.901931548170127e-05,
+      "step": 7296
+    },
+    {
+      "epoch": 0.75,
+      "gumbel_tau": 0.1788,
+      "image_masks_loss": 0.0072,
+      "llm_learning_rate": 7.895879983949653e-07,
+      "loss": 0.7718,
+      "output_text_mask_loss": 0.0088,
+      "predictor_learning_rate": 7.895879983949653e-05,
+      "step": 7297
+    },
+    {
+      "epoch": 0.75,
+      "gumbel_tau": 0.1788,
+      "image_masks_loss": 0.0059,
+      "llm_learning_rate": 7.889830303272963e-07,
+      "loss": 0.7012,
+      "output_text_mask_loss": 0.0096,
+      "predictor_learning_rate": 7.889830303272964e-05,
+      "step": 7298
+    },
+    {
+      "epoch": 0.75,
+      "gumbel_tau": 0.1788,
+      "image_masks_loss": 0.0124,
+      "llm_learning_rate": 7.883782506806278e-07,
+      "loss": 0.7637,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 7.883782506806278e-05,
+      "step": 7299
+    },
+    {
+      "epoch": 0.75,
+      "gumbel_tau": 0.1787,
+      "image_masks_loss": 0.0051,
+      "llm_learning_rate": 7.87773659521559e-07,
+      "loss": 0.7095,
+      "output_text_mask_loss": 0.0081,
+      "predictor_learning_rate": 7.87773659521559e-05,
+      "step": 7300
+    },
+    {
+      "epoch": 0.75,
+      "gumbel_tau": 0.1787,
+      "image_masks_loss": 0.0076,
+      "llm_learning_rate": 7.871692569166689e-07,
+      "loss": 0.7887,
+      "output_text_mask_loss": 0.0303,
+      "predictor_learning_rate": 7.871692569166688e-05,
+      "step": 7301
+    },
+    {
+      "epoch": 0.75,
+      "gumbel_tau": 0.1786,
+      "image_masks_loss": 0.0086,
+      "llm_learning_rate": 7.86565042932515e-07,
+      "loss": 0.7237,
+      "output_text_mask_loss": 0.0725,
+      "predictor_learning_rate": 7.86565042932515e-05,
+      "step": 7302
+    },
+    {
+      "epoch": 0.75,
+      "gumbel_tau": 0.1786,
+      "image_masks_loss": 0.0029,
+      "llm_learning_rate": 7.859610176356352e-07,
+      "loss": 0.6921,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 7.859610176356352e-05,
+      "step": 7303
+    },
+    {
+      "epoch": 0.75,
+      "gumbel_tau": 0.1785,
+      "image_masks_loss": 0.0113,
+      "llm_learning_rate": 7.853571810925461e-07,
+      "loss": 0.6538,
+      "output_text_mask_loss": 0.0355,
+      "predictor_learning_rate": 7.853571810925462e-05,
+      "step": 7304
+    },
+    {
+      "epoch": 0.75,
+      "gumbel_tau": 0.1785,
+      "image_masks_loss": 0.0042,
+      "llm_learning_rate": 7.847535333697426e-07,
+      "loss": 0.6841,
+      "output_text_mask_loss": 0.0039,
+      "predictor_learning_rate": 7.847535333697426e-05,
+      "step": 7305
+    },
+    {
+      "epoch": 0.75,
+      "gumbel_tau": 0.1785,
+      "image_masks_loss": 0.0084,
+      "llm_learning_rate": 7.841500745337022e-07,
+      "loss": 0.7475,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 7.841500745337021e-05,
+      "step": 7306
+    },
+    {
+      "epoch": 0.75,
+      "gumbel_tau": 0.1784,
+      "image_masks_loss": 0.0073,
+      "llm_learning_rate": 7.835468046508768e-07,
+      "loss": 0.7607,
+      "output_text_mask_loss": 0.038,
+      "predictor_learning_rate": 7.835468046508768e-05,
+      "step": 7307
+    },
+    {
+      "epoch": 0.75,
+      "gumbel_tau": 0.1784,
+      "image_masks_loss": 0.0118,
+      "llm_learning_rate": 7.829437237877002e-07,
+      "loss": 0.6834,
+      "output_text_mask_loss": 0.0202,
+      "predictor_learning_rate": 7.829437237877002e-05,
+      "step": 7308
+    },
+    {
+      "epoch": 0.75,
+      "gumbel_tau": 0.1783,
+      "image_masks_loss": 0.0095,
+      "llm_learning_rate": 7.823408320105855e-07,
+      "loss": 0.7734,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 7.823408320105855e-05,
+      "step": 7309
+    },
+    {
+      "epoch": 0.75,
+      "gumbel_tau": 0.1783,
+      "image_masks_loss": 0.0131,
+      "llm_learning_rate": 7.817381293859244e-07,
+      "loss": 0.6752,
+      "output_text_mask_loss": 0.021,
+      "predictor_learning_rate": 7.817381293859243e-05,
+      "step": 7310
+    },
+    {
+      "epoch": 0.75,
+      "gumbel_tau": 0.1782,
+      "image_masks_loss": 0.005,
+      "llm_learning_rate": 7.81135615980087e-07,
+      "loss": 0.6907,
+      "output_text_mask_loss": 0.0189,
+      "predictor_learning_rate": 7.81135615980087e-05,
+      "step": 7311
+    },
+    {
+      "epoch": 0.75,
+      "gumbel_tau": 0.1782,
+      "image_masks_loss": 0.005,
+      "llm_learning_rate": 7.805332918594247e-07,
+      "loss": 0.807,
+      "output_text_mask_loss": 0.0141,
+      "predictor_learning_rate": 7.805332918594247e-05,
+      "step": 7312
+    },
+    {
+      "epoch": 0.75,
+      "gumbel_tau": 0.1782,
+      "image_masks_loss": 0.0123,
+      "llm_learning_rate": 7.799311570902665e-07,
+      "loss": 0.6643,
+      "output_text_mask_loss": 0.0273,
+      "predictor_learning_rate": 7.799311570902664e-05,
+      "step": 7313
+    },
+    {
+      "epoch": 0.75,
+      "gumbel_tau": 0.1781,
+      "image_masks_loss": 0.0096,
+      "llm_learning_rate": 7.793292117389206e-07,
+      "loss": 0.7552,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 7.793292117389206e-05,
+      "step": 7314
+    },
+    {
+      "epoch": 0.75,
+      "gumbel_tau": 0.1781,
+      "image_masks_loss": 0.0067,
+      "llm_learning_rate": 7.787274558716745e-07,
+      "loss": 0.7167,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 7.787274558716745e-05,
+      "step": 7315
+    },
+    {
+      "epoch": 0.75,
+      "gumbel_tau": 0.178,
+      "image_masks_loss": 0.0042,
+      "llm_learning_rate": 7.781258895547949e-07,
+      "loss": 0.7771,
+      "output_text_mask_loss": 0.0185,
+      "predictor_learning_rate": 7.781258895547949e-05,
+      "step": 7316
+    },
+    {
+      "epoch": 0.75,
+      "gumbel_tau": 0.178,
+      "image_masks_loss": 0.0078,
+      "llm_learning_rate": 7.775245128545278e-07,
+      "loss": 0.7517,
+      "output_text_mask_loss": 0.0359,
+      "predictor_learning_rate": 7.775245128545277e-05,
+      "step": 7317
+    },
+    {
+      "epoch": 0.75,
+      "gumbel_tau": 0.178,
+      "image_masks_loss": 0.009,
+      "llm_learning_rate": 7.76923325837097e-07,
+      "loss": 0.7194,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 7.769233258370969e-05,
+      "step": 7318
+    },
+    {
+      "epoch": 0.75,
+      "gumbel_tau": 0.1779,
+      "image_masks_loss": 0.0083,
+      "llm_learning_rate": 7.763223285687085e-07,
+      "loss": 0.7744,
+      "output_text_mask_loss": 0.0299,
+      "predictor_learning_rate": 7.763223285687085e-05,
+      "step": 7319
+    },
+    {
+      "epoch": 0.75,
+      "gumbel_tau": 0.1779,
+      "image_masks_loss": 0.007,
+      "llm_learning_rate": 7.757215211155445e-07,
+      "loss": 0.7136,
+      "output_text_mask_loss": 0.0055,
+      "predictor_learning_rate": 7.757215211155444e-05,
+      "step": 7320
+    },
+    {
+      "epoch": 0.75,
+      "gumbel_tau": 0.1778,
+      "image_masks_loss": 0.0039,
+      "llm_learning_rate": 7.751209035437671e-07,
+      "loss": 0.7261,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 7.75120903543767e-05,
+      "step": 7321
+    },
+    {
+      "epoch": 0.75,
+      "gumbel_tau": 0.1778,
+      "image_masks_loss": 0.0194,
+      "llm_learning_rate": 7.745204759195179e-07,
+      "loss": 0.7377,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 7.745204759195179e-05,
+      "step": 7322
+    },
+    {
+      "epoch": 0.75,
+      "gumbel_tau": 0.1777,
+      "image_masks_loss": 0.0041,
+      "llm_learning_rate": 7.739202383089175e-07,
+      "loss": 0.7164,
+      "output_text_mask_loss": 0.0105,
+      "predictor_learning_rate": 7.739202383089175e-05,
+      "step": 7323
+    },
+    {
+      "epoch": 0.75,
+      "gumbel_tau": 0.1777,
+      "image_masks_loss": 0.0079,
+      "llm_learning_rate": 7.733201907780649e-07,
+      "loss": 0.7716,
+      "output_text_mask_loss": 0.0116,
+      "predictor_learning_rate": 7.733201907780648e-05,
+      "step": 7324
+    },
+    {
+      "epoch": 0.75,
+      "gumbel_tau": 0.1777,
+      "image_masks_loss": 0.0198,
+      "llm_learning_rate": 7.727203333930392e-07,
+      "loss": 0.7843,
+      "output_text_mask_loss": 0.0448,
+      "predictor_learning_rate": 7.727203333930391e-05,
+      "step": 7325
+    },
+    {
+      "epoch": 0.75,
+      "gumbel_tau": 0.1776,
+      "image_masks_loss": 0.0076,
+      "llm_learning_rate": 7.721206662198974e-07,
+      "loss": 0.6659,
+      "output_text_mask_loss": 0.0195,
+      "predictor_learning_rate": 7.721206662198974e-05,
+      "step": 7326
+    },
+    {
+      "epoch": 0.75,
+      "gumbel_tau": 0.1776,
+      "image_masks_loss": 0.0189,
+      "llm_learning_rate": 7.715211893246766e-07,
+      "loss": 0.748,
+      "output_text_mask_loss": 0.0126,
+      "predictor_learning_rate": 7.715211893246766e-05,
+      "step": 7327
+    },
+    {
+      "epoch": 0.75,
+      "gumbel_tau": 0.1775,
+      "image_masks_loss": 0.0061,
+      "llm_learning_rate": 7.709219027733928e-07,
+      "loss": 0.7922,
+      "output_text_mask_loss": 0.0809,
+      "predictor_learning_rate": 7.709219027733927e-05,
+      "step": 7328
+    },
+    {
+      "epoch": 0.75,
+      "gumbel_tau": 0.1775,
+      "image_masks_loss": 0.0057,
+      "llm_learning_rate": 7.703228066320401e-07,
+      "loss": 0.8225,
+      "output_text_mask_loss": 0.1221,
+      "predictor_learning_rate": 7.7032280663204e-05,
+      "step": 7329
+    },
+    {
+      "epoch": 0.75,
+      "gumbel_tau": 0.1775,
+      "image_masks_loss": 0.0038,
+      "llm_learning_rate": 7.697239009665919e-07,
+      "loss": 0.8406,
+      "output_text_mask_loss": 0.0115,
+      "predictor_learning_rate": 7.697239009665918e-05,
+      "step": 7330
+    },
+    {
+      "epoch": 0.75,
+      "gumbel_tau": 0.1774,
+      "image_masks_loss": 0.0033,
+      "llm_learning_rate": 7.691251858430026e-07,
+      "loss": 0.6926,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 7.691251858430026e-05,
+      "step": 7331
+    },
+    {
+      "epoch": 0.75,
+      "gumbel_tau": 0.1774,
+      "image_masks_loss": 0.0099,
+      "llm_learning_rate": 7.68526661327203e-07,
+      "loss": 0.6859,
+      "output_text_mask_loss": 0.0132,
+      "predictor_learning_rate": 7.685266613272029e-05,
+      "step": 7332
+    },
+    {
+      "epoch": 0.75,
+      "gumbel_tau": 0.1773,
+      "image_masks_loss": 0.0051,
+      "llm_learning_rate": 7.679283274851046e-07,
+      "loss": 0.7129,
+      "output_text_mask_loss": 0.0055,
+      "predictor_learning_rate": 7.679283274851045e-05,
+      "step": 7333
+    },
+    {
+      "epoch": 0.75,
+      "gumbel_tau": 0.1773,
+      "image_masks_loss": 0.0091,
+      "llm_learning_rate": 7.673301843825965e-07,
+      "loss": 0.8409,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 7.673301843825964e-05,
+      "step": 7334
+    },
+    {
+      "epoch": 0.75,
+      "gumbel_tau": 0.1772,
+      "image_masks_loss": 0.005,
+      "llm_learning_rate": 7.667322320855477e-07,
+      "loss": 0.7383,
+      "output_text_mask_loss": 0.0006,
+      "predictor_learning_rate": 7.667322320855476e-05,
+      "step": 7335
+    },
+    {
+      "epoch": 0.75,
+      "gumbel_tau": 0.1772,
+      "image_masks_loss": 0.0077,
+      "llm_learning_rate": 7.661344706598068e-07,
+      "loss": 0.7778,
+      "output_text_mask_loss": 0.0185,
+      "predictor_learning_rate": 7.661344706598066e-05,
+      "step": 7336
+    },
+    {
+      "epoch": 0.75,
+      "gumbel_tau": 0.1772,
+      "image_masks_loss": 0.0089,
+      "llm_learning_rate": 7.655369001711987e-07,
+      "loss": 0.6391,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 7.655369001711987e-05,
+      "step": 7337
+    },
+    {
+      "epoch": 0.75,
+      "gumbel_tau": 0.1771,
+      "image_masks_loss": 0.0056,
+      "llm_learning_rate": 7.649395206855315e-07,
+      "loss": 0.6933,
+      "output_text_mask_loss": 0.0117,
+      "predictor_learning_rate": 7.649395206855315e-05,
+      "step": 7338
+    },
+    {
+      "epoch": 0.75,
+      "gumbel_tau": 0.1771,
+      "image_masks_loss": 0.0054,
+      "llm_learning_rate": 7.643423322685889e-07,
+      "loss": 0.7229,
+      "output_text_mask_loss": 0.0012,
+      "predictor_learning_rate": 7.643423322685888e-05,
+      "step": 7339
+    },
+    {
+      "epoch": 0.75,
+      "gumbel_tau": 0.177,
+      "image_masks_loss": 0.0122,
+      "llm_learning_rate": 7.637453349861348e-07,
+      "loss": 0.7658,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 7.637453349861348e-05,
+      "step": 7340
+    },
+    {
+      "epoch": 0.75,
+      "gumbel_tau": 0.177,
+      "image_masks_loss": 0.0175,
+      "llm_learning_rate": 7.631485289039117e-07,
+      "loss": 0.7625,
+      "output_text_mask_loss": 0.0034,
+      "predictor_learning_rate": 7.631485289039117e-05,
+      "step": 7341
+    },
+    {
+      "epoch": 0.75,
+      "gumbel_tau": 0.1769,
+      "image_masks_loss": 0.0046,
+      "llm_learning_rate": 7.625519140876414e-07,
+      "loss": 0.7864,
+      "output_text_mask_loss": 0.0156,
+      "predictor_learning_rate": 7.625519140876413e-05,
+      "step": 7342
+    },
+    {
+      "epoch": 0.75,
+      "gumbel_tau": 0.1769,
+      "image_masks_loss": 0.0073,
+      "llm_learning_rate": 7.619554906030241e-07,
+      "loss": 0.7093,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 7.61955490603024e-05,
+      "step": 7343
+    },
+    {
+      "epoch": 0.75,
+      "gumbel_tau": 0.1769,
+      "image_masks_loss": 0.011,
+      "llm_learning_rate": 7.613592585157398e-07,
+      "loss": 0.5589,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 7.613592585157397e-05,
+      "step": 7344
+    },
+    {
+      "epoch": 0.75,
+      "gumbel_tau": 0.1768,
+      "image_masks_loss": 0.0089,
+      "llm_learning_rate": 7.607632178914464e-07,
+      "loss": 0.5571,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 7.607632178914464e-05,
+      "step": 7345
+    },
+    {
+      "epoch": 0.75,
+      "gumbel_tau": 0.1768,
+      "image_masks_loss": 0.01,
+      "llm_learning_rate": 7.601673687957817e-07,
+      "loss": 0.6598,
+      "output_text_mask_loss": 0.0408,
+      "predictor_learning_rate": 7.601673687957816e-05,
+      "step": 7346
+    },
+    {
+      "epoch": 0.75,
+      "gumbel_tau": 0.1767,
+      "image_masks_loss": 0.0121,
+      "llm_learning_rate": 7.595717112943615e-07,
+      "loss": 0.6896,
+      "output_text_mask_loss": 0.0113,
+      "predictor_learning_rate": 7.595717112943615e-05,
+      "step": 7347
+    },
+    {
+      "epoch": 0.75,
+      "gumbel_tau": 0.1767,
+      "image_masks_loss": 0.0073,
+      "llm_learning_rate": 7.58976245452781e-07,
+      "loss": 0.7768,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 7.589762454527809e-05,
+      "step": 7348
+    },
+    {
+      "epoch": 0.75,
+      "gumbel_tau": 0.1767,
+      "image_masks_loss": 0.0084,
+      "llm_learning_rate": 7.583809713366144e-07,
+      "loss": 0.7284,
+      "output_text_mask_loss": 0.025,
+      "predictor_learning_rate": 7.583809713366144e-05,
+      "step": 7349
+    },
+    {
+      "epoch": 0.75,
+      "gumbel_tau": 0.1766,
+      "image_masks_loss": 0.0024,
+      "llm_learning_rate": 7.577858890114137e-07,
+      "loss": 0.7021,
+      "output_text_mask_loss": 0.0137,
+      "predictor_learning_rate": 7.577858890114137e-05,
+      "step": 7350
+    },
+    {
+      "epoch": 0.75,
+      "gumbel_tau": 0.1766,
+      "image_masks_loss": 0.0159,
+      "llm_learning_rate": 7.571909985427125e-07,
+      "loss": 0.7933,
+      "output_text_mask_loss": 0.0025,
+      "predictor_learning_rate": 7.571909985427125e-05,
+      "step": 7351
+    },
+    {
+      "epoch": 0.75,
+      "gumbel_tau": 0.1765,
+      "image_masks_loss": 0.0105,
+      "llm_learning_rate": 7.565962999960205e-07,
+      "loss": 0.7544,
+      "output_text_mask_loss": 0.001,
+      "predictor_learning_rate": 7.565962999960205e-05,
+      "step": 7352
+    },
+    {
+      "epoch": 0.75,
+      "gumbel_tau": 0.1765,
+      "image_masks_loss": 0.0132,
+      "llm_learning_rate": 7.560017934368272e-07,
+      "loss": 0.6835,
+      "output_text_mask_loss": 0.0198,
+      "predictor_learning_rate": 7.560017934368271e-05,
+      "step": 7353
+    },
+    {
+      "epoch": 0.75,
+      "gumbel_tau": 0.1764,
+      "image_masks_loss": 0.0128,
+      "llm_learning_rate": 7.554074789306007e-07,
+      "loss": 0.6998,
+      "output_text_mask_loss": 0.0423,
+      "predictor_learning_rate": 7.554074789306006e-05,
+      "step": 7354
+    },
+    {
+      "epoch": 0.75,
+      "gumbel_tau": 0.1764,
+      "image_masks_loss": 0.0065,
+      "llm_learning_rate": 7.54813356542789e-07,
+      "loss": 0.7424,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 7.54813356542789e-05,
+      "step": 7355
+    },
+    {
+      "epoch": 0.75,
+      "gumbel_tau": 0.1764,
+      "image_masks_loss": 0.0136,
+      "llm_learning_rate": 7.542194263388175e-07,
+      "loss": 0.7442,
+      "output_text_mask_loss": 0.0429,
+      "predictor_learning_rate": 7.542194263388174e-05,
+      "step": 7356
+    },
+    {
+      "epoch": 0.75,
+      "gumbel_tau": 0.1763,
+      "image_masks_loss": 0.0086,
+      "llm_learning_rate": 7.536256883840904e-07,
+      "loss": 0.7297,
+      "output_text_mask_loss": 0.0046,
+      "predictor_learning_rate": 7.536256883840904e-05,
+      "step": 7357
+    },
+    {
+      "epoch": 0.75,
+      "gumbel_tau": 0.1763,
+      "image_masks_loss": 0.0051,
+      "llm_learning_rate": 7.530321427439935e-07,
+      "loss": 0.7598,
+      "output_text_mask_loss": 0.0162,
+      "predictor_learning_rate": 7.530321427439935e-05,
+      "step": 7358
+    },
+    {
+      "epoch": 0.75,
+      "gumbel_tau": 0.1762,
+      "image_masks_loss": 0.0089,
+      "llm_learning_rate": 7.524387894838881e-07,
+      "loss": 0.6963,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 7.52438789483888e-05,
+      "step": 7359
+    },
+    {
+      "epoch": 0.75,
+      "gumbel_tau": 0.1762,
+      "image_masks_loss": 0.0101,
+      "llm_learning_rate": 7.518456286691156e-07,
+      "loss": 0.7057,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 7.518456286691156e-05,
+      "step": 7360
+    },
+    {
+      "epoch": 0.75,
+      "gumbel_tau": 0.1762,
+      "image_masks_loss": 0.0091,
+      "llm_learning_rate": 7.512526603649975e-07,
+      "loss": 0.7081,
+      "output_text_mask_loss": 0.0206,
+      "predictor_learning_rate": 7.512526603649975e-05,
+      "step": 7361
+    },
+    {
+      "epoch": 0.75,
+      "gumbel_tau": 0.1761,
+      "image_masks_loss": 0.0071,
+      "llm_learning_rate": 7.506598846368296e-07,
+      "loss": 0.7053,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 7.506598846368295e-05,
+      "step": 7362
+    },
+    {
+      "epoch": 0.75,
+      "gumbel_tau": 0.1761,
+      "image_masks_loss": 0.0082,
+      "llm_learning_rate": 7.500673015498927e-07,
+      "loss": 0.6467,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 7.500673015498927e-05,
+      "step": 7363
+    },
+    {
+      "epoch": 0.75,
+      "gumbel_tau": 0.176,
+      "image_masks_loss": 0.0098,
+      "llm_learning_rate": 7.494749111694424e-07,
+      "loss": 0.7071,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 7.494749111694424e-05,
+      "step": 7364
+    },
+    {
+      "epoch": 0.75,
+      "gumbel_tau": 0.176,
+      "image_masks_loss": 0.0034,
+      "llm_learning_rate": 7.488827135607141e-07,
+      "loss": 0.7984,
+      "output_text_mask_loss": 0.0172,
+      "predictor_learning_rate": 7.488827135607141e-05,
+      "step": 7365
+    },
+    {
+      "epoch": 0.75,
+      "gumbel_tau": 0.176,
+      "image_masks_loss": 0.0086,
+      "llm_learning_rate": 7.482907087889219e-07,
+      "loss": 0.7637,
+      "output_text_mask_loss": 0.0139,
+      "predictor_learning_rate": 7.482907087889218e-05,
+      "step": 7366
+    },
+    {
+      "epoch": 0.75,
+      "gumbel_tau": 0.1759,
+      "image_masks_loss": 0.0133,
+      "llm_learning_rate": 7.476988969192584e-07,
+      "loss": 0.7488,
+      "output_text_mask_loss": 0.0042,
+      "predictor_learning_rate": 7.476988969192583e-05,
+      "step": 7367
+    },
+    {
+      "epoch": 0.75,
+      "gumbel_tau": 0.1759,
+      "image_masks_loss": 0.0059,
+      "llm_learning_rate": 7.471072780168956e-07,
+      "loss": 0.7033,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 7.471072780168956e-05,
+      "step": 7368
+    },
+    {
+      "epoch": 0.76,
+      "gumbel_tau": 0.1758,
+      "image_masks_loss": 0.0109,
+      "llm_learning_rate": 7.465158521469829e-07,
+      "loss": 0.6377,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 7.465158521469828e-05,
+      "step": 7369
+    },
+    {
+      "epoch": 0.76,
+      "gumbel_tau": 0.1758,
+      "image_masks_loss": 0.0082,
+      "llm_learning_rate": 7.459246193746511e-07,
+      "loss": 0.7097,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 7.45924619374651e-05,
+      "step": 7370
+    },
+    {
+      "epoch": 0.76,
+      "gumbel_tau": 0.1757,
+      "image_masks_loss": 0.0105,
+      "llm_learning_rate": 7.453335797650072e-07,
+      "loss": 0.6551,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 7.453335797650071e-05,
+      "step": 7371
+    },
+    {
+      "epoch": 0.76,
+      "gumbel_tau": 0.1757,
+      "image_masks_loss": 0.0105,
+      "llm_learning_rate": 7.44742733383138e-07,
+      "loss": 0.7334,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 7.447427333831381e-05,
+      "step": 7372
+    },
+    {
+      "epoch": 0.76,
+      "gumbel_tau": 0.1757,
+      "image_masks_loss": 0.0189,
+      "llm_learning_rate": 7.441520802941085e-07,
+      "loss": 0.7934,
+      "output_text_mask_loss": 0.0101,
+      "predictor_learning_rate": 7.441520802941085e-05,
+      "step": 7373
+    },
+    {
+      "epoch": 0.76,
+      "gumbel_tau": 0.1756,
+      "image_masks_loss": 0.0066,
+      "llm_learning_rate": 7.435616205629631e-07,
+      "loss": 0.763,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 7.435616205629631e-05,
+      "step": 7374
+    },
+    {
+      "epoch": 0.76,
+      "gumbel_tau": 0.1756,
+      "image_masks_loss": 0.0102,
+      "llm_learning_rate": 7.429713542547243e-07,
+      "loss": 0.6836,
+      "output_text_mask_loss": 0.0111,
+      "predictor_learning_rate": 7.429713542547243e-05,
+      "step": 7375
+    },
+    {
+      "epoch": 0.76,
+      "gumbel_tau": 0.1755,
+      "image_masks_loss": 0.005,
+      "llm_learning_rate": 7.423812814343928e-07,
+      "loss": 0.7227,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 7.423812814343928e-05,
+      "step": 7376
+    },
+    {
+      "epoch": 0.76,
+      "gumbel_tau": 0.1755,
+      "image_masks_loss": 0.012,
+      "llm_learning_rate": 7.417914021669504e-07,
+      "loss": 0.739,
+      "output_text_mask_loss": 0.0603,
+      "predictor_learning_rate": 7.417914021669503e-05,
+      "step": 7377
+    },
+    {
+      "epoch": 0.76,
+      "gumbel_tau": 0.1755,
+      "image_masks_loss": 0.0103,
+      "llm_learning_rate": 7.412017165173549e-07,
+      "loss": 0.8203,
+      "output_text_mask_loss": 0.0404,
+      "predictor_learning_rate": 7.412017165173548e-05,
+      "step": 7378
+    },
+    {
+      "epoch": 0.76,
+      "gumbel_tau": 0.1754,
+      "image_masks_loss": 0.0013,
+      "llm_learning_rate": 7.40612224550544e-07,
+      "loss": 0.7311,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 7.40612224550544e-05,
+      "step": 7379
+    },
+    {
+      "epoch": 0.76,
+      "gumbel_tau": 0.1754,
+      "image_masks_loss": 0.0082,
+      "llm_learning_rate": 7.400229263314346e-07,
+      "loss": 0.6091,
+      "output_text_mask_loss": 0.036,
+      "predictor_learning_rate": 7.400229263314345e-05,
+      "step": 7380
+    },
+    {
+      "epoch": 0.76,
+      "gumbel_tau": 0.1753,
+      "image_masks_loss": 0.0105,
+      "llm_learning_rate": 7.3943382192492e-07,
+      "loss": 0.5588,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 7.3943382192492e-05,
+      "step": 7381
+    },
+    {
+      "epoch": 0.76,
+      "gumbel_tau": 0.1753,
+      "image_masks_loss": 0.0126,
+      "llm_learning_rate": 7.388449113958738e-07,
+      "loss": 0.682,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 7.388449113958737e-05,
+      "step": 7382
+    },
+    {
+      "epoch": 0.76,
+      "gumbel_tau": 0.1752,
+      "image_masks_loss": 0.0079,
+      "llm_learning_rate": 7.382561948091493e-07,
+      "loss": 0.723,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 7.382561948091493e-05,
+      "step": 7383
+    },
+    {
+      "epoch": 0.76,
+      "gumbel_tau": 0.1752,
+      "image_masks_loss": 0.0106,
+      "llm_learning_rate": 7.376676722295767e-07,
+      "loss": 0.6785,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 7.376676722295766e-05,
+      "step": 7384
+    },
+    {
+      "epoch": 0.76,
+      "gumbel_tau": 0.1752,
+      "image_masks_loss": 0.0048,
+      "llm_learning_rate": 7.370793437219656e-07,
+      "loss": 0.762,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 7.370793437219655e-05,
+      "step": 7385
+    },
+    {
+      "epoch": 0.76,
+      "gumbel_tau": 0.1751,
+      "image_masks_loss": 0.0107,
+      "llm_learning_rate": 7.36491209351104e-07,
+      "loss": 0.6712,
+      "output_text_mask_loss": 0.0111,
+      "predictor_learning_rate": 7.36491209351104e-05,
+      "step": 7386
+    },
+    {
+      "epoch": 0.76,
+      "gumbel_tau": 0.1751,
+      "image_masks_loss": 0.0083,
+      "llm_learning_rate": 7.359032691817586e-07,
+      "loss": 0.747,
+      "output_text_mask_loss": 0.0347,
+      "predictor_learning_rate": 7.359032691817586e-05,
+      "step": 7387
+    },
+    {
+      "epoch": 0.76,
+      "gumbel_tau": 0.175,
+      "image_masks_loss": 0.0074,
+      "llm_learning_rate": 7.353155232786743e-07,
+      "loss": 0.7079,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 7.353155232786743e-05,
+      "step": 7388
+    },
+    {
+      "epoch": 0.76,
+      "gumbel_tau": 0.175,
+      "image_masks_loss": 0.0067,
+      "llm_learning_rate": 7.347279717065744e-07,
+      "loss": 0.766,
+      "output_text_mask_loss": 0.0307,
+      "predictor_learning_rate": 7.347279717065744e-05,
+      "step": 7389
+    },
+    {
+      "epoch": 0.76,
+      "gumbel_tau": 0.175,
+      "image_masks_loss": 0.0043,
+      "llm_learning_rate": 7.341406145301633e-07,
+      "loss": 0.585,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 7.341406145301633e-05,
+      "step": 7390
+    },
+    {
+      "epoch": 0.76,
+      "gumbel_tau": 0.1749,
+      "image_masks_loss": 0.0172,
+      "llm_learning_rate": 7.33553451814121e-07,
+      "loss": 0.7084,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 7.335534518141209e-05,
+      "step": 7391
+    },
+    {
+      "epoch": 0.76,
+      "gumbel_tau": 0.1749,
+      "image_masks_loss": 0.004,
+      "llm_learning_rate": 7.329664836231071e-07,
+      "loss": 0.7265,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 7.32966483623107e-05,
+      "step": 7392
+    },
+    {
+      "epoch": 0.76,
+      "gumbel_tau": 0.1748,
+      "image_masks_loss": 0.0034,
+      "llm_learning_rate": 7.3237971002176e-07,
+      "loss": 0.7273,
+      "output_text_mask_loss": 0.0193,
+      "predictor_learning_rate": 7.3237971002176e-05,
+      "step": 7393
+    },
+    {
+      "epoch": 0.76,
+      "gumbel_tau": 0.1748,
+      "image_masks_loss": 0.0091,
+      "llm_learning_rate": 7.317931310746964e-07,
+      "loss": 0.7843,
+      "output_text_mask_loss": 0.0454,
+      "predictor_learning_rate": 7.317931310746964e-05,
+      "step": 7394
+    },
+    {
+      "epoch": 0.76,
+      "gumbel_tau": 0.1748,
+      "image_masks_loss": 0.0037,
+      "llm_learning_rate": 7.312067468465117e-07,
+      "loss": 0.7101,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 7.312067468465117e-05,
+      "step": 7395
+    },
+    {
+      "epoch": 0.76,
+      "gumbel_tau": 0.1747,
+      "image_masks_loss": 0.0069,
+      "llm_learning_rate": 7.306205574017791e-07,
+      "loss": 0.7667,
+      "output_text_mask_loss": 0.0172,
+      "predictor_learning_rate": 7.306205574017791e-05,
+      "step": 7396
+    },
+    {
+      "epoch": 0.76,
+      "gumbel_tau": 0.1747,
+      "image_masks_loss": 0.0079,
+      "llm_learning_rate": 7.300345628050528e-07,
+      "loss": 0.7181,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 7.300345628050529e-05,
+      "step": 7397
+    },
+    {
+      "epoch": 0.76,
+      "gumbel_tau": 0.1746,
+      "image_masks_loss": 0.0039,
+      "llm_learning_rate": 7.294487631208636e-07,
+      "loss": 0.7555,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 7.294487631208636e-05,
+      "step": 7398
+    },
+    {
+      "epoch": 0.76,
+      "gumbel_tau": 0.1746,
+      "image_masks_loss": 0.0028,
+      "llm_learning_rate": 7.288631584137196e-07,
+      "loss": 0.8235,
+      "output_text_mask_loss": 0.0012,
+      "predictor_learning_rate": 7.288631584137196e-05,
+      "step": 7399
+    },
+    {
+      "epoch": 0.76,
+      "gumbel_tau": 0.1745,
+      "image_masks_loss": 0.0138,
+      "llm_learning_rate": 7.282777487481094e-07,
+      "loss": 0.7096,
+      "output_text_mask_loss": 0.033,
+      "predictor_learning_rate": 7.282777487481094e-05,
+      "step": 7400
+    },
+    {
+      "epoch": 0.76,
+      "gumbel_tau": 0.1745,
+      "image_masks_loss": 0.0046,
+      "llm_learning_rate": 7.276925341884991e-07,
+      "loss": 0.6881,
+      "output_text_mask_loss": 0.0078,
+      "predictor_learning_rate": 7.27692534188499e-05,
+      "step": 7401
+    },
+    {
+      "epoch": 0.76,
+      "gumbel_tau": 0.1745,
+      "image_masks_loss": 0.0097,
+      "llm_learning_rate": 7.271075147993356e-07,
+      "loss": 0.7994,
+      "output_text_mask_loss": 0.0051,
+      "predictor_learning_rate": 7.271075147993355e-05,
+      "step": 7402
+    },
+    {
+      "epoch": 0.76,
+      "gumbel_tau": 0.1744,
+      "image_masks_loss": 0.0098,
+      "llm_learning_rate": 7.265226906450415e-07,
+      "loss": 0.7207,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 7.265226906450415e-05,
+      "step": 7403
+    },
+    {
+      "epoch": 0.76,
+      "gumbel_tau": 0.1744,
+      "image_masks_loss": 0.0028,
+      "llm_learning_rate": 7.259380617900188e-07,
+      "loss": 0.7203,
+      "output_text_mask_loss": 0.001,
+      "predictor_learning_rate": 7.259380617900188e-05,
+      "step": 7404
+    },
+    {
+      "epoch": 0.76,
+      "gumbel_tau": 0.1743,
+      "image_masks_loss": 0.0077,
+      "llm_learning_rate": 7.253536282986487e-07,
+      "loss": 0.6836,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 7.253536282986486e-05,
+      "step": 7405
+    },
+    {
+      "epoch": 0.76,
+      "gumbel_tau": 0.1743,
+      "image_masks_loss": 0.0057,
+      "llm_learning_rate": 7.247693902352895e-07,
+      "loss": 0.749,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 7.247693902352895e-05,
+      "step": 7406
+    },
+    {
+      "epoch": 0.76,
+      "gumbel_tau": 0.1743,
+      "image_masks_loss": 0.0042,
+      "llm_learning_rate": 7.241853476642793e-07,
+      "loss": 0.6861,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 7.241853476642793e-05,
+      "step": 7407
+    },
+    {
+      "epoch": 0.76,
+      "gumbel_tau": 0.1742,
+      "image_masks_loss": 0.0031,
+      "llm_learning_rate": 7.236015006499337e-07,
+      "loss": 0.6111,
+      "output_text_mask_loss": NaN,
+      "predictor_learning_rate": 7.236015006499336e-05,
+      "step": 7408
+    },
+    {
+      "epoch": 0.76,
+      "gumbel_tau": 0.1742,
+      "image_masks_loss": 0.0049,
+      "llm_learning_rate": 7.230178492565479e-07,
+      "loss": 0.7452,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 7.230178492565479e-05,
+      "step": 7409
+    },
+    {
+      "epoch": 0.76,
+      "gumbel_tau": 0.1741,
+      "image_masks_loss": 0.0143,
+      "llm_learning_rate": 7.224343935483949e-07,
+      "loss": 0.7667,
+      "output_text_mask_loss": 0.0012,
+      "predictor_learning_rate": 7.224343935483949e-05,
+      "step": 7410
+    },
+    {
+      "epoch": 0.76,
+      "gumbel_tau": 0.1741,
+      "image_masks_loss": 0.0073,
+      "llm_learning_rate": 7.218511335897255e-07,
+      "loss": 0.6735,
+      "output_text_mask_loss": 0.0095,
+      "predictor_learning_rate": 7.218511335897255e-05,
+      "step": 7411
+    },
+    {
+      "epoch": 0.76,
+      "gumbel_tau": 0.1741,
+      "image_masks_loss": 0.0138,
+      "llm_learning_rate": 7.212680694447702e-07,
+      "loss": 0.7605,
+      "output_text_mask_loss": 0.0071,
+      "predictor_learning_rate": 7.212680694447701e-05,
+      "step": 7412
+    },
+    {
+      "epoch": 0.76,
+      "gumbel_tau": 0.174,
+      "image_masks_loss": 0.0105,
+      "llm_learning_rate": 7.206852011777371e-07,
+      "loss": 0.7186,
+      "output_text_mask_loss": 0.0374,
+      "predictor_learning_rate": 7.206852011777371e-05,
+      "step": 7413
+    },
+    {
+      "epoch": 0.76,
+      "gumbel_tau": 0.174,
+      "image_masks_loss": 0.0143,
+      "llm_learning_rate": 7.20102528852813e-07,
+      "loss": 0.7243,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 7.201025288528129e-05,
+      "step": 7414
+    },
+    {
+      "epoch": 0.76,
+      "gumbel_tau": 0.1739,
+      "image_masks_loss": 0.0118,
+      "llm_learning_rate": 7.19520052534162e-07,
+      "loss": 0.7887,
+      "output_text_mask_loss": 0.0179,
+      "predictor_learning_rate": 7.19520052534162e-05,
+      "step": 7415
+    },
+    {
+      "epoch": 0.76,
+      "gumbel_tau": 0.1739,
+      "image_masks_loss": 0.0067,
+      "llm_learning_rate": 7.189377722859294e-07,
+      "loss": 0.6343,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 7.189377722859294e-05,
+      "step": 7416
+    },
+    {
+      "epoch": 0.76,
+      "gumbel_tau": 0.1738,
+      "image_masks_loss": 0.0153,
+      "llm_learning_rate": 7.183556881722375e-07,
+      "loss": 0.8203,
+      "output_text_mask_loss": 0.0011,
+      "predictor_learning_rate": 7.183556881722376e-05,
+      "step": 7417
+    },
+    {
+      "epoch": 0.76,
+      "gumbel_tau": 0.1738,
+      "image_masks_loss": 0.0021,
+      "llm_learning_rate": 7.177738002571849e-07,
+      "loss": 0.7299,
+      "output_text_mask_loss": 0.0095,
+      "predictor_learning_rate": 7.177738002571848e-05,
+      "step": 7418
+    },
+    {
+      "epoch": 0.76,
+      "gumbel_tau": 0.1738,
+      "image_masks_loss": 0.0046,
+      "llm_learning_rate": 7.171921086048508e-07,
+      "loss": 0.7209,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 7.171921086048507e-05,
+      "step": 7419
+    },
+    {
+      "epoch": 0.76,
+      "gumbel_tau": 0.1737,
+      "image_masks_loss": 0.0062,
+      "llm_learning_rate": 7.166106132792932e-07,
+      "loss": 0.8125,
+      "output_text_mask_loss": 0.0381,
+      "predictor_learning_rate": 7.166106132792932e-05,
+      "step": 7420
+    },
+    {
+      "epoch": 0.76,
+      "gumbel_tau": 0.1737,
+      "image_masks_loss": 0.0141,
+      "llm_learning_rate": 7.160293143445463e-07,
+      "loss": 0.6742,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 7.160293143445463e-05,
+      "step": 7421
+    },
+    {
+      "epoch": 0.76,
+      "gumbel_tau": 0.1736,
+      "image_masks_loss": 0.0056,
+      "llm_learning_rate": 7.154482118646255e-07,
+      "loss": 0.7024,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 7.154482118646254e-05,
+      "step": 7422
+    },
+    {
+      "epoch": 0.76,
+      "gumbel_tau": 0.1736,
+      "image_masks_loss": 0.0087,
+      "llm_learning_rate": 7.148673059035228e-07,
+      "loss": 0.7252,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 7.148673059035227e-05,
+      "step": 7423
+    },
+    {
+      "epoch": 0.76,
+      "gumbel_tau": 0.1736,
+      "image_masks_loss": 0.0062,
+      "llm_learning_rate": 7.142865965252086e-07,
+      "loss": 0.6838,
+      "output_text_mask_loss": 0.002,
+      "predictor_learning_rate": 7.142865965252085e-05,
+      "step": 7424
+    },
+    {
+      "epoch": 0.76,
+      "gumbel_tau": 0.1735,
+      "image_masks_loss": 0.005,
+      "llm_learning_rate": 7.137060837936316e-07,
+      "loss": 0.6971,
+      "output_text_mask_loss": 0.0056,
+      "predictor_learning_rate": 7.137060837936315e-05,
+      "step": 7425
+    },
+    {
+      "epoch": 0.76,
+      "gumbel_tau": 0.1735,
+      "image_masks_loss": 0.0108,
+      "llm_learning_rate": 7.131257677727197e-07,
+      "loss": 0.7461,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 7.131257677727196e-05,
+      "step": 7426
+    },
+    {
+      "epoch": 0.76,
+      "gumbel_tau": 0.1734,
+      "image_masks_loss": 0.0072,
+      "llm_learning_rate": 7.125456485263785e-07,
+      "loss": 0.6653,
+      "output_text_mask_loss": 0.0058,
+      "predictor_learning_rate": 7.125456485263784e-05,
+      "step": 7427
+    },
+    {
+      "epoch": 0.76,
+      "gumbel_tau": 0.1734,
+      "image_masks_loss": 0.0052,
+      "llm_learning_rate": 7.119657261184909e-07,
+      "loss": 0.6689,
+      "output_text_mask_loss": 0.0011,
+      "predictor_learning_rate": 7.119657261184909e-05,
+      "step": 7428
+    },
+    {
+      "epoch": 0.76,
+      "gumbel_tau": 0.1734,
+      "image_masks_loss": 0.009,
+      "llm_learning_rate": 7.113860006129209e-07,
+      "loss": 0.8171,
+      "output_text_mask_loss": 0.0134,
+      "predictor_learning_rate": 7.113860006129209e-05,
+      "step": 7429
+    },
+    {
+      "epoch": 0.76,
+      "gumbel_tau": 0.1733,
+      "image_masks_loss": 0.01,
+      "llm_learning_rate": 7.108064720735086e-07,
+      "loss": 0.7515,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 7.108064720735085e-05,
+      "step": 7430
+    },
+    {
+      "epoch": 0.76,
+      "gumbel_tau": 0.1733,
+      "image_masks_loss": 0.0026,
+      "llm_learning_rate": 7.102271405640732e-07,
+      "loss": 0.8268,
+      "output_text_mask_loss": 0.0073,
+      "predictor_learning_rate": 7.102271405640732e-05,
+      "step": 7431
+    },
+    {
+      "epoch": 0.76,
+      "gumbel_tau": 0.1732,
+      "image_masks_loss": 0.0054,
+      "llm_learning_rate": 7.096480061484112e-07,
+      "loss": 0.7647,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 7.096480061484111e-05,
+      "step": 7432
+    },
+    {
+      "epoch": 0.76,
+      "gumbel_tau": 0.1732,
+      "image_masks_loss": 0.0069,
+      "llm_learning_rate": 7.090690688902988e-07,
+      "loss": 0.724,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 7.090690688902987e-05,
+      "step": 7433
+    },
+    {
+      "epoch": 0.76,
+      "gumbel_tau": 0.1731,
+      "image_masks_loss": 0.0068,
+      "llm_learning_rate": 7.084903288534897e-07,
+      "loss": 0.6336,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 7.084903288534897e-05,
+      "step": 7434
+    },
+    {
+      "epoch": 0.76,
+      "gumbel_tau": 0.1731,
+      "image_masks_loss": 0.0113,
+      "llm_learning_rate": 7.07911786101716e-07,
+      "loss": 0.7363,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 7.07911786101716e-05,
+      "step": 7435
+    },
+    {
+      "epoch": 0.76,
+      "gumbel_tau": 0.1731,
+      "image_masks_loss": 0.0043,
+      "llm_learning_rate": 7.073334406986879e-07,
+      "loss": 0.7558,
+      "output_text_mask_loss": 0.0065,
+      "predictor_learning_rate": 7.073334406986878e-05,
+      "step": 7436
+    },
+    {
+      "epoch": 0.76,
+      "gumbel_tau": 0.173,
+      "image_masks_loss": 0.0057,
+      "llm_learning_rate": 7.067552927080945e-07,
+      "loss": 0.7706,
+      "output_text_mask_loss": 0.0238,
+      "predictor_learning_rate": 7.067552927080944e-05,
+      "step": 7437
+    },
+    {
+      "epoch": 0.76,
+      "gumbel_tau": 0.173,
+      "image_masks_loss": 0.009,
+      "llm_learning_rate": 7.061773421936027e-07,
+      "loss": 0.7987,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 7.061773421936027e-05,
+      "step": 7438
+    },
+    {
+      "epoch": 0.76,
+      "gumbel_tau": 0.1729,
+      "image_masks_loss": 0.0071,
+      "llm_learning_rate": 7.055995892188575e-07,
+      "loss": 0.7216,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 7.055995892188574e-05,
+      "step": 7439
+    },
+    {
+      "epoch": 0.76,
+      "gumbel_tau": 0.1729,
+      "image_masks_loss": 0.008,
+      "llm_learning_rate": 7.050220338474815e-07,
+      "loss": 0.7664,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 7.050220338474814e-05,
+      "step": 7440
+    },
+    {
+      "epoch": 0.76,
+      "gumbel_tau": 0.1729,
+      "image_masks_loss": 0.003,
+      "llm_learning_rate": 7.044446761430782e-07,
+      "loss": 0.8561,
+      "output_text_mask_loss": 0.0345,
+      "predictor_learning_rate": 7.044446761430781e-05,
+      "step": 7441
+    },
+    {
+      "epoch": 0.76,
+      "gumbel_tau": 0.1728,
+      "image_masks_loss": 0.0111,
+      "llm_learning_rate": 7.038675161692265e-07,
+      "loss": 0.7483,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 7.038675161692264e-05,
+      "step": 7442
+    },
+    {
+      "epoch": 0.76,
+      "gumbel_tau": 0.1728,
+      "image_masks_loss": 0.0227,
+      "llm_learning_rate": 7.032905539894847e-07,
+      "loss": 0.8524,
+      "output_text_mask_loss": 0.0042,
+      "predictor_learning_rate": 7.032905539894846e-05,
+      "step": 7443
+    },
+    {
+      "epoch": 0.76,
+      "gumbel_tau": 0.1727,
+      "image_masks_loss": 0.0069,
+      "llm_learning_rate": 7.027137896673891e-07,
+      "loss": 0.6926,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 7.027137896673891e-05,
+      "step": 7444
+    },
+    {
+      "epoch": 0.76,
+      "gumbel_tau": 0.1727,
+      "image_masks_loss": 0.0032,
+      "llm_learning_rate": 7.021372232664542e-07,
+      "loss": 0.7277,
+      "output_text_mask_loss": 0.0177,
+      "predictor_learning_rate": 7.021372232664541e-05,
+      "step": 7445
+    },
+    {
+      "epoch": 0.76,
+      "gumbel_tau": 0.1727,
+      "image_masks_loss": 0.0059,
+      "llm_learning_rate": 7.015608548501729e-07,
+      "loss": 0.7402,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 7.015608548501728e-05,
+      "step": 7446
+    },
+    {
+      "epoch": 0.76,
+      "gumbel_tau": 0.1726,
+      "image_masks_loss": 0.0099,
+      "llm_learning_rate": 7.009846844820156e-07,
+      "loss": 0.684,
+      "output_text_mask_loss": 0.0027,
+      "predictor_learning_rate": 7.009846844820155e-05,
+      "step": 7447
+    },
+    {
+      "epoch": 0.76,
+      "gumbel_tau": 0.1726,
+      "image_masks_loss": 0.0046,
+      "llm_learning_rate": 7.004087122254325e-07,
+      "loss": 0.6868,
+      "output_text_mask_loss": 0.0088,
+      "predictor_learning_rate": 7.004087122254325e-05,
+      "step": 7448
+    },
+    {
+      "epoch": 0.76,
+      "gumbel_tau": 0.1725,
+      "image_masks_loss": 0.0028,
+      "llm_learning_rate": 6.998329381438504e-07,
+      "loss": 0.6897,
+      "output_text_mask_loss": 0.0095,
+      "predictor_learning_rate": 6.998329381438504e-05,
+      "step": 7449
+    },
+    {
+      "epoch": 0.76,
+      "gumbel_tau": 0.1725,
+      "image_masks_loss": 0.0113,
+      "llm_learning_rate": 6.992573623006751e-07,
+      "loss": 0.7596,
+      "output_text_mask_loss": 0.0004,
+      "predictor_learning_rate": 6.992573623006752e-05,
+      "step": 7450
+    },
+    {
+      "epoch": 0.76,
+      "gumbel_tau": 0.1725,
+      "image_masks_loss": 0.0089,
+      "llm_learning_rate": 6.986819847592899e-07,
+      "loss": 0.6665,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 6.986819847592899e-05,
+      "step": 7451
+    },
+    {
+      "epoch": 0.76,
+      "gumbel_tau": 0.1724,
+      "image_masks_loss": 0.0019,
+      "llm_learning_rate": 6.981068055830567e-07,
+      "loss": 0.7716,
+      "output_text_mask_loss": 0.0086,
+      "predictor_learning_rate": 6.981068055830567e-05,
+      "step": 7452
+    },
+    {
+      "epoch": 0.76,
+      "gumbel_tau": 0.1724,
+      "image_masks_loss": 0.0031,
+      "llm_learning_rate": 6.975318248353158e-07,
+      "loss": 0.7157,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 6.975318248353158e-05,
+      "step": 7453
+    },
+    {
+      "epoch": 0.76,
+      "gumbel_tau": 0.1723,
+      "image_masks_loss": 0.0023,
+      "llm_learning_rate": 6.969570425793853e-07,
+      "loss": 0.742,
+      "output_text_mask_loss": 0.021,
+      "predictor_learning_rate": 6.969570425793854e-05,
+      "step": 7454
+    },
+    {
+      "epoch": 0.76,
+      "gumbel_tau": 0.1723,
+      "image_masks_loss": 0.0087,
+      "llm_learning_rate": 6.963824588785611e-07,
+      "loss": 0.7294,
+      "output_text_mask_loss": 0.0479,
+      "predictor_learning_rate": 6.963824588785611e-05,
+      "step": 7455
+    },
+    {
+      "epoch": 0.76,
+      "gumbel_tau": 0.1723,
+      "image_masks_loss": 0.0054,
+      "llm_learning_rate": 6.958080737961181e-07,
+      "loss": 0.7465,
+      "output_text_mask_loss": 0.006,
+      "predictor_learning_rate": 6.958080737961181e-05,
+      "step": 7456
+    },
+    {
+      "epoch": 0.76,
+      "gumbel_tau": 0.1722,
+      "image_masks_loss": 0.0061,
+      "llm_learning_rate": 6.952338873953088e-07,
+      "loss": 0.764,
+      "output_text_mask_loss": 0.0038,
+      "predictor_learning_rate": 6.952338873953087e-05,
+      "step": 7457
+    },
+    {
+      "epoch": 0.76,
+      "gumbel_tau": 0.1722,
+      "image_masks_loss": 0.0034,
+      "llm_learning_rate": 6.946598997393636e-07,
+      "loss": 0.7052,
+      "output_text_mask_loss": 0.0486,
+      "predictor_learning_rate": 6.946598997393635e-05,
+      "step": 7458
+    },
+    {
+      "epoch": 0.76,
+      "gumbel_tau": 0.1721,
+      "image_masks_loss": 0.01,
+      "llm_learning_rate": 6.940861108914914e-07,
+      "loss": 0.7353,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 6.940861108914914e-05,
+      "step": 7459
+    },
+    {
+      "epoch": 0.76,
+      "gumbel_tau": 0.1721,
+      "image_masks_loss": 0.006,
+      "llm_learning_rate": 6.935125209148788e-07,
+      "loss": 0.703,
+      "output_text_mask_loss": 0.0372,
+      "predictor_learning_rate": 6.935125209148788e-05,
+      "step": 7460
+    },
+    {
+      "epoch": 0.76,
+      "gumbel_tau": 0.1721,
+      "image_masks_loss": 0.0133,
+      "llm_learning_rate": 6.929391298726917e-07,
+      "loss": 0.6793,
+      "output_text_mask_loss": 0.0107,
+      "predictor_learning_rate": 6.929391298726918e-05,
+      "step": 7461
+    },
+    {
+      "epoch": 0.76,
+      "gumbel_tau": 0.172,
+      "image_masks_loss": 0.0196,
+      "llm_learning_rate": 6.923659378280726e-07,
+      "loss": 0.7718,
+      "output_text_mask_loss": 0.0025,
+      "predictor_learning_rate": 6.923659378280725e-05,
+      "step": 7462
+    },
+    {
+      "epoch": 0.76,
+      "gumbel_tau": 0.172,
+      "image_masks_loss": 0.0078,
+      "llm_learning_rate": 6.91792944844143e-07,
+      "loss": 0.6899,
+      "output_text_mask_loss": 0.0012,
+      "predictor_learning_rate": 6.91792944844143e-05,
+      "step": 7463
+    },
+    {
+      "epoch": 0.76,
+      "gumbel_tau": 0.1719,
+      "image_masks_loss": 0.0033,
+      "llm_learning_rate": 6.912201509840016e-07,
+      "loss": 0.7909,
+      "output_text_mask_loss": 0.0511,
+      "predictor_learning_rate": 6.912201509840016e-05,
+      "step": 7464
+    },
+    {
+      "epoch": 0.76,
+      "gumbel_tau": 0.1719,
+      "image_masks_loss": 0.0042,
+      "llm_learning_rate": 6.906475563107264e-07,
+      "loss": 0.8361,
+      "output_text_mask_loss": 0.0278,
+      "predictor_learning_rate": 6.906475563107265e-05,
+      "step": 7465
+    },
+    {
+      "epoch": 0.76,
+      "gumbel_tau": 0.1718,
+      "image_masks_loss": 0.0087,
+      "llm_learning_rate": 6.900751608873723e-07,
+      "loss": 0.7643,
+      "output_text_mask_loss": 0.0041,
+      "predictor_learning_rate": 6.900751608873723e-05,
+      "step": 7466
+    },
+    {
+      "epoch": 0.77,
+      "gumbel_tau": 0.1718,
+      "image_masks_loss": 0.0175,
+      "llm_learning_rate": 6.895029647769724e-07,
+      "loss": 0.725,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 6.895029647769723e-05,
+      "step": 7467
+    },
+    {
+      "epoch": 0.77,
+      "gumbel_tau": 0.1718,
+      "image_masks_loss": 0.0093,
+      "llm_learning_rate": 6.889309680425394e-07,
+      "loss": 0.7026,
+      "output_text_mask_loss": 0.0019,
+      "predictor_learning_rate": 6.889309680425394e-05,
+      "step": 7468
+    },
+    {
+      "epoch": 0.77,
+      "gumbel_tau": 0.1717,
+      "image_masks_loss": 0.0149,
+      "llm_learning_rate": 6.883591707470618e-07,
+      "loss": 0.77,
+      "output_text_mask_loss": 0.0124,
+      "predictor_learning_rate": 6.883591707470619e-05,
+      "step": 7469
+    },
+    {
+      "epoch": 0.77,
+      "gumbel_tau": 0.1717,
+      "image_masks_loss": 0.0116,
+      "llm_learning_rate": 6.87787572953508e-07,
+      "loss": 0.7158,
+      "output_text_mask_loss": 0.0271,
+      "predictor_learning_rate": 6.87787572953508e-05,
+      "step": 7470
+    },
+    {
+      "epoch": 0.77,
+      "gumbel_tau": 0.1716,
+      "image_masks_loss": 0.0069,
+      "llm_learning_rate": 6.872161747248232e-07,
+      "loss": 0.7529,
+      "output_text_mask_loss": 0.0039,
+      "predictor_learning_rate": 6.872161747248232e-05,
+      "step": 7471
+    },
+    {
+      "epoch": 0.77,
+      "gumbel_tau": 0.1716,
+      "image_masks_loss": 0.0058,
+      "llm_learning_rate": 6.866449761239318e-07,
+      "loss": 0.774,
+      "output_text_mask_loss": 0.0219,
+      "predictor_learning_rate": 6.866449761239319e-05,
+      "step": 7472
+    },
+    {
+      "epoch": 0.77,
+      "gumbel_tau": 0.1716,
+      "image_masks_loss": 0.015,
+      "llm_learning_rate": 6.860739772137331e-07,
+      "loss": 0.7149,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 6.860739772137331e-05,
+      "step": 7473
+    },
+    {
+      "epoch": 0.77,
+      "gumbel_tau": 0.1715,
+      "image_masks_loss": 0.0061,
+      "llm_learning_rate": 6.855031780571092e-07,
+      "loss": 0.668,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 6.855031780571092e-05,
+      "step": 7474
+    },
+    {
+      "epoch": 0.77,
+      "gumbel_tau": 0.1715,
+      "image_masks_loss": 0.0046,
+      "llm_learning_rate": 6.84932578716917e-07,
+      "loss": 0.7476,
+      "output_text_mask_loss": 0.106,
+      "predictor_learning_rate": 6.849325787169169e-05,
+      "step": 7475
+    },
+    {
+      "epoch": 0.77,
+      "gumbel_tau": 0.1714,
+      "image_masks_loss": 0.0062,
+      "llm_learning_rate": 6.843621792559921e-07,
+      "loss": 0.7452,
+      "output_text_mask_loss": 0.0235,
+      "predictor_learning_rate": 6.84362179255992e-05,
+      "step": 7476
+    },
+    {
+      "epoch": 0.77,
+      "gumbel_tau": 0.1714,
+      "image_masks_loss": 0.0055,
+      "llm_learning_rate": 6.837919797371481e-07,
+      "loss": 0.6884,
+      "output_text_mask_loss": 0.0174,
+      "predictor_learning_rate": 6.83791979737148e-05,
+      "step": 7477
+    },
+    {
+      "epoch": 0.77,
+      "gumbel_tau": 0.1714,
+      "image_masks_loss": 0.0104,
+      "llm_learning_rate": 6.832219802231769e-07,
+      "loss": 0.6676,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 6.832219802231767e-05,
+      "step": 7478
+    },
+    {
+      "epoch": 0.77,
+      "gumbel_tau": 0.1713,
+      "image_masks_loss": 0.0093,
+      "llm_learning_rate": 6.82652180776848e-07,
+      "loss": 0.725,
+      "output_text_mask_loss": 0.0141,
+      "predictor_learning_rate": 6.826521807768479e-05,
+      "step": 7479
+    },
+    {
+      "epoch": 0.77,
+      "gumbel_tau": 0.1713,
+      "image_masks_loss": 0.0078,
+      "llm_learning_rate": 6.82082581460908e-07,
+      "loss": 0.6513,
+      "output_text_mask_loss": 0.0116,
+      "predictor_learning_rate": 6.82082581460908e-05,
+      "step": 7480
+    },
+    {
+      "epoch": 0.77,
+      "gumbel_tau": 0.1712,
+      "image_masks_loss": 0.0094,
+      "llm_learning_rate": 6.815131823380841e-07,
+      "loss": 0.7323,
+      "output_text_mask_loss": 0.0056,
+      "predictor_learning_rate": 6.815131823380841e-05,
+      "step": 7481
+    },
+    {
+      "epoch": 0.77,
+      "gumbel_tau": 0.1712,
+      "image_masks_loss": 0.0108,
+      "llm_learning_rate": 6.809439834710793e-07,
+      "loss": 0.7879,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 6.809439834710793e-05,
+      "step": 7482
+    },
+    {
+      "epoch": 0.77,
+      "gumbel_tau": 0.1712,
+      "image_masks_loss": 0.0082,
+      "llm_learning_rate": 6.803749849225746e-07,
+      "loss": 0.6935,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 6.803749849225745e-05,
+      "step": 7483
+    },
+    {
+      "epoch": 0.77,
+      "gumbel_tau": 0.1711,
+      "image_masks_loss": 0.0092,
+      "llm_learning_rate": 6.7980618675523e-07,
+      "loss": 0.7429,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 6.7980618675523e-05,
+      "step": 7484
+    },
+    {
+      "epoch": 0.77,
+      "gumbel_tau": 0.1711,
+      "image_masks_loss": 0.0048,
+      "llm_learning_rate": 6.792375890316822e-07,
+      "loss": 0.611,
+      "output_text_mask_loss": 0.0324,
+      "predictor_learning_rate": 6.79237589031682e-05,
+      "step": 7485
+    },
+    {
+      "epoch": 0.77,
+      "gumbel_tau": 0.171,
+      "image_masks_loss": 0.0075,
+      "llm_learning_rate": 6.786691918145463e-07,
+      "loss": 0.7944,
+      "output_text_mask_loss": 0.0131,
+      "predictor_learning_rate": 6.786691918145463e-05,
+      "step": 7486
+    },
+    {
+      "epoch": 0.77,
+      "gumbel_tau": 0.171,
+      "image_masks_loss": 0.0033,
+      "llm_learning_rate": 6.781009951664155e-07,
+      "loss": 0.7024,
+      "output_text_mask_loss": 0.0099,
+      "predictor_learning_rate": 6.781009951664155e-05,
+      "step": 7487
+    },
+    {
+      "epoch": 0.77,
+      "gumbel_tau": 0.171,
+      "image_masks_loss": 0.0059,
+      "llm_learning_rate": 6.77532999149862e-07,
+      "loss": 0.8019,
+      "output_text_mask_loss": 0.0046,
+      "predictor_learning_rate": 6.775329991498618e-05,
+      "step": 7488
+    },
+    {
+      "epoch": 0.77,
+      "gumbel_tau": 0.1709,
+      "image_masks_loss": 0.0043,
+      "llm_learning_rate": 6.769652038274341e-07,
+      "loss": 0.6784,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 6.76965203827434e-05,
+      "step": 7489
+    },
+    {
+      "epoch": 0.77,
+      "gumbel_tau": 0.1709,
+      "image_masks_loss": 0.0037,
+      "llm_learning_rate": 6.763976092616583e-07,
+      "loss": 0.7236,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 6.763976092616583e-05,
+      "step": 7490
+    },
+    {
+      "epoch": 0.77,
+      "gumbel_tau": 0.1708,
+      "image_masks_loss": 0.0094,
+      "llm_learning_rate": 6.758302155150407e-07,
+      "loss": 0.7516,
+      "output_text_mask_loss": 0.0035,
+      "predictor_learning_rate": 6.758302155150406e-05,
+      "step": 7491
+    },
+    {
+      "epoch": 0.77,
+      "gumbel_tau": 0.1708,
+      "image_masks_loss": 0.0158,
+      "llm_learning_rate": 6.752630226500614e-07,
+      "loss": 0.7859,
+      "output_text_mask_loss": 0.0174,
+      "predictor_learning_rate": 6.752630226500615e-05,
+      "step": 7492
+    },
+    {
+      "epoch": 0.77,
+      "gumbel_tau": 0.1708,
+      "image_masks_loss": 0.0041,
+      "llm_learning_rate": 6.74696030729183e-07,
+      "loss": 0.752,
+      "output_text_mask_loss": 0.0197,
+      "predictor_learning_rate": 6.74696030729183e-05,
+      "step": 7493
+    },
+    {
+      "epoch": 0.77,
+      "gumbel_tau": 0.1707,
+      "image_masks_loss": 0.0107,
+      "llm_learning_rate": 6.741292398148436e-07,
+      "loss": 0.6452,
+      "output_text_mask_loss": 0.0023,
+      "predictor_learning_rate": 6.741292398148435e-05,
+      "step": 7494
+    },
+    {
+      "epoch": 0.77,
+      "gumbel_tau": 0.1707,
+      "image_masks_loss": 0.0087,
+      "llm_learning_rate": 6.735626499694595e-07,
+      "loss": 0.8383,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 6.735626499694595e-05,
+      "step": 7495
+    },
+    {
+      "epoch": 0.77,
+      "gumbel_tau": 0.1706,
+      "image_masks_loss": 0.0076,
+      "llm_learning_rate": 6.729962612554242e-07,
+      "loss": 0.7636,
+      "output_text_mask_loss": 0.0595,
+      "predictor_learning_rate": 6.729962612554241e-05,
+      "step": 7496
+    },
+    {
+      "epoch": 0.77,
+      "gumbel_tau": 0.1706,
+      "image_masks_loss": 0.0048,
+      "llm_learning_rate": 6.724300737351103e-07,
+      "loss": 0.6333,
+      "output_text_mask_loss": 0.0322,
+      "predictor_learning_rate": 6.724300737351103e-05,
+      "step": 7497
+    },
+    {
+      "epoch": 0.77,
+      "gumbel_tau": 0.1706,
+      "image_masks_loss": 0.0078,
+      "llm_learning_rate": 6.718640874708673e-07,
+      "loss": 0.6372,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 6.718640874708673e-05,
+      "step": 7498
+    },
+    {
+      "epoch": 0.77,
+      "gumbel_tau": 0.1705,
+      "image_masks_loss": 0.0058,
+      "llm_learning_rate": 6.712983025250225e-07,
+      "loss": 0.7451,
+      "output_text_mask_loss": 0.0538,
+      "predictor_learning_rate": 6.712983025250225e-05,
+      "step": 7499
+    },
+    {
+      "epoch": 0.77,
+      "gumbel_tau": 0.1705,
+      "image_masks_loss": 0.0056,
+      "llm_learning_rate": 6.707327189598825e-07,
+      "loss": 0.7152,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 6.707327189598824e-05,
+      "step": 7500
+    },
+    {
+      "epoch": 0.77,
+      "gumbel_tau": 0.1704,
+      "image_masks_loss": 0.0047,
+      "llm_learning_rate": 6.701673368377301e-07,
+      "loss": 0.6591,
+      "output_text_mask_loss": 0.0137,
+      "predictor_learning_rate": 6.7016733683773e-05,
+      "step": 7501
+    },
+    {
+      "epoch": 0.77,
+      "gumbel_tau": 0.1704,
+      "image_masks_loss": 0.0085,
+      "llm_learning_rate": 6.696021562208263e-07,
+      "loss": 0.7989,
+      "output_text_mask_loss": 0.0012,
+      "predictor_learning_rate": 6.696021562208263e-05,
+      "step": 7502
+    },
+    {
+      "epoch": 0.77,
+      "gumbel_tau": 0.1704,
+      "image_masks_loss": 0.0088,
+      "llm_learning_rate": 6.6903717717141e-07,
+      "loss": 0.7821,
+      "output_text_mask_loss": 0.0019,
+      "predictor_learning_rate": 6.690371771714099e-05,
+      "step": 7503
+    },
+    {
+      "epoch": 0.77,
+      "gumbel_tau": 0.1703,
+      "image_masks_loss": 0.0068,
+      "llm_learning_rate": 6.684723997516982e-07,
+      "loss": 0.6889,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 6.684723997516981e-05,
+      "step": 7504
+    },
+    {
+      "epoch": 0.77,
+      "gumbel_tau": 0.1703,
+      "image_masks_loss": 0.0034,
+      "llm_learning_rate": 6.679078240238851e-07,
+      "loss": 0.7031,
+      "output_text_mask_loss": 0.0012,
+      "predictor_learning_rate": 6.679078240238851e-05,
+      "step": 7505
+    },
+    {
+      "epoch": 0.77,
+      "gumbel_tau": 0.1702,
+      "image_masks_loss": 0.0036,
+      "llm_learning_rate": 6.673434500501428e-07,
+      "loss": 0.8307,
+      "output_text_mask_loss": 0.0261,
+      "predictor_learning_rate": 6.673434500501426e-05,
+      "step": 7506
+    },
+    {
+      "epoch": 0.77,
+      "gumbel_tau": 0.1702,
+      "image_masks_loss": 0.0046,
+      "llm_learning_rate": 6.667792778926224e-07,
+      "loss": 0.6637,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 6.667792778926223e-05,
+      "step": 7507
+    },
+    {
+      "epoch": 0.77,
+      "gumbel_tau": 0.1702,
+      "image_masks_loss": 0.0142,
+      "llm_learning_rate": 6.662153076134515e-07,
+      "loss": 0.7271,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 6.662153076134514e-05,
+      "step": 7508
+    },
+    {
+      "epoch": 0.77,
+      "gumbel_tau": 0.1701,
+      "image_masks_loss": 0.0103,
+      "llm_learning_rate": 6.656515392747354e-07,
+      "loss": 0.6894,
+      "output_text_mask_loss": 0.0055,
+      "predictor_learning_rate": 6.656515392747353e-05,
+      "step": 7509
+    },
+    {
+      "epoch": 0.77,
+      "gumbel_tau": 0.1701,
+      "image_masks_loss": 0.0058,
+      "llm_learning_rate": 6.650879729385581e-07,
+      "loss": 0.7801,
+      "output_text_mask_loss": 0.0107,
+      "predictor_learning_rate": 6.650879729385581e-05,
+      "step": 7510
+    },
+    {
+      "epoch": 0.77,
+      "gumbel_tau": 0.17,
+      "image_masks_loss": 0.005,
+      "llm_learning_rate": 6.645246086669799e-07,
+      "loss": 0.7784,
+      "output_text_mask_loss": 0.034,
+      "predictor_learning_rate": 6.645246086669799e-05,
+      "step": 7511
+    },
+    {
+      "epoch": 0.77,
+      "gumbel_tau": 0.17,
+      "image_masks_loss": 0.0034,
+      "llm_learning_rate": 6.639614465220392e-07,
+      "loss": 0.7519,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 6.639614465220392e-05,
+      "step": 7512
+    },
+    {
+      "epoch": 0.77,
+      "gumbel_tau": 0.17,
+      "image_masks_loss": 0.0083,
+      "llm_learning_rate": 6.633984865657547e-07,
+      "loss": 0.8424,
+      "output_text_mask_loss": 0.0144,
+      "predictor_learning_rate": 6.633984865657546e-05,
+      "step": 7513
+    },
+    {
+      "epoch": 0.77,
+      "gumbel_tau": 0.1699,
+      "image_masks_loss": 0.0063,
+      "llm_learning_rate": 6.628357288601195e-07,
+      "loss": 0.6983,
+      "output_text_mask_loss": 0.0063,
+      "predictor_learning_rate": 6.628357288601195e-05,
+      "step": 7514
+    },
+    {
+      "epoch": 0.77,
+      "gumbel_tau": 0.1699,
+      "image_masks_loss": 0.0093,
+      "llm_learning_rate": 6.622731734671059e-07,
+      "loss": 0.8366,
+      "output_text_mask_loss": 0.0233,
+      "predictor_learning_rate": 6.622731734671058e-05,
+      "step": 7515
+    },
+    {
+      "epoch": 0.77,
+      "gumbel_tau": 0.1698,
+      "image_masks_loss": 0.0084,
+      "llm_learning_rate": 6.617108204486642e-07,
+      "loss": 0.7544,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 6.617108204486641e-05,
+      "step": 7516
+    },
+    {
+      "epoch": 0.77,
+      "gumbel_tau": 0.1698,
+      "image_masks_loss": 0.0053,
+      "llm_learning_rate": 6.61148669866721e-07,
+      "loss": 0.7561,
+      "output_text_mask_loss": 0.2151,
+      "predictor_learning_rate": 6.61148669866721e-05,
+      "step": 7517
+    },
+    {
+      "epoch": 0.77,
+      "gumbel_tau": 0.1698,
+      "image_masks_loss": 0.0075,
+      "llm_learning_rate": 6.605867217831826e-07,
+      "loss": 0.7066,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 6.605867217831827e-05,
+      "step": 7518
+    },
+    {
+      "epoch": 0.77,
+      "gumbel_tau": 0.1697,
+      "image_masks_loss": 0.0051,
+      "llm_learning_rate": 6.600249762599306e-07,
+      "loss": 0.7074,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 6.600249762599306e-05,
+      "step": 7519
+    },
+    {
+      "epoch": 0.77,
+      "gumbel_tau": 0.1697,
+      "image_masks_loss": 0.003,
+      "llm_learning_rate": 6.594634333588273e-07,
+      "loss": 0.8607,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 6.594634333588274e-05,
+      "step": 7520
+    },
+    {
+      "epoch": 0.77,
+      "gumbel_tau": 0.1696,
+      "image_masks_loss": 0.0071,
+      "llm_learning_rate": 6.589020931417106e-07,
+      "loss": 0.6941,
+      "output_text_mask_loss": 0.0004,
+      "predictor_learning_rate": 6.589020931417105e-05,
+      "step": 7521
+    },
+    {
+      "epoch": 0.77,
+      "gumbel_tau": 0.1696,
+      "image_masks_loss": 0.0057,
+      "llm_learning_rate": 6.58340955670396e-07,
+      "loss": 0.748,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 6.583409556703959e-05,
+      "step": 7522
+    },
+    {
+      "epoch": 0.77,
+      "gumbel_tau": 0.1696,
+      "image_masks_loss": 0.0092,
+      "llm_learning_rate": 6.577800210066776e-07,
+      "loss": 0.7049,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 6.577800210066775e-05,
+      "step": 7523
+    },
+    {
+      "epoch": 0.77,
+      "gumbel_tau": 0.1695,
+      "image_masks_loss": 0.0114,
+      "llm_learning_rate": 6.572192892123266e-07,
+      "loss": 0.7192,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 6.572192892123266e-05,
+      "step": 7524
+    },
+    {
+      "epoch": 0.77,
+      "gumbel_tau": 0.1695,
+      "image_masks_loss": 0.0113,
+      "llm_learning_rate": 6.566587603490924e-07,
+      "loss": 0.7979,
+      "output_text_mask_loss": 0.0156,
+      "predictor_learning_rate": 6.566587603490922e-05,
+      "step": 7525
+    },
+    {
+      "epoch": 0.77,
+      "gumbel_tau": 0.1694,
+      "image_masks_loss": 0.0089,
+      "llm_learning_rate": 6.560984344787002e-07,
+      "loss": 0.6204,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 6.560984344787002e-05,
+      "step": 7526
+    },
+    {
+      "epoch": 0.77,
+      "gumbel_tau": 0.1694,
+      "image_masks_loss": 0.0065,
+      "llm_learning_rate": 6.555383116628566e-07,
+      "loss": 0.7395,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 6.555383116628566e-05,
+      "step": 7527
+    },
+    {
+      "epoch": 0.77,
+      "gumbel_tau": 0.1694,
+      "image_masks_loss": 0.0085,
+      "llm_learning_rate": 6.549783919632433e-07,
+      "loss": 0.8148,
+      "output_text_mask_loss": 0.0162,
+      "predictor_learning_rate": 6.549783919632432e-05,
+      "step": 7528
+    },
+    {
+      "epoch": 0.77,
+      "gumbel_tau": 0.1693,
+      "image_masks_loss": 0.015,
+      "llm_learning_rate": 6.54418675441518e-07,
+      "loss": 0.7728,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 6.544186754415179e-05,
+      "step": 7529
+    },
+    {
+      "epoch": 0.77,
+      "gumbel_tau": 0.1693,
+      "image_masks_loss": 0.0028,
+      "llm_learning_rate": 6.538591621593194e-07,
+      "loss": 0.7384,
+      "output_text_mask_loss": 0.0162,
+      "predictor_learning_rate": 6.538591621593192e-05,
+      "step": 7530
+    },
+    {
+      "epoch": 0.77,
+      "gumbel_tau": 0.1692,
+      "image_masks_loss": 0.0093,
+      "llm_learning_rate": 6.53299852178261e-07,
+      "loss": 0.7391,
+      "output_text_mask_loss": 0.0286,
+      "predictor_learning_rate": 6.53299852178261e-05,
+      "step": 7531
+    },
+    {
+      "epoch": 0.77,
+      "gumbel_tau": 0.1692,
+      "image_masks_loss": 0.0042,
+      "llm_learning_rate": 6.527407455599372e-07,
+      "loss": 0.7342,
+      "output_text_mask_loss": 0.0261,
+      "predictor_learning_rate": 6.527407455599372e-05,
+      "step": 7532
+    },
+    {
+      "epoch": 0.77,
+      "gumbel_tau": 0.1692,
+      "image_masks_loss": 0.007,
+      "llm_learning_rate": 6.521818423659168e-07,
+      "loss": 0.7659,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 6.521818423659168e-05,
+      "step": 7533
+    },
+    {
+      "epoch": 0.77,
+      "gumbel_tau": 0.1691,
+      "image_masks_loss": 0.0066,
+      "llm_learning_rate": 6.51623142657748e-07,
+      "loss": 0.7342,
+      "output_text_mask_loss": 0.0141,
+      "predictor_learning_rate": 6.516231426577479e-05,
+      "step": 7534
+    },
+    {
+      "epoch": 0.77,
+      "gumbel_tau": 0.1691,
+      "image_masks_loss": 0.003,
+      "llm_learning_rate": 6.510646464969558e-07,
+      "loss": 0.7672,
+      "output_text_mask_loss": 0.005,
+      "predictor_learning_rate": 6.510646464969557e-05,
+      "step": 7535
+    },
+    {
+      "epoch": 0.77,
+      "gumbel_tau": 0.169,
+      "image_masks_loss": 0.0066,
+      "llm_learning_rate": 6.505063539450429e-07,
+      "loss": 0.7774,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 6.505063539450429e-05,
+      "step": 7536
+    },
+    {
+      "epoch": 0.77,
+      "gumbel_tau": 0.169,
+      "image_masks_loss": 0.0094,
+      "llm_learning_rate": 6.4994826506349e-07,
+      "loss": 0.8128,
+      "output_text_mask_loss": 0.0141,
+      "predictor_learning_rate": 6.499482650634899e-05,
+      "step": 7537
+    },
+    {
+      "epoch": 0.77,
+      "gumbel_tau": 0.169,
+      "image_masks_loss": 0.0095,
+      "llm_learning_rate": 6.493903799137544e-07,
+      "loss": 0.6781,
+      "output_text_mask_loss": 0.024,
+      "predictor_learning_rate": 6.493903799137544e-05,
+      "step": 7538
+    },
+    {
+      "epoch": 0.77,
+      "gumbel_tau": 0.1689,
+      "image_masks_loss": 0.0021,
+      "llm_learning_rate": 6.488326985572727e-07,
+      "loss": 0.708,
+      "output_text_mask_loss": 0.0376,
+      "predictor_learning_rate": 6.488326985572726e-05,
+      "step": 7539
+    },
+    {
+      "epoch": 0.77,
+      "gumbel_tau": 0.1689,
+      "image_masks_loss": 0.0047,
+      "llm_learning_rate": 6.482752210554577e-07,
+      "loss": 0.7408,
+      "output_text_mask_loss": 0.0055,
+      "predictor_learning_rate": 6.482752210554576e-05,
+      "step": 7540
+    },
+    {
+      "epoch": 0.77,
+      "gumbel_tau": 0.1688,
+      "image_masks_loss": 0.0032,
+      "llm_learning_rate": 6.477179474696998e-07,
+      "loss": 0.7559,
+      "output_text_mask_loss": 0.0202,
+      "predictor_learning_rate": 6.477179474696998e-05,
+      "step": 7541
+    },
+    {
+      "epoch": 0.77,
+      "gumbel_tau": 0.1688,
+      "image_masks_loss": 0.0027,
+      "llm_learning_rate": 6.471608778613675e-07,
+      "loss": 0.77,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 6.471608778613674e-05,
+      "step": 7542
+    },
+    {
+      "epoch": 0.77,
+      "gumbel_tau": 0.1688,
+      "image_masks_loss": 0.0041,
+      "llm_learning_rate": 6.466040122918063e-07,
+      "loss": 0.7149,
+      "output_text_mask_loss": 0.0801,
+      "predictor_learning_rate": 6.466040122918062e-05,
+      "step": 7543
+    },
+    {
+      "epoch": 0.77,
+      "gumbel_tau": 0.1687,
+      "image_masks_loss": 0.011,
+      "llm_learning_rate": 6.460473508223397e-07,
+      "loss": 0.7439,
+      "output_text_mask_loss": 0.0238,
+      "predictor_learning_rate": 6.460473508223397e-05,
+      "step": 7544
+    },
+    {
+      "epoch": 0.77,
+      "gumbel_tau": 0.1687,
+      "image_masks_loss": 0.0082,
+      "llm_learning_rate": 6.454908935142673e-07,
+      "loss": 0.6654,
+      "output_text_mask_loss": 0.0103,
+      "predictor_learning_rate": 6.454908935142672e-05,
+      "step": 7545
+    },
+    {
+      "epoch": 0.77,
+      "gumbel_tau": 0.1686,
+      "image_masks_loss": 0.0154,
+      "llm_learning_rate": 6.449346404288694e-07,
+      "loss": 0.7504,
+      "output_text_mask_loss": 0.0345,
+      "predictor_learning_rate": 6.449346404288694e-05,
+      "step": 7546
+    },
+    {
+      "epoch": 0.77,
+      "gumbel_tau": 0.1686,
+      "image_masks_loss": 0.0057,
+      "llm_learning_rate": 6.443785916274015e-07,
+      "loss": 0.7861,
+      "output_text_mask_loss": 0.0023,
+      "predictor_learning_rate": 6.443785916274015e-05,
+      "step": 7547
+    },
+    {
+      "epoch": 0.77,
+      "gumbel_tau": 0.1686,
+      "image_masks_loss": 0.0085,
+      "llm_learning_rate": 6.438227471710959e-07,
+      "loss": 0.7473,
+      "output_text_mask_loss": 0.0124,
+      "predictor_learning_rate": 6.438227471710958e-05,
+      "step": 7548
+    },
+    {
+      "epoch": 0.77,
+      "gumbel_tau": 0.1685,
+      "image_masks_loss": 0.0203,
+      "llm_learning_rate": 6.432671071211636e-07,
+      "loss": 0.67,
+      "output_text_mask_loss": 0.0092,
+      "predictor_learning_rate": 6.432671071211635e-05,
+      "step": 7549
+    },
+    {
+      "epoch": 0.77,
+      "gumbel_tau": 0.1685,
+      "image_masks_loss": 0.0059,
+      "llm_learning_rate": 6.42711671538793e-07,
+      "loss": 0.7504,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 6.427116715387929e-05,
+      "step": 7550
+    },
+    {
+      "epoch": 0.77,
+      "gumbel_tau": 0.1684,
+      "image_masks_loss": 0.0089,
+      "llm_learning_rate": 6.421564404851496e-07,
+      "loss": 0.7478,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 6.421564404851495e-05,
+      "step": 7551
+    },
+    {
+      "epoch": 0.77,
+      "gumbel_tau": 0.1684,
+      "image_masks_loss": 0.0037,
+      "llm_learning_rate": 6.416014140213778e-07,
+      "loss": 0.7149,
+      "output_text_mask_loss": 0.0071,
+      "predictor_learning_rate": 6.416014140213777e-05,
+      "step": 7552
+    },
+    {
+      "epoch": 0.77,
+      "gumbel_tau": 0.1684,
+      "image_masks_loss": 0.0038,
+      "llm_learning_rate": 6.410465922085979e-07,
+      "loss": 0.5885,
+      "output_text_mask_loss": 0.0032,
+      "predictor_learning_rate": 6.410465922085978e-05,
+      "step": 7553
+    },
+    {
+      "epoch": 0.77,
+      "gumbel_tau": 0.1683,
+      "image_masks_loss": 0.0071,
+      "llm_learning_rate": 6.404919751079075e-07,
+      "loss": 0.7388,
+      "output_text_mask_loss": 0.0116,
+      "predictor_learning_rate": 6.404919751079075e-05,
+      "step": 7554
+    },
+    {
+      "epoch": 0.77,
+      "gumbel_tau": 0.1683,
+      "image_masks_loss": 0.0148,
+      "llm_learning_rate": 6.399375627803831e-07,
+      "loss": 0.7602,
+      "output_text_mask_loss": 0.0252,
+      "predictor_learning_rate": 6.39937562780383e-05,
+      "step": 7555
+    },
+    {
+      "epoch": 0.77,
+      "gumbel_tau": 0.1682,
+      "image_masks_loss": 0.0062,
+      "llm_learning_rate": 6.393833552870773e-07,
+      "loss": 0.6236,
+      "output_text_mask_loss": 0.0198,
+      "predictor_learning_rate": 6.393833552870773e-05,
+      "step": 7556
+    },
+    {
+      "epoch": 0.77,
+      "gumbel_tau": 0.1682,
+      "image_masks_loss": 0.0069,
+      "llm_learning_rate": 6.388293526890207e-07,
+      "loss": 0.5855,
+      "output_text_mask_loss": 0.0162,
+      "predictor_learning_rate": 6.388293526890207e-05,
+      "step": 7557
+    },
+    {
+      "epoch": 0.77,
+      "gumbel_tau": 0.1682,
+      "image_masks_loss": 0.0115,
+      "llm_learning_rate": 6.382755550472205e-07,
+      "loss": 0.5395,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 6.382755550472205e-05,
+      "step": 7558
+    },
+    {
+      "epoch": 0.77,
+      "gumbel_tau": 0.1681,
+      "image_masks_loss": 0.007,
+      "llm_learning_rate": 6.377219624226641e-07,
+      "loss": 0.7322,
+      "output_text_mask_loss": 0.0294,
+      "predictor_learning_rate": 6.37721962422664e-05,
+      "step": 7559
+    },
+    {
+      "epoch": 0.77,
+      "gumbel_tau": 0.1681,
+      "image_masks_loss": 0.0089,
+      "llm_learning_rate": 6.371685748763131e-07,
+      "loss": 0.7192,
+      "output_text_mask_loss": 0.0193,
+      "predictor_learning_rate": 6.371685748763131e-05,
+      "step": 7560
+    },
+    {
+      "epoch": 0.77,
+      "gumbel_tau": 0.168,
+      "image_masks_loss": 0.0082,
+      "llm_learning_rate": 6.366153924691079e-07,
+      "loss": 0.7628,
+      "output_text_mask_loss": 0.0187,
+      "predictor_learning_rate": 6.366153924691078e-05,
+      "step": 7561
+    },
+    {
+      "epoch": 0.77,
+      "gumbel_tau": 0.168,
+      "image_masks_loss": 0.0173,
+      "llm_learning_rate": 6.360624152619663e-07,
+      "loss": 0.6995,
+      "output_text_mask_loss": 0.0263,
+      "predictor_learning_rate": 6.360624152619663e-05,
+      "step": 7562
+    },
+    {
+      "epoch": 0.77,
+      "gumbel_tau": 0.168,
+      "image_masks_loss": 0.0073,
+      "llm_learning_rate": 6.355096433157831e-07,
+      "loss": 0.64,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 6.355096433157831e-05,
+      "step": 7563
+    },
+    {
+      "epoch": 0.78,
+      "gumbel_tau": 0.1679,
+      "image_masks_loss": 0.0032,
+      "llm_learning_rate": 6.349570766914312e-07,
+      "loss": 0.7445,
+      "output_text_mask_loss": 0.0231,
+      "predictor_learning_rate": 6.349570766914312e-05,
+      "step": 7564
+    },
+    {
+      "epoch": 0.78,
+      "gumbel_tau": 0.1679,
+      "image_masks_loss": 0.0087,
+      "llm_learning_rate": 6.344047154497599e-07,
+      "loss": 0.7266,
+      "output_text_mask_loss": 0.0315,
+      "predictor_learning_rate": 6.344047154497597e-05,
+      "step": 7565
+    },
+    {
+      "epoch": 0.78,
+      "gumbel_tau": 0.1678,
+      "image_masks_loss": 0.0073,
+      "llm_learning_rate": 6.338525596515971e-07,
+      "loss": 0.6893,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 6.33852559651597e-05,
+      "step": 7566
+    },
+    {
+      "epoch": 0.78,
+      "gumbel_tau": 0.1678,
+      "image_masks_loss": 0.0173,
+      "llm_learning_rate": 6.33300609357747e-07,
+      "loss": 0.7927,
+      "output_text_mask_loss": 0.0042,
+      "predictor_learning_rate": 6.33300609357747e-05,
+      "step": 7567
+    },
+    {
+      "epoch": 0.78,
+      "gumbel_tau": 0.1678,
+      "image_masks_loss": 0.0073,
+      "llm_learning_rate": 6.327488646289917e-07,
+      "loss": 0.6471,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 6.327488646289917e-05,
+      "step": 7568
+    },
+    {
+      "epoch": 0.78,
+      "gumbel_tau": 0.1677,
+      "image_masks_loss": 0.0082,
+      "llm_learning_rate": 6.321973255260905e-07,
+      "loss": 0.7824,
+      "output_text_mask_loss": 0.0127,
+      "predictor_learning_rate": 6.321973255260905e-05,
+      "step": 7569
+    },
+    {
+      "epoch": 0.78,
+      "gumbel_tau": 0.1677,
+      "image_masks_loss": 0.0104,
+      "llm_learning_rate": 6.316459921097795e-07,
+      "loss": 0.723,
+      "output_text_mask_loss": 0.0052,
+      "predictor_learning_rate": 6.316459921097795e-05,
+      "step": 7570
+    },
+    {
+      "epoch": 0.78,
+      "gumbel_tau": 0.1676,
+      "image_masks_loss": 0.0094,
+      "llm_learning_rate": 6.310948644407744e-07,
+      "loss": 0.7495,
+      "output_text_mask_loss": 0.001,
+      "predictor_learning_rate": 6.310948644407744e-05,
+      "step": 7571
+    },
+    {
+      "epoch": 0.78,
+      "gumbel_tau": 0.1676,
+      "image_masks_loss": 0.0168,
+      "llm_learning_rate": 6.305439425797655e-07,
+      "loss": 0.6921,
+      "output_text_mask_loss": 0.0217,
+      "predictor_learning_rate": 6.305439425797655e-05,
+      "step": 7572
+    },
+    {
+      "epoch": 0.78,
+      "gumbel_tau": 0.1676,
+      "image_masks_loss": 0.0063,
+      "llm_learning_rate": 6.299932265874215e-07,
+      "loss": 0.7256,
+      "output_text_mask_loss": 0.0099,
+      "predictor_learning_rate": 6.299932265874214e-05,
+      "step": 7573
+    },
+    {
+      "epoch": 0.78,
+      "gumbel_tau": 0.1675,
+      "image_masks_loss": 0.0073,
+      "llm_learning_rate": 6.294427165243888e-07,
+      "loss": 0.7084,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 6.294427165243888e-05,
+      "step": 7574
+    },
+    {
+      "epoch": 0.78,
+      "gumbel_tau": 0.1675,
+      "image_masks_loss": 0.0023,
+      "llm_learning_rate": 6.288924124512907e-07,
+      "loss": 0.7558,
+      "output_text_mask_loss": 0.0015,
+      "predictor_learning_rate": 6.288924124512907e-05,
+      "step": 7575
+    },
+    {
+      "epoch": 0.78,
+      "gumbel_tau": 0.1674,
+      "image_masks_loss": 0.0117,
+      "llm_learning_rate": 6.283423144287279e-07,
+      "loss": 0.6346,
+      "output_text_mask_loss": 0.0433,
+      "predictor_learning_rate": 6.283423144287279e-05,
+      "step": 7576
+    },
+    {
+      "epoch": 0.78,
+      "gumbel_tau": 0.1674,
+      "image_masks_loss": 0.0033,
+      "llm_learning_rate": 6.277924225172777e-07,
+      "loss": 0.7362,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 6.277924225172777e-05,
+      "step": 7577
+    },
+    {
+      "epoch": 0.78,
+      "gumbel_tau": 0.1674,
+      "image_masks_loss": 0.0063,
+      "llm_learning_rate": 6.272427367774966e-07,
+      "loss": 0.7183,
+      "output_text_mask_loss": 0.001,
+      "predictor_learning_rate": 6.272427367774966e-05,
+      "step": 7578
+    },
+    {
+      "epoch": 0.78,
+      "gumbel_tau": 0.1673,
+      "image_masks_loss": 0.0132,
+      "llm_learning_rate": 6.26693257269917e-07,
+      "loss": 0.7278,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 6.266932572699169e-05,
+      "step": 7579
+    },
+    {
+      "epoch": 0.78,
+      "gumbel_tau": 0.1673,
+      "image_masks_loss": 0.0111,
+      "llm_learning_rate": 6.261439840550485e-07,
+      "loss": 0.7187,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 6.261439840550484e-05,
+      "step": 7580
+    },
+    {
+      "epoch": 0.78,
+      "gumbel_tau": 0.1672,
+      "image_masks_loss": 0.0064,
+      "llm_learning_rate": 6.255949171933781e-07,
+      "loss": 0.7518,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 6.255949171933781e-05,
+      "step": 7581
+    },
+    {
+      "epoch": 0.78,
+      "gumbel_tau": 0.1672,
+      "image_masks_loss": 0.0043,
+      "llm_learning_rate": 6.250460567453706e-07,
+      "loss": 0.7279,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 6.250460567453705e-05,
+      "step": 7582
+    },
+    {
+      "epoch": 0.78,
+      "gumbel_tau": 0.1672,
+      "image_masks_loss": 0.0104,
+      "llm_learning_rate": 6.244974027714676e-07,
+      "loss": 0.6381,
+      "output_text_mask_loss": 0.0519,
+      "predictor_learning_rate": 6.244974027714675e-05,
+      "step": 7583
+    },
+    {
+      "epoch": 0.78,
+      "gumbel_tau": 0.1671,
+      "image_masks_loss": 0.0014,
+      "llm_learning_rate": 6.239489553320883e-07,
+      "loss": 0.6679,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 6.239489553320881e-05,
+      "step": 7584
+    },
+    {
+      "epoch": 0.78,
+      "gumbel_tau": 0.1671,
+      "image_masks_loss": 0.0137,
+      "llm_learning_rate": 6.234007144876286e-07,
+      "loss": 0.6858,
+      "output_text_mask_loss": 0.0153,
+      "predictor_learning_rate": 6.234007144876286e-05,
+      "step": 7585
+    },
+    {
+      "epoch": 0.78,
+      "gumbel_tau": 0.1671,
+      "image_masks_loss": 0.0085,
+      "llm_learning_rate": 6.228526802984622e-07,
+      "loss": 0.768,
+      "output_text_mask_loss": 0.0286,
+      "predictor_learning_rate": 6.228526802984622e-05,
+      "step": 7586
+    },
+    {
+      "epoch": 0.78,
+      "gumbel_tau": 0.167,
+      "image_masks_loss": 0.0125,
+      "llm_learning_rate": 6.223048528249398e-07,
+      "loss": 0.6881,
+      "output_text_mask_loss": 0.0206,
+      "predictor_learning_rate": 6.223048528249398e-05,
+      "step": 7587
+    },
+    {
+      "epoch": 0.78,
+      "gumbel_tau": 0.167,
+      "image_masks_loss": 0.0122,
+      "llm_learning_rate": 6.217572321273896e-07,
+      "loss": 0.7015,
+      "output_text_mask_loss": 0.0299,
+      "predictor_learning_rate": 6.217572321273896e-05,
+      "step": 7588
+    },
+    {
+      "epoch": 0.78,
+      "gumbel_tau": 0.1669,
+      "image_masks_loss": 0.0098,
+      "llm_learning_rate": 6.212098182661164e-07,
+      "loss": 0.7494,
+      "output_text_mask_loss": 0.0058,
+      "predictor_learning_rate": 6.212098182661164e-05,
+      "step": 7589
+    },
+    {
+      "epoch": 0.78,
+      "gumbel_tau": 0.1669,
+      "image_masks_loss": 0.0152,
+      "llm_learning_rate": 6.206626113014022e-07,
+      "loss": 0.7342,
+      "output_text_mask_loss": 0.0107,
+      "predictor_learning_rate": 6.206626113014021e-05,
+      "step": 7590
+    },
+    {
+      "epoch": 0.78,
+      "gumbel_tau": 0.1669,
+      "image_masks_loss": 0.0034,
+      "llm_learning_rate": 6.20115611293508e-07,
+      "loss": 0.734,
+      "output_text_mask_loss": 0.0839,
+      "predictor_learning_rate": 6.20115611293508e-05,
+      "step": 7591
+    },
+    {
+      "epoch": 0.78,
+      "gumbel_tau": 0.1668,
+      "image_masks_loss": 0.0216,
+      "llm_learning_rate": 6.195688183026702e-07,
+      "loss": 0.666,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 6.195688183026701e-05,
+      "step": 7592
+    },
+    {
+      "epoch": 0.78,
+      "gumbel_tau": 0.1668,
+      "image_masks_loss": 0.0051,
+      "llm_learning_rate": 6.190222323891024e-07,
+      "loss": 0.7756,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 6.190222323891023e-05,
+      "step": 7593
+    },
+    {
+      "epoch": 0.78,
+      "gumbel_tau": 0.1667,
+      "image_masks_loss": 0.0074,
+      "llm_learning_rate": 6.184758536129962e-07,
+      "loss": 0.7395,
+      "output_text_mask_loss": 0.0012,
+      "predictor_learning_rate": 6.184758536129961e-05,
+      "step": 7594
+    },
+    {
+      "epoch": 0.78,
+      "gumbel_tau": 0.1667,
+      "image_masks_loss": 0.0045,
+      "llm_learning_rate": 6.179296820345199e-07,
+      "loss": 0.7661,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 6.179296820345198e-05,
+      "step": 7595
+    },
+    {
+      "epoch": 0.78,
+      "gumbel_tau": 0.1667,
+      "image_masks_loss": 0.0061,
+      "llm_learning_rate": 6.173837177138189e-07,
+      "loss": 0.7169,
+      "output_text_mask_loss": 0.0046,
+      "predictor_learning_rate": 6.173837177138189e-05,
+      "step": 7596
+    },
+    {
+      "epoch": 0.78,
+      "gumbel_tau": 0.1666,
+      "image_masks_loss": 0.0207,
+      "llm_learning_rate": 6.168379607110159e-07,
+      "loss": 0.7449,
+      "output_text_mask_loss": 0.0139,
+      "predictor_learning_rate": 6.168379607110159e-05,
+      "step": 7597
+    },
+    {
+      "epoch": 0.78,
+      "gumbel_tau": 0.1666,
+      "image_masks_loss": 0.0032,
+      "llm_learning_rate": 6.162924110862118e-07,
+      "loss": 0.6379,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 6.162924110862118e-05,
+      "step": 7598
+    },
+    {
+      "epoch": 0.78,
+      "gumbel_tau": 0.1665,
+      "image_masks_loss": 0.0049,
+      "llm_learning_rate": 6.157470688994832e-07,
+      "loss": 0.6984,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 6.157470688994831e-05,
+      "step": 7599
+    },
+    {
+      "epoch": 0.78,
+      "gumbel_tau": 0.1665,
+      "image_masks_loss": 0.004,
+      "llm_learning_rate": 6.152019342108845e-07,
+      "loss": 0.7625,
+      "output_text_mask_loss": 0.0111,
+      "predictor_learning_rate": 6.152019342108844e-05,
+      "step": 7600
+    },
+    {
+      "epoch": 0.78,
+      "gumbel_tau": 0.1665,
+      "image_masks_loss": 0.0047,
+      "llm_learning_rate": 6.14657007080447e-07,
+      "loss": 0.6659,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 6.14657007080447e-05,
+      "step": 7601
+    },
+    {
+      "epoch": 0.78,
+      "gumbel_tau": 0.1664,
+      "image_masks_loss": 0.0078,
+      "llm_learning_rate": 6.141122875681791e-07,
+      "loss": 0.716,
+      "output_text_mask_loss": 0.0081,
+      "predictor_learning_rate": 6.14112287568179e-05,
+      "step": 7602
+    },
+    {
+      "epoch": 0.78,
+      "gumbel_tau": 0.1664,
+      "image_masks_loss": 0.0084,
+      "llm_learning_rate": 6.135677757340668e-07,
+      "loss": 0.6911,
+      "output_text_mask_loss": 0.0755,
+      "predictor_learning_rate": 6.135677757340669e-05,
+      "step": 7603
+    },
+    {
+      "epoch": 0.78,
+      "gumbel_tau": 0.1663,
+      "image_masks_loss": 0.0042,
+      "llm_learning_rate": 6.130234716380731e-07,
+      "loss": 0.6967,
+      "output_text_mask_loss": 0.0248,
+      "predictor_learning_rate": 6.130234716380731e-05,
+      "step": 7604
+    },
+    {
+      "epoch": 0.78,
+      "gumbel_tau": 0.1663,
+      "image_masks_loss": 0.0115,
+      "llm_learning_rate": 6.124793753401378e-07,
+      "loss": 0.7588,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 6.124793753401378e-05,
+      "step": 7605
+    },
+    {
+      "epoch": 0.78,
+      "gumbel_tau": 0.1663,
+      "image_masks_loss": 0.0069,
+      "llm_learning_rate": 6.119354869001778e-07,
+      "loss": 0.763,
+      "output_text_mask_loss": 0.0347,
+      "predictor_learning_rate": 6.119354869001779e-05,
+      "step": 7606
+    },
+    {
+      "epoch": 0.78,
+      "gumbel_tau": 0.1662,
+      "image_masks_loss": 0.0088,
+      "llm_learning_rate": 6.113918063780875e-07,
+      "loss": 0.7147,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 6.113918063780873e-05,
+      "step": 7607
+    },
+    {
+      "epoch": 0.78,
+      "gumbel_tau": 0.1662,
+      "image_masks_loss": 0.0109,
+      "llm_learning_rate": 6.108483338337382e-07,
+      "loss": 0.7963,
+      "output_text_mask_loss": 0.002,
+      "predictor_learning_rate": 6.108483338337381e-05,
+      "step": 7608
+    },
+    {
+      "epoch": 0.78,
+      "gumbel_tau": 0.1661,
+      "image_masks_loss": 0.0059,
+      "llm_learning_rate": 6.103050693269777e-07,
+      "loss": 0.7781,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 6.103050693269777e-05,
+      "step": 7609
+    },
+    {
+      "epoch": 0.78,
+      "gumbel_tau": 0.1661,
+      "image_masks_loss": 0.0081,
+      "llm_learning_rate": 6.097620129176326e-07,
+      "loss": 0.7355,
+      "output_text_mask_loss": 0.0347,
+      "predictor_learning_rate": 6.097620129176326e-05,
+      "step": 7610
+    },
+    {
+      "epoch": 0.78,
+      "gumbel_tau": 0.1661,
+      "image_masks_loss": 0.0088,
+      "llm_learning_rate": 6.092191646655052e-07,
+      "loss": 0.7319,
+      "output_text_mask_loss": 0.0124,
+      "predictor_learning_rate": 6.092191646655051e-05,
+      "step": 7611
+    },
+    {
+      "epoch": 0.78,
+      "gumbel_tau": 0.166,
+      "image_masks_loss": 0.0149,
+      "llm_learning_rate": 6.086765246303747e-07,
+      "loss": 0.7162,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 6.086765246303747e-05,
+      "step": 7612
+    },
+    {
+      "epoch": 0.78,
+      "gumbel_tau": 0.166,
+      "image_masks_loss": 0.0083,
+      "llm_learning_rate": 6.081340928719981e-07,
+      "loss": 0.6634,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 6.08134092871998e-05,
+      "step": 7613
+    },
+    {
+      "epoch": 0.78,
+      "gumbel_tau": 0.166,
+      "image_masks_loss": 0.0013,
+      "llm_learning_rate": 6.075918694501092e-07,
+      "loss": 0.8305,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 6.075918694501092e-05,
+      "step": 7614
+    },
+    {
+      "epoch": 0.78,
+      "gumbel_tau": 0.1659,
+      "image_masks_loss": 0.0065,
+      "llm_learning_rate": 6.070498544244191e-07,
+      "loss": 0.7419,
+      "output_text_mask_loss": 0.0196,
+      "predictor_learning_rate": 6.070498544244191e-05,
+      "step": 7615
+    },
+    {
+      "epoch": 0.78,
+      "gumbel_tau": 0.1659,
+      "image_masks_loss": 0.0077,
+      "llm_learning_rate": 6.065080478546145e-07,
+      "loss": 0.7149,
+      "output_text_mask_loss": 0.0056,
+      "predictor_learning_rate": 6.065080478546145e-05,
+      "step": 7616
+    },
+    {
+      "epoch": 0.78,
+      "gumbel_tau": 0.1658,
+      "image_masks_loss": 0.0039,
+      "llm_learning_rate": 6.059664498003617e-07,
+      "loss": 0.64,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 6.059664498003617e-05,
+      "step": 7617
+    },
+    {
+      "epoch": 0.78,
+      "gumbel_tau": 0.1658,
+      "image_masks_loss": 0.0037,
+      "llm_learning_rate": 6.054250603213027e-07,
+      "loss": 0.724,
+      "output_text_mask_loss": 0.029,
+      "predictor_learning_rate": 6.0542506032130265e-05,
+      "step": 7618
+    },
+    {
+      "epoch": 0.78,
+      "gumbel_tau": 0.1658,
+      "image_masks_loss": 0.0089,
+      "llm_learning_rate": 6.048838794770556e-07,
+      "loss": 0.7862,
+      "output_text_mask_loss": 0.0071,
+      "predictor_learning_rate": 6.048838794770556e-05,
+      "step": 7619
+    },
+    {
+      "epoch": 0.78,
+      "gumbel_tau": 0.1657,
+      "image_masks_loss": 0.0142,
+      "llm_learning_rate": 6.043429073272173e-07,
+      "loss": 0.6572,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 6.043429073272172e-05,
+      "step": 7620
+    },
+    {
+      "epoch": 0.78,
+      "gumbel_tau": 0.1657,
+      "image_masks_loss": 0.0116,
+      "llm_learning_rate": 6.03802143931361e-07,
+      "loss": 0.6716,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 6.03802143931361e-05,
+      "step": 7621
+    },
+    {
+      "epoch": 0.78,
+      "gumbel_tau": 0.1656,
+      "image_masks_loss": 0.0084,
+      "llm_learning_rate": 6.032615893490349e-07,
+      "loss": 0.6914,
+      "output_text_mask_loss": 0.0032,
+      "predictor_learning_rate": 6.032615893490348e-05,
+      "step": 7622
+    },
+    {
+      "epoch": 0.78,
+      "gumbel_tau": 0.1656,
+      "image_masks_loss": 0.0045,
+      "llm_learning_rate": 6.027212436397678e-07,
+      "loss": 0.7213,
+      "output_text_mask_loss": 0.006,
+      "predictor_learning_rate": 6.0272124363976774e-05,
+      "step": 7623
+    },
+    {
+      "epoch": 0.78,
+      "gumbel_tau": 0.1656,
+      "image_masks_loss": 0.0188,
+      "llm_learning_rate": 6.021811068630637e-07,
+      "loss": 0.7591,
+      "output_text_mask_loss": 0.0032,
+      "predictor_learning_rate": 6.0218110686306365e-05,
+      "step": 7624
+    },
+    {
+      "epoch": 0.78,
+      "gumbel_tau": 0.1655,
+      "image_masks_loss": 0.0093,
+      "llm_learning_rate": 6.016411790784032e-07,
+      "loss": 0.7562,
+      "output_text_mask_loss": 0.0614,
+      "predictor_learning_rate": 6.016411790784032e-05,
+      "step": 7625
+    },
+    {
+      "epoch": 0.78,
+      "gumbel_tau": 0.1655,
+      "image_masks_loss": 0.0088,
+      "llm_learning_rate": 6.011014603452444e-07,
+      "loss": 0.7476,
+      "output_text_mask_loss": 0.0172,
+      "predictor_learning_rate": 6.0110146034524444e-05,
+      "step": 7626
+    },
+    {
+      "epoch": 0.78,
+      "gumbel_tau": 0.1654,
+      "image_masks_loss": 0.0104,
+      "llm_learning_rate": 6.005619507230226e-07,
+      "loss": 0.6999,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 6.0056195072302254e-05,
+      "step": 7627
+    },
+    {
+      "epoch": 0.78,
+      "gumbel_tau": 0.1654,
+      "image_masks_loss": 0.0193,
+      "llm_learning_rate": 6.000226502711495e-07,
+      "loss": 0.8067,
+      "output_text_mask_loss": 0.0162,
+      "predictor_learning_rate": 6.000226502711495e-05,
+      "step": 7628
+    },
+    {
+      "epoch": 0.78,
+      "gumbel_tau": 0.1654,
+      "image_masks_loss": 0.0039,
+      "llm_learning_rate": 5.994835590490136e-07,
+      "loss": 0.7151,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 5.994835590490136e-05,
+      "step": 7629
+    },
+    {
+      "epoch": 0.78,
+      "gumbel_tau": 0.1653,
+      "image_masks_loss": 0.0149,
+      "llm_learning_rate": 5.989446771159821e-07,
+      "loss": 0.7798,
+      "output_text_mask_loss": 0.0488,
+      "predictor_learning_rate": 5.9894467711598203e-05,
+      "step": 7630
+    },
+    {
+      "epoch": 0.78,
+      "gumbel_tau": 0.1653,
+      "image_masks_loss": 0.0061,
+      "llm_learning_rate": 5.984060045313972e-07,
+      "loss": 0.694,
+      "output_text_mask_loss": 0.0324,
+      "predictor_learning_rate": 5.9840600453139714e-05,
+      "step": 7631
+    },
+    {
+      "epoch": 0.78,
+      "gumbel_tau": 0.1652,
+      "image_masks_loss": 0.0062,
+      "llm_learning_rate": 5.978675413545784e-07,
+      "loss": 0.6402,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 5.978675413545784e-05,
+      "step": 7632
+    },
+    {
+      "epoch": 0.78,
+      "gumbel_tau": 0.1652,
+      "image_masks_loss": 0.003,
+      "llm_learning_rate": 5.97329287644823e-07,
+      "loss": 0.6949,
+      "output_text_mask_loss": 0.0065,
+      "predictor_learning_rate": 5.9732928764482304e-05,
+      "step": 7633
+    },
+    {
+      "epoch": 0.78,
+      "gumbel_tau": 0.1652,
+      "image_masks_loss": 0.0065,
+      "llm_learning_rate": 5.967912434614045e-07,
+      "loss": 0.6405,
+      "output_text_mask_loss": 0.0019,
+      "predictor_learning_rate": 5.967912434614045e-05,
+      "step": 7634
+    },
+    {
+      "epoch": 0.78,
+      "gumbel_tau": 0.1651,
+      "image_masks_loss": 0.008,
+      "llm_learning_rate": 5.962534088635738e-07,
+      "loss": 0.6929,
+      "output_text_mask_loss": 0.0078,
+      "predictor_learning_rate": 5.9625340886357376e-05,
+      "step": 7635
+    },
+    {
+      "epoch": 0.78,
+      "gumbel_tau": 0.1651,
+      "image_masks_loss": 0.0051,
+      "llm_learning_rate": 5.957157839105571e-07,
+      "loss": 0.7917,
+      "output_text_mask_loss": 0.0277,
+      "predictor_learning_rate": 5.95715783910557e-05,
+      "step": 7636
+    },
+    {
+      "epoch": 0.78,
+      "gumbel_tau": 0.1651,
+      "image_masks_loss": 0.0047,
+      "llm_learning_rate": 5.951783686615606e-07,
+      "loss": 0.7845,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 5.951783686615606e-05,
+      "step": 7637
+    },
+    {
+      "epoch": 0.78,
+      "gumbel_tau": 0.165,
+      "image_masks_loss": 0.0088,
+      "llm_learning_rate": 5.946411631757651e-07,
+      "loss": 0.6981,
+      "output_text_mask_loss": 0.0088,
+      "predictor_learning_rate": 5.946411631757651e-05,
+      "step": 7638
+    },
+    {
+      "epoch": 0.78,
+      "gumbel_tau": 0.165,
+      "image_masks_loss": 0.0061,
+      "llm_learning_rate": 5.941041675123294e-07,
+      "loss": 0.6792,
+      "output_text_mask_loss": 0.0248,
+      "predictor_learning_rate": 5.941041675123293e-05,
+      "step": 7639
+    },
+    {
+      "epoch": 0.78,
+      "gumbel_tau": 0.1649,
+      "image_masks_loss": 0.0048,
+      "llm_learning_rate": 5.935673817303872e-07,
+      "loss": 0.745,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 5.935673817303872e-05,
+      "step": 7640
+    },
+    {
+      "epoch": 0.78,
+      "gumbel_tau": 0.1649,
+      "image_masks_loss": 0.007,
+      "llm_learning_rate": 5.930308058890505e-07,
+      "loss": 0.7323,
+      "output_text_mask_loss": 0.0345,
+      "predictor_learning_rate": 5.930308058890505e-05,
+      "step": 7641
+    },
+    {
+      "epoch": 0.78,
+      "gumbel_tau": 0.1649,
+      "image_masks_loss": 0.0037,
+      "llm_learning_rate": 5.9249444004741e-07,
+      "loss": 0.7306,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 5.9249444004741e-05,
+      "step": 7642
+    },
+    {
+      "epoch": 0.78,
+      "gumbel_tau": 0.1648,
+      "image_masks_loss": 0.0038,
+      "llm_learning_rate": 5.919582842645305e-07,
+      "loss": 0.7655,
+      "output_text_mask_loss": 0.017,
+      "predictor_learning_rate": 5.9195828426453045e-05,
+      "step": 7643
+    },
+    {
+      "epoch": 0.78,
+      "gumbel_tau": 0.1648,
+      "image_masks_loss": 0.0149,
+      "llm_learning_rate": 5.914223385994547e-07,
+      "loss": 0.6732,
+      "output_text_mask_loss": 0.0111,
+      "predictor_learning_rate": 5.9142233859945464e-05,
+      "step": 7644
+    },
+    {
+      "epoch": 0.78,
+      "gumbel_tau": 0.1647,
+      "image_masks_loss": 0.007,
+      "llm_learning_rate": 5.908866031112023e-07,
+      "loss": 0.7419,
+      "output_text_mask_loss": 0.0229,
+      "predictor_learning_rate": 5.908866031112023e-05,
+      "step": 7645
+    },
+    {
+      "epoch": 0.78,
+      "gumbel_tau": 0.1647,
+      "image_masks_loss": 0.0065,
+      "llm_learning_rate": 5.903510778587695e-07,
+      "loss": 0.7029,
+      "output_text_mask_loss": 0.1175,
+      "predictor_learning_rate": 5.903510778587695e-05,
+      "step": 7646
+    },
+    {
+      "epoch": 0.78,
+      "gumbel_tau": 0.1647,
+      "image_masks_loss": 0.0084,
+      "llm_learning_rate": 5.898157629011295e-07,
+      "loss": 0.7657,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 5.8981576290112944e-05,
+      "step": 7647
+    },
+    {
+      "epoch": 0.78,
+      "gumbel_tau": 0.1646,
+      "image_masks_loss": 0.0029,
+      "llm_learning_rate": 5.892806582972319e-07,
+      "loss": 0.7476,
+      "output_text_mask_loss": 0.0153,
+      "predictor_learning_rate": 5.8928065829723185e-05,
+      "step": 7648
+    },
+    {
+      "epoch": 0.78,
+      "gumbel_tau": 0.1646,
+      "image_masks_loss": 0.0175,
+      "llm_learning_rate": 5.887457641060045e-07,
+      "loss": 0.6338,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 5.8874576410600456e-05,
+      "step": 7649
+    },
+    {
+      "epoch": 0.78,
+      "gumbel_tau": 0.1645,
+      "image_masks_loss": 0.0041,
+      "llm_learning_rate": 5.882110803863511e-07,
+      "loss": 0.7598,
+      "output_text_mask_loss": 0.0134,
+      "predictor_learning_rate": 5.8821108038635104e-05,
+      "step": 7650
+    },
+    {
+      "epoch": 0.78,
+      "gumbel_tau": 0.1645,
+      "image_masks_loss": 0.0066,
+      "llm_learning_rate": 5.876766071971518e-07,
+      "loss": 0.7017,
+      "output_text_mask_loss": 0.0181,
+      "predictor_learning_rate": 5.8767660719715175e-05,
+      "step": 7651
+    },
+    {
+      "epoch": 0.78,
+      "gumbel_tau": 0.1645,
+      "image_masks_loss": 0.0021,
+      "llm_learning_rate": 5.87142344597264e-07,
+      "loss": 0.7019,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 5.87142344597264e-05,
+      "step": 7652
+    },
+    {
+      "epoch": 0.78,
+      "gumbel_tau": 0.1644,
+      "image_masks_loss": 0.0112,
+      "llm_learning_rate": 5.866082926455218e-07,
+      "loss": 0.6755,
+      "output_text_mask_loss": 0.0092,
+      "predictor_learning_rate": 5.8660829264552186e-05,
+      "step": 7653
+    },
+    {
+      "epoch": 0.78,
+      "gumbel_tau": 0.1644,
+      "image_masks_loss": 0.0094,
+      "llm_learning_rate": 5.860744514007366e-07,
+      "loss": 0.7392,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 5.8607445140073656e-05,
+      "step": 7654
+    },
+    {
+      "epoch": 0.78,
+      "gumbel_tau": 0.1644,
+      "image_masks_loss": 0.0098,
+      "llm_learning_rate": 5.855408209216951e-07,
+      "loss": 0.7689,
+      "output_text_mask_loss": 0.0286,
+      "predictor_learning_rate": 5.8554082092169504e-05,
+      "step": 7655
+    },
+    {
+      "epoch": 0.78,
+      "gumbel_tau": 0.1643,
+      "image_masks_loss": 0.0099,
+      "llm_learning_rate": 5.850074012671634e-07,
+      "loss": 0.7021,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 5.850074012671633e-05,
+      "step": 7656
+    },
+    {
+      "epoch": 0.78,
+      "gumbel_tau": 0.1643,
+      "image_masks_loss": 0.0043,
+      "llm_learning_rate": 5.844741924958818e-07,
+      "loss": 0.7175,
+      "output_text_mask_loss": 0.001,
+      "predictor_learning_rate": 5.844741924958818e-05,
+      "step": 7657
+    },
+    {
+      "epoch": 0.78,
+      "gumbel_tau": 0.1642,
+      "image_masks_loss": 0.0148,
+      "llm_learning_rate": 5.839411946665699e-07,
+      "loss": 0.668,
+      "output_text_mask_loss": 0.0242,
+      "predictor_learning_rate": 5.839411946665699e-05,
+      "step": 7658
+    },
+    {
+      "epoch": 0.78,
+      "gumbel_tau": 0.1642,
+      "image_masks_loss": 0.0062,
+      "llm_learning_rate": 5.834084078379207e-07,
+      "loss": 0.7789,
+      "output_text_mask_loss": 0.0265,
+      "predictor_learning_rate": 5.834084078379206e-05,
+      "step": 7659
+    },
+    {
+      "epoch": 0.78,
+      "gumbel_tau": 0.1642,
+      "image_masks_loss": 0.0077,
+      "llm_learning_rate": 5.828758320686068e-07,
+      "loss": 0.7003,
+      "output_text_mask_loss": 0.0058,
+      "predictor_learning_rate": 5.828758320686067e-05,
+      "step": 7660
+    },
+    {
+      "epoch": 0.78,
+      "gumbel_tau": 0.1641,
+      "image_masks_loss": 0.0089,
+      "llm_learning_rate": 5.823434674172757e-07,
+      "loss": 0.8453,
+      "output_text_mask_loss": 0.0198,
+      "predictor_learning_rate": 5.823434674172756e-05,
+      "step": 7661
+    },
+    {
+      "epoch": 0.79,
+      "gumbel_tau": 0.1641,
+      "image_masks_loss": 0.0042,
+      "llm_learning_rate": 5.818113139425541e-07,
+      "loss": 0.7873,
+      "output_text_mask_loss": 0.0298,
+      "predictor_learning_rate": 5.8181131394255404e-05,
+      "step": 7662
+    },
+    {
+      "epoch": 0.79,
+      "gumbel_tau": 0.164,
+      "image_masks_loss": 0.0044,
+      "llm_learning_rate": 5.812793717030429e-07,
+      "loss": 0.7328,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 5.812793717030429e-05,
+      "step": 7663
+    },
+    {
+      "epoch": 0.79,
+      "gumbel_tau": 0.164,
+      "image_masks_loss": 0.0069,
+      "llm_learning_rate": 5.807476407573215e-07,
+      "loss": 0.6758,
+      "output_text_mask_loss": 0.0189,
+      "predictor_learning_rate": 5.807476407573214e-05,
+      "step": 7664
+    },
+    {
+      "epoch": 0.79,
+      "gumbel_tau": 0.164,
+      "image_masks_loss": 0.007,
+      "llm_learning_rate": 5.802161211639445e-07,
+      "loss": 0.7984,
+      "output_text_mask_loss": 0.005,
+      "predictor_learning_rate": 5.802161211639445e-05,
+      "step": 7665
+    },
+    {
+      "epoch": 0.79,
+      "gumbel_tau": 0.1639,
+      "image_masks_loss": 0.0051,
+      "llm_learning_rate": 5.796848129814448e-07,
+      "loss": 0.6605,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 5.796848129814447e-05,
+      "step": 7666
+    },
+    {
+      "epoch": 0.79,
+      "gumbel_tau": 0.1639,
+      "image_masks_loss": 0.0063,
+      "llm_learning_rate": 5.791537162683305e-07,
+      "loss": 0.67,
+      "output_text_mask_loss": 0.0126,
+      "predictor_learning_rate": 5.791537162683305e-05,
+      "step": 7667
+    },
+    {
+      "epoch": 0.79,
+      "gumbel_tau": 0.1639,
+      "image_masks_loss": 0.0021,
+      "llm_learning_rate": 5.786228310830869e-07,
+      "loss": 0.7434,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 5.7862283108308683e-05,
+      "step": 7668
+    },
+    {
+      "epoch": 0.79,
+      "gumbel_tau": 0.1638,
+      "image_masks_loss": 0.0044,
+      "llm_learning_rate": 5.780921574841774e-07,
+      "loss": 0.7401,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 5.780921574841774e-05,
+      "step": 7669
+    },
+    {
+      "epoch": 0.79,
+      "gumbel_tau": 0.1638,
+      "image_masks_loss": 0.0049,
+      "llm_learning_rate": 5.775616955300406e-07,
+      "loss": 0.7698,
+      "output_text_mask_loss": 0.0156,
+      "predictor_learning_rate": 5.775616955300406e-05,
+      "step": 7670
+    },
+    {
+      "epoch": 0.79,
+      "gumbel_tau": 0.1637,
+      "image_masks_loss": 0.0124,
+      "llm_learning_rate": 5.77031445279092e-07,
+      "loss": 0.7179,
+      "output_text_mask_loss": 0.0076,
+      "predictor_learning_rate": 5.77031445279092e-05,
+      "step": 7671
+    },
+    {
+      "epoch": 0.79,
+      "gumbel_tau": 0.1637,
+      "image_masks_loss": 0.0159,
+      "llm_learning_rate": 5.765014067897237e-07,
+      "loss": 0.7068,
+      "output_text_mask_loss": 0.0095,
+      "predictor_learning_rate": 5.765014067897237e-05,
+      "step": 7672
+    },
+    {
+      "epoch": 0.79,
+      "gumbel_tau": 0.1637,
+      "image_masks_loss": 0.008,
+      "llm_learning_rate": 5.75971580120305e-07,
+      "loss": 0.6699,
+      "output_text_mask_loss": 0.0223,
+      "predictor_learning_rate": 5.75971580120305e-05,
+      "step": 7673
+    },
+    {
+      "epoch": 0.79,
+      "gumbel_tau": 0.1636,
+      "image_masks_loss": 0.0057,
+      "llm_learning_rate": 5.754419653291815e-07,
+      "loss": 0.6126,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 5.7544196532918144e-05,
+      "step": 7674
+    },
+    {
+      "epoch": 0.79,
+      "gumbel_tau": 0.1636,
+      "image_masks_loss": 0.0093,
+      "llm_learning_rate": 5.749125624746749e-07,
+      "loss": 0.7506,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 5.749125624746748e-05,
+      "step": 7675
+    },
+    {
+      "epoch": 0.79,
+      "gumbel_tau": 0.1635,
+      "image_masks_loss": 0.0067,
+      "llm_learning_rate": 5.743833716150862e-07,
+      "loss": 0.7728,
+      "output_text_mask_loss": 0.0111,
+      "predictor_learning_rate": 5.743833716150862e-05,
+      "step": 7676
+    },
+    {
+      "epoch": 0.79,
+      "gumbel_tau": 0.1635,
+      "image_masks_loss": 0.0093,
+      "llm_learning_rate": 5.738543928086892e-07,
+      "loss": 0.6741,
+      "output_text_mask_loss": 0.0469,
+      "predictor_learning_rate": 5.738543928086892e-05,
+      "step": 7677
+    },
+    {
+      "epoch": 0.79,
+      "gumbel_tau": 0.1635,
+      "image_masks_loss": 0.0115,
+      "llm_learning_rate": 5.733256261137366e-07,
+      "loss": 0.709,
+      "output_text_mask_loss": 0.0423,
+      "predictor_learning_rate": 5.733256261137365e-05,
+      "step": 7678
+    },
+    {
+      "epoch": 0.79,
+      "gumbel_tau": 0.1634,
+      "image_masks_loss": 0.009,
+      "llm_learning_rate": 5.727970715884576e-07,
+      "loss": 0.6358,
+      "output_text_mask_loss": 0.0164,
+      "predictor_learning_rate": 5.7279707158845754e-05,
+      "step": 7679
+    },
+    {
+      "epoch": 0.79,
+      "gumbel_tau": 0.1634,
+      "image_masks_loss": 0.0068,
+      "llm_learning_rate": 5.722687292910572e-07,
+      "loss": 0.752,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 5.722687292910572e-05,
+      "step": 7680
+    },
+    {
+      "epoch": 0.79,
+      "gumbel_tau": 0.1633,
+      "image_masks_loss": 0.0129,
+      "llm_learning_rate": 5.717405992797187e-07,
+      "loss": 0.7032,
+      "output_text_mask_loss": 0.0076,
+      "predictor_learning_rate": 5.7174059927971865e-05,
+      "step": 7681
+    },
+    {
+      "epoch": 0.79,
+      "gumbel_tau": 0.1633,
+      "image_masks_loss": 0.0079,
+      "llm_learning_rate": 5.712126816126004e-07,
+      "loss": 0.6633,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 5.7121268161260034e-05,
+      "step": 7682
+    },
+    {
+      "epoch": 0.79,
+      "gumbel_tau": 0.1633,
+      "image_masks_loss": 0.0074,
+      "llm_learning_rate": 5.706849763478381e-07,
+      "loss": 0.7282,
+      "output_text_mask_loss": 0.0015,
+      "predictor_learning_rate": 5.706849763478381e-05,
+      "step": 7683
+    },
+    {
+      "epoch": 0.79,
+      "gumbel_tau": 0.1632,
+      "image_masks_loss": 0.0071,
+      "llm_learning_rate": 5.701574835435434e-07,
+      "loss": 0.7733,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 5.701574835435433e-05,
+      "step": 7684
+    },
+    {
+      "epoch": 0.79,
+      "gumbel_tau": 0.1632,
+      "image_masks_loss": 0.005,
+      "llm_learning_rate": 5.696302032578052e-07,
+      "loss": 0.7009,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 5.696302032578052e-05,
+      "step": 7685
+    },
+    {
+      "epoch": 0.79,
+      "gumbel_tau": 0.1632,
+      "image_masks_loss": 0.0087,
+      "llm_learning_rate": 5.691031355486887e-07,
+      "loss": 0.6225,
+      "output_text_mask_loss": 0.0116,
+      "predictor_learning_rate": 5.6910313554868866e-05,
+      "step": 7686
+    },
+    {
+      "epoch": 0.79,
+      "gumbel_tau": 0.1631,
+      "image_masks_loss": 0.006,
+      "llm_learning_rate": 5.685762804742356e-07,
+      "loss": 0.7355,
+      "output_text_mask_loss": 0.058,
+      "predictor_learning_rate": 5.685762804742356e-05,
+      "step": 7687
+    },
+    {
+      "epoch": 0.79,
+      "gumbel_tau": 0.1631,
+      "image_masks_loss": 0.0045,
+      "llm_learning_rate": 5.680496380924644e-07,
+      "loss": 0.7929,
+      "output_text_mask_loss": 0.0553,
+      "predictor_learning_rate": 5.6804963809246436e-05,
+      "step": 7688
+    },
+    {
+      "epoch": 0.79,
+      "gumbel_tau": 0.163,
+      "image_masks_loss": 0.0122,
+      "llm_learning_rate": 5.675232084613705e-07,
+      "loss": 0.7953,
+      "output_text_mask_loss": 0.0085,
+      "predictor_learning_rate": 5.675232084613705e-05,
+      "step": 7689
+    },
+    {
+      "epoch": 0.79,
+      "gumbel_tau": 0.163,
+      "image_masks_loss": 0.0154,
+      "llm_learning_rate": 5.669969916389257e-07,
+      "loss": 0.7183,
+      "output_text_mask_loss": 0.0075,
+      "predictor_learning_rate": 5.669969916389256e-05,
+      "step": 7690
+    },
+    {
+      "epoch": 0.79,
+      "gumbel_tau": 0.163,
+      "image_masks_loss": 0.0125,
+      "llm_learning_rate": 5.664709876830773e-07,
+      "loss": 0.8218,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 5.664709876830773e-05,
+      "step": 7691
+    },
+    {
+      "epoch": 0.79,
+      "gumbel_tau": 0.1629,
+      "image_masks_loss": 0.0107,
+      "llm_learning_rate": 5.65945196651751e-07,
+      "loss": 0.6869,
+      "output_text_mask_loss": 0.0412,
+      "predictor_learning_rate": 5.6594519665175096e-05,
+      "step": 7692
+    },
+    {
+      "epoch": 0.79,
+      "gumbel_tau": 0.1629,
+      "image_masks_loss": 0.0174,
+      "llm_learning_rate": 5.654196186028474e-07,
+      "loss": 0.6507,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 5.654196186028473e-05,
+      "step": 7693
+    },
+    {
+      "epoch": 0.79,
+      "gumbel_tau": 0.1628,
+      "image_masks_loss": 0.0114,
+      "llm_learning_rate": 5.648942535942445e-07,
+      "loss": 0.7238,
+      "output_text_mask_loss": 0.0587,
+      "predictor_learning_rate": 5.648942535942445e-05,
+      "step": 7694
+    },
+    {
+      "epoch": 0.79,
+      "gumbel_tau": 0.1628,
+      "image_masks_loss": 0.0174,
+      "llm_learning_rate": 5.643691016837966e-07,
+      "loss": 0.7075,
+      "output_text_mask_loss": 0.0538,
+      "predictor_learning_rate": 5.643691016837965e-05,
+      "step": 7695
+    },
+    {
+      "epoch": 0.79,
+      "gumbel_tau": 0.1628,
+      "image_masks_loss": 0.0041,
+      "llm_learning_rate": 5.638441629293348e-07,
+      "loss": 0.7447,
+      "output_text_mask_loss": 0.0374,
+      "predictor_learning_rate": 5.6384416292933476e-05,
+      "step": 7696
+    },
+    {
+      "epoch": 0.79,
+      "gumbel_tau": 0.1627,
+      "image_masks_loss": 0.0086,
+      "llm_learning_rate": 5.633194373886663e-07,
+      "loss": 0.7532,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 5.633194373886663e-05,
+      "step": 7697
+    },
+    {
+      "epoch": 0.79,
+      "gumbel_tau": 0.1627,
+      "image_masks_loss": 0.0184,
+      "llm_learning_rate": 5.627949251195753e-07,
+      "loss": 0.714,
+      "output_text_mask_loss": 0.02,
+      "predictor_learning_rate": 5.627949251195752e-05,
+      "step": 7698
+    },
+    {
+      "epoch": 0.79,
+      "gumbel_tau": 0.1627,
+      "image_masks_loss": 0.0067,
+      "llm_learning_rate": 5.622706261798217e-07,
+      "loss": 0.7366,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 5.622706261798216e-05,
+      "step": 7699
+    },
+    {
+      "epoch": 0.79,
+      "gumbel_tau": 0.1626,
+      "image_masks_loss": 0.0078,
+      "llm_learning_rate": 5.617465406271425e-07,
+      "loss": 0.7806,
+      "output_text_mask_loss": 0.0416,
+      "predictor_learning_rate": 5.617465406271424e-05,
+      "step": 7700
+    },
+    {
+      "epoch": 0.79,
+      "gumbel_tau": 0.1626,
+      "image_masks_loss": 0.0087,
+      "llm_learning_rate": 5.61222668519252e-07,
+      "loss": 0.7174,
+      "output_text_mask_loss": 0.0042,
+      "predictor_learning_rate": 5.61222668519252e-05,
+      "step": 7701
+    },
+    {
+      "epoch": 0.79,
+      "gumbel_tau": 0.1625,
+      "image_masks_loss": 0.008,
+      "llm_learning_rate": 5.606990099138401e-07,
+      "loss": 0.7315,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 5.606990099138401e-05,
+      "step": 7702
+    },
+    {
+      "epoch": 0.79,
+      "gumbel_tau": 0.1625,
+      "image_masks_loss": 0.0083,
+      "llm_learning_rate": 5.601755648685725e-07,
+      "loss": 0.8205,
+      "output_text_mask_loss": 0.0488,
+      "predictor_learning_rate": 5.601755648685725e-05,
+      "step": 7703
+    },
+    {
+      "epoch": 0.79,
+      "gumbel_tau": 0.1625,
+      "image_masks_loss": 0.0135,
+      "llm_learning_rate": 5.596523334410927e-07,
+      "loss": 0.7535,
+      "output_text_mask_loss": 0.0227,
+      "predictor_learning_rate": 5.596523334410927e-05,
+      "step": 7704
+    },
+    {
+      "epoch": 0.79,
+      "gumbel_tau": 0.1624,
+      "image_masks_loss": 0.0054,
+      "llm_learning_rate": 5.591293156890198e-07,
+      "loss": 0.6465,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 5.5912931568901984e-05,
+      "step": 7705
+    },
+    {
+      "epoch": 0.79,
+      "gumbel_tau": 0.1624,
+      "image_masks_loss": 0.009,
+      "llm_learning_rate": 5.586065116699499e-07,
+      "loss": 0.6916,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 5.586065116699498e-05,
+      "step": 7706
+    },
+    {
+      "epoch": 0.79,
+      "gumbel_tau": 0.1623,
+      "image_masks_loss": 0.0171,
+      "llm_learning_rate": 5.580839214414543e-07,
+      "loss": 0.6612,
+      "output_text_mask_loss": 0.0032,
+      "predictor_learning_rate": 5.580839214414543e-05,
+      "step": 7707
+    },
+    {
+      "epoch": 0.79,
+      "gumbel_tau": 0.1623,
+      "image_masks_loss": 0.0068,
+      "llm_learning_rate": 5.575615450610836e-07,
+      "loss": 0.6819,
+      "output_text_mask_loss": 0.0101,
+      "predictor_learning_rate": 5.575615450610835e-05,
+      "step": 7708
+    },
+    {
+      "epoch": 0.79,
+      "gumbel_tau": 0.1623,
+      "image_masks_loss": 0.009,
+      "llm_learning_rate": 5.570393825863621e-07,
+      "loss": 0.709,
+      "output_text_mask_loss": 0.0238,
+      "predictor_learning_rate": 5.57039382586362e-05,
+      "step": 7709
+    },
+    {
+      "epoch": 0.79,
+      "gumbel_tau": 0.1622,
+      "image_masks_loss": 0.0059,
+      "llm_learning_rate": 5.565174340747914e-07,
+      "loss": 0.72,
+      "output_text_mask_loss": 0.0353,
+      "predictor_learning_rate": 5.565174340747914e-05,
+      "step": 7710
+    },
+    {
+      "epoch": 0.79,
+      "gumbel_tau": 0.1622,
+      "image_masks_loss": 0.0052,
+      "llm_learning_rate": 5.5599569958385e-07,
+      "loss": 0.7291,
+      "output_text_mask_loss": 0.0107,
+      "predictor_learning_rate": 5.5599569958385e-05,
+      "step": 7711
+    },
+    {
+      "epoch": 0.79,
+      "gumbel_tau": 0.1622,
+      "image_masks_loss": 0.0174,
+      "llm_learning_rate": 5.554741791709922e-07,
+      "loss": 0.692,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 5.554741791709922e-05,
+      "step": 7712
+    },
+    {
+      "epoch": 0.79,
+      "gumbel_tau": 0.1621,
+      "image_masks_loss": 0.0026,
+      "llm_learning_rate": 5.54952872893649e-07,
+      "loss": 0.7361,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 5.54952872893649e-05,
+      "step": 7713
+    },
+    {
+      "epoch": 0.79,
+      "gumbel_tau": 0.1621,
+      "image_masks_loss": 0.0032,
+      "llm_learning_rate": 5.544317808092281e-07,
+      "loss": 0.6632,
+      "output_text_mask_loss": 0.0381,
+      "predictor_learning_rate": 5.544317808092281e-05,
+      "step": 7714
+    },
+    {
+      "epoch": 0.79,
+      "gumbel_tau": 0.162,
+      "image_masks_loss": 0.0017,
+      "llm_learning_rate": 5.539109029751131e-07,
+      "loss": 0.7516,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 5.53910902975113e-05,
+      "step": 7715
+    },
+    {
+      "epoch": 0.79,
+      "gumbel_tau": 0.162,
+      "image_masks_loss": 0.0114,
+      "llm_learning_rate": 5.533902394486637e-07,
+      "loss": 0.7182,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 5.533902394486637e-05,
+      "step": 7716
+    },
+    {
+      "epoch": 0.79,
+      "gumbel_tau": 0.162,
+      "image_masks_loss": 0.0075,
+      "llm_learning_rate": 5.528697902872174e-07,
+      "loss": 0.6993,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 5.528697902872173e-05,
+      "step": 7717
+    },
+    {
+      "epoch": 0.79,
+      "gumbel_tau": 0.1619,
+      "image_masks_loss": 0.0237,
+      "llm_learning_rate": 5.523495555480868e-07,
+      "loss": 0.7764,
+      "output_text_mask_loss": 0.0109,
+      "predictor_learning_rate": 5.5234955554808673e-05,
+      "step": 7718
+    },
+    {
+      "epoch": 0.79,
+      "gumbel_tau": 0.1619,
+      "image_masks_loss": 0.003,
+      "llm_learning_rate": 5.518295352885614e-07,
+      "loss": 0.7847,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 5.5182953528856135e-05,
+      "step": 7719
+    },
+    {
+      "epoch": 0.79,
+      "gumbel_tau": 0.1619,
+      "image_masks_loss": 0.0052,
+      "llm_learning_rate": 5.513097295659062e-07,
+      "loss": 0.7136,
+      "output_text_mask_loss": 0.0141,
+      "predictor_learning_rate": 5.5130972956590616e-05,
+      "step": 7720
+    },
+    {
+      "epoch": 0.79,
+      "gumbel_tau": 0.1618,
+      "image_masks_loss": 0.0109,
+      "llm_learning_rate": 5.507901384373648e-07,
+      "loss": 0.7476,
+      "output_text_mask_loss": 0.0248,
+      "predictor_learning_rate": 5.507901384373648e-05,
+      "step": 7721
+    },
+    {
+      "epoch": 0.79,
+      "gumbel_tau": 0.1618,
+      "image_masks_loss": 0.005,
+      "llm_learning_rate": 5.502707619601552e-07,
+      "loss": 0.7743,
+      "output_text_mask_loss": 0.0023,
+      "predictor_learning_rate": 5.502707619601552e-05,
+      "step": 7722
+    },
+    {
+      "epoch": 0.79,
+      "gumbel_tau": 0.1617,
+      "image_masks_loss": 0.007,
+      "llm_learning_rate": 5.497516001914721e-07,
+      "loss": 0.7256,
+      "output_text_mask_loss": 0.0103,
+      "predictor_learning_rate": 5.4975160019147206e-05,
+      "step": 7723
+    },
+    {
+      "epoch": 0.79,
+      "gumbel_tau": 0.1617,
+      "image_masks_loss": 0.0094,
+      "llm_learning_rate": 5.492326531884867e-07,
+      "loss": 0.8054,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 5.492326531884867e-05,
+      "step": 7724
+    },
+    {
+      "epoch": 0.79,
+      "gumbel_tau": 0.1617,
+      "image_masks_loss": 0.0102,
+      "llm_learning_rate": 5.487139210083467e-07,
+      "loss": 0.7883,
+      "output_text_mask_loss": 0.0322,
+      "predictor_learning_rate": 5.487139210083467e-05,
+      "step": 7725
+    },
+    {
+      "epoch": 0.79,
+      "gumbel_tau": 0.1616,
+      "image_masks_loss": 0.003,
+      "llm_learning_rate": 5.48195403708176e-07,
+      "loss": 0.7242,
+      "output_text_mask_loss": 0.001,
+      "predictor_learning_rate": 5.48195403708176e-05,
+      "step": 7726
+    },
+    {
+      "epoch": 0.79,
+      "gumbel_tau": 0.1616,
+      "image_masks_loss": 0.0066,
+      "llm_learning_rate": 5.476771013450744e-07,
+      "loss": 0.7955,
+      "output_text_mask_loss": 0.0221,
+      "predictor_learning_rate": 5.476771013450743e-05,
+      "step": 7727
+    },
+    {
+      "epoch": 0.79,
+      "gumbel_tau": 0.1615,
+      "image_masks_loss": 0.0224,
+      "llm_learning_rate": 5.471590139761196e-07,
+      "loss": 0.7952,
+      "output_text_mask_loss": 0.0032,
+      "predictor_learning_rate": 5.4715901397611955e-05,
+      "step": 7728
+    },
+    {
+      "epoch": 0.79,
+      "gumbel_tau": 0.1615,
+      "image_masks_loss": 0.0154,
+      "llm_learning_rate": 5.466411416583639e-07,
+      "loss": 0.7156,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 5.466411416583639e-05,
+      "step": 7729
+    },
+    {
+      "epoch": 0.79,
+      "gumbel_tau": 0.1615,
+      "image_masks_loss": 0.0128,
+      "llm_learning_rate": 5.461234844488366e-07,
+      "loss": 0.7891,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 5.461234844488366e-05,
+      "step": 7730
+    },
+    {
+      "epoch": 0.79,
+      "gumbel_tau": 0.1614,
+      "image_masks_loss": 0.0166,
+      "llm_learning_rate": 5.456060424045437e-07,
+      "loss": 0.718,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 5.456060424045436e-05,
+      "step": 7731
+    },
+    {
+      "epoch": 0.79,
+      "gumbel_tau": 0.1614,
+      "image_masks_loss": 0.0169,
+      "llm_learning_rate": 5.450888155824663e-07,
+      "loss": 0.7374,
+      "output_text_mask_loss": 0.0322,
+      "predictor_learning_rate": 5.450888155824663e-05,
+      "step": 7732
+    },
+    {
+      "epoch": 0.79,
+      "gumbel_tau": 0.1614,
+      "image_masks_loss": 0.0204,
+      "llm_learning_rate": 5.445718040395628e-07,
+      "loss": 0.8058,
+      "output_text_mask_loss": 0.0086,
+      "predictor_learning_rate": 5.445718040395628e-05,
+      "step": 7733
+    },
+    {
+      "epoch": 0.79,
+      "gumbel_tau": 0.1613,
+      "image_masks_loss": 0.0221,
+      "llm_learning_rate": 5.44055007832768e-07,
+      "loss": 0.752,
+      "output_text_mask_loss": 0.0071,
+      "predictor_learning_rate": 5.4405500783276804e-05,
+      "step": 7734
+    },
+    {
+      "epoch": 0.79,
+      "gumbel_tau": 0.1613,
+      "image_masks_loss": 0.0233,
+      "llm_learning_rate": 5.435384270189922e-07,
+      "loss": 0.7236,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 5.435384270189922e-05,
+      "step": 7735
+    },
+    {
+      "epoch": 0.79,
+      "gumbel_tau": 0.1612,
+      "image_masks_loss": 0.0072,
+      "llm_learning_rate": 5.430220616551229e-07,
+      "loss": 0.6566,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 5.430220616551229e-05,
+      "step": 7736
+    },
+    {
+      "epoch": 0.79,
+      "gumbel_tau": 0.1612,
+      "image_masks_loss": 0.0194,
+      "llm_learning_rate": 5.425059117980234e-07,
+      "loss": 0.679,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 5.425059117980233e-05,
+      "step": 7737
+    },
+    {
+      "epoch": 0.79,
+      "gumbel_tau": 0.1612,
+      "image_masks_loss": 0.0162,
+      "llm_learning_rate": 5.419899775045325e-07,
+      "loss": 0.7681,
+      "output_text_mask_loss": 0.0298,
+      "predictor_learning_rate": 5.4198997750453254e-05,
+      "step": 7738
+    },
+    {
+      "epoch": 0.79,
+      "gumbel_tau": 0.1611,
+      "image_masks_loss": 0.0135,
+      "llm_learning_rate": 5.414742588314661e-07,
+      "loss": 0.702,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 5.4147425883146615e-05,
+      "step": 7739
+    },
+    {
+      "epoch": 0.79,
+      "gumbel_tau": 0.1611,
+      "image_masks_loss": 0.0265,
+      "llm_learning_rate": 5.409587558356175e-07,
+      "loss": 0.7104,
+      "output_text_mask_loss": 0.0072,
+      "predictor_learning_rate": 5.4095875583561746e-05,
+      "step": 7740
+    },
+    {
+      "epoch": 0.79,
+      "gumbel_tau": 0.1611,
+      "image_masks_loss": 0.0094,
+      "llm_learning_rate": 5.404434685737545e-07,
+      "loss": 0.7487,
+      "output_text_mask_loss": 0.0296,
+      "predictor_learning_rate": 5.404434685737544e-05,
+      "step": 7741
+    },
+    {
+      "epoch": 0.79,
+      "gumbel_tau": 0.161,
+      "image_masks_loss": 0.0135,
+      "llm_learning_rate": 5.399283971026212e-07,
+      "loss": 0.7379,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 5.399283971026211e-05,
+      "step": 7742
+    },
+    {
+      "epoch": 0.79,
+      "gumbel_tau": 0.161,
+      "image_masks_loss": 0.0158,
+      "llm_learning_rate": 5.394135414789387e-07,
+      "loss": 0.7151,
+      "output_text_mask_loss": 0.0124,
+      "predictor_learning_rate": 5.394135414789386e-05,
+      "step": 7743
+    },
+    {
+      "epoch": 0.79,
+      "gumbel_tau": 0.1609,
+      "image_masks_loss": 0.0082,
+      "llm_learning_rate": 5.38898901759404e-07,
+      "loss": 0.7485,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 5.38898901759404e-05,
+      "step": 7744
+    },
+    {
+      "epoch": 0.79,
+      "gumbel_tau": 0.1609,
+      "image_masks_loss": 0.0125,
+      "llm_learning_rate": 5.383844780006906e-07,
+      "loss": 0.7072,
+      "output_text_mask_loss": 0.0124,
+      "predictor_learning_rate": 5.3838447800069054e-05,
+      "step": 7745
+    },
+    {
+      "epoch": 0.79,
+      "gumbel_tau": 0.1609,
+      "image_masks_loss": 0.0034,
+      "llm_learning_rate": 5.37870270259447e-07,
+      "loss": 0.6923,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 5.3787027025944706e-05,
+      "step": 7746
+    },
+    {
+      "epoch": 0.79,
+      "gumbel_tau": 0.1608,
+      "image_masks_loss": 0.0044,
+      "llm_learning_rate": 5.373562785923006e-07,
+      "loss": 0.7247,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 5.373562785923006e-05,
+      "step": 7747
+    },
+    {
+      "epoch": 0.79,
+      "gumbel_tau": 0.1608,
+      "image_masks_loss": 0.013,
+      "llm_learning_rate": 5.368425030558522e-07,
+      "loss": 0.6317,
+      "output_text_mask_loss": 0.0296,
+      "predictor_learning_rate": 5.3684250305585215e-05,
+      "step": 7748
+    },
+    {
+      "epoch": 0.79,
+      "gumbel_tau": 0.1607,
+      "image_masks_loss": 0.0125,
+      "llm_learning_rate": 5.363289437066801e-07,
+      "loss": 0.6962,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 5.363289437066801e-05,
+      "step": 7749
+    },
+    {
+      "epoch": 0.79,
+      "gumbel_tau": 0.1607,
+      "image_masks_loss": 0.0188,
+      "llm_learning_rate": 5.358156006013396e-07,
+      "loss": 0.7715,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 5.358156006013395e-05,
+      "step": 7750
+    },
+    {
+      "epoch": 0.79,
+      "gumbel_tau": 0.1607,
+      "image_masks_loss": 0.0261,
+      "llm_learning_rate": 5.353024737963592e-07,
+      "loss": 0.6942,
+      "output_text_mask_loss": 0.0027,
+      "predictor_learning_rate": 5.353024737963591e-05,
+      "step": 7751
+    },
+    {
+      "epoch": 0.79,
+      "gumbel_tau": 0.1606,
+      "image_masks_loss": 0.038,
+      "llm_learning_rate": 5.347895633482462e-07,
+      "loss": 0.6947,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 5.347895633482461e-05,
+      "step": 7752
+    },
+    {
+      "epoch": 0.79,
+      "gumbel_tau": 0.1606,
+      "image_masks_loss": 0.0234,
+      "llm_learning_rate": 5.342768693134842e-07,
+      "loss": 0.749,
+      "output_text_mask_loss": 0.0622,
+      "predictor_learning_rate": 5.342768693134842e-05,
+      "step": 7753
+    },
+    {
+      "epoch": 0.79,
+      "gumbel_tau": 0.1606,
+      "image_masks_loss": 0.0092,
+      "llm_learning_rate": 5.337643917485319e-07,
+      "loss": 0.7347,
+      "output_text_mask_loss": 0.0147,
+      "predictor_learning_rate": 5.337643917485319e-05,
+      "step": 7754
+    },
+    {
+      "epoch": 0.79,
+      "gumbel_tau": 0.1605,
+      "image_masks_loss": 0.0114,
+      "llm_learning_rate": 5.332521307098246e-07,
+      "loss": 0.7749,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 5.3325213070982457e-05,
+      "step": 7755
+    },
+    {
+      "epoch": 0.79,
+      "gumbel_tau": 0.1605,
+      "image_masks_loss": 0.0075,
+      "llm_learning_rate": 5.327400862537732e-07,
+      "loss": 0.7917,
+      "output_text_mask_loss": 0.0246,
+      "predictor_learning_rate": 5.327400862537732e-05,
+      "step": 7756
+    },
+    {
+      "epoch": 0.79,
+      "gumbel_tau": 0.1604,
+      "image_masks_loss": 0.0228,
+      "llm_learning_rate": 5.322282584367655e-07,
+      "loss": 0.6131,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 5.3222825843676545e-05,
+      "step": 7757
+    },
+    {
+      "epoch": 0.79,
+      "gumbel_tau": 0.1604,
+      "image_masks_loss": 0.0088,
+      "llm_learning_rate": 5.317166473151649e-07,
+      "loss": 0.734,
+      "output_text_mask_loss": 0.0591,
+      "predictor_learning_rate": 5.31716647315165e-05,
+      "step": 7758
+    },
+    {
+      "epoch": 0.79,
+      "gumbel_tau": 0.1604,
+      "image_masks_loss": 0.0198,
+      "llm_learning_rate": 5.312052529453107e-07,
+      "loss": 0.8049,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 5.312052529453107e-05,
+      "step": 7759
+    },
+    {
+      "epoch": 0.8,
+      "gumbel_tau": 0.1603,
+      "image_masks_loss": 0.0307,
+      "llm_learning_rate": 5.306940753835202e-07,
+      "loss": 0.7066,
+      "output_text_mask_loss": 0.0284,
+      "predictor_learning_rate": 5.306940753835201e-05,
+      "step": 7760
+    },
+    {
+      "epoch": 0.8,
+      "gumbel_tau": 0.1603,
+      "image_masks_loss": 0.0154,
+      "llm_learning_rate": 5.301831146860842e-07,
+      "loss": 0.6992,
+      "output_text_mask_loss": 0.0126,
+      "predictor_learning_rate": 5.301831146860842e-05,
+      "step": 7761
+    },
+    {
+      "epoch": 0.8,
+      "gumbel_tau": 0.1603,
+      "image_masks_loss": 0.0237,
+      "llm_learning_rate": 5.296723709092716e-07,
+      "loss": 0.7293,
+      "output_text_mask_loss": 0.1144,
+      "predictor_learning_rate": 5.2967237090927155e-05,
+      "step": 7762
+    },
+    {
+      "epoch": 0.8,
+      "gumbel_tau": 0.1602,
+      "image_masks_loss": 0.0077,
+      "llm_learning_rate": 5.29161844109326e-07,
+      "loss": 0.7718,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 5.29161844109326e-05,
+      "step": 7763
+    },
+    {
+      "epoch": 0.8,
+      "gumbel_tau": 0.1602,
+      "image_masks_loss": 0.0196,
+      "llm_learning_rate": 5.286515343424681e-07,
+      "loss": 0.7303,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 5.2865153434246804e-05,
+      "step": 7764
+    },
+    {
+      "epoch": 0.8,
+      "gumbel_tau": 0.1601,
+      "image_masks_loss": 0.0138,
+      "llm_learning_rate": 5.281414416648941e-07,
+      "loss": 0.8089,
+      "output_text_mask_loss": 0.0231,
+      "predictor_learning_rate": 5.2814144166489415e-05,
+      "step": 7765
+    },
+    {
+      "epoch": 0.8,
+      "gumbel_tau": 0.1601,
+      "image_masks_loss": 0.0169,
+      "llm_learning_rate": 5.276315661327763e-07,
+      "loss": 0.7234,
+      "output_text_mask_loss": 0.0038,
+      "predictor_learning_rate": 5.276315661327763e-05,
+      "step": 7766
+    },
+    {
+      "epoch": 0.8,
+      "gumbel_tau": 0.1601,
+      "image_masks_loss": 0.0126,
+      "llm_learning_rate": 5.27121907802264e-07,
+      "loss": 0.679,
+      "output_text_mask_loss": 0.0086,
+      "predictor_learning_rate": 5.271219078022641e-05,
+      "step": 7767
+    },
+    {
+      "epoch": 0.8,
+      "gumbel_tau": 0.16,
+      "image_masks_loss": 0.0126,
+      "llm_learning_rate": 5.26612466729482e-07,
+      "loss": 0.7155,
+      "output_text_mask_loss": 0.0149,
+      "predictor_learning_rate": 5.266124667294819e-05,
+      "step": 7768
+    },
+    {
+      "epoch": 0.8,
+      "gumbel_tau": 0.16,
+      "image_masks_loss": 0.0081,
+      "llm_learning_rate": 5.261032429705312e-07,
+      "loss": 0.729,
+      "output_text_mask_loss": 0.0149,
+      "predictor_learning_rate": 5.261032429705312e-05,
+      "step": 7769
+    },
+    {
+      "epoch": 0.8,
+      "gumbel_tau": 0.16,
+      "image_masks_loss": 0.0108,
+      "llm_learning_rate": 5.255942365814873e-07,
+      "loss": 0.6899,
+      "output_text_mask_loss": 0.0202,
+      "predictor_learning_rate": 5.255942365814873e-05,
+      "step": 7770
+    },
+    {
+      "epoch": 0.8,
+      "gumbel_tau": 0.1599,
+      "image_masks_loss": 0.0189,
+      "llm_learning_rate": 5.250854476184034e-07,
+      "loss": 0.7061,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 5.250854476184033e-05,
+      "step": 7771
+    },
+    {
+      "epoch": 0.8,
+      "gumbel_tau": 0.1599,
+      "image_masks_loss": 0.0212,
+      "llm_learning_rate": 5.245768761373093e-07,
+      "loss": 0.6159,
+      "output_text_mask_loss": 0.0029,
+      "predictor_learning_rate": 5.245768761373093e-05,
+      "step": 7772
+    },
+    {
+      "epoch": 0.8,
+      "gumbel_tau": 0.1598,
+      "image_masks_loss": 0.008,
+      "llm_learning_rate": 5.240685221942099e-07,
+      "loss": 0.8171,
+      "output_text_mask_loss": 0.02,
+      "predictor_learning_rate": 5.240685221942099e-05,
+      "step": 7773
+    },
+    {
+      "epoch": 0.8,
+      "gumbel_tau": 0.1598,
+      "image_masks_loss": 0.003,
+      "llm_learning_rate": 5.235603858450861e-07,
+      "loss": 0.7207,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 5.2356038584508606e-05,
+      "step": 7774
+    },
+    {
+      "epoch": 0.8,
+      "gumbel_tau": 0.1598,
+      "image_masks_loss": 0.0102,
+      "llm_learning_rate": 5.230524671458947e-07,
+      "loss": 0.7455,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 5.2305246714589474e-05,
+      "step": 7775
+    },
+    {
+      "epoch": 0.8,
+      "gumbel_tau": 0.1597,
+      "image_masks_loss": 0.0033,
+      "llm_learning_rate": 5.225447661525693e-07,
+      "loss": 0.6537,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 5.225447661525692e-05,
+      "step": 7776
+    },
+    {
+      "epoch": 0.8,
+      "gumbel_tau": 0.1597,
+      "image_masks_loss": 0.0094,
+      "llm_learning_rate": 5.220372829210188e-07,
+      "loss": 0.7456,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 5.220372829210188e-05,
+      "step": 7777
+    },
+    {
+      "epoch": 0.8,
+      "gumbel_tau": 0.1597,
+      "image_masks_loss": 0.0175,
+      "llm_learning_rate": 5.215300175071278e-07,
+      "loss": 0.7646,
+      "output_text_mask_loss": 0.0038,
+      "predictor_learning_rate": 5.2153001750712775e-05,
+      "step": 7778
+    },
+    {
+      "epoch": 0.8,
+      "gumbel_tau": 0.1596,
+      "image_masks_loss": 0.0057,
+      "llm_learning_rate": 5.210229699667588e-07,
+      "loss": 0.7515,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 5.210229699667587e-05,
+      "step": 7779
+    },
+    {
+      "epoch": 0.8,
+      "gumbel_tau": 0.1596,
+      "image_masks_loss": 0.0101,
+      "llm_learning_rate": 5.205161403557483e-07,
+      "loss": 0.7547,
+      "output_text_mask_loss": 0.0071,
+      "predictor_learning_rate": 5.205161403557482e-05,
+      "step": 7780
+    },
+    {
+      "epoch": 0.8,
+      "gumbel_tau": 0.1595,
+      "image_masks_loss": 0.0036,
+      "llm_learning_rate": 5.200095287299095e-07,
+      "loss": 0.6763,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 5.200095287299095e-05,
+      "step": 7781
+    },
+    {
+      "epoch": 0.8,
+      "gumbel_tau": 0.1595,
+      "image_masks_loss": 0.0307,
+      "llm_learning_rate": 5.19503135145032e-07,
+      "loss": 0.7153,
+      "output_text_mask_loss": 0.0383,
+      "predictor_learning_rate": 5.1950313514503186e-05,
+      "step": 7782
+    },
+    {
+      "epoch": 0.8,
+      "gumbel_tau": 0.1595,
+      "image_masks_loss": 0.0123,
+      "llm_learning_rate": 5.189969596568803e-07,
+      "loss": 0.6826,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 5.1899695965688026e-05,
+      "step": 7783
+    },
+    {
+      "epoch": 0.8,
+      "gumbel_tau": 0.1594,
+      "image_masks_loss": 0.0114,
+      "llm_learning_rate": 5.184910023211959e-07,
+      "loss": 0.7145,
+      "output_text_mask_loss": 0.0004,
+      "predictor_learning_rate": 5.1849100232119585e-05,
+      "step": 7784
+    },
+    {
+      "epoch": 0.8,
+      "gumbel_tau": 0.1594,
+      "image_masks_loss": 0.0083,
+      "llm_learning_rate": 5.179852631936954e-07,
+      "loss": 0.614,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 5.1798526319369544e-05,
+      "step": 7785
+    },
+    {
+      "epoch": 0.8,
+      "gumbel_tau": 0.1594,
+      "image_masks_loss": 0.0343,
+      "llm_learning_rate": 5.174797423300732e-07,
+      "loss": 0.7558,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 5.1747974233007314e-05,
+      "step": 7786
+    },
+    {
+      "epoch": 0.8,
+      "gumbel_tau": 0.1593,
+      "image_masks_loss": 0.0161,
+      "llm_learning_rate": 5.169744397859983e-07,
+      "loss": 0.7437,
+      "output_text_mask_loss": 0.028,
+      "predictor_learning_rate": 5.1697443978599826e-05,
+      "step": 7787
+    },
+    {
+      "epoch": 0.8,
+      "gumbel_tau": 0.1593,
+      "image_masks_loss": 0.0123,
+      "llm_learning_rate": 5.164693556171144e-07,
+      "loss": 0.654,
+      "output_text_mask_loss": 0.001,
+      "predictor_learning_rate": 5.1646935561711435e-05,
+      "step": 7788
+    },
+    {
+      "epoch": 0.8,
+      "gumbel_tau": 0.1592,
+      "image_masks_loss": 0.0317,
+      "llm_learning_rate": 5.159644898790433e-07,
+      "loss": 0.8172,
+      "output_text_mask_loss": 0.0107,
+      "predictor_learning_rate": 5.159644898790433e-05,
+      "step": 7789
+    },
+    {
+      "epoch": 0.8,
+      "gumbel_tau": 0.1592,
+      "image_masks_loss": 0.0233,
+      "llm_learning_rate": 5.154598426273816e-07,
+      "loss": 0.7553,
+      "output_text_mask_loss": 0.0056,
+      "predictor_learning_rate": 5.1545984262738155e-05,
+      "step": 7790
+    },
+    {
+      "epoch": 0.8,
+      "gumbel_tau": 0.1592,
+      "image_masks_loss": 0.0151,
+      "llm_learning_rate": 5.149554139177021e-07,
+      "loss": 0.785,
+      "output_text_mask_loss": 0.0077,
+      "predictor_learning_rate": 5.149554139177021e-05,
+      "step": 7791
+    },
+    {
+      "epoch": 0.8,
+      "gumbel_tau": 0.1591,
+      "image_masks_loss": 0.0078,
+      "llm_learning_rate": 5.144512038055544e-07,
+      "loss": 0.7158,
+      "output_text_mask_loss": 0.0034,
+      "predictor_learning_rate": 5.144512038055543e-05,
+      "step": 7792
+    },
+    {
+      "epoch": 0.8,
+      "gumbel_tau": 0.1591,
+      "image_masks_loss": 0.0109,
+      "llm_learning_rate": 5.139472123464629e-07,
+      "loss": 0.6907,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 5.139472123464628e-05,
+      "step": 7793
+    },
+    {
+      "epoch": 0.8,
+      "gumbel_tau": 0.1591,
+      "image_masks_loss": 0.0121,
+      "llm_learning_rate": 5.134434395959279e-07,
+      "loss": 0.7514,
+      "output_text_mask_loss": 0.0156,
+      "predictor_learning_rate": 5.1344343959592784e-05,
+      "step": 7794
+    },
+    {
+      "epoch": 0.8,
+      "gumbel_tau": 0.159,
+      "image_masks_loss": 0.0212,
+      "llm_learning_rate": 5.129398856094262e-07,
+      "loss": 0.6701,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 5.129398856094261e-05,
+      "step": 7795
+    },
+    {
+      "epoch": 0.8,
+      "gumbel_tau": 0.159,
+      "image_masks_loss": 0.0172,
+      "llm_learning_rate": 5.124365504424106e-07,
+      "loss": 0.7653,
+      "output_text_mask_loss": 0.013,
+      "predictor_learning_rate": 5.1243655044241055e-05,
+      "step": 7796
+    },
+    {
+      "epoch": 0.8,
+      "gumbel_tau": 0.1589,
+      "image_masks_loss": 0.011,
+      "llm_learning_rate": 5.11933434150309e-07,
+      "loss": 0.7753,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 5.1193343415030905e-05,
+      "step": 7797
+    },
+    {
+      "epoch": 0.8,
+      "gumbel_tau": 0.1589,
+      "image_masks_loss": 0.0135,
+      "llm_learning_rate": 5.114305367885253e-07,
+      "loss": 0.6907,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 5.114305367885252e-05,
+      "step": 7798
+    },
+    {
+      "epoch": 0.8,
+      "gumbel_tau": 0.1589,
+      "image_masks_loss": 0.0098,
+      "llm_learning_rate": 5.109278584124408e-07,
+      "loss": 0.6908,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 5.109278584124408e-05,
+      "step": 7799
+    },
+    {
+      "epoch": 0.8,
+      "gumbel_tau": 0.1588,
+      "image_masks_loss": 0.0311,
+      "llm_learning_rate": 5.10425399077411e-07,
+      "loss": 0.6997,
+      "output_text_mask_loss": 0.0058,
+      "predictor_learning_rate": 5.1042539907741104e-05,
+      "step": 7800
+    },
+    {
+      "epoch": 0.8,
+      "gumbel_tau": 0.1588,
+      "image_masks_loss": 0.0154,
+      "llm_learning_rate": 5.099231588387679e-07,
+      "loss": 0.7244,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 5.099231588387679e-05,
+      "step": 7801
+    },
+    {
+      "epoch": 0.8,
+      "gumbel_tau": 0.1588,
+      "image_masks_loss": 0.0092,
+      "llm_learning_rate": 5.094211377518193e-07,
+      "loss": 0.6668,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 5.094211377518193e-05,
+      "step": 7802
+    },
+    {
+      "epoch": 0.8,
+      "gumbel_tau": 0.1587,
+      "image_masks_loss": 0.0113,
+      "llm_learning_rate": 5.089193358718488e-07,
+      "loss": 0.7514,
+      "output_text_mask_loss": 0.0275,
+      "predictor_learning_rate": 5.0891933587184886e-05,
+      "step": 7803
+    },
+    {
+      "epoch": 0.8,
+      "gumbel_tau": 0.1587,
+      "image_masks_loss": 0.0161,
+      "llm_learning_rate": 5.084177532541162e-07,
+      "loss": 0.5697,
+      "output_text_mask_loss": 0.0081,
+      "predictor_learning_rate": 5.0841775325411623e-05,
+      "step": 7804
+    },
+    {
+      "epoch": 0.8,
+      "gumbel_tau": 0.1586,
+      "image_masks_loss": 0.0032,
+      "llm_learning_rate": 5.079163899538559e-07,
+      "loss": 0.6326,
+      "output_text_mask_loss": 0.0012,
+      "predictor_learning_rate": 5.079163899538558e-05,
+      "step": 7805
+    },
+    {
+      "epoch": 0.8,
+      "gumbel_tau": 0.1586,
+      "image_masks_loss": 0.018,
+      "llm_learning_rate": 5.07415246026281e-07,
+      "loss": 0.7326,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 5.0741524602628096e-05,
+      "step": 7806
+    },
+    {
+      "epoch": 0.8,
+      "gumbel_tau": 0.1586,
+      "image_masks_loss": 0.0087,
+      "llm_learning_rate": 5.06914321526577e-07,
+      "loss": 0.783,
+      "output_text_mask_loss": 0.0221,
+      "predictor_learning_rate": 5.06914321526577e-05,
+      "step": 7807
+    },
+    {
+      "epoch": 0.8,
+      "gumbel_tau": 0.1585,
+      "image_masks_loss": 0.0128,
+      "llm_learning_rate": 5.064136165099076e-07,
+      "loss": 0.7539,
+      "output_text_mask_loss": 0.0095,
+      "predictor_learning_rate": 5.0641361650990756e-05,
+      "step": 7808
+    },
+    {
+      "epoch": 0.8,
+      "gumbel_tau": 0.1585,
+      "image_masks_loss": 0.0202,
+      "llm_learning_rate": 5.059131310314111e-07,
+      "loss": 0.7151,
+      "output_text_mask_loss": 0.0185,
+      "predictor_learning_rate": 5.059131310314111e-05,
+      "step": 7809
+    },
+    {
+      "epoch": 0.8,
+      "gumbel_tau": 0.1585,
+      "image_masks_loss": 0.013,
+      "llm_learning_rate": 5.054128651462017e-07,
+      "loss": 0.7711,
+      "output_text_mask_loss": 0.0322,
+      "predictor_learning_rate": 5.054128651462017e-05,
+      "step": 7810
+    },
+    {
+      "epoch": 0.8,
+      "gumbel_tau": 0.1584,
+      "image_masks_loss": 0.0324,
+      "llm_learning_rate": 5.049128189093708e-07,
+      "loss": 0.6563,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 5.0491281890937074e-05,
+      "step": 7811
+    },
+    {
+      "epoch": 0.8,
+      "gumbel_tau": 0.1584,
+      "image_masks_loss": 0.0135,
+      "llm_learning_rate": 5.044129923759844e-07,
+      "loss": 0.575,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 5.044129923759844e-05,
+      "step": 7812
+    },
+    {
+      "epoch": 0.8,
+      "gumbel_tau": 0.1583,
+      "image_masks_loss": 0.0092,
+      "llm_learning_rate": 5.039133856010841e-07,
+      "loss": 0.7626,
+      "output_text_mask_loss": 0.0156,
+      "predictor_learning_rate": 5.0391338560108414e-05,
+      "step": 7813
+    },
+    {
+      "epoch": 0.8,
+      "gumbel_tau": 0.1583,
+      "image_masks_loss": 0.0216,
+      "llm_learning_rate": 5.034139986396879e-07,
+      "loss": 0.7653,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 5.03413998639688e-05,
+      "step": 7814
+    },
+    {
+      "epoch": 0.8,
+      "gumbel_tau": 0.1583,
+      "image_masks_loss": 0.0093,
+      "llm_learning_rate": 5.029148315467894e-07,
+      "loss": 0.683,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 5.0291483154678934e-05,
+      "step": 7815
+    },
+    {
+      "epoch": 0.8,
+      "gumbel_tau": 0.1582,
+      "image_masks_loss": 0.0157,
+      "llm_learning_rate": 5.024158843773582e-07,
+      "loss": 0.7804,
+      "output_text_mask_loss": 0.0231,
+      "predictor_learning_rate": 5.024158843773582e-05,
+      "step": 7816
+    },
+    {
+      "epoch": 0.8,
+      "gumbel_tau": 0.1582,
+      "image_masks_loss": 0.0112,
+      "llm_learning_rate": 5.019171571863385e-07,
+      "loss": 0.767,
+      "output_text_mask_loss": 0.0668,
+      "predictor_learning_rate": 5.019171571863385e-05,
+      "step": 7817
+    },
+    {
+      "epoch": 0.8,
+      "gumbel_tau": 0.1582,
+      "image_masks_loss": 0.0366,
+      "llm_learning_rate": 5.014186500286525e-07,
+      "loss": 0.8307,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 5.014186500286525e-05,
+      "step": 7818
+    },
+    {
+      "epoch": 0.8,
+      "gumbel_tau": 0.1581,
+      "image_masks_loss": 0.0056,
+      "llm_learning_rate": 5.009203629591966e-07,
+      "loss": 0.7465,
+      "output_text_mask_loss": 0.0139,
+      "predictor_learning_rate": 5.009203629591966e-05,
+      "step": 7819
+    },
+    {
+      "epoch": 0.8,
+      "gumbel_tau": 0.1581,
+      "image_masks_loss": 0.0081,
+      "llm_learning_rate": 5.004222960328428e-07,
+      "loss": 0.7075,
+      "output_text_mask_loss": 0.033,
+      "predictor_learning_rate": 5.004222960328428e-05,
+      "step": 7820
+    },
+    {
+      "epoch": 0.8,
+      "gumbel_tau": 0.158,
+      "image_masks_loss": 0.0166,
+      "llm_learning_rate": 4.9992444930444e-07,
+      "loss": 0.8363,
+      "output_text_mask_loss": 0.0345,
+      "predictor_learning_rate": 4.9992444930443995e-05,
+      "step": 7821
+    },
+    {
+      "epoch": 0.8,
+      "gumbel_tau": 0.158,
+      "image_masks_loss": 0.0099,
+      "llm_learning_rate": 4.994268228288115e-07,
+      "loss": 0.8092,
+      "output_text_mask_loss": 0.0026,
+      "predictor_learning_rate": 4.994268228288115e-05,
+      "step": 7822
+    },
+    {
+      "epoch": 0.8,
+      "gumbel_tau": 0.158,
+      "image_masks_loss": 0.0163,
+      "llm_learning_rate": 4.989294166607578e-07,
+      "loss": 0.702,
+      "output_text_mask_loss": 0.0519,
+      "predictor_learning_rate": 4.989294166607577e-05,
+      "step": 7823
+    },
+    {
+      "epoch": 0.8,
+      "gumbel_tau": 0.1579,
+      "image_masks_loss": 0.0087,
+      "llm_learning_rate": 4.984322308550538e-07,
+      "loss": 0.6187,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 4.9843223085505374e-05,
+      "step": 7824
+    },
+    {
+      "epoch": 0.8,
+      "gumbel_tau": 0.1579,
+      "image_masks_loss": 0.0055,
+      "llm_learning_rate": 4.979352654664507e-07,
+      "loss": 0.788,
+      "output_text_mask_loss": 0.006,
+      "predictor_learning_rate": 4.979352654664507e-05,
+      "step": 7825
+    },
+    {
+      "epoch": 0.8,
+      "gumbel_tau": 0.1579,
+      "image_masks_loss": 0.0024,
+      "llm_learning_rate": 4.974385205496756e-07,
+      "loss": 0.7505,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 4.9743852054967565e-05,
+      "step": 7826
+    },
+    {
+      "epoch": 0.8,
+      "gumbel_tau": 0.1578,
+      "image_masks_loss": 0.0048,
+      "llm_learning_rate": 4.969419961594313e-07,
+      "loss": 0.7453,
+      "output_text_mask_loss": 0.0004,
+      "predictor_learning_rate": 4.969419961594313e-05,
+      "step": 7827
+    },
+    {
+      "epoch": 0.8,
+      "gumbel_tau": 0.1578,
+      "image_masks_loss": 0.0091,
+      "llm_learning_rate": 4.96445692350396e-07,
+      "loss": 0.7327,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 4.96445692350396e-05,
+      "step": 7828
+    },
+    {
+      "epoch": 0.8,
+      "gumbel_tau": 0.1577,
+      "image_masks_loss": 0.0136,
+      "llm_learning_rate": 4.959496091772237e-07,
+      "loss": 0.7294,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 4.9594960917722366e-05,
+      "step": 7829
+    },
+    {
+      "epoch": 0.8,
+      "gumbel_tau": 0.1577,
+      "image_masks_loss": 0.0117,
+      "llm_learning_rate": 4.954537466945436e-07,
+      "loss": 0.7664,
+      "output_text_mask_loss": 0.0139,
+      "predictor_learning_rate": 4.954537466945436e-05,
+      "step": 7830
+    },
+    {
+      "epoch": 0.8,
+      "gumbel_tau": 0.1577,
+      "image_masks_loss": 0.0166,
+      "llm_learning_rate": 4.949581049569627e-07,
+      "loss": 0.7315,
+      "output_text_mask_loss": 0.0618,
+      "predictor_learning_rate": 4.9495810495696266e-05,
+      "step": 7831
+    },
+    {
+      "epoch": 0.8,
+      "gumbel_tau": 0.1576,
+      "image_masks_loss": 0.0121,
+      "llm_learning_rate": 4.944626840190611e-07,
+      "loss": 0.6851,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 4.944626840190611e-05,
+      "step": 7832
+    },
+    {
+      "epoch": 0.8,
+      "gumbel_tau": 0.1576,
+      "image_masks_loss": 0.0269,
+      "llm_learning_rate": 4.939674839353961e-07,
+      "loss": 0.7219,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 4.939674839353961e-05,
+      "step": 7833
+    },
+    {
+      "epoch": 0.8,
+      "gumbel_tau": 0.1576,
+      "image_masks_loss": 0.0056,
+      "llm_learning_rate": 4.934725047604999e-07,
+      "loss": 0.7482,
+      "output_text_mask_loss": 0.0265,
+      "predictor_learning_rate": 4.934725047604999e-05,
+      "step": 7834
+    },
+    {
+      "epoch": 0.8,
+      "gumbel_tau": 0.1575,
+      "image_masks_loss": 0.0092,
+      "llm_learning_rate": 4.929777465488808e-07,
+      "loss": 0.6175,
+      "output_text_mask_loss": 0.0004,
+      "predictor_learning_rate": 4.929777465488808e-05,
+      "step": 7835
+    },
+    {
+      "epoch": 0.8,
+      "gumbel_tau": 0.1575,
+      "image_masks_loss": 0.0059,
+      "llm_learning_rate": 4.924832093550231e-07,
+      "loss": 0.7016,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 4.92483209355023e-05,
+      "step": 7836
+    },
+    {
+      "epoch": 0.8,
+      "gumbel_tau": 0.1574,
+      "image_masks_loss": 0.0121,
+      "llm_learning_rate": 4.91988893233385e-07,
+      "loss": 0.801,
+      "output_text_mask_loss": 0.0041,
+      "predictor_learning_rate": 4.9198889323338494e-05,
+      "step": 7837
+    },
+    {
+      "epoch": 0.8,
+      "gumbel_tau": 0.1574,
+      "image_masks_loss": 0.0029,
+      "llm_learning_rate": 4.91494798238403e-07,
+      "loss": 0.7307,
+      "output_text_mask_loss": 0.005,
+      "predictor_learning_rate": 4.91494798238403e-05,
+      "step": 7838
+    },
+    {
+      "epoch": 0.8,
+      "gumbel_tau": 0.1574,
+      "image_masks_loss": 0.0045,
+      "llm_learning_rate": 4.910009244244881e-07,
+      "loss": 0.8029,
+      "output_text_mask_loss": 0.0073,
+      "predictor_learning_rate": 4.910009244244881e-05,
+      "step": 7839
+    },
+    {
+      "epoch": 0.8,
+      "gumbel_tau": 0.1573,
+      "image_masks_loss": 0.0086,
+      "llm_learning_rate": 4.905072718460261e-07,
+      "loss": 0.7149,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 4.90507271846026e-05,
+      "step": 7840
+    },
+    {
+      "epoch": 0.8,
+      "gumbel_tau": 0.1573,
+      "image_masks_loss": 0.0135,
+      "llm_learning_rate": 4.900138405573793e-07,
+      "loss": 0.7413,
+      "output_text_mask_loss": 0.006,
+      "predictor_learning_rate": 4.900138405573792e-05,
+      "step": 7841
+    },
+    {
+      "epoch": 0.8,
+      "gumbel_tau": 0.1573,
+      "image_masks_loss": 0.0132,
+      "llm_learning_rate": 4.895206306128852e-07,
+      "loss": 0.8022,
+      "output_text_mask_loss": 0.0034,
+      "predictor_learning_rate": 4.895206306128852e-05,
+      "step": 7842
+    },
+    {
+      "epoch": 0.8,
+      "gumbel_tau": 0.1572,
+      "image_masks_loss": 0.0194,
+      "llm_learning_rate": 4.890276420668577e-07,
+      "loss": 0.6776,
+      "output_text_mask_loss": 0.0082,
+      "predictor_learning_rate": 4.890276420668577e-05,
+      "step": 7843
+    },
+    {
+      "epoch": 0.8,
+      "gumbel_tau": 0.1572,
+      "image_masks_loss": 0.0248,
+      "llm_learning_rate": 4.885348749735854e-07,
+      "loss": 0.6789,
+      "output_text_mask_loss": 0.001,
+      "predictor_learning_rate": 4.885348749735854e-05,
+      "step": 7844
+    },
+    {
+      "epoch": 0.8,
+      "gumbel_tau": 0.1571,
+      "image_masks_loss": 0.0186,
+      "llm_learning_rate": 4.880423293873332e-07,
+      "loss": 0.701,
+      "output_text_mask_loss": 0.0185,
+      "predictor_learning_rate": 4.880423293873332e-05,
+      "step": 7845
+    },
+    {
+      "epoch": 0.8,
+      "gumbel_tau": 0.1571,
+      "image_masks_loss": 0.0065,
+      "llm_learning_rate": 4.875500053623411e-07,
+      "loss": 0.8105,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 4.8755000536234106e-05,
+      "step": 7846
+    },
+    {
+      "epoch": 0.8,
+      "gumbel_tau": 0.1571,
+      "image_masks_loss": 0.0126,
+      "llm_learning_rate": 4.87057902952825e-07,
+      "loss": 0.6128,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 4.870579029528249e-05,
+      "step": 7847
+    },
+    {
+      "epoch": 0.8,
+      "gumbel_tau": 0.157,
+      "image_masks_loss": 0.0134,
+      "llm_learning_rate": 4.865660222129759e-07,
+      "loss": 0.7411,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 4.8656602221297586e-05,
+      "step": 7848
+    },
+    {
+      "epoch": 0.8,
+      "gumbel_tau": 0.157,
+      "image_masks_loss": 0.0151,
+      "llm_learning_rate": 4.860743631969608e-07,
+      "loss": 0.7568,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 4.860743631969608e-05,
+      "step": 7849
+    },
+    {
+      "epoch": 0.8,
+      "gumbel_tau": 0.157,
+      "image_masks_loss": 0.0145,
+      "llm_learning_rate": 4.855829259589234e-07,
+      "loss": 0.7483,
+      "output_text_mask_loss": 0.0023,
+      "predictor_learning_rate": 4.855829259589234e-05,
+      "step": 7850
+    },
+    {
+      "epoch": 0.8,
+      "gumbel_tau": 0.1569,
+      "image_masks_loss": 0.0082,
+      "llm_learning_rate": 4.850917105529812e-07,
+      "loss": 0.6408,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 4.850917105529812e-05,
+      "step": 7851
+    },
+    {
+      "epoch": 0.8,
+      "gumbel_tau": 0.1569,
+      "image_masks_loss": 0.0087,
+      "llm_learning_rate": 4.846007170332276e-07,
+      "loss": 0.72,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 4.846007170332276e-05,
+      "step": 7852
+    },
+    {
+      "epoch": 0.8,
+      "gumbel_tau": 0.1569,
+      "image_masks_loss": 0.0082,
+      "llm_learning_rate": 4.841099454537324e-07,
+      "loss": 0.5783,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 4.8410994545373234e-05,
+      "step": 7853
+    },
+    {
+      "epoch": 0.8,
+      "gumbel_tau": 0.1568,
+      "image_masks_loss": 0.0062,
+      "llm_learning_rate": 4.8361939586854e-07,
+      "loss": 0.7295,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 4.8361939586854006e-05,
+      "step": 7854
+    },
+    {
+      "epoch": 0.8,
+      "gumbel_tau": 0.1568,
+      "image_masks_loss": 0.007,
+      "llm_learning_rate": 4.831290683316714e-07,
+      "loss": 0.6194,
+      "output_text_mask_loss": 0.0039,
+      "predictor_learning_rate": 4.831290683316714e-05,
+      "step": 7855
+    },
+    {
+      "epoch": 0.8,
+      "gumbel_tau": 0.1567,
+      "image_masks_loss": 0.0252,
+      "llm_learning_rate": 4.826389628971212e-07,
+      "loss": 0.6751,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 4.8263896289712115e-05,
+      "step": 7856
+    },
+    {
+      "epoch": 0.81,
+      "gumbel_tau": 0.1567,
+      "image_masks_loss": 0.009,
+      "llm_learning_rate": 4.821490796188627e-07,
+      "loss": 0.7022,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 4.821490796188627e-05,
+      "step": 7857
+    },
+    {
+      "epoch": 0.81,
+      "gumbel_tau": 0.1567,
+      "image_masks_loss": 0.015,
+      "llm_learning_rate": 4.816594185508419e-07,
+      "loss": 0.716,
+      "output_text_mask_loss": 0.0019,
+      "predictor_learning_rate": 4.816594185508419e-05,
+      "step": 7858
+    },
+    {
+      "epoch": 0.81,
+      "gumbel_tau": 0.1566,
+      "image_masks_loss": 0.0056,
+      "llm_learning_rate": 4.811699797469818e-07,
+      "loss": 0.7535,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 4.811699797469818e-05,
+      "step": 7859
+    },
+    {
+      "epoch": 0.81,
+      "gumbel_tau": 0.1566,
+      "image_masks_loss": 0.0282,
+      "llm_learning_rate": 4.806807632611804e-07,
+      "loss": 0.7595,
+      "output_text_mask_loss": 0.0071,
+      "predictor_learning_rate": 4.8068076326118035e-05,
+      "step": 7860
+    },
+    {
+      "epoch": 0.81,
+      "gumbel_tau": 0.1566,
+      "image_masks_loss": 0.0148,
+      "llm_learning_rate": 4.801917691473113e-07,
+      "loss": 0.6758,
+      "output_text_mask_loss": 0.0271,
+      "predictor_learning_rate": 4.8019176914731124e-05,
+      "step": 7861
+    },
+    {
+      "epoch": 0.81,
+      "gumbel_tau": 0.1565,
+      "image_masks_loss": 0.0107,
+      "llm_learning_rate": 4.797029974592233e-07,
+      "loss": 0.6438,
+      "output_text_mask_loss": 0.0141,
+      "predictor_learning_rate": 4.797029974592232e-05,
+      "step": 7862
+    },
+    {
+      "epoch": 0.81,
+      "gumbel_tau": 0.1565,
+      "image_masks_loss": 0.0062,
+      "llm_learning_rate": 4.792144482507413e-07,
+      "loss": 0.7671,
+      "output_text_mask_loss": 0.0208,
+      "predictor_learning_rate": 4.792144482507413e-05,
+      "step": 7863
+    },
+    {
+      "epoch": 0.81,
+      "gumbel_tau": 0.1564,
+      "image_masks_loss": 0.0063,
+      "llm_learning_rate": 4.787261215756653e-07,
+      "loss": 0.7993,
+      "output_text_mask_loss": 0.0345,
+      "predictor_learning_rate": 4.7872612157566525e-05,
+      "step": 7864
+    },
+    {
+      "epoch": 0.81,
+      "gumbel_tau": 0.1564,
+      "image_masks_loss": 0.0059,
+      "llm_learning_rate": 4.78238017487771e-07,
+      "loss": 0.7098,
+      "output_text_mask_loss": 0.0467,
+      "predictor_learning_rate": 4.7823801748777095e-05,
+      "step": 7865
+    },
+    {
+      "epoch": 0.81,
+      "gumbel_tau": 0.1564,
+      "image_masks_loss": 0.0071,
+      "llm_learning_rate": 4.777501360408096e-07,
+      "loss": 0.7694,
+      "output_text_mask_loss": 0.0151,
+      "predictor_learning_rate": 4.777501360408096e-05,
+      "step": 7866
+    },
+    {
+      "epoch": 0.81,
+      "gumbel_tau": 0.1563,
+      "image_masks_loss": 0.0139,
+      "llm_learning_rate": 4.772624772885076e-07,
+      "loss": 0.6399,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 4.7726247728850756e-05,
+      "step": 7867
+    },
+    {
+      "epoch": 0.81,
+      "gumbel_tau": 0.1563,
+      "image_masks_loss": 0.0771,
+      "llm_learning_rate": 4.7677504128456704e-07,
+      "loss": 0.7855,
+      "output_text_mask_loss": 0.013,
+      "predictor_learning_rate": 4.7677504128456703e-05,
+      "step": 7868
+    },
+    {
+      "epoch": 0.81,
+      "gumbel_tau": 0.1563,
+      "image_masks_loss": 0.0022,
+      "llm_learning_rate": 4.7628782808266496e-07,
+      "loss": 0.6545,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 4.762878280826649e-05,
+      "step": 7869
+    },
+    {
+      "epoch": 0.81,
+      "gumbel_tau": 0.1562,
+      "image_masks_loss": 0.0128,
+      "llm_learning_rate": 4.7580083773645563e-07,
+      "loss": 0.6898,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 4.758008377364556e-05,
+      "step": 7870
+    },
+    {
+      "epoch": 0.81,
+      "gumbel_tau": 0.1562,
+      "image_masks_loss": 0.012,
+      "llm_learning_rate": 4.753140702995665e-07,
+      "loss": 0.7307,
+      "output_text_mask_loss": 0.0288,
+      "predictor_learning_rate": 4.7531407029956644e-05,
+      "step": 7871
+    },
+    {
+      "epoch": 0.81,
+      "gumbel_tau": 0.1562,
+      "image_masks_loss": 0.0098,
+      "llm_learning_rate": 4.748275258256019e-07,
+      "loss": 0.7445,
+      "output_text_mask_loss": 0.0226,
+      "predictor_learning_rate": 4.7482752582560185e-05,
+      "step": 7872
+    },
+    {
+      "epoch": 0.81,
+      "gumbel_tau": 0.1561,
+      "image_masks_loss": 0.0149,
+      "llm_learning_rate": 4.7434120436814087e-07,
+      "loss": 0.6704,
+      "output_text_mask_loss": 0.0193,
+      "predictor_learning_rate": 4.743412043681408e-05,
+      "step": 7873
+    },
+    {
+      "epoch": 0.81,
+      "gumbel_tau": 0.1561,
+      "image_masks_loss": 0.0081,
+      "llm_learning_rate": 4.7385510598073865e-07,
+      "loss": 0.7405,
+      "output_text_mask_loss": 0.0151,
+      "predictor_learning_rate": 4.738551059807386e-05,
+      "step": 7874
+    },
+    {
+      "epoch": 0.81,
+      "gumbel_tau": 0.156,
+      "image_masks_loss": 0.0043,
+      "llm_learning_rate": 4.733692307169252e-07,
+      "loss": 0.6437,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 4.7336923071692517e-05,
+      "step": 7875
+    },
+    {
+      "epoch": 0.81,
+      "gumbel_tau": 0.156,
+      "image_masks_loss": 0.0087,
+      "llm_learning_rate": 4.728835786302052e-07,
+      "loss": 0.7363,
+      "output_text_mask_loss": 0.0057,
+      "predictor_learning_rate": 4.728835786302052e-05,
+      "step": 7876
+    },
+    {
+      "epoch": 0.81,
+      "gumbel_tau": 0.156,
+      "image_masks_loss": 0.0117,
+      "llm_learning_rate": 4.723981497740618e-07,
+      "loss": 0.7622,
+      "output_text_mask_loss": 0.0153,
+      "predictor_learning_rate": 4.723981497740618e-05,
+      "step": 7877
+    },
+    {
+      "epoch": 0.81,
+      "gumbel_tau": 0.1559,
+      "image_masks_loss": 0.0034,
+      "llm_learning_rate": 4.719129442019504e-07,
+      "loss": 0.8053,
+      "output_text_mask_loss": 0.0225,
+      "predictor_learning_rate": 4.719129442019504e-05,
+      "step": 7878
+    },
+    {
+      "epoch": 0.81,
+      "gumbel_tau": 0.1559,
+      "image_masks_loss": 0.023,
+      "llm_learning_rate": 4.714279619673029e-07,
+      "loss": 0.8005,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 4.714279619673029e-05,
+      "step": 7879
+    },
+    {
+      "epoch": 0.81,
+      "gumbel_tau": 0.1559,
+      "image_masks_loss": 0.0035,
+      "llm_learning_rate": 4.7094320312352724e-07,
+      "loss": 0.6947,
+      "output_text_mask_loss": 0.0519,
+      "predictor_learning_rate": 4.7094320312352726e-05,
+      "step": 7880
+    },
+    {
+      "epoch": 0.81,
+      "gumbel_tau": 0.1558,
+      "image_masks_loss": 0.0062,
+      "llm_learning_rate": 4.7045866772400446e-07,
+      "loss": 0.868,
+      "output_text_mask_loss": 0.0086,
+      "predictor_learning_rate": 4.7045866772400444e-05,
+      "step": 7881
+    },
+    {
+      "epoch": 0.81,
+      "gumbel_tau": 0.1558,
+      "image_masks_loss": 0.018,
+      "llm_learning_rate": 4.6997435582209427e-07,
+      "loss": 0.766,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 4.6997435582209425e-05,
+      "step": 7882
+    },
+    {
+      "epoch": 0.81,
+      "gumbel_tau": 0.1557,
+      "image_masks_loss": 0.0071,
+      "llm_learning_rate": 4.694902674711299e-07,
+      "loss": 0.7399,
+      "output_text_mask_loss": 0.0553,
+      "predictor_learning_rate": 4.694902674711299e-05,
+      "step": 7883
+    },
+    {
+      "epoch": 0.81,
+      "gumbel_tau": 0.1557,
+      "image_masks_loss": 0.0074,
+      "llm_learning_rate": 4.690064027244201e-07,
+      "loss": 0.7402,
+      "output_text_mask_loss": 0.001,
+      "predictor_learning_rate": 4.690064027244201e-05,
+      "step": 7884
+    },
+    {
+      "epoch": 0.81,
+      "gumbel_tau": 0.1557,
+      "image_masks_loss": 0.0138,
+      "llm_learning_rate": 4.685227616352489e-07,
+      "loss": 0.7733,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 4.685227616352489e-05,
+      "step": 7885
+    },
+    {
+      "epoch": 0.81,
+      "gumbel_tau": 0.1556,
+      "image_masks_loss": 0.0175,
+      "llm_learning_rate": 4.680393442568762e-07,
+      "loss": 0.6583,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 4.6803934425687616e-05,
+      "step": 7886
+    },
+    {
+      "epoch": 0.81,
+      "gumbel_tau": 0.1556,
+      "image_masks_loss": 0.0092,
+      "llm_learning_rate": 4.675561506425369e-07,
+      "loss": 0.6919,
+      "output_text_mask_loss": 0.0081,
+      "predictor_learning_rate": 4.675561506425369e-05,
+      "step": 7887
+    },
+    {
+      "epoch": 0.81,
+      "gumbel_tau": 0.1556,
+      "image_masks_loss": 0.0184,
+      "llm_learning_rate": 4.6707318084544096e-07,
+      "loss": 0.7324,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 4.6707318084544095e-05,
+      "step": 7888
+    },
+    {
+      "epoch": 0.81,
+      "gumbel_tau": 0.1555,
+      "image_masks_loss": 0.0288,
+      "llm_learning_rate": 4.66590434918775e-07,
+      "loss": 0.7099,
+      "output_text_mask_loss": 0.0038,
+      "predictor_learning_rate": 4.6659043491877493e-05,
+      "step": 7889
+    },
+    {
+      "epoch": 0.81,
+      "gumbel_tau": 0.1555,
+      "image_masks_loss": 0.0091,
+      "llm_learning_rate": 4.661079129156995e-07,
+      "loss": 0.6975,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 4.661079129156995e-05,
+      "step": 7890
+    },
+    {
+      "epoch": 0.81,
+      "gumbel_tau": 0.1555,
+      "image_masks_loss": 0.0155,
+      "llm_learning_rate": 4.656256148893509e-07,
+      "loss": 0.7054,
+      "output_text_mask_loss": 0.0231,
+      "predictor_learning_rate": 4.656256148893509e-05,
+      "step": 7891
+    },
+    {
+      "epoch": 0.81,
+      "gumbel_tau": 0.1554,
+      "image_masks_loss": 0.0107,
+      "llm_learning_rate": 4.651435408928412e-07,
+      "loss": 0.7278,
+      "output_text_mask_loss": 0.0078,
+      "predictor_learning_rate": 4.6514354089284114e-05,
+      "step": 7892
+    },
+    {
+      "epoch": 0.81,
+      "gumbel_tau": 0.1554,
+      "image_masks_loss": 0.0237,
+      "llm_learning_rate": 4.646616909792573e-07,
+      "loss": 0.7665,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 4.6466169097925724e-05,
+      "step": 7893
+    },
+    {
+      "epoch": 0.81,
+      "gumbel_tau": 0.1553,
+      "image_masks_loss": 0.0084,
+      "llm_learning_rate": 4.641800652016612e-07,
+      "loss": 0.7122,
+      "output_text_mask_loss": 0.0042,
+      "predictor_learning_rate": 4.6418006520166116e-05,
+      "step": 7894
+    },
+    {
+      "epoch": 0.81,
+      "gumbel_tau": 0.1553,
+      "image_masks_loss": 0.0103,
+      "llm_learning_rate": 4.636986636130905e-07,
+      "loss": 0.6797,
+      "output_text_mask_loss": 0.0037,
+      "predictor_learning_rate": 4.636986636130905e-05,
+      "step": 7895
+    },
+    {
+      "epoch": 0.81,
+      "gumbel_tau": 0.1553,
+      "image_masks_loss": 0.0037,
+      "llm_learning_rate": 4.6321748626655923e-07,
+      "loss": 0.7271,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 4.632174862665592e-05,
+      "step": 7896
+    },
+    {
+      "epoch": 0.81,
+      "gumbel_tau": 0.1552,
+      "image_masks_loss": 0.0066,
+      "llm_learning_rate": 4.627365332150552e-07,
+      "loss": 0.7146,
+      "output_text_mask_loss": 0.0095,
+      "predictor_learning_rate": 4.627365332150551e-05,
+      "step": 7897
+    },
+    {
+      "epoch": 0.81,
+      "gumbel_tau": 0.1552,
+      "image_masks_loss": 0.0177,
+      "llm_learning_rate": 4.6225580451154187e-07,
+      "loss": 0.7041,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 4.622558045115419e-05,
+      "step": 7898
+    },
+    {
+      "epoch": 0.81,
+      "gumbel_tau": 0.1552,
+      "image_masks_loss": 0.0165,
+      "llm_learning_rate": 4.6177530020895895e-07,
+      "loss": 0.7148,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 4.6177530020895895e-05,
+      "step": 7899
+    },
+    {
+      "epoch": 0.81,
+      "gumbel_tau": 0.1551,
+      "image_masks_loss": 0.0041,
+      "llm_learning_rate": 4.6129502036021926e-07,
+      "loss": 0.7443,
+      "output_text_mask_loss": 0.008,
+      "predictor_learning_rate": 4.6129502036021924e-05,
+      "step": 7900
+    },
+    {
+      "epoch": 0.81,
+      "gumbel_tau": 0.1551,
+      "image_masks_loss": 0.0067,
+      "llm_learning_rate": 4.608149650182123e-07,
+      "loss": 0.7735,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 4.6081496501821225e-05,
+      "step": 7901
+    },
+    {
+      "epoch": 0.81,
+      "gumbel_tau": 0.155,
+      "image_masks_loss": 0.0113,
+      "llm_learning_rate": 4.603351342358042e-07,
+      "loss": 0.6544,
+      "output_text_mask_loss": 0.002,
+      "predictor_learning_rate": 4.603351342358042e-05,
+      "step": 7902
+    },
+    {
+      "epoch": 0.81,
+      "gumbel_tau": 0.155,
+      "image_masks_loss": 0.008,
+      "llm_learning_rate": 4.598555280658343e-07,
+      "loss": 0.6723,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 4.598555280658342e-05,
+      "step": 7903
+    },
+    {
+      "epoch": 0.81,
+      "gumbel_tau": 0.155,
+      "image_masks_loss": 0.011,
+      "llm_learning_rate": 4.5937614656111795e-07,
+      "loss": 0.7437,
+      "output_text_mask_loss": 0.0172,
+      "predictor_learning_rate": 4.593761465611179e-05,
+      "step": 7904
+    },
+    {
+      "epoch": 0.81,
+      "gumbel_tau": 0.1549,
+      "image_masks_loss": 0.0206,
+      "llm_learning_rate": 4.588969897744455e-07,
+      "loss": 0.7049,
+      "output_text_mask_loss": 0.0058,
+      "predictor_learning_rate": 4.588969897744455e-05,
+      "step": 7905
+    },
+    {
+      "epoch": 0.81,
+      "gumbel_tau": 0.1549,
+      "image_masks_loss": 0.0107,
+      "llm_learning_rate": 4.5841805775858285e-07,
+      "loss": 0.7276,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 4.584180577585828e-05,
+      "step": 7906
+    },
+    {
+      "epoch": 0.81,
+      "gumbel_tau": 0.1549,
+      "image_masks_loss": 0.0125,
+      "llm_learning_rate": 4.5793935056627134e-07,
+      "loss": 0.8029,
+      "output_text_mask_loss": 0.0391,
+      "predictor_learning_rate": 4.5793935056627125e-05,
+      "step": 7907
+    },
+    {
+      "epoch": 0.81,
+      "gumbel_tau": 0.1548,
+      "image_masks_loss": 0.0075,
+      "llm_learning_rate": 4.574608682502263e-07,
+      "loss": 0.6611,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 4.574608682502263e-05,
+      "step": 7908
+    },
+    {
+      "epoch": 0.81,
+      "gumbel_tau": 0.1548,
+      "image_masks_loss": 0.0053,
+      "llm_learning_rate": 4.569826108631406e-07,
+      "loss": 0.7809,
+      "output_text_mask_loss": 0.0126,
+      "predictor_learning_rate": 4.5698261086314056e-05,
+      "step": 7909
+    },
+    {
+      "epoch": 0.81,
+      "gumbel_tau": 0.1548,
+      "image_masks_loss": 0.0143,
+      "llm_learning_rate": 4.565045784576808e-07,
+      "loss": 0.747,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 4.565045784576807e-05,
+      "step": 7910
+    },
+    {
+      "epoch": 0.81,
+      "gumbel_tau": 0.1547,
+      "image_masks_loss": 0.0052,
+      "llm_learning_rate": 4.5602677108648823e-07,
+      "loss": 0.7016,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 4.560267710864882e-05,
+      "step": 7911
+    },
+    {
+      "epoch": 0.81,
+      "gumbel_tau": 0.1547,
+      "image_masks_loss": 0.0164,
+      "llm_learning_rate": 4.5554918880218043e-07,
+      "loss": 0.7113,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 4.555491888021804e-05,
+      "step": 7912
+    },
+    {
+      "epoch": 0.81,
+      "gumbel_tau": 0.1546,
+      "image_masks_loss": 0.0135,
+      "llm_learning_rate": 4.5507183165734964e-07,
+      "loss": 0.8133,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 4.550718316573496e-05,
+      "step": 7913
+    },
+    {
+      "epoch": 0.81,
+      "gumbel_tau": 0.1546,
+      "image_masks_loss": 0.0148,
+      "llm_learning_rate": 4.5459469970456385e-07,
+      "loss": 0.7757,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 4.5459469970456384e-05,
+      "step": 7914
+    },
+    {
+      "epoch": 0.81,
+      "gumbel_tau": 0.1546,
+      "image_masks_loss": 0.0058,
+      "llm_learning_rate": 4.54117792996365e-07,
+      "loss": 0.7518,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 4.5411779299636496e-05,
+      "step": 7915
+    },
+    {
+      "epoch": 0.81,
+      "gumbel_tau": 0.1545,
+      "image_masks_loss": 0.0106,
+      "llm_learning_rate": 4.5364111158527233e-07,
+      "loss": 0.7247,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 4.5364111158527236e-05,
+      "step": 7916
+    },
+    {
+      "epoch": 0.81,
+      "gumbel_tau": 0.1545,
+      "image_masks_loss": 0.0163,
+      "llm_learning_rate": 4.531646555237795e-07,
+      "loss": 0.7513,
+      "output_text_mask_loss": 0.0071,
+      "predictor_learning_rate": 4.531646555237795e-05,
+      "step": 7917
+    },
+    {
+      "epoch": 0.81,
+      "gumbel_tau": 0.1545,
+      "image_masks_loss": 0.0101,
+      "llm_learning_rate": 4.526884248643529e-07,
+      "loss": 0.6865,
+      "output_text_mask_loss": 0.0132,
+      "predictor_learning_rate": 4.5268842486435285e-05,
+      "step": 7918
+    },
+    {
+      "epoch": 0.81,
+      "gumbel_tau": 0.1544,
+      "image_masks_loss": 0.011,
+      "llm_learning_rate": 4.522124196594374e-07,
+      "loss": 0.6852,
+      "output_text_mask_loss": 0.0023,
+      "predictor_learning_rate": 4.522124196594374e-05,
+      "step": 7919
+    },
+    {
+      "epoch": 0.81,
+      "gumbel_tau": 0.1544,
+      "image_masks_loss": 0.0085,
+      "llm_learning_rate": 4.5173663996145146e-07,
+      "loss": 0.7887,
+      "output_text_mask_loss": 0.0071,
+      "predictor_learning_rate": 4.5173663996145145e-05,
+      "step": 7920
+    },
+    {
+      "epoch": 0.81,
+      "gumbel_tau": 0.1544,
+      "image_masks_loss": 0.0054,
+      "llm_learning_rate": 4.5126108582278864e-07,
+      "loss": 0.7758,
+      "output_text_mask_loss": 0.0143,
+      "predictor_learning_rate": 4.5126108582278864e-05,
+      "step": 7921
+    },
+    {
+      "epoch": 0.81,
+      "gumbel_tau": 0.1543,
+      "image_masks_loss": 0.0051,
+      "llm_learning_rate": 4.507857572958188e-07,
+      "loss": 0.7735,
+      "output_text_mask_loss": 0.0018,
+      "predictor_learning_rate": 4.507857572958188e-05,
+      "step": 7922
+    },
+    {
+      "epoch": 0.81,
+      "gumbel_tau": 0.1543,
+      "image_masks_loss": 0.0112,
+      "llm_learning_rate": 4.503106544328858e-07,
+      "loss": 0.6748,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 4.503106544328858e-05,
+      "step": 7923
+    },
+    {
+      "epoch": 0.81,
+      "gumbel_tau": 0.1542,
+      "image_masks_loss": 0.0124,
+      "llm_learning_rate": 4.4983577728630936e-07,
+      "loss": 0.7319,
+      "output_text_mask_loss": 0.0202,
+      "predictor_learning_rate": 4.498357772863093e-05,
+      "step": 7924
+    },
+    {
+      "epoch": 0.81,
+      "gumbel_tau": 0.1542,
+      "image_masks_loss": 0.0239,
+      "llm_learning_rate": 4.493611259083833e-07,
+      "loss": 0.7535,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 4.493611259083832e-05,
+      "step": 7925
+    },
+    {
+      "epoch": 0.81,
+      "gumbel_tau": 0.1542,
+      "image_masks_loss": 0.0145,
+      "llm_learning_rate": 4.488867003513778e-07,
+      "loss": 0.7007,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 4.488867003513777e-05,
+      "step": 7926
+    },
+    {
+      "epoch": 0.81,
+      "gumbel_tau": 0.1541,
+      "image_masks_loss": 0.0112,
+      "llm_learning_rate": 4.484125006675377e-07,
+      "loss": 0.6827,
+      "output_text_mask_loss": 0.0275,
+      "predictor_learning_rate": 4.484125006675377e-05,
+      "step": 7927
+    },
+    {
+      "epoch": 0.81,
+      "gumbel_tau": 0.1541,
+      "image_masks_loss": 0.0024,
+      "llm_learning_rate": 4.479385269090819e-07,
+      "loss": 0.7318,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 4.479385269090819e-05,
+      "step": 7928
+    },
+    {
+      "epoch": 0.81,
+      "gumbel_tau": 0.1541,
+      "image_masks_loss": 0.0118,
+      "llm_learning_rate": 4.474647791282069e-07,
+      "loss": 0.7328,
+      "output_text_mask_loss": 0.0221,
+      "predictor_learning_rate": 4.474647791282069e-05,
+      "step": 7929
+    },
+    {
+      "epoch": 0.81,
+      "gumbel_tau": 0.154,
+      "image_masks_loss": 0.0093,
+      "llm_learning_rate": 4.4699125737708267e-07,
+      "loss": 0.7754,
+      "output_text_mask_loss": 0.0248,
+      "predictor_learning_rate": 4.4699125737708264e-05,
+      "step": 7930
+    },
+    {
+      "epoch": 0.81,
+      "gumbel_tau": 0.154,
+      "image_masks_loss": 0.0164,
+      "llm_learning_rate": 4.4651796170785354e-07,
+      "loss": 0.7109,
+      "output_text_mask_loss": 0.0058,
+      "predictor_learning_rate": 4.465179617078535e-05,
+      "step": 7931
+    },
+    {
+      "epoch": 0.81,
+      "gumbel_tau": 0.154,
+      "image_masks_loss": 0.0197,
+      "llm_learning_rate": 4.4604489217264067e-07,
+      "loss": 0.8083,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 4.4604489217264066e-05,
+      "step": 7932
+    },
+    {
+      "epoch": 0.81,
+      "gumbel_tau": 0.1539,
+      "image_masks_loss": 0.0045,
+      "llm_learning_rate": 4.4557204882353933e-07,
+      "loss": 0.7731,
+      "output_text_mask_loss": 0.001,
+      "predictor_learning_rate": 4.455720488235393e-05,
+      "step": 7933
+    },
+    {
+      "epoch": 0.81,
+      "gumbel_tau": 0.1539,
+      "image_masks_loss": 0.0133,
+      "llm_learning_rate": 4.4509943171262e-07,
+      "loss": 0.7184,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 4.4509943171262e-05,
+      "step": 7934
+    },
+    {
+      "epoch": 0.81,
+      "gumbel_tau": 0.1538,
+      "image_masks_loss": 0.0066,
+      "llm_learning_rate": 4.4462704089192785e-07,
+      "loss": 0.7779,
+      "output_text_mask_loss": 0.0097,
+      "predictor_learning_rate": 4.446270408919279e-05,
+      "step": 7935
+    },
+    {
+      "epoch": 0.81,
+      "gumbel_tau": 0.1538,
+      "image_masks_loss": 0.003,
+      "llm_learning_rate": 4.441548764134851e-07,
+      "loss": 0.8447,
+      "output_text_mask_loss": 0.0034,
+      "predictor_learning_rate": 4.4415487641348506e-05,
+      "step": 7936
+    },
+    {
+      "epoch": 0.81,
+      "gumbel_tau": 0.1538,
+      "image_masks_loss": 0.0064,
+      "llm_learning_rate": 4.436829383292862e-07,
+      "loss": 0.6595,
+      "output_text_mask_loss": 0.0014,
+      "predictor_learning_rate": 4.4368293832928616e-05,
+      "step": 7937
+    },
+    {
+      "epoch": 0.81,
+      "gumbel_tau": 0.1537,
+      "image_masks_loss": 0.0102,
+      "llm_learning_rate": 4.432112266913022e-07,
+      "loss": 0.6895,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 4.4321122669130216e-05,
+      "step": 7938
+    },
+    {
+      "epoch": 0.81,
+      "gumbel_tau": 0.1537,
+      "image_masks_loss": 0.0118,
+      "llm_learning_rate": 4.4273974155147945e-07,
+      "loss": 0.7469,
+      "output_text_mask_loss": 0.0006,
+      "predictor_learning_rate": 4.427397415514795e-05,
+      "step": 7939
+    },
+    {
+      "epoch": 0.81,
+      "gumbel_tau": 0.1537,
+      "image_masks_loss": 0.0076,
+      "llm_learning_rate": 4.42268482961738e-07,
+      "loss": 0.7343,
+      "output_text_mask_loss": 0.021,
+      "predictor_learning_rate": 4.42268482961738e-05,
+      "step": 7940
+    },
+    {
+      "epoch": 0.81,
+      "gumbel_tau": 0.1536,
+      "image_masks_loss": 0.0072,
+      "llm_learning_rate": 4.4179745097397504e-07,
+      "loss": 0.6771,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 4.41797450973975e-05,
+      "step": 7941
+    },
+    {
+      "epoch": 0.81,
+      "gumbel_tau": 0.1536,
+      "image_masks_loss": 0.009,
+      "llm_learning_rate": 4.4132664564006167e-07,
+      "loss": 0.7588,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 4.4132664564006165e-05,
+      "step": 7942
+    },
+    {
+      "epoch": 0.81,
+      "gumbel_tau": 0.1536,
+      "image_masks_loss": 0.0362,
+      "llm_learning_rate": 4.4085606701184334e-07,
+      "loss": 0.8492,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 4.4085606701184336e-05,
+      "step": 7943
+    },
+    {
+      "epoch": 0.81,
+      "gumbel_tau": 0.1535,
+      "image_masks_loss": 0.011,
+      "llm_learning_rate": 4.403857151411414e-07,
+      "loss": 0.6967,
+      "output_text_mask_loss": 0.0587,
+      "predictor_learning_rate": 4.403857151411414e-05,
+      "step": 7944
+    },
+    {
+      "epoch": 0.81,
+      "gumbel_tau": 0.1535,
+      "image_masks_loss": 0.0074,
+      "llm_learning_rate": 4.399155900797522e-07,
+      "loss": 0.6614,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 4.399155900797522e-05,
+      "step": 7945
+    },
+    {
+      "epoch": 0.81,
+      "gumbel_tau": 0.1534,
+      "image_masks_loss": 0.0082,
+      "llm_learning_rate": 4.394456918794468e-07,
+      "loss": 0.6491,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 4.394456918794468e-05,
+      "step": 7946
+    },
+    {
+      "epoch": 0.81,
+      "gumbel_tau": 0.1534,
+      "image_masks_loss": 0.0079,
+      "llm_learning_rate": 4.3897602059197104e-07,
+      "loss": 0.7886,
+      "output_text_mask_loss": 0.0271,
+      "predictor_learning_rate": 4.38976020591971e-05,
+      "step": 7947
+    },
+    {
+      "epoch": 0.81,
+      "gumbel_tau": 0.1534,
+      "image_masks_loss": 0.0052,
+      "llm_learning_rate": 4.3850657626904715e-07,
+      "loss": 0.7388,
+      "output_text_mask_loss": 0.0162,
+      "predictor_learning_rate": 4.385065762690471e-05,
+      "step": 7948
+    },
+    {
+      "epoch": 0.81,
+      "gumbel_tau": 0.1533,
+      "image_masks_loss": 0.0154,
+      "llm_learning_rate": 4.38037358962371e-07,
+      "loss": 0.7306,
+      "output_text_mask_loss": 0.0483,
+      "predictor_learning_rate": 4.3803735896237096e-05,
+      "step": 7949
+    },
+    {
+      "epoch": 0.81,
+      "gumbel_tau": 0.1533,
+      "image_masks_loss": 0.0089,
+      "llm_learning_rate": 4.3756836872361376e-07,
+      "loss": 0.6507,
+      "output_text_mask_loss": 0.0021,
+      "predictor_learning_rate": 4.375683687236137e-05,
+      "step": 7950
+    },
+    {
+      "epoch": 0.81,
+      "gumbel_tau": 0.1533,
+      "image_masks_loss": 0.0134,
+      "llm_learning_rate": 4.370996056044216e-07,
+      "loss": 0.7248,
+      "output_text_mask_loss": 0.0015,
+      "predictor_learning_rate": 4.370996056044216e-05,
+      "step": 7951
+    },
+    {
+      "epoch": 0.81,
+      "gumbel_tau": 0.1532,
+      "image_masks_loss": 0.0125,
+      "llm_learning_rate": 4.366310696564158e-07,
+      "loss": 0.7226,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 4.366310696564158e-05,
+      "step": 7952
+    },
+    {
+      "epoch": 0.81,
+      "gumbel_tau": 0.1532,
+      "image_masks_loss": 0.0101,
+      "llm_learning_rate": 4.3616276093119246e-07,
+      "loss": 0.6681,
+      "output_text_mask_loss": 0.0162,
+      "predictor_learning_rate": 4.3616276093119243e-05,
+      "step": 7953
+    },
+    {
+      "epoch": 0.81,
+      "gumbel_tau": 0.1532,
+      "image_masks_loss": 0.0151,
+      "llm_learning_rate": 4.3569467948032294e-07,
+      "loss": 0.7769,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 4.3569467948032295e-05,
+      "step": 7954
+    },
+    {
+      "epoch": 0.82,
+      "gumbel_tau": 0.1531,
+      "image_masks_loss": 0.0081,
+      "llm_learning_rate": 4.352268253553535e-07,
+      "loss": 0.6469,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 4.352268253553535e-05,
+      "step": 7955
+    },
+    {
+      "epoch": 0.82,
+      "gumbel_tau": 0.1531,
+      "image_masks_loss": 0.0116,
+      "llm_learning_rate": 4.347591986078048e-07,
+      "loss": 0.7647,
+      "output_text_mask_loss": 0.0221,
+      "predictor_learning_rate": 4.3475919860780475e-05,
+      "step": 7956
+    },
+    {
+      "epoch": 0.82,
+      "gumbel_tau": 0.1531,
+      "image_masks_loss": 0.0035,
+      "llm_learning_rate": 4.3429179928917373e-07,
+      "loss": 0.7505,
+      "output_text_mask_loss": 0.0076,
+      "predictor_learning_rate": 4.3429179928917374e-05,
+      "step": 7957
+    },
+    {
+      "epoch": 0.82,
+      "gumbel_tau": 0.153,
+      "image_masks_loss": 0.0199,
+      "llm_learning_rate": 4.338246274509306e-07,
+      "loss": 0.6718,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 4.338246274509306e-05,
+      "step": 7958
+    },
+    {
+      "epoch": 0.82,
+      "gumbel_tau": 0.153,
+      "image_masks_loss": 0.0257,
+      "llm_learning_rate": 4.3335768314452184e-07,
+      "loss": 0.7279,
+      "output_text_mask_loss": 0.005,
+      "predictor_learning_rate": 4.3335768314452185e-05,
+      "step": 7959
+    },
+    {
+      "epoch": 0.82,
+      "gumbel_tau": 0.1529,
+      "image_masks_loss": 0.0058,
+      "llm_learning_rate": 4.328909664213679e-07,
+      "loss": 0.7494,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 4.328909664213679e-05,
+      "step": 7960
+    },
+    {
+      "epoch": 0.82,
+      "gumbel_tau": 0.1529,
+      "image_masks_loss": 0.0061,
+      "llm_learning_rate": 4.3242447733286536e-07,
+      "loss": 0.5738,
+      "output_text_mask_loss": 0.0014,
+      "predictor_learning_rate": 4.324244773328653e-05,
+      "step": 7961
+    },
+    {
+      "epoch": 0.82,
+      "gumbel_tau": 0.1529,
+      "image_masks_loss": 0.008,
+      "llm_learning_rate": 4.319582159303848e-07,
+      "loss": 0.7814,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 4.319582159303848e-05,
+      "step": 7962
+    },
+    {
+      "epoch": 0.82,
+      "gumbel_tau": 0.1528,
+      "image_masks_loss": 0.0161,
+      "llm_learning_rate": 4.314921822652718e-07,
+      "loss": 0.8222,
+      "output_text_mask_loss": 0.0261,
+      "predictor_learning_rate": 4.314921822652718e-05,
+      "step": 7963
+    },
+    {
+      "epoch": 0.82,
+      "gumbel_tau": 0.1528,
+      "image_masks_loss": 0.0178,
+      "llm_learning_rate": 4.3102637638884725e-07,
+      "loss": 0.7248,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 4.310263763888472e-05,
+      "step": 7964
+    },
+    {
+      "epoch": 0.82,
+      "gumbel_tau": 0.1528,
+      "image_masks_loss": 0.02,
+      "llm_learning_rate": 4.3056079835240657e-07,
+      "loss": 0.6836,
+      "output_text_mask_loss": 0.0111,
+      "predictor_learning_rate": 4.3056079835240655e-05,
+      "step": 7965
+    },
+    {
+      "epoch": 0.82,
+      "gumbel_tau": 0.1527,
+      "image_masks_loss": 0.0168,
+      "llm_learning_rate": 4.3009544820722007e-07,
+      "loss": 0.7474,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 4.3009544820722e-05,
+      "step": 7966
+    },
+    {
+      "epoch": 0.82,
+      "gumbel_tau": 0.1527,
+      "image_masks_loss": 0.0143,
+      "llm_learning_rate": 4.296303260045329e-07,
+      "loss": 0.7918,
+      "output_text_mask_loss": 0.0079,
+      "predictor_learning_rate": 4.2963032600453284e-05,
+      "step": 7967
+    },
+    {
+      "epoch": 0.82,
+      "gumbel_tau": 0.1527,
+      "image_masks_loss": 0.0045,
+      "llm_learning_rate": 4.291654317955665e-07,
+      "loss": 0.7508,
+      "output_text_mask_loss": 0.0032,
+      "predictor_learning_rate": 4.291654317955665e-05,
+      "step": 7968
+    },
+    {
+      "epoch": 0.82,
+      "gumbel_tau": 0.1526,
+      "image_masks_loss": 0.02,
+      "llm_learning_rate": 4.2870076563151536e-07,
+      "loss": 0.7417,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 4.287007656315153e-05,
+      "step": 7969
+    },
+    {
+      "epoch": 0.82,
+      "gumbel_tau": 0.1526,
+      "image_masks_loss": 0.0103,
+      "llm_learning_rate": 4.282363275635493e-07,
+      "loss": 0.8024,
+      "output_text_mask_loss": 0.0368,
+      "predictor_learning_rate": 4.282363275635492e-05,
+      "step": 7970
+    },
+    {
+      "epoch": 0.82,
+      "gumbel_tau": 0.1525,
+      "image_masks_loss": 0.0218,
+      "llm_learning_rate": 4.2777211764281376e-07,
+      "loss": 0.7518,
+      "output_text_mask_loss": 0.0115,
+      "predictor_learning_rate": 4.277721176428137e-05,
+      "step": 7971
+    },
+    {
+      "epoch": 0.82,
+      "gumbel_tau": 0.1525,
+      "image_masks_loss": 0.0179,
+      "llm_learning_rate": 4.2730813592042813e-07,
+      "loss": 0.7501,
+      "output_text_mask_loss": 0.0015,
+      "predictor_learning_rate": 4.273081359204281e-05,
+      "step": 7972
+    },
+    {
+      "epoch": 0.82,
+      "gumbel_tau": 0.1525,
+      "image_masks_loss": 0.0075,
+      "llm_learning_rate": 4.2684438244748744e-07,
+      "loss": 0.8467,
+      "output_text_mask_loss": 0.0195,
+      "predictor_learning_rate": 4.268443824474874e-05,
+      "step": 7973
+    },
+    {
+      "epoch": 0.82,
+      "gumbel_tau": 0.1524,
+      "image_masks_loss": 0.0107,
+      "llm_learning_rate": 4.2638085727506105e-07,
+      "loss": 0.7261,
+      "output_text_mask_loss": 0.0406,
+      "predictor_learning_rate": 4.26380857275061e-05,
+      "step": 7974
+    },
+    {
+      "epoch": 0.82,
+      "gumbel_tau": 0.1524,
+      "image_masks_loss": 0.0039,
+      "llm_learning_rate": 4.2591756045419345e-07,
+      "loss": 0.7079,
+      "output_text_mask_loss": 0.0102,
+      "predictor_learning_rate": 4.2591756045419344e-05,
+      "step": 7975
+    },
+    {
+      "epoch": 0.82,
+      "gumbel_tau": 0.1524,
+      "image_masks_loss": 0.0202,
+      "llm_learning_rate": 4.254544920359041e-07,
+      "loss": 0.719,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 4.254544920359041e-05,
+      "step": 7976
+    },
+    {
+      "epoch": 0.82,
+      "gumbel_tau": 0.1523,
+      "image_masks_loss": 0.0123,
+      "llm_learning_rate": 4.2499165207118693e-07,
+      "loss": 0.7458,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 4.249916520711869e-05,
+      "step": 7977
+    },
+    {
+      "epoch": 0.82,
+      "gumbel_tau": 0.1523,
+      "image_masks_loss": 0.0097,
+      "llm_learning_rate": 4.2452904061101083e-07,
+      "loss": 0.6628,
+      "output_text_mask_loss": 0.0259,
+      "predictor_learning_rate": 4.245290406110108e-05,
+      "step": 7978
+    },
+    {
+      "epoch": 0.82,
+      "gumbel_tau": 0.1523,
+      "image_masks_loss": 0.0054,
+      "llm_learning_rate": 4.240666577063191e-07,
+      "loss": 0.6357,
+      "output_text_mask_loss": 0.0107,
+      "predictor_learning_rate": 4.24066657706319e-05,
+      "step": 7979
+    },
+    {
+      "epoch": 0.82,
+      "gumbel_tau": 0.1522,
+      "image_masks_loss": 0.0296,
+      "llm_learning_rate": 4.236045034080316e-07,
+      "loss": 0.7386,
+      "output_text_mask_loss": 0.0345,
+      "predictor_learning_rate": 4.2360450340803156e-05,
+      "step": 7980
+    },
+    {
+      "epoch": 0.82,
+      "gumbel_tau": 0.1522,
+      "image_masks_loss": 0.0101,
+      "llm_learning_rate": 4.2314257776704116e-07,
+      "loss": 0.6495,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 4.231425777670411e-05,
+      "step": 7981
+    },
+    {
+      "epoch": 0.82,
+      "gumbel_tau": 0.1522,
+      "image_masks_loss": 0.0192,
+      "llm_learning_rate": 4.2268088083421616e-07,
+      "loss": 0.8183,
+      "output_text_mask_loss": 0.0196,
+      "predictor_learning_rate": 4.2268088083421615e-05,
+      "step": 7982
+    },
+    {
+      "epoch": 0.82,
+      "gumbel_tau": 0.1521,
+      "image_masks_loss": 0.01,
+      "llm_learning_rate": 4.222194126603996e-07,
+      "loss": 0.7077,
+      "output_text_mask_loss": 0.0423,
+      "predictor_learning_rate": 4.222194126603995e-05,
+      "step": 7983
+    },
+    {
+      "epoch": 0.82,
+      "gumbel_tau": 0.1521,
+      "image_masks_loss": 0.0092,
+      "llm_learning_rate": 4.217581732964093e-07,
+      "loss": 0.7011,
+      "output_text_mask_loss": 0.0038,
+      "predictor_learning_rate": 4.2175817329640934e-05,
+      "step": 7984
+    },
+    {
+      "epoch": 0.82,
+      "gumbel_tau": 0.152,
+      "image_masks_loss": 0.0121,
+      "llm_learning_rate": 4.2129716279303817e-07,
+      "loss": 0.7284,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 4.212971627930382e-05,
+      "step": 7985
+    },
+    {
+      "epoch": 0.82,
+      "gumbel_tau": 0.152,
+      "image_masks_loss": 0.0115,
+      "llm_learning_rate": 4.2083638120105296e-07,
+      "loss": 0.6975,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 4.2083638120105296e-05,
+      "step": 7986
+    },
+    {
+      "epoch": 0.82,
+      "gumbel_tau": 0.152,
+      "image_masks_loss": 0.0094,
+      "llm_learning_rate": 4.203758285711973e-07,
+      "loss": 0.735,
+      "output_text_mask_loss": 0.0179,
+      "predictor_learning_rate": 4.203758285711973e-05,
+      "step": 7987
+    },
+    {
+      "epoch": 0.82,
+      "gumbel_tau": 0.1519,
+      "image_masks_loss": 0.0309,
+      "llm_learning_rate": 4.1991550495418723e-07,
+      "loss": 0.7188,
+      "output_text_mask_loss": 0.0252,
+      "predictor_learning_rate": 4.1991550495418725e-05,
+      "step": 7988
+    },
+    {
+      "epoch": 0.82,
+      "gumbel_tau": 0.1519,
+      "image_masks_loss": 0.0124,
+      "llm_learning_rate": 4.1945541040071535e-07,
+      "loss": 0.7653,
+      "output_text_mask_loss": 0.0099,
+      "predictor_learning_rate": 4.1945541040071535e-05,
+      "step": 7989
+    },
+    {
+      "epoch": 0.82,
+      "gumbel_tau": 0.1519,
+      "image_masks_loss": 0.025,
+      "llm_learning_rate": 4.1899554496144774e-07,
+      "loss": 0.8018,
+      "output_text_mask_loss": 0.0322,
+      "predictor_learning_rate": 4.1899554496144774e-05,
+      "step": 7990
+    },
+    {
+      "epoch": 0.82,
+      "gumbel_tau": 0.1518,
+      "image_masks_loss": 0.0073,
+      "llm_learning_rate": 4.1853590868702664e-07,
+      "loss": 0.6507,
+      "output_text_mask_loss": 0.0004,
+      "predictor_learning_rate": 4.185359086870266e-05,
+      "step": 7991
+    },
+    {
+      "epoch": 0.82,
+      "gumbel_tau": 0.1518,
+      "image_masks_loss": 0.0226,
+      "llm_learning_rate": 4.1807650162806636e-07,
+      "loss": 0.6804,
+      "output_text_mask_loss": 0.0107,
+      "predictor_learning_rate": 4.180765016280663e-05,
+      "step": 7992
+    },
+    {
+      "epoch": 0.82,
+      "gumbel_tau": 0.1518,
+      "image_masks_loss": 0.0227,
+      "llm_learning_rate": 4.176173238351594e-07,
+      "loss": 0.7539,
+      "output_text_mask_loss": 0.0345,
+      "predictor_learning_rate": 4.176173238351594e-05,
+      "step": 7993
+    },
+    {
+      "epoch": 0.82,
+      "gumbel_tau": 0.1517,
+      "image_masks_loss": 0.0064,
+      "llm_learning_rate": 4.1715837535887123e-07,
+      "loss": 0.7149,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 4.1715837535887116e-05,
+      "step": 7994
+    },
+    {
+      "epoch": 0.82,
+      "gumbel_tau": 0.1517,
+      "image_masks_loss": 0.0169,
+      "llm_learning_rate": 4.166996562497422e-07,
+      "loss": 0.717,
+      "output_text_mask_loss": 0.005,
+      "predictor_learning_rate": 4.166996562497422e-05,
+      "step": 7995
+    },
+    {
+      "epoch": 0.82,
+      "gumbel_tau": 0.1516,
+      "image_masks_loss": 0.0163,
+      "llm_learning_rate": 4.1624116655828743e-07,
+      "loss": 0.7389,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 4.162411665582874e-05,
+      "step": 7996
+    },
+    {
+      "epoch": 0.82,
+      "gumbel_tau": 0.1516,
+      "image_masks_loss": 0.0066,
+      "llm_learning_rate": 4.1578290633499665e-07,
+      "loss": 0.7511,
+      "output_text_mask_loss": 0.037,
+      "predictor_learning_rate": 4.157829063349966e-05,
+      "step": 7997
+    },
+    {
+      "epoch": 0.82,
+      "gumbel_tau": 0.1516,
+      "image_masks_loss": 0.0188,
+      "llm_learning_rate": 4.1532487563033465e-07,
+      "loss": 0.7714,
+      "output_text_mask_loss": 0.0568,
+      "predictor_learning_rate": 4.153248756303346e-05,
+      "step": 7998
+    },
+    {
+      "epoch": 0.82,
+      "gumbel_tau": 0.1515,
+      "image_masks_loss": 0.0163,
+      "llm_learning_rate": 4.1486707449474e-07,
+      "loss": 0.6152,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 4.1486707449474e-05,
+      "step": 7999
+    },
+    {
+      "epoch": 0.82,
+      "gumbel_tau": 0.1515,
+      "image_masks_loss": 0.0129,
+      "llm_learning_rate": 4.144095029786285e-07,
+      "loss": 0.7395,
+      "output_text_mask_loss": 0.0496,
+      "predictor_learning_rate": 4.144095029786285e-05,
+      "step": 8000
+    },
+    {
+      "epoch": 0.82,
+      "gumbel_tau": 0.1515,
+      "image_masks_loss": 0.0412,
+      "llm_learning_rate": 4.1395216113238794e-07,
+      "loss": 0.7534,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 4.1395216113238795e-05,
+      "step": 8001
+    },
+    {
+      "epoch": 0.82,
+      "gumbel_tau": 0.1514,
+      "image_masks_loss": 0.0311,
+      "llm_learning_rate": 4.134950490063816e-07,
+      "loss": 0.7603,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 4.134950490063816e-05,
+      "step": 8002
+    },
+    {
+      "epoch": 0.82,
+      "gumbel_tau": 0.1514,
+      "image_masks_loss": 0.0061,
+      "llm_learning_rate": 4.130381666509484e-07,
+      "loss": 0.79,
+      "output_text_mask_loss": 0.0065,
+      "predictor_learning_rate": 4.130381666509484e-05,
+      "step": 8003
+    },
+    {
+      "epoch": 0.82,
+      "gumbel_tau": 0.1514,
+      "image_masks_loss": 0.0154,
+      "llm_learning_rate": 4.1258151411640053e-07,
+      "loss": 0.8268,
+      "output_text_mask_loss": 0.0233,
+      "predictor_learning_rate": 4.125815141164005e-05,
+      "step": 8004
+    },
+    {
+      "epoch": 0.82,
+      "gumbel_tau": 0.1513,
+      "image_masks_loss": 0.0103,
+      "llm_learning_rate": 4.121250914530259e-07,
+      "loss": 0.7157,
+      "output_text_mask_loss": 0.0671,
+      "predictor_learning_rate": 4.121250914530258e-05,
+      "step": 8005
+    },
+    {
+      "epoch": 0.82,
+      "gumbel_tau": 0.1513,
+      "image_masks_loss": 0.025,
+      "llm_learning_rate": 4.116688987110862e-07,
+      "loss": 0.8052,
+      "output_text_mask_loss": 0.0282,
+      "predictor_learning_rate": 4.1166889871108616e-05,
+      "step": 8006
+    },
+    {
+      "epoch": 0.82,
+      "gumbel_tau": 0.1513,
+      "image_masks_loss": 0.0357,
+      "llm_learning_rate": 4.112129359408193e-07,
+      "loss": 0.7258,
+      "output_text_mask_loss": 0.0107,
+      "predictor_learning_rate": 4.112129359408193e-05,
+      "step": 8007
+    },
+    {
+      "epoch": 0.82,
+      "gumbel_tau": 0.1512,
+      "image_masks_loss": 0.0378,
+      "llm_learning_rate": 4.1075720319243675e-07,
+      "loss": 0.7615,
+      "output_text_mask_loss": 0.1938,
+      "predictor_learning_rate": 4.107572031924367e-05,
+      "step": 8008
+    },
+    {
+      "epoch": 0.82,
+      "gumbel_tau": 0.1512,
+      "image_masks_loss": 0.0115,
+      "llm_learning_rate": 4.1030170051612474e-07,
+      "loss": 0.7012,
+      "output_text_mask_loss": 0.0198,
+      "predictor_learning_rate": 4.103017005161247e-05,
+      "step": 8009
+    },
+    {
+      "epoch": 0.82,
+      "gumbel_tau": 0.1511,
+      "image_masks_loss": 0.0071,
+      "llm_learning_rate": 4.098464279620443e-07,
+      "loss": 0.6662,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 4.098464279620443e-05,
+      "step": 8010
+    },
+    {
+      "epoch": 0.82,
+      "gumbel_tau": 0.1511,
+      "image_masks_loss": 0.0094,
+      "llm_learning_rate": 4.0939138558033e-07,
+      "loss": 0.7484,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 4.0939138558032996e-05,
+      "step": 8011
+    },
+    {
+      "epoch": 0.82,
+      "gumbel_tau": 0.1511,
+      "image_masks_loss": 0.0165,
+      "llm_learning_rate": 4.0893657342109345e-07,
+      "loss": 0.6749,
+      "output_text_mask_loss": 0.0046,
+      "predictor_learning_rate": 4.0893657342109343e-05,
+      "step": 8012
+    },
+    {
+      "epoch": 0.82,
+      "gumbel_tau": 0.151,
+      "image_masks_loss": 0.0311,
+      "llm_learning_rate": 4.08481991534419e-07,
+      "loss": 0.7105,
+      "output_text_mask_loss": 0.0124,
+      "predictor_learning_rate": 4.08481991534419e-05,
+      "step": 8013
+    },
+    {
+      "epoch": 0.82,
+      "gumbel_tau": 0.151,
+      "image_masks_loss": 0.0136,
+      "llm_learning_rate": 4.0802763997036644e-07,
+      "loss": 0.7184,
+      "output_text_mask_loss": 0.0187,
+      "predictor_learning_rate": 4.080276399703664e-05,
+      "step": 8014
+    },
+    {
+      "epoch": 0.82,
+      "gumbel_tau": 0.151,
+      "image_masks_loss": 0.0116,
+      "llm_learning_rate": 4.0757351877896977e-07,
+      "loss": 0.6572,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 4.0757351877896974e-05,
+      "step": 8015
+    },
+    {
+      "epoch": 0.82,
+      "gumbel_tau": 0.1509,
+      "image_masks_loss": 0.0101,
+      "llm_learning_rate": 4.0711962801023796e-07,
+      "loss": 0.7401,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 4.071196280102379e-05,
+      "step": 8016
+    },
+    {
+      "epoch": 0.82,
+      "gumbel_tau": 0.1509,
+      "image_masks_loss": 0.0125,
+      "llm_learning_rate": 4.0666596771415435e-07,
+      "loss": 0.6255,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 4.066659677141543e-05,
+      "step": 8017
+    },
+    {
+      "epoch": 0.82,
+      "gumbel_tau": 0.1509,
+      "image_masks_loss": 0.0069,
+      "llm_learning_rate": 4.0621253794067644e-07,
+      "loss": 0.7642,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 4.062125379406764e-05,
+      "step": 8018
+    },
+    {
+      "epoch": 0.82,
+      "gumbel_tau": 0.1508,
+      "image_masks_loss": 0.006,
+      "llm_learning_rate": 4.057593387397382e-07,
+      "loss": 0.7474,
+      "output_text_mask_loss": 0.0324,
+      "predictor_learning_rate": 4.057593387397382e-05,
+      "step": 8019
+    },
+    {
+      "epoch": 0.82,
+      "gumbel_tau": 0.1508,
+      "image_masks_loss": 0.0294,
+      "llm_learning_rate": 4.053063701612464e-07,
+      "loss": 0.7466,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 4.053063701612464e-05,
+      "step": 8020
+    },
+    {
+      "epoch": 0.82,
+      "gumbel_tau": 0.1508,
+      "image_masks_loss": 0.0094,
+      "llm_learning_rate": 4.0485363225508304e-07,
+      "loss": 0.7766,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 4.04853632255083e-05,
+      "step": 8021
+    },
+    {
+      "epoch": 0.82,
+      "gumbel_tau": 0.1507,
+      "image_masks_loss": 0.0122,
+      "llm_learning_rate": 4.0440112507110413e-07,
+      "loss": 0.8023,
+      "output_text_mask_loss": 0.0324,
+      "predictor_learning_rate": 4.044011250711041e-05,
+      "step": 8022
+    },
+    {
+      "epoch": 0.82,
+      "gumbel_tau": 0.1507,
+      "image_masks_loss": 0.019,
+      "llm_learning_rate": 4.039488486591414e-07,
+      "loss": 0.7185,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 4.0394884865914136e-05,
+      "step": 8023
+    },
+    {
+      "epoch": 0.82,
+      "gumbel_tau": 0.1507,
+      "image_masks_loss": 0.0067,
+      "llm_learning_rate": 4.034968030690001e-07,
+      "loss": 0.7804,
+      "output_text_mask_loss": 0.0098,
+      "predictor_learning_rate": 4.03496803069e-05,
+      "step": 8024
+    },
+    {
+      "epoch": 0.82,
+      "gumbel_tau": 0.1506,
+      "image_masks_loss": 0.0078,
+      "llm_learning_rate": 4.0304498835045975e-07,
+      "loss": 0.7517,
+      "output_text_mask_loss": 0.005,
+      "predictor_learning_rate": 4.030449883504597e-05,
+      "step": 8025
+    },
+    {
+      "epoch": 0.82,
+      "gumbel_tau": 0.1506,
+      "image_masks_loss": 0.0061,
+      "llm_learning_rate": 4.02593404553277e-07,
+      "loss": 0.7764,
+      "output_text_mask_loss": 0.0134,
+      "predictor_learning_rate": 4.0259340455327696e-05,
+      "step": 8026
+    },
+    {
+      "epoch": 0.82,
+      "gumbel_tau": 0.1505,
+      "image_masks_loss": 0.0137,
+      "llm_learning_rate": 4.021420517271801e-07,
+      "loss": 0.6961,
+      "output_text_mask_loss": 0.0023,
+      "predictor_learning_rate": 4.021420517271801e-05,
+      "step": 8027
+    },
+    {
+      "epoch": 0.82,
+      "gumbel_tau": 0.1505,
+      "image_masks_loss": 0.0103,
+      "llm_learning_rate": 4.016909299218741e-07,
+      "loss": 0.6664,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 4.0169092992187404e-05,
+      "step": 8028
+    },
+    {
+      "epoch": 0.82,
+      "gumbel_tau": 0.1505,
+      "image_masks_loss": 0.0064,
+      "llm_learning_rate": 4.012400391870361e-07,
+      "loss": 0.7098,
+      "output_text_mask_loss": 0.0116,
+      "predictor_learning_rate": 4.01240039187036e-05,
+      "step": 8029
+    },
+    {
+      "epoch": 0.82,
+      "gumbel_tau": 0.1504,
+      "image_masks_loss": 0.0121,
+      "llm_learning_rate": 4.0078937957231973e-07,
+      "loss": 0.6088,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 4.007893795723197e-05,
+      "step": 8030
+    },
+    {
+      "epoch": 0.82,
+      "gumbel_tau": 0.1504,
+      "image_masks_loss": 0.01,
+      "llm_learning_rate": 4.00338951127352e-07,
+      "loss": 0.7691,
+      "output_text_mask_loss": 0.0071,
+      "predictor_learning_rate": 4.0033895112735196e-05,
+      "step": 8031
+    },
+    {
+      "epoch": 0.82,
+      "gumbel_tau": 0.1504,
+      "image_masks_loss": 0.0267,
+      "llm_learning_rate": 3.9988875390173623e-07,
+      "loss": 0.723,
+      "output_text_mask_loss": 0.0038,
+      "predictor_learning_rate": 3.998887539017362e-05,
+      "step": 8032
+    },
+    {
+      "epoch": 0.82,
+      "gumbel_tau": 0.1503,
+      "image_masks_loss": 0.0104,
+      "llm_learning_rate": 3.99438787945049e-07,
+      "loss": 0.7377,
+      "output_text_mask_loss": 0.0347,
+      "predictor_learning_rate": 3.994387879450489e-05,
+      "step": 8033
+    },
+    {
+      "epoch": 0.82,
+      "gumbel_tau": 0.1503,
+      "image_masks_loss": 0.0056,
+      "llm_learning_rate": 3.989890533068411e-07,
+      "loss": 0.6649,
+      "output_text_mask_loss": 0.0004,
+      "predictor_learning_rate": 3.989890533068411e-05,
+      "step": 8034
+    },
+    {
+      "epoch": 0.82,
+      "gumbel_tau": 0.1503,
+      "image_masks_loss": 0.022,
+      "llm_learning_rate": 3.985395500366385e-07,
+      "loss": 0.7524,
+      "output_text_mask_loss": 0.0429,
+      "predictor_learning_rate": 3.985395500366384e-05,
+      "step": 8035
+    },
+    {
+      "epoch": 0.82,
+      "gumbel_tau": 0.1502,
+      "image_masks_loss": 0.0094,
+      "llm_learning_rate": 3.9809027818394156e-07,
+      "loss": 0.6859,
+      "output_text_mask_loss": 0.001,
+      "predictor_learning_rate": 3.980902781839416e-05,
+      "step": 8036
+    },
+    {
+      "epoch": 0.82,
+      "gumbel_tau": 0.1502,
+      "image_masks_loss": 0.0079,
+      "llm_learning_rate": 3.97641237798225e-07,
+      "loss": 0.7122,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 3.97641237798225e-05,
+      "step": 8037
+    },
+    {
+      "epoch": 0.82,
+      "gumbel_tau": 0.1502,
+      "image_masks_loss": 0.0122,
+      "llm_learning_rate": 3.971924289289375e-07,
+      "loss": 0.7267,
+      "output_text_mask_loss": 0.0153,
+      "predictor_learning_rate": 3.971924289289375e-05,
+      "step": 8038
+    },
+    {
+      "epoch": 0.82,
+      "gumbel_tau": 0.1501,
+      "image_masks_loss": 0.0116,
+      "llm_learning_rate": 3.96743851625504e-07,
+      "loss": 0.7705,
+      "output_text_mask_loss": 0.0086,
+      "predictor_learning_rate": 3.96743851625504e-05,
+      "step": 8039
+    },
+    {
+      "epoch": 0.82,
+      "gumbel_tau": 0.1501,
+      "image_masks_loss": 0.0104,
+      "llm_learning_rate": 3.9629550593732245e-07,
+      "loss": 0.7466,
+      "output_text_mask_loss": 0.0039,
+      "predictor_learning_rate": 3.9629550593732245e-05,
+      "step": 8040
+    },
+    {
+      "epoch": 0.82,
+      "gumbel_tau": 0.15,
+      "image_masks_loss": 0.0136,
+      "llm_learning_rate": 3.9584739191376564e-07,
+      "loss": 0.6753,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 3.9584739191376565e-05,
+      "step": 8041
+    },
+    {
+      "epoch": 0.82,
+      "gumbel_tau": 0.15,
+      "image_masks_loss": 0.0035,
+      "llm_learning_rate": 3.95399509604181e-07,
+      "loss": 0.7176,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 3.95399509604181e-05,
+      "step": 8042
+    },
+    {
+      "epoch": 0.82,
+      "gumbel_tau": 0.15,
+      "image_masks_loss": 0.0028,
+      "llm_learning_rate": 3.9495185905788993e-07,
+      "loss": 0.6976,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 3.949518590578899e-05,
+      "step": 8043
+    },
+    {
+      "epoch": 0.82,
+      "gumbel_tau": 0.1499,
+      "image_masks_loss": 0.0259,
+      "llm_learning_rate": 3.9450444032418905e-07,
+      "loss": 0.7775,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 3.94504440324189e-05,
+      "step": 8044
+    },
+    {
+      "epoch": 0.82,
+      "gumbel_tau": 0.1499,
+      "image_masks_loss": 0.009,
+      "llm_learning_rate": 3.940572534523485e-07,
+      "loss": 0.7578,
+      "output_text_mask_loss": 0.0809,
+      "predictor_learning_rate": 3.9405725345234844e-05,
+      "step": 8045
+    },
+    {
+      "epoch": 0.82,
+      "gumbel_tau": 0.1499,
+      "image_masks_loss": 0.0209,
+      "llm_learning_rate": 3.936102984916146e-07,
+      "loss": 0.7476,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 3.936102984916146e-05,
+      "step": 8046
+    },
+    {
+      "epoch": 0.82,
+      "gumbel_tau": 0.1498,
+      "image_masks_loss": 0.017,
+      "llm_learning_rate": 3.931635754912069e-07,
+      "loss": 0.6293,
+      "output_text_mask_loss": 0.017,
+      "predictor_learning_rate": 3.931635754912069e-05,
+      "step": 8047
+    },
+    {
+      "epoch": 0.82,
+      "gumbel_tau": 0.1498,
+      "image_masks_loss": 0.0205,
+      "llm_learning_rate": 3.9271708450031864e-07,
+      "loss": 0.7328,
+      "output_text_mask_loss": 0.0032,
+      "predictor_learning_rate": 3.9271708450031865e-05,
+      "step": 8048
+    },
+    {
+      "epoch": 0.82,
+      "gumbel_tau": 0.1498,
+      "image_masks_loss": 0.019,
+      "llm_learning_rate": 3.9227082556811874e-07,
+      "loss": 0.6609,
+      "output_text_mask_loss": 0.0231,
+      "predictor_learning_rate": 3.922708255681187e-05,
+      "step": 8049
+    },
+    {
+      "epoch": 0.82,
+      "gumbel_tau": 0.1497,
+      "image_masks_loss": 0.0129,
+      "llm_learning_rate": 3.9182479874374957e-07,
+      "loss": 0.6343,
+      "output_text_mask_loss": 0.0183,
+      "predictor_learning_rate": 3.9182479874374956e-05,
+      "step": 8050
+    },
+    {
+      "epoch": 0.82,
+      "gumbel_tau": 0.1497,
+      "image_masks_loss": 0.0154,
+      "llm_learning_rate": 3.913790040763299e-07,
+      "loss": 0.6761,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 3.913790040763299e-05,
+      "step": 8051
+    },
+    {
+      "epoch": 0.82,
+      "gumbel_tau": 0.1497,
+      "image_masks_loss": 0.0307,
+      "llm_learning_rate": 3.909334416149513e-07,
+      "loss": 0.6629,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 3.9093344161495126e-05,
+      "step": 8052
+    },
+    {
+      "epoch": 0.83,
+      "gumbel_tau": 0.1496,
+      "image_masks_loss": 0.0194,
+      "llm_learning_rate": 3.9048811140867977e-07,
+      "loss": 0.7354,
+      "output_text_mask_loss": 0.0248,
+      "predictor_learning_rate": 3.904881114086797e-05,
+      "step": 8053
+    },
+    {
+      "epoch": 0.83,
+      "gumbel_tau": 0.1496,
+      "image_masks_loss": 0.0224,
+      "llm_learning_rate": 3.900430135065564e-07,
+      "loss": 0.6888,
+      "output_text_mask_loss": 0.0076,
+      "predictor_learning_rate": 3.900430135065563e-05,
+      "step": 8054
+    },
+    {
+      "epoch": 0.83,
+      "gumbel_tau": 0.1496,
+      "image_masks_loss": 0.0134,
+      "llm_learning_rate": 3.8959814795759584e-07,
+      "loss": 0.7226,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 3.895981479575958e-05,
+      "step": 8055
+    },
+    {
+      "epoch": 0.83,
+      "gumbel_tau": 0.1495,
+      "image_masks_loss": 0.0218,
+      "llm_learning_rate": 3.8915351481078804e-07,
+      "loss": 0.7703,
+      "output_text_mask_loss": 0.0039,
+      "predictor_learning_rate": 3.89153514810788e-05,
+      "step": 8056
+    },
+    {
+      "epoch": 0.83,
+      "gumbel_tau": 0.1495,
+      "image_masks_loss": 0.008,
+      "llm_learning_rate": 3.8870911411509617e-07,
+      "loss": 0.673,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 3.887091141150961e-05,
+      "step": 8057
+    },
+    {
+      "epoch": 0.83,
+      "gumbel_tau": 0.1494,
+      "image_masks_loss": 0.0061,
+      "llm_learning_rate": 3.8826494591946014e-07,
+      "loss": 0.6441,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 3.882649459194601e-05,
+      "step": 8058
+    },
+    {
+      "epoch": 0.83,
+      "gumbel_tau": 0.1494,
+      "image_masks_loss": 0.0082,
+      "llm_learning_rate": 3.8782101027279194e-07,
+      "loss": 0.6264,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 3.8782101027279196e-05,
+      "step": 8059
+    },
+    {
+      "epoch": 0.83,
+      "gumbel_tau": 0.1494,
+      "image_masks_loss": 0.0117,
+      "llm_learning_rate": 3.873773072239789e-07,
+      "loss": 0.6134,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 3.873773072239789e-05,
+      "step": 8060
+    },
+    {
+      "epoch": 0.83,
+      "gumbel_tau": 0.1493,
+      "image_masks_loss": 0.0143,
+      "llm_learning_rate": 3.8693383682188224e-07,
+      "loss": 0.7367,
+      "output_text_mask_loss": 0.0435,
+      "predictor_learning_rate": 3.8693383682188225e-05,
+      "step": 8061
+    },
+    {
+      "epoch": 0.83,
+      "gumbel_tau": 0.1493,
+      "image_masks_loss": 0.0107,
+      "llm_learning_rate": 3.864905991153384e-07,
+      "loss": 0.7833,
+      "output_text_mask_loss": 0.0193,
+      "predictor_learning_rate": 3.864905991153383e-05,
+      "step": 8062
+    },
+    {
+      "epoch": 0.83,
+      "gumbel_tau": 0.1493,
+      "image_masks_loss": 0.005,
+      "llm_learning_rate": 3.8604759415315716e-07,
+      "loss": 0.6573,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 3.860475941531572e-05,
+      "step": 8063
+    },
+    {
+      "epoch": 0.83,
+      "gumbel_tau": 0.1492,
+      "image_masks_loss": 0.0191,
+      "llm_learning_rate": 3.85604821984123e-07,
+      "loss": 0.7483,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 3.856048219841229e-05,
+      "step": 8064
+    },
+    {
+      "epoch": 0.83,
+      "gumbel_tau": 0.1492,
+      "image_masks_loss": 0.0047,
+      "llm_learning_rate": 3.851622826569967e-07,
+      "loss": 0.7241,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 3.851622826569967e-05,
+      "step": 8065
+    },
+    {
+      "epoch": 0.83,
+      "gumbel_tau": 0.1492,
+      "image_masks_loss": 0.0165,
+      "llm_learning_rate": 3.8471997622050977e-07,
+      "loss": 0.7395,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 3.847199762205097e-05,
+      "step": 8066
+    },
+    {
+      "epoch": 0.83,
+      "gumbel_tau": 0.1491,
+      "image_masks_loss": 0.0118,
+      "llm_learning_rate": 3.842779027233709e-07,
+      "loss": 0.7469,
+      "output_text_mask_loss": 0.0025,
+      "predictor_learning_rate": 3.842779027233709e-05,
+      "step": 8067
+    },
+    {
+      "epoch": 0.83,
+      "gumbel_tau": 0.1491,
+      "image_masks_loss": 0.0132,
+      "llm_learning_rate": 3.8383606221426173e-07,
+      "loss": 0.8092,
+      "output_text_mask_loss": 0.0391,
+      "predictor_learning_rate": 3.838360622142617e-05,
+      "step": 8068
+    },
+    {
+      "epoch": 0.83,
+      "gumbel_tau": 0.1491,
+      "image_masks_loss": 0.0123,
+      "llm_learning_rate": 3.833944547418392e-07,
+      "loss": 0.7332,
+      "output_text_mask_loss": 0.0039,
+      "predictor_learning_rate": 3.8339445474183915e-05,
+      "step": 8069
+    },
+    {
+      "epoch": 0.83,
+      "gumbel_tau": 0.149,
+      "image_masks_loss": 0.009,
+      "llm_learning_rate": 3.829530803547329e-07,
+      "loss": 0.7617,
+      "output_text_mask_loss": 0.0484,
+      "predictor_learning_rate": 3.8295308035473294e-05,
+      "step": 8070
+    },
+    {
+      "epoch": 0.83,
+      "gumbel_tau": 0.149,
+      "image_masks_loss": 0.0065,
+      "llm_learning_rate": 3.8251193910155e-07,
+      "loss": 0.6914,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 3.8251193910155e-05,
+      "step": 8071
+    },
+    {
+      "epoch": 0.83,
+      "gumbel_tau": 0.149,
+      "image_masks_loss": 0.0119,
+      "llm_learning_rate": 3.820710310308684e-07,
+      "loss": 0.6414,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 3.820710310308684e-05,
+      "step": 8072
+    },
+    {
+      "epoch": 0.83,
+      "gumbel_tau": 0.1489,
+      "image_masks_loss": 0.0181,
+      "llm_learning_rate": 3.8163035619124265e-07,
+      "loss": 0.7291,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 3.816303561912426e-05,
+      "step": 8073
+    },
+    {
+      "epoch": 0.83,
+      "gumbel_tau": 0.1489,
+      "image_masks_loss": 0.0129,
+      "llm_learning_rate": 3.8118991463120063e-07,
+      "loss": 0.5774,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 3.811899146312006e-05,
+      "step": 8074
+    },
+    {
+      "epoch": 0.83,
+      "gumbel_tau": 0.1488,
+      "image_masks_loss": 0.0091,
+      "llm_learning_rate": 3.8074970639924433e-07,
+      "loss": 0.7926,
+      "output_text_mask_loss": 0.0254,
+      "predictor_learning_rate": 3.8074970639924435e-05,
+      "step": 8075
+    },
+    {
+      "epoch": 0.83,
+      "gumbel_tau": 0.1488,
+      "image_masks_loss": 0.0059,
+      "llm_learning_rate": 3.8030973154385113e-07,
+      "loss": 0.7768,
+      "output_text_mask_loss": 0.0116,
+      "predictor_learning_rate": 3.803097315438511e-05,
+      "step": 8076
+    },
+    {
+      "epoch": 0.83,
+      "gumbel_tau": 0.1488,
+      "image_masks_loss": 0.0042,
+      "llm_learning_rate": 3.798699901134711e-07,
+      "loss": 0.7577,
+      "output_text_mask_loss": 0.0004,
+      "predictor_learning_rate": 3.798699901134711e-05,
+      "step": 8077
+    },
+    {
+      "epoch": 0.83,
+      "gumbel_tau": 0.1487,
+      "image_masks_loss": 0.017,
+      "llm_learning_rate": 3.7943048215653063e-07,
+      "loss": 0.7064,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 3.794304821565306e-05,
+      "step": 8078
+    },
+    {
+      "epoch": 0.83,
+      "gumbel_tau": 0.1487,
+      "image_masks_loss": 0.0095,
+      "llm_learning_rate": 3.789912077214289e-07,
+      "loss": 0.6785,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 3.789912077214289e-05,
+      "step": 8079
+    },
+    {
+      "epoch": 0.83,
+      "gumbel_tau": 0.1487,
+      "image_masks_loss": 0.0106,
+      "llm_learning_rate": 3.7855216685653976e-07,
+      "loss": 0.6789,
+      "output_text_mask_loss": 0.0086,
+      "predictor_learning_rate": 3.7855216685653974e-05,
+      "step": 8080
+    },
+    {
+      "epoch": 0.83,
+      "gumbel_tau": 0.1486,
+      "image_masks_loss": 0.0123,
+      "llm_learning_rate": 3.7811335961021127e-07,
+      "loss": 0.8166,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 3.781133596102112e-05,
+      "step": 8081
+    },
+    {
+      "epoch": 0.83,
+      "gumbel_tau": 0.1486,
+      "image_masks_loss": 0.0215,
+      "llm_learning_rate": 3.776747860307661e-07,
+      "loss": 0.6366,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 3.776747860307661e-05,
+      "step": 8082
+    },
+    {
+      "epoch": 0.83,
+      "gumbel_tau": 0.1486,
+      "image_masks_loss": 0.0115,
+      "llm_learning_rate": 3.7723644616650066e-07,
+      "loss": 0.7055,
+      "output_text_mask_loss": 0.0124,
+      "predictor_learning_rate": 3.772364461665007e-05,
+      "step": 8083
+    },
+    {
+      "epoch": 0.83,
+      "gumbel_tau": 0.1485,
+      "image_masks_loss": 0.0246,
+      "llm_learning_rate": 3.7679834006568603e-07,
+      "loss": 0.6387,
+      "output_text_mask_loss": 0.0126,
+      "predictor_learning_rate": 3.76798340065686e-05,
+      "step": 8084
+    },
+    {
+      "epoch": 0.83,
+      "gumbel_tau": 0.1485,
+      "image_masks_loss": 0.0108,
+      "llm_learning_rate": 3.763604677765673e-07,
+      "loss": 0.7065,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 3.763604677765672e-05,
+      "step": 8085
+    },
+    {
+      "epoch": 0.83,
+      "gumbel_tau": 0.1485,
+      "image_masks_loss": 0.0081,
+      "llm_learning_rate": 3.7592282934736433e-07,
+      "loss": 0.6734,
+      "output_text_mask_loss": 0.0292,
+      "predictor_learning_rate": 3.759228293473643e-05,
+      "step": 8086
+    },
+    {
+      "epoch": 0.83,
+      "gumbel_tau": 0.1484,
+      "image_masks_loss": 0.0091,
+      "llm_learning_rate": 3.7548542482627043e-07,
+      "loss": 0.8088,
+      "output_text_mask_loss": 0.0275,
+      "predictor_learning_rate": 3.7548542482627044e-05,
+      "step": 8087
+    },
+    {
+      "epoch": 0.83,
+      "gumbel_tau": 0.1484,
+      "image_masks_loss": 0.0136,
+      "llm_learning_rate": 3.750482542614539e-07,
+      "loss": 0.6993,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 3.750482542614539e-05,
+      "step": 8088
+    },
+    {
+      "epoch": 0.83,
+      "gumbel_tau": 0.1484,
+      "image_masks_loss": 0.0105,
+      "llm_learning_rate": 3.7461131770105603e-07,
+      "loss": 0.7466,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 3.74611317701056e-05,
+      "step": 8089
+    },
+    {
+      "epoch": 0.83,
+      "gumbel_tau": 0.1483,
+      "image_masks_loss": 0.009,
+      "llm_learning_rate": 3.7417461519319464e-07,
+      "loss": 0.7609,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 3.741746151931946e-05,
+      "step": 8090
+    },
+    {
+      "epoch": 0.83,
+      "gumbel_tau": 0.1483,
+      "image_masks_loss": 0.0102,
+      "llm_learning_rate": 3.7373814678595966e-07,
+      "loss": 0.6405,
+      "output_text_mask_loss": 0.0081,
+      "predictor_learning_rate": 3.737381467859596e-05,
+      "step": 8091
+    },
+    {
+      "epoch": 0.83,
+      "gumbel_tau": 0.1483,
+      "image_masks_loss": 0.0039,
+      "llm_learning_rate": 3.7330191252741596e-07,
+      "loss": 0.6256,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 3.733019125274159e-05,
+      "step": 8092
+    },
+    {
+      "epoch": 0.83,
+      "gumbel_tau": 0.1482,
+      "image_masks_loss": 0.0026,
+      "llm_learning_rate": 3.728659124656028e-07,
+      "loss": 0.7816,
+      "output_text_mask_loss": 0.0075,
+      "predictor_learning_rate": 3.728659124656028e-05,
+      "step": 8093
+    },
+    {
+      "epoch": 0.83,
+      "gumbel_tau": 0.1482,
+      "image_masks_loss": 0.0152,
+      "llm_learning_rate": 3.724301466485333e-07,
+      "loss": 0.7945,
+      "output_text_mask_loss": 0.0035,
+      "predictor_learning_rate": 3.724301466485333e-05,
+      "step": 8094
+    },
+    {
+      "epoch": 0.83,
+      "gumbel_tau": 0.1481,
+      "image_masks_loss": 0.0116,
+      "llm_learning_rate": 3.719946151241949e-07,
+      "loss": 0.6798,
+      "output_text_mask_loss": 0.0235,
+      "predictor_learning_rate": 3.719946151241949e-05,
+      "step": 8095
+    },
+    {
+      "epoch": 0.83,
+      "gumbel_tau": 0.1481,
+      "image_masks_loss": 0.0116,
+      "llm_learning_rate": 3.7155931794054913e-07,
+      "loss": 0.7596,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 3.715593179405491e-05,
+      "step": 8096
+    },
+    {
+      "epoch": 0.83,
+      "gumbel_tau": 0.1481,
+      "image_masks_loss": 0.0147,
+      "llm_learning_rate": 3.711242551455324e-07,
+      "loss": 0.7631,
+      "output_text_mask_loss": 0.0033,
+      "predictor_learning_rate": 3.711242551455324e-05,
+      "step": 8097
+    },
+    {
+      "epoch": 0.83,
+      "gumbel_tau": 0.148,
+      "image_masks_loss": 0.018,
+      "llm_learning_rate": 3.7068942678705505e-07,
+      "loss": 0.7961,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 3.70689426787055e-05,
+      "step": 8098
+    },
+    {
+      "epoch": 0.83,
+      "gumbel_tau": 0.148,
+      "image_masks_loss": 0.0177,
+      "llm_learning_rate": 3.7025483291300043e-07,
+      "loss": 0.6821,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 3.702548329130004e-05,
+      "step": 8099
+    },
+    {
+      "epoch": 0.83,
+      "gumbel_tau": 0.148,
+      "image_masks_loss": 0.0129,
+      "llm_learning_rate": 3.6982047357122767e-07,
+      "loss": 0.7633,
+      "output_text_mask_loss": 0.001,
+      "predictor_learning_rate": 3.6982047357122766e-05,
+      "step": 8100
+    },
+    {
+      "epoch": 0.83,
+      "gumbel_tau": 0.1479,
+      "image_masks_loss": 0.0056,
+      "llm_learning_rate": 3.693863488095692e-07,
+      "loss": 0.7123,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 3.693863488095692e-05,
+      "step": 8101
+    },
+    {
+      "epoch": 0.83,
+      "gumbel_tau": 0.1479,
+      "image_masks_loss": 0.012,
+      "llm_learning_rate": 3.689524586758317e-07,
+      "loss": 0.7593,
+      "output_text_mask_loss": 0.0467,
+      "predictor_learning_rate": 3.689524586758317e-05,
+      "step": 8102
+    },
+    {
+      "epoch": 0.83,
+      "gumbel_tau": 0.1479,
+      "image_masks_loss": 0.012,
+      "llm_learning_rate": 3.685188032177961e-07,
+      "loss": 0.6043,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 3.685188032177961e-05,
+      "step": 8103
+    },
+    {
+      "epoch": 0.83,
+      "gumbel_tau": 0.1478,
+      "image_masks_loss": 0.0091,
+      "llm_learning_rate": 3.6808538248321757e-07,
+      "loss": 0.7145,
+      "output_text_mask_loss": 0.0024,
+      "predictor_learning_rate": 3.680853824832175e-05,
+      "step": 8104
+    },
+    {
+      "epoch": 0.83,
+      "gumbel_tau": 0.1478,
+      "image_masks_loss": 0.0093,
+      "llm_learning_rate": 3.676521965198254e-07,
+      "loss": 0.7053,
+      "output_text_mask_loss": 0.0123,
+      "predictor_learning_rate": 3.676521965198254e-05,
+      "step": 8105
+    },
+    {
+      "epoch": 0.83,
+      "gumbel_tau": 0.1478,
+      "image_masks_loss": 0.0192,
+      "llm_learning_rate": 3.672192453753229e-07,
+      "loss": 0.8305,
+      "output_text_mask_loss": 0.0081,
+      "predictor_learning_rate": 3.672192453753229e-05,
+      "step": 8106
+    },
+    {
+      "epoch": 0.83,
+      "gumbel_tau": 0.1477,
+      "image_masks_loss": 0.0061,
+      "llm_learning_rate": 3.6678652909738796e-07,
+      "loss": 0.7742,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 3.667865290973879e-05,
+      "step": 8107
+    },
+    {
+      "epoch": 0.83,
+      "gumbel_tau": 0.1477,
+      "image_masks_loss": 0.0151,
+      "llm_learning_rate": 3.6635404773367163e-07,
+      "loss": 0.7548,
+      "output_text_mask_loss": 0.0038,
+      "predictor_learning_rate": 3.663540477336716e-05,
+      "step": 8108
+    },
+    {
+      "epoch": 0.83,
+      "gumbel_tau": 0.1477,
+      "image_masks_loss": 0.0219,
+      "llm_learning_rate": 3.659218013317997e-07,
+      "loss": 0.6933,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 3.659218013317997e-05,
+      "step": 8109
+    },
+    {
+      "epoch": 0.83,
+      "gumbel_tau": 0.1476,
+      "image_masks_loss": 0.016,
+      "llm_learning_rate": 3.6548978993937304e-07,
+      "loss": 0.6572,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 3.6548978993937305e-05,
+      "step": 8110
+    },
+    {
+      "epoch": 0.83,
+      "gumbel_tau": 0.1476,
+      "image_masks_loss": 0.013,
+      "llm_learning_rate": 3.65058013603965e-07,
+      "loss": 0.736,
+      "output_text_mask_loss": 0.0277,
+      "predictor_learning_rate": 3.6505801360396495e-05,
+      "step": 8111
+    },
+    {
+      "epoch": 0.83,
+      "gumbel_tau": 0.1476,
+      "image_masks_loss": 0.0121,
+      "llm_learning_rate": 3.646264723731241e-07,
+      "loss": 0.7994,
+      "output_text_mask_loss": 0.0187,
+      "predictor_learning_rate": 3.646264723731241e-05,
+      "step": 8112
+    },
+    {
+      "epoch": 0.83,
+      "gumbel_tau": 0.1475,
+      "image_masks_loss": 0.0094,
+      "llm_learning_rate": 3.641951662943724e-07,
+      "loss": 0.7464,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 3.641951662943724e-05,
+      "step": 8113
+    },
+    {
+      "epoch": 0.83,
+      "gumbel_tau": 0.1475,
+      "image_masks_loss": 0.0122,
+      "llm_learning_rate": 3.637640954152061e-07,
+      "loss": 0.769,
+      "output_text_mask_loss": 0.0233,
+      "predictor_learning_rate": 3.637640954152061e-05,
+      "step": 8114
+    },
+    {
+      "epoch": 0.83,
+      "gumbel_tau": 0.1475,
+      "image_masks_loss": 0.0109,
+      "llm_learning_rate": 3.63333259783096e-07,
+      "loss": 0.748,
+      "output_text_mask_loss": 0.0198,
+      "predictor_learning_rate": 3.63333259783096e-05,
+      "step": 8115
+    },
+    {
+      "epoch": 0.83,
+      "gumbel_tau": 0.1474,
+      "image_masks_loss": 0.0167,
+      "llm_learning_rate": 3.62902659445486e-07,
+      "loss": 0.6961,
+      "output_text_mask_loss": 0.0035,
+      "predictor_learning_rate": 3.6290265944548596e-05,
+      "step": 8116
+    },
+    {
+      "epoch": 0.83,
+      "gumbel_tau": 0.1474,
+      "image_masks_loss": 0.0103,
+      "llm_learning_rate": 3.6247229444979604e-07,
+      "loss": 0.7294,
+      "output_text_mask_loss": 0.0082,
+      "predictor_learning_rate": 3.6247229444979604e-05,
+      "step": 8117
+    },
+    {
+      "epoch": 0.83,
+      "gumbel_tau": 0.1473,
+      "image_masks_loss": 0.004,
+      "llm_learning_rate": 3.62042164843418e-07,
+      "loss": 0.6677,
+      "output_text_mask_loss": 0.0032,
+      "predictor_learning_rate": 3.62042164843418e-05,
+      "step": 8118
+    },
+    {
+      "epoch": 0.83,
+      "gumbel_tau": 0.1473,
+      "image_masks_loss": 0.0177,
+      "llm_learning_rate": 3.6161227067371904e-07,
+      "loss": 0.7351,
+      "output_text_mask_loss": 0.0307,
+      "predictor_learning_rate": 3.6161227067371905e-05,
+      "step": 8119
+    },
+    {
+      "epoch": 0.83,
+      "gumbel_tau": 0.1473,
+      "image_masks_loss": 0.0134,
+      "llm_learning_rate": 3.611826119880399e-07,
+      "loss": 0.7205,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 3.611826119880399e-05,
+      "step": 8120
+    },
+    {
+      "epoch": 0.83,
+      "gumbel_tau": 0.1472,
+      "image_masks_loss": 0.025,
+      "llm_learning_rate": 3.6075318883369543e-07,
+      "loss": 0.7022,
+      "output_text_mask_loss": 0.0061,
+      "predictor_learning_rate": 3.607531888336954e-05,
+      "step": 8121
+    },
+    {
+      "epoch": 0.83,
+      "gumbel_tau": 0.1472,
+      "image_masks_loss": 0.0095,
+      "llm_learning_rate": 3.603240012579748e-07,
+      "loss": 0.7728,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 3.6032400125797475e-05,
+      "step": 8122
+    },
+    {
+      "epoch": 0.83,
+      "gumbel_tau": 0.1472,
+      "image_masks_loss": 0.0515,
+      "llm_learning_rate": 3.598950493081413e-07,
+      "loss": 0.7152,
+      "output_text_mask_loss": 0.0055,
+      "predictor_learning_rate": 3.5989504930814124e-05,
+      "step": 8123
+    },
+    {
+      "epoch": 0.83,
+      "gumbel_tau": 0.1471,
+      "image_masks_loss": 0.0095,
+      "llm_learning_rate": 3.5946633303143155e-07,
+      "loss": 0.7058,
+      "output_text_mask_loss": 0.017,
+      "predictor_learning_rate": 3.5946633303143154e-05,
+      "step": 8124
+    },
+    {
+      "epoch": 0.83,
+      "gumbel_tau": 0.1471,
+      "image_masks_loss": 0.007,
+      "llm_learning_rate": 3.590378524750573e-07,
+      "loss": 0.7832,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 3.590378524750573e-05,
+      "step": 8125
+    },
+    {
+      "epoch": 0.83,
+      "gumbel_tau": 0.1471,
+      "image_masks_loss": 0.0074,
+      "llm_learning_rate": 3.586096076862033e-07,
+      "loss": 0.7588,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 3.5860960768620334e-05,
+      "step": 8126
+    },
+    {
+      "epoch": 0.83,
+      "gumbel_tau": 0.147,
+      "image_masks_loss": 0.0155,
+      "llm_learning_rate": 3.5818159871202936e-07,
+      "loss": 0.7574,
+      "output_text_mask_loss": 0.0141,
+      "predictor_learning_rate": 3.5818159871202934e-05,
+      "step": 8127
+    },
+    {
+      "epoch": 0.83,
+      "gumbel_tau": 0.147,
+      "image_masks_loss": 0.0114,
+      "llm_learning_rate": 3.5775382559966766e-07,
+      "loss": 0.5915,
+      "output_text_mask_loss": 0.0023,
+      "predictor_learning_rate": 3.577538255996676e-05,
+      "step": 8128
+    },
+    {
+      "epoch": 0.83,
+      "gumbel_tau": 0.147,
+      "image_masks_loss": 0.0056,
+      "llm_learning_rate": 3.5732628839622694e-07,
+      "loss": 0.6642,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 3.57326288396227e-05,
+      "step": 8129
+    },
+    {
+      "epoch": 0.83,
+      "gumbel_tau": 0.1469,
+      "image_masks_loss": 0.0082,
+      "llm_learning_rate": 3.568989871487882e-07,
+      "loss": 0.6805,
+      "output_text_mask_loss": 0.0096,
+      "predictor_learning_rate": 3.568989871487882e-05,
+      "step": 8130
+    },
+    {
+      "epoch": 0.83,
+      "gumbel_tau": 0.1469,
+      "image_masks_loss": 0.009,
+      "llm_learning_rate": 3.5647192190440635e-07,
+      "loss": 0.7146,
+      "output_text_mask_loss": 0.0038,
+      "predictor_learning_rate": 3.564719219044063e-05,
+      "step": 8131
+    },
+    {
+      "epoch": 0.83,
+      "gumbel_tau": 0.1469,
+      "image_masks_loss": 0.0116,
+      "llm_learning_rate": 3.5604509271011113e-07,
+      "loss": 0.6786,
+      "output_text_mask_loss": 0.0189,
+      "predictor_learning_rate": 3.560450927101111e-05,
+      "step": 8132
+    },
+    {
+      "epoch": 0.83,
+      "gumbel_tau": 0.1468,
+      "image_masks_loss": 0.0156,
+      "llm_learning_rate": 3.556184996129056e-07,
+      "loss": 0.6084,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 3.556184996129055e-05,
+      "step": 8133
+    },
+    {
+      "epoch": 0.83,
+      "gumbel_tau": 0.1468,
+      "image_masks_loss": 0.006,
+      "llm_learning_rate": 3.551921426597674e-07,
+      "loss": 0.7796,
+      "output_text_mask_loss": 0.0139,
+      "predictor_learning_rate": 3.5519214265976736e-05,
+      "step": 8134
+    },
+    {
+      "epoch": 0.83,
+      "gumbel_tau": 0.1468,
+      "image_masks_loss": 0.0101,
+      "llm_learning_rate": 3.547660218976473e-07,
+      "loss": 0.7461,
+      "output_text_mask_loss": 0.0072,
+      "predictor_learning_rate": 3.5476602189764725e-05,
+      "step": 8135
+    },
+    {
+      "epoch": 0.83,
+      "gumbel_tau": 0.1467,
+      "image_masks_loss": 0.0164,
+      "llm_learning_rate": 3.543401373734717e-07,
+      "loss": 0.755,
+      "output_text_mask_loss": 0.0652,
+      "predictor_learning_rate": 3.543401373734717e-05,
+      "step": 8136
+    },
+    {
+      "epoch": 0.83,
+      "gumbel_tau": 0.1467,
+      "image_masks_loss": 0.004,
+      "llm_learning_rate": 3.5391448913413953e-07,
+      "loss": 0.6667,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 3.539144891341395e-05,
+      "step": 8137
+    },
+    {
+      "epoch": 0.83,
+      "gumbel_tau": 0.1467,
+      "image_masks_loss": 0.029,
+      "llm_learning_rate": 3.53489077226524e-07,
+      "loss": 0.7381,
+      "output_text_mask_loss": 0.0044,
+      "predictor_learning_rate": 3.5348907722652395e-05,
+      "step": 8138
+    },
+    {
+      "epoch": 0.83,
+      "gumbel_tau": 0.1466,
+      "image_masks_loss": 0.0113,
+      "llm_learning_rate": 3.5306390169747295e-07,
+      "loss": 0.6928,
+      "output_text_mask_loss": 0.016,
+      "predictor_learning_rate": 3.530639016974729e-05,
+      "step": 8139
+    },
+    {
+      "epoch": 0.83,
+      "gumbel_tau": 0.1466,
+      "image_masks_loss": 0.0072,
+      "llm_learning_rate": 3.5263896259380647e-07,
+      "loss": 0.6528,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 3.526389625938065e-05,
+      "step": 8140
+    },
+    {
+      "epoch": 0.83,
+      "gumbel_tau": 0.1465,
+      "image_masks_loss": 0.0086,
+      "llm_learning_rate": 3.5221425996232016e-07,
+      "loss": 0.685,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 3.522142599623202e-05,
+      "step": 8141
+    },
+    {
+      "epoch": 0.83,
+      "gumbel_tau": 0.1465,
+      "image_masks_loss": 0.0067,
+      "llm_learning_rate": 3.517897938497841e-07,
+      "loss": 0.678,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 3.517897938497841e-05,
+      "step": 8142
+    },
+    {
+      "epoch": 0.83,
+      "gumbel_tau": 0.1465,
+      "image_masks_loss": 0.0267,
+      "llm_learning_rate": 3.5136556430294097e-07,
+      "loss": 0.7424,
+      "output_text_mask_loss": 0.0233,
+      "predictor_learning_rate": 3.513655643029409e-05,
+      "step": 8143
+    },
+    {
+      "epoch": 0.83,
+      "gumbel_tau": 0.1464,
+      "image_masks_loss": 0.0096,
+      "llm_learning_rate": 3.5094157136850743e-07,
+      "loss": 0.7458,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 3.509415713685074e-05,
+      "step": 8144
+    },
+    {
+      "epoch": 0.83,
+      "gumbel_tau": 0.1464,
+      "image_masks_loss": 0.0073,
+      "llm_learning_rate": 3.505178150931751e-07,
+      "loss": 0.7293,
+      "output_text_mask_loss": 0.0322,
+      "predictor_learning_rate": 3.505178150931751e-05,
+      "step": 8145
+    },
+    {
+      "epoch": 0.83,
+      "gumbel_tau": 0.1464,
+      "image_masks_loss": 0.0106,
+      "llm_learning_rate": 3.500942955236089e-07,
+      "loss": 0.6595,
+      "output_text_mask_loss": 0.0006,
+      "predictor_learning_rate": 3.500942955236089e-05,
+      "step": 8146
+    },
+    {
+      "epoch": 0.83,
+      "gumbel_tau": 0.1463,
+      "image_masks_loss": 0.0064,
+      "llm_learning_rate": 3.4967101270644734e-07,
+      "loss": 0.7651,
+      "output_text_mask_loss": 0.0086,
+      "predictor_learning_rate": 3.496710127064473e-05,
+      "step": 8147
+    },
+    {
+      "epoch": 0.83,
+      "gumbel_tau": 0.1463,
+      "image_masks_loss": 0.0094,
+      "llm_learning_rate": 3.4924796668830306e-07,
+      "loss": 0.719,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 3.4924796668830304e-05,
+      "step": 8148
+    },
+    {
+      "epoch": 0.83,
+      "gumbel_tau": 0.1463,
+      "image_masks_loss": 0.0036,
+      "llm_learning_rate": 3.488251575157639e-07,
+      "loss": 0.7294,
+      "output_text_mask_loss": 0.001,
+      "predictor_learning_rate": 3.488251575157639e-05,
+      "step": 8149
+    },
+    {
+      "epoch": 0.84,
+      "gumbel_tau": 0.1462,
+      "image_masks_loss": 0.0169,
+      "llm_learning_rate": 3.4840258523539006e-07,
+      "loss": 0.8495,
+      "output_text_mask_loss": 0.0156,
+      "predictor_learning_rate": 3.4840258523539e-05,
+      "step": 8150
+    },
+    {
+      "epoch": 0.84,
+      "gumbel_tau": 0.1462,
+      "image_masks_loss": 0.0106,
+      "llm_learning_rate": 3.47980249893716e-07,
+      "loss": 0.7538,
+      "output_text_mask_loss": 0.0035,
+      "predictor_learning_rate": 3.4798024989371596e-05,
+      "step": 8151
+    },
+    {
+      "epoch": 0.84,
+      "gumbel_tau": 0.1462,
+      "image_masks_loss": 0.0068,
+      "llm_learning_rate": 3.4755815153725e-07,
+      "loss": 0.7022,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 3.4755815153725e-05,
+      "step": 8152
+    },
+    {
+      "epoch": 0.84,
+      "gumbel_tau": 0.1461,
+      "image_masks_loss": 0.0047,
+      "llm_learning_rate": 3.4713629021247505e-07,
+      "loss": 0.7003,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 3.47136290212475e-05,
+      "step": 8153
+    },
+    {
+      "epoch": 0.84,
+      "gumbel_tau": 0.1461,
+      "image_masks_loss": 0.0083,
+      "llm_learning_rate": 3.4671466596584683e-07,
+      "loss": 0.8209,
+      "output_text_mask_loss": 0.0086,
+      "predictor_learning_rate": 3.467146659658468e-05,
+      "step": 8154
+    },
+    {
+      "epoch": 0.84,
+      "gumbel_tau": 0.1461,
+      "image_masks_loss": 0.0081,
+      "llm_learning_rate": 3.462932788437956e-07,
+      "loss": 0.7254,
+      "output_text_mask_loss": 0.011,
+      "predictor_learning_rate": 3.4629327884379556e-05,
+      "step": 8155
+    },
+    {
+      "epoch": 0.84,
+      "gumbel_tau": 0.146,
+      "image_masks_loss": 0.0046,
+      "llm_learning_rate": 3.458721288927261e-07,
+      "loss": 0.7248,
+      "output_text_mask_loss": 0.0256,
+      "predictor_learning_rate": 3.458721288927261e-05,
+      "step": 8156
+    },
+    {
+      "epoch": 0.84,
+      "gumbel_tau": 0.146,
+      "image_masks_loss": 0.0117,
+      "llm_learning_rate": 3.4545121615901596e-07,
+      "loss": 0.7473,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 3.454512161590159e-05,
+      "step": 8157
+    },
+    {
+      "epoch": 0.84,
+      "gumbel_tau": 0.146,
+      "image_masks_loss": 0.0134,
+      "llm_learning_rate": 3.4503054068901753e-07,
+      "loss": 0.7718,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 3.450305406890175e-05,
+      "step": 8158
+    },
+    {
+      "epoch": 0.84,
+      "gumbel_tau": 0.1459,
+      "image_masks_loss": 0.0048,
+      "llm_learning_rate": 3.446101025290552e-07,
+      "loss": 0.8358,
+      "output_text_mask_loss": 0.0058,
+      "predictor_learning_rate": 3.446101025290552e-05,
+      "step": 8159
+    },
+    {
+      "epoch": 0.84,
+      "gumbel_tau": 0.1459,
+      "image_masks_loss": 0.0167,
+      "llm_learning_rate": 3.441899017254294e-07,
+      "loss": 0.7314,
+      "output_text_mask_loss": 0.0019,
+      "predictor_learning_rate": 3.441899017254294e-05,
+      "step": 8160
+    },
+    {
+      "epoch": 0.84,
+      "gumbel_tau": 0.1459,
+      "image_masks_loss": 0.0062,
+      "llm_learning_rate": 3.437699383244131e-07,
+      "loss": 0.7227,
+      "output_text_mask_loss": 0.0343,
+      "predictor_learning_rate": 3.437699383244131e-05,
+      "step": 8161
+    },
+    {
+      "epoch": 0.84,
+      "gumbel_tau": 0.1458,
+      "image_masks_loss": 0.0213,
+      "llm_learning_rate": 3.4335021237225456e-07,
+      "loss": 0.8354,
+      "output_text_mask_loss": 0.0025,
+      "predictor_learning_rate": 3.4335021237225455e-05,
+      "step": 8162
+    },
+    {
+      "epoch": 0.84,
+      "gumbel_tau": 0.1458,
+      "image_masks_loss": 0.0097,
+      "llm_learning_rate": 3.4293072391517414e-07,
+      "loss": 0.7088,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 3.429307239151741e-05,
+      "step": 8163
+    },
+    {
+      "epoch": 0.84,
+      "gumbel_tau": 0.1458,
+      "image_masks_loss": 0.0128,
+      "llm_learning_rate": 3.425114729993673e-07,
+      "loss": 0.7411,
+      "output_text_mask_loss": 0.0275,
+      "predictor_learning_rate": 3.425114729993672e-05,
+      "step": 8164
+    },
+    {
+      "epoch": 0.84,
+      "gumbel_tau": 0.1457,
+      "image_masks_loss": 0.02,
+      "llm_learning_rate": 3.4209245967100265e-07,
+      "loss": 0.7297,
+      "output_text_mask_loss": 0.0012,
+      "predictor_learning_rate": 3.420924596710026e-05,
+      "step": 8165
+    },
+    {
+      "epoch": 0.84,
+      "gumbel_tau": 0.1457,
+      "image_masks_loss": 0.0085,
+      "llm_learning_rate": 3.4167368397622276e-07,
+      "loss": 0.7779,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 3.416736839762227e-05,
+      "step": 8166
+    },
+    {
+      "epoch": 0.84,
+      "gumbel_tau": 0.1457,
+      "image_masks_loss": 0.0094,
+      "llm_learning_rate": 3.4125514596114463e-07,
+      "loss": 0.8044,
+      "output_text_mask_loss": 0.0504,
+      "predictor_learning_rate": 3.412551459611446e-05,
+      "step": 8167
+    },
+    {
+      "epoch": 0.84,
+      "gumbel_tau": 0.1456,
+      "image_masks_loss": 0.0071,
+      "llm_learning_rate": 3.4083684567185723e-07,
+      "loss": 0.6908,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 3.408368456718572e-05,
+      "step": 8168
+    },
+    {
+      "epoch": 0.84,
+      "gumbel_tau": 0.1456,
+      "image_masks_loss": 0.0114,
+      "llm_learning_rate": 3.4041878315442645e-07,
+      "loss": 0.7081,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 3.404187831544264e-05,
+      "step": 8169
+    },
+    {
+      "epoch": 0.84,
+      "gumbel_tau": 0.1456,
+      "image_masks_loss": 0.0101,
+      "llm_learning_rate": 3.400009584548897e-07,
+      "loss": 0.6697,
+      "output_text_mask_loss": 0.0078,
+      "predictor_learning_rate": 3.400009584548897e-05,
+      "step": 8170
+    },
+    {
+      "epoch": 0.84,
+      "gumbel_tau": 0.1455,
+      "image_masks_loss": 0.024,
+      "llm_learning_rate": 3.395833716192584e-07,
+      "loss": 0.7398,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 3.395833716192584e-05,
+      "step": 8171
+    },
+    {
+      "epoch": 0.84,
+      "gumbel_tau": 0.1455,
+      "image_masks_loss": 0.0037,
+      "llm_learning_rate": 3.3916602269351866e-07,
+      "loss": 0.7342,
+      "output_text_mask_loss": 0.0098,
+      "predictor_learning_rate": 3.391660226935187e-05,
+      "step": 8172
+    },
+    {
+      "epoch": 0.84,
+      "gumbel_tau": 0.1454,
+      "image_masks_loss": 0.0113,
+      "llm_learning_rate": 3.387489117236292e-07,
+      "loss": 0.7057,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 3.387489117236292e-05,
+      "step": 8173
+    },
+    {
+      "epoch": 0.84,
+      "gumbel_tau": 0.1454,
+      "image_masks_loss": 0.0145,
+      "llm_learning_rate": 3.3833203875552387e-07,
+      "loss": 0.7175,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 3.3833203875552385e-05,
+      "step": 8174
+    },
+    {
+      "epoch": 0.84,
+      "gumbel_tau": 0.1454,
+      "image_masks_loss": 0.0142,
+      "llm_learning_rate": 3.3791540383510873e-07,
+      "loss": 0.7081,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 3.379154038351087e-05,
+      "step": 8175
+    },
+    {
+      "epoch": 0.84,
+      "gumbel_tau": 0.1453,
+      "image_masks_loss": 0.0082,
+      "llm_learning_rate": 3.374990070082662e-07,
+      "loss": 0.7794,
+      "output_text_mask_loss": 0.0126,
+      "predictor_learning_rate": 3.374990070082662e-05,
+      "step": 8176
+    },
+    {
+      "epoch": 0.84,
+      "gumbel_tau": 0.1453,
+      "image_masks_loss": 0.0106,
+      "llm_learning_rate": 3.370828483208491e-07,
+      "loss": 0.6548,
+      "output_text_mask_loss": 0.0187,
+      "predictor_learning_rate": 3.370828483208491e-05,
+      "step": 8177
+    },
+    {
+      "epoch": 0.84,
+      "gumbel_tau": 0.1453,
+      "image_masks_loss": 0.0096,
+      "llm_learning_rate": 3.366669278186868e-07,
+      "loss": 0.7856,
+      "output_text_mask_loss": 0.0267,
+      "predictor_learning_rate": 3.366669278186868e-05,
+      "step": 8178
+    },
+    {
+      "epoch": 0.84,
+      "gumbel_tau": 0.1452,
+      "image_masks_loss": 0.0154,
+      "llm_learning_rate": 3.3625124554758083e-07,
+      "loss": 0.736,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 3.362512455475808e-05,
+      "step": 8179
+    },
+    {
+      "epoch": 0.84,
+      "gumbel_tau": 0.1452,
+      "image_masks_loss": 0.0134,
+      "llm_learning_rate": 3.358358015533067e-07,
+      "loss": 0.7965,
+      "output_text_mask_loss": 0.009,
+      "predictor_learning_rate": 3.358358015533067e-05,
+      "step": 8180
+    },
+    {
+      "epoch": 0.84,
+      "gumbel_tau": 0.1452,
+      "image_masks_loss": 0.0078,
+      "llm_learning_rate": 3.3542059588161535e-07,
+      "loss": 0.7646,
+      "output_text_mask_loss": 0.0078,
+      "predictor_learning_rate": 3.354205958816153e-05,
+      "step": 8181
+    },
+    {
+      "epoch": 0.84,
+      "gumbel_tau": 0.1451,
+      "image_masks_loss": 0.0086,
+      "llm_learning_rate": 3.350056285782291e-07,
+      "loss": 0.7308,
+      "output_text_mask_loss": 0.0496,
+      "predictor_learning_rate": 3.3500562857822905e-05,
+      "step": 8182
+    },
+    {
+      "epoch": 0.84,
+      "gumbel_tau": 0.1451,
+      "image_masks_loss": 0.0123,
+      "llm_learning_rate": 3.345908996888456e-07,
+      "loss": 0.7262,
+      "output_text_mask_loss": 0.0063,
+      "predictor_learning_rate": 3.3459089968884555e-05,
+      "step": 8183
+    },
+    {
+      "epoch": 0.84,
+      "gumbel_tau": 0.1451,
+      "image_masks_loss": 0.0082,
+      "llm_learning_rate": 3.341764092591354e-07,
+      "loss": 0.7212,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 3.341764092591354e-05,
+      "step": 8184
+    },
+    {
+      "epoch": 0.84,
+      "gumbel_tau": 0.145,
+      "image_masks_loss": 0.0065,
+      "llm_learning_rate": 3.3376215733474327e-07,
+      "loss": 0.6792,
+      "output_text_mask_loss": 0.0591,
+      "predictor_learning_rate": 3.3376215733474325e-05,
+      "step": 8185
+    },
+    {
+      "epoch": 0.84,
+      "gumbel_tau": 0.145,
+      "image_masks_loss": 0.0131,
+      "llm_learning_rate": 3.3334814396128736e-07,
+      "loss": 0.7322,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 3.333481439612873e-05,
+      "step": 8186
+    },
+    {
+      "epoch": 0.84,
+      "gumbel_tau": 0.145,
+      "image_masks_loss": 0.005,
+      "llm_learning_rate": 3.3293436918435923e-07,
+      "loss": 0.7508,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 3.329343691843592e-05,
+      "step": 8187
+    },
+    {
+      "epoch": 0.84,
+      "gumbel_tau": 0.1449,
+      "image_masks_loss": 0.0186,
+      "llm_learning_rate": 3.325208330495261e-07,
+      "loss": 0.7257,
+      "output_text_mask_loss": 0.0078,
+      "predictor_learning_rate": 3.325208330495261e-05,
+      "step": 8188
+    },
+    {
+      "epoch": 0.84,
+      "gumbel_tau": 0.1449,
+      "image_masks_loss": 0.0084,
+      "llm_learning_rate": 3.321075356023265e-07,
+      "loss": 0.66,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 3.321075356023265e-05,
+      "step": 8189
+    },
+    {
+      "epoch": 0.84,
+      "gumbel_tau": 0.1449,
+      "image_masks_loss": 0.0123,
+      "llm_learning_rate": 3.3169447688827404e-07,
+      "loss": 0.7062,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 3.31694476888274e-05,
+      "step": 8190
+    },
+    {
+      "epoch": 0.84,
+      "gumbel_tau": 0.1448,
+      "image_masks_loss": 0.0187,
+      "llm_learning_rate": 3.3128165695285543e-07,
+      "loss": 0.7648,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 3.312816569528554e-05,
+      "step": 8191
+    },
+    {
+      "epoch": 0.84,
+      "gumbel_tau": 0.1448,
+      "image_masks_loss": 0.0113,
+      "llm_learning_rate": 3.3086907584153145e-07,
+      "loss": 0.7251,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 3.3086907584153144e-05,
+      "step": 8192
+    },
+    {
+      "epoch": 0.84,
+      "gumbel_tau": 0.1448,
+      "image_masks_loss": 0.0114,
+      "llm_learning_rate": 3.304567335997361e-07,
+      "loss": 0.7058,
+      "output_text_mask_loss": 0.0174,
+      "predictor_learning_rate": 3.304567335997361e-05,
+      "step": 8193
+    },
+    {
+      "epoch": 0.84,
+      "gumbel_tau": 0.1447,
+      "image_masks_loss": 0.0267,
+      "llm_learning_rate": 3.300446302728774e-07,
+      "loss": 0.8069,
+      "output_text_mask_loss": 0.0366,
+      "predictor_learning_rate": 3.300446302728774e-05,
+      "step": 8194
+    },
+    {
+      "epoch": 0.84,
+      "gumbel_tau": 0.1447,
+      "image_masks_loss": 0.0042,
+      "llm_learning_rate": 3.2963276590633864e-07,
+      "loss": 0.6874,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 3.296327659063386e-05,
+      "step": 8195
+    },
+    {
+      "epoch": 0.84,
+      "gumbel_tau": 0.1447,
+      "image_masks_loss": 0.0052,
+      "llm_learning_rate": 3.2922114054547316e-07,
+      "loss": 0.7817,
+      "output_text_mask_loss": 0.0111,
+      "predictor_learning_rate": 3.2922114054547315e-05,
+      "step": 8196
+    },
+    {
+      "epoch": 0.84,
+      "gumbel_tau": 0.1446,
+      "image_masks_loss": 0.0049,
+      "llm_learning_rate": 3.288097542356111e-07,
+      "loss": 0.709,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 3.288097542356111e-05,
+      "step": 8197
+    },
+    {
+      "epoch": 0.84,
+      "gumbel_tau": 0.1446,
+      "image_masks_loss": 0.0133,
+      "llm_learning_rate": 3.283986070220549e-07,
+      "loss": 0.7379,
+      "output_text_mask_loss": 0.0181,
+      "predictor_learning_rate": 3.2839860702205486e-05,
+      "step": 8198
+    },
+    {
+      "epoch": 0.84,
+      "gumbel_tau": 0.1446,
+      "image_masks_loss": 0.0077,
+      "llm_learning_rate": 3.279876989500813e-07,
+      "loss": 0.8031,
+      "output_text_mask_loss": 0.0038,
+      "predictor_learning_rate": 3.279876989500813e-05,
+      "step": 8199
+    },
+    {
+      "epoch": 0.84,
+      "gumbel_tau": 0.1445,
+      "image_masks_loss": 0.007,
+      "llm_learning_rate": 3.2757703006493974e-07,
+      "loss": 0.7298,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 3.275770300649397e-05,
+      "step": 8200
+    },
+    {
+      "epoch": 0.84,
+      "gumbel_tau": 0.1445,
+      "image_masks_loss": 0.0164,
+      "llm_learning_rate": 3.2716660041185483e-07,
+      "loss": 0.7067,
+      "output_text_mask_loss": 0.0198,
+      "predictor_learning_rate": 3.271666004118548e-05,
+      "step": 8201
+    },
+    {
+      "epoch": 0.84,
+      "gumbel_tau": 0.1445,
+      "image_masks_loss": 0.0093,
+      "llm_learning_rate": 3.267564100360238e-07,
+      "loss": 0.7472,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 3.267564100360237e-05,
+      "step": 8202
+    },
+    {
+      "epoch": 0.84,
+      "gumbel_tau": 0.1444,
+      "image_masks_loss": 0.0157,
+      "llm_learning_rate": 3.2634645898261764e-07,
+      "loss": 0.7664,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 3.263464589826176e-05,
+      "step": 8203
+    },
+    {
+      "epoch": 0.84,
+      "gumbel_tau": 0.1444,
+      "image_masks_loss": 0.0035,
+      "llm_learning_rate": 3.2593674729678115e-07,
+      "loss": 0.7088,
+      "output_text_mask_loss": 0.0078,
+      "predictor_learning_rate": 3.259367472967811e-05,
+      "step": 8204
+    },
+    {
+      "epoch": 0.84,
+      "gumbel_tau": 0.1444,
+      "image_masks_loss": 0.0147,
+      "llm_learning_rate": 3.255272750236324e-07,
+      "loss": 0.5608,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 3.255272750236324e-05,
+      "step": 8205
+    },
+    {
+      "epoch": 0.84,
+      "gumbel_tau": 0.1443,
+      "image_masks_loss": 0.0122,
+      "llm_learning_rate": 3.251180422082637e-07,
+      "loss": 0.7927,
+      "output_text_mask_loss": 0.0041,
+      "predictor_learning_rate": 3.251180422082636e-05,
+      "step": 8206
+    },
+    {
+      "epoch": 0.84,
+      "gumbel_tau": 0.1443,
+      "image_masks_loss": 0.0296,
+      "llm_learning_rate": 3.2470904889573966e-07,
+      "loss": 0.7728,
+      "output_text_mask_loss": 0.0139,
+      "predictor_learning_rate": 3.247090488957396e-05,
+      "step": 8207
+    },
+    {
+      "epoch": 0.84,
+      "gumbel_tau": 0.1443,
+      "image_masks_loss": 0.0073,
+      "llm_learning_rate": 3.243002951311014e-07,
+      "loss": 0.7144,
+      "output_text_mask_loss": 0.0025,
+      "predictor_learning_rate": 3.2430029513110136e-05,
+      "step": 8208
+    },
+    {
+      "epoch": 0.84,
+      "gumbel_tau": 0.1442,
+      "image_masks_loss": 0.0078,
+      "llm_learning_rate": 3.238917809593606e-07,
+      "loss": 0.6971,
+      "output_text_mask_loss": 0.0023,
+      "predictor_learning_rate": 3.2389178095936054e-05,
+      "step": 8209
+    },
+    {
+      "epoch": 0.84,
+      "gumbel_tau": 0.1442,
+      "image_masks_loss": 0.0037,
+      "llm_learning_rate": 3.2348350642550403e-07,
+      "loss": 0.7483,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 3.23483506425504e-05,
+      "step": 8210
+    },
+    {
+      "epoch": 0.84,
+      "gumbel_tau": 0.1441,
+      "image_masks_loss": 0.0063,
+      "llm_learning_rate": 3.230754715744919e-07,
+      "loss": 0.8513,
+      "output_text_mask_loss": 0.0151,
+      "predictor_learning_rate": 3.2307547157449185e-05,
+      "step": 8211
+    },
+    {
+      "epoch": 0.84,
+      "gumbel_tau": 0.1441,
+      "image_masks_loss": 0.016,
+      "llm_learning_rate": 3.2266767645125734e-07,
+      "loss": 0.656,
+      "output_text_mask_loss": 0.0067,
+      "predictor_learning_rate": 3.2266767645125735e-05,
+      "step": 8212
+    },
+    {
+      "epoch": 0.84,
+      "gumbel_tau": 0.1441,
+      "image_masks_loss": 0.0039,
+      "llm_learning_rate": 3.222601211007084e-07,
+      "loss": 0.792,
+      "output_text_mask_loss": 0.0099,
+      "predictor_learning_rate": 3.222601211007084e-05,
+      "step": 8213
+    },
+    {
+      "epoch": 0.84,
+      "gumbel_tau": 0.144,
+      "image_masks_loss": 0.0121,
+      "llm_learning_rate": 3.218528055677256e-07,
+      "loss": 0.7767,
+      "output_text_mask_loss": 0.0111,
+      "predictor_learning_rate": 3.2185280556772554e-05,
+      "step": 8214
+    },
+    {
+      "epoch": 0.84,
+      "gumbel_tau": 0.144,
+      "image_masks_loss": 0.0112,
+      "llm_learning_rate": 3.214457298971635e-07,
+      "loss": 0.7741,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 3.2144572989716346e-05,
+      "step": 8215
+    },
+    {
+      "epoch": 0.84,
+      "gumbel_tau": 0.144,
+      "image_masks_loss": 0.0068,
+      "llm_learning_rate": 3.210388941338499e-07,
+      "loss": 0.7913,
+      "output_text_mask_loss": 0.0158,
+      "predictor_learning_rate": 3.210388941338499e-05,
+      "step": 8216
+    },
+    {
+      "epoch": 0.84,
+      "gumbel_tau": 0.1439,
+      "image_masks_loss": 0.0124,
+      "llm_learning_rate": 3.206322983225871e-07,
+      "loss": 0.7099,
+      "output_text_mask_loss": 0.0374,
+      "predictor_learning_rate": 3.206322983225871e-05,
+      "step": 8217
+    },
+    {
+      "epoch": 0.84,
+      "gumbel_tau": 0.1439,
+      "image_masks_loss": 0.0063,
+      "llm_learning_rate": 3.202259425081497e-07,
+      "loss": 0.7214,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 3.202259425081497e-05,
+      "step": 8218
+    },
+    {
+      "epoch": 0.84,
+      "gumbel_tau": 0.1439,
+      "image_masks_loss": 0.0153,
+      "llm_learning_rate": 3.1981982673528645e-07,
+      "loss": 0.7362,
+      "output_text_mask_loss": 0.0125,
+      "predictor_learning_rate": 3.198198267352864e-05,
+      "step": 8219
+    },
+    {
+      "epoch": 0.84,
+      "gumbel_tau": 0.1438,
+      "image_masks_loss": 0.0092,
+      "llm_learning_rate": 3.194139510487204e-07,
+      "loss": 0.6683,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 3.194139510487204e-05,
+      "step": 8220
+    },
+    {
+      "epoch": 0.84,
+      "gumbel_tau": 0.1438,
+      "image_masks_loss": 0.0136,
+      "llm_learning_rate": 3.1900831549314727e-07,
+      "loss": 0.787,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 3.1900831549314727e-05,
+      "step": 8221
+    },
+    {
+      "epoch": 0.84,
+      "gumbel_tau": 0.1438,
+      "image_masks_loss": 0.0071,
+      "llm_learning_rate": 3.1860292011323625e-07,
+      "loss": 0.6521,
+      "output_text_mask_loss": NaN,
+      "predictor_learning_rate": 3.186029201132362e-05,
+      "step": 8222
+    },
+    {
+      "epoch": 0.84,
+      "gumbel_tau": 0.1437,
+      "image_masks_loss": 0.0121,
+      "llm_learning_rate": 3.1819776495363053e-07,
+      "loss": 0.7855,
+      "output_text_mask_loss": 0.0094,
+      "predictor_learning_rate": 3.181977649536305e-05,
+      "step": 8223
+    },
+    {
+      "epoch": 0.84,
+      "gumbel_tau": 0.1437,
+      "image_masks_loss": 0.0087,
+      "llm_learning_rate": 3.1779285005894693e-07,
+      "loss": 0.7557,
+      "output_text_mask_loss": 0.0152,
+      "predictor_learning_rate": 3.1779285005894693e-05,
+      "step": 8224
+    },
+    {
+      "epoch": 0.84,
+      "gumbel_tau": 0.1437,
+      "image_masks_loss": 0.0133,
+      "llm_learning_rate": 3.173881754737751e-07,
+      "loss": 0.7167,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 3.173881754737751e-05,
+      "step": 8225
+    },
+    {
+      "epoch": 0.84,
+      "gumbel_tau": 0.1436,
+      "image_masks_loss": 0.0104,
+      "llm_learning_rate": 3.169837412426785e-07,
+      "loss": 0.7132,
+      "output_text_mask_loss": 0.0109,
+      "predictor_learning_rate": 3.169837412426785e-05,
+      "step": 8226
+    },
+    {
+      "epoch": 0.84,
+      "gumbel_tau": 0.1436,
+      "image_masks_loss": 0.0116,
+      "llm_learning_rate": 3.1657954741019544e-07,
+      "loss": 0.6535,
+      "output_text_mask_loss": 0.0023,
+      "predictor_learning_rate": 3.165795474101954e-05,
+      "step": 8227
+    },
+    {
+      "epoch": 0.84,
+      "gumbel_tau": 0.1436,
+      "image_masks_loss": 0.0167,
+      "llm_learning_rate": 3.161755940208358e-07,
+      "loss": 0.7271,
+      "output_text_mask_loss": 0.0259,
+      "predictor_learning_rate": 3.161755940208358e-05,
+      "step": 8228
+    },
+    {
+      "epoch": 0.84,
+      "gumbel_tau": 0.1435,
+      "image_masks_loss": 0.0155,
+      "llm_learning_rate": 3.1577188111908407e-07,
+      "loss": 0.7743,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 3.157718811190841e-05,
+      "step": 8229
+    },
+    {
+      "epoch": 0.84,
+      "gumbel_tau": 0.1435,
+      "image_masks_loss": 0.0282,
+      "llm_learning_rate": 3.1536840874939826e-07,
+      "loss": 0.7437,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 3.1536840874939825e-05,
+      "step": 8230
+    },
+    {
+      "epoch": 0.84,
+      "gumbel_tau": 0.1435,
+      "image_masks_loss": 0.0143,
+      "llm_learning_rate": 3.14965176956209e-07,
+      "loss": 0.7426,
+      "output_text_mask_loss": 0.0648,
+      "predictor_learning_rate": 3.1496517695620894e-05,
+      "step": 8231
+    },
+    {
+      "epoch": 0.84,
+      "gumbel_tau": 0.1434,
+      "image_masks_loss": 0.0094,
+      "llm_learning_rate": 3.145621857839215e-07,
+      "loss": 0.7108,
+      "output_text_mask_loss": 0.0908,
+      "predictor_learning_rate": 3.1456218578392146e-05,
+      "step": 8232
+    },
+    {
+      "epoch": 0.84,
+      "gumbel_tau": 0.1434,
+      "image_masks_loss": 0.0138,
+      "llm_learning_rate": 3.1415943527691425e-07,
+      "loss": 0.686,
+      "output_text_mask_loss": 0.0187,
+      "predictor_learning_rate": 3.141594352769142e-05,
+      "step": 8233
+    },
+    {
+      "epoch": 0.84,
+      "gumbel_tau": 0.1434,
+      "image_masks_loss": 0.007,
+      "llm_learning_rate": 3.137569254795383e-07,
+      "loss": 0.7865,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 3.137569254795383e-05,
+      "step": 8234
+    },
+    {
+      "epoch": 0.84,
+      "gumbel_tau": 0.1433,
+      "image_masks_loss": 0.0036,
+      "llm_learning_rate": 3.133546564361198e-07,
+      "loss": 0.7039,
+      "output_text_mask_loss": 0.0134,
+      "predictor_learning_rate": 3.133546564361198e-05,
+      "step": 8235
+    },
+    {
+      "epoch": 0.84,
+      "gumbel_tau": 0.1433,
+      "image_masks_loss": 0.005,
+      "llm_learning_rate": 3.12952628190957e-07,
+      "loss": 0.6831,
+      "output_text_mask_loss": 0.0104,
+      "predictor_learning_rate": 3.12952628190957e-05,
+      "step": 8236
+    },
+    {
+      "epoch": 0.84,
+      "gumbel_tau": 0.1433,
+      "image_masks_loss": 0.0099,
+      "llm_learning_rate": 3.125508407883221e-07,
+      "loss": 0.7026,
+      "output_text_mask_loss": 0.001,
+      "predictor_learning_rate": 3.1255084078832206e-05,
+      "step": 8237
+    },
+    {
+      "epoch": 0.84,
+      "gumbel_tau": 0.1432,
+      "image_masks_loss": 0.0108,
+      "llm_learning_rate": 3.1214929427246104e-07,
+      "loss": 0.6778,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 3.12149294272461e-05,
+      "step": 8238
+    },
+    {
+      "epoch": 0.84,
+      "gumbel_tau": 0.1432,
+      "image_masks_loss": 0.0058,
+      "llm_learning_rate": 3.1174798868759237e-07,
+      "loss": 0.7137,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 3.1174798868759234e-05,
+      "step": 8239
+    },
+    {
+      "epoch": 0.84,
+      "gumbel_tau": 0.1432,
+      "image_masks_loss": 0.0091,
+      "llm_learning_rate": 3.113469240779099e-07,
+      "loss": 0.8276,
+      "output_text_mask_loss": 0.0095,
+      "predictor_learning_rate": 3.1134692407790985e-05,
+      "step": 8240
+    },
+    {
+      "epoch": 0.84,
+      "gumbel_tau": 0.1431,
+      "image_masks_loss": 0.0074,
+      "llm_learning_rate": 3.109461004875791e-07,
+      "loss": 0.7041,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 3.109461004875791e-05,
+      "step": 8241
+    },
+    {
+      "epoch": 0.84,
+      "gumbel_tau": 0.1431,
+      "image_masks_loss": 0.0136,
+      "llm_learning_rate": 3.1054551796073974e-07,
+      "loss": 0.7194,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 3.105455179607397e-05,
+      "step": 8242
+    },
+    {
+      "epoch": 0.84,
+      "gumbel_tau": 0.1431,
+      "image_masks_loss": 0.0057,
+      "llm_learning_rate": 3.1014517654150475e-07,
+      "loss": 0.7561,
+      "output_text_mask_loss": 0.0353,
+      "predictor_learning_rate": 3.101451765415047e-05,
+      "step": 8243
+    },
+    {
+      "epoch": 0.84,
+      "gumbel_tau": 0.143,
+      "image_masks_loss": 0.0195,
+      "llm_learning_rate": 3.0974507627396064e-07,
+      "loss": 0.8011,
+      "output_text_mask_loss": 0.0004,
+      "predictor_learning_rate": 3.097450762739606e-05,
+      "step": 8244
+    },
+    {
+      "epoch": 0.84,
+      "gumbel_tau": 0.143,
+      "image_masks_loss": 0.0165,
+      "llm_learning_rate": 3.093452172021674e-07,
+      "loss": 0.7106,
+      "output_text_mask_loss": 0.0039,
+      "predictor_learning_rate": 3.0934521720216744e-05,
+      "step": 8245
+    },
+    {
+      "epoch": 0.84,
+      "gumbel_tau": 0.143,
+      "image_masks_loss": 0.009,
+      "llm_learning_rate": 3.0894559937015763e-07,
+      "loss": 0.7526,
+      "output_text_mask_loss": 0.0046,
+      "predictor_learning_rate": 3.089455993701576e-05,
+      "step": 8246
+    },
+    {
+      "epoch": 0.84,
+      "gumbel_tau": 0.1429,
+      "image_masks_loss": 0.018,
+      "llm_learning_rate": 3.0854622282193933e-07,
+      "loss": 0.7227,
+      "output_text_mask_loss": 0.0179,
+      "predictor_learning_rate": 3.0854622282193934e-05,
+      "step": 8247
+    },
+    {
+      "epoch": 0.85,
+      "gumbel_tau": 0.1429,
+      "image_masks_loss": 0.0092,
+      "llm_learning_rate": 3.081470876014925e-07,
+      "loss": 0.5913,
+      "output_text_mask_loss": 0.0116,
+      "predictor_learning_rate": 3.0814708760149244e-05,
+      "step": 8248
+    },
+    {
+      "epoch": 0.85,
+      "gumbel_tau": 0.1429,
+      "image_masks_loss": 0.0102,
+      "llm_learning_rate": 3.077481937527704e-07,
+      "loss": 0.7635,
+      "output_text_mask_loss": 0.0094,
+      "predictor_learning_rate": 3.077481937527704e-05,
+      "step": 8249
+    },
+    {
+      "epoch": 0.85,
+      "gumbel_tau": 0.1428,
+      "image_masks_loss": 0.0165,
+      "llm_learning_rate": 3.073495413197003e-07,
+      "loss": 0.7087,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 3.073495413197003e-05,
+      "step": 8250
+    },
+    {
+      "epoch": 0.85,
+      "gumbel_tau": 0.1428,
+      "image_masks_loss": 0.0114,
+      "llm_learning_rate": 3.069511303461825e-07,
+      "loss": 0.6567,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 3.069511303461825e-05,
+      "step": 8251
+    },
+    {
+      "epoch": 0.85,
+      "gumbel_tau": 0.1428,
+      "image_masks_loss": 0.0007,
+      "llm_learning_rate": 3.0655296087609115e-07,
+      "loss": 0.6852,
+      "output_text_mask_loss": 0.0038,
+      "predictor_learning_rate": 3.065529608760911e-05,
+      "step": 8252
+    },
+    {
+      "epoch": 0.85,
+      "gumbel_tau": 0.1427,
+      "image_masks_loss": 0.0164,
+      "llm_learning_rate": 3.0615503295327366e-07,
+      "loss": 0.6677,
+      "output_text_mask_loss": 0.041,
+      "predictor_learning_rate": 3.061550329532736e-05,
+      "step": 8253
+    },
+    {
+      "epoch": 0.85,
+      "gumbel_tau": 0.1427,
+      "image_masks_loss": 0.0122,
+      "llm_learning_rate": 3.057573466215502e-07,
+      "loss": 0.7526,
+      "output_text_mask_loss": 0.005,
+      "predictor_learning_rate": 3.057573466215502e-05,
+      "step": 8254
+    },
+    {
+      "epoch": 0.85,
+      "gumbel_tau": 0.1427,
+      "image_masks_loss": 0.0041,
+      "llm_learning_rate": 3.053599019247153e-07,
+      "loss": 0.6776,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 3.053599019247153e-05,
+      "step": 8255
+    },
+    {
+      "epoch": 0.85,
+      "gumbel_tau": 0.1426,
+      "image_masks_loss": 0.0065,
+      "llm_learning_rate": 3.049626989065366e-07,
+      "loss": 0.7457,
+      "output_text_mask_loss": 0.0387,
+      "predictor_learning_rate": 3.0496269890653657e-05,
+      "step": 8256
+    },
+    {
+      "epoch": 0.85,
+      "gumbel_tau": 0.1426,
+      "image_masks_loss": 0.0106,
+      "llm_learning_rate": 3.045657376107544e-07,
+      "loss": 0.754,
+      "output_text_mask_loss": 0.0099,
+      "predictor_learning_rate": 3.045657376107544e-05,
+      "step": 8257
+    },
+    {
+      "epoch": 0.85,
+      "gumbel_tau": 0.1426,
+      "image_masks_loss": 0.0074,
+      "llm_learning_rate": 3.04169018081083e-07,
+      "loss": 0.7419,
+      "output_text_mask_loss": 0.0216,
+      "predictor_learning_rate": 3.04169018081083e-05,
+      "step": 8258
+    },
+    {
+      "epoch": 0.85,
+      "gumbel_tau": 0.1425,
+      "image_masks_loss": 0.0108,
+      "llm_learning_rate": 3.0377254036121103e-07,
+      "loss": 0.8382,
+      "output_text_mask_loss": 0.0542,
+      "predictor_learning_rate": 3.03772540361211e-05,
+      "step": 8259
+    },
+    {
+      "epoch": 0.85,
+      "gumbel_tau": 0.1425,
+      "image_masks_loss": 0.0113,
+      "llm_learning_rate": 3.0337630449479847e-07,
+      "loss": 0.6738,
+      "output_text_mask_loss": 0.0095,
+      "predictor_learning_rate": 3.0337630449479847e-05,
+      "step": 8260
+    },
+    {
+      "epoch": 0.85,
+      "gumbel_tau": 0.1425,
+      "image_masks_loss": 0.0101,
+      "llm_learning_rate": 3.029803105254803e-07,
+      "loss": 0.6851,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 3.0298031052548025e-05,
+      "step": 8261
+    },
+    {
+      "epoch": 0.85,
+      "gumbel_tau": 0.1424,
+      "image_masks_loss": 0.0161,
+      "llm_learning_rate": 3.02584558496864e-07,
+      "loss": 0.6913,
+      "output_text_mask_loss": 0.0065,
+      "predictor_learning_rate": 3.0258455849686395e-05,
+      "step": 8262
+    },
+    {
+      "epoch": 0.85,
+      "gumbel_tau": 0.1424,
+      "image_masks_loss": 0.0041,
+      "llm_learning_rate": 3.0218904845253055e-07,
+      "loss": 0.6424,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 3.0218904845253054e-05,
+      "step": 8263
+    },
+    {
+      "epoch": 0.85,
+      "gumbel_tau": 0.1424,
+      "image_masks_loss": 0.0149,
+      "llm_learning_rate": 3.0179378043603436e-07,
+      "loss": 0.6953,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 3.0179378043603434e-05,
+      "step": 8264
+    },
+    {
+      "epoch": 0.85,
+      "gumbel_tau": 0.1423,
+      "image_masks_loss": 0.0059,
+      "llm_learning_rate": 3.0139875449090305e-07,
+      "loss": 0.7422,
+      "output_text_mask_loss": 0.0332,
+      "predictor_learning_rate": 3.0139875449090303e-05,
+      "step": 8265
+    },
+    {
+      "epoch": 0.85,
+      "gumbel_tau": 0.1423,
+      "image_masks_loss": 0.0063,
+      "llm_learning_rate": 3.010039706606388e-07,
+      "loss": 0.7223,
+      "output_text_mask_loss": 0.0057,
+      "predictor_learning_rate": 3.010039706606388e-05,
+      "step": 8266
+    },
+    {
+      "epoch": 0.85,
+      "gumbel_tau": 0.1423,
+      "image_masks_loss": 0.0057,
+      "llm_learning_rate": 3.006094289887151e-07,
+      "loss": 0.6905,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 3.0060942898871513e-05,
+      "step": 8267
+    },
+    {
+      "epoch": 0.85,
+      "gumbel_tau": 0.1422,
+      "image_masks_loss": 0.0087,
+      "llm_learning_rate": 3.0021512951858013e-07,
+      "loss": 0.7362,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 3.002151295185801e-05,
+      "step": 8268
+    },
+    {
+      "epoch": 0.85,
+      "gumbel_tau": 0.1422,
+      "image_masks_loss": 0.0055,
+      "llm_learning_rate": 2.998210722936554e-07,
+      "loss": 0.7445,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 2.9982107229365536e-05,
+      "step": 8269
+    },
+    {
+      "epoch": 0.85,
+      "gumbel_tau": 0.1422,
+      "image_masks_loss": 0.0062,
+      "llm_learning_rate": 2.9942725735733454e-07,
+      "loss": 0.6375,
+      "output_text_mask_loss": 0.0019,
+      "predictor_learning_rate": 2.994272573573345e-05,
+      "step": 8270
+    },
+    {
+      "epoch": 0.85,
+      "gumbel_tau": 0.1421,
+      "image_masks_loss": 0.014,
+      "llm_learning_rate": 2.9903368475298484e-07,
+      "loss": 0.7224,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 2.9903368475298483e-05,
+      "step": 8271
+    },
+    {
+      "epoch": 0.85,
+      "gumbel_tau": 0.1421,
+      "image_masks_loss": 0.0017,
+      "llm_learning_rate": 2.986403545239491e-07,
+      "loss": 0.7627,
+      "output_text_mask_loss": 0.0139,
+      "predictor_learning_rate": 2.986403545239491e-05,
+      "step": 8272
+    },
+    {
+      "epoch": 0.85,
+      "gumbel_tau": 0.1421,
+      "image_masks_loss": 0.0135,
+      "llm_learning_rate": 2.9824726671354055e-07,
+      "loss": 0.6724,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 2.9824726671354052e-05,
+      "step": 8273
+    },
+    {
+      "epoch": 0.85,
+      "gumbel_tau": 0.142,
+      "image_masks_loss": 0.0132,
+      "llm_learning_rate": 2.978544213650475e-07,
+      "loss": 0.6687,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 2.978544213650475e-05,
+      "step": 8274
+    },
+    {
+      "epoch": 0.85,
+      "gumbel_tau": 0.142,
+      "image_masks_loss": 0.0114,
+      "llm_learning_rate": 2.974618185217304e-07,
+      "loss": 0.8003,
+      "output_text_mask_loss": 0.0175,
+      "predictor_learning_rate": 2.974618185217304e-05,
+      "step": 8275
+    },
+    {
+      "epoch": 0.85,
+      "gumbel_tau": 0.142,
+      "image_masks_loss": 0.0111,
+      "llm_learning_rate": 2.970694582268241e-07,
+      "loss": 0.7562,
+      "output_text_mask_loss": 0.0162,
+      "predictor_learning_rate": 2.970694582268241e-05,
+      "step": 8276
+    },
+    {
+      "epoch": 0.85,
+      "gumbel_tau": 0.1419,
+      "image_masks_loss": 0.0139,
+      "llm_learning_rate": 2.966773405235357e-07,
+      "loss": 0.7411,
+      "output_text_mask_loss": 0.0063,
+      "predictor_learning_rate": 2.966773405235357e-05,
+      "step": 8277
+    },
+    {
+      "epoch": 0.85,
+      "gumbel_tau": 0.1419,
+      "image_masks_loss": 0.0129,
+      "llm_learning_rate": 2.9628546545504566e-07,
+      "loss": 0.7181,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 2.9628546545504565e-05,
+      "step": 8278
+    },
+    {
+      "epoch": 0.85,
+      "gumbel_tau": 0.1419,
+      "image_masks_loss": 0.0154,
+      "llm_learning_rate": 2.9589383306450947e-07,
+      "loss": 0.7068,
+      "output_text_mask_loss": 0.0679,
+      "predictor_learning_rate": 2.9589383306450945e-05,
+      "step": 8279
+    },
+    {
+      "epoch": 0.85,
+      "gumbel_tau": 0.1418,
+      "image_masks_loss": 0.0078,
+      "llm_learning_rate": 2.955024433950537e-07,
+      "loss": 0.753,
+      "output_text_mask_loss": 0.0322,
+      "predictor_learning_rate": 2.955024433950537e-05,
+      "step": 8280
+    },
+    {
+      "epoch": 0.85,
+      "gumbel_tau": 0.1418,
+      "image_masks_loss": 0.0172,
+      "llm_learning_rate": 2.9511129648977927e-07,
+      "loss": 0.7753,
+      "output_text_mask_loss": 0.0259,
+      "predictor_learning_rate": 2.9511129648977923e-05,
+      "step": 8281
+    },
+    {
+      "epoch": 0.85,
+      "gumbel_tau": 0.1418,
+      "image_masks_loss": 0.013,
+      "llm_learning_rate": 2.947203923917599e-07,
+      "loss": 0.7024,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 2.947203923917599e-05,
+      "step": 8282
+    },
+    {
+      "epoch": 0.85,
+      "gumbel_tau": 0.1417,
+      "image_masks_loss": 0.0278,
+      "llm_learning_rate": 2.9432973114404324e-07,
+      "loss": 0.6565,
+      "output_text_mask_loss": 0.0437,
+      "predictor_learning_rate": 2.9432973114404322e-05,
+      "step": 8283
+    },
+    {
+      "epoch": 0.85,
+      "gumbel_tau": 0.1417,
+      "image_masks_loss": 0.0134,
+      "llm_learning_rate": 2.939393127896492e-07,
+      "loss": 0.7967,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 2.939393127896492e-05,
+      "step": 8284
+    },
+    {
+      "epoch": 0.85,
+      "gumbel_tau": 0.1417,
+      "image_masks_loss": 0.0043,
+      "llm_learning_rate": 2.935491373715715e-07,
+      "loss": 0.7335,
+      "output_text_mask_loss": 0.0035,
+      "predictor_learning_rate": 2.9354913737157153e-05,
+      "step": 8285
+    },
+    {
+      "epoch": 0.85,
+      "gumbel_tau": 0.1416,
+      "image_masks_loss": 0.0088,
+      "llm_learning_rate": 2.931592049327781e-07,
+      "loss": 0.7372,
+      "output_text_mask_loss": 0.0012,
+      "predictor_learning_rate": 2.9315920493277802e-05,
+      "step": 8286
+    },
+    {
+      "epoch": 0.85,
+      "gumbel_tau": 0.1416,
+      "image_masks_loss": 0.0038,
+      "llm_learning_rate": 2.927695155162086e-07,
+      "loss": 0.6499,
+      "output_text_mask_loss": 0.0124,
+      "predictor_learning_rate": 2.9276951551620856e-05,
+      "step": 8287
+    },
+    {
+      "epoch": 0.85,
+      "gumbel_tau": 0.1416,
+      "image_masks_loss": 0.0057,
+      "llm_learning_rate": 2.9238006916477683e-07,
+      "loss": 0.7123,
+      "output_text_mask_loss": 0.0038,
+      "predictor_learning_rate": 2.9238006916477682e-05,
+      "step": 8288
+    },
+    {
+      "epoch": 0.85,
+      "gumbel_tau": 0.1415,
+      "image_masks_loss": 0.0139,
+      "llm_learning_rate": 2.9199086592136853e-07,
+      "loss": 0.7172,
+      "output_text_mask_loss": 0.0107,
+      "predictor_learning_rate": 2.9199086592136852e-05,
+      "step": 8289
+    },
+    {
+      "epoch": 0.85,
+      "gumbel_tau": 0.1415,
+      "image_masks_loss": 0.0061,
+      "llm_learning_rate": 2.9160190582884383e-07,
+      "loss": 0.6909,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 2.9160190582884383e-05,
+      "step": 8290
+    },
+    {
+      "epoch": 0.85,
+      "gumbel_tau": 0.1415,
+      "image_masks_loss": 0.003,
+      "llm_learning_rate": 2.912131889300368e-07,
+      "loss": 0.8316,
+      "output_text_mask_loss": 0.0561,
+      "predictor_learning_rate": 2.912131889300368e-05,
+      "step": 8291
+    },
+    {
+      "epoch": 0.85,
+      "gumbel_tau": 0.1414,
+      "image_masks_loss": 0.0121,
+      "llm_learning_rate": 2.908247152677535e-07,
+      "loss": 0.7236,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 2.9082471526775345e-05,
+      "step": 8292
+    },
+    {
+      "epoch": 0.85,
+      "gumbel_tau": 0.1414,
+      "image_masks_loss": 0.0102,
+      "llm_learning_rate": 2.9043648488477293e-07,
+      "loss": 0.7388,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 2.9043648488477293e-05,
+      "step": 8293
+    },
+    {
+      "epoch": 0.85,
+      "gumbel_tau": 0.1414,
+      "image_masks_loss": 0.0161,
+      "llm_learning_rate": 2.900484978238485e-07,
+      "loss": 0.7886,
+      "output_text_mask_loss": 0.0082,
+      "predictor_learning_rate": 2.900484978238485e-05,
+      "step": 8294
+    },
+    {
+      "epoch": 0.85,
+      "gumbel_tau": 0.1413,
+      "image_masks_loss": 0.0067,
+      "llm_learning_rate": 2.896607541277058e-07,
+      "loss": 0.6946,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 2.8966075412770575e-05,
+      "step": 8295
+    },
+    {
+      "epoch": 0.85,
+      "gumbel_tau": 0.1413,
+      "image_masks_loss": 0.0147,
+      "llm_learning_rate": 2.892732538390444e-07,
+      "loss": 0.7523,
+      "output_text_mask_loss": 0.005,
+      "predictor_learning_rate": 2.892732538390444e-05,
+      "step": 8296
+    },
+    {
+      "epoch": 0.85,
+      "gumbel_tau": 0.1413,
+      "image_masks_loss": 0.01,
+      "llm_learning_rate": 2.888859970005356e-07,
+      "loss": 0.7747,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 2.8888599700053558e-05,
+      "step": 8297
+    },
+    {
+      "epoch": 0.85,
+      "gumbel_tau": 0.1412,
+      "image_masks_loss": 0.0026,
+      "llm_learning_rate": 2.8849898365482686e-07,
+      "loss": 0.6194,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 2.8849898365482686e-05,
+      "step": 8298
+    },
+    {
+      "epoch": 0.85,
+      "gumbel_tau": 0.1412,
+      "image_masks_loss": 0.0032,
+      "llm_learning_rate": 2.881122138445355e-07,
+      "loss": 0.7123,
+      "output_text_mask_loss": 0.0039,
+      "predictor_learning_rate": 2.881122138445355e-05,
+      "step": 8299
+    },
+    {
+      "epoch": 0.85,
+      "gumbel_tau": 0.1412,
+      "image_masks_loss": 0.0079,
+      "llm_learning_rate": 2.877256876122542e-07,
+      "loss": 0.7131,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 2.8772568761225415e-05,
+      "step": 8300
+    },
+    {
+      "epoch": 0.85,
+      "gumbel_tau": 0.1411,
+      "image_masks_loss": 0.0123,
+      "llm_learning_rate": 2.873394050005479e-07,
+      "loss": 0.7006,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 2.873394050005479e-05,
+      "step": 8301
+    },
+    {
+      "epoch": 0.85,
+      "gumbel_tau": 0.1411,
+      "image_masks_loss": 0.014,
+      "llm_learning_rate": 2.8695336605195466e-07,
+      "loss": 0.7336,
+      "output_text_mask_loss": 0.0077,
+      "predictor_learning_rate": 2.8695336605195467e-05,
+      "step": 8302
+    },
+    {
+      "epoch": 0.85,
+      "gumbel_tau": 0.1411,
+      "image_masks_loss": 0.007,
+      "llm_learning_rate": 2.8656757080898627e-07,
+      "loss": 0.6964,
+      "output_text_mask_loss": 0.0088,
+      "predictor_learning_rate": 2.8656757080898628e-05,
+      "step": 8303
+    },
+    {
+      "epoch": 0.85,
+      "gumbel_tau": 0.141,
+      "image_masks_loss": 0.0091,
+      "llm_learning_rate": 2.861820193141265e-07,
+      "loss": 0.8073,
+      "output_text_mask_loss": 0.0065,
+      "predictor_learning_rate": 2.861820193141265e-05,
+      "step": 8304
+    },
+    {
+      "epoch": 0.85,
+      "gumbel_tau": 0.141,
+      "image_masks_loss": 0.0175,
+      "llm_learning_rate": 2.857967116098345e-07,
+      "loss": 0.7488,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 2.857967116098345e-05,
+      "step": 8305
+    },
+    {
+      "epoch": 0.85,
+      "gumbel_tau": 0.141,
+      "image_masks_loss": 0.0093,
+      "llm_learning_rate": 2.85411647738541e-07,
+      "loss": 0.7228,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 2.85411647738541e-05,
+      "step": 8306
+    },
+    {
+      "epoch": 0.85,
+      "gumbel_tau": 0.1409,
+      "image_masks_loss": 0.0144,
+      "llm_learning_rate": 2.8502682774264935e-07,
+      "loss": 0.5893,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 2.850268277426493e-05,
+      "step": 8307
+    },
+    {
+      "epoch": 0.85,
+      "gumbel_tau": 0.1409,
+      "image_masks_loss": 0.012,
+      "llm_learning_rate": 2.8464225166453695e-07,
+      "loss": 0.7369,
+      "output_text_mask_loss": 0.0019,
+      "predictor_learning_rate": 2.8464225166453693e-05,
+      "step": 8308
+    },
+    {
+      "epoch": 0.85,
+      "gumbel_tau": 0.1409,
+      "image_masks_loss": 0.0096,
+      "llm_learning_rate": 2.842579195465542e-07,
+      "loss": 0.7593,
+      "output_text_mask_loss": 0.0557,
+      "predictor_learning_rate": 2.8425791954655416e-05,
+      "step": 8309
+    },
+    {
+      "epoch": 0.85,
+      "gumbel_tau": 0.1408,
+      "image_masks_loss": 0.0117,
+      "llm_learning_rate": 2.8387383143102435e-07,
+      "loss": 0.696,
+      "output_text_mask_loss": 0.0305,
+      "predictor_learning_rate": 2.8387383143102434e-05,
+      "step": 8310
+    },
+    {
+      "epoch": 0.85,
+      "gumbel_tau": 0.1408,
+      "image_masks_loss": 0.0139,
+      "llm_learning_rate": 2.834899873602448e-07,
+      "loss": 0.775,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 2.834899873602448e-05,
+      "step": 8311
+    },
+    {
+      "epoch": 0.85,
+      "gumbel_tau": 0.1408,
+      "image_masks_loss": 0.0079,
+      "llm_learning_rate": 2.83106387376485e-07,
+      "loss": 0.6836,
+      "output_text_mask_loss": 0.0341,
+      "predictor_learning_rate": 2.83106387376485e-05,
+      "step": 8312
+    },
+    {
+      "epoch": 0.85,
+      "gumbel_tau": 0.1407,
+      "image_masks_loss": 0.0134,
+      "llm_learning_rate": 2.8272303152198783e-07,
+      "loss": 0.8006,
+      "output_text_mask_loss": 0.0305,
+      "predictor_learning_rate": 2.827230315219878e-05,
+      "step": 8313
+    },
+    {
+      "epoch": 0.85,
+      "gumbel_tau": 0.1407,
+      "image_masks_loss": 0.0182,
+      "llm_learning_rate": 2.8233991983896924e-07,
+      "loss": 0.7679,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 2.8233991983896924e-05,
+      "step": 8314
+    },
+    {
+      "epoch": 0.85,
+      "gumbel_tau": 0.1407,
+      "image_masks_loss": 0.0182,
+      "llm_learning_rate": 2.819570523696183e-07,
+      "loss": 0.7727,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 2.8195705236961826e-05,
+      "step": 8315
+    },
+    {
+      "epoch": 0.85,
+      "gumbel_tau": 0.1406,
+      "image_masks_loss": 0.0064,
+      "llm_learning_rate": 2.815744291560971e-07,
+      "loss": 0.6754,
+      "output_text_mask_loss": 0.0071,
+      "predictor_learning_rate": 2.8157442915609704e-05,
+      "step": 8316
+    },
+    {
+      "epoch": 0.85,
+      "gumbel_tau": 0.1406,
+      "image_masks_loss": 0.0055,
+      "llm_learning_rate": 2.8119205024054075e-07,
+      "loss": 0.6629,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 2.8119205024054073e-05,
+      "step": 8317
+    },
+    {
+      "epoch": 0.85,
+      "gumbel_tau": 0.1406,
+      "image_masks_loss": 0.0257,
+      "llm_learning_rate": 2.8080991566505845e-07,
+      "loss": 0.6732,
+      "output_text_mask_loss": 0.0041,
+      "predictor_learning_rate": 2.8080991566505843e-05,
+      "step": 8318
+    },
+    {
+      "epoch": 0.85,
+      "gumbel_tau": 0.1405,
+      "image_masks_loss": 0.0034,
+      "llm_learning_rate": 2.804280254717312e-07,
+      "loss": 0.7283,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 2.8042802547173117e-05,
+      "step": 8319
+    },
+    {
+      "epoch": 0.85,
+      "gumbel_tau": 0.1405,
+      "image_masks_loss": 0.0032,
+      "llm_learning_rate": 2.8004637970261374e-07,
+      "loss": 0.7485,
+      "output_text_mask_loss": 0.0117,
+      "predictor_learning_rate": 2.800463797026137e-05,
+      "step": 8320
+    },
+    {
+      "epoch": 0.85,
+      "gumbel_tau": 0.1405,
+      "image_masks_loss": 0.013,
+      "llm_learning_rate": 2.7966497839973376e-07,
+      "loss": 0.6801,
+      "output_text_mask_loss": 0.0231,
+      "predictor_learning_rate": 2.7966497839973377e-05,
+      "step": 8321
+    },
+    {
+      "epoch": 0.85,
+      "gumbel_tau": 0.1404,
+      "image_masks_loss": 0.0106,
+      "llm_learning_rate": 2.7928382160509195e-07,
+      "loss": 0.6933,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 2.7928382160509192e-05,
+      "step": 8322
+    },
+    {
+      "epoch": 0.85,
+      "gumbel_tau": 0.1404,
+      "image_masks_loss": 0.0191,
+      "llm_learning_rate": 2.789029093606624e-07,
+      "loss": 0.7346,
+      "output_text_mask_loss": 0.0012,
+      "predictor_learning_rate": 2.7890290936066236e-05,
+      "step": 8323
+    },
+    {
+      "epoch": 0.85,
+      "gumbel_tau": 0.1404,
+      "image_masks_loss": 0.0091,
+      "llm_learning_rate": 2.7852224170839076e-07,
+      "loss": 0.7866,
+      "output_text_mask_loss": 0.0156,
+      "predictor_learning_rate": 2.7852224170839074e-05,
+      "step": 8324
+    },
+    {
+      "epoch": 0.85,
+      "gumbel_tau": 0.1403,
+      "image_masks_loss": 0.0081,
+      "llm_learning_rate": 2.7814181869019935e-07,
+      "loss": 0.696,
+      "output_text_mask_loss": 0.0137,
+      "predictor_learning_rate": 2.7814181869019934e-05,
+      "step": 8325
+    },
+    {
+      "epoch": 0.85,
+      "gumbel_tau": 0.1403,
+      "image_masks_loss": 0.0077,
+      "llm_learning_rate": 2.777616403479794e-07,
+      "loss": 0.7432,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 2.7776164034797935e-05,
+      "step": 8326
+    },
+    {
+      "epoch": 0.85,
+      "gumbel_tau": 0.1403,
+      "image_masks_loss": 0.0053,
+      "llm_learning_rate": 2.7738170672359733e-07,
+      "loss": 0.791,
+      "output_text_mask_loss": 0.0326,
+      "predictor_learning_rate": 2.773817067235973e-05,
+      "step": 8327
+    },
+    {
+      "epoch": 0.85,
+      "gumbel_tau": 0.1402,
+      "image_masks_loss": 0.0045,
+      "llm_learning_rate": 2.7700201785889254e-07,
+      "loss": 0.6893,
+      "output_text_mask_loss": 0.001,
+      "predictor_learning_rate": 2.770020178588925e-05,
+      "step": 8328
+    },
+    {
+      "epoch": 0.85,
+      "gumbel_tau": 0.1402,
+      "image_masks_loss": 0.0083,
+      "llm_learning_rate": 2.7662257379567624e-07,
+      "loss": 0.7148,
+      "output_text_mask_loss": 0.0019,
+      "predictor_learning_rate": 2.7662257379567624e-05,
+      "step": 8329
+    },
+    {
+      "epoch": 0.85,
+      "gumbel_tau": 0.1402,
+      "image_masks_loss": 0.0104,
+      "llm_learning_rate": 2.7624337457573534e-07,
+      "loss": 0.6459,
+      "output_text_mask_loss": 0.018,
+      "predictor_learning_rate": 2.7624337457573534e-05,
+      "step": 8330
+    },
+    {
+      "epoch": 0.85,
+      "gumbel_tau": 0.1401,
+      "image_masks_loss": 0.0118,
+      "llm_learning_rate": 2.75864420240827e-07,
+      "loss": 0.7471,
+      "output_text_mask_loss": 0.0187,
+      "predictor_learning_rate": 2.7586442024082697e-05,
+      "step": 8331
+    },
+    {
+      "epoch": 0.85,
+      "gumbel_tau": 0.1401,
+      "image_masks_loss": 0.0132,
+      "llm_learning_rate": 2.7548571083268303e-07,
+      "loss": 0.7775,
+      "output_text_mask_loss": 0.0664,
+      "predictor_learning_rate": 2.7548571083268302e-05,
+      "step": 8332
+    },
+    {
+      "epoch": 0.85,
+      "gumbel_tau": 0.1401,
+      "image_masks_loss": 0.0207,
+      "llm_learning_rate": 2.751072463930074e-07,
+      "loss": 0.731,
+      "output_text_mask_loss": 0.0538,
+      "predictor_learning_rate": 2.751072463930074e-05,
+      "step": 8333
+    },
+    {
+      "epoch": 0.85,
+      "gumbel_tau": 0.14,
+      "image_masks_loss": 0.0126,
+      "llm_learning_rate": 2.747290269634775e-07,
+      "loss": 0.7523,
+      "output_text_mask_loss": 0.0058,
+      "predictor_learning_rate": 2.7472902696347752e-05,
+      "step": 8334
+    },
+    {
+      "epoch": 0.85,
+      "gumbel_tau": 0.14,
+      "image_masks_loss": 0.0125,
+      "llm_learning_rate": 2.7435105258574374e-07,
+      "loss": 0.7157,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 2.7435105258574373e-05,
+      "step": 8335
+    },
+    {
+      "epoch": 0.85,
+      "gumbel_tau": 0.14,
+      "image_masks_loss": 0.0131,
+      "llm_learning_rate": 2.739733233014291e-07,
+      "loss": 0.7877,
+      "output_text_mask_loss": 0.0446,
+      "predictor_learning_rate": 2.739733233014291e-05,
+      "step": 8336
+    },
+    {
+      "epoch": 0.85,
+      "gumbel_tau": 0.1399,
+      "image_masks_loss": 0.0192,
+      "llm_learning_rate": 2.7359583915213067e-07,
+      "loss": 0.7247,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 2.7359583915213067e-05,
+      "step": 8337
+    },
+    {
+      "epoch": 0.85,
+      "gumbel_tau": 0.1399,
+      "image_masks_loss": 0.0086,
+      "llm_learning_rate": 2.7321860017941795e-07,
+      "loss": 0.594,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 2.732186001794179e-05,
+      "step": 8338
+    },
+    {
+      "epoch": 0.85,
+      "gumbel_tau": 0.1399,
+      "image_masks_loss": 0.0093,
+      "llm_learning_rate": 2.7284160642483266e-07,
+      "loss": 0.7385,
+      "output_text_mask_loss": 0.0149,
+      "predictor_learning_rate": 2.7284160642483262e-05,
+      "step": 8339
+    },
+    {
+      "epoch": 0.85,
+      "gumbel_tau": 0.1398,
+      "image_masks_loss": 0.0103,
+      "llm_learning_rate": 2.724648579298902e-07,
+      "loss": 0.7395,
+      "output_text_mask_loss": 0.0767,
+      "predictor_learning_rate": 2.7246485792989023e-05,
+      "step": 8340
+    },
+    {
+      "epoch": 0.85,
+      "gumbel_tau": 0.1398,
+      "image_masks_loss": 0.0069,
+      "llm_learning_rate": 2.7208835473607974e-07,
+      "loss": 0.6879,
+      "output_text_mask_loss": 0.0071,
+      "predictor_learning_rate": 2.7208835473607975e-05,
+      "step": 8341
+    },
+    {
+      "epoch": 0.85,
+      "gumbel_tau": 0.1398,
+      "image_masks_loss": 0.0097,
+      "llm_learning_rate": 2.717120968848619e-07,
+      "loss": 0.7226,
+      "output_text_mask_loss": 0.0212,
+      "predictor_learning_rate": 2.7171209688486187e-05,
+      "step": 8342
+    },
+    {
+      "epoch": 0.85,
+      "gumbel_tau": 0.1397,
+      "image_masks_loss": 0.0119,
+      "llm_learning_rate": 2.713360844176713e-07,
+      "loss": 0.7701,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 2.7133608441767126e-05,
+      "step": 8343
+    },
+    {
+      "epoch": 0.85,
+      "gumbel_tau": 0.1397,
+      "image_masks_loss": 0.0177,
+      "llm_learning_rate": 2.709603173759151e-07,
+      "loss": 0.7706,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 2.7096031737591507e-05,
+      "step": 8344
+    },
+    {
+      "epoch": 0.86,
+      "gumbel_tau": 0.1397,
+      "image_masks_loss": 0.0034,
+      "llm_learning_rate": 2.7058479580097384e-07,
+      "loss": 0.7657,
+      "output_text_mask_loss": 0.0034,
+      "predictor_learning_rate": 2.7058479580097383e-05,
+      "step": 8345
+    },
+    {
+      "epoch": 0.86,
+      "gumbel_tau": 0.1396,
+      "image_masks_loss": 0.0203,
+      "llm_learning_rate": 2.7020951973420084e-07,
+      "loss": 0.6748,
+      "output_text_mask_loss": 0.0034,
+      "predictor_learning_rate": 2.7020951973420082e-05,
+      "step": 8346
+    },
+    {
+      "epoch": 0.86,
+      "gumbel_tau": 0.1396,
+      "image_masks_loss": 0.013,
+      "llm_learning_rate": 2.698344892169222e-07,
+      "loss": 0.7752,
+      "output_text_mask_loss": 0.0248,
+      "predictor_learning_rate": 2.6983448921692218e-05,
+      "step": 8347
+    },
+    {
+      "epoch": 0.86,
+      "gumbel_tau": 0.1396,
+      "image_masks_loss": 0.0252,
+      "llm_learning_rate": 2.6945970429043676e-07,
+      "loss": 0.7051,
+      "output_text_mask_loss": 0.007,
+      "predictor_learning_rate": 2.6945970429043676e-05,
+      "step": 8348
+    },
+    {
+      "epoch": 0.86,
+      "gumbel_tau": 0.1395,
+      "image_masks_loss": 0.0139,
+      "llm_learning_rate": 2.690851649960169e-07,
+      "loss": 0.6804,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 2.6908516499601688e-05,
+      "step": 8349
+    },
+    {
+      "epoch": 0.86,
+      "gumbel_tau": 0.1395,
+      "image_masks_loss": 0.0099,
+      "llm_learning_rate": 2.687108713749081e-07,
+      "loss": 0.738,
+      "output_text_mask_loss": 0.0012,
+      "predictor_learning_rate": 2.6871087137490814e-05,
+      "step": 8350
+    },
+    {
+      "epoch": 0.86,
+      "gumbel_tau": 0.1395,
+      "image_masks_loss": 0.0088,
+      "llm_learning_rate": 2.683368234683281e-07,
+      "loss": 0.715,
+      "output_text_mask_loss": 0.0126,
+      "predictor_learning_rate": 2.683368234683281e-05,
+      "step": 8351
+    },
+    {
+      "epoch": 0.86,
+      "gumbel_tau": 0.1394,
+      "image_masks_loss": 0.0223,
+      "llm_learning_rate": 2.6796302131746834e-07,
+      "loss": 0.8032,
+      "output_text_mask_loss": 0.0238,
+      "predictor_learning_rate": 2.679630213174683e-05,
+      "step": 8352
+    },
+    {
+      "epoch": 0.86,
+      "gumbel_tau": 0.1394,
+      "image_masks_loss": 0.0119,
+      "llm_learning_rate": 2.675894649634919e-07,
+      "loss": 0.8079,
+      "output_text_mask_loss": 0.0019,
+      "predictor_learning_rate": 2.675894649634919e-05,
+      "step": 8353
+    },
+    {
+      "epoch": 0.86,
+      "gumbel_tau": 0.1394,
+      "image_masks_loss": 0.0078,
+      "llm_learning_rate": 2.6721615444753627e-07,
+      "loss": 0.7982,
+      "output_text_mask_loss": 0.0113,
+      "predictor_learning_rate": 2.672161544475363e-05,
+      "step": 8354
+    },
+    {
+      "epoch": 0.86,
+      "gumbel_tau": 0.1393,
+      "image_masks_loss": 0.0074,
+      "llm_learning_rate": 2.6684308981071107e-07,
+      "loss": 0.7414,
+      "output_text_mask_loss": 0.0198,
+      "predictor_learning_rate": 2.6684308981071103e-05,
+      "step": 8355
+    },
+    {
+      "epoch": 0.86,
+      "gumbel_tau": 0.1393,
+      "image_masks_loss": 0.0174,
+      "llm_learning_rate": 2.664702710940983e-07,
+      "loss": 0.714,
+      "output_text_mask_loss": 0.0343,
+      "predictor_learning_rate": 2.6647027109409834e-05,
+      "step": 8356
+    },
+    {
+      "epoch": 0.86,
+      "gumbel_tau": 0.1393,
+      "image_masks_loss": 0.0165,
+      "llm_learning_rate": 2.6609769833875505e-07,
+      "loss": 0.7769,
+      "output_text_mask_loss": 0.0121,
+      "predictor_learning_rate": 2.6609769833875502e-05,
+      "step": 8357
+    },
+    {
+      "epoch": 0.86,
+      "gumbel_tau": 0.1392,
+      "image_masks_loss": 0.0326,
+      "llm_learning_rate": 2.6572537158570886e-07,
+      "loss": 0.8198,
+      "output_text_mask_loss": 0.053,
+      "predictor_learning_rate": 2.6572537158570888e-05,
+      "step": 8358
+    },
+    {
+      "epoch": 0.86,
+      "gumbel_tau": 0.1392,
+      "image_masks_loss": 0.0195,
+      "llm_learning_rate": 2.653532908759615e-07,
+      "loss": 0.8036,
+      "output_text_mask_loss": 0.0511,
+      "predictor_learning_rate": 2.653532908759615e-05,
+      "step": 8359
+    },
+    {
+      "epoch": 0.86,
+      "gumbel_tau": 0.1392,
+      "image_masks_loss": 0.0067,
+      "llm_learning_rate": 2.649814562504871e-07,
+      "loss": 0.7288,
+      "output_text_mask_loss": 0.0055,
+      "predictor_learning_rate": 2.6498145625048708e-05,
+      "step": 8360
+    },
+    {
+      "epoch": 0.86,
+      "gumbel_tau": 0.1391,
+      "image_masks_loss": 0.0102,
+      "llm_learning_rate": 2.6460986775023315e-07,
+      "loss": 0.791,
+      "output_text_mask_loss": 0.0071,
+      "predictor_learning_rate": 2.6460986775023317e-05,
+      "step": 8361
+    },
+    {
+      "epoch": 0.86,
+      "gumbel_tau": 0.1391,
+      "image_masks_loss": 0.007,
+      "llm_learning_rate": 2.642385254161198e-07,
+      "loss": 0.7272,
+      "output_text_mask_loss": 0.0202,
+      "predictor_learning_rate": 2.6423852541611977e-05,
+      "step": 8362
+    },
+    {
+      "epoch": 0.86,
+      "gumbel_tau": 0.1391,
+      "image_masks_loss": 0.0015,
+      "llm_learning_rate": 2.6386742928903976e-07,
+      "loss": 0.6484,
+      "output_text_mask_loss": 0.0172,
+      "predictor_learning_rate": 2.638674292890397e-05,
+      "step": 8363
+    },
+    {
+      "epoch": 0.86,
+      "gumbel_tau": 0.139,
+      "image_masks_loss": 0.0068,
+      "llm_learning_rate": 2.634965794098593e-07,
+      "loss": 0.7641,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 2.634965794098593e-05,
+      "step": 8364
+    },
+    {
+      "epoch": 0.86,
+      "gumbel_tau": 0.139,
+      "image_masks_loss": 0.0032,
+      "llm_learning_rate": 2.6312597581941714e-07,
+      "loss": 0.73,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 2.631259758194171e-05,
+      "step": 8365
+    },
+    {
+      "epoch": 0.86,
+      "gumbel_tau": 0.139,
+      "image_masks_loss": 0.0179,
+      "llm_learning_rate": 2.6275561855852505e-07,
+      "loss": 0.7442,
+      "output_text_mask_loss": 0.0004,
+      "predictor_learning_rate": 2.6275561855852505e-05,
+      "step": 8366
+    },
+    {
+      "epoch": 0.86,
+      "gumbel_tau": 0.1389,
+      "image_masks_loss": 0.0128,
+      "llm_learning_rate": 2.623855076679674e-07,
+      "loss": 0.7484,
+      "output_text_mask_loss": 0.0143,
+      "predictor_learning_rate": 2.6238550766796738e-05,
+      "step": 8367
+    },
+    {
+      "epoch": 0.86,
+      "gumbel_tau": 0.1389,
+      "image_masks_loss": 0.0116,
+      "llm_learning_rate": 2.6201564318850124e-07,
+      "loss": 0.6772,
+      "output_text_mask_loss": 0.071,
+      "predictor_learning_rate": 2.6201564318850125e-05,
+      "step": 8368
+    },
+    {
+      "epoch": 0.86,
+      "gumbel_tau": 0.1389,
+      "image_masks_loss": 0.0161,
+      "llm_learning_rate": 2.6164602516085796e-07,
+      "loss": 0.6779,
+      "output_text_mask_loss": 0.006,
+      "predictor_learning_rate": 2.6164602516085796e-05,
+      "step": 8369
+    },
+    {
+      "epoch": 0.86,
+      "gumbel_tau": 0.1388,
+      "image_masks_loss": 0.0051,
+      "llm_learning_rate": 2.6127665362573993e-07,
+      "loss": 0.6802,
+      "output_text_mask_loss": 0.0044,
+      "predictor_learning_rate": 2.6127665362573993e-05,
+      "step": 8370
+    },
+    {
+      "epoch": 0.86,
+      "gumbel_tau": 0.1388,
+      "image_masks_loss": 0.0073,
+      "llm_learning_rate": 2.6090752862382357e-07,
+      "loss": 0.7044,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 2.6090752862382356e-05,
+      "step": 8371
+    },
+    {
+      "epoch": 0.86,
+      "gumbel_tau": 0.1388,
+      "image_masks_loss": 0.02,
+      "llm_learning_rate": 2.605386501957577e-07,
+      "loss": 0.7284,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 2.6053865019575767e-05,
+      "step": 8372
+    },
+    {
+      "epoch": 0.86,
+      "gumbel_tau": 0.1387,
+      "image_masks_loss": 0.0077,
+      "llm_learning_rate": 2.601700183821634e-07,
+      "loss": 0.6729,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 2.6017001838216342e-05,
+      "step": 8373
+    },
+    {
+      "epoch": 0.86,
+      "gumbel_tau": 0.1387,
+      "image_masks_loss": 0.0194,
+      "llm_learning_rate": 2.598016332236361e-07,
+      "loss": 0.7331,
+      "output_text_mask_loss": 0.0111,
+      "predictor_learning_rate": 2.598016332236361e-05,
+      "step": 8374
+    },
+    {
+      "epoch": 0.86,
+      "gumbel_tau": 0.1387,
+      "image_masks_loss": 0.0235,
+      "llm_learning_rate": 2.594334947607427e-07,
+      "loss": 0.7137,
+      "output_text_mask_loss": 0.0238,
+      "predictor_learning_rate": 2.594334947607427e-05,
+      "step": 8375
+    },
+    {
+      "epoch": 0.86,
+      "gumbel_tau": 0.1386,
+      "image_masks_loss": 0.0212,
+      "llm_learning_rate": 2.59065603034023e-07,
+      "loss": 0.6997,
+      "output_text_mask_loss": 0.0012,
+      "predictor_learning_rate": 2.59065603034023e-05,
+      "step": 8376
+    },
+    {
+      "epoch": 0.86,
+      "gumbel_tau": 0.1386,
+      "image_masks_loss": 0.0107,
+      "llm_learning_rate": 2.58697958083991e-07,
+      "loss": 0.6763,
+      "output_text_mask_loss": 0.0052,
+      "predictor_learning_rate": 2.58697958083991e-05,
+      "step": 8377
+    },
+    {
+      "epoch": 0.86,
+      "gumbel_tau": 0.1386,
+      "image_masks_loss": 0.0088,
+      "llm_learning_rate": 2.583305599511321e-07,
+      "loss": 0.7176,
+      "output_text_mask_loss": 0.0028,
+      "predictor_learning_rate": 2.5833055995113207e-05,
+      "step": 8378
+    },
+    {
+      "epoch": 0.86,
+      "gumbel_tau": 0.1385,
+      "image_masks_loss": 0.0125,
+      "llm_learning_rate": 2.579634086759053e-07,
+      "loss": 0.833,
+      "output_text_mask_loss": 0.0378,
+      "predictor_learning_rate": 2.579634086759053e-05,
+      "step": 8379
+    },
+    {
+      "epoch": 0.86,
+      "gumbel_tau": 0.1385,
+      "image_masks_loss": 0.011,
+      "llm_learning_rate": 2.5759650429874217e-07,
+      "loss": 0.7434,
+      "output_text_mask_loss": 0.0328,
+      "predictor_learning_rate": 2.5759650429874215e-05,
+      "step": 8380
+    },
+    {
+      "epoch": 0.86,
+      "gumbel_tau": 0.1385,
+      "image_masks_loss": 0.0131,
+      "llm_learning_rate": 2.572298468600459e-07,
+      "loss": 0.7955,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 2.572298468600459e-05,
+      "step": 8381
+    },
+    {
+      "epoch": 0.86,
+      "gumbel_tau": 0.1384,
+      "image_masks_loss": 0.0216,
+      "llm_learning_rate": 2.5686343640019486e-07,
+      "loss": 0.8206,
+      "output_text_mask_loss": 0.0137,
+      "predictor_learning_rate": 2.5686343640019483e-05,
+      "step": 8382
+    },
+    {
+      "epoch": 0.86,
+      "gumbel_tau": 0.1384,
+      "image_masks_loss": 0.0118,
+      "llm_learning_rate": 2.5649727295953857e-07,
+      "loss": 0.6935,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 2.5649727295953857e-05,
+      "step": 8383
+    },
+    {
+      "epoch": 0.86,
+      "gumbel_tau": 0.1384,
+      "image_masks_loss": 0.0248,
+      "llm_learning_rate": 2.561313565783999e-07,
+      "loss": 0.72,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 2.5613135657839986e-05,
+      "step": 8384
+    },
+    {
+      "epoch": 0.86,
+      "gumbel_tau": 0.1384,
+      "image_masks_loss": 0.003,
+      "llm_learning_rate": 2.557656872970743e-07,
+      "loss": 0.563,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 2.5576568729707426e-05,
+      "step": 8385
+    },
+    {
+      "epoch": 0.86,
+      "gumbel_tau": 0.1383,
+      "image_masks_loss": 0.0107,
+      "llm_learning_rate": 2.554002651558299e-07,
+      "loss": 0.7162,
+      "output_text_mask_loss": 0.0038,
+      "predictor_learning_rate": 2.5540026515582986e-05,
+      "step": 8386
+    },
+    {
+      "epoch": 0.86,
+      "gumbel_tau": 0.1383,
+      "image_masks_loss": 0.0078,
+      "llm_learning_rate": 2.550350901949078e-07,
+      "loss": 0.6927,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 2.550350901949078e-05,
+      "step": 8387
+    },
+    {
+      "epoch": 0.86,
+      "gumbel_tau": 0.1383,
+      "image_masks_loss": 0.0125,
+      "llm_learning_rate": 2.5467016245452195e-07,
+      "loss": 0.723,
+      "output_text_mask_loss": 0.0035,
+      "predictor_learning_rate": 2.5467016245452197e-05,
+      "step": 8388
+    },
+    {
+      "epoch": 0.86,
+      "gumbel_tau": 0.1382,
+      "image_masks_loss": 0.0062,
+      "llm_learning_rate": 2.5430548197485945e-07,
+      "loss": 0.7384,
+      "output_text_mask_loss": 0.0032,
+      "predictor_learning_rate": 2.5430548197485943e-05,
+      "step": 8389
+    },
+    {
+      "epoch": 0.86,
+      "gumbel_tau": 0.1382,
+      "image_masks_loss": 0.0158,
+      "llm_learning_rate": 2.5394104879607915e-07,
+      "loss": 0.7861,
+      "output_text_mask_loss": 0.0275,
+      "predictor_learning_rate": 2.5394104879607916e-05,
+      "step": 8390
+    },
+    {
+      "epoch": 0.86,
+      "gumbel_tau": 0.1382,
+      "image_masks_loss": 0.014,
+      "llm_learning_rate": 2.535768629583135e-07,
+      "loss": 0.7187,
+      "output_text_mask_loss": 0.0576,
+      "predictor_learning_rate": 2.535768629583135e-05,
+      "step": 8391
+    },
+    {
+      "epoch": 0.86,
+      "gumbel_tau": 0.1381,
+      "image_masks_loss": 0.0085,
+      "llm_learning_rate": 2.5321292450166765e-07,
+      "loss": 0.7027,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 2.5321292450166762e-05,
+      "step": 8392
+    },
+    {
+      "epoch": 0.86,
+      "gumbel_tau": 0.1381,
+      "image_masks_loss": 0.0056,
+      "llm_learning_rate": 2.5284923346621886e-07,
+      "loss": 0.7962,
+      "output_text_mask_loss": 0.0039,
+      "predictor_learning_rate": 2.5284923346621886e-05,
+      "step": 8393
+    },
+    {
+      "epoch": 0.86,
+      "gumbel_tau": 0.1381,
+      "image_masks_loss": 0.0083,
+      "llm_learning_rate": 2.52485789892018e-07,
+      "loss": 0.7086,
+      "output_text_mask_loss": 0.0038,
+      "predictor_learning_rate": 2.5248578989201797e-05,
+      "step": 8394
+    },
+    {
+      "epoch": 0.86,
+      "gumbel_tau": 0.138,
+      "image_masks_loss": 0.0073,
+      "llm_learning_rate": 2.521225938190874e-07,
+      "loss": 0.7808,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 2.5212259381908742e-05,
+      "step": 8395
+    },
+    {
+      "epoch": 0.86,
+      "gumbel_tau": 0.138,
+      "image_masks_loss": 0.0063,
+      "llm_learning_rate": 2.5175964528742433e-07,
+      "loss": 0.7812,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 2.5175964528742434e-05,
+      "step": 8396
+    },
+    {
+      "epoch": 0.86,
+      "gumbel_tau": 0.138,
+      "image_masks_loss": 0.0056,
+      "llm_learning_rate": 2.513969443369968e-07,
+      "loss": 0.816,
+      "output_text_mask_loss": 0.0082,
+      "predictor_learning_rate": 2.513969443369968e-05,
+      "step": 8397
+    },
+    {
+      "epoch": 0.86,
+      "gumbel_tau": 0.1379,
+      "image_masks_loss": 0.006,
+      "llm_learning_rate": 2.510344910077464e-07,
+      "loss": 0.7249,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 2.510344910077464e-05,
+      "step": 8398
+    },
+    {
+      "epoch": 0.86,
+      "gumbel_tau": 0.1379,
+      "image_masks_loss": 0.0109,
+      "llm_learning_rate": 2.5067228533958767e-07,
+      "loss": 0.7197,
+      "output_text_mask_loss": 0.001,
+      "predictor_learning_rate": 2.5067228533958765e-05,
+      "step": 8399
+    },
+    {
+      "epoch": 0.86,
+      "gumbel_tau": 0.1379,
+      "image_masks_loss": 0.0091,
+      "llm_learning_rate": 2.5031032737240666e-07,
+      "loss": 0.6895,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 2.5031032737240668e-05,
+      "step": 8400
+    },
+    {
+      "epoch": 0.86,
+      "gumbel_tau": 0.1378,
+      "image_masks_loss": 0.0037,
+      "llm_learning_rate": 2.49948617146063e-07,
+      "loss": 0.7456,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 2.49948617146063e-05,
+      "step": 8401
+    },
+    {
+      "epoch": 0.86,
+      "gumbel_tau": 0.1378,
+      "image_masks_loss": 0.0177,
+      "llm_learning_rate": 2.4958715470038973e-07,
+      "loss": 0.8032,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 2.4958715470038972e-05,
+      "step": 8402
+    },
+    {
+      "epoch": 0.86,
+      "gumbel_tau": 0.1378,
+      "image_masks_loss": 0.0053,
+      "llm_learning_rate": 2.4922594007519144e-07,
+      "loss": 0.6288,
+      "output_text_mask_loss": 0.0004,
+      "predictor_learning_rate": 2.4922594007519147e-05,
+      "step": 8403
+    },
+    {
+      "epoch": 0.86,
+      "gumbel_tau": 0.1377,
+      "image_masks_loss": 0.0039,
+      "llm_learning_rate": 2.4886497331024636e-07,
+      "loss": 0.7678,
+      "output_text_mask_loss": 0.0124,
+      "predictor_learning_rate": 2.488649733102463e-05,
+      "step": 8404
+    },
+    {
+      "epoch": 0.86,
+      "gumbel_tau": 0.1377,
+      "image_masks_loss": 0.0112,
+      "llm_learning_rate": 2.4850425444530423e-07,
+      "loss": 0.6992,
+      "output_text_mask_loss": 0.0244,
+      "predictor_learning_rate": 2.4850425444530423e-05,
+      "step": 8405
+    },
+    {
+      "epoch": 0.86,
+      "gumbel_tau": 0.1377,
+      "image_masks_loss": 0.0141,
+      "llm_learning_rate": 2.4814378352008893e-07,
+      "loss": 0.7439,
+      "output_text_mask_loss": 0.0381,
+      "predictor_learning_rate": 2.4814378352008892e-05,
+      "step": 8406
+    },
+    {
+      "epoch": 0.86,
+      "gumbel_tau": 0.1376,
+      "image_masks_loss": 0.0105,
+      "llm_learning_rate": 2.477835605742959e-07,
+      "loss": 0.6887,
+      "output_text_mask_loss": 0.0088,
+      "predictor_learning_rate": 2.4778356057429592e-05,
+      "step": 8407
+    },
+    {
+      "epoch": 0.86,
+      "gumbel_tau": 0.1376,
+      "image_masks_loss": 0.0044,
+      "llm_learning_rate": 2.4742358564759344e-07,
+      "loss": 0.7087,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 2.4742358564759344e-05,
+      "step": 8408
+    },
+    {
+      "epoch": 0.86,
+      "gumbel_tau": 0.1376,
+      "image_masks_loss": 0.0134,
+      "llm_learning_rate": 2.4706385877962343e-07,
+      "loss": 0.7327,
+      "output_text_mask_loss": 0.016,
+      "predictor_learning_rate": 2.470638587796234e-05,
+      "step": 8409
+    },
+    {
+      "epoch": 0.86,
+      "gumbel_tau": 0.1375,
+      "image_masks_loss": 0.0015,
+      "llm_learning_rate": 2.467043800099994e-07,
+      "loss": 0.676,
+      "output_text_mask_loss": 0.0303,
+      "predictor_learning_rate": 2.467043800099994e-05,
+      "step": 8410
+    },
+    {
+      "epoch": 0.86,
+      "gumbel_tau": 0.1375,
+      "image_masks_loss": 0.0393,
+      "llm_learning_rate": 2.463451493783084e-07,
+      "loss": 0.7511,
+      "output_text_mask_loss": 0.0324,
+      "predictor_learning_rate": 2.4634514937830838e-05,
+      "step": 8411
+    },
+    {
+      "epoch": 0.86,
+      "gumbel_tau": 0.1375,
+      "image_masks_loss": 0.0089,
+      "llm_learning_rate": 2.4598616692410925e-07,
+      "loss": 0.7621,
+      "output_text_mask_loss": 0.0246,
+      "predictor_learning_rate": 2.4598616692410924e-05,
+      "step": 8412
+    },
+    {
+      "epoch": 0.86,
+      "gumbel_tau": 0.1374,
+      "image_masks_loss": 0.0097,
+      "llm_learning_rate": 2.45627432686934e-07,
+      "loss": 0.7665,
+      "output_text_mask_loss": 0.0305,
+      "predictor_learning_rate": 2.4562743268693398e-05,
+      "step": 8413
+    },
+    {
+      "epoch": 0.86,
+      "gumbel_tau": 0.1374,
+      "image_masks_loss": 0.0088,
+      "llm_learning_rate": 2.4526894670628716e-07,
+      "loss": 0.7287,
+      "output_text_mask_loss": 0.0202,
+      "predictor_learning_rate": 2.4526894670628715e-05,
+      "step": 8414
+    },
+    {
+      "epoch": 0.86,
+      "gumbel_tau": 0.1374,
+      "image_masks_loss": 0.0156,
+      "llm_learning_rate": 2.4491070902164573e-07,
+      "loss": 0.7982,
+      "output_text_mask_loss": 0.0132,
+      "predictor_learning_rate": 2.4491070902164568e-05,
+      "step": 8415
+    },
+    {
+      "epoch": 0.86,
+      "gumbel_tau": 0.1373,
+      "image_masks_loss": 0.018,
+      "llm_learning_rate": 2.4455271967246064e-07,
+      "loss": 0.6503,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 2.4455271967246062e-05,
+      "step": 8416
+    },
+    {
+      "epoch": 0.86,
+      "gumbel_tau": 0.1373,
+      "image_masks_loss": 0.0057,
+      "llm_learning_rate": 2.441949786981543e-07,
+      "loss": 0.6494,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 2.4419497869815427e-05,
+      "step": 8417
+    },
+    {
+      "epoch": 0.86,
+      "gumbel_tau": 0.1373,
+      "image_masks_loss": 0.0102,
+      "llm_learning_rate": 2.43837486138121e-07,
+      "loss": 0.7935,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 2.4383748613812095e-05,
+      "step": 8418
+    },
+    {
+      "epoch": 0.86,
+      "gumbel_tau": 0.1372,
+      "image_masks_loss": 0.0084,
+      "llm_learning_rate": 2.4348024203172894e-07,
+      "loss": 0.7434,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 2.434802420317289e-05,
+      "step": 8419
+    },
+    {
+      "epoch": 0.86,
+      "gumbel_tau": 0.1372,
+      "image_masks_loss": 0.0121,
+      "llm_learning_rate": 2.4312324641831865e-07,
+      "loss": 0.8018,
+      "output_text_mask_loss": 0.037,
+      "predictor_learning_rate": 2.4312324641831863e-05,
+      "step": 8420
+    },
+    {
+      "epoch": 0.86,
+      "gumbel_tau": 0.1372,
+      "image_masks_loss": 0.0071,
+      "llm_learning_rate": 2.4276649933720366e-07,
+      "loss": 0.6599,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 2.4276649933720364e-05,
+      "step": 8421
+    },
+    {
+      "epoch": 0.86,
+      "gumbel_tau": 0.1371,
+      "image_masks_loss": 0.0145,
+      "llm_learning_rate": 2.4241000082766973e-07,
+      "loss": 0.7002,
+      "output_text_mask_loss": 0.0004,
+      "predictor_learning_rate": 2.4241000082766975e-05,
+      "step": 8422
+    },
+    {
+      "epoch": 0.86,
+      "gumbel_tau": 0.1371,
+      "image_masks_loss": 0.0234,
+      "llm_learning_rate": 2.4205375092897526e-07,
+      "loss": 0.7396,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 2.4205375092897524e-05,
+      "step": 8423
+    },
+    {
+      "epoch": 0.86,
+      "gumbel_tau": 0.1371,
+      "image_masks_loss": 0.0037,
+      "llm_learning_rate": 2.416977496803508e-07,
+      "loss": 0.6562,
+      "output_text_mask_loss": 0.0004,
+      "predictor_learning_rate": 2.4169774968035076e-05,
+      "step": 8424
+    },
+    {
+      "epoch": 0.86,
+      "gumbel_tau": 0.1371,
+      "image_masks_loss": 0.0112,
+      "llm_learning_rate": 2.413419971210004e-07,
+      "loss": 0.7322,
+      "output_text_mask_loss": 0.0244,
+      "predictor_learning_rate": 2.413419971210004e-05,
+      "step": 8425
+    },
+    {
+      "epoch": 0.86,
+      "gumbel_tau": 0.137,
+      "image_masks_loss": 0.0051,
+      "llm_learning_rate": 2.4098649329010045e-07,
+      "loss": 0.6166,
+      "output_text_mask_loss": 0.0162,
+      "predictor_learning_rate": 2.409864932901004e-05,
+      "step": 8426
+    },
+    {
+      "epoch": 0.86,
+      "gumbel_tau": 0.137,
+      "image_masks_loss": 0.0175,
+      "llm_learning_rate": 2.406312382267989e-07,
+      "loss": 0.6737,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 2.406312382267989e-05,
+      "step": 8427
+    },
+    {
+      "epoch": 0.86,
+      "gumbel_tau": 0.137,
+      "image_masks_loss": 0.0244,
+      "llm_learning_rate": 2.4027623197021845e-07,
+      "loss": 0.7106,
+      "output_text_mask_loss": 0.0656,
+      "predictor_learning_rate": 2.4027623197021842e-05,
+      "step": 8428
+    },
+    {
+      "epoch": 0.86,
+      "gumbel_tau": 0.1369,
+      "image_masks_loss": 0.0069,
+      "llm_learning_rate": 2.3992147455945276e-07,
+      "loss": 0.6943,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 2.3992147455945274e-05,
+      "step": 8429
+    },
+    {
+      "epoch": 0.86,
+      "gumbel_tau": 0.1369,
+      "image_masks_loss": 0.0131,
+      "llm_learning_rate": 2.3956696603356837e-07,
+      "loss": 0.6505,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 2.3956696603356838e-05,
+      "step": 8430
+    },
+    {
+      "epoch": 0.86,
+      "gumbel_tau": 0.1369,
+      "image_masks_loss": 0.0037,
+      "llm_learning_rate": 2.3921270643160464e-07,
+      "loss": 0.7332,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 2.3921270643160464e-05,
+      "step": 8431
+    },
+    {
+      "epoch": 0.86,
+      "gumbel_tau": 0.1368,
+      "image_masks_loss": 0.0047,
+      "llm_learning_rate": 2.3885869579257316e-07,
+      "loss": 0.7128,
+      "output_text_mask_loss": 0.0504,
+      "predictor_learning_rate": 2.3885869579257314e-05,
+      "step": 8432
+    },
+    {
+      "epoch": 0.86,
+      "gumbel_tau": 0.1368,
+      "image_masks_loss": 0.0086,
+      "llm_learning_rate": 2.385049341554588e-07,
+      "loss": 0.7046,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 2.3850493415545876e-05,
+      "step": 8433
+    },
+    {
+      "epoch": 0.86,
+      "gumbel_tau": 0.1368,
+      "image_masks_loss": 0.0078,
+      "llm_learning_rate": 2.381514215592179e-07,
+      "loss": 0.764,
+      "output_text_mask_loss": 0.0072,
+      "predictor_learning_rate": 2.3815142155921788e-05,
+      "step": 8434
+    },
+    {
+      "epoch": 0.86,
+      "gumbel_tau": 0.1367,
+      "image_masks_loss": 0.0097,
+      "llm_learning_rate": 2.37798158042781e-07,
+      "loss": 0.7581,
+      "output_text_mask_loss": 0.0078,
+      "predictor_learning_rate": 2.37798158042781e-05,
+      "step": 8435
+    },
+    {
+      "epoch": 0.86,
+      "gumbel_tau": 0.1367,
+      "image_masks_loss": 0.0106,
+      "llm_learning_rate": 2.3744514364505033e-07,
+      "loss": 0.7751,
+      "output_text_mask_loss": 0.0904,
+      "predictor_learning_rate": 2.374451436450503e-05,
+      "step": 8436
+    },
+    {
+      "epoch": 0.86,
+      "gumbel_tau": 0.1367,
+      "image_masks_loss": 0.0225,
+      "llm_learning_rate": 2.3709237840489952e-07,
+      "loss": 0.7344,
+      "output_text_mask_loss": 0.007,
+      "predictor_learning_rate": 2.370923784048995e-05,
+      "step": 8437
+    },
+    {
+      "epoch": 0.86,
+      "gumbel_tau": 0.1366,
+      "image_masks_loss": 0.009,
+      "llm_learning_rate": 2.3673986236117635e-07,
+      "loss": 0.6685,
+      "output_text_mask_loss": 0.0231,
+      "predictor_learning_rate": 2.3673986236117633e-05,
+      "step": 8438
+    },
+    {
+      "epoch": 0.86,
+      "gumbel_tau": 0.1366,
+      "image_masks_loss": 0.0184,
+      "llm_learning_rate": 2.363875955527009e-07,
+      "loss": 0.785,
+      "output_text_mask_loss": 0.0107,
+      "predictor_learning_rate": 2.363875955527009e-05,
+      "step": 8439
+    },
+    {
+      "epoch": 0.86,
+      "gumbel_tau": 0.1366,
+      "image_masks_loss": 0.0183,
+      "llm_learning_rate": 2.3603557801826493e-07,
+      "loss": 0.6787,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 2.3603557801826492e-05,
+      "step": 8440
+    },
+    {
+      "epoch": 0.86,
+      "gumbel_tau": 0.1365,
+      "image_masks_loss": 0.012,
+      "llm_learning_rate": 2.3568380979663436e-07,
+      "loss": 0.7392,
+      "output_text_mask_loss": 0.0139,
+      "predictor_learning_rate": 2.3568380979663434e-05,
+      "step": 8441
+    },
+    {
+      "epoch": 0.86,
+      "gumbel_tau": 0.1365,
+      "image_masks_loss": 0.0042,
+      "llm_learning_rate": 2.3533229092654626e-07,
+      "loss": 0.626,
+      "output_text_mask_loss": 0.0231,
+      "predictor_learning_rate": 2.3533229092654623e-05,
+      "step": 8442
+    },
+    {
+      "epoch": 0.87,
+      "gumbel_tau": 0.1365,
+      "image_masks_loss": 0.0111,
+      "llm_learning_rate": 2.349810214467105e-07,
+      "loss": 0.6442,
+      "output_text_mask_loss": 0.0041,
+      "predictor_learning_rate": 2.3498102144671045e-05,
+      "step": 8443
+    },
+    {
+      "epoch": 0.87,
+      "gumbel_tau": 0.1364,
+      "image_masks_loss": 0.0077,
+      "llm_learning_rate": 2.3463000139580972e-07,
+      "loss": 0.7067,
+      "output_text_mask_loss": 0.0111,
+      "predictor_learning_rate": 2.346300013958097e-05,
+      "step": 8444
+    },
+    {
+      "epoch": 0.87,
+      "gumbel_tau": 0.1364,
+      "image_masks_loss": 0.01,
+      "llm_learning_rate": 2.3427923081249943e-07,
+      "loss": 0.7326,
+      "output_text_mask_loss": 0.1259,
+      "predictor_learning_rate": 2.342792308124994e-05,
+      "step": 8445
+    },
+    {
+      "epoch": 0.87,
+      "gumbel_tau": 0.1364,
+      "image_masks_loss": 0.0075,
+      "llm_learning_rate": 2.3392870973540648e-07,
+      "loss": 0.796,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 2.3392870973540647e-05,
+      "step": 8446
+    },
+    {
+      "epoch": 0.87,
+      "gumbel_tau": 0.1363,
+      "image_masks_loss": 0.0051,
+      "llm_learning_rate": 2.3357843820313142e-07,
+      "loss": 0.6544,
+      "output_text_mask_loss": 0.0032,
+      "predictor_learning_rate": 2.335784382031314e-05,
+      "step": 8447
+    },
+    {
+      "epoch": 0.87,
+      "gumbel_tau": 0.1363,
+      "image_masks_loss": 0.0123,
+      "llm_learning_rate": 2.3322841625424724e-07,
+      "loss": 0.6141,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 2.3322841625424722e-05,
+      "step": 8448
+    },
+    {
+      "epoch": 0.87,
+      "gumbel_tau": 0.1363,
+      "image_masks_loss": 0.0174,
+      "llm_learning_rate": 2.3287864392729899e-07,
+      "loss": 0.7305,
+      "output_text_mask_loss": 0.0244,
+      "predictor_learning_rate": 2.3287864392729895e-05,
+      "step": 8449
+    },
+    {
+      "epoch": 0.87,
+      "gumbel_tau": 0.1362,
+      "image_masks_loss": 0.0165,
+      "llm_learning_rate": 2.325291212608044e-07,
+      "loss": 0.6984,
+      "output_text_mask_loss": 0.0736,
+      "predictor_learning_rate": 2.325291212608044e-05,
+      "step": 8450
+    },
+    {
+      "epoch": 0.87,
+      "gumbel_tau": 0.1362,
+      "image_masks_loss": 0.0065,
+      "llm_learning_rate": 2.3217984829325356e-07,
+      "loss": 0.7805,
+      "output_text_mask_loss": 0.0025,
+      "predictor_learning_rate": 2.3217984829325357e-05,
+      "step": 8451
+    },
+    {
+      "epoch": 0.87,
+      "gumbel_tau": 0.1362,
+      "image_masks_loss": 0.0026,
+      "llm_learning_rate": 2.3183082506310904e-07,
+      "loss": 0.727,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 2.3183082506310904e-05,
+      "step": 8452
+    },
+    {
+      "epoch": 0.87,
+      "gumbel_tau": 0.1361,
+      "image_masks_loss": 0.0103,
+      "llm_learning_rate": 2.314820516088062e-07,
+      "loss": 0.7445,
+      "output_text_mask_loss": 0.0041,
+      "predictor_learning_rate": 2.314820516088062e-05,
+      "step": 8453
+    },
+    {
+      "epoch": 0.87,
+      "gumbel_tau": 0.1361,
+      "image_masks_loss": 0.0207,
+      "llm_learning_rate": 2.3113352796875295e-07,
+      "loss": 0.7989,
+      "output_text_mask_loss": 0.0099,
+      "predictor_learning_rate": 2.311335279687529e-05,
+      "step": 8454
+    },
+    {
+      "epoch": 0.87,
+      "gumbel_tau": 0.1361,
+      "image_masks_loss": 0.0146,
+      "llm_learning_rate": 2.307852541813291e-07,
+      "loss": 0.7883,
+      "output_text_mask_loss": 0.0156,
+      "predictor_learning_rate": 2.3078525418132906e-05,
+      "step": 8455
+    },
+    {
+      "epoch": 0.87,
+      "gumbel_tau": 0.1361,
+      "image_masks_loss": 0.0022,
+      "llm_learning_rate": 2.3043723028488757e-07,
+      "loss": 0.7999,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 2.3043723028488754e-05,
+      "step": 8456
+    },
+    {
+      "epoch": 0.87,
+      "gumbel_tau": 0.136,
+      "image_masks_loss": 0.0168,
+      "llm_learning_rate": 2.300894563177536e-07,
+      "loss": 0.7556,
+      "output_text_mask_loss": 0.006,
+      "predictor_learning_rate": 2.300894563177536e-05,
+      "step": 8457
+    },
+    {
+      "epoch": 0.87,
+      "gumbel_tau": 0.136,
+      "image_masks_loss": 0.0171,
+      "llm_learning_rate": 2.2974193231822455e-07,
+      "loss": 0.8229,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 2.2974193231822454e-05,
+      "step": 8458
+    },
+    {
+      "epoch": 0.87,
+      "gumbel_tau": 0.136,
+      "image_masks_loss": 0.0093,
+      "llm_learning_rate": 2.2939465832457043e-07,
+      "loss": 0.7594,
+      "output_text_mask_loss": 0.001,
+      "predictor_learning_rate": 2.2939465832457042e-05,
+      "step": 8459
+    },
+    {
+      "epoch": 0.87,
+      "gumbel_tau": 0.1359,
+      "image_masks_loss": 0.0053,
+      "llm_learning_rate": 2.2904763437503447e-07,
+      "loss": 0.6743,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 2.2904763437503444e-05,
+      "step": 8460
+    },
+    {
+      "epoch": 0.87,
+      "gumbel_tau": 0.1359,
+      "image_masks_loss": 0.0045,
+      "llm_learning_rate": 2.287008605078314e-07,
+      "loss": 0.7374,
+      "output_text_mask_loss": 0.0004,
+      "predictor_learning_rate": 2.287008605078314e-05,
+      "step": 8461
+    },
+    {
+      "epoch": 0.87,
+      "gumbel_tau": 0.1359,
+      "image_masks_loss": 0.0144,
+      "llm_learning_rate": 2.2835433676114877e-07,
+      "loss": 0.7158,
+      "output_text_mask_loss": 0.038,
+      "predictor_learning_rate": 2.2835433676114873e-05,
+      "step": 8462
+    },
+    {
+      "epoch": 0.87,
+      "gumbel_tau": 0.1358,
+      "image_masks_loss": 0.0085,
+      "llm_learning_rate": 2.2800806317314627e-07,
+      "loss": 0.6989,
+      "output_text_mask_loss": 0.02,
+      "predictor_learning_rate": 2.2800806317314626e-05,
+      "step": 8463
+    },
+    {
+      "epoch": 0.87,
+      "gumbel_tau": 0.1358,
+      "image_masks_loss": 0.0068,
+      "llm_learning_rate": 2.2766203978195678e-07,
+      "loss": 0.7444,
+      "output_text_mask_loss": 0.132,
+      "predictor_learning_rate": 2.2766203978195678e-05,
+      "step": 8464
+    },
+    {
+      "epoch": 0.87,
+      "gumbel_tau": 0.1358,
+      "image_masks_loss": 0.0036,
+      "llm_learning_rate": 2.273162666256848e-07,
+      "loss": 0.7578,
+      "output_text_mask_loss": 0.0231,
+      "predictor_learning_rate": 2.273162666256848e-05,
+      "step": 8465
+    },
+    {
+      "epoch": 0.87,
+      "gumbel_tau": 0.1357,
+      "image_masks_loss": 0.009,
+      "llm_learning_rate": 2.2697074374240735e-07,
+      "loss": 0.7325,
+      "output_text_mask_loss": 0.0162,
+      "predictor_learning_rate": 2.2697074374240734e-05,
+      "step": 8466
+    },
+    {
+      "epoch": 0.87,
+      "gumbel_tau": 0.1357,
+      "image_masks_loss": 0.0211,
+      "llm_learning_rate": 2.2662547117017487e-07,
+      "loss": 0.7313,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 2.2662547117017484e-05,
+      "step": 8467
+    },
+    {
+      "epoch": 0.87,
+      "gumbel_tau": 0.1357,
+      "image_masks_loss": 0.0088,
+      "llm_learning_rate": 2.2628044894700912e-07,
+      "loss": 0.6161,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 2.262804489470091e-05,
+      "step": 8468
+    },
+    {
+      "epoch": 0.87,
+      "gumbel_tau": 0.1356,
+      "image_masks_loss": 0.0169,
+      "llm_learning_rate": 2.2593567711090532e-07,
+      "loss": 0.6225,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 2.259356771109053e-05,
+      "step": 8469
+    },
+    {
+      "epoch": 0.87,
+      "gumbel_tau": 0.1356,
+      "image_masks_loss": 0.0164,
+      "llm_learning_rate": 2.255911556998297e-07,
+      "loss": 0.7587,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 2.2559115569982968e-05,
+      "step": 8470
+    },
+    {
+      "epoch": 0.87,
+      "gumbel_tau": 0.1356,
+      "image_masks_loss": 0.0024,
+      "llm_learning_rate": 2.2524688475172195e-07,
+      "loss": 0.7389,
+      "output_text_mask_loss": 0.0014,
+      "predictor_learning_rate": 2.2524688475172194e-05,
+      "step": 8471
+    },
+    {
+      "epoch": 0.87,
+      "gumbel_tau": 0.1355,
+      "image_masks_loss": 0.0062,
+      "llm_learning_rate": 2.2490286430449426e-07,
+      "loss": 0.718,
+      "output_text_mask_loss": 0.0145,
+      "predictor_learning_rate": 2.2490286430449424e-05,
+      "step": 8472
+    },
+    {
+      "epoch": 0.87,
+      "gumbel_tau": 0.1355,
+      "image_masks_loss": 0.0113,
+      "llm_learning_rate": 2.2455909439603075e-07,
+      "loss": 0.7343,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 2.2455909439603074e-05,
+      "step": 8473
+    },
+    {
+      "epoch": 0.87,
+      "gumbel_tau": 0.1355,
+      "image_masks_loss": 0.0083,
+      "llm_learning_rate": 2.2421557506418807e-07,
+      "loss": 0.8059,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 2.2421557506418805e-05,
+      "step": 8474
+    },
+    {
+      "epoch": 0.87,
+      "gumbel_tau": 0.1354,
+      "image_masks_loss": 0.0214,
+      "llm_learning_rate": 2.238723063467954e-07,
+      "loss": 0.7368,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 2.238723063467954e-05,
+      "step": 8475
+    },
+    {
+      "epoch": 0.87,
+      "gumbel_tau": 0.1354,
+      "image_masks_loss": 0.0221,
+      "llm_learning_rate": 2.235292882816542e-07,
+      "loss": 0.7366,
+      "output_text_mask_loss": 0.0198,
+      "predictor_learning_rate": 2.235292882816542e-05,
+      "step": 8476
+    },
+    {
+      "epoch": 0.87,
+      "gumbel_tau": 0.1354,
+      "image_masks_loss": 0.0071,
+      "llm_learning_rate": 2.2318652090653842e-07,
+      "loss": 0.6998,
+      "output_text_mask_loss": 0.0216,
+      "predictor_learning_rate": 2.231865209065384e-05,
+      "step": 8477
+    },
+    {
+      "epoch": 0.87,
+      "gumbel_tau": 0.1353,
+      "image_masks_loss": 0.005,
+      "llm_learning_rate": 2.2284400425919422e-07,
+      "loss": 0.6849,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 2.228440042591942e-05,
+      "step": 8478
+    },
+    {
+      "epoch": 0.87,
+      "gumbel_tau": 0.1353,
+      "image_masks_loss": 0.0023,
+      "llm_learning_rate": 2.2250173837734036e-07,
+      "loss": 0.7176,
+      "output_text_mask_loss": 0.0118,
+      "predictor_learning_rate": 2.2250173837734032e-05,
+      "step": 8479
+    },
+    {
+      "epoch": 0.87,
+      "gumbel_tau": 0.1353,
+      "image_masks_loss": 0.013,
+      "llm_learning_rate": 2.2215972329866803e-07,
+      "loss": 0.614,
+      "output_text_mask_loss": 0.0275,
+      "predictor_learning_rate": 2.2215972329866802e-05,
+      "step": 8480
+    },
+    {
+      "epoch": 0.87,
+      "gumbel_tau": 0.1353,
+      "image_masks_loss": 0.0156,
+      "llm_learning_rate": 2.2181795906084102e-07,
+      "loss": 0.6191,
+      "output_text_mask_loss": 0.021,
+      "predictor_learning_rate": 2.21817959060841e-05,
+      "step": 8481
+    },
+    {
+      "epoch": 0.87,
+      "gumbel_tau": 0.1352,
+      "image_masks_loss": 0.0103,
+      "llm_learning_rate": 2.2147644570149452e-07,
+      "loss": 0.7104,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 2.214764457014945e-05,
+      "step": 8482
+    },
+    {
+      "epoch": 0.87,
+      "gumbel_tau": 0.1352,
+      "image_masks_loss": 0.0018,
+      "llm_learning_rate": 2.2113518325823703e-07,
+      "loss": 0.7024,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 2.2113518325823702e-05,
+      "step": 8483
+    },
+    {
+      "epoch": 0.87,
+      "gumbel_tau": 0.1352,
+      "image_masks_loss": 0.0089,
+      "llm_learning_rate": 2.207941717686493e-07,
+      "loss": 0.6967,
+      "output_text_mask_loss": 0.025,
+      "predictor_learning_rate": 2.207941717686493e-05,
+      "step": 8484
+    },
+    {
+      "epoch": 0.87,
+      "gumbel_tau": 0.1351,
+      "image_masks_loss": 0.0044,
+      "llm_learning_rate": 2.2045341127028414e-07,
+      "loss": 0.7536,
+      "output_text_mask_loss": 0.0065,
+      "predictor_learning_rate": 2.2045341127028412e-05,
+      "step": 8485
+    },
+    {
+      "epoch": 0.87,
+      "gumbel_tau": 0.1351,
+      "image_masks_loss": 0.0204,
+      "llm_learning_rate": 2.201129018006662e-07,
+      "loss": 0.7796,
+      "output_text_mask_loss": 0.0039,
+      "predictor_learning_rate": 2.2011290180066617e-05,
+      "step": 8486
+    },
+    {
+      "epoch": 0.87,
+      "gumbel_tau": 0.1351,
+      "image_masks_loss": 0.01,
+      "llm_learning_rate": 2.1977264339729437e-07,
+      "loss": 0.7938,
+      "output_text_mask_loss": 0.0357,
+      "predictor_learning_rate": 2.1977264339729437e-05,
+      "step": 8487
+    },
+    {
+      "epoch": 0.87,
+      "gumbel_tau": 0.135,
+      "image_masks_loss": 0.0172,
+      "llm_learning_rate": 2.1943263609763787e-07,
+      "loss": 0.722,
+      "output_text_mask_loss": 0.0238,
+      "predictor_learning_rate": 2.1943263609763785e-05,
+      "step": 8488
+    },
+    {
+      "epoch": 0.87,
+      "gumbel_tau": 0.135,
+      "image_masks_loss": 0.015,
+      "llm_learning_rate": 2.1909287993913952e-07,
+      "loss": 0.7582,
+      "output_text_mask_loss": 0.0021,
+      "predictor_learning_rate": 2.190928799391395e-05,
+      "step": 8489
+    },
+    {
+      "epoch": 0.87,
+      "gumbel_tau": 0.135,
+      "image_masks_loss": 0.0107,
+      "llm_learning_rate": 2.1875337495921355e-07,
+      "loss": 0.716,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 2.1875337495921353e-05,
+      "step": 8490
+    },
+    {
+      "epoch": 0.87,
+      "gumbel_tau": 0.1349,
+      "image_masks_loss": 0.01,
+      "llm_learning_rate": 2.184141211952473e-07,
+      "loss": 0.7475,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 2.1841412119524728e-05,
+      "step": 8491
+    },
+    {
+      "epoch": 0.87,
+      "gumbel_tau": 0.1349,
+      "image_masks_loss": 0.008,
+      "llm_learning_rate": 2.180751186846e-07,
+      "loss": 0.7322,
+      "output_text_mask_loss": 0.0374,
+      "predictor_learning_rate": 2.1807511868460002e-05,
+      "step": 8492
+    },
+    {
+      "epoch": 0.87,
+      "gumbel_tau": 0.1349,
+      "image_masks_loss": 0.0052,
+      "llm_learning_rate": 2.1773636746460326e-07,
+      "loss": 0.7481,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 2.1773636746460324e-05,
+      "step": 8493
+    },
+    {
+      "epoch": 0.87,
+      "gumbel_tau": 0.1348,
+      "image_masks_loss": 0.0154,
+      "llm_learning_rate": 2.173978675725616e-07,
+      "loss": 0.6935,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 2.173978675725616e-05,
+      "step": 8494
+    },
+    {
+      "epoch": 0.87,
+      "gumbel_tau": 0.1348,
+      "image_masks_loss": 0.0038,
+      "llm_learning_rate": 2.170596190457508e-07,
+      "loss": 0.6724,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 2.170596190457508e-05,
+      "step": 8495
+    },
+    {
+      "epoch": 0.87,
+      "gumbel_tau": 0.1348,
+      "image_masks_loss": 0.0073,
+      "llm_learning_rate": 2.1672162192141994e-07,
+      "loss": 0.7865,
+      "output_text_mask_loss": 0.0381,
+      "predictor_learning_rate": 2.1672162192141992e-05,
+      "step": 8496
+    },
+    {
+      "epoch": 0.87,
+      "gumbel_tau": 0.1347,
+      "image_masks_loss": 0.0159,
+      "llm_learning_rate": 2.163838762367898e-07,
+      "loss": 0.766,
+      "output_text_mask_loss": 0.0109,
+      "predictor_learning_rate": 2.163838762367898e-05,
+      "step": 8497
+    },
+    {
+      "epoch": 0.87,
+      "gumbel_tau": 0.1347,
+      "image_masks_loss": 0.0092,
+      "llm_learning_rate": 2.1604638202905343e-07,
+      "loss": 0.7433,
+      "output_text_mask_loss": 0.0595,
+      "predictor_learning_rate": 2.160463820290534e-05,
+      "step": 8498
+    },
+    {
+      "epoch": 0.87,
+      "gumbel_tau": 0.1347,
+      "image_masks_loss": 0.0123,
+      "llm_learning_rate": 2.1570913933537718e-07,
+      "loss": 0.7706,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 2.1570913933537718e-05,
+      "step": 8499
+    },
+    {
+      "epoch": 0.87,
+      "gumbel_tau": 0.1346,
+      "image_masks_loss": 0.0146,
+      "llm_learning_rate": 2.1537214819289858e-07,
+      "loss": 0.7394,
+      "output_text_mask_loss": 0.0038,
+      "predictor_learning_rate": 2.1537214819289856e-05,
+      "step": 8500
+    },
+    {
+      "epoch": 0.87,
+      "gumbel_tau": 0.1346,
+      "image_masks_loss": 0.0049,
+      "llm_learning_rate": 2.1503540863872768e-07,
+      "loss": 0.8229,
+      "output_text_mask_loss": 0.0282,
+      "predictor_learning_rate": 2.1503540863872766e-05,
+      "step": 8501
+    },
+    {
+      "epoch": 0.87,
+      "gumbel_tau": 0.1346,
+      "image_masks_loss": 0.0039,
+      "llm_learning_rate": 2.146989207099473e-07,
+      "loss": 0.5965,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 2.1469892070994728e-05,
+      "step": 8502
+    },
+    {
+      "epoch": 0.87,
+      "gumbel_tau": 0.1346,
+      "image_masks_loss": 0.0079,
+      "llm_learning_rate": 2.1436268444361198e-07,
+      "loss": 0.7453,
+      "output_text_mask_loss": 0.0273,
+      "predictor_learning_rate": 2.1436268444361195e-05,
+      "step": 8503
+    },
+    {
+      "epoch": 0.87,
+      "gumbel_tau": 0.1345,
+      "image_masks_loss": 0.0153,
+      "llm_learning_rate": 2.1402669987674902e-07,
+      "loss": 0.6974,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 2.14026699876749e-05,
+      "step": 8504
+    },
+    {
+      "epoch": 0.87,
+      "gumbel_tau": 0.1345,
+      "image_masks_loss": 0.0061,
+      "llm_learning_rate": 2.1369096704635718e-07,
+      "loss": 0.7681,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 2.1369096704635715e-05,
+      "step": 8505
+    },
+    {
+      "epoch": 0.87,
+      "gumbel_tau": 0.1345,
+      "image_masks_loss": 0.0079,
+      "llm_learning_rate": 2.1335548598940908e-07,
+      "loss": 0.7124,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 2.1335548598940908e-05,
+      "step": 8506
+    },
+    {
+      "epoch": 0.87,
+      "gumbel_tau": 0.1344,
+      "image_masks_loss": 0.0099,
+      "llm_learning_rate": 2.13020256742848e-07,
+      "loss": 0.7256,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 2.1302025674284798e-05,
+      "step": 8507
+    },
+    {
+      "epoch": 0.87,
+      "gumbel_tau": 0.1344,
+      "image_masks_loss": 0.01,
+      "llm_learning_rate": 2.1268527934359045e-07,
+      "loss": 0.7637,
+      "output_text_mask_loss": 0.016,
+      "predictor_learning_rate": 2.1268527934359044e-05,
+      "step": 8508
+    },
+    {
+      "epoch": 0.87,
+      "gumbel_tau": 0.1344,
+      "image_masks_loss": 0.011,
+      "llm_learning_rate": 2.1235055382852476e-07,
+      "loss": 0.6907,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 2.1235055382852475e-05,
+      "step": 8509
+    },
+    {
+      "epoch": 0.87,
+      "gumbel_tau": 0.1343,
+      "image_masks_loss": 0.0071,
+      "llm_learning_rate": 2.1201608023451199e-07,
+      "loss": 0.799,
+      "output_text_mask_loss": 0.0322,
+      "predictor_learning_rate": 2.1201608023451197e-05,
+      "step": 8510
+    },
+    {
+      "epoch": 0.87,
+      "gumbel_tau": 0.1343,
+      "image_masks_loss": 0.0035,
+      "llm_learning_rate": 2.1168185859838374e-07,
+      "loss": 0.6836,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 2.1168185859838375e-05,
+      "step": 8511
+    },
+    {
+      "epoch": 0.87,
+      "gumbel_tau": 0.1343,
+      "image_masks_loss": 0.007,
+      "llm_learning_rate": 2.1134788895694702e-07,
+      "loss": 0.6554,
+      "output_text_mask_loss": 0.005,
+      "predictor_learning_rate": 2.1134788895694702e-05,
+      "step": 8512
+    },
+    {
+      "epoch": 0.87,
+      "gumbel_tau": 0.1342,
+      "image_masks_loss": 0.0054,
+      "llm_learning_rate": 2.110141713469782e-07,
+      "loss": 0.8546,
+      "output_text_mask_loss": 0.0706,
+      "predictor_learning_rate": 2.110141713469782e-05,
+      "step": 8513
+    },
+    {
+      "epoch": 0.87,
+      "gumbel_tau": 0.1342,
+      "image_masks_loss": 0.0097,
+      "llm_learning_rate": 2.1068070580522766e-07,
+      "loss": 0.7418,
+      "output_text_mask_loss": 0.0231,
+      "predictor_learning_rate": 2.1068070580522762e-05,
+      "step": 8514
+    },
+    {
+      "epoch": 0.87,
+      "gumbel_tau": 0.1342,
+      "image_masks_loss": 0.0085,
+      "llm_learning_rate": 2.1034749236841678e-07,
+      "loss": 0.7881,
+      "output_text_mask_loss": 0.0165,
+      "predictor_learning_rate": 2.103474923684168e-05,
+      "step": 8515
+    },
+    {
+      "epoch": 0.87,
+      "gumbel_tau": 0.1341,
+      "image_masks_loss": 0.0051,
+      "llm_learning_rate": 2.1001453107324017e-07,
+      "loss": 0.6378,
+      "output_text_mask_loss": 0.0277,
+      "predictor_learning_rate": 2.1001453107324015e-05,
+      "step": 8516
+    },
+    {
+      "epoch": 0.87,
+      "gumbel_tau": 0.1341,
+      "image_masks_loss": 0.006,
+      "llm_learning_rate": 2.0968182195636428e-07,
+      "loss": 0.6528,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 2.0968182195636426e-05,
+      "step": 8517
+    },
+    {
+      "epoch": 0.87,
+      "gumbel_tau": 0.1341,
+      "image_masks_loss": 0.0032,
+      "llm_learning_rate": 2.0934936505442703e-07,
+      "loss": 0.7801,
+      "output_text_mask_loss": 0.0179,
+      "predictor_learning_rate": 2.0934936505442703e-05,
+      "step": 8518
+    },
+    {
+      "epoch": 0.87,
+      "gumbel_tau": 0.134,
+      "image_masks_loss": 0.0155,
+      "llm_learning_rate": 2.090171604040403e-07,
+      "loss": 0.7642,
+      "output_text_mask_loss": 0.0012,
+      "predictor_learning_rate": 2.0901716040404028e-05,
+      "step": 8519
+    },
+    {
+      "epoch": 0.87,
+      "gumbel_tau": 0.134,
+      "image_masks_loss": 0.0136,
+      "llm_learning_rate": 2.08685208041787e-07,
+      "loss": 0.6122,
+      "output_text_mask_loss": 0.0193,
+      "predictor_learning_rate": 2.0868520804178698e-05,
+      "step": 8520
+    },
+    {
+      "epoch": 0.87,
+      "gumbel_tau": 0.134,
+      "image_masks_loss": 0.0062,
+      "llm_learning_rate": 2.0835350800422204e-07,
+      "loss": 0.6901,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 2.0835350800422205e-05,
+      "step": 8521
+    },
+    {
+      "epoch": 0.87,
+      "gumbel_tau": 0.134,
+      "image_masks_loss": 0.012,
+      "llm_learning_rate": 2.0802206032787347e-07,
+      "loss": 0.7246,
+      "output_text_mask_loss": 0.0193,
+      "predictor_learning_rate": 2.0802206032787346e-05,
+      "step": 8522
+    },
+    {
+      "epoch": 0.87,
+      "gumbel_tau": 0.1339,
+      "image_masks_loss": 0.0138,
+      "llm_learning_rate": 2.076908650492407e-07,
+      "loss": 0.7182,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 2.0769086504924067e-05,
+      "step": 8523
+    },
+    {
+      "epoch": 0.87,
+      "gumbel_tau": 0.1339,
+      "image_masks_loss": 0.0228,
+      "llm_learning_rate": 2.073599222047959e-07,
+      "loss": 0.7618,
+      "output_text_mask_loss": 0.0145,
+      "predictor_learning_rate": 2.0735992220479587e-05,
+      "step": 8524
+    },
+    {
+      "epoch": 0.87,
+      "gumbel_tau": 0.1339,
+      "image_masks_loss": 0.0044,
+      "llm_learning_rate": 2.0702923183098245e-07,
+      "loss": 0.7478,
+      "output_text_mask_loss": 0.0039,
+      "predictor_learning_rate": 2.0702923183098244e-05,
+      "step": 8525
+    },
+    {
+      "epoch": 0.87,
+      "gumbel_tau": 0.1338,
+      "image_masks_loss": 0.0037,
+      "llm_learning_rate": 2.0669879396421787e-07,
+      "loss": 0.6236,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 2.0669879396421786e-05,
+      "step": 8526
+    },
+    {
+      "epoch": 0.87,
+      "gumbel_tau": 0.1338,
+      "image_masks_loss": 0.0093,
+      "llm_learning_rate": 2.0636860864089059e-07,
+      "loss": 0.7572,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 2.0636860864089056e-05,
+      "step": 8527
+    },
+    {
+      "epoch": 0.87,
+      "gumbel_tau": 0.1338,
+      "image_masks_loss": 0.0031,
+      "llm_learning_rate": 2.0603867589736094e-07,
+      "loss": 0.6509,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 2.0603867589736092e-05,
+      "step": 8528
+    },
+    {
+      "epoch": 0.87,
+      "gumbel_tau": 0.1337,
+      "image_masks_loss": 0.0122,
+      "llm_learning_rate": 2.0570899576996178e-07,
+      "loss": 0.583,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 2.0570899576996178e-05,
+      "step": 8529
+    },
+    {
+      "epoch": 0.87,
+      "gumbel_tau": 0.1337,
+      "image_masks_loss": 0.0033,
+      "llm_learning_rate": 2.0537956829499772e-07,
+      "loss": 0.7114,
+      "output_text_mask_loss": 0.0189,
+      "predictor_learning_rate": 2.053795682949977e-05,
+      "step": 8530
+    },
+    {
+      "epoch": 0.87,
+      "gumbel_tau": 0.1337,
+      "image_masks_loss": 0.005,
+      "llm_learning_rate": 2.0505039350874723e-07,
+      "loss": 0.6625,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 2.0505039350874722e-05,
+      "step": 8531
+    },
+    {
+      "epoch": 0.87,
+      "gumbel_tau": 0.1336,
+      "image_masks_loss": 0.004,
+      "llm_learning_rate": 2.0472147144745906e-07,
+      "loss": 0.6573,
+      "output_text_mask_loss": 0.005,
+      "predictor_learning_rate": 2.0472147144745905e-05,
+      "step": 8532
+    },
+    {
+      "epoch": 0.87,
+      "gumbel_tau": 0.1336,
+      "image_masks_loss": 0.0096,
+      "llm_learning_rate": 2.0439280214735514e-07,
+      "loss": 0.8661,
+      "output_text_mask_loss": 0.0248,
+      "predictor_learning_rate": 2.043928021473551e-05,
+      "step": 8533
+    },
+    {
+      "epoch": 0.87,
+      "gumbel_tau": 0.1336,
+      "image_masks_loss": 0.0162,
+      "llm_learning_rate": 2.0406438564462922e-07,
+      "loss": 0.7139,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 2.040643856446292e-05,
+      "step": 8534
+    },
+    {
+      "epoch": 0.87,
+      "gumbel_tau": 0.1335,
+      "image_masks_loss": 0.0168,
+      "llm_learning_rate": 2.0373622197544685e-07,
+      "loss": 0.7593,
+      "output_text_mask_loss": 0.0057,
+      "predictor_learning_rate": 2.0373622197544683e-05,
+      "step": 8535
+    },
+    {
+      "epoch": 0.87,
+      "gumbel_tau": 0.1335,
+      "image_masks_loss": 0.0154,
+      "llm_learning_rate": 2.0340831117594662e-07,
+      "loss": 0.7592,
+      "output_text_mask_loss": 0.0343,
+      "predictor_learning_rate": 2.034083111759466e-05,
+      "step": 8536
+    },
+    {
+      "epoch": 0.87,
+      "gumbel_tau": 0.1335,
+      "image_masks_loss": 0.0273,
+      "llm_learning_rate": 2.0308065328223824e-07,
+      "loss": 0.7358,
+      "output_text_mask_loss": 0.0238,
+      "predictor_learning_rate": 2.0308065328223824e-05,
+      "step": 8537
+    },
+    {
+      "epoch": 0.87,
+      "gumbel_tau": 0.1334,
+      "image_masks_loss": 0.012,
+      "llm_learning_rate": 2.0275324833040483e-07,
+      "loss": 0.7774,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 2.0275324833040483e-05,
+      "step": 8538
+    },
+    {
+      "epoch": 0.87,
+      "gumbel_tau": 0.1334,
+      "image_masks_loss": 0.0095,
+      "llm_learning_rate": 2.0242609635650057e-07,
+      "loss": 0.7225,
+      "output_text_mask_loss": 0.0081,
+      "predictor_learning_rate": 2.0242609635650054e-05,
+      "step": 8539
+    },
+    {
+      "epoch": 0.88,
+      "gumbel_tau": 0.1334,
+      "image_masks_loss": 0.0041,
+      "llm_learning_rate": 2.020991973965522e-07,
+      "loss": 0.7968,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 2.020991973965522e-05,
+      "step": 8540
+    },
+    {
+      "epoch": 0.88,
+      "gumbel_tau": 0.1334,
+      "image_masks_loss": 0.0127,
+      "llm_learning_rate": 2.0177255148655873e-07,
+      "loss": 0.6866,
+      "output_text_mask_loss": 0.0054,
+      "predictor_learning_rate": 2.0177255148655872e-05,
+      "step": 8541
+    },
+    {
+      "epoch": 0.88,
+      "gumbel_tau": 0.1333,
+      "image_masks_loss": 0.0106,
+      "llm_learning_rate": 2.0144615866249108e-07,
+      "loss": 0.7187,
+      "output_text_mask_loss": 0.0004,
+      "predictor_learning_rate": 2.0144615866249104e-05,
+      "step": 8542
+    },
+    {
+      "epoch": 0.88,
+      "gumbel_tau": 0.1333,
+      "image_masks_loss": 0.015,
+      "llm_learning_rate": 2.0112001896029216e-07,
+      "loss": 0.679,
+      "output_text_mask_loss": 0.0004,
+      "predictor_learning_rate": 2.0112001896029215e-05,
+      "step": 8543
+    },
+    {
+      "epoch": 0.88,
+      "gumbel_tau": 0.1333,
+      "image_masks_loss": 0.0153,
+      "llm_learning_rate": 2.0079413241587687e-07,
+      "loss": 0.7067,
+      "output_text_mask_loss": 0.0134,
+      "predictor_learning_rate": 2.0079413241587686e-05,
+      "step": 8544
+    },
+    {
+      "epoch": 0.88,
+      "gumbel_tau": 0.1332,
+      "image_masks_loss": 0.0073,
+      "llm_learning_rate": 2.0046849906513316e-07,
+      "loss": 0.6873,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 2.0046849906513315e-05,
+      "step": 8545
+    },
+    {
+      "epoch": 0.88,
+      "gumbel_tau": 0.1332,
+      "image_masks_loss": 0.0144,
+      "llm_learning_rate": 2.0014311894392068e-07,
+      "loss": 0.7429,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 2.0014311894392066e-05,
+      "step": 8546
+    },
+    {
+      "epoch": 0.88,
+      "gumbel_tau": 0.1332,
+      "image_masks_loss": 0.0063,
+      "llm_learning_rate": 1.9981799208807107e-07,
+      "loss": 0.7092,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 1.9981799208807106e-05,
+      "step": 8547
+    },
+    {
+      "epoch": 0.88,
+      "gumbel_tau": 0.1331,
+      "image_masks_loss": 0.0099,
+      "llm_learning_rate": 1.9949311853338705e-07,
+      "loss": 0.744,
+      "output_text_mask_loss": 0.0071,
+      "predictor_learning_rate": 1.9949311853338703e-05,
+      "step": 8548
+    },
+    {
+      "epoch": 0.88,
+      "gumbel_tau": 0.1331,
+      "image_masks_loss": 0.0057,
+      "llm_learning_rate": 1.9916849831564527e-07,
+      "loss": 0.7438,
+      "output_text_mask_loss": 0.0125,
+      "predictor_learning_rate": 1.9916849831564526e-05,
+      "step": 8549
+    },
+    {
+      "epoch": 0.88,
+      "gumbel_tau": 0.1331,
+      "image_masks_loss": 0.0128,
+      "llm_learning_rate": 1.9884413147059271e-07,
+      "loss": 0.6834,
+      "output_text_mask_loss": 0.0004,
+      "predictor_learning_rate": 1.9884413147059273e-05,
+      "step": 8550
+    },
+    {
+      "epoch": 0.88,
+      "gumbel_tau": 0.133,
+      "image_masks_loss": 0.0091,
+      "llm_learning_rate": 1.9852001803395026e-07,
+      "loss": 0.7851,
+      "output_text_mask_loss": 0.0046,
+      "predictor_learning_rate": 1.9852001803395025e-05,
+      "step": 8551
+    },
+    {
+      "epoch": 0.88,
+      "gumbel_tau": 0.133,
+      "image_masks_loss": 0.005,
+      "llm_learning_rate": 1.9819615804141017e-07,
+      "loss": 0.7318,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 1.981961580414102e-05,
+      "step": 8552
+    },
+    {
+      "epoch": 0.88,
+      "gumbel_tau": 0.133,
+      "image_masks_loss": 0.0093,
+      "llm_learning_rate": 1.9787255152863588e-07,
+      "loss": 0.6814,
+      "output_text_mask_loss": 0.0299,
+      "predictor_learning_rate": 1.978725515286359e-05,
+      "step": 8553
+    },
+    {
+      "epoch": 0.88,
+      "gumbel_tau": 0.1329,
+      "image_masks_loss": 0.0182,
+      "llm_learning_rate": 1.9754919853126414e-07,
+      "loss": 0.7319,
+      "output_text_mask_loss": 0.005,
+      "predictor_learning_rate": 1.9754919853126412e-05,
+      "step": 8554
+    },
+    {
+      "epoch": 0.88,
+      "gumbel_tau": 0.1329,
+      "image_masks_loss": 0.0182,
+      "llm_learning_rate": 1.9722609908490342e-07,
+      "loss": 0.7643,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 1.9722609908490342e-05,
+      "step": 8555
+    },
+    {
+      "epoch": 0.88,
+      "gumbel_tau": 0.1329,
+      "image_masks_loss": 0.0096,
+      "llm_learning_rate": 1.969032532251336e-07,
+      "loss": 0.7494,
+      "output_text_mask_loss": 0.0248,
+      "predictor_learning_rate": 1.969032532251336e-05,
+      "step": 8556
+    },
+    {
+      "epoch": 0.88,
+      "gumbel_tau": 0.1328,
+      "image_masks_loss": 0.012,
+      "llm_learning_rate": 1.9658066098750705e-07,
+      "loss": 0.7298,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 1.9658066098750704e-05,
+      "step": 8557
+    },
+    {
+      "epoch": 0.88,
+      "gumbel_tau": 0.1328,
+      "image_masks_loss": 0.0175,
+      "llm_learning_rate": 1.9625832240754928e-07,
+      "loss": 0.8302,
+      "output_text_mask_loss": 0.0103,
+      "predictor_learning_rate": 1.9625832240754928e-05,
+      "step": 8558
+    },
+    {
+      "epoch": 0.88,
+      "gumbel_tau": 0.1328,
+      "image_masks_loss": 0.0077,
+      "llm_learning_rate": 1.9593623752075657e-07,
+      "loss": 0.7927,
+      "output_text_mask_loss": 0.0156,
+      "predictor_learning_rate": 1.9593623752075656e-05,
+      "step": 8559
+    },
+    {
+      "epoch": 0.88,
+      "gumbel_tau": 0.1328,
+      "image_masks_loss": 0.0103,
+      "llm_learning_rate": 1.9561440636259727e-07,
+      "loss": 0.7245,
+      "output_text_mask_loss": 0.0052,
+      "predictor_learning_rate": 1.9561440636259727e-05,
+      "step": 8560
+    },
+    {
+      "epoch": 0.88,
+      "gumbel_tau": 0.1327,
+      "image_masks_loss": 0.0042,
+      "llm_learning_rate": 1.9529282896851244e-07,
+      "loss": 0.7726,
+      "output_text_mask_loss": 0.0305,
+      "predictor_learning_rate": 1.952928289685124e-05,
+      "step": 8561
+    },
+    {
+      "epoch": 0.88,
+      "gumbel_tau": 0.1327,
+      "image_masks_loss": 0.0156,
+      "llm_learning_rate": 1.9497150537391484e-07,
+      "loss": 0.744,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 1.9497150537391485e-05,
+      "step": 8562
+    },
+    {
+      "epoch": 0.88,
+      "gumbel_tau": 0.1327,
+      "image_masks_loss": 0.0118,
+      "llm_learning_rate": 1.9465043561418928e-07,
+      "loss": 0.7201,
+      "output_text_mask_loss": 0.0224,
+      "predictor_learning_rate": 1.9465043561418928e-05,
+      "step": 8563
+    },
+    {
+      "epoch": 0.88,
+      "gumbel_tau": 0.1326,
+      "image_masks_loss": 0.0224,
+      "llm_learning_rate": 1.9432961972469216e-07,
+      "loss": 0.7388,
+      "output_text_mask_loss": 0.0193,
+      "predictor_learning_rate": 1.9432961972469215e-05,
+      "step": 8564
+    },
+    {
+      "epoch": 0.88,
+      "gumbel_tau": 0.1326,
+      "image_masks_loss": 0.0101,
+      "llm_learning_rate": 1.9400905774075357e-07,
+      "loss": 0.6332,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 1.9400905774075356e-05,
+      "step": 8565
+    },
+    {
+      "epoch": 0.88,
+      "gumbel_tau": 0.1326,
+      "image_masks_loss": 0.0089,
+      "llm_learning_rate": 1.9368874969767443e-07,
+      "loss": 0.7918,
+      "output_text_mask_loss": 0.0355,
+      "predictor_learning_rate": 1.936887496976744e-05,
+      "step": 8566
+    },
+    {
+      "epoch": 0.88,
+      "gumbel_tau": 0.1325,
+      "image_masks_loss": 0.0119,
+      "llm_learning_rate": 1.933686956307268e-07,
+      "loss": 0.7099,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 1.933686956307268e-05,
+      "step": 8567
+    },
+    {
+      "epoch": 0.88,
+      "gumbel_tau": 0.1325,
+      "image_masks_loss": 0.0232,
+      "llm_learning_rate": 1.930488955751561e-07,
+      "loss": 0.7268,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 1.930488955751561e-05,
+      "step": 8568
+    },
+    {
+      "epoch": 0.88,
+      "gumbel_tau": 0.1325,
+      "image_masks_loss": 0.0043,
+      "llm_learning_rate": 1.9272934956617916e-07,
+      "loss": 0.7639,
+      "output_text_mask_loss": 0.0139,
+      "predictor_learning_rate": 1.9272934956617916e-05,
+      "step": 8569
+    },
+    {
+      "epoch": 0.88,
+      "gumbel_tau": 0.1324,
+      "image_masks_loss": 0.0114,
+      "llm_learning_rate": 1.9241005763898618e-07,
+      "loss": 0.745,
+      "output_text_mask_loss": 0.0212,
+      "predictor_learning_rate": 1.9241005763898616e-05,
+      "step": 8570
+    },
+    {
+      "epoch": 0.88,
+      "gumbel_tau": 0.1324,
+      "image_masks_loss": 0.0099,
+      "llm_learning_rate": 1.9209101982873734e-07,
+      "loss": 0.686,
+      "output_text_mask_loss": 0.0012,
+      "predictor_learning_rate": 1.9209101982873733e-05,
+      "step": 8571
+    },
+    {
+      "epoch": 0.88,
+      "gumbel_tau": 0.1324,
+      "image_masks_loss": 0.005,
+      "llm_learning_rate": 1.9177223617056622e-07,
+      "loss": 0.7101,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 1.917722361705662e-05,
+      "step": 8572
+    },
+    {
+      "epoch": 0.88,
+      "gumbel_tau": 0.1323,
+      "image_masks_loss": 0.0107,
+      "llm_learning_rate": 1.9145370669957751e-07,
+      "loss": 0.713,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 1.914537066995775e-05,
+      "step": 8573
+    },
+    {
+      "epoch": 0.88,
+      "gumbel_tau": 0.1323,
+      "image_masks_loss": 0.0137,
+      "llm_learning_rate": 1.91135431450849e-07,
+      "loss": 0.8207,
+      "output_text_mask_loss": 0.0467,
+      "predictor_learning_rate": 1.91135431450849e-05,
+      "step": 8574
+    },
+    {
+      "epoch": 0.88,
+      "gumbel_tau": 0.1323,
+      "image_masks_loss": 0.0298,
+      "llm_learning_rate": 1.908174104594296e-07,
+      "loss": 0.786,
+      "output_text_mask_loss": 0.0038,
+      "predictor_learning_rate": 1.9081741045942958e-05,
+      "step": 8575
+    },
+    {
+      "epoch": 0.88,
+      "gumbel_tau": 0.1323,
+      "image_masks_loss": 0.0104,
+      "llm_learning_rate": 1.9049964376033986e-07,
+      "loss": 0.7594,
+      "output_text_mask_loss": 0.0095,
+      "predictor_learning_rate": 1.9049964376033986e-05,
+      "step": 8576
+    },
+    {
+      "epoch": 0.88,
+      "gumbel_tau": 0.1322,
+      "image_masks_loss": 0.0048,
+      "llm_learning_rate": 1.9018213138857406e-07,
+      "loss": 0.7243,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 1.9018213138857403e-05,
+      "step": 8577
+    },
+    {
+      "epoch": 0.88,
+      "gumbel_tau": 0.1322,
+      "image_masks_loss": 0.009,
+      "llm_learning_rate": 1.8986487337909698e-07,
+      "loss": 0.71,
+      "output_text_mask_loss": 0.0077,
+      "predictor_learning_rate": 1.8986487337909696e-05,
+      "step": 8578
+    },
+    {
+      "epoch": 0.88,
+      "gumbel_tau": 0.1322,
+      "image_masks_loss": 0.0144,
+      "llm_learning_rate": 1.8954786976684542e-07,
+      "loss": 0.8143,
+      "output_text_mask_loss": 0.0164,
+      "predictor_learning_rate": 1.8954786976684542e-05,
+      "step": 8579
+    },
+    {
+      "epoch": 0.88,
+      "gumbel_tau": 0.1321,
+      "image_masks_loss": 0.0126,
+      "llm_learning_rate": 1.8923112058672894e-07,
+      "loss": 0.7759,
+      "output_text_mask_loss": 0.0193,
+      "predictor_learning_rate": 1.8923112058672893e-05,
+      "step": 8580
+    },
+    {
+      "epoch": 0.88,
+      "gumbel_tau": 0.1321,
+      "image_masks_loss": 0.018,
+      "llm_learning_rate": 1.8891462587362852e-07,
+      "loss": 0.7473,
+      "output_text_mask_loss": 0.016,
+      "predictor_learning_rate": 1.889146258736285e-05,
+      "step": 8581
+    },
+    {
+      "epoch": 0.88,
+      "gumbel_tau": 0.1321,
+      "image_masks_loss": 0.0083,
+      "llm_learning_rate": 1.8859838566239713e-07,
+      "loss": 0.7227,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 1.8859838566239712e-05,
+      "step": 8582
+    },
+    {
+      "epoch": 0.88,
+      "gumbel_tau": 0.132,
+      "image_masks_loss": 0.0107,
+      "llm_learning_rate": 1.8828239998785968e-07,
+      "loss": 0.6544,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 1.8828239998785968e-05,
+      "step": 8583
+    },
+    {
+      "epoch": 0.88,
+      "gumbel_tau": 0.132,
+      "image_masks_loss": 0.0049,
+      "llm_learning_rate": 1.8796666888481418e-07,
+      "loss": 0.736,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 1.8796666888481416e-05,
+      "step": 8584
+    },
+    {
+      "epoch": 0.88,
+      "gumbel_tau": 0.132,
+      "image_masks_loss": 0.007,
+      "llm_learning_rate": 1.8765119238802865e-07,
+      "loss": 0.7563,
+      "output_text_mask_loss": 0.036,
+      "predictor_learning_rate": 1.8765119238802863e-05,
+      "step": 8585
+    },
+    {
+      "epoch": 0.88,
+      "gumbel_tau": 0.1319,
+      "image_masks_loss": 0.0073,
+      "llm_learning_rate": 1.8733597053224446e-07,
+      "loss": 0.7376,
+      "output_text_mask_loss": 0.0039,
+      "predictor_learning_rate": 1.8733597053224445e-05,
+      "step": 8586
+    },
+    {
+      "epoch": 0.88,
+      "gumbel_tau": 0.1319,
+      "image_masks_loss": 0.0149,
+      "llm_learning_rate": 1.8702100335217438e-07,
+      "loss": 0.6779,
+      "output_text_mask_loss": 0.0107,
+      "predictor_learning_rate": 1.8702100335217438e-05,
+      "step": 8587
+    },
+    {
+      "epoch": 0.88,
+      "gumbel_tau": 0.1319,
+      "image_masks_loss": 0.0055,
+      "llm_learning_rate": 1.867062908825032e-07,
+      "loss": 0.7532,
+      "output_text_mask_loss": 0.0368,
+      "predictor_learning_rate": 1.867062908825032e-05,
+      "step": 8588
+    },
+    {
+      "epoch": 0.88,
+      "gumbel_tau": 0.1319,
+      "image_masks_loss": 0.0046,
+      "llm_learning_rate": 1.8639183315788763e-07,
+      "loss": 0.7513,
+      "output_text_mask_loss": 0.0185,
+      "predictor_learning_rate": 1.863918331578876e-05,
+      "step": 8589
+    },
+    {
+      "epoch": 0.88,
+      "gumbel_tau": 0.1318,
+      "image_masks_loss": 0.0153,
+      "llm_learning_rate": 1.8607763021295717e-07,
+      "loss": 0.7924,
+      "output_text_mask_loss": 0.0238,
+      "predictor_learning_rate": 1.8607763021295717e-05,
+      "step": 8590
+    },
+    {
+      "epoch": 0.88,
+      "gumbel_tau": 0.1318,
+      "image_masks_loss": 0.0041,
+      "llm_learning_rate": 1.857636820823122e-07,
+      "loss": 0.8279,
+      "output_text_mask_loss": 0.0087,
+      "predictor_learning_rate": 1.8576368208231222e-05,
+      "step": 8591
+    },
+    {
+      "epoch": 0.88,
+      "gumbel_tau": 0.1318,
+      "image_masks_loss": 0.004,
+      "llm_learning_rate": 1.8544998880052511e-07,
+      "loss": 0.7852,
+      "output_text_mask_loss": 0.0278,
+      "predictor_learning_rate": 1.854499888005251e-05,
+      "step": 8592
+    },
+    {
+      "epoch": 0.88,
+      "gumbel_tau": 0.1317,
+      "image_masks_loss": 0.0094,
+      "llm_learning_rate": 1.8513655040214102e-07,
+      "loss": 0.7266,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 1.85136550402141e-05,
+      "step": 8593
+    },
+    {
+      "epoch": 0.88,
+      "gumbel_tau": 0.1317,
+      "image_masks_loss": 0.014,
+      "llm_learning_rate": 1.8482336692167592e-07,
+      "loss": 0.7231,
+      "output_text_mask_loss": 0.0452,
+      "predictor_learning_rate": 1.848233669216759e-05,
+      "step": 8594
+    },
+    {
+      "epoch": 0.88,
+      "gumbel_tau": 0.1317,
+      "image_masks_loss": 0.0271,
+      "llm_learning_rate": 1.845104383936186e-07,
+      "loss": 0.6305,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 1.845104383936186e-05,
+      "step": 8595
+    },
+    {
+      "epoch": 0.88,
+      "gumbel_tau": 0.1316,
+      "image_masks_loss": 0.0083,
+      "llm_learning_rate": 1.8419776485242874e-07,
+      "loss": 0.7223,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 1.8419776485242872e-05,
+      "step": 8596
+    },
+    {
+      "epoch": 0.88,
+      "gumbel_tau": 0.1316,
+      "image_masks_loss": 0.0079,
+      "llm_learning_rate": 1.8388534633253991e-07,
+      "loss": 0.7008,
+      "output_text_mask_loss": 0.0078,
+      "predictor_learning_rate": 1.838853463325399e-05,
+      "step": 8597
+    },
+    {
+      "epoch": 0.88,
+      "gumbel_tau": 0.1316,
+      "image_masks_loss": 0.0064,
+      "llm_learning_rate": 1.8357318286835568e-07,
+      "loss": 0.6963,
+      "output_text_mask_loss": 0.006,
+      "predictor_learning_rate": 1.835731828683557e-05,
+      "step": 8598
+    },
+    {
+      "epoch": 0.88,
+      "gumbel_tau": 0.1315,
+      "image_masks_loss": 0.0055,
+      "llm_learning_rate": 1.8326127449425217e-07,
+      "loss": 0.7262,
+      "output_text_mask_loss": 0.0006,
+      "predictor_learning_rate": 1.8326127449425217e-05,
+      "step": 8599
+    },
+    {
+      "epoch": 0.88,
+      "gumbel_tau": 0.1315,
+      "image_masks_loss": 0.0073,
+      "llm_learning_rate": 1.8294962124457723e-07,
+      "loss": 0.6761,
+      "output_text_mask_loss": 0.0153,
+      "predictor_learning_rate": 1.8294962124457722e-05,
+      "step": 8600
+    },
+    {
+      "epoch": 0.88,
+      "gumbel_tau": 0.1315,
+      "image_masks_loss": 0.0146,
+      "llm_learning_rate": 1.8263822315365137e-07,
+      "loss": 0.616,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 1.8263822315365135e-05,
+      "step": 8601
+    },
+    {
+      "epoch": 0.88,
+      "gumbel_tau": 0.1314,
+      "image_masks_loss": 0.0138,
+      "llm_learning_rate": 1.8232708025576585e-07,
+      "loss": 0.7916,
+      "output_text_mask_loss": 0.001,
+      "predictor_learning_rate": 1.8232708025576584e-05,
+      "step": 8602
+    },
+    {
+      "epoch": 0.88,
+      "gumbel_tau": 0.1314,
+      "image_masks_loss": 0.0093,
+      "llm_learning_rate": 1.820161925851846e-07,
+      "loss": 0.7119,
+      "output_text_mask_loss": 0.001,
+      "predictor_learning_rate": 1.8201619258518458e-05,
+      "step": 8603
+    },
+    {
+      "epoch": 0.88,
+      "gumbel_tau": 0.1314,
+      "image_masks_loss": 0.0068,
+      "llm_learning_rate": 1.817055601761436e-07,
+      "loss": 0.7758,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 1.817055601761436e-05,
+      "step": 8604
+    },
+    {
+      "epoch": 0.88,
+      "gumbel_tau": 0.1314,
+      "image_masks_loss": 0.0037,
+      "llm_learning_rate": 1.8139518306284992e-07,
+      "loss": 0.7329,
+      "output_text_mask_loss": 0.0235,
+      "predictor_learning_rate": 1.813951830628499e-05,
+      "step": 8605
+    },
+    {
+      "epoch": 0.88,
+      "gumbel_tau": 0.1313,
+      "image_masks_loss": 0.0055,
+      "llm_learning_rate": 1.8108506127948317e-07,
+      "loss": 0.7362,
+      "output_text_mask_loss": 0.0072,
+      "predictor_learning_rate": 1.8108506127948315e-05,
+      "step": 8606
+    },
+    {
+      "epoch": 0.88,
+      "gumbel_tau": 0.1313,
+      "image_masks_loss": 0.0063,
+      "llm_learning_rate": 1.807751948601949e-07,
+      "loss": 0.7545,
+      "output_text_mask_loss": 0.0047,
+      "predictor_learning_rate": 1.807751948601949e-05,
+      "step": 8607
+    },
+    {
+      "epoch": 0.88,
+      "gumbel_tau": 0.1313,
+      "image_masks_loss": 0.0138,
+      "llm_learning_rate": 1.8046558383910783e-07,
+      "loss": 0.7725,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 1.804655838391078e-05,
+      "step": 8608
+    },
+    {
+      "epoch": 0.88,
+      "gumbel_tau": 0.1312,
+      "image_masks_loss": 0.0207,
+      "llm_learning_rate": 1.8015622825031692e-07,
+      "loss": 0.7217,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 1.801562282503169e-05,
+      "step": 8609
+    },
+    {
+      "epoch": 0.88,
+      "gumbel_tau": 0.1312,
+      "image_masks_loss": 0.0063,
+      "llm_learning_rate": 1.7984712812788964e-07,
+      "loss": 0.6658,
+      "output_text_mask_loss": 0.017,
+      "predictor_learning_rate": 1.798471281278896e-05,
+      "step": 8610
+    },
+    {
+      "epoch": 0.88,
+      "gumbel_tau": 0.1312,
+      "image_masks_loss": 0.0086,
+      "llm_learning_rate": 1.7953828350586456e-07,
+      "loss": 0.7332,
+      "output_text_mask_loss": 0.0261,
+      "predictor_learning_rate": 1.7953828350586453e-05,
+      "step": 8611
+    },
+    {
+      "epoch": 0.88,
+      "gumbel_tau": 0.1311,
+      "image_masks_loss": 0.0102,
+      "llm_learning_rate": 1.7922969441825256e-07,
+      "loss": 0.7214,
+      "output_text_mask_loss": 0.0096,
+      "predictor_learning_rate": 1.7922969441825253e-05,
+      "step": 8612
+    },
+    {
+      "epoch": 0.88,
+      "gumbel_tau": 0.1311,
+      "image_masks_loss": 0.0128,
+      "llm_learning_rate": 1.789213608990359e-07,
+      "loss": 0.6941,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 1.789213608990359e-05,
+      "step": 8613
+    },
+    {
+      "epoch": 0.88,
+      "gumbel_tau": 0.1311,
+      "image_masks_loss": 0.006,
+      "llm_learning_rate": 1.786132829821688e-07,
+      "loss": 0.6639,
+      "output_text_mask_loss": 0.0025,
+      "predictor_learning_rate": 1.7861328298216876e-05,
+      "step": 8614
+    },
+    {
+      "epoch": 0.88,
+      "gumbel_tau": 0.131,
+      "image_masks_loss": 0.012,
+      "llm_learning_rate": 1.7830546070157773e-07,
+      "loss": 0.7239,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 1.783054607015777e-05,
+      "step": 8615
+    },
+    {
+      "epoch": 0.88,
+      "gumbel_tau": 0.131,
+      "image_masks_loss": 0.024,
+      "llm_learning_rate": 1.779978940911603e-07,
+      "loss": 0.6968,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 1.779978940911603e-05,
+      "step": 8616
+    },
+    {
+      "epoch": 0.88,
+      "gumbel_tau": 0.131,
+      "image_masks_loss": 0.0118,
+      "llm_learning_rate": 1.7769058318478722e-07,
+      "loss": 0.7611,
+      "output_text_mask_loss": 0.0124,
+      "predictor_learning_rate": 1.776905831847872e-05,
+      "step": 8617
+    },
+    {
+      "epoch": 0.88,
+      "gumbel_tau": 0.131,
+      "image_masks_loss": 0.0122,
+      "llm_learning_rate": 1.7738352801629972e-07,
+      "loss": 0.7129,
+      "output_text_mask_loss": 0.0515,
+      "predictor_learning_rate": 1.773835280162997e-05,
+      "step": 8618
+    },
+    {
+      "epoch": 0.88,
+      "gumbel_tau": 0.1309,
+      "image_masks_loss": 0.0082,
+      "llm_learning_rate": 1.7707672861951186e-07,
+      "loss": 0.7524,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 1.7707672861951186e-05,
+      "step": 8619
+    },
+    {
+      "epoch": 0.88,
+      "gumbel_tau": 0.1309,
+      "image_masks_loss": 0.0212,
+      "llm_learning_rate": 1.7677018502820858e-07,
+      "loss": 0.787,
+      "output_text_mask_loss": 0.0698,
+      "predictor_learning_rate": 1.7677018502820858e-05,
+      "step": 8620
+    },
+    {
+      "epoch": 0.88,
+      "gumbel_tau": 0.1309,
+      "image_masks_loss": 0.0054,
+      "llm_learning_rate": 1.7646389727614733e-07,
+      "loss": 0.7239,
+      "output_text_mask_loss": 0.001,
+      "predictor_learning_rate": 1.764638972761473e-05,
+      "step": 8621
+    },
+    {
+      "epoch": 0.88,
+      "gumbel_tau": 0.1308,
+      "image_masks_loss": 0.003,
+      "llm_learning_rate": 1.7615786539705693e-07,
+      "loss": 0.7097,
+      "output_text_mask_loss": 0.0282,
+      "predictor_learning_rate": 1.761578653970569e-05,
+      "step": 8622
+    },
+    {
+      "epoch": 0.88,
+      "gumbel_tau": 0.1308,
+      "image_masks_loss": 0.0108,
+      "llm_learning_rate": 1.7585208942463878e-07,
+      "loss": 0.7309,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 1.7585208942463878e-05,
+      "step": 8623
+    },
+    {
+      "epoch": 0.88,
+      "gumbel_tau": 0.1308,
+      "image_masks_loss": 0.0045,
+      "llm_learning_rate": 1.7554656939256565e-07,
+      "loss": 0.5805,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 1.7554656939256562e-05,
+      "step": 8624
+    },
+    {
+      "epoch": 0.88,
+      "gumbel_tau": 0.1307,
+      "image_masks_loss": 0.008,
+      "llm_learning_rate": 1.7524130533448147e-07,
+      "loss": 0.6891,
+      "output_text_mask_loss": 0.0338,
+      "predictor_learning_rate": 1.7524130533448146e-05,
+      "step": 8625
+    },
+    {
+      "epoch": 0.88,
+      "gumbel_tau": 0.1307,
+      "image_masks_loss": 0.0059,
+      "llm_learning_rate": 1.7493629728400296e-07,
+      "loss": 0.7804,
+      "output_text_mask_loss": 0.0086,
+      "predictor_learning_rate": 1.7493629728400295e-05,
+      "step": 8626
+    },
+    {
+      "epoch": 0.88,
+      "gumbel_tau": 0.1307,
+      "image_masks_loss": 0.0179,
+      "llm_learning_rate": 1.7463154527471826e-07,
+      "loss": 0.7631,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 1.7463154527471824e-05,
+      "step": 8627
+    },
+    {
+      "epoch": 0.88,
+      "gumbel_tau": 0.1306,
+      "image_masks_loss": 0.0067,
+      "llm_learning_rate": 1.743270493401872e-07,
+      "loss": 0.6993,
+      "output_text_mask_loss": 0.0557,
+      "predictor_learning_rate": 1.7432704934018715e-05,
+      "step": 8628
+    },
+    {
+      "epoch": 0.88,
+      "gumbel_tau": 0.1306,
+      "image_masks_loss": 0.0121,
+      "llm_learning_rate": 1.7402280951394178e-07,
+      "loss": 0.7011,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 1.740228095139418e-05,
+      "step": 8629
+    },
+    {
+      "epoch": 0.88,
+      "gumbel_tau": 0.1306,
+      "image_masks_loss": 0.0068,
+      "llm_learning_rate": 1.7371882582948557e-07,
+      "loss": 0.6949,
+      "output_text_mask_loss": 0.0056,
+      "predictor_learning_rate": 1.7371882582948555e-05,
+      "step": 8630
+    },
+    {
+      "epoch": 0.88,
+      "gumbel_tau": 0.1306,
+      "image_masks_loss": 0.0146,
+      "llm_learning_rate": 1.7341509832029397e-07,
+      "loss": 0.7075,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 1.7341509832029396e-05,
+      "step": 8631
+    },
+    {
+      "epoch": 0.88,
+      "gumbel_tau": 0.1305,
+      "image_masks_loss": 0.0127,
+      "llm_learning_rate": 1.731116270198141e-07,
+      "loss": 0.7315,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 1.731116270198141e-05,
+      "step": 8632
+    },
+    {
+      "epoch": 0.88,
+      "gumbel_tau": 0.1305,
+      "image_masks_loss": 0.0024,
+      "llm_learning_rate": 1.7280841196146453e-07,
+      "loss": 0.7044,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 1.7280841196146452e-05,
+      "step": 8633
+    },
+    {
+      "epoch": 0.88,
+      "gumbel_tau": 0.1305,
+      "image_masks_loss": 0.0093,
+      "llm_learning_rate": 1.725054531786363e-07,
+      "loss": 0.7468,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 1.725054531786363e-05,
+      "step": 8634
+    },
+    {
+      "epoch": 0.88,
+      "gumbel_tau": 0.1304,
+      "image_masks_loss": 0.0175,
+      "llm_learning_rate": 1.7220275070469167e-07,
+      "loss": 0.6988,
+      "output_text_mask_loss": 0.0004,
+      "predictor_learning_rate": 1.7220275070469167e-05,
+      "step": 8635
+    },
+    {
+      "epoch": 0.88,
+      "gumbel_tau": 0.1304,
+      "image_masks_loss": 0.0202,
+      "llm_learning_rate": 1.7190030457296531e-07,
+      "loss": 0.741,
+      "output_text_mask_loss": 0.0088,
+      "predictor_learning_rate": 1.719003045729653e-05,
+      "step": 8636
+    },
+    {
+      "epoch": 0.88,
+      "gumbel_tau": 0.1304,
+      "image_masks_loss": 0.0077,
+      "llm_learning_rate": 1.7159811481676308e-07,
+      "loss": 0.6932,
+      "output_text_mask_loss": 0.0103,
+      "predictor_learning_rate": 1.715981148167631e-05,
+      "step": 8637
+    },
+    {
+      "epoch": 0.89,
+      "gumbel_tau": 0.1303,
+      "image_masks_loss": 0.007,
+      "llm_learning_rate": 1.7129618146936284e-07,
+      "loss": 0.8359,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 1.712961814693628e-05,
+      "step": 8638
+    },
+    {
+      "epoch": 0.89,
+      "gumbel_tau": 0.1303,
+      "image_masks_loss": 0.0088,
+      "llm_learning_rate": 1.7099450456401407e-07,
+      "loss": 0.7899,
+      "output_text_mask_loss": 0.036,
+      "predictor_learning_rate": 1.7099450456401404e-05,
+      "step": 8639
+    },
+    {
+      "epoch": 0.89,
+      "gumbel_tau": 0.1303,
+      "image_masks_loss": 0.011,
+      "llm_learning_rate": 1.7069308413393853e-07,
+      "loss": 0.6801,
+      "output_text_mask_loss": 0.0039,
+      "predictor_learning_rate": 1.7069308413393852e-05,
+      "step": 8640
+    },
+    {
+      "epoch": 0.89,
+      "gumbel_tau": 0.1302,
+      "image_masks_loss": 0.0069,
+      "llm_learning_rate": 1.7039192021232803e-07,
+      "loss": 0.7484,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 1.70391920212328e-05,
+      "step": 8641
+    },
+    {
+      "epoch": 0.89,
+      "gumbel_tau": 0.1302,
+      "image_masks_loss": 0.0104,
+      "llm_learning_rate": 1.7009101283234852e-07,
+      "loss": 0.8015,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 1.700910128323485e-05,
+      "step": 8642
+    },
+    {
+      "epoch": 0.89,
+      "gumbel_tau": 0.1302,
+      "image_masks_loss": 0.0049,
+      "llm_learning_rate": 1.6979036202713656e-07,
+      "loss": 0.7536,
+      "output_text_mask_loss": 0.0231,
+      "predictor_learning_rate": 1.6979036202713656e-05,
+      "step": 8643
+    },
+    {
+      "epoch": 0.89,
+      "gumbel_tau": 0.1302,
+      "image_masks_loss": 0.0043,
+      "llm_learning_rate": 1.694899678298001e-07,
+      "loss": 0.7102,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 1.694899678298001e-05,
+      "step": 8644
+    },
+    {
+      "epoch": 0.89,
+      "gumbel_tau": 0.1301,
+      "image_masks_loss": 0.0065,
+      "llm_learning_rate": 1.6918983027341963e-07,
+      "loss": 0.7751,
+      "output_text_mask_loss": 0.0307,
+      "predictor_learning_rate": 1.691898302734196e-05,
+      "step": 8645
+    },
+    {
+      "epoch": 0.89,
+      "gumbel_tau": 0.1301,
+      "image_masks_loss": 0.0131,
+      "llm_learning_rate": 1.6888994939104647e-07,
+      "loss": 0.7023,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 1.6888994939104647e-05,
+      "step": 8646
+    },
+    {
+      "epoch": 0.89,
+      "gumbel_tau": 0.1301,
+      "image_masks_loss": 0.0037,
+      "llm_learning_rate": 1.685903252157045e-07,
+      "loss": 0.7675,
+      "output_text_mask_loss": 0.0278,
+      "predictor_learning_rate": 1.685903252157045e-05,
+      "step": 8647
+    },
+    {
+      "epoch": 0.89,
+      "gumbel_tau": 0.13,
+      "image_masks_loss": 0.0057,
+      "llm_learning_rate": 1.6829095778038872e-07,
+      "loss": 0.6863,
+      "output_text_mask_loss": 0.0076,
+      "predictor_learning_rate": 1.6829095778038872e-05,
+      "step": 8648
+    },
+    {
+      "epoch": 0.89,
+      "gumbel_tau": 0.13,
+      "image_masks_loss": 0.0038,
+      "llm_learning_rate": 1.6799184711806667e-07,
+      "loss": 0.6964,
+      "output_text_mask_loss": 0.0019,
+      "predictor_learning_rate": 1.6799184711806664e-05,
+      "step": 8649
+    },
+    {
+      "epoch": 0.89,
+      "gumbel_tau": 0.13,
+      "image_masks_loss": 0.0201,
+      "llm_learning_rate": 1.6769299326167666e-07,
+      "loss": 0.7222,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 1.6769299326167664e-05,
+      "step": 8650
+    },
+    {
+      "epoch": 0.89,
+      "gumbel_tau": 0.1299,
+      "image_masks_loss": 0.0041,
+      "llm_learning_rate": 1.6739439624412935e-07,
+      "loss": 0.7474,
+      "output_text_mask_loss": 0.0418,
+      "predictor_learning_rate": 1.6739439624412934e-05,
+      "step": 8651
+    },
+    {
+      "epoch": 0.89,
+      "gumbel_tau": 0.1299,
+      "image_masks_loss": 0.0059,
+      "llm_learning_rate": 1.6709605609830677e-07,
+      "loss": 0.7276,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 1.6709605609830675e-05,
+      "step": 8652
+    },
+    {
+      "epoch": 0.89,
+      "gumbel_tau": 0.1299,
+      "image_masks_loss": 0.0047,
+      "llm_learning_rate": 1.6679797285706262e-07,
+      "loss": 0.7684,
+      "output_text_mask_loss": 0.0557,
+      "predictor_learning_rate": 1.667979728570626e-05,
+      "step": 8653
+    },
+    {
+      "epoch": 0.89,
+      "gumbel_tau": 0.1298,
+      "image_masks_loss": 0.0196,
+      "llm_learning_rate": 1.6650014655322288e-07,
+      "loss": 0.728,
+      "output_text_mask_loss": 0.0437,
+      "predictor_learning_rate": 1.6650014655322286e-05,
+      "step": 8654
+    },
+    {
+      "epoch": 0.89,
+      "gumbel_tau": 0.1298,
+      "image_masks_loss": 0.0083,
+      "llm_learning_rate": 1.662025772195844e-07,
+      "loss": 0.7038,
+      "output_text_mask_loss": 0.0076,
+      "predictor_learning_rate": 1.6620257721958436e-05,
+      "step": 8655
+    },
+    {
+      "epoch": 0.89,
+      "gumbel_tau": 0.1298,
+      "image_masks_loss": 0.003,
+      "llm_learning_rate": 1.6590526488891672e-07,
+      "loss": 0.8053,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 1.6590526488891673e-05,
+      "step": 8656
+    },
+    {
+      "epoch": 0.89,
+      "gumbel_tau": 0.1298,
+      "image_masks_loss": 0.0059,
+      "llm_learning_rate": 1.6560820959396012e-07,
+      "loss": 0.7448,
+      "output_text_mask_loss": 0.001,
+      "predictor_learning_rate": 1.6560820959396012e-05,
+      "step": 8657
+    },
+    {
+      "epoch": 0.89,
+      "gumbel_tau": 0.1297,
+      "image_masks_loss": 0.0109,
+      "llm_learning_rate": 1.6531141136742783e-07,
+      "loss": 0.7074,
+      "output_text_mask_loss": 0.0206,
+      "predictor_learning_rate": 1.6531141136742784e-05,
+      "step": 8658
+    },
+    {
+      "epoch": 0.89,
+      "gumbel_tau": 0.1297,
+      "image_masks_loss": 0.0078,
+      "llm_learning_rate": 1.6501487024200264e-07,
+      "loss": 0.7162,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 1.6501487024200262e-05,
+      "step": 8659
+    },
+    {
+      "epoch": 0.89,
+      "gumbel_tau": 0.1297,
+      "image_masks_loss": 0.0252,
+      "llm_learning_rate": 1.647185862503406e-07,
+      "loss": 0.8028,
+      "output_text_mask_loss": 0.0248,
+      "predictor_learning_rate": 1.6471858625034057e-05,
+      "step": 8660
+    },
+    {
+      "epoch": 0.89,
+      "gumbel_tau": 0.1296,
+      "image_masks_loss": 0.0027,
+      "llm_learning_rate": 1.644225594250698e-07,
+      "loss": 0.7483,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 1.6442255942506978e-05,
+      "step": 8661
+    },
+    {
+      "epoch": 0.89,
+      "gumbel_tau": 0.1296,
+      "image_masks_loss": 0.0206,
+      "llm_learning_rate": 1.641267897987889e-07,
+      "loss": 0.702,
+      "output_text_mask_loss": 0.0132,
+      "predictor_learning_rate": 1.6412678979878887e-05,
+      "step": 8662
+    },
+    {
+      "epoch": 0.89,
+      "gumbel_tau": 0.1296,
+      "image_masks_loss": 0.0416,
+      "llm_learning_rate": 1.63831277404069e-07,
+      "loss": 0.7653,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 1.63831277404069e-05,
+      "step": 8663
+    },
+    {
+      "epoch": 0.89,
+      "gumbel_tau": 0.1295,
+      "image_masks_loss": 0.0278,
+      "llm_learning_rate": 1.635360222734525e-07,
+      "loss": 0.6712,
+      "output_text_mask_loss": 0.0101,
+      "predictor_learning_rate": 1.635360222734525e-05,
+      "step": 8664
+    },
+    {
+      "epoch": 0.89,
+      "gumbel_tau": 0.1295,
+      "image_masks_loss": 0.0085,
+      "llm_learning_rate": 1.6324102443945362e-07,
+      "loss": 0.7047,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 1.632410244394536e-05,
+      "step": 8665
+    },
+    {
+      "epoch": 0.89,
+      "gumbel_tau": 0.1295,
+      "image_masks_loss": 0.0116,
+      "llm_learning_rate": 1.6294628393455803e-07,
+      "loss": 0.7628,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 1.62946283934558e-05,
+      "step": 8666
+    },
+    {
+      "epoch": 0.89,
+      "gumbel_tau": 0.1294,
+      "image_masks_loss": 0.0046,
+      "llm_learning_rate": 1.6265180079122256e-07,
+      "loss": 0.7209,
+      "output_text_mask_loss": 0.0058,
+      "predictor_learning_rate": 1.6265180079122256e-05,
+      "step": 8667
+    },
+    {
+      "epoch": 0.89,
+      "gumbel_tau": 0.1294,
+      "image_masks_loss": 0.0099,
+      "llm_learning_rate": 1.6235757504187765e-07,
+      "loss": 0.7208,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 1.6235757504187763e-05,
+      "step": 8668
+    },
+    {
+      "epoch": 0.89,
+      "gumbel_tau": 0.1294,
+      "image_masks_loss": 0.0229,
+      "llm_learning_rate": 1.6206360671892352e-07,
+      "loss": 0.7646,
+      "output_text_mask_loss": 0.037,
+      "predictor_learning_rate": 1.620636067189235e-05,
+      "step": 8669
+    },
+    {
+      "epoch": 0.89,
+      "gumbel_tau": 0.1294,
+      "image_masks_loss": 0.0107,
+      "llm_learning_rate": 1.617698958547323e-07,
+      "loss": 0.8163,
+      "output_text_mask_loss": 0.0315,
+      "predictor_learning_rate": 1.617698958547323e-05,
+      "step": 8670
+    },
+    {
+      "epoch": 0.89,
+      "gumbel_tau": 0.1293,
+      "image_masks_loss": 0.0105,
+      "llm_learning_rate": 1.614764424816484e-07,
+      "loss": 0.7462,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 1.614764424816484e-05,
+      "step": 8671
+    },
+    {
+      "epoch": 0.89,
+      "gumbel_tau": 0.1293,
+      "image_masks_loss": 0.0244,
+      "llm_learning_rate": 1.6118324663198765e-07,
+      "loss": 0.7463,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 1.6118324663198764e-05,
+      "step": 8672
+    },
+    {
+      "epoch": 0.89,
+      "gumbel_tau": 0.1293,
+      "image_masks_loss": 0.0092,
+      "llm_learning_rate": 1.6089030833803726e-07,
+      "loss": 0.7468,
+      "output_text_mask_loss": 0.013,
+      "predictor_learning_rate": 1.6089030833803726e-05,
+      "step": 8673
+    },
+    {
+      "epoch": 0.89,
+      "gumbel_tau": 0.1292,
+      "image_masks_loss": 0.0199,
+      "llm_learning_rate": 1.605976276320559e-07,
+      "loss": 0.7127,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 1.6059762763205587e-05,
+      "step": 8674
+    },
+    {
+      "epoch": 0.89,
+      "gumbel_tau": 0.1292,
+      "image_masks_loss": 0.0055,
+      "llm_learning_rate": 1.6030520454627495e-07,
+      "loss": 0.7348,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 1.6030520454627496e-05,
+      "step": 8675
+    },
+    {
+      "epoch": 0.89,
+      "gumbel_tau": 0.1292,
+      "image_masks_loss": 0.013,
+      "llm_learning_rate": 1.600130391128965e-07,
+      "loss": 0.6031,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 1.600130391128965e-05,
+      "step": 8676
+    },
+    {
+      "epoch": 0.89,
+      "gumbel_tau": 0.1291,
+      "image_masks_loss": 0.0192,
+      "llm_learning_rate": 1.597211313640945e-07,
+      "loss": 0.6933,
+      "output_text_mask_loss": 0.0374,
+      "predictor_learning_rate": 1.597211313640945e-05,
+      "step": 8677
+    },
+    {
+      "epoch": 0.89,
+      "gumbel_tau": 0.1291,
+      "image_masks_loss": 0.0063,
+      "llm_learning_rate": 1.594294813320141e-07,
+      "loss": 0.723,
+      "output_text_mask_loss": 0.017,
+      "predictor_learning_rate": 1.5942948133201407e-05,
+      "step": 8678
+    },
+    {
+      "epoch": 0.89,
+      "gumbel_tau": 0.1291,
+      "image_masks_loss": 0.0182,
+      "llm_learning_rate": 1.5913808904877233e-07,
+      "loss": 0.7199,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 1.591380890487723e-05,
+      "step": 8679
+    },
+    {
+      "epoch": 0.89,
+      "gumbel_tau": 0.1291,
+      "image_masks_loss": 0.0028,
+      "llm_learning_rate": 1.58846954546458e-07,
+      "loss": 0.755,
+      "output_text_mask_loss": 0.021,
+      "predictor_learning_rate": 1.58846954546458e-05,
+      "step": 8680
+    },
+    {
+      "epoch": 0.89,
+      "gumbel_tau": 0.129,
+      "image_masks_loss": 0.0072,
+      "llm_learning_rate": 1.5855607785713212e-07,
+      "loss": 0.7197,
+      "output_text_mask_loss": 0.0208,
+      "predictor_learning_rate": 1.585560778571321e-05,
+      "step": 8681
+    },
+    {
+      "epoch": 0.89,
+      "gumbel_tau": 0.129,
+      "image_masks_loss": 0.0343,
+      "llm_learning_rate": 1.5826545901282659e-07,
+      "loss": 0.6895,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 1.5826545901282658e-05,
+      "step": 8682
+    },
+    {
+      "epoch": 0.89,
+      "gumbel_tau": 0.129,
+      "image_masks_loss": 0.0083,
+      "llm_learning_rate": 1.579750980455444e-07,
+      "loss": 0.7814,
+      "output_text_mask_loss": 0.013,
+      "predictor_learning_rate": 1.579750980455444e-05,
+      "step": 8683
+    },
+    {
+      "epoch": 0.89,
+      "gumbel_tau": 0.1289,
+      "image_masks_loss": 0.0064,
+      "llm_learning_rate": 1.5768499498726108e-07,
+      "loss": 0.7366,
+      "output_text_mask_loss": 0.0391,
+      "predictor_learning_rate": 1.576849949872611e-05,
+      "step": 8684
+    },
+    {
+      "epoch": 0.89,
+      "gumbel_tau": 0.1289,
+      "image_masks_loss": 0.0051,
+      "llm_learning_rate": 1.5739514986992361e-07,
+      "loss": 0.707,
+      "output_text_mask_loss": 0.0124,
+      "predictor_learning_rate": 1.573951498699236e-05,
+      "step": 8685
+    },
+    {
+      "epoch": 0.89,
+      "gumbel_tau": 0.1289,
+      "image_masks_loss": 0.0119,
+      "llm_learning_rate": 1.5710556272545003e-07,
+      "loss": 0.7113,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 1.5710556272545002e-05,
+      "step": 8686
+    },
+    {
+      "epoch": 0.89,
+      "gumbel_tau": 0.1288,
+      "image_masks_loss": 0.0074,
+      "llm_learning_rate": 1.568162335857301e-07,
+      "loss": 0.8055,
+      "output_text_mask_loss": 0.0015,
+      "predictor_learning_rate": 1.568162335857301e-05,
+      "step": 8687
+    },
+    {
+      "epoch": 0.89,
+      "gumbel_tau": 0.1288,
+      "image_masks_loss": 0.009,
+      "llm_learning_rate": 1.5652716248262618e-07,
+      "loss": 0.828,
+      "output_text_mask_loss": 0.0174,
+      "predictor_learning_rate": 1.5652716248262616e-05,
+      "step": 8688
+    },
+    {
+      "epoch": 0.89,
+      "gumbel_tau": 0.1288,
+      "image_masks_loss": 0.0077,
+      "llm_learning_rate": 1.5623834944797105e-07,
+      "loss": 0.7554,
+      "output_text_mask_loss": 0.0089,
+      "predictor_learning_rate": 1.5623834944797105e-05,
+      "step": 8689
+    },
+    {
+      "epoch": 0.89,
+      "gumbel_tau": 0.1287,
+      "image_masks_loss": 0.0259,
+      "llm_learning_rate": 1.5594979451356935e-07,
+      "loss": 0.7625,
+      "output_text_mask_loss": 0.0429,
+      "predictor_learning_rate": 1.5594979451356934e-05,
+      "step": 8690
+    },
+    {
+      "epoch": 0.89,
+      "gumbel_tau": 0.1287,
+      "image_masks_loss": 0.0087,
+      "llm_learning_rate": 1.5566149771119727e-07,
+      "loss": 0.7291,
+      "output_text_mask_loss": 0.0046,
+      "predictor_learning_rate": 1.5566149771119724e-05,
+      "step": 8691
+    },
+    {
+      "epoch": 0.89,
+      "gumbel_tau": 0.1287,
+      "image_masks_loss": 0.007,
+      "llm_learning_rate": 1.5537345907260303e-07,
+      "loss": 0.7401,
+      "output_text_mask_loss": 0.0023,
+      "predictor_learning_rate": 1.5537345907260304e-05,
+      "step": 8692
+    },
+    {
+      "epoch": 0.89,
+      "gumbel_tau": 0.1287,
+      "image_masks_loss": 0.0173,
+      "llm_learning_rate": 1.5508567862950603e-07,
+      "loss": 0.6947,
+      "output_text_mask_loss": 0.0023,
+      "predictor_learning_rate": 1.55085678629506e-05,
+      "step": 8693
+    },
+    {
+      "epoch": 0.89,
+      "gumbel_tau": 0.1286,
+      "image_masks_loss": 0.0267,
+      "llm_learning_rate": 1.5479815641359668e-07,
+      "loss": 0.7143,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 1.5479815641359668e-05,
+      "step": 8694
+    },
+    {
+      "epoch": 0.89,
+      "gumbel_tau": 0.1286,
+      "image_masks_loss": 0.0097,
+      "llm_learning_rate": 1.5451089245653888e-07,
+      "loss": 0.6541,
+      "output_text_mask_loss": 0.0124,
+      "predictor_learning_rate": 1.5451089245653888e-05,
+      "step": 8695
+    },
+    {
+      "epoch": 0.89,
+      "gumbel_tau": 0.1286,
+      "image_masks_loss": 0.0078,
+      "llm_learning_rate": 1.5422388678996564e-07,
+      "loss": 0.6938,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 1.5422388678996563e-05,
+      "step": 8696
+    },
+    {
+      "epoch": 0.89,
+      "gumbel_tau": 0.1285,
+      "image_masks_loss": 0.0111,
+      "llm_learning_rate": 1.5393713944548304e-07,
+      "loss": 0.7884,
+      "output_text_mask_loss": 0.0014,
+      "predictor_learning_rate": 1.5393713944548303e-05,
+      "step": 8697
+    },
+    {
+      "epoch": 0.89,
+      "gumbel_tau": 0.1285,
+      "image_masks_loss": 0.0269,
+      "llm_learning_rate": 1.536506504546681e-07,
+      "loss": 0.7801,
+      "output_text_mask_loss": 0.0248,
+      "predictor_learning_rate": 1.5365065045466808e-05,
+      "step": 8698
+    },
+    {
+      "epoch": 0.89,
+      "gumbel_tau": 0.1285,
+      "image_masks_loss": 0.0116,
+      "llm_learning_rate": 1.533644198490697e-07,
+      "loss": 0.7326,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 1.5336441984906967e-05,
+      "step": 8699
+    },
+    {
+      "epoch": 0.89,
+      "gumbel_tau": 0.1284,
+      "image_masks_loss": 0.0067,
+      "llm_learning_rate": 1.5307844766020847e-07,
+      "loss": 0.8424,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 1.5307844766020846e-05,
+      "step": 8700
+    },
+    {
+      "epoch": 0.89,
+      "gumbel_tau": 0.1284,
+      "image_masks_loss": 0.0098,
+      "llm_learning_rate": 1.5279273391957615e-07,
+      "loss": 0.6903,
+      "output_text_mask_loss": 0.0488,
+      "predictor_learning_rate": 1.5279273391957616e-05,
+      "step": 8701
+    },
+    {
+      "epoch": 0.89,
+      "gumbel_tau": 0.1284,
+      "image_masks_loss": 0.0052,
+      "llm_learning_rate": 1.5250727865863623e-07,
+      "loss": 0.733,
+      "output_text_mask_loss": 0.0023,
+      "predictor_learning_rate": 1.5250727865863623e-05,
+      "step": 8702
+    },
+    {
+      "epoch": 0.89,
+      "gumbel_tau": 0.1284,
+      "image_masks_loss": 0.0089,
+      "llm_learning_rate": 1.5222208190882382e-07,
+      "loss": 0.7415,
+      "output_text_mask_loss": 0.0095,
+      "predictor_learning_rate": 1.5222208190882381e-05,
+      "step": 8703
+    },
+    {
+      "epoch": 0.89,
+      "gumbel_tau": 0.1283,
+      "image_masks_loss": 0.0265,
+      "llm_learning_rate": 1.5193714370154488e-07,
+      "loss": 0.7415,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 1.5193714370154487e-05,
+      "step": 8704
+    },
+    {
+      "epoch": 0.89,
+      "gumbel_tau": 0.1283,
+      "image_masks_loss": 0.007,
+      "llm_learning_rate": 1.5165246406817798e-07,
+      "loss": 0.7059,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 1.5165246406817796e-05,
+      "step": 8705
+    },
+    {
+      "epoch": 0.89,
+      "gumbel_tau": 0.1283,
+      "image_masks_loss": 0.0178,
+      "llm_learning_rate": 1.5136804304007213e-07,
+      "loss": 0.7722,
+      "output_text_mask_loss": 0.0103,
+      "predictor_learning_rate": 1.5136804304007212e-05,
+      "step": 8706
+    },
+    {
+      "epoch": 0.89,
+      "gumbel_tau": 0.1282,
+      "image_masks_loss": 0.0214,
+      "llm_learning_rate": 1.51083880648549e-07,
+      "loss": 0.6699,
+      "output_text_mask_loss": 0.0141,
+      "predictor_learning_rate": 1.51083880648549e-05,
+      "step": 8707
+    },
+    {
+      "epoch": 0.89,
+      "gumbel_tau": 0.1282,
+      "image_masks_loss": 0.0039,
+      "llm_learning_rate": 1.5079997692490105e-07,
+      "loss": 0.7648,
+      "output_text_mask_loss": 0.0246,
+      "predictor_learning_rate": 1.5079997692490105e-05,
+      "step": 8708
+    },
+    {
+      "epoch": 0.89,
+      "gumbel_tau": 0.1282,
+      "image_masks_loss": 0.0064,
+      "llm_learning_rate": 1.5051633190039216e-07,
+      "loss": 0.6681,
+      "output_text_mask_loss": 0.0025,
+      "predictor_learning_rate": 1.5051633190039216e-05,
+      "step": 8709
+    },
+    {
+      "epoch": 0.89,
+      "gumbel_tau": 0.1281,
+      "image_masks_loss": 0.0022,
+      "llm_learning_rate": 1.5023294560625817e-07,
+      "loss": 0.7461,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 1.5023294560625816e-05,
+      "step": 8710
+    },
+    {
+      "epoch": 0.89,
+      "gumbel_tau": 0.1281,
+      "image_masks_loss": 0.0086,
+      "llm_learning_rate": 1.4994981807370607e-07,
+      "loss": 0.7997,
+      "output_text_mask_loss": 0.0062,
+      "predictor_learning_rate": 1.4994981807370605e-05,
+      "step": 8711
+    },
+    {
+      "epoch": 0.89,
+      "gumbel_tau": 0.1281,
+      "image_masks_loss": 0.0028,
+      "llm_learning_rate": 1.496669493339145e-07,
+      "loss": 0.6333,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 1.496669493339145e-05,
+      "step": 8712
+    },
+    {
+      "epoch": 0.89,
+      "gumbel_tau": 0.128,
+      "image_masks_loss": 0.0075,
+      "llm_learning_rate": 1.493843394180333e-07,
+      "loss": 0.697,
+      "output_text_mask_loss": 0.001,
+      "predictor_learning_rate": 1.4938433941803331e-05,
+      "step": 8713
+    },
+    {
+      "epoch": 0.89,
+      "gumbel_tau": 0.128,
+      "image_masks_loss": 0.0121,
+      "llm_learning_rate": 1.491019883571851e-07,
+      "loss": 0.718,
+      "output_text_mask_loss": 0.0256,
+      "predictor_learning_rate": 1.4910198835718508e-05,
+      "step": 8714
+    },
+    {
+      "epoch": 0.89,
+      "gumbel_tau": 0.128,
+      "image_masks_loss": 0.0084,
+      "llm_learning_rate": 1.4881989618246217e-07,
+      "loss": 0.6742,
+      "output_text_mask_loss": 0.001,
+      "predictor_learning_rate": 1.4881989618246217e-05,
+      "step": 8715
+    },
+    {
+      "epoch": 0.89,
+      "gumbel_tau": 0.128,
+      "image_masks_loss": 0.0128,
+      "llm_learning_rate": 1.4853806292492918e-07,
+      "loss": 0.7287,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 1.4853806292492916e-05,
+      "step": 8716
+    },
+    {
+      "epoch": 0.89,
+      "gumbel_tau": 0.1279,
+      "image_masks_loss": 0.0125,
+      "llm_learning_rate": 1.4825648861562235e-07,
+      "loss": 0.8626,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 1.4825648861562236e-05,
+      "step": 8717
+    },
+    {
+      "epoch": 0.89,
+      "gumbel_tau": 0.1279,
+      "image_masks_loss": 0.0159,
+      "llm_learning_rate": 1.4797517328554944e-07,
+      "loss": 0.7269,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 1.4797517328554944e-05,
+      "step": 8718
+    },
+    {
+      "epoch": 0.89,
+      "gumbel_tau": 0.1279,
+      "image_masks_loss": 0.0062,
+      "llm_learning_rate": 1.4769411696568897e-07,
+      "loss": 0.7244,
+      "output_text_mask_loss": 0.0168,
+      "predictor_learning_rate": 1.4769411696568896e-05,
+      "step": 8719
+    },
+    {
+      "epoch": 0.89,
+      "gumbel_tau": 0.1278,
+      "image_masks_loss": 0.0123,
+      "llm_learning_rate": 1.4741331968699228e-07,
+      "loss": 0.6273,
+      "output_text_mask_loss": 0.0039,
+      "predictor_learning_rate": 1.4741331968699229e-05,
+      "step": 8720
+    },
+    {
+      "epoch": 0.89,
+      "gumbel_tau": 0.1278,
+      "image_masks_loss": 0.0106,
+      "llm_learning_rate": 1.4713278148038078e-07,
+      "loss": 0.7203,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 1.4713278148038078e-05,
+      "step": 8721
+    },
+    {
+      "epoch": 0.89,
+      "gumbel_tau": 0.1278,
+      "image_masks_loss": 0.0143,
+      "llm_learning_rate": 1.4685250237674836e-07,
+      "loss": 0.7757,
+      "output_text_mask_loss": 0.005,
+      "predictor_learning_rate": 1.4685250237674836e-05,
+      "step": 8722
+    },
+    {
+      "epoch": 0.89,
+      "gumbel_tau": 0.1277,
+      "image_masks_loss": 0.0108,
+      "llm_learning_rate": 1.4657248240696005e-07,
+      "loss": 0.739,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 1.4657248240696003e-05,
+      "step": 8723
+    },
+    {
+      "epoch": 0.89,
+      "gumbel_tau": 0.1277,
+      "image_masks_loss": 0.0136,
+      "llm_learning_rate": 1.462927216018517e-07,
+      "loss": 0.7695,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 1.462927216018517e-05,
+      "step": 8724
+    },
+    {
+      "epoch": 0.89,
+      "gumbel_tau": 0.1277,
+      "image_masks_loss": 0.0267,
+      "llm_learning_rate": 1.460132199922315e-07,
+      "loss": 0.7247,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 1.460132199922315e-05,
+      "step": 8725
+    },
+    {
+      "epoch": 0.89,
+      "gumbel_tau": 0.1277,
+      "image_masks_loss": 0.023,
+      "llm_learning_rate": 1.457339776088787e-07,
+      "loss": 0.6929,
+      "output_text_mask_loss": 0.0174,
+      "predictor_learning_rate": 1.4573397760887869e-05,
+      "step": 8726
+    },
+    {
+      "epoch": 0.89,
+      "gumbel_tau": 0.1276,
+      "image_masks_loss": 0.0064,
+      "llm_learning_rate": 1.4545499448254425e-07,
+      "loss": 0.8211,
+      "output_text_mask_loss": 0.0195,
+      "predictor_learning_rate": 1.4545499448254424e-05,
+      "step": 8727
+    },
+    {
+      "epoch": 0.89,
+      "gumbel_tau": 0.1276,
+      "image_masks_loss": 0.0204,
+      "llm_learning_rate": 1.451762706439508e-07,
+      "loss": 0.7512,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 1.4517627064395079e-05,
+      "step": 8728
+    },
+    {
+      "epoch": 0.89,
+      "gumbel_tau": 0.1276,
+      "image_masks_loss": 0.009,
+      "llm_learning_rate": 1.448978061237913e-07,
+      "loss": 0.7652,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 1.4489780612379127e-05,
+      "step": 8729
+    },
+    {
+      "epoch": 0.89,
+      "gumbel_tau": 0.1275,
+      "image_masks_loss": 0.0111,
+      "llm_learning_rate": 1.446196009527312e-07,
+      "loss": 0.7684,
+      "output_text_mask_loss": 0.0139,
+      "predictor_learning_rate": 1.446196009527312e-05,
+      "step": 8730
+    },
+    {
+      "epoch": 0.89,
+      "gumbel_tau": 0.1275,
+      "image_masks_loss": 0.0161,
+      "llm_learning_rate": 1.4434165516140687e-07,
+      "loss": 0.7412,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 1.4434165516140684e-05,
+      "step": 8731
+    },
+    {
+      "epoch": 0.89,
+      "gumbel_tau": 0.1275,
+      "image_masks_loss": 0.0075,
+      "llm_learning_rate": 1.4406396878042683e-07,
+      "loss": 0.7533,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 1.4406396878042682e-05,
+      "step": 8732
+    },
+    {
+      "epoch": 0.89,
+      "gumbel_tau": 0.1274,
+      "image_masks_loss": 0.0133,
+      "llm_learning_rate": 1.4378654184037e-07,
+      "loss": 0.7272,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 1.4378654184036999e-05,
+      "step": 8733
+    },
+    {
+      "epoch": 0.89,
+      "gumbel_tau": 0.1274,
+      "image_masks_loss": 0.0122,
+      "llm_learning_rate": 1.4350937437178748e-07,
+      "loss": 0.7562,
+      "output_text_mask_loss": 0.0038,
+      "predictor_learning_rate": 1.4350937437178745e-05,
+      "step": 8734
+    },
+    {
+      "epoch": 0.89,
+      "gumbel_tau": 0.1274,
+      "image_masks_loss": 0.0072,
+      "llm_learning_rate": 1.432324664052015e-07,
+      "loss": 0.7231,
+      "output_text_mask_loss": 0.0065,
+      "predictor_learning_rate": 1.432324664052015e-05,
+      "step": 8735
+    },
+    {
+      "epoch": 0.9,
+      "gumbel_tau": 0.1274,
+      "image_masks_loss": 0.0108,
+      "llm_learning_rate": 1.4295581797110575e-07,
+      "loss": 0.7295,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 1.4295581797110574e-05,
+      "step": 8736
+    },
+    {
+      "epoch": 0.9,
+      "gumbel_tau": 0.1273,
+      "image_masks_loss": 0.0033,
+      "llm_learning_rate": 1.4267942909996558e-07,
+      "loss": 0.6541,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 1.4267942909996556e-05,
+      "step": 8737
+    },
+    {
+      "epoch": 0.9,
+      "gumbel_tau": 0.1273,
+      "image_masks_loss": 0.013,
+      "llm_learning_rate": 1.4240329982221717e-07,
+      "loss": 0.7271,
+      "output_text_mask_loss": 0.0058,
+      "predictor_learning_rate": 1.4240329982221717e-05,
+      "step": 8738
+    },
+    {
+      "epoch": 0.9,
+      "gumbel_tau": 0.1273,
+      "image_masks_loss": 0.0152,
+      "llm_learning_rate": 1.4212743016826874e-07,
+      "loss": 0.7323,
+      "output_text_mask_loss": 0.0073,
+      "predictor_learning_rate": 1.4212743016826873e-05,
+      "step": 8739
+    },
+    {
+      "epoch": 0.9,
+      "gumbel_tau": 0.1272,
+      "image_masks_loss": 0.0218,
+      "llm_learning_rate": 1.4185182016849984e-07,
+      "loss": 0.7726,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 1.4185182016849984e-05,
+      "step": 8740
+    },
+    {
+      "epoch": 0.9,
+      "gumbel_tau": 0.1272,
+      "image_masks_loss": 0.006,
+      "llm_learning_rate": 1.4157646985326124e-07,
+      "loss": 0.764,
+      "output_text_mask_loss": 0.0077,
+      "predictor_learning_rate": 1.4157646985326123e-05,
+      "step": 8741
+    },
+    {
+      "epoch": 0.9,
+      "gumbel_tau": 0.1272,
+      "image_masks_loss": 0.0153,
+      "llm_learning_rate": 1.4130137925287473e-07,
+      "loss": 0.6991,
+      "output_text_mask_loss": 0.0355,
+      "predictor_learning_rate": 1.4130137925287474e-05,
+      "step": 8742
+    },
+    {
+      "epoch": 0.9,
+      "gumbel_tau": 0.1271,
+      "image_masks_loss": 0.0209,
+      "llm_learning_rate": 1.410265483976342e-07,
+      "loss": 0.7799,
+      "output_text_mask_loss": 0.015,
+      "predictor_learning_rate": 1.4102654839763418e-05,
+      "step": 8743
+    },
+    {
+      "epoch": 0.9,
+      "gumbel_tau": 0.1271,
+      "image_masks_loss": 0.0085,
+      "llm_learning_rate": 1.4075197731780455e-07,
+      "loss": 0.6548,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 1.4075197731780454e-05,
+      "step": 8744
+    },
+    {
+      "epoch": 0.9,
+      "gumbel_tau": 0.1271,
+      "image_masks_loss": 0.0047,
+      "llm_learning_rate": 1.4047766604362162e-07,
+      "loss": 0.6861,
+      "output_text_mask_loss": 0.0137,
+      "predictor_learning_rate": 1.4047766604362162e-05,
+      "step": 8745
+    },
+    {
+      "epoch": 0.9,
+      "gumbel_tau": 0.1271,
+      "image_masks_loss": 0.012,
+      "llm_learning_rate": 1.402036146052943e-07,
+      "loss": 0.7025,
+      "output_text_mask_loss": 0.0063,
+      "predictor_learning_rate": 1.4020361460529429e-05,
+      "step": 8746
+    },
+    {
+      "epoch": 0.9,
+      "gumbel_tau": 0.127,
+      "image_masks_loss": 0.0056,
+      "llm_learning_rate": 1.3992982303300123e-07,
+      "loss": 0.6575,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 1.3992982303300122e-05,
+      "step": 8747
+    },
+    {
+      "epoch": 0.9,
+      "gumbel_tau": 0.127,
+      "image_masks_loss": 0.0156,
+      "llm_learning_rate": 1.39656291356893e-07,
+      "loss": 0.697,
+      "output_text_mask_loss": 0.0389,
+      "predictor_learning_rate": 1.3965629135689301e-05,
+      "step": 8748
+    },
+    {
+      "epoch": 0.9,
+      "gumbel_tau": 0.127,
+      "image_masks_loss": 0.0075,
+      "llm_learning_rate": 1.393830196070911e-07,
+      "loss": 0.7585,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 1.393830196070911e-05,
+      "step": 8749
+    },
+    {
+      "epoch": 0.9,
+      "gumbel_tau": 0.1269,
+      "image_masks_loss": 0.0038,
+      "llm_learning_rate": 1.391100078136895e-07,
+      "loss": 0.7458,
+      "output_text_mask_loss": 0.0118,
+      "predictor_learning_rate": 1.3911000781368949e-05,
+      "step": 8750
+    },
+    {
+      "epoch": 0.9,
+      "gumbel_tau": 0.1269,
+      "image_masks_loss": 0.0035,
+      "llm_learning_rate": 1.388372560067522e-07,
+      "loss": 0.7449,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 1.3883725600675218e-05,
+      "step": 8751
+    },
+    {
+      "epoch": 0.9,
+      "gumbel_tau": 0.1269,
+      "image_masks_loss": 0.0058,
+      "llm_learning_rate": 1.385647642163157e-07,
+      "loss": 0.7873,
+      "output_text_mask_loss": 0.0029,
+      "predictor_learning_rate": 1.3856476421631569e-05,
+      "step": 8752
+    },
+    {
+      "epoch": 0.9,
+      "gumbel_tau": 0.1268,
+      "image_masks_loss": 0.0098,
+      "llm_learning_rate": 1.3829253247238715e-07,
+      "loss": 0.8259,
+      "output_text_mask_loss": 0.0208,
+      "predictor_learning_rate": 1.3829253247238715e-05,
+      "step": 8753
+    },
+    {
+      "epoch": 0.9,
+      "gumbel_tau": 0.1268,
+      "image_masks_loss": 0.0132,
+      "llm_learning_rate": 1.3802056080494563e-07,
+      "loss": 0.625,
+      "output_text_mask_loss": 0.0137,
+      "predictor_learning_rate": 1.3802056080494563e-05,
+      "step": 8754
+    },
+    {
+      "epoch": 0.9,
+      "gumbel_tau": 0.1268,
+      "image_masks_loss": 0.0109,
+      "llm_learning_rate": 1.3774884924394078e-07,
+      "loss": 0.7088,
+      "output_text_mask_loss": 0.0292,
+      "predictor_learning_rate": 1.3774884924394076e-05,
+      "step": 8755
+    },
+    {
+      "epoch": 0.9,
+      "gumbel_tau": 0.1268,
+      "image_masks_loss": 0.0145,
+      "llm_learning_rate": 1.3747739781929475e-07,
+      "loss": 0.7452,
+      "output_text_mask_loss": 0.0214,
+      "predictor_learning_rate": 1.3747739781929475e-05,
+      "step": 8756
+    },
+    {
+      "epoch": 0.9,
+      "gumbel_tau": 0.1267,
+      "image_masks_loss": 0.018,
+      "llm_learning_rate": 1.372062065608998e-07,
+      "loss": 0.6263,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 1.3720620656089978e-05,
+      "step": 8757
+    },
+    {
+      "epoch": 0.9,
+      "gumbel_tau": 0.1267,
+      "image_masks_loss": 0.0092,
+      "llm_learning_rate": 1.3693527549861973e-07,
+      "loss": 0.6954,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 1.3693527549861972e-05,
+      "step": 8758
+    },
+    {
+      "epoch": 0.9,
+      "gumbel_tau": 0.1267,
+      "image_masks_loss": 0.0061,
+      "llm_learning_rate": 1.366646046622913e-07,
+      "loss": 0.7323,
+      "output_text_mask_loss": 0.0023,
+      "predictor_learning_rate": 1.366646046622913e-05,
+      "step": 8759
+    },
+    {
+      "epoch": 0.9,
+      "gumbel_tau": 0.1266,
+      "image_masks_loss": 0.0054,
+      "llm_learning_rate": 1.3639419408172037e-07,
+      "loss": 0.794,
+      "output_text_mask_loss": 0.0012,
+      "predictor_learning_rate": 1.3639419408172039e-05,
+      "step": 8760
+    },
+    {
+      "epoch": 0.9,
+      "gumbel_tau": 0.1266,
+      "image_masks_loss": 0.0082,
+      "llm_learning_rate": 1.3612404378668593e-07,
+      "loss": 0.709,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 1.3612404378668591e-05,
+      "step": 8761
+    },
+    {
+      "epoch": 0.9,
+      "gumbel_tau": 0.1266,
+      "image_masks_loss": 0.0146,
+      "llm_learning_rate": 1.358541538069369e-07,
+      "loss": 0.6829,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 1.358541538069369e-05,
+      "step": 8762
+    },
+    {
+      "epoch": 0.9,
+      "gumbel_tau": 0.1265,
+      "image_masks_loss": 0.0072,
+      "llm_learning_rate": 1.3558452417219425e-07,
+      "loss": 0.6852,
+      "output_text_mask_loss": 0.0248,
+      "predictor_learning_rate": 1.3558452417219425e-05,
+      "step": 8763
+    },
+    {
+      "epoch": 0.9,
+      "gumbel_tau": 0.1265,
+      "image_masks_loss": 0.0146,
+      "llm_learning_rate": 1.3531515491215064e-07,
+      "loss": 0.8323,
+      "output_text_mask_loss": 0.0056,
+      "predictor_learning_rate": 1.3531515491215064e-05,
+      "step": 8764
+    },
+    {
+      "epoch": 0.9,
+      "gumbel_tau": 0.1265,
+      "image_masks_loss": 0.0061,
+      "llm_learning_rate": 1.350460460564687e-07,
+      "loss": 0.7113,
+      "output_text_mask_loss": 0.0015,
+      "predictor_learning_rate": 1.3504604605646869e-05,
+      "step": 8765
+    },
+    {
+      "epoch": 0.9,
+      "gumbel_tau": 0.1265,
+      "image_masks_loss": 0.0071,
+      "llm_learning_rate": 1.3477719763478448e-07,
+      "loss": 0.7826,
+      "output_text_mask_loss": 0.0319,
+      "predictor_learning_rate": 1.3477719763478446e-05,
+      "step": 8766
+    },
+    {
+      "epoch": 0.9,
+      "gumbel_tau": 0.1264,
+      "image_masks_loss": 0.0058,
+      "llm_learning_rate": 1.3450860967670342e-07,
+      "loss": 0.7245,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 1.3450860967670342e-05,
+      "step": 8767
+    },
+    {
+      "epoch": 0.9,
+      "gumbel_tau": 0.1264,
+      "image_masks_loss": 0.0174,
+      "llm_learning_rate": 1.3424028221180303e-07,
+      "loss": 0.7493,
+      "output_text_mask_loss": 0.0042,
+      "predictor_learning_rate": 1.3424028221180301e-05,
+      "step": 8768
+    },
+    {
+      "epoch": 0.9,
+      "gumbel_tau": 0.1264,
+      "image_masks_loss": 0.0271,
+      "llm_learning_rate": 1.3397221526963294e-07,
+      "loss": 0.6928,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 1.3397221526963293e-05,
+      "step": 8769
+    },
+    {
+      "epoch": 0.9,
+      "gumbel_tau": 0.1263,
+      "image_masks_loss": 0.0186,
+      "llm_learning_rate": 1.337044088797118e-07,
+      "loss": 0.6984,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 1.337044088797118e-05,
+      "step": 8770
+    },
+    {
+      "epoch": 0.9,
+      "gumbel_tau": 0.1263,
+      "image_masks_loss": 0.0213,
+      "llm_learning_rate": 1.3343686307153213e-07,
+      "loss": 0.7227,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 1.3343686307153214e-05,
+      "step": 8771
+    },
+    {
+      "epoch": 0.9,
+      "gumbel_tau": 0.1263,
+      "image_masks_loss": 0.0103,
+      "llm_learning_rate": 1.3316957787455648e-07,
+      "loss": 0.7726,
+      "output_text_mask_loss": 0.0381,
+      "predictor_learning_rate": 1.3316957787455647e-05,
+      "step": 8772
+    },
+    {
+      "epoch": 0.9,
+      "gumbel_tau": 0.1262,
+      "image_masks_loss": 0.0041,
+      "llm_learning_rate": 1.3290255331821906e-07,
+      "loss": 0.7847,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 1.3290255331821905e-05,
+      "step": 8773
+    },
+    {
+      "epoch": 0.9,
+      "gumbel_tau": 0.1262,
+      "image_masks_loss": 0.0154,
+      "llm_learning_rate": 1.3263578943192468e-07,
+      "loss": 0.6793,
+      "output_text_mask_loss": 0.0391,
+      "predictor_learning_rate": 1.3263578943192467e-05,
+      "step": 8774
+    },
+    {
+      "epoch": 0.9,
+      "gumbel_tau": 0.1262,
+      "image_masks_loss": 0.0039,
+      "llm_learning_rate": 1.323692862450504e-07,
+      "loss": 0.7829,
+      "output_text_mask_loss": 0.0038,
+      "predictor_learning_rate": 1.323692862450504e-05,
+      "step": 8775
+    },
+    {
+      "epoch": 0.9,
+      "gumbel_tau": 0.1262,
+      "image_masks_loss": 0.0077,
+      "llm_learning_rate": 1.3210304378694387e-07,
+      "loss": 0.7368,
+      "output_text_mask_loss": 0.0368,
+      "predictor_learning_rate": 1.3210304378694387e-05,
+      "step": 8776
+    },
+    {
+      "epoch": 0.9,
+      "gumbel_tau": 0.1261,
+      "image_masks_loss": 0.0229,
+      "llm_learning_rate": 1.3183706208692414e-07,
+      "loss": 0.608,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 1.3183706208692415e-05,
+      "step": 8777
+    },
+    {
+      "epoch": 0.9,
+      "gumbel_tau": 0.1261,
+      "image_masks_loss": 0.0117,
+      "llm_learning_rate": 1.315713411742825e-07,
+      "loss": 0.771,
+      "output_text_mask_loss": 0.0039,
+      "predictor_learning_rate": 1.315713411742825e-05,
+      "step": 8778
+    },
+    {
+      "epoch": 0.9,
+      "gumbel_tau": 0.1261,
+      "image_masks_loss": 0.0151,
+      "llm_learning_rate": 1.3130588107827997e-07,
+      "loss": 0.7973,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 1.3130588107827995e-05,
+      "step": 8779
+    },
+    {
+      "epoch": 0.9,
+      "gumbel_tau": 0.126,
+      "image_masks_loss": 0.0235,
+      "llm_learning_rate": 1.3104068182815012e-07,
+      "loss": 0.7273,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 1.310406818281501e-05,
+      "step": 8780
+    },
+    {
+      "epoch": 0.9,
+      "gumbel_tau": 0.126,
+      "image_masks_loss": 0.0161,
+      "llm_learning_rate": 1.307757434530968e-07,
+      "loss": 0.706,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 1.3077574345309678e-05,
+      "step": 8781
+    },
+    {
+      "epoch": 0.9,
+      "gumbel_tau": 0.126,
+      "image_masks_loss": 0.0088,
+      "llm_learning_rate": 1.3051106598229614e-07,
+      "loss": 0.7958,
+      "output_text_mask_loss": 0.0072,
+      "predictor_learning_rate": 1.3051106598229612e-05,
+      "step": 8782
+    },
+    {
+      "epoch": 0.9,
+      "gumbel_tau": 0.126,
+      "image_masks_loss": 0.0087,
+      "llm_learning_rate": 1.302466494448945e-07,
+      "loss": 0.6995,
+      "output_text_mask_loss": 0.0141,
+      "predictor_learning_rate": 1.302466494448945e-05,
+      "step": 8783
+    },
+    {
+      "epoch": 0.9,
+      "gumbel_tau": 0.1259,
+      "image_masks_loss": 0.0126,
+      "llm_learning_rate": 1.2998249387001006e-07,
+      "loss": 0.684,
+      "output_text_mask_loss": 0.0132,
+      "predictor_learning_rate": 1.2998249387001005e-05,
+      "step": 8784
+    },
+    {
+      "epoch": 0.9,
+      "gumbel_tau": 0.1259,
+      "image_masks_loss": 0.0071,
+      "llm_learning_rate": 1.2971859928673254e-07,
+      "loss": 0.7203,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 1.2971859928673252e-05,
+      "step": 8785
+    },
+    {
+      "epoch": 0.9,
+      "gumbel_tau": 0.1259,
+      "image_masks_loss": 0.028,
+      "llm_learning_rate": 1.2945496572412259e-07,
+      "loss": 0.7403,
+      "output_text_mask_loss": 0.0124,
+      "predictor_learning_rate": 1.2945496572412258e-05,
+      "step": 8786
+    },
+    {
+      "epoch": 0.9,
+      "gumbel_tau": 0.1258,
+      "image_masks_loss": 0.0019,
+      "llm_learning_rate": 1.29191593211212e-07,
+      "loss": 0.7358,
+      "output_text_mask_loss": 0.0019,
+      "predictor_learning_rate": 1.29191593211212e-05,
+      "step": 8787
+    },
+    {
+      "epoch": 0.9,
+      "gumbel_tau": 0.1258,
+      "image_masks_loss": 0.0053,
+      "llm_learning_rate": 1.289284817770045e-07,
+      "loss": 0.6563,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 1.289284817770045e-05,
+      "step": 8788
+    },
+    {
+      "epoch": 0.9,
+      "gumbel_tau": 0.1258,
+      "image_masks_loss": 0.0082,
+      "llm_learning_rate": 1.2866563145047333e-07,
+      "loss": 0.7862,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 1.2866563145047333e-05,
+      "step": 8789
+    },
+    {
+      "epoch": 0.9,
+      "gumbel_tau": 0.1257,
+      "image_masks_loss": 0.0075,
+      "llm_learning_rate": 1.284030422605645e-07,
+      "loss": 0.7871,
+      "output_text_mask_loss": 0.0246,
+      "predictor_learning_rate": 1.2840304226056448e-05,
+      "step": 8790
+    },
+    {
+      "epoch": 0.9,
+      "gumbel_tau": 0.1257,
+      "image_masks_loss": 0.0087,
+      "llm_learning_rate": 1.281407142361954e-07,
+      "loss": 0.7705,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 1.281407142361954e-05,
+      "step": 8791
+    },
+    {
+      "epoch": 0.9,
+      "gumbel_tau": 0.1257,
+      "image_masks_loss": 0.0076,
+      "llm_learning_rate": 1.2787864740625404e-07,
+      "loss": 0.6448,
+      "output_text_mask_loss": 0.0259,
+      "predictor_learning_rate": 1.2787864740625404e-05,
+      "step": 8792
+    },
+    {
+      "epoch": 0.9,
+      "gumbel_tau": 0.1257,
+      "image_masks_loss": 0.0101,
+      "llm_learning_rate": 1.2761684179959987e-07,
+      "loss": 0.6793,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 1.2761684179959988e-05,
+      "step": 8793
+    },
+    {
+      "epoch": 0.9,
+      "gumbel_tau": 0.1256,
+      "image_masks_loss": 0.0101,
+      "llm_learning_rate": 1.2735529744506315e-07,
+      "loss": 0.7669,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 1.2735529744506314e-05,
+      "step": 8794
+    },
+    {
+      "epoch": 0.9,
+      "gumbel_tau": 0.1256,
+      "image_masks_loss": 0.0359,
+      "llm_learning_rate": 1.270940143714461e-07,
+      "loss": 0.7493,
+      "output_text_mask_loss": 0.0235,
+      "predictor_learning_rate": 1.2709401437144607e-05,
+      "step": 8795
+    },
+    {
+      "epoch": 0.9,
+      "gumbel_tau": 0.1256,
+      "image_masks_loss": 0.0143,
+      "llm_learning_rate": 1.2683299260752152e-07,
+      "loss": 0.8104,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 1.2683299260752152e-05,
+      "step": 8796
+    },
+    {
+      "epoch": 0.9,
+      "gumbel_tau": 0.1255,
+      "image_masks_loss": 0.0123,
+      "llm_learning_rate": 1.2657223218203341e-07,
+      "loss": 0.7517,
+      "output_text_mask_loss": 0.0107,
+      "predictor_learning_rate": 1.265722321820334e-05,
+      "step": 8797
+    },
+    {
+      "epoch": 0.9,
+      "gumbel_tau": 0.1255,
+      "image_masks_loss": 0.0029,
+      "llm_learning_rate": 1.2631173312369821e-07,
+      "loss": 0.6626,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 1.263117331236982e-05,
+      "step": 8798
+    },
+    {
+      "epoch": 0.9,
+      "gumbel_tau": 0.1255,
+      "image_masks_loss": 0.0109,
+      "llm_learning_rate": 1.2605149546120189e-07,
+      "loss": 0.8359,
+      "output_text_mask_loss": 0.0231,
+      "predictor_learning_rate": 1.2605149546120188e-05,
+      "step": 8799
+    },
+    {
+      "epoch": 0.9,
+      "gumbel_tau": 0.1254,
+      "image_masks_loss": 0.0195,
+      "llm_learning_rate": 1.2579151922320286e-07,
+      "loss": 0.7824,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 1.2579151922320287e-05,
+      "step": 8800
+    },
+    {
+      "epoch": 0.9,
+      "gumbel_tau": 0.1254,
+      "image_masks_loss": 0.0142,
+      "llm_learning_rate": 1.2553180443832996e-07,
+      "loss": 0.6469,
+      "output_text_mask_loss": 0.0034,
+      "predictor_learning_rate": 1.2553180443832995e-05,
+      "step": 8801
+    },
+    {
+      "epoch": 0.9,
+      "gumbel_tau": 0.1254,
+      "image_masks_loss": 0.0115,
+      "llm_learning_rate": 1.2527235113518388e-07,
+      "loss": 0.7615,
+      "output_text_mask_loss": 0.0076,
+      "predictor_learning_rate": 1.2527235113518387e-05,
+      "step": 8802
+    },
+    {
+      "epoch": 0.9,
+      "gumbel_tau": 0.1254,
+      "image_masks_loss": 0.0059,
+      "llm_learning_rate": 1.2501315934233565e-07,
+      "loss": 0.7876,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 1.2501315934233564e-05,
+      "step": 8803
+    },
+    {
+      "epoch": 0.9,
+      "gumbel_tau": 0.1253,
+      "image_masks_loss": 0.0103,
+      "llm_learning_rate": 1.247542290883283e-07,
+      "loss": 0.8305,
+      "output_text_mask_loss": 0.0397,
+      "predictor_learning_rate": 1.2475422908832828e-05,
+      "step": 8804
+    },
+    {
+      "epoch": 0.9,
+      "gumbel_tau": 0.1253,
+      "image_masks_loss": 0.0051,
+      "llm_learning_rate": 1.244955604016762e-07,
+      "loss": 0.725,
+      "output_text_mask_loss": 0.013,
+      "predictor_learning_rate": 1.244955604016762e-05,
+      "step": 8805
+    },
+    {
+      "epoch": 0.9,
+      "gumbel_tau": 0.1253,
+      "image_masks_loss": 0.0058,
+      "llm_learning_rate": 1.2423715331086467e-07,
+      "loss": 0.6641,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 1.2423715331086465e-05,
+      "step": 8806
+    },
+    {
+      "epoch": 0.9,
+      "gumbel_tau": 0.1252,
+      "image_masks_loss": 0.0063,
+      "llm_learning_rate": 1.2397900784434896e-07,
+      "loss": 0.6688,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 1.2397900784434895e-05,
+      "step": 8807
+    },
+    {
+      "epoch": 0.9,
+      "gumbel_tau": 0.1252,
+      "image_masks_loss": 0.0088,
+      "llm_learning_rate": 1.2372112403055748e-07,
+      "loss": 0.7677,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 1.2372112403055747e-05,
+      "step": 8808
+    },
+    {
+      "epoch": 0.9,
+      "gumbel_tau": 0.1252,
+      "image_masks_loss": 0.0122,
+      "llm_learning_rate": 1.2346350189788858e-07,
+      "loss": 0.7361,
+      "output_text_mask_loss": 0.0113,
+      "predictor_learning_rate": 1.2346350189788858e-05,
+      "step": 8809
+    },
+    {
+      "epoch": 0.9,
+      "gumbel_tau": 0.1252,
+      "image_masks_loss": 0.0066,
+      "llm_learning_rate": 1.2320614147471266e-07,
+      "loss": 0.7454,
+      "output_text_mask_loss": 0.0056,
+      "predictor_learning_rate": 1.2320614147471265e-05,
+      "step": 8810
+    },
+    {
+      "epoch": 0.9,
+      "gumbel_tau": 0.1251,
+      "image_masks_loss": 0.0044,
+      "llm_learning_rate": 1.2294904278937038e-07,
+      "loss": 0.6886,
+      "output_text_mask_loss": 0.005,
+      "predictor_learning_rate": 1.2294904278937036e-05,
+      "step": 8811
+    },
+    {
+      "epoch": 0.9,
+      "gumbel_tau": 0.1251,
+      "image_masks_loss": 0.0117,
+      "llm_learning_rate": 1.2269220587017434e-07,
+      "loss": 0.8163,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 1.2269220587017432e-05,
+      "step": 8812
+    },
+    {
+      "epoch": 0.9,
+      "gumbel_tau": 0.1251,
+      "image_masks_loss": 0.0119,
+      "llm_learning_rate": 1.2243563074540777e-07,
+      "loss": 0.7488,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 1.2243563074540776e-05,
+      "step": 8813
+    },
+    {
+      "epoch": 0.9,
+      "gumbel_tau": 0.125,
+      "image_masks_loss": 0.0048,
+      "llm_learning_rate": 1.2217931744332528e-07,
+      "loss": 0.7743,
+      "output_text_mask_loss": 0.0187,
+      "predictor_learning_rate": 1.2217931744332529e-05,
+      "step": 8814
+    },
+    {
+      "epoch": 0.9,
+      "gumbel_tau": 0.125,
+      "image_masks_loss": 0.0133,
+      "llm_learning_rate": 1.2192326599215265e-07,
+      "loss": 0.7145,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 1.2192326599215264e-05,
+      "step": 8815
+    },
+    {
+      "epoch": 0.9,
+      "gumbel_tau": 0.125,
+      "image_masks_loss": 0.0117,
+      "llm_learning_rate": 1.2166747642008674e-07,
+      "loss": 0.7092,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 1.2166747642008675e-05,
+      "step": 8816
+    },
+    {
+      "epoch": 0.9,
+      "gumbel_tau": 0.1249,
+      "image_masks_loss": 0.0098,
+      "llm_learning_rate": 1.2141194875529615e-07,
+      "loss": 0.8035,
+      "output_text_mask_loss": 0.0307,
+      "predictor_learning_rate": 1.2141194875529616e-05,
+      "step": 8817
+    },
+    {
+      "epoch": 0.9,
+      "gumbel_tau": 0.1249,
+      "image_masks_loss": 0.0224,
+      "llm_learning_rate": 1.2115668302592005e-07,
+      "loss": 0.6984,
+      "output_text_mask_loss": 0.0082,
+      "predictor_learning_rate": 1.2115668302592003e-05,
+      "step": 8818
+    },
+    {
+      "epoch": 0.9,
+      "gumbel_tau": 0.1249,
+      "image_masks_loss": 0.0134,
+      "llm_learning_rate": 1.209016792600684e-07,
+      "loss": 0.8197,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 1.209016792600684e-05,
+      "step": 8819
+    },
+    {
+      "epoch": 0.9,
+      "gumbel_tau": 0.1249,
+      "image_masks_loss": 0.0048,
+      "llm_learning_rate": 1.2064693748582323e-07,
+      "loss": 0.7434,
+      "output_text_mask_loss": 0.01,
+      "predictor_learning_rate": 1.2064693748582322e-05,
+      "step": 8820
+    },
+    {
+      "epoch": 0.9,
+      "gumbel_tau": 0.1248,
+      "image_masks_loss": 0.0173,
+      "llm_learning_rate": 1.2039245773123709e-07,
+      "loss": 0.7838,
+      "output_text_mask_loss": 0.0172,
+      "predictor_learning_rate": 1.2039245773123708e-05,
+      "step": 8821
+    },
+    {
+      "epoch": 0.9,
+      "gumbel_tau": 0.1248,
+      "image_masks_loss": 0.0159,
+      "llm_learning_rate": 1.2013824002433394e-07,
+      "loss": 0.8008,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 1.2013824002433393e-05,
+      "step": 8822
+    },
+    {
+      "epoch": 0.9,
+      "gumbel_tau": 0.1248,
+      "image_masks_loss": 0.0066,
+      "llm_learning_rate": 1.1988428439310835e-07,
+      "loss": 0.6374,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 1.1988428439310833e-05,
+      "step": 8823
+    },
+    {
+      "epoch": 0.9,
+      "gumbel_tau": 0.1247,
+      "image_masks_loss": 0.0155,
+      "llm_learning_rate": 1.1963059086552708e-07,
+      "loss": 0.7304,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 1.1963059086552708e-05,
+      "step": 8824
+    },
+    {
+      "epoch": 0.9,
+      "gumbel_tau": 0.1247,
+      "image_masks_loss": 0.0151,
+      "llm_learning_rate": 1.1937715946952756e-07,
+      "loss": 0.7899,
+      "output_text_mask_loss": 0.0162,
+      "predictor_learning_rate": 1.1937715946952754e-05,
+      "step": 8825
+    },
+    {
+      "epoch": 0.9,
+      "gumbel_tau": 0.1247,
+      "image_masks_loss": 0.0141,
+      "llm_learning_rate": 1.1912399023301768e-07,
+      "loss": 0.5995,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 1.1912399023301768e-05,
+      "step": 8826
+    },
+    {
+      "epoch": 0.9,
+      "gumbel_tau": 0.1247,
+      "image_masks_loss": 0.0275,
+      "llm_learning_rate": 1.1887108318387713e-07,
+      "loss": 0.7276,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 1.1887108318387713e-05,
+      "step": 8827
+    },
+    {
+      "epoch": 0.9,
+      "gumbel_tau": 0.1246,
+      "image_masks_loss": 0.004,
+      "llm_learning_rate": 1.1861843834995668e-07,
+      "loss": 0.6798,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 1.1861843834995667e-05,
+      "step": 8828
+    },
+    {
+      "epoch": 0.9,
+      "gumbel_tau": 0.1246,
+      "image_masks_loss": 0.0114,
+      "llm_learning_rate": 1.1836605575907767e-07,
+      "loss": 0.6673,
+      "output_text_mask_loss": 0.037,
+      "predictor_learning_rate": 1.1836605575907767e-05,
+      "step": 8829
+    },
+    {
+      "epoch": 0.9,
+      "gumbel_tau": 0.1246,
+      "image_masks_loss": 0.0094,
+      "llm_learning_rate": 1.18113935439034e-07,
+      "loss": 0.7793,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 1.1811393543903399e-05,
+      "step": 8830
+    },
+    {
+      "epoch": 0.9,
+      "gumbel_tau": 0.1245,
+      "image_masks_loss": 0.014,
+      "llm_learning_rate": 1.1786207741758898e-07,
+      "loss": 0.7336,
+      "output_text_mask_loss": 0.0128,
+      "predictor_learning_rate": 1.1786207741758898e-05,
+      "step": 8831
+    },
+    {
+      "epoch": 0.9,
+      "gumbel_tau": 0.1245,
+      "image_masks_loss": 0.0104,
+      "llm_learning_rate": 1.1761048172247824e-07,
+      "loss": 0.6731,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 1.1761048172247824e-05,
+      "step": 8832
+    },
+    {
+      "epoch": 0.91,
+      "gumbel_tau": 0.1245,
+      "image_masks_loss": 0.0062,
+      "llm_learning_rate": 1.1735914838140767e-07,
+      "loss": 0.7194,
+      "output_text_mask_loss": 0.0275,
+      "predictor_learning_rate": 1.1735914838140765e-05,
+      "step": 8833
+    },
+    {
+      "epoch": 0.91,
+      "gumbel_tau": 0.1244,
+      "image_masks_loss": 0.0033,
+      "llm_learning_rate": 1.1710807742205482e-07,
+      "loss": 0.7293,
+      "output_text_mask_loss": 0.0091,
+      "predictor_learning_rate": 1.1710807742205482e-05,
+      "step": 8834
+    },
+    {
+      "epoch": 0.91,
+      "gumbel_tau": 0.1244,
+      "image_masks_loss": 0.0145,
+      "llm_learning_rate": 1.1685726887206816e-07,
+      "loss": 0.7541,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 1.1685726887206815e-05,
+      "step": 8835
+    },
+    {
+      "epoch": 0.91,
+      "gumbel_tau": 0.1244,
+      "image_masks_loss": 0.0121,
+      "llm_learning_rate": 1.1660672275906698e-07,
+      "loss": 0.7322,
+      "output_text_mask_loss": 0.0141,
+      "predictor_learning_rate": 1.1660672275906697e-05,
+      "step": 8836
+    },
+    {
+      "epoch": 0.91,
+      "gumbel_tau": 0.1244,
+      "image_masks_loss": 0.0155,
+      "llm_learning_rate": 1.1635643911064254e-07,
+      "loss": 0.7554,
+      "output_text_mask_loss": 0.0132,
+      "predictor_learning_rate": 1.1635643911064252e-05,
+      "step": 8837
+    },
+    {
+      "epoch": 0.91,
+      "gumbel_tau": 0.1243,
+      "image_masks_loss": 0.0136,
+      "llm_learning_rate": 1.161064179543564e-07,
+      "loss": 0.8314,
+      "output_text_mask_loss": 0.0018,
+      "predictor_learning_rate": 1.161064179543564e-05,
+      "step": 8838
+    },
+    {
+      "epoch": 0.91,
+      "gumbel_tau": 0.1243,
+      "image_masks_loss": 0.0102,
+      "llm_learning_rate": 1.1585665931774154e-07,
+      "loss": 0.754,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 1.1585665931774153e-05,
+      "step": 8839
+    },
+    {
+      "epoch": 0.91,
+      "gumbel_tau": 0.1243,
+      "image_masks_loss": 0.0128,
+      "llm_learning_rate": 1.1560716322830179e-07,
+      "loss": 0.7534,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 1.1560716322830179e-05,
+      "step": 8840
+    },
+    {
+      "epoch": 0.91,
+      "gumbel_tau": 0.1242,
+      "image_masks_loss": 0.0046,
+      "llm_learning_rate": 1.1535792971351239e-07,
+      "loss": 0.7852,
+      "output_text_mask_loss": 0.0038,
+      "predictor_learning_rate": 1.1535792971351238e-05,
+      "step": 8841
+    },
+    {
+      "epoch": 0.91,
+      "gumbel_tau": 0.1242,
+      "image_masks_loss": 0.005,
+      "llm_learning_rate": 1.1510895880081918e-07,
+      "loss": 0.8167,
+      "output_text_mask_loss": 0.0101,
+      "predictor_learning_rate": 1.1510895880081918e-05,
+      "step": 8842
+    },
+    {
+      "epoch": 0.91,
+      "gumbel_tau": 0.1242,
+      "image_masks_loss": 0.0199,
+      "llm_learning_rate": 1.1486025051763939e-07,
+      "loss": 0.6352,
+      "output_text_mask_loss": 0.0231,
+      "predictor_learning_rate": 1.1486025051763938e-05,
+      "step": 8843
+    },
+    {
+      "epoch": 0.91,
+      "gumbel_tau": 0.1242,
+      "image_masks_loss": 0.0182,
+      "llm_learning_rate": 1.1461180489136192e-07,
+      "loss": 0.7702,
+      "output_text_mask_loss": 0.0196,
+      "predictor_learning_rate": 1.1461180489136192e-05,
+      "step": 8844
+    },
+    {
+      "epoch": 0.91,
+      "gumbel_tau": 0.1241,
+      "image_masks_loss": 0.0048,
+      "llm_learning_rate": 1.1436362194934547e-07,
+      "loss": 0.7511,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 1.1436362194934547e-05,
+      "step": 8845
+    },
+    {
+      "epoch": 0.91,
+      "gumbel_tau": 0.1241,
+      "image_masks_loss": 0.0042,
+      "llm_learning_rate": 1.1411570171892067e-07,
+      "loss": 0.7522,
+      "output_text_mask_loss": 0.0591,
+      "predictor_learning_rate": 1.1411570171892066e-05,
+      "step": 8846
+    },
+    {
+      "epoch": 0.91,
+      "gumbel_tau": 0.1241,
+      "image_masks_loss": 0.012,
+      "llm_learning_rate": 1.1386804422738929e-07,
+      "loss": 0.6745,
+      "output_text_mask_loss": 0.0021,
+      "predictor_learning_rate": 1.1386804422738928e-05,
+      "step": 8847
+    },
+    {
+      "epoch": 0.91,
+      "gumbel_tau": 0.124,
+      "image_masks_loss": 0.0091,
+      "llm_learning_rate": 1.1362064950202395e-07,
+      "loss": 0.8084,
+      "output_text_mask_loss": 0.0277,
+      "predictor_learning_rate": 1.1362064950202394e-05,
+      "step": 8848
+    },
+    {
+      "epoch": 0.91,
+      "gumbel_tau": 0.124,
+      "image_masks_loss": 0.0236,
+      "llm_learning_rate": 1.1337351757006759e-07,
+      "loss": 0.6657,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 1.1337351757006759e-05,
+      "step": 8849
+    },
+    {
+      "epoch": 0.91,
+      "gumbel_tau": 0.124,
+      "image_masks_loss": 0.0026,
+      "llm_learning_rate": 1.1312664845873566e-07,
+      "loss": 0.6803,
+      "output_text_mask_loss": 0.0221,
+      "predictor_learning_rate": 1.1312664845873567e-05,
+      "step": 8850
+    },
+    {
+      "epoch": 0.91,
+      "gumbel_tau": 0.1239,
+      "image_masks_loss": 0.0123,
+      "llm_learning_rate": 1.128800421952142e-07,
+      "loss": 0.792,
+      "output_text_mask_loss": 0.0341,
+      "predictor_learning_rate": 1.128800421952142e-05,
+      "step": 8851
+    },
+    {
+      "epoch": 0.91,
+      "gumbel_tau": 0.1239,
+      "image_masks_loss": 0.0106,
+      "llm_learning_rate": 1.1263369880665926e-07,
+      "loss": 0.6635,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 1.1263369880665925e-05,
+      "step": 8852
+    },
+    {
+      "epoch": 0.91,
+      "gumbel_tau": 0.1239,
+      "image_masks_loss": 0.0231,
+      "llm_learning_rate": 1.1238761832019912e-07,
+      "loss": 0.8101,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 1.1238761832019912e-05,
+      "step": 8853
+    },
+    {
+      "epoch": 0.91,
+      "gumbel_tau": 0.1239,
+      "image_masks_loss": 0.0125,
+      "llm_learning_rate": 1.1214180076293269e-07,
+      "loss": 0.7402,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 1.1214180076293268e-05,
+      "step": 8854
+    },
+    {
+      "epoch": 0.91,
+      "gumbel_tau": 0.1238,
+      "image_masks_loss": 0.0064,
+      "llm_learning_rate": 1.1189624616192995e-07,
+      "loss": 0.7664,
+      "output_text_mask_loss": 0.0076,
+      "predictor_learning_rate": 1.1189624616192996e-05,
+      "step": 8855
+    },
+    {
+      "epoch": 0.91,
+      "gumbel_tau": 0.1238,
+      "image_masks_loss": 0.0124,
+      "llm_learning_rate": 1.1165095454423152e-07,
+      "loss": 0.6622,
+      "output_text_mask_loss": 0.0018,
+      "predictor_learning_rate": 1.1165095454423151e-05,
+      "step": 8856
+    },
+    {
+      "epoch": 0.91,
+      "gumbel_tau": 0.1238,
+      "image_masks_loss": 0.0144,
+      "llm_learning_rate": 1.1140592593684996e-07,
+      "loss": 0.6798,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 1.1140592593684996e-05,
+      "step": 8857
+    },
+    {
+      "epoch": 0.91,
+      "gumbel_tau": 0.1237,
+      "image_masks_loss": 0.004,
+      "llm_learning_rate": 1.1116116036676839e-07,
+      "loss": 0.6305,
+      "output_text_mask_loss": 0.0465,
+      "predictor_learning_rate": 1.1116116036676837e-05,
+      "step": 8858
+    },
+    {
+      "epoch": 0.91,
+      "gumbel_tau": 0.1237,
+      "image_masks_loss": 0.0109,
+      "llm_learning_rate": 1.109166578609408e-07,
+      "loss": 0.7412,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 1.109166578609408e-05,
+      "step": 8859
+    },
+    {
+      "epoch": 0.91,
+      "gumbel_tau": 0.1237,
+      "image_masks_loss": 0.0194,
+      "llm_learning_rate": 1.1067241844629234e-07,
+      "loss": 0.8228,
+      "output_text_mask_loss": 0.0134,
+      "predictor_learning_rate": 1.1067241844629234e-05,
+      "step": 8860
+    },
+    {
+      "epoch": 0.91,
+      "gumbel_tau": 0.1237,
+      "image_masks_loss": 0.0045,
+      "llm_learning_rate": 1.1042844214971898e-07,
+      "loss": 0.7109,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 1.1042844214971898e-05,
+      "step": 8861
+    },
+    {
+      "epoch": 0.91,
+      "gumbel_tau": 0.1236,
+      "image_masks_loss": 0.0154,
+      "llm_learning_rate": 1.101847289980884e-07,
+      "loss": 0.6877,
+      "output_text_mask_loss": 0.0012,
+      "predictor_learning_rate": 1.101847289980884e-05,
+      "step": 8862
+    },
+    {
+      "epoch": 0.91,
+      "gumbel_tau": 0.1236,
+      "image_masks_loss": 0.0168,
+      "llm_learning_rate": 1.0994127901823831e-07,
+      "loss": 0.7084,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 1.099412790182383e-05,
+      "step": 8863
+    },
+    {
+      "epoch": 0.91,
+      "gumbel_tau": 0.1236,
+      "image_masks_loss": 0.0054,
+      "llm_learning_rate": 1.0969809223697836e-07,
+      "loss": 0.7857,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 1.0969809223697835e-05,
+      "step": 8864
+    },
+    {
+      "epoch": 0.91,
+      "gumbel_tau": 0.1235,
+      "image_masks_loss": 0.0192,
+      "llm_learning_rate": 1.0945516868108852e-07,
+      "loss": 0.8212,
+      "output_text_mask_loss": 0.004,
+      "predictor_learning_rate": 1.094551686810885e-05,
+      "step": 8865
+    },
+    {
+      "epoch": 0.91,
+      "gumbel_tau": 0.1235,
+      "image_masks_loss": 0.0117,
+      "llm_learning_rate": 1.0921250837732045e-07,
+      "loss": 0.7489,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 1.0921250837732044e-05,
+      "step": 8866
+    },
+    {
+      "epoch": 0.91,
+      "gumbel_tau": 0.1235,
+      "image_masks_loss": 0.0125,
+      "llm_learning_rate": 1.089701113523961e-07,
+      "loss": 0.655,
+      "output_text_mask_loss": 0.0132,
+      "predictor_learning_rate": 1.0897011135239609e-05,
+      "step": 8867
+    },
+    {
+      "epoch": 0.91,
+      "gumbel_tau": 0.1235,
+      "image_masks_loss": 0.0132,
+      "llm_learning_rate": 1.0872797763300885e-07,
+      "loss": 0.7689,
+      "output_text_mask_loss": 0.0233,
+      "predictor_learning_rate": 1.0872797763300885e-05,
+      "step": 8868
+    },
+    {
+      "epoch": 0.91,
+      "gumbel_tau": 0.1234,
+      "image_masks_loss": 0.0028,
+      "llm_learning_rate": 1.0848610724582292e-07,
+      "loss": 0.7694,
+      "output_text_mask_loss": 0.0141,
+      "predictor_learning_rate": 1.0848610724582292e-05,
+      "step": 8869
+    },
+    {
+      "epoch": 0.91,
+      "gumbel_tau": 0.1234,
+      "image_masks_loss": 0.0077,
+      "llm_learning_rate": 1.0824450021747424e-07,
+      "loss": 0.7944,
+      "output_text_mask_loss": 0.0025,
+      "predictor_learning_rate": 1.0824450021747423e-05,
+      "step": 8870
+    },
+    {
+      "epoch": 0.91,
+      "gumbel_tau": 0.1234,
+      "image_masks_loss": 0.0123,
+      "llm_learning_rate": 1.0800315657456816e-07,
+      "loss": 0.7491,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 1.0800315657456816e-05,
+      "step": 8871
+    },
+    {
+      "epoch": 0.91,
+      "gumbel_tau": 0.1233,
+      "image_masks_loss": 0.0052,
+      "llm_learning_rate": 1.0776207634368263e-07,
+      "loss": 0.6967,
+      "output_text_mask_loss": 0.006,
+      "predictor_learning_rate": 1.0776207634368263e-05,
+      "step": 8872
+    },
+    {
+      "epoch": 0.91,
+      "gumbel_tau": 0.1233,
+      "image_masks_loss": 0.0092,
+      "llm_learning_rate": 1.0752125955136583e-07,
+      "loss": 0.7,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 1.0752125955136583e-05,
+      "step": 8873
+    },
+    {
+      "epoch": 0.91,
+      "gumbel_tau": 0.1233,
+      "image_masks_loss": 0.0138,
+      "llm_learning_rate": 1.0728070622413683e-07,
+      "loss": 0.7304,
+      "output_text_mask_loss": 0.0128,
+      "predictor_learning_rate": 1.0728070622413682e-05,
+      "step": 8874
+    },
+    {
+      "epoch": 0.91,
+      "gumbel_tau": 0.1232,
+      "image_masks_loss": 0.0151,
+      "llm_learning_rate": 1.0704041638848584e-07,
+      "loss": 0.7377,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 1.0704041638848583e-05,
+      "step": 8875
+    },
+    {
+      "epoch": 0.91,
+      "gumbel_tau": 0.1232,
+      "image_masks_loss": 0.0025,
+      "llm_learning_rate": 1.0680039007087473e-07,
+      "loss": 0.7401,
+      "output_text_mask_loss": 0.0058,
+      "predictor_learning_rate": 1.0680039007087471e-05,
+      "step": 8876
+    },
+    {
+      "epoch": 0.91,
+      "gumbel_tau": 0.1232,
+      "image_masks_loss": 0.0188,
+      "llm_learning_rate": 1.0656062729773515e-07,
+      "loss": 0.7327,
+      "output_text_mask_loss": 0.0322,
+      "predictor_learning_rate": 1.0656062729773513e-05,
+      "step": 8877
+    },
+    {
+      "epoch": 0.91,
+      "gumbel_tau": 0.1232,
+      "image_masks_loss": 0.0141,
+      "llm_learning_rate": 1.0632112809547041e-07,
+      "loss": 0.7228,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 1.063211280954704e-05,
+      "step": 8878
+    },
+    {
+      "epoch": 0.91,
+      "gumbel_tau": 0.1231,
+      "image_masks_loss": 0.0129,
+      "llm_learning_rate": 1.060818924904547e-07,
+      "loss": 0.6908,
+      "output_text_mask_loss": 0.0095,
+      "predictor_learning_rate": 1.0608189249045469e-05,
+      "step": 8879
+    },
+    {
+      "epoch": 0.91,
+      "gumbel_tau": 0.1231,
+      "image_masks_loss": 0.0095,
+      "llm_learning_rate": 1.0584292050903389e-07,
+      "loss": 0.6903,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 1.0584292050903388e-05,
+      "step": 8880
+    },
+    {
+      "epoch": 0.91,
+      "gumbel_tau": 0.1231,
+      "image_masks_loss": 0.0069,
+      "llm_learning_rate": 1.0560421217752247e-07,
+      "loss": 0.7345,
+      "output_text_mask_loss": 0.0039,
+      "predictor_learning_rate": 1.0560421217752247e-05,
+      "step": 8881
+    },
+    {
+      "epoch": 0.91,
+      "gumbel_tau": 0.123,
+      "image_masks_loss": 0.0389,
+      "llm_learning_rate": 1.053657675222089e-07,
+      "loss": 0.65,
+      "output_text_mask_loss": 0.0081,
+      "predictor_learning_rate": 1.053657675222089e-05,
+      "step": 8882
+    },
+    {
+      "epoch": 0.91,
+      "gumbel_tau": 0.123,
+      "image_masks_loss": 0.0103,
+      "llm_learning_rate": 1.0512758656935102e-07,
+      "loss": 0.7619,
+      "output_text_mask_loss": 0.0101,
+      "predictor_learning_rate": 1.05127586569351e-05,
+      "step": 8883
+    },
+    {
+      "epoch": 0.91,
+      "gumbel_tau": 0.123,
+      "image_masks_loss": 0.0093,
+      "llm_learning_rate": 1.0488966934517758e-07,
+      "loss": 0.7478,
+      "output_text_mask_loss": 0.0124,
+      "predictor_learning_rate": 1.0488966934517757e-05,
+      "step": 8884
+    },
+    {
+      "epoch": 0.91,
+      "gumbel_tau": 0.123,
+      "image_masks_loss": 0.0106,
+      "llm_learning_rate": 1.0465201587588847e-07,
+      "loss": 0.7197,
+      "output_text_mask_loss": 0.0221,
+      "predictor_learning_rate": 1.0465201587588846e-05,
+      "step": 8885
+    },
+    {
+      "epoch": 0.91,
+      "gumbel_tau": 0.1229,
+      "image_masks_loss": 0.0133,
+      "llm_learning_rate": 1.0441462618765497e-07,
+      "loss": 0.8277,
+      "output_text_mask_loss": 0.0227,
+      "predictor_learning_rate": 1.0441462618765495e-05,
+      "step": 8886
+    },
+    {
+      "epoch": 0.91,
+      "gumbel_tau": 0.1229,
+      "image_masks_loss": 0.0156,
+      "llm_learning_rate": 1.0417750030661893e-07,
+      "loss": 0.7288,
+      "output_text_mask_loss": 0.0116,
+      "predictor_learning_rate": 1.0417750030661893e-05,
+      "step": 8887
+    },
+    {
+      "epoch": 0.91,
+      "gumbel_tau": 0.1229,
+      "image_masks_loss": 0.009,
+      "llm_learning_rate": 1.0394063825889255e-07,
+      "loss": 0.6973,
+      "output_text_mask_loss": 0.0081,
+      "predictor_learning_rate": 1.0394063825889255e-05,
+      "step": 8888
+    },
+    {
+      "epoch": 0.91,
+      "gumbel_tau": 0.1228,
+      "image_masks_loss": 0.0153,
+      "llm_learning_rate": 1.037040400705605e-07,
+      "loss": 0.7645,
+      "output_text_mask_loss": 0.0006,
+      "predictor_learning_rate": 1.037040400705605e-05,
+      "step": 8889
+    },
+    {
+      "epoch": 0.91,
+      "gumbel_tau": 0.1228,
+      "image_masks_loss": 0.0132,
+      "llm_learning_rate": 1.0346770576767696e-07,
+      "loss": 0.7367,
+      "output_text_mask_loss": 0.0298,
+      "predictor_learning_rate": 1.0346770576767694e-05,
+      "step": 8890
+    },
+    {
+      "epoch": 0.91,
+      "gumbel_tau": 0.1228,
+      "image_masks_loss": 0.014,
+      "llm_learning_rate": 1.0323163537626774e-07,
+      "loss": 0.7317,
+      "output_text_mask_loss": 0.0145,
+      "predictor_learning_rate": 1.0323163537626774e-05,
+      "step": 8891
+    },
+    {
+      "epoch": 0.91,
+      "gumbel_tau": 0.1228,
+      "image_masks_loss": 0.0141,
+      "llm_learning_rate": 1.0299582892232956e-07,
+      "loss": 0.6251,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 1.0299582892232956e-05,
+      "step": 8892
+    },
+    {
+      "epoch": 0.91,
+      "gumbel_tau": 0.1227,
+      "image_masks_loss": 0.0069,
+      "llm_learning_rate": 1.0276028643182972e-07,
+      "loss": 0.6347,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 1.0276028643182972e-05,
+      "step": 8893
+    },
+    {
+      "epoch": 0.91,
+      "gumbel_tau": 0.1227,
+      "image_masks_loss": 0.01,
+      "llm_learning_rate": 1.0252500793070691e-07,
+      "loss": 0.7969,
+      "output_text_mask_loss": 0.0101,
+      "predictor_learning_rate": 1.0252500793070691e-05,
+      "step": 8894
+    },
+    {
+      "epoch": 0.91,
+      "gumbel_tau": 0.1227,
+      "image_masks_loss": 0.0118,
+      "llm_learning_rate": 1.0228999344486984e-07,
+      "loss": 0.7038,
+      "output_text_mask_loss": 0.0189,
+      "predictor_learning_rate": 1.0228999344486984e-05,
+      "step": 8895
+    },
+    {
+      "epoch": 0.91,
+      "gumbel_tau": 0.1226,
+      "image_masks_loss": 0.0148,
+      "llm_learning_rate": 1.0205524300020004e-07,
+      "loss": 0.7351,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 1.0205524300020003e-05,
+      "step": 8896
+    },
+    {
+      "epoch": 0.91,
+      "gumbel_tau": 0.1226,
+      "image_masks_loss": 0.0122,
+      "llm_learning_rate": 1.0182075662254792e-07,
+      "loss": 0.7054,
+      "output_text_mask_loss": 0.0015,
+      "predictor_learning_rate": 1.0182075662254792e-05,
+      "step": 8897
+    },
+    {
+      "epoch": 0.91,
+      "gumbel_tau": 0.1226,
+      "image_masks_loss": 0.0165,
+      "llm_learning_rate": 1.0158653433773563e-07,
+      "loss": 0.7168,
+      "output_text_mask_loss": 0.0143,
+      "predictor_learning_rate": 1.0158653433773563e-05,
+      "step": 8898
+    },
+    {
+      "epoch": 0.91,
+      "gumbel_tau": 0.1226,
+      "image_masks_loss": 0.013,
+      "llm_learning_rate": 1.0135257617155697e-07,
+      "loss": 0.7135,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 1.0135257617155697e-05,
+      "step": 8899
+    },
+    {
+      "epoch": 0.91,
+      "gumbel_tau": 0.1225,
+      "image_masks_loss": 0.0126,
+      "llm_learning_rate": 1.0111888214977467e-07,
+      "loss": 0.7436,
+      "output_text_mask_loss": 0.005,
+      "predictor_learning_rate": 1.0111888214977466e-05,
+      "step": 8900
+    },
+    {
+      "epoch": 0.91,
+      "gumbel_tau": 0.1225,
+      "image_masks_loss": 0.0088,
+      "llm_learning_rate": 1.0088545229812479e-07,
+      "loss": 0.6568,
+      "output_text_mask_loss": 0.0061,
+      "predictor_learning_rate": 1.008854522981248e-05,
+      "step": 8901
+    },
+    {
+      "epoch": 0.91,
+      "gumbel_tau": 0.1225,
+      "image_masks_loss": 0.0086,
+      "llm_learning_rate": 1.0065228664231264e-07,
+      "loss": 0.7334,
+      "output_text_mask_loss": 0.0124,
+      "predictor_learning_rate": 1.0065228664231263e-05,
+      "step": 8902
+    },
+    {
+      "epoch": 0.91,
+      "gumbel_tau": 0.1224,
+      "image_masks_loss": 0.006,
+      "llm_learning_rate": 1.0041938520801487e-07,
+      "loss": 0.7626,
+      "output_text_mask_loss": 0.001,
+      "predictor_learning_rate": 1.0041938520801486e-05,
+      "step": 8903
+    },
+    {
+      "epoch": 0.91,
+      "gumbel_tau": 0.1224,
+      "image_masks_loss": 0.0084,
+      "llm_learning_rate": 1.0018674802087958e-07,
+      "loss": 0.665,
+      "output_text_mask_loss": 0.0094,
+      "predictor_learning_rate": 1.0018674802087958e-05,
+      "step": 8904
+    },
+    {
+      "epoch": 0.91,
+      "gumbel_tau": 0.1224,
+      "image_masks_loss": 0.007,
+      "llm_learning_rate": 9.995437510652462e-08,
+      "loss": 0.7879,
+      "output_text_mask_loss": 0.033,
+      "predictor_learning_rate": 9.995437510652461e-06,
+      "step": 8905
+    },
+    {
+      "epoch": 0.91,
+      "gumbel_tau": 0.1223,
+      "image_masks_loss": 0.0219,
+      "llm_learning_rate": 9.972226649054007e-08,
+      "loss": 0.7162,
+      "output_text_mask_loss": 0.0099,
+      "predictor_learning_rate": 9.972226649054006e-06,
+      "step": 8906
+    },
+    {
+      "epoch": 0.91,
+      "gumbel_tau": 0.1223,
+      "image_masks_loss": 0.0179,
+      "llm_learning_rate": 9.949042219848548e-08,
+      "loss": 0.7404,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 9.949042219848547e-06,
+      "step": 8907
+    },
+    {
+      "epoch": 0.91,
+      "gumbel_tau": 0.1223,
+      "image_masks_loss": 0.0114,
+      "llm_learning_rate": 9.925884225589294e-08,
+      "loss": 0.6469,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 9.925884225589293e-06,
+      "step": 8908
+    },
+    {
+      "epoch": 0.91,
+      "gumbel_tau": 0.1223,
+      "image_masks_loss": 0.0341,
+      "llm_learning_rate": 9.902752668826399e-08,
+      "loss": 0.7651,
+      "output_text_mask_loss": 0.0172,
+      "predictor_learning_rate": 9.902752668826398e-06,
+      "step": 8909
+    },
+    {
+      "epoch": 0.91,
+      "gumbel_tau": 0.1222,
+      "image_masks_loss": 0.0103,
+      "llm_learning_rate": 9.879647552107186e-08,
+      "loss": 0.6822,
+      "output_text_mask_loss": 0.0076,
+      "predictor_learning_rate": 9.879647552107184e-06,
+      "step": 8910
+    },
+    {
+      "epoch": 0.91,
+      "gumbel_tau": 0.1222,
+      "image_masks_loss": 0.0025,
+      "llm_learning_rate": 9.856568877976036e-08,
+      "loss": 0.7859,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 9.856568877976036e-06,
+      "step": 8911
+    },
+    {
+      "epoch": 0.91,
+      "gumbel_tau": 0.1222,
+      "image_masks_loss": 0.0052,
+      "llm_learning_rate": 9.833516648974417e-08,
+      "loss": 0.7856,
+      "output_text_mask_loss": 0.0025,
+      "predictor_learning_rate": 9.833516648974417e-06,
+      "step": 8912
+    },
+    {
+      "epoch": 0.91,
+      "gumbel_tau": 0.1221,
+      "image_masks_loss": 0.0146,
+      "llm_learning_rate": 9.810490867640882e-08,
+      "loss": 0.7129,
+      "output_text_mask_loss": 0.0071,
+      "predictor_learning_rate": 9.81049086764088e-06,
+      "step": 8913
+    },
+    {
+      "epoch": 0.91,
+      "gumbel_tau": 0.1221,
+      "image_masks_loss": 0.0106,
+      "llm_learning_rate": 9.787491536511068e-08,
+      "loss": 0.736,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 9.787491536511068e-06,
+      "step": 8914
+    },
+    {
+      "epoch": 0.91,
+      "gumbel_tau": 0.1221,
+      "image_masks_loss": 0.0068,
+      "llm_learning_rate": 9.764518658117782e-08,
+      "loss": 0.7149,
+      "output_text_mask_loss": 0.0185,
+      "predictor_learning_rate": 9.764518658117783e-06,
+      "step": 8915
+    },
+    {
+      "epoch": 0.91,
+      "gumbel_tau": 0.1221,
+      "image_masks_loss": 0.0046,
+      "llm_learning_rate": 9.741572234990782e-08,
+      "loss": 0.7737,
+      "output_text_mask_loss": 0.0145,
+      "predictor_learning_rate": 9.74157223499078e-06,
+      "step": 8916
+    },
+    {
+      "epoch": 0.91,
+      "gumbel_tau": 0.122,
+      "image_masks_loss": 0.0085,
+      "llm_learning_rate": 9.718652269657014e-08,
+      "loss": 0.7025,
+      "output_text_mask_loss": 0.0058,
+      "predictor_learning_rate": 9.718652269657014e-06,
+      "step": 8917
+    },
+    {
+      "epoch": 0.91,
+      "gumbel_tau": 0.122,
+      "image_masks_loss": 0.0122,
+      "llm_learning_rate": 9.695758764640517e-08,
+      "loss": 0.7019,
+      "output_text_mask_loss": 0.0305,
+      "predictor_learning_rate": 9.695758764640517e-06,
+      "step": 8918
+    },
+    {
+      "epoch": 0.91,
+      "gumbel_tau": 0.122,
+      "image_masks_loss": 0.0039,
+      "llm_learning_rate": 9.672891722462247e-08,
+      "loss": 0.7207,
+      "output_text_mask_loss": 0.0027,
+      "predictor_learning_rate": 9.672891722462247e-06,
+      "step": 8919
+    },
+    {
+      "epoch": 0.91,
+      "gumbel_tau": 0.1219,
+      "image_masks_loss": 0.0018,
+      "llm_learning_rate": 9.650051145640438e-08,
+      "loss": 0.7198,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 9.650051145640438e-06,
+      "step": 8920
+    },
+    {
+      "epoch": 0.91,
+      "gumbel_tau": 0.1219,
+      "image_masks_loss": 0.0055,
+      "llm_learning_rate": 9.627237036690412e-08,
+      "loss": 0.7294,
+      "output_text_mask_loss": 0.0223,
+      "predictor_learning_rate": 9.627237036690412e-06,
+      "step": 8921
+    },
+    {
+      "epoch": 0.91,
+      "gumbel_tau": 0.1219,
+      "image_masks_loss": 0.0071,
+      "llm_learning_rate": 9.60444939812441e-08,
+      "loss": 0.7548,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 9.60444939812441e-06,
+      "step": 8922
+    },
+    {
+      "epoch": 0.91,
+      "gumbel_tau": 0.1219,
+      "image_masks_loss": 0.0049,
+      "llm_learning_rate": 9.581688232451925e-08,
+      "loss": 0.7704,
+      "output_text_mask_loss": 0.0023,
+      "predictor_learning_rate": 9.581688232451923e-06,
+      "step": 8923
+    },
+    {
+      "epoch": 0.91,
+      "gumbel_tau": 0.1218,
+      "image_masks_loss": 0.0165,
+      "llm_learning_rate": 9.55895354217945e-08,
+      "loss": 0.7533,
+      "output_text_mask_loss": 0.0132,
+      "predictor_learning_rate": 9.55895354217945e-06,
+      "step": 8924
+    },
+    {
+      "epoch": 0.91,
+      "gumbel_tau": 0.1218,
+      "image_masks_loss": 0.0072,
+      "llm_learning_rate": 9.536245329810539e-08,
+      "loss": 0.8152,
+      "output_text_mask_loss": 0.0463,
+      "predictor_learning_rate": 9.536245329810538e-06,
+      "step": 8925
+    },
+    {
+      "epoch": 0.91,
+      "gumbel_tau": 0.1218,
+      "image_masks_loss": 0.0035,
+      "llm_learning_rate": 9.513563597845942e-08,
+      "loss": 0.5717,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 9.513563597845943e-06,
+      "step": 8926
+    },
+    {
+      "epoch": 0.91,
+      "gumbel_tau": 0.1217,
+      "image_masks_loss": 0.0042,
+      "llm_learning_rate": 9.490908348783329e-08,
+      "loss": 0.7056,
+      "output_text_mask_loss": 0.0687,
+      "predictor_learning_rate": 9.490908348783329e-06,
+      "step": 8927
+    },
+    {
+      "epoch": 0.91,
+      "gumbel_tau": 0.1217,
+      "image_masks_loss": 0.0109,
+      "llm_learning_rate": 9.468279585117646e-08,
+      "loss": 0.6814,
+      "output_text_mask_loss": 0.0467,
+      "predictor_learning_rate": 9.468279585117646e-06,
+      "step": 8928
+    },
+    {
+      "epoch": 0.91,
+      "gumbel_tau": 0.1217,
+      "image_masks_loss": 0.0041,
+      "llm_learning_rate": 9.445677309340761e-08,
+      "loss": 0.6793,
+      "output_text_mask_loss": 0.0101,
+      "predictor_learning_rate": 9.445677309340761e-06,
+      "step": 8929
+    },
+    {
+      "epoch": 0.91,
+      "gumbel_tau": 0.1217,
+      "image_masks_loss": 0.0048,
+      "llm_learning_rate": 9.423101523941714e-08,
+      "loss": 0.7147,
+      "output_text_mask_loss": 0.006,
+      "predictor_learning_rate": 9.423101523941714e-06,
+      "step": 8930
+    },
+    {
+      "epoch": 0.92,
+      "gumbel_tau": 0.1216,
+      "image_masks_loss": 0.0152,
+      "llm_learning_rate": 9.400552231406596e-08,
+      "loss": 0.736,
+      "output_text_mask_loss": 0.0107,
+      "predictor_learning_rate": 9.400552231406596e-06,
+      "step": 8931
+    },
+    {
+      "epoch": 0.92,
+      "gumbel_tau": 0.1216,
+      "image_masks_loss": 0.0027,
+      "llm_learning_rate": 9.378029434218588e-08,
+      "loss": 0.7747,
+      "output_text_mask_loss": 0.0618,
+      "predictor_learning_rate": 9.378029434218588e-06,
+      "step": 8932
+    },
+    {
+      "epoch": 0.92,
+      "gumbel_tau": 0.1216,
+      "image_masks_loss": 0.0059,
+      "llm_learning_rate": 9.35553313485793e-08,
+      "loss": 0.7278,
+      "output_text_mask_loss": 0.0032,
+      "predictor_learning_rate": 9.35553313485793e-06,
+      "step": 8933
+    },
+    {
+      "epoch": 0.92,
+      "gumbel_tau": 0.1215,
+      "image_masks_loss": 0.0078,
+      "llm_learning_rate": 9.333063335801945e-08,
+      "loss": 0.7435,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 9.333063335801944e-06,
+      "step": 8934
+    },
+    {
+      "epoch": 0.92,
+      "gumbel_tau": 0.1215,
+      "image_masks_loss": 0.0032,
+      "llm_learning_rate": 9.310620039525126e-08,
+      "loss": 0.6406,
+      "output_text_mask_loss": 0.0269,
+      "predictor_learning_rate": 9.310620039525125e-06,
+      "step": 8935
+    },
+    {
+      "epoch": 0.92,
+      "gumbel_tau": 0.1215,
+      "image_masks_loss": 0.0105,
+      "llm_learning_rate": 9.288203248498967e-08,
+      "loss": 0.7356,
+      "output_text_mask_loss": 0.0126,
+      "predictor_learning_rate": 9.288203248498967e-06,
+      "step": 8936
+    },
+    {
+      "epoch": 0.92,
+      "gumbel_tau": 0.1215,
+      "image_masks_loss": 0.0029,
+      "llm_learning_rate": 9.265812965192023e-08,
+      "loss": 0.7792,
+      "output_text_mask_loss": 0.0248,
+      "predictor_learning_rate": 9.265812965192022e-06,
+      "step": 8937
+    },
+    {
+      "epoch": 0.92,
+      "gumbel_tau": 0.1214,
+      "image_masks_loss": 0.0082,
+      "llm_learning_rate": 9.243449192069931e-08,
+      "loss": 0.6704,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 9.243449192069931e-06,
+      "step": 8938
+    },
+    {
+      "epoch": 0.92,
+      "gumbel_tau": 0.1214,
+      "image_masks_loss": 0.0202,
+      "llm_learning_rate": 9.221111931595472e-08,
+      "loss": 0.7393,
+      "output_text_mask_loss": 0.0215,
+      "predictor_learning_rate": 9.221111931595473e-06,
+      "step": 8939
+    },
+    {
+      "epoch": 0.92,
+      "gumbel_tau": 0.1214,
+      "image_masks_loss": 0.0095,
+      "llm_learning_rate": 9.198801186228512e-08,
+      "loss": 0.7763,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 9.198801186228511e-06,
+      "step": 8940
+    },
+    {
+      "epoch": 0.92,
+      "gumbel_tau": 0.1213,
+      "image_masks_loss": 0.0124,
+      "llm_learning_rate": 9.176516958425918e-08,
+      "loss": 0.8728,
+      "output_text_mask_loss": 0.0042,
+      "predictor_learning_rate": 9.176516958425918e-06,
+      "step": 8941
+    },
+    {
+      "epoch": 0.92,
+      "gumbel_tau": 0.1213,
+      "image_masks_loss": 0.0127,
+      "llm_learning_rate": 9.154259250641701e-08,
+      "loss": 0.6856,
+      "output_text_mask_loss": 0.0322,
+      "predictor_learning_rate": 9.1542592506417e-06,
+      "step": 8942
+    },
+    {
+      "epoch": 0.92,
+      "gumbel_tau": 0.1213,
+      "image_masks_loss": 0.0118,
+      "llm_learning_rate": 9.132028065326926e-08,
+      "loss": 0.6553,
+      "output_text_mask_loss": 0.0204,
+      "predictor_learning_rate": 9.132028065326925e-06,
+      "step": 8943
+    },
+    {
+      "epoch": 0.92,
+      "gumbel_tau": 0.1213,
+      "image_masks_loss": 0.0153,
+      "llm_learning_rate": 9.109823404929719e-08,
+      "loss": 0.6804,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 9.109823404929718e-06,
+      "step": 8944
+    },
+    {
+      "epoch": 0.92,
+      "gumbel_tau": 0.1212,
+      "image_masks_loss": 0.0198,
+      "llm_learning_rate": 9.087645271895317e-08,
+      "loss": 0.7546,
+      "output_text_mask_loss": 0.0366,
+      "predictor_learning_rate": 9.087645271895317e-06,
+      "step": 8945
+    },
+    {
+      "epoch": 0.92,
+      "gumbel_tau": 0.1212,
+      "image_masks_loss": 0.0074,
+      "llm_learning_rate": 9.06549366866602e-08,
+      "loss": 0.6267,
+      "output_text_mask_loss": 0.0055,
+      "predictor_learning_rate": 9.065493668666019e-06,
+      "step": 8946
+    },
+    {
+      "epoch": 0.92,
+      "gumbel_tau": 0.1212,
+      "image_masks_loss": 0.0284,
+      "llm_learning_rate": 9.043368597681234e-08,
+      "loss": 0.6369,
+      "output_text_mask_loss": 0.0024,
+      "predictor_learning_rate": 9.043368597681235e-06,
+      "step": 8947
+    },
+    {
+      "epoch": 0.92,
+      "gumbel_tau": 0.1211,
+      "image_masks_loss": 0.0072,
+      "llm_learning_rate": 9.021270061377402e-08,
+      "loss": 0.7434,
+      "output_text_mask_loss": 0.0071,
+      "predictor_learning_rate": 9.0212700613774e-06,
+      "step": 8948
+    },
+    {
+      "epoch": 0.92,
+      "gumbel_tau": 0.1211,
+      "image_masks_loss": 0.0109,
+      "llm_learning_rate": 8.999198062188075e-08,
+      "loss": 0.8181,
+      "output_text_mask_loss": 0.0111,
+      "predictor_learning_rate": 8.999198062188074e-06,
+      "step": 8949
+    },
+    {
+      "epoch": 0.92,
+      "gumbel_tau": 0.1211,
+      "image_masks_loss": 0.0099,
+      "llm_learning_rate": 8.977152602543866e-08,
+      "loss": 0.688,
+      "output_text_mask_loss": 0.0004,
+      "predictor_learning_rate": 8.977152602543865e-06,
+      "step": 8950
+    },
+    {
+      "epoch": 0.92,
+      "gumbel_tau": 0.1211,
+      "image_masks_loss": 0.0192,
+      "llm_learning_rate": 8.955133684872469e-08,
+      "loss": 0.7099,
+      "output_text_mask_loss": 0.0322,
+      "predictor_learning_rate": 8.95513368487247e-06,
+      "step": 8951
+    },
+    {
+      "epoch": 0.92,
+      "gumbel_tau": 0.121,
+      "image_masks_loss": 0.006,
+      "llm_learning_rate": 8.933141311598669e-08,
+      "loss": 0.7287,
+      "output_text_mask_loss": 0.0282,
+      "predictor_learning_rate": 8.933141311598669e-06,
+      "step": 8952
+    },
+    {
+      "epoch": 0.92,
+      "gumbel_tau": 0.121,
+      "image_masks_loss": 0.0044,
+      "llm_learning_rate": 8.91117548514428e-08,
+      "loss": 0.6992,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 8.911175485144279e-06,
+      "step": 8953
+    },
+    {
+      "epoch": 0.92,
+      "gumbel_tau": 0.121,
+      "image_masks_loss": 0.0079,
+      "llm_learning_rate": 8.889236207928254e-08,
+      "loss": 0.7302,
+      "output_text_mask_loss": 0.0068,
+      "predictor_learning_rate": 8.889236207928253e-06,
+      "step": 8954
+    },
+    {
+      "epoch": 0.92,
+      "gumbel_tau": 0.1209,
+      "image_masks_loss": 0.0073,
+      "llm_learning_rate": 8.86732348236663e-08,
+      "loss": 0.7484,
+      "output_text_mask_loss": 0.0408,
+      "predictor_learning_rate": 8.867323482366629e-06,
+      "step": 8955
+    },
+    {
+      "epoch": 0.92,
+      "gumbel_tau": 0.1209,
+      "image_masks_loss": 0.0209,
+      "llm_learning_rate": 8.845437310872423e-08,
+      "loss": 0.6774,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 8.845437310872423e-06,
+      "step": 8956
+    },
+    {
+      "epoch": 0.92,
+      "gumbel_tau": 0.1209,
+      "image_masks_loss": 0.0067,
+      "llm_learning_rate": 8.823577695855789e-08,
+      "loss": 0.706,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 8.82357769585579e-06,
+      "step": 8957
+    },
+    {
+      "epoch": 0.92,
+      "gumbel_tau": 0.1209,
+      "image_masks_loss": 0.0063,
+      "llm_learning_rate": 8.801744639723997e-08,
+      "loss": 0.6792,
+      "output_text_mask_loss": 0.0391,
+      "predictor_learning_rate": 8.801744639723996e-06,
+      "step": 8958
+    },
+    {
+      "epoch": 0.92,
+      "gumbel_tau": 0.1208,
+      "image_masks_loss": 0.0109,
+      "llm_learning_rate": 8.779938144881318e-08,
+      "loss": 0.7813,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 8.779938144881316e-06,
+      "step": 8959
+    },
+    {
+      "epoch": 0.92,
+      "gumbel_tau": 0.1208,
+      "image_masks_loss": 0.0158,
+      "llm_learning_rate": 8.758158213729162e-08,
+      "loss": 0.6738,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 8.758158213729161e-06,
+      "step": 8960
+    },
+    {
+      "epoch": 0.92,
+      "gumbel_tau": 0.1208,
+      "image_masks_loss": 0.0238,
+      "llm_learning_rate": 8.736404848665947e-08,
+      "loss": 0.7324,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 8.736404848665946e-06,
+      "step": 8961
+    },
+    {
+      "epoch": 0.92,
+      "gumbel_tau": 0.1207,
+      "image_masks_loss": 0.0071,
+      "llm_learning_rate": 8.714678052087255e-08,
+      "loss": 0.7033,
+      "output_text_mask_loss": 0.071,
+      "predictor_learning_rate": 8.714678052087256e-06,
+      "step": 8962
+    },
+    {
+      "epoch": 0.92,
+      "gumbel_tau": 0.1207,
+      "image_masks_loss": 0.0107,
+      "llm_learning_rate": 8.692977826385645e-08,
+      "loss": 0.619,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 8.692977826385645e-06,
+      "step": 8963
+    },
+    {
+      "epoch": 0.92,
+      "gumbel_tau": 0.1207,
+      "image_masks_loss": 0.0051,
+      "llm_learning_rate": 8.671304173950789e-08,
+      "loss": 0.7645,
+      "output_text_mask_loss": 0.0012,
+      "predictor_learning_rate": 8.67130417395079e-06,
+      "step": 8964
+    },
+    {
+      "epoch": 0.92,
+      "gumbel_tau": 0.1207,
+      "image_masks_loss": 0.0034,
+      "llm_learning_rate": 8.649657097169501e-08,
+      "loss": 0.7193,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 8.6496570971695e-06,
+      "step": 8965
+    },
+    {
+      "epoch": 0.92,
+      "gumbel_tau": 0.1206,
+      "image_masks_loss": 0.0146,
+      "llm_learning_rate": 8.628036598425482e-08,
+      "loss": 0.6927,
+      "output_text_mask_loss": 0.0311,
+      "predictor_learning_rate": 8.628036598425482e-06,
+      "step": 8966
+    },
+    {
+      "epoch": 0.92,
+      "gumbel_tau": 0.1206,
+      "image_masks_loss": 0.0073,
+      "llm_learning_rate": 8.606442680099774e-08,
+      "loss": 0.6276,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 8.606442680099775e-06,
+      "step": 8967
+    },
+    {
+      "epoch": 0.92,
+      "gumbel_tau": 0.1206,
+      "image_masks_loss": 0.0033,
+      "llm_learning_rate": 8.584875344570254e-08,
+      "loss": 0.6767,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 8.584875344570253e-06,
+      "step": 8968
+    },
+    {
+      "epoch": 0.92,
+      "gumbel_tau": 0.1205,
+      "image_masks_loss": 0.0085,
+      "llm_learning_rate": 8.56333459421202e-08,
+      "loss": 0.7027,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 8.56333459421202e-06,
+      "step": 8969
+    },
+    {
+      "epoch": 0.92,
+      "gumbel_tau": 0.1205,
+      "image_masks_loss": 0.0075,
+      "llm_learning_rate": 8.541820431397147e-08,
+      "loss": 0.6984,
+      "output_text_mask_loss": 0.0195,
+      "predictor_learning_rate": 8.541820431397146e-06,
+      "step": 8970
+    },
+    {
+      "epoch": 0.92,
+      "gumbel_tau": 0.1205,
+      "image_masks_loss": 0.0115,
+      "llm_learning_rate": 8.520332858494823e-08,
+      "loss": 0.7153,
+      "output_text_mask_loss": 0.0202,
+      "predictor_learning_rate": 8.520332858494822e-06,
+      "step": 8971
+    },
+    {
+      "epoch": 0.92,
+      "gumbel_tau": 0.1205,
+      "image_masks_loss": 0.0044,
+      "llm_learning_rate": 8.498871877871323e-08,
+      "loss": 0.7227,
+      "output_text_mask_loss": 0.058,
+      "predictor_learning_rate": 8.498871877871323e-06,
+      "step": 8972
+    },
+    {
+      "epoch": 0.92,
+      "gumbel_tau": 0.1204,
+      "image_masks_loss": 0.0116,
+      "llm_learning_rate": 8.477437491889978e-08,
+      "loss": 0.7245,
+      "output_text_mask_loss": 0.0041,
+      "predictor_learning_rate": 8.477437491889978e-06,
+      "step": 8973
+    },
+    {
+      "epoch": 0.92,
+      "gumbel_tau": 0.1204,
+      "image_masks_loss": 0.003,
+      "llm_learning_rate": 8.456029702911206e-08,
+      "loss": 0.7336,
+      "output_text_mask_loss": 0.0081,
+      "predictor_learning_rate": 8.456029702911205e-06,
+      "step": 8974
+    },
+    {
+      "epoch": 0.92,
+      "gumbel_tau": 0.1204,
+      "image_masks_loss": 0.0041,
+      "llm_learning_rate": 8.434648513292454e-08,
+      "loss": 0.7696,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 8.434648513292454e-06,
+      "step": 8975
+    },
+    {
+      "epoch": 0.92,
+      "gumbel_tau": 0.1203,
+      "image_masks_loss": 0.0087,
+      "llm_learning_rate": 8.413293925388256e-08,
+      "loss": 0.817,
+      "output_text_mask_loss": 0.0158,
+      "predictor_learning_rate": 8.413293925388254e-06,
+      "step": 8976
+    },
+    {
+      "epoch": 0.92,
+      "gumbel_tau": 0.1203,
+      "image_masks_loss": 0.0149,
+      "llm_learning_rate": 8.391965941550257e-08,
+      "loss": 0.7361,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 8.391965941550257e-06,
+      "step": 8977
+    },
+    {
+      "epoch": 0.92,
+      "gumbel_tau": 0.1203,
+      "image_masks_loss": 0.0078,
+      "llm_learning_rate": 8.370664564127078e-08,
+      "loss": 0.7476,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 8.370664564127078e-06,
+      "step": 8978
+    },
+    {
+      "epoch": 0.92,
+      "gumbel_tau": 0.1203,
+      "image_masks_loss": 0.0153,
+      "llm_learning_rate": 8.349389795464568e-08,
+      "loss": 0.6909,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 8.349389795464568e-06,
+      "step": 8979
+    },
+    {
+      "epoch": 0.92,
+      "gumbel_tau": 0.1202,
+      "image_masks_loss": 0.0103,
+      "llm_learning_rate": 8.328141637905517e-08,
+      "loss": 0.7545,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 8.328141637905518e-06,
+      "step": 8980
+    },
+    {
+      "epoch": 0.92,
+      "gumbel_tau": 0.1202,
+      "image_masks_loss": 0.0064,
+      "llm_learning_rate": 8.306920093789805e-08,
+      "loss": 0.7408,
+      "output_text_mask_loss": 0.0357,
+      "predictor_learning_rate": 8.306920093789805e-06,
+      "step": 8981
+    },
+    {
+      "epoch": 0.92,
+      "gumbel_tau": 0.1202,
+      "image_masks_loss": 0.0144,
+      "llm_learning_rate": 8.285725165454423e-08,
+      "loss": 0.7107,
+      "output_text_mask_loss": 0.0039,
+      "predictor_learning_rate": 8.285725165454421e-06,
+      "step": 8982
+    },
+    {
+      "epoch": 0.92,
+      "gumbel_tau": 0.1201,
+      "image_masks_loss": 0.0066,
+      "llm_learning_rate": 8.264556855233363e-08,
+      "loss": 0.7417,
+      "output_text_mask_loss": 0.0298,
+      "predictor_learning_rate": 8.264556855233363e-06,
+      "step": 8983
+    },
+    {
+      "epoch": 0.92,
+      "gumbel_tau": 0.1201,
+      "image_masks_loss": 0.0094,
+      "llm_learning_rate": 8.243415165457736e-08,
+      "loss": 0.6777,
+      "output_text_mask_loss": 0.0012,
+      "predictor_learning_rate": 8.243415165457734e-06,
+      "step": 8984
+    },
+    {
+      "epoch": 0.92,
+      "gumbel_tau": 0.1201,
+      "image_masks_loss": 0.0197,
+      "llm_learning_rate": 8.222300098455732e-08,
+      "loss": 0.7587,
+      "output_text_mask_loss": 0.0015,
+      "predictor_learning_rate": 8.222300098455731e-06,
+      "step": 8985
+    },
+    {
+      "epoch": 0.92,
+      "gumbel_tau": 0.1201,
+      "image_masks_loss": 0.0079,
+      "llm_learning_rate": 8.201211656552577e-08,
+      "loss": 0.638,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 8.201211656552577e-06,
+      "step": 8986
+    },
+    {
+      "epoch": 0.92,
+      "gumbel_tau": 0.12,
+      "image_masks_loss": 0.0064,
+      "llm_learning_rate": 8.180149842070578e-08,
+      "loss": 0.7858,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 8.180149842070578e-06,
+      "step": 8987
+    },
+    {
+      "epoch": 0.92,
+      "gumbel_tau": 0.12,
+      "image_masks_loss": 0.0036,
+      "llm_learning_rate": 8.159114657329131e-08,
+      "loss": 0.7422,
+      "output_text_mask_loss": 0.1007,
+      "predictor_learning_rate": 8.15911465732913e-06,
+      "step": 8988
+    },
+    {
+      "epoch": 0.92,
+      "gumbel_tau": 0.12,
+      "image_masks_loss": 0.0216,
+      "llm_learning_rate": 8.138106104644633e-08,
+      "loss": 0.6936,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 8.138106104644633e-06,
+      "step": 8989
+    },
+    {
+      "epoch": 0.92,
+      "gumbel_tau": 0.1199,
+      "image_masks_loss": 0.01,
+      "llm_learning_rate": 8.11712418633065e-08,
+      "loss": 0.7955,
+      "output_text_mask_loss": 0.0208,
+      "predictor_learning_rate": 8.11712418633065e-06,
+      "step": 8990
+    },
+    {
+      "epoch": 0.92,
+      "gumbel_tau": 0.1199,
+      "image_masks_loss": 0.0108,
+      "llm_learning_rate": 8.096168904697693e-08,
+      "loss": 0.6553,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 8.096168904697693e-06,
+      "step": 8991
+    },
+    {
+      "epoch": 0.92,
+      "gumbel_tau": 0.1199,
+      "image_masks_loss": 0.0056,
+      "llm_learning_rate": 8.075240262053447e-08,
+      "loss": 0.7626,
+      "output_text_mask_loss": 0.0058,
+      "predictor_learning_rate": 8.075240262053445e-06,
+      "step": 8992
+    },
+    {
+      "epoch": 0.92,
+      "gumbel_tau": 0.1199,
+      "image_masks_loss": 0.0062,
+      "llm_learning_rate": 8.05433826070262e-08,
+      "loss": 0.6469,
+      "output_text_mask_loss": 0.0038,
+      "predictor_learning_rate": 8.05433826070262e-06,
+      "step": 8993
+    },
+    {
+      "epoch": 0.92,
+      "gumbel_tau": 0.1198,
+      "image_masks_loss": 0.0054,
+      "llm_learning_rate": 8.033462902946959e-08,
+      "loss": 0.6988,
+      "output_text_mask_loss": 0.0748,
+      "predictor_learning_rate": 8.033462902946959e-06,
+      "step": 8994
+    },
+    {
+      "epoch": 0.92,
+      "gumbel_tau": 0.1198,
+      "image_masks_loss": 0.0066,
+      "llm_learning_rate": 8.012614191085343e-08,
+      "loss": 0.7601,
+      "output_text_mask_loss": 0.0126,
+      "predictor_learning_rate": 8.012614191085344e-06,
+      "step": 8995
+    },
+    {
+      "epoch": 0.92,
+      "gumbel_tau": 0.1198,
+      "image_masks_loss": 0.0073,
+      "llm_learning_rate": 7.991792127413633e-08,
+      "loss": 0.673,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 7.991792127413634e-06,
+      "step": 8996
+    },
+    {
+      "epoch": 0.92,
+      "gumbel_tau": 0.1198,
+      "image_masks_loss": 0.0163,
+      "llm_learning_rate": 7.970996714224855e-08,
+      "loss": 0.7627,
+      "output_text_mask_loss": 0.005,
+      "predictor_learning_rate": 7.970996714224854e-06,
+      "step": 8997
+    },
+    {
+      "epoch": 0.92,
+      "gumbel_tau": 0.1197,
+      "image_masks_loss": 0.0236,
+      "llm_learning_rate": 7.950227953808954e-08,
+      "loss": 0.7048,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 7.950227953808953e-06,
+      "step": 8998
+    },
+    {
+      "epoch": 0.92,
+      "gumbel_tau": 0.1197,
+      "image_masks_loss": 0.0069,
+      "llm_learning_rate": 7.929485848453155e-08,
+      "loss": 0.6516,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 7.929485848453155e-06,
+      "step": 8999
+    },
+    {
+      "epoch": 0.92,
+      "gumbel_tau": 0.1197,
+      "image_masks_loss": 0.0165,
+      "llm_learning_rate": 7.908770400441551e-08,
+      "loss": 0.8558,
+      "output_text_mask_loss": 0.0278,
+      "predictor_learning_rate": 7.90877040044155e-06,
+      "step": 9000
+    },
+    {
+      "epoch": 0.92,
+      "gumbel_tau": 0.1196,
+      "image_masks_loss": 0.0089,
+      "llm_learning_rate": 7.888081612055399e-08,
+      "loss": 0.7101,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 7.888081612055397e-06,
+      "step": 9001
+    },
+    {
+      "epoch": 0.92,
+      "gumbel_tau": 0.1196,
+      "image_masks_loss": 0.0189,
+      "llm_learning_rate": 7.867419485572958e-08,
+      "loss": 0.7136,
+      "output_text_mask_loss": 0.0192,
+      "predictor_learning_rate": 7.867419485572957e-06,
+      "step": 9002
+    },
+    {
+      "epoch": 0.92,
+      "gumbel_tau": 0.1196,
+      "image_masks_loss": 0.0037,
+      "llm_learning_rate": 7.846784023269632e-08,
+      "loss": 0.7195,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 7.846784023269632e-06,
+      "step": 9003
+    },
+    {
+      "epoch": 0.92,
+      "gumbel_tau": 0.1196,
+      "image_masks_loss": 0.0094,
+      "llm_learning_rate": 7.826175227417826e-08,
+      "loss": 0.6997,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 7.826175227417826e-06,
+      "step": 9004
+    },
+    {
+      "epoch": 0.92,
+      "gumbel_tau": 0.1195,
+      "image_masks_loss": 0.0101,
+      "llm_learning_rate": 7.805593100287001e-08,
+      "loss": 0.6835,
+      "output_text_mask_loss": 0.0322,
+      "predictor_learning_rate": 7.805593100287001e-06,
+      "step": 9005
+    },
+    {
+      "epoch": 0.92,
+      "gumbel_tau": 0.1195,
+      "image_masks_loss": 0.0132,
+      "llm_learning_rate": 7.785037644143733e-08,
+      "loss": 0.7326,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 7.785037644143734e-06,
+      "step": 9006
+    },
+    {
+      "epoch": 0.92,
+      "gumbel_tau": 0.1195,
+      "image_masks_loss": 0.0074,
+      "llm_learning_rate": 7.764508861251657e-08,
+      "loss": 0.6782,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 7.764508861251658e-06,
+      "step": 9007
+    },
+    {
+      "epoch": 0.92,
+      "gumbel_tau": 0.1194,
+      "image_masks_loss": 0.0233,
+      "llm_learning_rate": 7.744006753871408e-08,
+      "loss": 0.7508,
+      "output_text_mask_loss": 0.0231,
+      "predictor_learning_rate": 7.744006753871407e-06,
+      "step": 9008
+    },
+    {
+      "epoch": 0.92,
+      "gumbel_tau": 0.1194,
+      "image_masks_loss": 0.0094,
+      "llm_learning_rate": 7.723531324260764e-08,
+      "loss": 0.74,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 7.723531324260763e-06,
+      "step": 9009
+    },
+    {
+      "epoch": 0.92,
+      "gumbel_tau": 0.1194,
+      "image_masks_loss": 0.0081,
+      "llm_learning_rate": 7.703082574674475e-08,
+      "loss": 0.6463,
+      "output_text_mask_loss": 0.0786,
+      "predictor_learning_rate": 7.703082574674475e-06,
+      "step": 9010
+    },
+    {
+      "epoch": 0.92,
+      "gumbel_tau": 0.1194,
+      "image_masks_loss": 0.0092,
+      "llm_learning_rate": 7.682660507364437e-08,
+      "loss": 0.8038,
+      "output_text_mask_loss": 0.028,
+      "predictor_learning_rate": 7.682660507364437e-06,
+      "step": 9011
+    },
+    {
+      "epoch": 0.92,
+      "gumbel_tau": 0.1193,
+      "image_masks_loss": 0.0175,
+      "llm_learning_rate": 7.662265124579572e-08,
+      "loss": 0.6606,
+      "output_text_mask_loss": 0.0034,
+      "predictor_learning_rate": 7.66226512457957e-06,
+      "step": 9012
+    },
+    {
+      "epoch": 0.92,
+      "gumbel_tau": 0.1193,
+      "image_masks_loss": 0.0161,
+      "llm_learning_rate": 7.641896428565832e-08,
+      "loss": 0.6672,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 7.641896428565831e-06,
+      "step": 9013
+    },
+    {
+      "epoch": 0.92,
+      "gumbel_tau": 0.1193,
+      "image_masks_loss": 0.0068,
+      "llm_learning_rate": 7.621554421566312e-08,
+      "loss": 0.6858,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 7.621554421566312e-06,
+      "step": 9014
+    },
+    {
+      "epoch": 0.92,
+      "gumbel_tau": 0.1192,
+      "image_masks_loss": 0.0074,
+      "llm_learning_rate": 7.601239105821084e-08,
+      "loss": 0.8153,
+      "output_text_mask_loss": 0.0065,
+      "predictor_learning_rate": 7.601239105821084e-06,
+      "step": 9015
+    },
+    {
+      "epoch": 0.92,
+      "gumbel_tau": 0.1192,
+      "image_masks_loss": 0.0143,
+      "llm_learning_rate": 7.580950483567329e-08,
+      "loss": 0.7663,
+      "output_text_mask_loss": 0.0156,
+      "predictor_learning_rate": 7.580950483567328e-06,
+      "step": 9016
+    },
+    {
+      "epoch": 0.92,
+      "gumbel_tau": 0.1192,
+      "image_masks_loss": 0.0056,
+      "llm_learning_rate": 7.560688557039259e-08,
+      "loss": 0.7882,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 7.5606885570392585e-06,
+      "step": 9017
+    },
+    {
+      "epoch": 0.92,
+      "gumbel_tau": 0.1192,
+      "image_masks_loss": 0.023,
+      "llm_learning_rate": 7.5404533284682e-08,
+      "loss": 0.6748,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 7.5404533284682e-06,
+      "step": 9018
+    },
+    {
+      "epoch": 0.92,
+      "gumbel_tau": 0.1191,
+      "image_masks_loss": 0.0051,
+      "llm_learning_rate": 7.520244800082482e-08,
+      "loss": 0.7816,
+      "output_text_mask_loss": 0.0238,
+      "predictor_learning_rate": 7.520244800082482e-06,
+      "step": 9019
+    },
+    {
+      "epoch": 0.92,
+      "gumbel_tau": 0.1191,
+      "image_masks_loss": 0.0077,
+      "llm_learning_rate": 7.500062974107491e-08,
+      "loss": 0.6479,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 7.500062974107491e-06,
+      "step": 9020
+    },
+    {
+      "epoch": 0.92,
+      "gumbel_tau": 0.1191,
+      "image_masks_loss": 0.0115,
+      "llm_learning_rate": 7.479907852765728e-08,
+      "loss": 0.7706,
+      "output_text_mask_loss": 0.0099,
+      "predictor_learning_rate": 7.479907852765727e-06,
+      "step": 9021
+    },
+    {
+      "epoch": 0.92,
+      "gumbel_tau": 0.119,
+      "image_masks_loss": 0.0049,
+      "llm_learning_rate": 7.459779438276693e-08,
+      "loss": 0.7145,
+      "output_text_mask_loss": 0.001,
+      "predictor_learning_rate": 7.459779438276693e-06,
+      "step": 9022
+    },
+    {
+      "epoch": 0.92,
+      "gumbel_tau": 0.119,
+      "image_masks_loss": 0.0134,
+      "llm_learning_rate": 7.439677732857004e-08,
+      "loss": 0.8398,
+      "output_text_mask_loss": 0.0014,
+      "predictor_learning_rate": 7.439677732857003e-06,
+      "step": 9023
+    },
+    {
+      "epoch": 0.92,
+      "gumbel_tau": 0.119,
+      "image_masks_loss": 0.0117,
+      "llm_learning_rate": 7.41960273872025e-08,
+      "loss": 0.7062,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 7.4196027387202495e-06,
+      "step": 9024
+    },
+    {
+      "epoch": 0.92,
+      "gumbel_tau": 0.119,
+      "image_masks_loss": 0.0094,
+      "llm_learning_rate": 7.399554458077191e-08,
+      "loss": 0.6988,
+      "output_text_mask_loss": 0.0099,
+      "predictor_learning_rate": 7.39955445807719e-06,
+      "step": 9025
+    },
+    {
+      "epoch": 0.92,
+      "gumbel_tau": 0.1189,
+      "image_masks_loss": 0.008,
+      "llm_learning_rate": 7.37953289313556e-08,
+      "loss": 0.7492,
+      "output_text_mask_loss": 0.0082,
+      "predictor_learning_rate": 7.37953289313556e-06,
+      "step": 9026
+    },
+    {
+      "epoch": 0.92,
+      "gumbel_tau": 0.1189,
+      "image_masks_loss": 0.0063,
+      "llm_learning_rate": 7.359538046100207e-08,
+      "loss": 0.6897,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 7.359538046100206e-06,
+      "step": 9027
+    },
+    {
+      "epoch": 0.93,
+      "gumbel_tau": 0.1189,
+      "image_masks_loss": 0.0181,
+      "llm_learning_rate": 7.33956991917295e-08,
+      "loss": 0.7751,
+      "output_text_mask_loss": 0.0319,
+      "predictor_learning_rate": 7.33956991917295e-06,
+      "step": 9028
+    },
+    {
+      "epoch": 0.93,
+      "gumbel_tau": 0.1189,
+      "image_masks_loss": 0.0028,
+      "llm_learning_rate": 7.319628514552813e-08,
+      "loss": 0.7067,
+      "output_text_mask_loss": 0.0338,
+      "predictor_learning_rate": 7.319628514552812e-06,
+      "step": 9029
+    },
+    {
+      "epoch": 0.93,
+      "gumbel_tau": 0.1188,
+      "image_masks_loss": 0.0135,
+      "llm_learning_rate": 7.299713834435673e-08,
+      "loss": 0.7289,
+      "output_text_mask_loss": 0.0402,
+      "predictor_learning_rate": 7.299713834435673e-06,
+      "step": 9030
+    },
+    {
+      "epoch": 0.93,
+      "gumbel_tau": 0.1188,
+      "image_masks_loss": 0.0073,
+      "llm_learning_rate": 7.279825881014641e-08,
+      "loss": 0.6982,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 7.2798258810146405e-06,
+      "step": 9031
+    },
+    {
+      "epoch": 0.93,
+      "gumbel_tau": 0.1188,
+      "image_masks_loss": 0.0055,
+      "llm_learning_rate": 7.259964656479851e-08,
+      "loss": 0.7108,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 7.259964656479851e-06,
+      "step": 9032
+    },
+    {
+      "epoch": 0.93,
+      "gumbel_tau": 0.1187,
+      "image_masks_loss": 0.0172,
+      "llm_learning_rate": 7.240130163018443e-08,
+      "loss": 0.6489,
+      "output_text_mask_loss": 0.0058,
+      "predictor_learning_rate": 7.240130163018444e-06,
+      "step": 9033
+    },
+    {
+      "epoch": 0.93,
+      "gumbel_tau": 0.1187,
+      "image_masks_loss": 0.0096,
+      "llm_learning_rate": 7.220322402814589e-08,
+      "loss": 0.7793,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 7.220322402814589e-06,
+      "step": 9034
+    },
+    {
+      "epoch": 0.93,
+      "gumbel_tau": 0.1187,
+      "image_masks_loss": 0.0034,
+      "llm_learning_rate": 7.200541378049625e-08,
+      "loss": 0.7443,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 7.200541378049624e-06,
+      "step": 9035
+    },
+    {
+      "epoch": 0.93,
+      "gumbel_tau": 0.1187,
+      "image_masks_loss": 0.0147,
+      "llm_learning_rate": 7.180787090901864e-08,
+      "loss": 0.6049,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 7.180787090901863e-06,
+      "step": 9036
+    },
+    {
+      "epoch": 0.93,
+      "gumbel_tau": 0.1186,
+      "image_masks_loss": 0.0127,
+      "llm_learning_rate": 7.16105954354665e-08,
+      "loss": 0.7075,
+      "output_text_mask_loss": 0.0025,
+      "predictor_learning_rate": 7.16105954354665e-06,
+      "step": 9037
+    },
+    {
+      "epoch": 0.93,
+      "gumbel_tau": 0.1186,
+      "image_masks_loss": 0.0091,
+      "llm_learning_rate": 7.141358738156495e-08,
+      "loss": 0.7168,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 7.141358738156495e-06,
+      "step": 9038
+    },
+    {
+      "epoch": 0.93,
+      "gumbel_tau": 0.1186,
+      "image_masks_loss": 0.0046,
+      "llm_learning_rate": 7.121684676900858e-08,
+      "loss": 0.6822,
+      "output_text_mask_loss": 0.0032,
+      "predictor_learning_rate": 7.121684676900858e-06,
+      "step": 9039
+    },
+    {
+      "epoch": 0.93,
+      "gumbel_tau": 0.1185,
+      "image_masks_loss": 0.0023,
+      "llm_learning_rate": 7.102037361946312e-08,
+      "loss": 0.7399,
+      "output_text_mask_loss": 0.0496,
+      "predictor_learning_rate": 7.102037361946312e-06,
+      "step": 9040
+    },
+    {
+      "epoch": 0.93,
+      "gumbel_tau": 0.1185,
+      "image_masks_loss": 0.0096,
+      "llm_learning_rate": 7.082416795456432e-08,
+      "loss": 0.6829,
+      "output_text_mask_loss": 0.0023,
+      "predictor_learning_rate": 7.0824167954564315e-06,
+      "step": 9041
+    },
+    {
+      "epoch": 0.93,
+      "gumbel_tau": 0.1185,
+      "image_masks_loss": 0.0036,
+      "llm_learning_rate": 7.062822979591905e-08,
+      "loss": 0.7059,
+      "output_text_mask_loss": 0.0095,
+      "predictor_learning_rate": 7.062822979591904e-06,
+      "step": 9042
+    },
+    {
+      "epoch": 0.93,
+      "gumbel_tau": 0.1185,
+      "image_masks_loss": 0.0165,
+      "llm_learning_rate": 7.043255916510395e-08,
+      "loss": 0.7864,
+      "output_text_mask_loss": 0.0077,
+      "predictor_learning_rate": 7.043255916510394e-06,
+      "step": 9043
+    },
+    {
+      "epoch": 0.93,
+      "gumbel_tau": 0.1184,
+      "image_masks_loss": 0.0025,
+      "llm_learning_rate": 7.023715608366705e-08,
+      "loss": 0.6694,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 7.023715608366704e-06,
+      "step": 9044
+    },
+    {
+      "epoch": 0.93,
+      "gumbel_tau": 0.1184,
+      "image_masks_loss": 0.0057,
+      "llm_learning_rate": 7.004202057312698e-08,
+      "loss": 0.7028,
+      "output_text_mask_loss": 0.0086,
+      "predictor_learning_rate": 7.004202057312697e-06,
+      "step": 9045
+    },
+    {
+      "epoch": 0.93,
+      "gumbel_tau": 0.1184,
+      "image_masks_loss": 0.0096,
+      "llm_learning_rate": 6.98471526549721e-08,
+      "loss": 0.7193,
+      "output_text_mask_loss": 0.0038,
+      "predictor_learning_rate": 6.98471526549721e-06,
+      "step": 9046
+    },
+    {
+      "epoch": 0.93,
+      "gumbel_tau": 0.1183,
+      "image_masks_loss": 0.008,
+      "llm_learning_rate": 6.965255235066192e-08,
+      "loss": 0.7558,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 6.965255235066193e-06,
+      "step": 9047
+    },
+    {
+      "epoch": 0.93,
+      "gumbel_tau": 0.1183,
+      "image_masks_loss": 0.0083,
+      "llm_learning_rate": 6.945821968162569e-08,
+      "loss": 0.7642,
+      "output_text_mask_loss": 0.0261,
+      "predictor_learning_rate": 6.945821968162569e-06,
+      "step": 9048
+    },
+    {
+      "epoch": 0.93,
+      "gumbel_tau": 0.1183,
+      "image_masks_loss": 0.0081,
+      "llm_learning_rate": 6.926415466926406e-08,
+      "loss": 0.6933,
+      "output_text_mask_loss": 0.0275,
+      "predictor_learning_rate": 6.926415466926405e-06,
+      "step": 9049
+    },
+    {
+      "epoch": 0.93,
+      "gumbel_tau": 0.1183,
+      "image_masks_loss": 0.0062,
+      "llm_learning_rate": 6.907035733494799e-08,
+      "loss": 0.7925,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 6.9070357334947986e-06,
+      "step": 9050
+    },
+    {
+      "epoch": 0.93,
+      "gumbel_tau": 0.1182,
+      "image_masks_loss": 0.0134,
+      "llm_learning_rate": 6.887682770001902e-08,
+      "loss": 0.7723,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 6.887682770001902e-06,
+      "step": 9051
+    },
+    {
+      "epoch": 0.93,
+      "gumbel_tau": 0.1182,
+      "image_masks_loss": 0.0124,
+      "llm_learning_rate": 6.8683565785789e-08,
+      "loss": 0.8105,
+      "output_text_mask_loss": 0.0694,
+      "predictor_learning_rate": 6.8683565785789e-06,
+      "step": 9052
+    },
+    {
+      "epoch": 0.93,
+      "gumbel_tau": 0.1182,
+      "image_masks_loss": 0.0112,
+      "llm_learning_rate": 6.849057161354006e-08,
+      "loss": 0.7389,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 6.849057161354006e-06,
+      "step": 9053
+    },
+    {
+      "epoch": 0.93,
+      "gumbel_tau": 0.1182,
+      "image_masks_loss": 0.0078,
+      "llm_learning_rate": 6.829784520452548e-08,
+      "loss": 0.7765,
+      "output_text_mask_loss": 0.0216,
+      "predictor_learning_rate": 6.829784520452548e-06,
+      "step": 9054
+    },
+    {
+      "epoch": 0.93,
+      "gumbel_tau": 0.1181,
+      "image_masks_loss": 0.0103,
+      "llm_learning_rate": 6.810538657996857e-08,
+      "loss": 0.739,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 6.810538657996856e-06,
+      "step": 9055
+    },
+    {
+      "epoch": 0.93,
+      "gumbel_tau": 0.1181,
+      "image_masks_loss": 0.0078,
+      "llm_learning_rate": 6.791319576106293e-08,
+      "loss": 0.7238,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 6.791319576106292e-06,
+      "step": 9056
+    },
+    {
+      "epoch": 0.93,
+      "gumbel_tau": 0.1181,
+      "image_masks_loss": 0.0139,
+      "llm_learning_rate": 6.772127276897383e-08,
+      "loss": 0.7779,
+      "output_text_mask_loss": 0.0162,
+      "predictor_learning_rate": 6.772127276897383e-06,
+      "step": 9057
+    },
+    {
+      "epoch": 0.93,
+      "gumbel_tau": 0.118,
+      "image_masks_loss": 0.0105,
+      "llm_learning_rate": 6.752961762483606e-08,
+      "loss": 0.7557,
+      "output_text_mask_loss": 0.0193,
+      "predictor_learning_rate": 6.752961762483606e-06,
+      "step": 9058
+    },
+    {
+      "epoch": 0.93,
+      "gumbel_tau": 0.118,
+      "image_masks_loss": 0.0106,
+      "llm_learning_rate": 6.733823034975468e-08,
+      "loss": 0.6367,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 6.733823034975467e-06,
+      "step": 9059
+    },
+    {
+      "epoch": 0.93,
+      "gumbel_tau": 0.118,
+      "image_masks_loss": 0.0079,
+      "llm_learning_rate": 6.714711096480587e-08,
+      "loss": 0.7314,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 6.714711096480586e-06,
+      "step": 9060
+    },
+    {
+      "epoch": 0.93,
+      "gumbel_tau": 0.118,
+      "image_masks_loss": 0.0089,
+      "llm_learning_rate": 6.695625949103612e-08,
+      "loss": 0.6151,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 6.695625949103612e-06,
+      "step": 9061
+    },
+    {
+      "epoch": 0.93,
+      "gumbel_tau": 0.1179,
+      "image_masks_loss": 0.0134,
+      "llm_learning_rate": 6.676567594946254e-08,
+      "loss": 0.6742,
+      "output_text_mask_loss": 0.0151,
+      "predictor_learning_rate": 6.676567594946253e-06,
+      "step": 9062
+    },
+    {
+      "epoch": 0.93,
+      "gumbel_tau": 0.1179,
+      "image_masks_loss": 0.0103,
+      "llm_learning_rate": 6.657536036107248e-08,
+      "loss": 0.6795,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 6.657536036107248e-06,
+      "step": 9063
+    },
+    {
+      "epoch": 0.93,
+      "gumbel_tau": 0.1179,
+      "image_masks_loss": 0.0149,
+      "llm_learning_rate": 6.638531274682391e-08,
+      "loss": 0.7219,
+      "output_text_mask_loss": 0.0231,
+      "predictor_learning_rate": 6.6385312746823915e-06,
+      "step": 9064
+    },
+    {
+      "epoch": 0.93,
+      "gumbel_tau": 0.1178,
+      "image_masks_loss": 0.0047,
+      "llm_learning_rate": 6.619553312764565e-08,
+      "loss": 0.7312,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 6.619553312764564e-06,
+      "step": 9065
+    },
+    {
+      "epoch": 0.93,
+      "gumbel_tau": 0.1178,
+      "image_masks_loss": 0.0062,
+      "llm_learning_rate": 6.600602152443653e-08,
+      "loss": 0.6687,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 6.600602152443652e-06,
+      "step": 9066
+    },
+    {
+      "epoch": 0.93,
+      "gumbel_tau": 0.1178,
+      "image_masks_loss": 0.0213,
+      "llm_learning_rate": 6.58167779580654e-08,
+      "loss": 0.7372,
+      "output_text_mask_loss": 0.0233,
+      "predictor_learning_rate": 6.58167779580654e-06,
+      "step": 9067
+    },
+    {
+      "epoch": 0.93,
+      "gumbel_tau": 0.1178,
+      "image_masks_loss": 0.0213,
+      "llm_learning_rate": 6.562780244937284e-08,
+      "loss": 0.6767,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 6.5627802449372835e-06,
+      "step": 9068
+    },
+    {
+      "epoch": 0.93,
+      "gumbel_tau": 0.1177,
+      "image_masks_loss": 0.0131,
+      "llm_learning_rate": 6.543909501916856e-08,
+      "loss": 0.6717,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 6.543909501916856e-06,
+      "step": 9069
+    },
+    {
+      "epoch": 0.93,
+      "gumbel_tau": 0.1177,
+      "image_masks_loss": 0.0115,
+      "llm_learning_rate": 6.525065568823429e-08,
+      "loss": 0.8235,
+      "output_text_mask_loss": 0.0208,
+      "predictor_learning_rate": 6.5250655688234296e-06,
+      "step": 9070
+    },
+    {
+      "epoch": 0.93,
+      "gumbel_tau": 0.1177,
+      "image_masks_loss": 0.0129,
+      "llm_learning_rate": 6.506248447732095e-08,
+      "loss": 0.7845,
+      "output_text_mask_loss": 0.0229,
+      "predictor_learning_rate": 6.506248447732094e-06,
+      "step": 9071
+    },
+    {
+      "epoch": 0.93,
+      "gumbel_tau": 0.1177,
+      "image_masks_loss": 0.0124,
+      "llm_learning_rate": 6.487458140715053e-08,
+      "loss": 0.7427,
+      "output_text_mask_loss": 0.0065,
+      "predictor_learning_rate": 6.487458140715052e-06,
+      "step": 9072
+    },
+    {
+      "epoch": 0.93,
+      "gumbel_tau": 0.1176,
+      "image_masks_loss": 0.0075,
+      "llm_learning_rate": 6.468694649841539e-08,
+      "loss": 0.6919,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 6.468694649841539e-06,
+      "step": 9073
+    },
+    {
+      "epoch": 0.93,
+      "gumbel_tau": 0.1176,
+      "image_masks_loss": 0.0101,
+      "llm_learning_rate": 6.44995797717779e-08,
+      "loss": 0.7036,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 6.449957977177789e-06,
+      "step": 9074
+    },
+    {
+      "epoch": 0.93,
+      "gumbel_tau": 0.1176,
+      "image_masks_loss": 0.0049,
+      "llm_learning_rate": 6.431248124787154e-08,
+      "loss": 0.6943,
+      "output_text_mask_loss": 0.0141,
+      "predictor_learning_rate": 6.431248124787154e-06,
+      "step": 9075
+    },
+    {
+      "epoch": 0.93,
+      "gumbel_tau": 0.1175,
+      "image_masks_loss": 0.0195,
+      "llm_learning_rate": 6.412565094729984e-08,
+      "loss": 0.7596,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 6.412565094729983e-06,
+      "step": 9076
+    },
+    {
+      "epoch": 0.93,
+      "gumbel_tau": 0.1175,
+      "image_masks_loss": 0.0204,
+      "llm_learning_rate": 6.393908889063744e-08,
+      "loss": 0.5465,
+      "output_text_mask_loss": 0.0055,
+      "predictor_learning_rate": 6.393908889063743e-06,
+      "step": 9077
+    },
+    {
+      "epoch": 0.93,
+      "gumbel_tau": 0.1175,
+      "image_masks_loss": 0.0307,
+      "llm_learning_rate": 6.375279509842847e-08,
+      "loss": 0.7336,
+      "output_text_mask_loss": 0.0261,
+      "predictor_learning_rate": 6.375279509842846e-06,
+      "step": 9078
+    },
+    {
+      "epoch": 0.93,
+      "gumbel_tau": 0.1175,
+      "image_masks_loss": 0.0045,
+      "llm_learning_rate": 6.356676959118846e-08,
+      "loss": 0.7602,
+      "output_text_mask_loss": 0.0488,
+      "predictor_learning_rate": 6.356676959118846e-06,
+      "step": 9079
+    },
+    {
+      "epoch": 0.93,
+      "gumbel_tau": 0.1174,
+      "image_masks_loss": 0.0041,
+      "llm_learning_rate": 6.338101238940269e-08,
+      "loss": 0.7329,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 6.338101238940269e-06,
+      "step": 9080
+    },
+    {
+      "epoch": 0.93,
+      "gumbel_tau": 0.1174,
+      "image_masks_loss": 0.0072,
+      "llm_learning_rate": 6.319552351352731e-08,
+      "loss": 0.7585,
+      "output_text_mask_loss": 0.0041,
+      "predictor_learning_rate": 6.319552351352731e-06,
+      "step": 9081
+    },
+    {
+      "epoch": 0.93,
+      "gumbel_tau": 0.1174,
+      "image_masks_loss": 0.028,
+      "llm_learning_rate": 6.301030298398846e-08,
+      "loss": 0.7045,
+      "output_text_mask_loss": 0.0376,
+      "predictor_learning_rate": 6.3010302983988465e-06,
+      "step": 9082
+    },
+    {
+      "epoch": 0.93,
+      "gumbel_tau": 0.1173,
+      "image_masks_loss": 0.0076,
+      "llm_learning_rate": 6.282535082118291e-08,
+      "loss": 0.7759,
+      "output_text_mask_loss": 0.0019,
+      "predictor_learning_rate": 6.28253508211829e-06,
+      "step": 9083
+    },
+    {
+      "epoch": 0.93,
+      "gumbel_tau": 0.1173,
+      "image_masks_loss": 0.0103,
+      "llm_learning_rate": 6.264066704547906e-08,
+      "loss": 0.7967,
+      "output_text_mask_loss": 0.0181,
+      "predictor_learning_rate": 6.264066704547905e-06,
+      "step": 9084
+    },
+    {
+      "epoch": 0.93,
+      "gumbel_tau": 0.1173,
+      "image_masks_loss": 0.0107,
+      "llm_learning_rate": 6.245625167721341e-08,
+      "loss": 0.678,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 6.245625167721341e-06,
+      "step": 9085
+    },
+    {
+      "epoch": 0.93,
+      "gumbel_tau": 0.1173,
+      "image_masks_loss": 0.0199,
+      "llm_learning_rate": 6.227210473669504e-08,
+      "loss": 0.7975,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 6.227210473669504e-06,
+      "step": 9086
+    },
+    {
+      "epoch": 0.93,
+      "gumbel_tau": 0.1172,
+      "image_masks_loss": 0.0099,
+      "llm_learning_rate": 6.208822624420213e-08,
+      "loss": 0.7218,
+      "output_text_mask_loss": 0.0016,
+      "predictor_learning_rate": 6.208822624420213e-06,
+      "step": 9087
+    },
+    {
+      "epoch": 0.93,
+      "gumbel_tau": 0.1172,
+      "image_masks_loss": 0.0057,
+      "llm_learning_rate": 6.190461621998406e-08,
+      "loss": 0.6292,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 6.190461621998405e-06,
+      "step": 9088
+    },
+    {
+      "epoch": 0.93,
+      "gumbel_tau": 0.1172,
+      "image_masks_loss": 0.0092,
+      "llm_learning_rate": 6.17212746842602e-08,
+      "loss": 0.7026,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 6.17212746842602e-06,
+      "step": 9089
+    },
+    {
+      "epoch": 0.93,
+      "gumbel_tau": 0.1172,
+      "image_masks_loss": 0.0181,
+      "llm_learning_rate": 6.153820165722052e-08,
+      "loss": 0.7175,
+      "output_text_mask_loss": 0.0294,
+      "predictor_learning_rate": 6.153820165722052e-06,
+      "step": 9090
+    },
+    {
+      "epoch": 0.93,
+      "gumbel_tau": 0.1171,
+      "image_masks_loss": 0.0292,
+      "llm_learning_rate": 6.135539715902584e-08,
+      "loss": 0.7378,
+      "output_text_mask_loss": 0.0116,
+      "predictor_learning_rate": 6.135539715902582e-06,
+      "step": 9091
+    },
+    {
+      "epoch": 0.93,
+      "gumbel_tau": 0.1171,
+      "image_masks_loss": 0.0075,
+      "llm_learning_rate": 6.11728612098067e-08,
+      "loss": 0.8,
+      "output_text_mask_loss": 0.0243,
+      "predictor_learning_rate": 6.11728612098067e-06,
+      "step": 9092
+    },
+    {
+      "epoch": 0.93,
+      "gumbel_tau": 0.1171,
+      "image_masks_loss": 0.0036,
+      "llm_learning_rate": 6.099059382966399e-08,
+      "loss": 0.7141,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 6.0990593829663985e-06,
+      "step": 9093
+    },
+    {
+      "epoch": 0.93,
+      "gumbel_tau": 0.117,
+      "image_masks_loss": 0.0185,
+      "llm_learning_rate": 6.080859503867026e-08,
+      "loss": 0.7929,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 6.080859503867025e-06,
+      "step": 9094
+    },
+    {
+      "epoch": 0.93,
+      "gumbel_tau": 0.117,
+      "image_masks_loss": 0.0052,
+      "llm_learning_rate": 6.06268648568667e-08,
+      "loss": 0.6395,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 6.062686485686669e-06,
+      "step": 9095
+    },
+    {
+      "epoch": 0.93,
+      "gumbel_tau": 0.117,
+      "image_masks_loss": 0.0128,
+      "llm_learning_rate": 6.044540330426617e-08,
+      "loss": 0.7277,
+      "output_text_mask_loss": 0.0179,
+      "predictor_learning_rate": 6.044540330426618e-06,
+      "step": 9096
+    },
+    {
+      "epoch": 0.93,
+      "gumbel_tau": 0.117,
+      "image_masks_loss": 0.0074,
+      "llm_learning_rate": 6.026421040085217e-08,
+      "loss": 0.7535,
+      "output_text_mask_loss": 0.0208,
+      "predictor_learning_rate": 6.026421040085217e-06,
+      "step": 9097
+    },
+    {
+      "epoch": 0.93,
+      "gumbel_tau": 0.1169,
+      "image_masks_loss": 0.0052,
+      "llm_learning_rate": 6.008328616657733e-08,
+      "loss": 0.7106,
+      "output_text_mask_loss": 0.0137,
+      "predictor_learning_rate": 6.008328616657732e-06,
+      "step": 9098
+    },
+    {
+      "epoch": 0.93,
+      "gumbel_tau": 0.1169,
+      "image_masks_loss": 0.0114,
+      "llm_learning_rate": 5.990263062136598e-08,
+      "loss": 0.6319,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 5.990263062136597e-06,
+      "step": 9099
+    },
+    {
+      "epoch": 0.93,
+      "gumbel_tau": 0.1169,
+      "image_masks_loss": 0.0046,
+      "llm_learning_rate": 5.972224378511193e-08,
+      "loss": 0.6641,
+      "output_text_mask_loss": 0.0035,
+      "predictor_learning_rate": 5.972224378511193e-06,
+      "step": 9100
+    },
+    {
+      "epoch": 0.93,
+      "gumbel_tau": 0.1168,
+      "image_masks_loss": 0.0029,
+      "llm_learning_rate": 5.954212567767986e-08,
+      "loss": 0.7259,
+      "output_text_mask_loss": 0.0277,
+      "predictor_learning_rate": 5.954212567767986e-06,
+      "step": 9101
+    },
+    {
+      "epoch": 0.93,
+      "gumbel_tau": 0.1168,
+      "image_masks_loss": 0.0036,
+      "llm_learning_rate": 5.936227631890501e-08,
+      "loss": 0.7935,
+      "output_text_mask_loss": 0.0132,
+      "predictor_learning_rate": 5.936227631890501e-06,
+      "step": 9102
+    },
+    {
+      "epoch": 0.93,
+      "gumbel_tau": 0.1168,
+      "image_masks_loss": 0.0069,
+      "llm_learning_rate": 5.9182695728592363e-08,
+      "loss": 0.7564,
+      "output_text_mask_loss": 0.0073,
+      "predictor_learning_rate": 5.918269572859236e-06,
+      "step": 9103
+    },
+    {
+      "epoch": 0.93,
+      "gumbel_tau": 0.1168,
+      "image_masks_loss": 0.0139,
+      "llm_learning_rate": 5.900338392651805e-08,
+      "loss": 0.8102,
+      "output_text_mask_loss": 0.0271,
+      "predictor_learning_rate": 5.900338392651805e-06,
+      "step": 9104
+    },
+    {
+      "epoch": 0.93,
+      "gumbel_tau": 0.1167,
+      "image_masks_loss": 0.0047,
+      "llm_learning_rate": 5.882434093242823e-08,
+      "loss": 0.8192,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 5.882434093242822e-06,
+      "step": 9105
+    },
+    {
+      "epoch": 0.93,
+      "gumbel_tau": 0.1167,
+      "image_masks_loss": 0.0071,
+      "llm_learning_rate": 5.864556676603933e-08,
+      "loss": 0.7277,
+      "output_text_mask_loss": 0.0345,
+      "predictor_learning_rate": 5.8645566766039325e-06,
+      "step": 9106
+    },
+    {
+      "epoch": 0.93,
+      "gumbel_tau": 0.1167,
+      "image_masks_loss": 0.0043,
+      "llm_learning_rate": 5.846706144703868e-08,
+      "loss": 0.7496,
+      "output_text_mask_loss": 0.0164,
+      "predictor_learning_rate": 5.846706144703867e-06,
+      "step": 9107
+    },
+    {
+      "epoch": 0.93,
+      "gumbel_tau": 0.1167,
+      "image_masks_loss": 0.0228,
+      "llm_learning_rate": 5.828882499508332e-08,
+      "loss": 0.6322,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 5.828882499508331e-06,
+      "step": 9108
+    },
+    {
+      "epoch": 0.93,
+      "gumbel_tau": 0.1166,
+      "image_masks_loss": 0.0141,
+      "llm_learning_rate": 5.811085742980144e-08,
+      "loss": 0.7483,
+      "output_text_mask_loss": 0.0633,
+      "predictor_learning_rate": 5.811085742980144e-06,
+      "step": 9109
+    },
+    {
+      "epoch": 0.93,
+      "gumbel_tau": 0.1166,
+      "image_masks_loss": 0.0092,
+      "llm_learning_rate": 5.793315877079097e-08,
+      "loss": 0.7024,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 5.793315877079098e-06,
+      "step": 9110
+    },
+    {
+      "epoch": 0.93,
+      "gumbel_tau": 0.1166,
+      "image_masks_loss": 0.0185,
+      "llm_learning_rate": 5.7755729037620724e-08,
+      "loss": 0.6963,
+      "output_text_mask_loss": 0.0019,
+      "predictor_learning_rate": 5.775572903762072e-06,
+      "step": 9111
+    },
+    {
+      "epoch": 0.93,
+      "gumbel_tau": 0.1165,
+      "image_masks_loss": 0.0078,
+      "llm_learning_rate": 5.7578568249829215e-08,
+      "loss": 0.7156,
+      "output_text_mask_loss": 0.0374,
+      "predictor_learning_rate": 5.757856824982921e-06,
+      "step": 9112
+    },
+    {
+      "epoch": 0.93,
+      "gumbel_tau": 0.1165,
+      "image_masks_loss": 0.009,
+      "llm_learning_rate": 5.74016764269264e-08,
+      "loss": 0.7459,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 5.74016764269264e-06,
+      "step": 9113
+    },
+    {
+      "epoch": 0.93,
+      "gumbel_tau": 0.1165,
+      "image_masks_loss": 0.0094,
+      "llm_learning_rate": 5.722505358839142e-08,
+      "loss": 0.7653,
+      "output_text_mask_loss": 0.0019,
+      "predictor_learning_rate": 5.722505358839142e-06,
+      "step": 9114
+    },
+    {
+      "epoch": 0.93,
+      "gumbel_tau": 0.1165,
+      "image_masks_loss": 0.0132,
+      "llm_learning_rate": 5.7048699753674276e-08,
+      "loss": 0.6556,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 5.704869975367427e-06,
+      "step": 9115
+    },
+    {
+      "epoch": 0.93,
+      "gumbel_tau": 0.1164,
+      "image_masks_loss": 0.0109,
+      "llm_learning_rate": 5.687261494219609e-08,
+      "loss": 0.7194,
+      "output_text_mask_loss": 0.0284,
+      "predictor_learning_rate": 5.6872614942196086e-06,
+      "step": 9116
+    },
+    {
+      "epoch": 0.93,
+      "gumbel_tau": 0.1164,
+      "image_masks_loss": 0.0263,
+      "llm_learning_rate": 5.669679917334747e-08,
+      "loss": 0.6553,
+      "output_text_mask_loss": 0.0275,
+      "predictor_learning_rate": 5.6696799173347466e-06,
+      "step": 9117
+    },
+    {
+      "epoch": 0.93,
+      "gumbel_tau": 0.1164,
+      "image_masks_loss": 0.0054,
+      "llm_learning_rate": 5.652125246648959e-08,
+      "loss": 0.7914,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 5.652125246648959e-06,
+      "step": 9118
+    },
+    {
+      "epoch": 0.93,
+      "gumbel_tau": 0.1164,
+      "image_masks_loss": 0.012,
+      "llm_learning_rate": 5.634597484095394e-08,
+      "loss": 0.7317,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 5.634597484095394e-06,
+      "step": 9119
+    },
+    {
+      "epoch": 0.93,
+      "gumbel_tau": 0.1163,
+      "image_masks_loss": 0.0046,
+      "llm_learning_rate": 5.6170966316042574e-08,
+      "loss": 0.7306,
+      "output_text_mask_loss": 0.0071,
+      "predictor_learning_rate": 5.617096631604257e-06,
+      "step": 9120
+    },
+    {
+      "epoch": 0.93,
+      "gumbel_tau": 0.1163,
+      "image_masks_loss": 0.0042,
+      "llm_learning_rate": 5.599622691102785e-08,
+      "loss": 0.7813,
+      "output_text_mask_loss": 0.0109,
+      "predictor_learning_rate": 5.599622691102785e-06,
+      "step": 9121
+    },
+    {
+      "epoch": 0.93,
+      "gumbel_tau": 0.1163,
+      "image_masks_loss": 0.0142,
+      "llm_learning_rate": 5.5821756645152435e-08,
+      "loss": 0.7552,
+      "output_text_mask_loss": 0.0025,
+      "predictor_learning_rate": 5.582175664515243e-06,
+      "step": 9122
+    },
+    {
+      "epoch": 0.93,
+      "gumbel_tau": 0.1162,
+      "image_masks_loss": 0.0021,
+      "llm_learning_rate": 5.564755553762929e-08,
+      "loss": 0.6987,
+      "output_text_mask_loss": 0.0055,
+      "predictor_learning_rate": 5.564755553762929e-06,
+      "step": 9123
+    },
+    {
+      "epoch": 0.93,
+      "gumbel_tau": 0.1162,
+      "image_masks_loss": 0.0057,
+      "llm_learning_rate": 5.547362360764197e-08,
+      "loss": 0.775,
+      "output_text_mask_loss": 0.0016,
+      "predictor_learning_rate": 5.547362360764197e-06,
+      "step": 9124
+    },
+    {
+      "epoch": 0.93,
+      "gumbel_tau": 0.1162,
+      "image_masks_loss": 0.0055,
+      "llm_learning_rate": 5.5299960874344306e-08,
+      "loss": 0.7563,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 5.52999608743443e-06,
+      "step": 9125
+    },
+    {
+      "epoch": 0.94,
+      "gumbel_tau": 0.1162,
+      "image_masks_loss": 0.0083,
+      "llm_learning_rate": 5.5126567356860174e-08,
+      "loss": 0.7359,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 5.512656735686017e-06,
+      "step": 9126
+    },
+    {
+      "epoch": 0.94,
+      "gumbel_tau": 0.1161,
+      "image_masks_loss": 0.0113,
+      "llm_learning_rate": 5.4953443074284295e-08,
+      "loss": 0.7028,
+      "output_text_mask_loss": 0.0599,
+      "predictor_learning_rate": 5.495344307428429e-06,
+      "step": 9127
+    },
+    {
+      "epoch": 0.94,
+      "gumbel_tau": 0.1161,
+      "image_masks_loss": 0.0023,
+      "llm_learning_rate": 5.478058804568115e-08,
+      "loss": 0.7715,
+      "output_text_mask_loss": 0.001,
+      "predictor_learning_rate": 5.4780588045681145e-06,
+      "step": 9128
+    },
+    {
+      "epoch": 0.94,
+      "gumbel_tau": 0.1161,
+      "image_masks_loss": 0.019,
+      "llm_learning_rate": 5.460800229008634e-08,
+      "loss": 0.7815,
+      "output_text_mask_loss": 0.0353,
+      "predictor_learning_rate": 5.460800229008633e-06,
+      "step": 9129
+    },
+    {
+      "epoch": 0.94,
+      "gumbel_tau": 0.1161,
+      "image_masks_loss": 0.0171,
+      "llm_learning_rate": 5.443568582650494e-08,
+      "loss": 0.6459,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 5.443568582650493e-06,
+      "step": 9130
+    },
+    {
+      "epoch": 0.94,
+      "gumbel_tau": 0.116,
+      "image_masks_loss": 0.0144,
+      "llm_learning_rate": 5.4263638673913434e-08,
+      "loss": 0.703,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 5.4263638673913425e-06,
+      "step": 9131
+    },
+    {
+      "epoch": 0.94,
+      "gumbel_tau": 0.116,
+      "image_masks_loss": 0.0175,
+      "llm_learning_rate": 5.409186085125778e-08,
+      "loss": 0.7991,
+      "output_text_mask_loss": 0.0504,
+      "predictor_learning_rate": 5.409186085125778e-06,
+      "step": 9132
+    },
+    {
+      "epoch": 0.94,
+      "gumbel_tau": 0.116,
+      "image_masks_loss": 0.02,
+      "llm_learning_rate": 5.392035237745424e-08,
+      "loss": 0.645,
+      "output_text_mask_loss": 0.0174,
+      "predictor_learning_rate": 5.392035237745424e-06,
+      "step": 9133
+    },
+    {
+      "epoch": 0.94,
+      "gumbel_tau": 0.1159,
+      "image_masks_loss": 0.0056,
+      "llm_learning_rate": 5.374911327138993e-08,
+      "loss": 0.6562,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 5.374911327138992e-06,
+      "step": 9134
+    },
+    {
+      "epoch": 0.94,
+      "gumbel_tau": 0.1159,
+      "image_masks_loss": 0.0102,
+      "llm_learning_rate": 5.357814355192198e-08,
+      "loss": 0.6904,
+      "output_text_mask_loss": 0.032,
+      "predictor_learning_rate": 5.357814355192198e-06,
+      "step": 9135
+    },
+    {
+      "epoch": 0.94,
+      "gumbel_tau": 0.1159,
+      "image_masks_loss": 0.006,
+      "llm_learning_rate": 5.340744323787811e-08,
+      "loss": 0.6827,
+      "output_text_mask_loss": 0.0572,
+      "predictor_learning_rate": 5.3407443237878104e-06,
+      "step": 9136
+    },
+    {
+      "epoch": 0.94,
+      "gumbel_tau": 0.1159,
+      "image_masks_loss": 0.0198,
+      "llm_learning_rate": 5.323701234805634e-08,
+      "loss": 0.7637,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 5.323701234805634e-06,
+      "step": 9137
+    },
+    {
+      "epoch": 0.94,
+      "gumbel_tau": 0.1158,
+      "image_masks_loss": 0.0071,
+      "llm_learning_rate": 5.3066850901224434e-08,
+      "loss": 0.6789,
+      "output_text_mask_loss": 0.0095,
+      "predictor_learning_rate": 5.306685090122443e-06,
+      "step": 9138
+    },
+    {
+      "epoch": 0.94,
+      "gumbel_tau": 0.1158,
+      "image_masks_loss": 0.0041,
+      "llm_learning_rate": 5.289695891612157e-08,
+      "loss": 0.7185,
+      "output_text_mask_loss": 0.024,
+      "predictor_learning_rate": 5.289695891612157e-06,
+      "step": 9139
+    },
+    {
+      "epoch": 0.94,
+      "gumbel_tau": 0.1158,
+      "image_masks_loss": 0.0079,
+      "llm_learning_rate": 5.272733641145611e-08,
+      "loss": 0.6208,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 5.272733641145611e-06,
+      "step": 9140
+    },
+    {
+      "epoch": 0.94,
+      "gumbel_tau": 0.1158,
+      "image_masks_loss": 0.0034,
+      "llm_learning_rate": 5.255798340590756e-08,
+      "loss": 0.7268,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 5.255798340590756e-06,
+      "step": 9141
+    },
+    {
+      "epoch": 0.94,
+      "gumbel_tau": 0.1157,
+      "image_masks_loss": 0.0074,
+      "llm_learning_rate": 5.238889991812518e-08,
+      "loss": 0.7307,
+      "output_text_mask_loss": 0.0193,
+      "predictor_learning_rate": 5.238889991812518e-06,
+      "step": 9142
+    },
+    {
+      "epoch": 0.94,
+      "gumbel_tau": 0.1157,
+      "image_masks_loss": 0.0078,
+      "llm_learning_rate": 5.2220085966729063e-08,
+      "loss": 0.7547,
+      "output_text_mask_loss": 0.0072,
+      "predictor_learning_rate": 5.222008596672906e-06,
+      "step": 9143
+    },
+    {
+      "epoch": 0.94,
+      "gumbel_tau": 0.1157,
+      "image_masks_loss": 0.0103,
+      "llm_learning_rate": 5.205154157030934e-08,
+      "loss": 0.671,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 5.205154157030934e-06,
+      "step": 9144
+    },
+    {
+      "epoch": 0.94,
+      "gumbel_tau": 0.1156,
+      "image_masks_loss": 0.0089,
+      "llm_learning_rate": 5.188326674742644e-08,
+      "loss": 0.5764,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 5.188326674742644e-06,
+      "step": 9145
+    },
+    {
+      "epoch": 0.94,
+      "gumbel_tau": 0.1156,
+      "image_masks_loss": 0.0055,
+      "llm_learning_rate": 5.17152615166111e-08,
+      "loss": 0.6811,
+      "output_text_mask_loss": 0.0065,
+      "predictor_learning_rate": 5.17152615166111e-06,
+      "step": 9146
+    },
+    {
+      "epoch": 0.94,
+      "gumbel_tau": 0.1156,
+      "image_masks_loss": 0.0068,
+      "llm_learning_rate": 5.154752589636436e-08,
+      "loss": 0.7599,
+      "output_text_mask_loss": 0.0433,
+      "predictor_learning_rate": 5.154752589636435e-06,
+      "step": 9147
+    },
+    {
+      "epoch": 0.94,
+      "gumbel_tau": 0.1156,
+      "image_masks_loss": 0.0112,
+      "llm_learning_rate": 5.13800599051581e-08,
+      "loss": 0.7866,
+      "output_text_mask_loss": 0.0797,
+      "predictor_learning_rate": 5.13800599051581e-06,
+      "step": 9148
+    },
+    {
+      "epoch": 0.94,
+      "gumbel_tau": 0.1155,
+      "image_masks_loss": 0.0105,
+      "llm_learning_rate": 5.1212863561433676e-08,
+      "loss": 0.7738,
+      "output_text_mask_loss": 0.0275,
+      "predictor_learning_rate": 5.121286356143368e-06,
+      "step": 9149
+    },
+    {
+      "epoch": 0.94,
+      "gumbel_tau": 0.1155,
+      "image_masks_loss": 0.0182,
+      "llm_learning_rate": 5.1045936883603044e-08,
+      "loss": 0.6689,
+      "output_text_mask_loss": 0.0065,
+      "predictor_learning_rate": 5.104593688360304e-06,
+      "step": 9150
+    },
+    {
+      "epoch": 0.94,
+      "gumbel_tau": 0.1155,
+      "image_masks_loss": 0.0133,
+      "llm_learning_rate": 5.08792798900487e-08,
+      "loss": 0.7828,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 5.08792798900487e-06,
+      "step": 9151
+    },
+    {
+      "epoch": 0.94,
+      "gumbel_tau": 0.1155,
+      "image_masks_loss": 0.0074,
+      "llm_learning_rate": 5.071289259912349e-08,
+      "loss": 0.6977,
+      "output_text_mask_loss": 0.0156,
+      "predictor_learning_rate": 5.071289259912349e-06,
+      "step": 9152
+    },
+    {
+      "epoch": 0.94,
+      "gumbel_tau": 0.1154,
+      "image_masks_loss": 0.0022,
+      "llm_learning_rate": 5.054677502914995e-08,
+      "loss": 0.814,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 5.054677502914995e-06,
+      "step": 9153
+    },
+    {
+      "epoch": 0.94,
+      "gumbel_tau": 0.1154,
+      "image_masks_loss": 0.0099,
+      "llm_learning_rate": 5.0380927198421246e-08,
+      "loss": 0.6556,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 5.038092719842125e-06,
+      "step": 9154
+    },
+    {
+      "epoch": 0.94,
+      "gumbel_tau": 0.1154,
+      "image_masks_loss": 0.0074,
+      "llm_learning_rate": 5.0215349125201375e-08,
+      "loss": 0.7181,
+      "output_text_mask_loss": 0.0035,
+      "predictor_learning_rate": 5.021534912520137e-06,
+      "step": 9155
+    },
+    {
+      "epoch": 0.94,
+      "gumbel_tau": 0.1153,
+      "image_masks_loss": 0.0108,
+      "llm_learning_rate": 5.00500408277238e-08,
+      "loss": 0.7284,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 5.005004082772379e-06,
+      "step": 9156
+    },
+    {
+      "epoch": 0.94,
+      "gumbel_tau": 0.1153,
+      "image_masks_loss": 0.0044,
+      "llm_learning_rate": 4.9885002324192855e-08,
+      "loss": 0.7449,
+      "output_text_mask_loss": 0.0116,
+      "predictor_learning_rate": 4.988500232419285e-06,
+      "step": 9157
+    },
+    {
+      "epoch": 0.94,
+      "gumbel_tau": 0.1153,
+      "image_masks_loss": 0.0161,
+      "llm_learning_rate": 4.9720233632783166e-08,
+      "loss": 0.8,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 4.9720233632783165e-06,
+      "step": 9158
+    },
+    {
+      "epoch": 0.94,
+      "gumbel_tau": 0.1153,
+      "image_masks_loss": 0.0013,
+      "llm_learning_rate": 4.9555734771639106e-08,
+      "loss": 0.7011,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 4.955573477163911e-06,
+      "step": 9159
+    },
+    {
+      "epoch": 0.94,
+      "gumbel_tau": 0.1152,
+      "image_masks_loss": 0.0068,
+      "llm_learning_rate": 4.939150575887508e-08,
+      "loss": 0.7411,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 4.939150575887508e-06,
+      "step": 9160
+    },
+    {
+      "epoch": 0.94,
+      "gumbel_tau": 0.1152,
+      "image_masks_loss": 0.0078,
+      "llm_learning_rate": 4.922754661257745e-08,
+      "loss": 0.7153,
+      "output_text_mask_loss": 0.0416,
+      "predictor_learning_rate": 4.922754661257745e-06,
+      "step": 9161
+    },
+    {
+      "epoch": 0.94,
+      "gumbel_tau": 0.1152,
+      "image_masks_loss": 0.0156,
+      "llm_learning_rate": 4.9063857350801216e-08,
+      "loss": 0.7399,
+      "output_text_mask_loss": 0.0056,
+      "predictor_learning_rate": 4.9063857350801215e-06,
+      "step": 9162
+    },
+    {
+      "epoch": 0.94,
+      "gumbel_tau": 0.1152,
+      "image_masks_loss": 0.0239,
+      "llm_learning_rate": 4.8900437991572236e-08,
+      "loss": 0.7195,
+      "output_text_mask_loss": 0.0076,
+      "predictor_learning_rate": 4.890043799157224e-06,
+      "step": 9163
+    },
+    {
+      "epoch": 0.94,
+      "gumbel_tau": 0.1151,
+      "image_masks_loss": 0.0167,
+      "llm_learning_rate": 4.87372885528864e-08,
+      "loss": 0.7062,
+      "output_text_mask_loss": 0.0153,
+      "predictor_learning_rate": 4.87372885528864e-06,
+      "step": 9164
+    },
+    {
+      "epoch": 0.94,
+      "gumbel_tau": 0.1151,
+      "image_masks_loss": 0.0046,
+      "llm_learning_rate": 4.857440905271071e-08,
+      "loss": 0.7286,
+      "output_text_mask_loss": 0.058,
+      "predictor_learning_rate": 4.857440905271071e-06,
+      "step": 9165
+    },
+    {
+      "epoch": 0.94,
+      "gumbel_tau": 0.1151,
+      "image_masks_loss": 0.0121,
+      "llm_learning_rate": 4.8411799508981114e-08,
+      "loss": 0.7918,
+      "output_text_mask_loss": 0.0095,
+      "predictor_learning_rate": 4.841179950898111e-06,
+      "step": 9166
+    },
+    {
+      "epoch": 0.94,
+      "gumbel_tau": 0.115,
+      "image_masks_loss": 0.0121,
+      "llm_learning_rate": 4.824945993960495e-08,
+      "loss": 0.7014,
+      "output_text_mask_loss": 0.0052,
+      "predictor_learning_rate": 4.824945993960494e-06,
+      "step": 9167
+    },
+    {
+      "epoch": 0.94,
+      "gumbel_tau": 0.115,
+      "image_masks_loss": 0.0183,
+      "llm_learning_rate": 4.80873903624593e-08,
+      "loss": 0.7425,
+      "output_text_mask_loss": 0.0277,
+      "predictor_learning_rate": 4.808739036245929e-06,
+      "step": 9168
+    },
+    {
+      "epoch": 0.94,
+      "gumbel_tau": 0.115,
+      "image_masks_loss": 0.0056,
+      "llm_learning_rate": 4.792559079539183e-08,
+      "loss": 0.7945,
+      "output_text_mask_loss": 0.0124,
+      "predictor_learning_rate": 4.792559079539183e-06,
+      "step": 9169
+    },
+    {
+      "epoch": 0.94,
+      "gumbel_tau": 0.115,
+      "image_masks_loss": 0.0101,
+      "llm_learning_rate": 4.776406125621996e-08,
+      "loss": 0.697,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 4.7764061256219956e-06,
+      "step": 9170
+    },
+    {
+      "epoch": 0.94,
+      "gumbel_tau": 0.1149,
+      "image_masks_loss": 0.0178,
+      "llm_learning_rate": 4.760280176273169e-08,
+      "loss": 0.6644,
+      "output_text_mask_loss": 0.0156,
+      "predictor_learning_rate": 4.760280176273168e-06,
+      "step": 9171
+    },
+    {
+      "epoch": 0.94,
+      "gumbel_tau": 0.1149,
+      "image_masks_loss": 0.0074,
+      "llm_learning_rate": 4.744181233268558e-08,
+      "loss": 0.6734,
+      "output_text_mask_loss": 0.0081,
+      "predictor_learning_rate": 4.744181233268557e-06,
+      "step": 9172
+    },
+    {
+      "epoch": 0.94,
+      "gumbel_tau": 0.1149,
+      "image_masks_loss": 0.0165,
+      "llm_learning_rate": 4.728109298380995e-08,
+      "loss": 0.7534,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 4.728109298380995e-06,
+      "step": 9173
+    },
+    {
+      "epoch": 0.94,
+      "gumbel_tau": 0.1149,
+      "image_masks_loss": 0.0065,
+      "llm_learning_rate": 4.712064373380315e-08,
+      "loss": 0.7953,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 4.712064373380315e-06,
+      "step": 9174
+    },
+    {
+      "epoch": 0.94,
+      "gumbel_tau": 0.1148,
+      "image_masks_loss": 0.0065,
+      "llm_learning_rate": 4.696046460033521e-08,
+      "loss": 0.8046,
+      "output_text_mask_loss": 0.0023,
+      "predictor_learning_rate": 4.6960464600335205e-06,
+      "step": 9175
+    },
+    {
+      "epoch": 0.94,
+      "gumbel_tau": 0.1148,
+      "image_masks_loss": 0.0096,
+      "llm_learning_rate": 4.680055560104452e-08,
+      "loss": 0.731,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 4.680055560104451e-06,
+      "step": 9176
+    },
+    {
+      "epoch": 0.94,
+      "gumbel_tau": 0.1148,
+      "image_masks_loss": 0.0034,
+      "llm_learning_rate": 4.6640916753541157e-08,
+      "loss": 0.5997,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 4.664091675354115e-06,
+      "step": 9177
+    },
+    {
+      "epoch": 0.94,
+      "gumbel_tau": 0.1147,
+      "image_masks_loss": 0.0104,
+      "llm_learning_rate": 4.648154807540467e-08,
+      "loss": 0.7627,
+      "output_text_mask_loss": 0.0208,
+      "predictor_learning_rate": 4.648154807540467e-06,
+      "step": 9178
+    },
+    {
+      "epoch": 0.94,
+      "gumbel_tau": 0.1147,
+      "image_masks_loss": 0.0127,
+      "llm_learning_rate": 4.6322449584184916e-08,
+      "loss": 0.7173,
+      "output_text_mask_loss": 0.0011,
+      "predictor_learning_rate": 4.632244958418491e-06,
+      "step": 9179
+    },
+    {
+      "epoch": 0.94,
+      "gumbel_tau": 0.1147,
+      "image_masks_loss": 0.0169,
+      "llm_learning_rate": 4.616362129740232e-08,
+      "loss": 0.6581,
+      "output_text_mask_loss": 0.0073,
+      "predictor_learning_rate": 4.616362129740231e-06,
+      "step": 9180
+    },
+    {
+      "epoch": 0.94,
+      "gumbel_tau": 0.1147,
+      "image_masks_loss": 0.0122,
+      "llm_learning_rate": 4.6005063232547606e-08,
+      "loss": 0.6704,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 4.600506323254761e-06,
+      "step": 9181
+    },
+    {
+      "epoch": 0.94,
+      "gumbel_tau": 0.1146,
+      "image_masks_loss": 0.0095,
+      "llm_learning_rate": 4.5846775407081266e-08,
+      "loss": 0.7582,
+      "output_text_mask_loss": 0.0233,
+      "predictor_learning_rate": 4.584677540708126e-06,
+      "step": 9182
+    },
+    {
+      "epoch": 0.94,
+      "gumbel_tau": 0.1146,
+      "image_masks_loss": 0.0103,
+      "llm_learning_rate": 4.568875783843463e-08,
+      "loss": 0.7659,
+      "output_text_mask_loss": 0.0368,
+      "predictor_learning_rate": 4.568875783843462e-06,
+      "step": 9183
+    },
+    {
+      "epoch": 0.94,
+      "gumbel_tau": 0.1146,
+      "image_masks_loss": 0.014,
+      "llm_learning_rate": 4.553101054400877e-08,
+      "loss": 0.6877,
+      "output_text_mask_loss": 0.0081,
+      "predictor_learning_rate": 4.5531010544008764e-06,
+      "step": 9184
+    },
+    {
+      "epoch": 0.94,
+      "gumbel_tau": 0.1146,
+      "image_masks_loss": 0.0129,
+      "llm_learning_rate": 4.5373533541175074e-08,
+      "loss": 0.7045,
+      "output_text_mask_loss": 0.0065,
+      "predictor_learning_rate": 4.537353354117507e-06,
+      "step": 9185
+    },
+    {
+      "epoch": 0.94,
+      "gumbel_tau": 0.1145,
+      "image_masks_loss": 0.0046,
+      "llm_learning_rate": 4.521632684727495e-08,
+      "loss": 0.689,
+      "output_text_mask_loss": 0.0435,
+      "predictor_learning_rate": 4.521632684727495e-06,
+      "step": 9186
+    },
+    {
+      "epoch": 0.94,
+      "gumbel_tau": 0.1145,
+      "image_masks_loss": 0.0082,
+      "llm_learning_rate": 4.5059390479621215e-08,
+      "loss": 0.7226,
+      "output_text_mask_loss": 0.0107,
+      "predictor_learning_rate": 4.505939047962121e-06,
+      "step": 9187
+    },
+    {
+      "epoch": 0.94,
+      "gumbel_tau": 0.1145,
+      "image_masks_loss": 0.0121,
+      "llm_learning_rate": 4.4902724455495316e-08,
+      "loss": 0.7118,
+      "output_text_mask_loss": 0.0603,
+      "predictor_learning_rate": 4.490272445549531e-06,
+      "step": 9188
+    },
+    {
+      "epoch": 0.94,
+      "gumbel_tau": 0.1144,
+      "image_masks_loss": 0.0086,
+      "llm_learning_rate": 4.474632879215013e-08,
+      "loss": 0.7427,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 4.474632879215012e-06,
+      "step": 9189
+    },
+    {
+      "epoch": 0.94,
+      "gumbel_tau": 0.1144,
+      "image_masks_loss": 0.0021,
+      "llm_learning_rate": 4.459020350680798e-08,
+      "loss": 0.8011,
+      "output_text_mask_loss": 0.0288,
+      "predictor_learning_rate": 4.459020350680798e-06,
+      "step": 9190
+    },
+    {
+      "epoch": 0.94,
+      "gumbel_tau": 0.1144,
+      "image_masks_loss": 0.014,
+      "llm_learning_rate": 4.4434348616661795e-08,
+      "loss": 0.7352,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 4.443434861666179e-06,
+      "step": 9191
+    },
+    {
+      "epoch": 0.94,
+      "gumbel_tau": 0.1144,
+      "image_masks_loss": 0.0113,
+      "llm_learning_rate": 4.42787641388745e-08,
+      "loss": 0.6892,
+      "output_text_mask_loss": 0.0053,
+      "predictor_learning_rate": 4.427876413887449e-06,
+      "step": 9192
+    },
+    {
+      "epoch": 0.94,
+      "gumbel_tau": 0.1143,
+      "image_masks_loss": 0.003,
+      "llm_learning_rate": 4.4123450090579614e-08,
+      "loss": 0.7628,
+      "output_text_mask_loss": 0.0004,
+      "predictor_learning_rate": 4.412345009057961e-06,
+      "step": 9193
+    },
+    {
+      "epoch": 0.94,
+      "gumbel_tau": 0.1143,
+      "image_masks_loss": 0.005,
+      "llm_learning_rate": 4.396840648888096e-08,
+      "loss": 0.7262,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 4.3968406488880955e-06,
+      "step": 9194
+    },
+    {
+      "epoch": 0.94,
+      "gumbel_tau": 0.1143,
+      "image_masks_loss": 0.0084,
+      "llm_learning_rate": 4.3813633350852095e-08,
+      "loss": 0.7094,
+      "output_text_mask_loss": 0.0162,
+      "predictor_learning_rate": 4.381363335085209e-06,
+      "step": 9195
+    },
+    {
+      "epoch": 0.94,
+      "gumbel_tau": 0.1143,
+      "image_masks_loss": 0.0074,
+      "llm_learning_rate": 4.365913069353661e-08,
+      "loss": 0.7799,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 4.365913069353661e-06,
+      "step": 9196
+    },
+    {
+      "epoch": 0.94,
+      "gumbel_tau": 0.1142,
+      "image_masks_loss": 0.007,
+      "llm_learning_rate": 4.3504898533948955e-08,
+      "loss": 0.7956,
+      "output_text_mask_loss": 0.0023,
+      "predictor_learning_rate": 4.350489853394896e-06,
+      "step": 9197
+    },
+    {
+      "epoch": 0.94,
+      "gumbel_tau": 0.1142,
+      "image_masks_loss": 0.0042,
+      "llm_learning_rate": 4.33509368890736e-08,
+      "loss": 0.7174,
+      "output_text_mask_loss": 0.0254,
+      "predictor_learning_rate": 4.33509368890736e-06,
+      "step": 9198
+    },
+    {
+      "epoch": 0.94,
+      "gumbel_tau": 0.1142,
+      "image_masks_loss": 0.0037,
+      "llm_learning_rate": 4.319724577586476e-08,
+      "loss": 0.6995,
+      "output_text_mask_loss": 0.0177,
+      "predictor_learning_rate": 4.3197245775864755e-06,
+      "step": 9199
+    },
+    {
+      "epoch": 0.94,
+      "gumbel_tau": 0.1142,
+      "image_masks_loss": 0.0054,
+      "llm_learning_rate": 4.304382521124778e-08,
+      "loss": 0.7269,
+      "output_text_mask_loss": 0.0206,
+      "predictor_learning_rate": 4.3043825211247775e-06,
+      "step": 9200
+    },
+    {
+      "epoch": 0.94,
+      "gumbel_tau": 0.1141,
+      "image_masks_loss": 0.0032,
+      "llm_learning_rate": 4.289067521211748e-08,
+      "loss": 0.7516,
+      "output_text_mask_loss": 0.0004,
+      "predictor_learning_rate": 4.289067521211748e-06,
+      "step": 9201
+    },
+    {
+      "epoch": 0.94,
+      "gumbel_tau": 0.1141,
+      "image_masks_loss": 0.0085,
+      "llm_learning_rate": 4.273779579533899e-08,
+      "loss": 0.752,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 4.273779579533899e-06,
+      "step": 9202
+    },
+    {
+      "epoch": 0.94,
+      "gumbel_tau": 0.1141,
+      "image_masks_loss": 0.0073,
+      "llm_learning_rate": 4.2585186977747714e-08,
+      "loss": 0.7331,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 4.258518697774771e-06,
+      "step": 9203
+    },
+    {
+      "epoch": 0.94,
+      "gumbel_tau": 0.114,
+      "image_masks_loss": 0.012,
+      "llm_learning_rate": 4.243284877614912e-08,
+      "loss": 0.7395,
+      "output_text_mask_loss": 0.0086,
+      "predictor_learning_rate": 4.243284877614911e-06,
+      "step": 9204
+    },
+    {
+      "epoch": 0.94,
+      "gumbel_tau": 0.114,
+      "image_masks_loss": 0.0118,
+      "llm_learning_rate": 4.2280781207319486e-08,
+      "loss": 0.7384,
+      "output_text_mask_loss": 0.0231,
+      "predictor_learning_rate": 4.228078120731949e-06,
+      "step": 9205
+    },
+    {
+      "epoch": 0.94,
+      "gumbel_tau": 0.114,
+      "image_masks_loss": 0.0132,
+      "llm_learning_rate": 4.2128984288004326e-08,
+      "loss": 0.652,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 4.212898428800432e-06,
+      "step": 9206
+    },
+    {
+      "epoch": 0.94,
+      "gumbel_tau": 0.114,
+      "image_masks_loss": 0.0046,
+      "llm_learning_rate": 4.197745803492026e-08,
+      "loss": 0.7241,
+      "output_text_mask_loss": 0.005,
+      "predictor_learning_rate": 4.197745803492026e-06,
+      "step": 9207
+    },
+    {
+      "epoch": 0.94,
+      "gumbel_tau": 0.1139,
+      "image_masks_loss": 0.0108,
+      "llm_learning_rate": 4.182620246475366e-08,
+      "loss": 0.6875,
+      "output_text_mask_loss": 0.0131,
+      "predictor_learning_rate": 4.182620246475366e-06,
+      "step": 9208
+    },
+    {
+      "epoch": 0.94,
+      "gumbel_tau": 0.1139,
+      "image_masks_loss": 0.0059,
+      "llm_learning_rate": 4.167521759416093e-08,
+      "loss": 0.8266,
+      "output_text_mask_loss": 0.0066,
+      "predictor_learning_rate": 4.167521759416093e-06,
+      "step": 9209
+    },
+    {
+      "epoch": 0.94,
+      "gumbel_tau": 0.1139,
+      "image_masks_loss": 0.0112,
+      "llm_learning_rate": 4.152450343976905e-08,
+      "loss": 0.6968,
+      "output_text_mask_loss": 0.0162,
+      "predictor_learning_rate": 4.152450343976905e-06,
+      "step": 9210
+    },
+    {
+      "epoch": 0.94,
+      "gumbel_tau": 0.1139,
+      "image_masks_loss": 0.0141,
+      "llm_learning_rate": 4.1374060018174736e-08,
+      "loss": 0.6443,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 4.137406001817473e-06,
+      "step": 9211
+    },
+    {
+      "epoch": 0.94,
+      "gumbel_tau": 0.1138,
+      "image_masks_loss": 0.0144,
+      "llm_learning_rate": 4.122388734594529e-08,
+      "loss": 0.7677,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 4.122388734594529e-06,
+      "step": 9212
+    },
+    {
+      "epoch": 0.94,
+      "gumbel_tau": 0.1138,
+      "image_masks_loss": 0.0156,
+      "llm_learning_rate": 4.107398543961805e-08,
+      "loss": 0.7184,
+      "output_text_mask_loss": 0.042,
+      "predictor_learning_rate": 4.107398543961804e-06,
+      "step": 9213
+    },
+    {
+      "epoch": 0.94,
+      "gumbel_tau": 0.1138,
+      "image_masks_loss": 0.0053,
+      "llm_learning_rate": 4.092435431570091e-08,
+      "loss": 0.7447,
+      "output_text_mask_loss": 0.0198,
+      "predictor_learning_rate": 4.092435431570091e-06,
+      "step": 9214
+    },
+    {
+      "epoch": 0.94,
+      "gumbel_tau": 0.1137,
+      "image_masks_loss": 0.0034,
+      "llm_learning_rate": 4.0774993990670975e-08,
+      "loss": 0.6533,
+      "output_text_mask_loss": 0.0124,
+      "predictor_learning_rate": 4.077499399067097e-06,
+      "step": 9215
+    },
+    {
+      "epoch": 0.94,
+      "gumbel_tau": 0.1137,
+      "image_masks_loss": 0.0154,
+      "llm_learning_rate": 4.0625904480976474e-08,
+      "loss": 0.6642,
+      "output_text_mask_loss": 0.0158,
+      "predictor_learning_rate": 4.062590448097648e-06,
+      "step": 9216
+    },
+    {
+      "epoch": 0.94,
+      "gumbel_tau": 0.1137,
+      "image_masks_loss": 0.0165,
+      "llm_learning_rate": 4.0477085803035386e-08,
+      "loss": 0.6589,
+      "output_text_mask_loss": 0.0614,
+      "predictor_learning_rate": 4.047708580303538e-06,
+      "step": 9217
+    },
+    {
+      "epoch": 0.94,
+      "gumbel_tau": 0.1137,
+      "image_masks_loss": 0.0134,
+      "llm_learning_rate": 4.032853797323599e-08,
+      "loss": 0.7854,
+      "output_text_mask_loss": 0.046,
+      "predictor_learning_rate": 4.032853797323599e-06,
+      "step": 9218
+    },
+    {
+      "epoch": 0.94,
+      "gumbel_tau": 0.1136,
+      "image_masks_loss": 0.0116,
+      "llm_learning_rate": 4.018026100793659e-08,
+      "loss": 0.7586,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 4.018026100793659e-06,
+      "step": 9219
+    },
+    {
+      "epoch": 0.94,
+      "gumbel_tau": 0.1136,
+      "image_masks_loss": 0.0174,
+      "llm_learning_rate": 4.0032254923466076e-08,
+      "loss": 0.7301,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 4.003225492346607e-06,
+      "step": 9220
+    },
+    {
+      "epoch": 0.94,
+      "gumbel_tau": 0.1136,
+      "image_masks_loss": 0.0072,
+      "llm_learning_rate": 3.988451973612278e-08,
+      "loss": 0.7134,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 3.988451973612278e-06,
+      "step": 9221
+    },
+    {
+      "epoch": 0.94,
+      "gumbel_tau": 0.1136,
+      "image_masks_loss": 0.0195,
+      "llm_learning_rate": 3.973705546217621e-08,
+      "loss": 0.793,
+      "output_text_mask_loss": 0.0165,
+      "predictor_learning_rate": 3.973705546217621e-06,
+      "step": 9222
+    },
+    {
+      "epoch": 0.94,
+      "gumbel_tau": 0.1135,
+      "image_masks_loss": 0.0064,
+      "llm_learning_rate": 3.9589862117864755e-08,
+      "loss": 0.7362,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 3.958986211786475e-06,
+      "step": 9223
+    },
+    {
+      "epoch": 0.95,
+      "gumbel_tau": 0.1135,
+      "image_masks_loss": 0.0126,
+      "llm_learning_rate": 3.944293971939822e-08,
+      "loss": 0.6837,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 3.944293971939822e-06,
+      "step": 9224
+    },
+    {
+      "epoch": 0.95,
+      "gumbel_tau": 0.1135,
+      "image_masks_loss": 0.0122,
+      "llm_learning_rate": 3.929628828295562e-08,
+      "loss": 0.6991,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 3.9296288282955616e-06,
+      "step": 9225
+    },
+    {
+      "epoch": 0.95,
+      "gumbel_tau": 0.1135,
+      "image_masks_loss": 0.005,
+      "llm_learning_rate": 3.914990782468708e-08,
+      "loss": 0.6528,
+      "output_text_mask_loss": 0.0065,
+      "predictor_learning_rate": 3.914990782468708e-06,
+      "step": 9226
+    },
+    {
+      "epoch": 0.95,
+      "gumbel_tau": 0.1134,
+      "image_masks_loss": 0.0054,
+      "llm_learning_rate": 3.9003798360711943e-08,
+      "loss": 0.7761,
+      "output_text_mask_loss": 0.0023,
+      "predictor_learning_rate": 3.900379836071194e-06,
+      "step": 9227
+    },
+    {
+      "epoch": 0.95,
+      "gumbel_tau": 0.1134,
+      "image_masks_loss": 0.0142,
+      "llm_learning_rate": 3.88579599071201e-08,
+      "loss": 0.8142,
+      "output_text_mask_loss": 0.0345,
+      "predictor_learning_rate": 3.88579599071201e-06,
+      "step": 9228
+    },
+    {
+      "epoch": 0.95,
+      "gumbel_tau": 0.1134,
+      "image_masks_loss": 0.0085,
+      "llm_learning_rate": 3.871239247997177e-08,
+      "loss": 0.7546,
+      "output_text_mask_loss": 0.037,
+      "predictor_learning_rate": 3.871239247997177e-06,
+      "step": 9229
+    },
+    {
+      "epoch": 0.95,
+      "gumbel_tau": 0.1133,
+      "image_masks_loss": 0.0191,
+      "llm_learning_rate": 3.8567096095296906e-08,
+      "loss": 0.672,
+      "output_text_mask_loss": 0.0051,
+      "predictor_learning_rate": 3.856709609529691e-06,
+      "step": 9230
+    },
+    {
+      "epoch": 0.95,
+      "gumbel_tau": 0.1133,
+      "image_masks_loss": 0.0045,
+      "llm_learning_rate": 3.842207076909632e-08,
+      "loss": 0.8356,
+      "output_text_mask_loss": 0.0124,
+      "predictor_learning_rate": 3.842207076909632e-06,
+      "step": 9231
+    },
+    {
+      "epoch": 0.95,
+      "gumbel_tau": 0.1133,
+      "image_masks_loss": 0.006,
+      "llm_learning_rate": 3.827731651734001e-08,
+      "loss": 0.8006,
+      "output_text_mask_loss": 0.0039,
+      "predictor_learning_rate": 3.827731651734001e-06,
+      "step": 9232
+    },
+    {
+      "epoch": 0.95,
+      "gumbel_tau": 0.1133,
+      "image_masks_loss": 0.0131,
+      "llm_learning_rate": 3.813283335596884e-08,
+      "loss": 0.6808,
+      "output_text_mask_loss": 0.061,
+      "predictor_learning_rate": 3.8132833355968844e-06,
+      "step": 9233
+    },
+    {
+      "epoch": 0.95,
+      "gumbel_tau": 0.1132,
+      "image_masks_loss": 0.0097,
+      "llm_learning_rate": 3.798862130089398e-08,
+      "loss": 0.726,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 3.7988621300893978e-06,
+      "step": 9234
+    },
+    {
+      "epoch": 0.95,
+      "gumbel_tau": 0.1132,
+      "image_masks_loss": 0.0201,
+      "llm_learning_rate": 3.784468036799577e-08,
+      "loss": 0.7942,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 3.7844680367995767e-06,
+      "step": 9235
+    },
+    {
+      "epoch": 0.95,
+      "gumbel_tau": 0.1132,
+      "image_masks_loss": 0.0094,
+      "llm_learning_rate": 3.77010105731257e-08,
+      "loss": 0.6348,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 3.770101057312569e-06,
+      "step": 9236
+    },
+    {
+      "epoch": 0.95,
+      "gumbel_tau": 0.1132,
+      "image_masks_loss": 0.0082,
+      "llm_learning_rate": 3.7557611932104995e-08,
+      "loss": 0.7441,
+      "output_text_mask_loss": 0.0231,
+      "predictor_learning_rate": 3.7557611932104994e-06,
+      "step": 9237
+    },
+    {
+      "epoch": 0.95,
+      "gumbel_tau": 0.1131,
+      "image_masks_loss": 0.0075,
+      "llm_learning_rate": 3.7414484460724655e-08,
+      "loss": 0.7021,
+      "output_text_mask_loss": 0.0086,
+      "predictor_learning_rate": 3.741448446072465e-06,
+      "step": 9238
+    },
+    {
+      "epoch": 0.95,
+      "gumbel_tau": 0.1131,
+      "image_masks_loss": 0.0106,
+      "llm_learning_rate": 3.72716281747465e-08,
+      "loss": 0.6841,
+      "output_text_mask_loss": 0.036,
+      "predictor_learning_rate": 3.7271628174746495e-06,
+      "step": 9239
+    },
+    {
+      "epoch": 0.95,
+      "gumbel_tau": 0.1131,
+      "image_masks_loss": 0.0033,
+      "llm_learning_rate": 3.7129043089902394e-08,
+      "loss": 0.6753,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 3.712904308990239e-06,
+      "step": 9240
+    },
+    {
+      "epoch": 0.95,
+      "gumbel_tau": 0.1131,
+      "image_masks_loss": 0.0144,
+      "llm_learning_rate": 3.698672922189395e-08,
+      "loss": 0.6463,
+      "output_text_mask_loss": 0.0105,
+      "predictor_learning_rate": 3.6986729221893945e-06,
+      "step": 9241
+    },
+    {
+      "epoch": 0.95,
+      "gumbel_tau": 0.113,
+      "image_masks_loss": 0.009,
+      "llm_learning_rate": 3.6844686586392786e-08,
+      "loss": 0.7187,
+      "output_text_mask_loss": 0.0195,
+      "predictor_learning_rate": 3.6844686586392783e-06,
+      "step": 9242
+    },
+    {
+      "epoch": 0.95,
+      "gumbel_tau": 0.113,
+      "image_masks_loss": 0.0066,
+      "llm_learning_rate": 3.67029151990414e-08,
+      "loss": 0.7131,
+      "output_text_mask_loss": 0.0063,
+      "predictor_learning_rate": 3.67029151990414e-06,
+      "step": 9243
+    },
+    {
+      "epoch": 0.95,
+      "gumbel_tau": 0.113,
+      "image_masks_loss": 0.0049,
+      "llm_learning_rate": 3.6561415075451745e-08,
+      "loss": 0.7338,
+      "output_text_mask_loss": 0.0641,
+      "predictor_learning_rate": 3.6561415075451745e-06,
+      "step": 9244
+    },
+    {
+      "epoch": 0.95,
+      "gumbel_tau": 0.1129,
+      "image_masks_loss": 0.0113,
+      "llm_learning_rate": 3.64201862312058e-08,
+      "loss": 0.7604,
+      "output_text_mask_loss": 0.0012,
+      "predictor_learning_rate": 3.6420186231205797e-06,
+      "step": 9245
+    },
+    {
+      "epoch": 0.95,
+      "gumbel_tau": 0.1129,
+      "image_masks_loss": 0.0032,
+      "llm_learning_rate": 3.627922868185668e-08,
+      "loss": 0.7467,
+      "output_text_mask_loss": 0.0174,
+      "predictor_learning_rate": 3.627922868185668e-06,
+      "step": 9246
+    },
+    {
+      "epoch": 0.95,
+      "gumbel_tau": 0.1129,
+      "image_masks_loss": 0.0036,
+      "llm_learning_rate": 3.6138542442926693e-08,
+      "loss": 0.654,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 3.6138542442926692e-06,
+      "step": 9247
+    },
+    {
+      "epoch": 0.95,
+      "gumbel_tau": 0.1129,
+      "image_masks_loss": 0.0046,
+      "llm_learning_rate": 3.5998127529908444e-08,
+      "loss": 0.709,
+      "output_text_mask_loss": 0.0174,
+      "predictor_learning_rate": 3.5998127529908442e-06,
+      "step": 9248
+    },
+    {
+      "epoch": 0.95,
+      "gumbel_tau": 0.1128,
+      "image_masks_loss": 0.0079,
+      "llm_learning_rate": 3.5857983958264567e-08,
+      "loss": 0.8009,
+      "output_text_mask_loss": 0.0006,
+      "predictor_learning_rate": 3.5857983958264562e-06,
+      "step": 9249
+    },
+    {
+      "epoch": 0.95,
+      "gumbel_tau": 0.1128,
+      "image_masks_loss": 0.0122,
+      "llm_learning_rate": 3.5718111743428266e-08,
+      "loss": 0.6475,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 3.571811174342826e-06,
+      "step": 9250
+    },
+    {
+      "epoch": 0.95,
+      "gumbel_tau": 0.1128,
+      "image_masks_loss": 0.0174,
+      "llm_learning_rate": 3.55785109008025e-08,
+      "loss": 0.7709,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 3.5578510900802497e-06,
+      "step": 9251
+    },
+    {
+      "epoch": 0.95,
+      "gumbel_tau": 0.1128,
+      "image_masks_loss": 0.0167,
+      "llm_learning_rate": 3.543918144576025e-08,
+      "loss": 0.6402,
+      "output_text_mask_loss": 0.0004,
+      "predictor_learning_rate": 3.543918144576025e-06,
+      "step": 9252
+    },
+    {
+      "epoch": 0.95,
+      "gumbel_tau": 0.1127,
+      "image_masks_loss": 0.0067,
+      "llm_learning_rate": 3.530012339364508e-08,
+      "loss": 0.7753,
+      "output_text_mask_loss": 0.0233,
+      "predictor_learning_rate": 3.530012339364508e-06,
+      "step": 9253
+    },
+    {
+      "epoch": 0.95,
+      "gumbel_tau": 0.1127,
+      "image_masks_loss": 0.0079,
+      "llm_learning_rate": 3.516133675977001e-08,
+      "loss": 0.7701,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 3.5161336759770012e-06,
+      "step": 9254
+    },
+    {
+      "epoch": 0.95,
+      "gumbel_tau": 0.1127,
+      "image_masks_loss": 0.0252,
+      "llm_learning_rate": 3.502282155941894e-08,
+      "loss": 0.775,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 3.502282155941894e-06,
+      "step": 9255
+    },
+    {
+      "epoch": 0.95,
+      "gumbel_tau": 0.1127,
+      "image_masks_loss": 0.0134,
+      "llm_learning_rate": 3.488457780784549e-08,
+      "loss": 0.7248,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 3.4884577807845485e-06,
+      "step": 9256
+    },
+    {
+      "epoch": 0.95,
+      "gumbel_tau": 0.1126,
+      "image_masks_loss": 0.0102,
+      "llm_learning_rate": 3.474660552027276e-08,
+      "loss": 0.7281,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 3.4746605520272757e-06,
+      "step": 9257
+    },
+    {
+      "epoch": 0.95,
+      "gumbel_tau": 0.1126,
+      "image_masks_loss": 0.0046,
+      "llm_learning_rate": 3.4608904711894986e-08,
+      "loss": 0.7213,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 3.4608904711894982e-06,
+      "step": 9258
+    },
+    {
+      "epoch": 0.95,
+      "gumbel_tau": 0.1126,
+      "image_masks_loss": 0.0168,
+      "llm_learning_rate": 3.447147539787615e-08,
+      "loss": 0.6934,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 3.447147539787615e-06,
+      "step": 9259
+    },
+    {
+      "epoch": 0.95,
+      "gumbel_tau": 0.1125,
+      "image_masks_loss": 0.0091,
+      "llm_learning_rate": 3.4334317593350266e-08,
+      "loss": 0.7063,
+      "output_text_mask_loss": 0.028,
+      "predictor_learning_rate": 3.4334317593350262e-06,
+      "step": 9260
+    },
+    {
+      "epoch": 0.95,
+      "gumbel_tau": 0.1125,
+      "image_masks_loss": 0.0177,
+      "llm_learning_rate": 3.419743131342135e-08,
+      "loss": 0.7366,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 3.4197431313421347e-06,
+      "step": 9261
+    },
+    {
+      "epoch": 0.95,
+      "gumbel_tau": 0.1125,
+      "image_masks_loss": 0.0054,
+      "llm_learning_rate": 3.4060816573163734e-08,
+      "loss": 0.7166,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 3.4060816573163733e-06,
+      "step": 9262
+    },
+    {
+      "epoch": 0.95,
+      "gumbel_tau": 0.1125,
+      "image_masks_loss": 0.0066,
+      "llm_learning_rate": 3.392447338762178e-08,
+      "loss": 0.6101,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 3.392447338762178e-06,
+      "step": 9263
+    },
+    {
+      "epoch": 0.95,
+      "gumbel_tau": 0.1124,
+      "image_masks_loss": 0.0123,
+      "llm_learning_rate": 3.378840177180931e-08,
+      "loss": 0.6623,
+      "output_text_mask_loss": 0.0004,
+      "predictor_learning_rate": 3.3788401771809306e-06,
+      "step": 9264
+    },
+    {
+      "epoch": 0.95,
+      "gumbel_tau": 0.1124,
+      "image_masks_loss": 0.0204,
+      "llm_learning_rate": 3.365260174071156e-08,
+      "loss": 0.6788,
+      "output_text_mask_loss": 0.0488,
+      "predictor_learning_rate": 3.3652601740711554e-06,
+      "step": 9265
+    },
+    {
+      "epoch": 0.95,
+      "gumbel_tau": 0.1124,
+      "image_masks_loss": 0.0106,
+      "llm_learning_rate": 3.3517073309282945e-08,
+      "loss": 0.6759,
+      "output_text_mask_loss": 0.0082,
+      "predictor_learning_rate": 3.351707330928294e-06,
+      "step": 9266
+    },
+    {
+      "epoch": 0.95,
+      "gumbel_tau": 0.1124,
+      "image_masks_loss": 0.0075,
+      "llm_learning_rate": 3.33818164924482e-08,
+      "loss": 0.7616,
+      "output_text_mask_loss": 0.001,
+      "predictor_learning_rate": 3.3381816492448204e-06,
+      "step": 9267
+    },
+    {
+      "epoch": 0.95,
+      "gumbel_tau": 0.1123,
+      "image_masks_loss": 0.0079,
+      "llm_learning_rate": 3.324683130510181e-08,
+      "loss": 0.7529,
+      "output_text_mask_loss": 0.0469,
+      "predictor_learning_rate": 3.324683130510181e-06,
+      "step": 9268
+    },
+    {
+      "epoch": 0.95,
+      "gumbel_tau": 0.1123,
+      "image_masks_loss": 0.0043,
+      "llm_learning_rate": 3.31121177621091e-08,
+      "loss": 0.7752,
+      "output_text_mask_loss": 0.0286,
+      "predictor_learning_rate": 3.3112117762109096e-06,
+      "step": 9269
+    },
+    {
+      "epoch": 0.95,
+      "gumbel_tau": 0.1123,
+      "image_masks_loss": 0.0093,
+      "llm_learning_rate": 3.297767587830458e-08,
+      "loss": 0.6491,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 3.297767587830458e-06,
+      "step": 9270
+    },
+    {
+      "epoch": 0.95,
+      "gumbel_tau": 0.1123,
+      "image_masks_loss": 0.017,
+      "llm_learning_rate": 3.284350566849337e-08,
+      "loss": 0.7362,
+      "output_text_mask_loss": 0.0015,
+      "predictor_learning_rate": 3.2843505668493368e-06,
+      "step": 9271
+    },
+    {
+      "epoch": 0.95,
+      "gumbel_tau": 0.1122,
+      "image_masks_loss": 0.0054,
+      "llm_learning_rate": 3.2709607147450864e-08,
+      "loss": 0.5654,
+      "output_text_mask_loss": NaN,
+      "predictor_learning_rate": 3.270960714745086e-06,
+      "step": 9272
+    },
+    {
+      "epoch": 0.95,
+      "gumbel_tau": 0.1122,
+      "image_masks_loss": 0.028,
+      "llm_learning_rate": 3.257598032992193e-08,
+      "loss": 0.7225,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 3.257598032992193e-06,
+      "step": 9273
+    },
+    {
+      "epoch": 0.95,
+      "gumbel_tau": 0.1122,
+      "image_masks_loss": 0.0132,
+      "llm_learning_rate": 3.244262523062203e-08,
+      "loss": 0.7675,
+      "output_text_mask_loss": 0.0086,
+      "predictor_learning_rate": 3.2442625230622024e-06,
+      "step": 9274
+    },
+    {
+      "epoch": 0.95,
+      "gumbel_tau": 0.1121,
+      "image_masks_loss": 0.0072,
+      "llm_learning_rate": 3.230954186423663e-08,
+      "loss": 0.7971,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 3.2309541864236623e-06,
+      "step": 9275
+    },
+    {
+      "epoch": 0.95,
+      "gumbel_tau": 0.1121,
+      "image_masks_loss": 0.0097,
+      "llm_learning_rate": 3.217673024542095e-08,
+      "loss": 0.6964,
+      "output_text_mask_loss": 0.0429,
+      "predictor_learning_rate": 3.217673024542095e-06,
+      "step": 9276
+    },
+    {
+      "epoch": 0.95,
+      "gumbel_tau": 0.1121,
+      "image_masks_loss": 0.0204,
+      "llm_learning_rate": 3.2044190388800524e-08,
+      "loss": 0.7619,
+      "output_text_mask_loss": 0.0088,
+      "predictor_learning_rate": 3.204419038880052e-06,
+      "step": 9277
+    },
+    {
+      "epoch": 0.95,
+      "gumbel_tau": 0.1121,
+      "image_masks_loss": 0.009,
+      "llm_learning_rate": 3.1911922308971166e-08,
+      "loss": 0.6821,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 3.1911922308971163e-06,
+      "step": 9278
+    },
+    {
+      "epoch": 0.95,
+      "gumbel_tau": 0.112,
+      "image_masks_loss": 0.0102,
+      "llm_learning_rate": 3.177992602049818e-08,
+      "loss": 0.7318,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 3.1779926020498173e-06,
+      "step": 9279
+    },
+    {
+      "epoch": 0.95,
+      "gumbel_tau": 0.112,
+      "image_masks_loss": 0.0033,
+      "llm_learning_rate": 3.16482015379177e-08,
+      "loss": 0.6955,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 3.16482015379177e-06,
+      "step": 9280
+    },
+    {
+      "epoch": 0.95,
+      "gumbel_tau": 0.112,
+      "image_masks_loss": 0.005,
+      "llm_learning_rate": 3.151674887573508e-08,
+      "loss": 0.8466,
+      "output_text_mask_loss": 0.0088,
+      "predictor_learning_rate": 3.151674887573508e-06,
+      "step": 9281
+    },
+    {
+      "epoch": 0.95,
+      "gumbel_tau": 0.112,
+      "image_masks_loss": 0.0084,
+      "llm_learning_rate": 3.1385568048426795e-08,
+      "loss": 0.6969,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 3.1385568048426795e-06,
+      "step": 9282
+    },
+    {
+      "epoch": 0.95,
+      "gumbel_tau": 0.1119,
+      "image_masks_loss": 0.0051,
+      "llm_learning_rate": 3.1254659070437954e-08,
+      "loss": 0.78,
+      "output_text_mask_loss": 0.0271,
+      "predictor_learning_rate": 3.125465907043795e-06,
+      "step": 9283
+    },
+    {
+      "epoch": 0.95,
+      "gumbel_tau": 0.1119,
+      "image_masks_loss": 0.0099,
+      "llm_learning_rate": 3.112402195618508e-08,
+      "loss": 0.7681,
+      "output_text_mask_loss": 0.0374,
+      "predictor_learning_rate": 3.1124021956185078e-06,
+      "step": 9284
+    },
+    {
+      "epoch": 0.95,
+      "gumbel_tau": 0.1119,
+      "image_masks_loss": 0.0069,
+      "llm_learning_rate": 3.099365672005417e-08,
+      "loss": 0.7317,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 3.0993656720054163e-06,
+      "step": 9285
+    },
+    {
+      "epoch": 0.95,
+      "gumbel_tau": 0.1119,
+      "image_masks_loss": 0.0125,
+      "llm_learning_rate": 3.0863563376401226e-08,
+      "loss": 0.7279,
+      "output_text_mask_loss": 0.0259,
+      "predictor_learning_rate": 3.086356337640123e-06,
+      "step": 9286
+    },
+    {
+      "epoch": 0.95,
+      "gumbel_tau": 0.1118,
+      "image_masks_loss": 0.0134,
+      "llm_learning_rate": 3.0733741939552585e-08,
+      "loss": 0.6944,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 3.0733741939552583e-06,
+      "step": 9287
+    },
+    {
+      "epoch": 0.95,
+      "gumbel_tau": 0.1118,
+      "image_masks_loss": 0.0044,
+      "llm_learning_rate": 3.060419242380458e-08,
+      "loss": 0.8289,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 3.0604192423804576e-06,
+      "step": 9288
+    },
+    {
+      "epoch": 0.95,
+      "gumbel_tau": 0.1118,
+      "image_masks_loss": 0.0125,
+      "llm_learning_rate": 3.047491484342274e-08,
+      "loss": 0.6693,
+      "output_text_mask_loss": 0.0086,
+      "predictor_learning_rate": 3.0474914843422742e-06,
+      "step": 9289
+    },
+    {
+      "epoch": 0.95,
+      "gumbel_tau": 0.1118,
+      "image_masks_loss": 0.0107,
+      "llm_learning_rate": 3.03459092126443e-08,
+      "loss": 0.6774,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 3.0345909212644294e-06,
+      "step": 9290
+    },
+    {
+      "epoch": 0.95,
+      "gumbel_tau": 0.1117,
+      "image_masks_loss": 0.0169,
+      "llm_learning_rate": 3.021717554567511e-08,
+      "loss": 0.7396,
+      "output_text_mask_loss": 0.0078,
+      "predictor_learning_rate": 3.02171755456751e-06,
+      "step": 9291
+    },
+    {
+      "epoch": 0.95,
+      "gumbel_tau": 0.1117,
+      "image_masks_loss": 0.0045,
+      "llm_learning_rate": 3.008871385669187e-08,
+      "loss": 0.7483,
+      "output_text_mask_loss": 0.021,
+      "predictor_learning_rate": 3.0088713856691873e-06,
+      "step": 9292
+    },
+    {
+      "epoch": 0.95,
+      "gumbel_tau": 0.1117,
+      "image_masks_loss": 0.0195,
+      "llm_learning_rate": 2.996052415984107e-08,
+      "loss": 0.7091,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 2.9960524159841073e-06,
+      "step": 9293
+    },
+    {
+      "epoch": 0.95,
+      "gumbel_tau": 0.1116,
+      "image_masks_loss": 0.0073,
+      "llm_learning_rate": 2.983260646923891e-08,
+      "loss": 0.6942,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 2.983260646923891e-06,
+      "step": 9294
+    },
+    {
+      "epoch": 0.95,
+      "gumbel_tau": 0.1116,
+      "image_masks_loss": 0.0049,
+      "llm_learning_rate": 2.970496079897217e-08,
+      "loss": 0.7359,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 2.9704960798972167e-06,
+      "step": 9295
+    },
+    {
+      "epoch": 0.95,
+      "gumbel_tau": 0.1116,
+      "image_masks_loss": 0.006,
+      "llm_learning_rate": 2.9577587163097666e-08,
+      "loss": 0.694,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 2.9577587163097663e-06,
+      "step": 9296
+    },
+    {
+      "epoch": 0.95,
+      "gumbel_tau": 0.1116,
+      "image_masks_loss": 0.0065,
+      "llm_learning_rate": 2.945048557564195e-08,
+      "loss": 0.6785,
+      "output_text_mask_loss": 0.0111,
+      "predictor_learning_rate": 2.945048557564195e-06,
+      "step": 9297
+    },
+    {
+      "epoch": 0.95,
+      "gumbel_tau": 0.1115,
+      "image_masks_loss": 0.0067,
+      "llm_learning_rate": 2.932365605060161e-08,
+      "loss": 0.7359,
+      "output_text_mask_loss": 0.0328,
+      "predictor_learning_rate": 2.932365605060161e-06,
+      "step": 9298
+    },
+    {
+      "epoch": 0.95,
+      "gumbel_tau": 0.1115,
+      "image_masks_loss": 0.0145,
+      "llm_learning_rate": 2.919709860194353e-08,
+      "loss": 0.8141,
+      "output_text_mask_loss": 0.0082,
+      "predictor_learning_rate": 2.919709860194353e-06,
+      "step": 9299
+    },
+    {
+      "epoch": 0.95,
+      "gumbel_tau": 0.1115,
+      "image_masks_loss": 0.0094,
+      "llm_learning_rate": 2.9070813243604336e-08,
+      "loss": 0.6539,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 2.9070813243604334e-06,
+      "step": 9300
+    },
+    {
+      "epoch": 0.95,
+      "gumbel_tau": 0.1115,
+      "image_masks_loss": 0.0065,
+      "llm_learning_rate": 2.8944799989491234e-08,
+      "loss": 0.7214,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 2.8944799989491233e-06,
+      "step": 9301
+    },
+    {
+      "epoch": 0.95,
+      "gumbel_tau": 0.1114,
+      "image_masks_loss": 0.0053,
+      "llm_learning_rate": 2.8819058853480632e-08,
+      "loss": 0.6894,
+      "output_text_mask_loss": 0.0076,
+      "predictor_learning_rate": 2.8819058853480627e-06,
+      "step": 9302
+    },
+    {
+      "epoch": 0.95,
+      "gumbel_tau": 0.1114,
+      "image_masks_loss": 0.0157,
+      "llm_learning_rate": 2.8693589849419502e-08,
+      "loss": 0.8297,
+      "output_text_mask_loss": 0.0322,
+      "predictor_learning_rate": 2.86935898494195e-06,
+      "step": 9303
+    },
+    {
+      "epoch": 0.95,
+      "gumbel_tau": 0.1114,
+      "image_masks_loss": 0.0071,
+      "llm_learning_rate": 2.856839299112485e-08,
+      "loss": 0.6982,
+      "output_text_mask_loss": 0.0025,
+      "predictor_learning_rate": 2.8568392991124845e-06,
+      "step": 9304
+    },
+    {
+      "epoch": 0.95,
+      "gumbel_tau": 0.1114,
+      "image_masks_loss": 0.0172,
+      "llm_learning_rate": 2.844346829238398e-08,
+      "loss": 0.6867,
+      "output_text_mask_loss": 0.0755,
+      "predictor_learning_rate": 2.844346829238398e-06,
+      "step": 9305
+    },
+    {
+      "epoch": 0.95,
+      "gumbel_tau": 0.1113,
+      "image_masks_loss": 0.0054,
+      "llm_learning_rate": 2.8318815766953388e-08,
+      "loss": 0.741,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 2.8318815766953387e-06,
+      "step": 9306
+    },
+    {
+      "epoch": 0.95,
+      "gumbel_tau": 0.1113,
+      "image_masks_loss": 0.0085,
+      "llm_learning_rate": 2.819443542856043e-08,
+      "loss": 0.7235,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 2.8194435428560427e-06,
+      "step": 9307
+    },
+    {
+      "epoch": 0.95,
+      "gumbel_tau": 0.1113,
+      "image_masks_loss": 0.0234,
+      "llm_learning_rate": 2.8070327290901645e-08,
+      "loss": 0.6736,
+      "output_text_mask_loss": 0.0275,
+      "predictor_learning_rate": 2.807032729090164e-06,
+      "step": 9308
+    },
+    {
+      "epoch": 0.95,
+      "gumbel_tau": 0.1113,
+      "image_masks_loss": 0.0114,
+      "llm_learning_rate": 2.7946491367644435e-08,
+      "loss": 0.7525,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 2.7946491367644433e-06,
+      "step": 9309
+    },
+    {
+      "epoch": 0.95,
+      "gumbel_tau": 0.1112,
+      "image_masks_loss": 0.0114,
+      "llm_learning_rate": 2.7822927672426224e-08,
+      "loss": 0.6763,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 2.7822927672426226e-06,
+      "step": 9310
+    },
+    {
+      "epoch": 0.95,
+      "gumbel_tau": 0.1112,
+      "image_masks_loss": 0.0081,
+      "llm_learning_rate": 2.7699636218853355e-08,
+      "loss": 0.6985,
+      "output_text_mask_loss": 0.0177,
+      "predictor_learning_rate": 2.769963621885335e-06,
+      "step": 9311
+    },
+    {
+      "epoch": 0.95,
+      "gumbel_tau": 0.1112,
+      "image_masks_loss": 0.0081,
+      "llm_learning_rate": 2.7576617020503848e-08,
+      "loss": 0.7253,
+      "output_text_mask_loss": 0.0081,
+      "predictor_learning_rate": 2.7576617020503846e-06,
+      "step": 9312
+    },
+    {
+      "epoch": 0.95,
+      "gumbel_tau": 0.1111,
+      "image_masks_loss": 0.0128,
+      "llm_learning_rate": 2.7453870090924094e-08,
+      "loss": 0.7319,
+      "output_text_mask_loss": 0.0486,
+      "predictor_learning_rate": 2.7453870090924094e-06,
+      "step": 9313
+    },
+    {
+      "epoch": 0.95,
+      "gumbel_tau": 0.1111,
+      "image_masks_loss": 0.0094,
+      "llm_learning_rate": 2.7331395443631615e-08,
+      "loss": 0.6934,
+      "output_text_mask_loss": 0.0128,
+      "predictor_learning_rate": 2.733139544363161e-06,
+      "step": 9314
+    },
+    {
+      "epoch": 0.95,
+      "gumbel_tau": 0.1111,
+      "image_masks_loss": 0.0045,
+      "llm_learning_rate": 2.7209193092113395e-08,
+      "loss": 0.7229,
+      "output_text_mask_loss": 0.0175,
+      "predictor_learning_rate": 2.7209193092113395e-06,
+      "step": 9315
+    },
+    {
+      "epoch": 0.95,
+      "gumbel_tau": 0.1111,
+      "image_masks_loss": 0.0107,
+      "llm_learning_rate": 2.7087263049826727e-08,
+      "loss": 0.6906,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 2.7087263049826727e-06,
+      "step": 9316
+    },
+    {
+      "epoch": 0.95,
+      "gumbel_tau": 0.111,
+      "image_masks_loss": 0.0292,
+      "llm_learning_rate": 2.696560533019893e-08,
+      "loss": 0.724,
+      "output_text_mask_loss": 0.0355,
+      "predictor_learning_rate": 2.6965605330198926e-06,
+      "step": 9317
+    },
+    {
+      "epoch": 0.95,
+      "gumbel_tau": 0.111,
+      "image_masks_loss": 0.0136,
+      "llm_learning_rate": 2.6844219946627338e-08,
+      "loss": 0.7647,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 2.6844219946627334e-06,
+      "step": 9318
+    },
+    {
+      "epoch": 0.95,
+      "gumbel_tau": 0.111,
+      "image_masks_loss": 0.0061,
+      "llm_learning_rate": 2.6723106912478757e-08,
+      "loss": 0.7075,
+      "output_text_mask_loss": 0.009,
+      "predictor_learning_rate": 2.6723106912478755e-06,
+      "step": 9319
+    },
+    {
+      "epoch": 0.95,
+      "gumbel_tau": 0.111,
+      "image_masks_loss": 0.0055,
+      "llm_learning_rate": 2.6602266241090856e-08,
+      "loss": 0.6443,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 2.6602266241090855e-06,
+      "step": 9320
+    },
+    {
+      "epoch": 0.96,
+      "gumbel_tau": 0.1109,
+      "image_masks_loss": 0.008,
+      "llm_learning_rate": 2.648169794577049e-08,
+      "loss": 0.7559,
+      "output_text_mask_loss": 0.0305,
+      "predictor_learning_rate": 2.648169794577049e-06,
+      "step": 9321
+    },
+    {
+      "epoch": 0.96,
+      "gumbel_tau": 0.1109,
+      "image_masks_loss": 0.0186,
+      "llm_learning_rate": 2.6361402039795093e-08,
+      "loss": 0.7436,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 2.636140203979509e-06,
+      "step": 9322
+    },
+    {
+      "epoch": 0.96,
+      "gumbel_tau": 0.1109,
+      "image_masks_loss": 0.005,
+      "llm_learning_rate": 2.6241378536411842e-08,
+      "loss": 0.7197,
+      "output_text_mask_loss": 0.0143,
+      "predictor_learning_rate": 2.624137853641184e-06,
+      "step": 9323
+    },
+    {
+      "epoch": 0.96,
+      "gumbel_tau": 0.1109,
+      "image_masks_loss": 0.017,
+      "llm_learning_rate": 2.6121627448838227e-08,
+      "loss": 0.716,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 2.6121627448838227e-06,
+      "step": 9324
+    },
+    {
+      "epoch": 0.96,
+      "gumbel_tau": 0.1108,
+      "image_masks_loss": 0.0139,
+      "llm_learning_rate": 2.600214879026147e-08,
+      "loss": 0.678,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 2.600214879026147e-06,
+      "step": 9325
+    },
+    {
+      "epoch": 0.96,
+      "gumbel_tau": 0.1108,
+      "image_masks_loss": 0.0113,
+      "llm_learning_rate": 2.5882942573838553e-08,
+      "loss": 0.8117,
+      "output_text_mask_loss": 0.0051,
+      "predictor_learning_rate": 2.588294257383855e-06,
+      "step": 9326
+    },
+    {
+      "epoch": 0.96,
+      "gumbel_tau": 0.1108,
+      "image_masks_loss": 0.0177,
+      "llm_learning_rate": 2.5764008812697028e-08,
+      "loss": 0.6775,
+      "output_text_mask_loss": 0.0322,
+      "predictor_learning_rate": 2.5764008812697026e-06,
+      "step": 9327
+    },
+    {
+      "epoch": 0.96,
+      "gumbel_tau": 0.1108,
+      "image_masks_loss": 0.0076,
+      "llm_learning_rate": 2.5645347519933917e-08,
+      "loss": 0.7695,
+      "output_text_mask_loss": 0.009,
+      "predictor_learning_rate": 2.5645347519933915e-06,
+      "step": 9328
+    },
+    {
+      "epoch": 0.96,
+      "gumbel_tau": 0.1107,
+      "image_masks_loss": 0.007,
+      "llm_learning_rate": 2.5526958708616267e-08,
+      "loss": 0.7787,
+      "output_text_mask_loss": 0.0374,
+      "predictor_learning_rate": 2.5526958708616267e-06,
+      "step": 9329
+    },
+    {
+      "epoch": 0.96,
+      "gumbel_tau": 0.1107,
+      "image_masks_loss": 0.0093,
+      "llm_learning_rate": 2.5408842391781708e-08,
+      "loss": 0.8236,
+      "output_text_mask_loss": 0.02,
+      "predictor_learning_rate": 2.5408842391781706e-06,
+      "step": 9330
+    },
+    {
+      "epoch": 0.96,
+      "gumbel_tau": 0.1107,
+      "image_masks_loss": 0.012,
+      "llm_learning_rate": 2.529099858243761e-08,
+      "loss": 0.7529,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 2.5290998582437606e-06,
+      "step": 9331
+    },
+    {
+      "epoch": 0.96,
+      "gumbel_tau": 0.1107,
+      "image_masks_loss": 0.0099,
+      "llm_learning_rate": 2.5173427293560815e-08,
+      "loss": 0.7744,
+      "output_text_mask_loss": 0.0147,
+      "predictor_learning_rate": 2.517342729356081e-06,
+      "step": 9332
+    },
+    {
+      "epoch": 0.96,
+      "gumbel_tau": 0.1106,
+      "image_masks_loss": 0.0222,
+      "llm_learning_rate": 2.505612853809847e-08,
+      "loss": 0.7398,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 2.5056128538098466e-06,
+      "step": 9333
+    },
+    {
+      "epoch": 0.96,
+      "gumbel_tau": 0.1106,
+      "image_masks_loss": 0.0186,
+      "llm_learning_rate": 2.493910232896829e-08,
+      "loss": 0.7469,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 2.493910232896829e-06,
+      "step": 9334
+    },
+    {
+      "epoch": 0.96,
+      "gumbel_tau": 0.1106,
+      "image_masks_loss": 0.0062,
+      "llm_learning_rate": 2.4822348679056918e-08,
+      "loss": 0.6775,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 2.4822348679056918e-06,
+      "step": 9335
+    },
+    {
+      "epoch": 0.96,
+      "gumbel_tau": 0.1105,
+      "image_masks_loss": 0.0151,
+      "llm_learning_rate": 2.4705867601221846e-08,
+      "loss": 0.6834,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 2.4705867601221844e-06,
+      "step": 9336
+    },
+    {
+      "epoch": 0.96,
+      "gumbel_tau": 0.1105,
+      "image_masks_loss": 0.017,
+      "llm_learning_rate": 2.458965910829003e-08,
+      "loss": 0.8305,
+      "output_text_mask_loss": 0.0012,
+      "predictor_learning_rate": 2.458965910829003e-06,
+      "step": 9337
+    },
+    {
+      "epoch": 0.96,
+      "gumbel_tau": 0.1105,
+      "image_masks_loss": 0.018,
+      "llm_learning_rate": 2.4473723213058742e-08,
+      "loss": 0.7064,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 2.447372321305874e-06,
+      "step": 9338
+    },
+    {
+      "epoch": 0.96,
+      "gumbel_tau": 0.1105,
+      "image_masks_loss": 0.0142,
+      "llm_learning_rate": 2.4358059928295264e-08,
+      "loss": 0.7345,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 2.4358059928295262e-06,
+      "step": 9339
+    },
+    {
+      "epoch": 0.96,
+      "gumbel_tau": 0.1104,
+      "image_masks_loss": 0.022,
+      "llm_learning_rate": 2.424266926673635e-08,
+      "loss": 0.7119,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 2.424266926673635e-06,
+      "step": 9340
+    },
+    {
+      "epoch": 0.96,
+      "gumbel_tau": 0.1104,
+      "image_masks_loss": 0.0062,
+      "llm_learning_rate": 2.4127551241089342e-08,
+      "loss": 0.6749,
+      "output_text_mask_loss": 0.0025,
+      "predictor_learning_rate": 2.4127551241089342e-06,
+      "step": 9341
+    },
+    {
+      "epoch": 0.96,
+      "gumbel_tau": 0.1104,
+      "image_masks_loss": 0.0156,
+      "llm_learning_rate": 2.401270586403104e-08,
+      "loss": 0.7816,
+      "output_text_mask_loss": 0.0492,
+      "predictor_learning_rate": 2.401270586403104e-06,
+      "step": 9342
+    },
+    {
+      "epoch": 0.96,
+      "gumbel_tau": 0.1104,
+      "image_masks_loss": 0.0056,
+      "llm_learning_rate": 2.3898133148208547e-08,
+      "loss": 0.7731,
+      "output_text_mask_loss": 0.0103,
+      "predictor_learning_rate": 2.3898133148208543e-06,
+      "step": 9343
+    },
+    {
+      "epoch": 0.96,
+      "gumbel_tau": 0.1103,
+      "image_masks_loss": 0.0093,
+      "llm_learning_rate": 2.3783833106239272e-08,
+      "loss": 0.7472,
+      "output_text_mask_loss": 0.0208,
+      "predictor_learning_rate": 2.378383310623927e-06,
+      "step": 9344
+    },
+    {
+      "epoch": 0.96,
+      "gumbel_tau": 0.1103,
+      "image_masks_loss": 0.0052,
+      "llm_learning_rate": 2.366980575070954e-08,
+      "loss": 0.8269,
+      "output_text_mask_loss": 0.0027,
+      "predictor_learning_rate": 2.3669805750709537e-06,
+      "step": 9345
+    },
+    {
+      "epoch": 0.96,
+      "gumbel_tau": 0.1103,
+      "image_masks_loss": 0.0083,
+      "llm_learning_rate": 2.35560510941768e-08,
+      "loss": 0.81,
+      "output_text_mask_loss": 0.0301,
+      "predictor_learning_rate": 2.35560510941768e-06,
+      "step": 9346
+    },
+    {
+      "epoch": 0.96,
+      "gumbel_tau": 0.1103,
+      "image_masks_loss": 0.0056,
+      "llm_learning_rate": 2.3442569149167705e-08,
+      "loss": 0.6957,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 2.3442569149167704e-06,
+      "step": 9347
+    },
+    {
+      "epoch": 0.96,
+      "gumbel_tau": 0.1102,
+      "image_masks_loss": 0.0061,
+      "llm_learning_rate": 2.3329359928179195e-08,
+      "loss": 0.8074,
+      "output_text_mask_loss": 0.0263,
+      "predictor_learning_rate": 2.3329359928179194e-06,
+      "step": 9348
+    },
+    {
+      "epoch": 0.96,
+      "gumbel_tau": 0.1102,
+      "image_masks_loss": 0.0152,
+      "llm_learning_rate": 2.3216423443677972e-08,
+      "loss": 0.7769,
+      "output_text_mask_loss": 0.0139,
+      "predictor_learning_rate": 2.321642344367797e-06,
+      "step": 9349
+    },
+    {
+      "epoch": 0.96,
+      "gumbel_tau": 0.1102,
+      "image_masks_loss": 0.0129,
+      "llm_learning_rate": 2.3103759708101302e-08,
+      "loss": 0.6976,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 2.31037597081013e-06,
+      "step": 9350
+    },
+    {
+      "epoch": 0.96,
+      "gumbel_tau": 0.1102,
+      "image_masks_loss": 0.0058,
+      "llm_learning_rate": 2.2991368733855656e-08,
+      "loss": 0.8144,
+      "output_text_mask_loss": 0.0687,
+      "predictor_learning_rate": 2.2991368733855656e-06,
+      "step": 9351
+    },
+    {
+      "epoch": 0.96,
+      "gumbel_tau": 0.1101,
+      "image_masks_loss": 0.0071,
+      "llm_learning_rate": 2.287925053331752e-08,
+      "loss": 0.7324,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 2.2879250533317518e-06,
+      "step": 9352
+    },
+    {
+      "epoch": 0.96,
+      "gumbel_tau": 0.1101,
+      "image_masks_loss": 0.0143,
+      "llm_learning_rate": 2.276740511883424e-08,
+      "loss": 0.7862,
+      "output_text_mask_loss": 0.0057,
+      "predictor_learning_rate": 2.276740511883424e-06,
+      "step": 9353
+    },
+    {
+      "epoch": 0.96,
+      "gumbel_tau": 0.1101,
+      "image_masks_loss": 0.0176,
+      "llm_learning_rate": 2.2655832502721797e-08,
+      "loss": 0.7499,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 2.2655832502721796e-06,
+      "step": 9354
+    },
+    {
+      "epoch": 0.96,
+      "gumbel_tau": 0.1101,
+      "image_masks_loss": 0.0129,
+      "llm_learning_rate": 2.254453269726731e-08,
+      "loss": 0.7074,
+      "output_text_mask_loss": 0.0721,
+      "predictor_learning_rate": 2.2544532697267305e-06,
+      "step": 9355
+    },
+    {
+      "epoch": 0.96,
+      "gumbel_tau": 0.11,
+      "image_masks_loss": 0.0061,
+      "llm_learning_rate": 2.2433505714727078e-08,
+      "loss": 0.6822,
+      "output_text_mask_loss": 0.0021,
+      "predictor_learning_rate": 2.2433505714727074e-06,
+      "step": 9356
+    },
+    {
+      "epoch": 0.96,
+      "gumbel_tau": 0.11,
+      "image_masks_loss": 0.0155,
+      "llm_learning_rate": 2.2322751567327995e-08,
+      "loss": 0.7618,
+      "output_text_mask_loss": 0.033,
+      "predictor_learning_rate": 2.2322751567327994e-06,
+      "step": 9357
+    },
+    {
+      "epoch": 0.96,
+      "gumbel_tau": 0.11,
+      "image_masks_loss": 0.015,
+      "llm_learning_rate": 2.2212270267266133e-08,
+      "loss": 0.7429,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 2.2212270267266132e-06,
+      "step": 9358
+    },
+    {
+      "epoch": 0.96,
+      "gumbel_tau": 0.1099,
+      "image_masks_loss": 0.0093,
+      "llm_learning_rate": 2.2102061826708155e-08,
+      "loss": 0.7254,
+      "output_text_mask_loss": 0.0433,
+      "predictor_learning_rate": 2.210206182670815e-06,
+      "step": 9359
+    },
+    {
+      "epoch": 0.96,
+      "gumbel_tau": 0.1099,
+      "image_masks_loss": 0.0194,
+      "llm_learning_rate": 2.1992126257790735e-08,
+      "loss": 0.7416,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 2.1992126257790734e-06,
+      "step": 9360
+    },
+    {
+      "epoch": 0.96,
+      "gumbel_tau": 0.1099,
+      "image_masks_loss": 0.0136,
+      "llm_learning_rate": 2.1882463572619474e-08,
+      "loss": 0.7242,
+      "output_text_mask_loss": 0.0275,
+      "predictor_learning_rate": 2.188246357261947e-06,
+      "step": 9361
+    },
+    {
+      "epoch": 0.96,
+      "gumbel_tau": 0.1099,
+      "image_masks_loss": 0.0134,
+      "llm_learning_rate": 2.177307378327137e-08,
+      "loss": 0.7381,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 2.177307378327137e-06,
+      "step": 9362
+    },
+    {
+      "epoch": 0.96,
+      "gumbel_tau": 0.1098,
+      "image_masks_loss": 0.004,
+      "llm_learning_rate": 2.166395690179235e-08,
+      "loss": 0.6858,
+      "output_text_mask_loss": 0.0137,
+      "predictor_learning_rate": 2.1663956901792347e-06,
+      "step": 9363
+    },
+    {
+      "epoch": 0.96,
+      "gumbel_tau": 0.1098,
+      "image_masks_loss": 0.003,
+      "llm_learning_rate": 2.155511294019863e-08,
+      "loss": 0.7571,
+      "output_text_mask_loss": 0.0078,
+      "predictor_learning_rate": 2.155511294019863e-06,
+      "step": 9364
+    },
+    {
+      "epoch": 0.96,
+      "gumbel_tau": 0.1098,
+      "image_masks_loss": 0.0078,
+      "llm_learning_rate": 2.144654191047646e-08,
+      "loss": 0.6995,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 2.144654191047646e-06,
+      "step": 9365
+    },
+    {
+      "epoch": 0.96,
+      "gumbel_tau": 0.1098,
+      "image_masks_loss": 0.0103,
+      "llm_learning_rate": 2.1338243824581826e-08,
+      "loss": 0.7045,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 2.1338243824581825e-06,
+      "step": 9366
+    },
+    {
+      "epoch": 0.96,
+      "gumbel_tau": 0.1097,
+      "image_masks_loss": 0.0066,
+      "llm_learning_rate": 2.123021869444075e-08,
+      "loss": 0.6752,
+      "output_text_mask_loss": 0.0065,
+      "predictor_learning_rate": 2.123021869444075e-06,
+      "step": 9367
+    },
+    {
+      "epoch": 0.96,
+      "gumbel_tau": 0.1097,
+      "image_masks_loss": 0.0142,
+      "llm_learning_rate": 2.112246653194927e-08,
+      "loss": 0.7147,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 2.112246653194927e-06,
+      "step": 9368
+    },
+    {
+      "epoch": 0.96,
+      "gumbel_tau": 0.1097,
+      "image_masks_loss": 0.0051,
+      "llm_learning_rate": 2.101498734897317e-08,
+      "loss": 0.5722,
+      "output_text_mask_loss": 0.0326,
+      "predictor_learning_rate": 2.101498734897317e-06,
+      "step": 9369
+    },
+    {
+      "epoch": 0.96,
+      "gumbel_tau": 0.1097,
+      "image_masks_loss": 0.0064,
+      "llm_learning_rate": 2.090778115734854e-08,
+      "loss": 0.7527,
+      "output_text_mask_loss": 0.0422,
+      "predictor_learning_rate": 2.090778115734854e-06,
+      "step": 9370
+    },
+    {
+      "epoch": 0.96,
+      "gumbel_tau": 0.1096,
+      "image_masks_loss": 0.01,
+      "llm_learning_rate": 2.0800847968880933e-08,
+      "loss": 0.6953,
+      "output_text_mask_loss": 0.0172,
+      "predictor_learning_rate": 2.080084796888093e-06,
+      "step": 9371
+    },
+    {
+      "epoch": 0.96,
+      "gumbel_tau": 0.1096,
+      "image_masks_loss": 0.0092,
+      "llm_learning_rate": 2.069418779534621e-08,
+      "loss": 0.7176,
+      "output_text_mask_loss": 0.0124,
+      "predictor_learning_rate": 2.069418779534621e-06,
+      "step": 9372
+    },
+    {
+      "epoch": 0.96,
+      "gumbel_tau": 0.1096,
+      "image_masks_loss": 0.0113,
+      "llm_learning_rate": 2.0587800648489974e-08,
+      "loss": 0.8225,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 2.0587800648489974e-06,
+      "step": 9373
+    },
+    {
+      "epoch": 0.96,
+      "gumbel_tau": 0.1096,
+      "image_masks_loss": 0.0092,
+      "llm_learning_rate": 2.048168654002758e-08,
+      "loss": 0.7476,
+      "output_text_mask_loss": 0.0032,
+      "predictor_learning_rate": 2.0481686540027577e-06,
+      "step": 9374
+    },
+    {
+      "epoch": 0.96,
+      "gumbel_tau": 0.1095,
+      "image_masks_loss": 0.0047,
+      "llm_learning_rate": 2.037584548164495e-08,
+      "loss": 0.7852,
+      "output_text_mask_loss": 0.0519,
+      "predictor_learning_rate": 2.037584548164495e-06,
+      "step": 9375
+    },
+    {
+      "epoch": 0.96,
+      "gumbel_tau": 0.1095,
+      "image_masks_loss": 0.0139,
+      "llm_learning_rate": 2.0270277484997213e-08,
+      "loss": 0.7065,
+      "output_text_mask_loss": 0.001,
+      "predictor_learning_rate": 2.027027748499721e-06,
+      "step": 9376
+    },
+    {
+      "epoch": 0.96,
+      "gumbel_tau": 0.1095,
+      "image_masks_loss": 0.0091,
+      "llm_learning_rate": 2.0164982561710068e-08,
+      "loss": 0.8018,
+      "output_text_mask_loss": 0.025,
+      "predictor_learning_rate": 2.0164982561710068e-06,
+      "step": 9377
+    },
+    {
+      "epoch": 0.96,
+      "gumbel_tau": 0.1095,
+      "image_masks_loss": 0.0034,
+      "llm_learning_rate": 2.005996072337868e-08,
+      "loss": 0.7467,
+      "output_text_mask_loss": 0.0062,
+      "predictor_learning_rate": 2.005996072337868e-06,
+      "step": 9378
+    },
+    {
+      "epoch": 0.96,
+      "gumbel_tau": 0.1094,
+      "image_masks_loss": 0.0063,
+      "llm_learning_rate": 1.9955211981568246e-08,
+      "loss": 0.7256,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 1.9955211981568246e-06,
+      "step": 9379
+    },
+    {
+      "epoch": 0.96,
+      "gumbel_tau": 0.1094,
+      "image_masks_loss": 0.0128,
+      "llm_learning_rate": 1.985073634781398e-08,
+      "loss": 0.6944,
+      "output_text_mask_loss": 0.0052,
+      "predictor_learning_rate": 1.985073634781398e-06,
+      "step": 9380
+    },
+    {
+      "epoch": 0.96,
+      "gumbel_tau": 0.1094,
+      "image_masks_loss": 0.0216,
+      "llm_learning_rate": 1.974653383362085e-08,
+      "loss": 0.7329,
+      "output_text_mask_loss": 0.0214,
+      "predictor_learning_rate": 1.974653383362085e-06,
+      "step": 9381
+    },
+    {
+      "epoch": 0.96,
+      "gumbel_tau": 0.1094,
+      "image_masks_loss": 0.01,
+      "llm_learning_rate": 1.964260445046412e-08,
+      "loss": 0.7502,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 1.964260445046412e-06,
+      "step": 9382
+    },
+    {
+      "epoch": 0.96,
+      "gumbel_tau": 0.1093,
+      "image_masks_loss": 0.0102,
+      "llm_learning_rate": 1.9538948209788522e-08,
+      "loss": 0.7288,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 1.953894820978852e-06,
+      "step": 9383
+    },
+    {
+      "epoch": 0.96,
+      "gumbel_tau": 0.1093,
+      "image_masks_loss": 0.0075,
+      "llm_learning_rate": 1.9435565123008814e-08,
+      "loss": 0.8246,
+      "output_text_mask_loss": 0.0332,
+      "predictor_learning_rate": 1.9435565123008815e-06,
+      "step": 9384
+    },
+    {
+      "epoch": 0.96,
+      "gumbel_tau": 0.1093,
+      "image_masks_loss": 0.0077,
+      "llm_learning_rate": 1.9332455201510058e-08,
+      "loss": 0.7314,
+      "output_text_mask_loss": 0.0254,
+      "predictor_learning_rate": 1.9332455201510056e-06,
+      "step": 9385
+    },
+    {
+      "epoch": 0.96,
+      "gumbel_tau": 0.1093,
+      "image_masks_loss": 0.0164,
+      "llm_learning_rate": 1.9229618456646505e-08,
+      "loss": 0.7343,
+      "output_text_mask_loss": 0.0024,
+      "predictor_learning_rate": 1.9229618456646503e-06,
+      "step": 9386
+    },
+    {
+      "epoch": 0.96,
+      "gumbel_tau": 0.1092,
+      "image_masks_loss": 0.0125,
+      "llm_learning_rate": 1.912705489974326e-08,
+      "loss": 0.6863,
+      "output_text_mask_loss": 0.0376,
+      "predictor_learning_rate": 1.9127054899743257e-06,
+      "step": 9387
+    },
+    {
+      "epoch": 0.96,
+      "gumbel_tau": 0.1092,
+      "image_masks_loss": 0.007,
+      "llm_learning_rate": 1.9024764542094622e-08,
+      "loss": 0.7832,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 1.9024764542094619e-06,
+      "step": 9388
+    },
+    {
+      "epoch": 0.96,
+      "gumbel_tau": 0.1092,
+      "image_masks_loss": 0.0119,
+      "llm_learning_rate": 1.8922747394965192e-08,
+      "loss": 0.7445,
+      "output_text_mask_loss": 0.0301,
+      "predictor_learning_rate": 1.8922747394965189e-06,
+      "step": 9389
+    },
+    {
+      "epoch": 0.96,
+      "gumbel_tau": 0.1091,
+      "image_masks_loss": 0.0075,
+      "llm_learning_rate": 1.8821003469589317e-08,
+      "loss": 0.7245,
+      "output_text_mask_loss": 0.0084,
+      "predictor_learning_rate": 1.8821003469589314e-06,
+      "step": 9390
+    },
+    {
+      "epoch": 0.96,
+      "gumbel_tau": 0.1091,
+      "image_masks_loss": 0.0087,
+      "llm_learning_rate": 1.8719532777170813e-08,
+      "loss": 0.8129,
+      "output_text_mask_loss": 0.0088,
+      "predictor_learning_rate": 1.8719532777170812e-06,
+      "step": 9391
+    },
+    {
+      "epoch": 0.96,
+      "gumbel_tau": 0.1091,
+      "image_masks_loss": 0.0039,
+      "llm_learning_rate": 1.861833532888435e-08,
+      "loss": 0.7495,
+      "output_text_mask_loss": 0.0235,
+      "predictor_learning_rate": 1.861833532888435e-06,
+      "step": 9392
+    },
+    {
+      "epoch": 0.96,
+      "gumbel_tau": 0.1091,
+      "image_masks_loss": 0.0064,
+      "llm_learning_rate": 1.8517411135874075e-08,
+      "loss": 0.7142,
+      "output_text_mask_loss": 0.0025,
+      "predictor_learning_rate": 1.8517411135874073e-06,
+      "step": 9393
+    },
+    {
+      "epoch": 0.96,
+      "gumbel_tau": 0.109,
+      "image_masks_loss": 0.0076,
+      "llm_learning_rate": 1.8416760209253314e-08,
+      "loss": 0.7389,
+      "output_text_mask_loss": 0.0208,
+      "predictor_learning_rate": 1.8416760209253313e-06,
+      "step": 9394
+    },
+    {
+      "epoch": 0.96,
+      "gumbel_tau": 0.109,
+      "image_masks_loss": 0.0065,
+      "llm_learning_rate": 1.8316382560106816e-08,
+      "loss": 0.7524,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 1.8316382560106815e-06,
+      "step": 9395
+    },
+    {
+      "epoch": 0.96,
+      "gumbel_tau": 0.109,
+      "image_masks_loss": 0.0031,
+      "llm_learning_rate": 1.821627819948768e-08,
+      "loss": 0.7219,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 1.8216278199487678e-06,
+      "step": 9396
+    },
+    {
+      "epoch": 0.96,
+      "gumbel_tau": 0.109,
+      "image_masks_loss": 0.0071,
+      "llm_learning_rate": 1.8116447138420422e-08,
+      "loss": 0.552,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 1.8116447138420423e-06,
+      "step": 9397
+    },
+    {
+      "epoch": 0.96,
+      "gumbel_tau": 0.1089,
+      "image_masks_loss": 0.0115,
+      "llm_learning_rate": 1.801688938789764e-08,
+      "loss": 0.7849,
+      "output_text_mask_loss": 0.0019,
+      "predictor_learning_rate": 1.801688938789764e-06,
+      "step": 9398
+    },
+    {
+      "epoch": 0.96,
+      "gumbel_tau": 0.1089,
+      "image_masks_loss": 0.0116,
+      "llm_learning_rate": 1.7917604958883894e-08,
+      "loss": 0.676,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 1.7917604958883892e-06,
+      "step": 9399
+    },
+    {
+      "epoch": 0.96,
+      "gumbel_tau": 0.1089,
+      "image_masks_loss": 0.0026,
+      "llm_learning_rate": 1.781859386231183e-08,
+      "loss": 0.7555,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 1.7818593862311827e-06,
+      "step": 9400
+    },
+    {
+      "epoch": 0.96,
+      "gumbel_tau": 0.1089,
+      "image_masks_loss": 0.0165,
+      "llm_learning_rate": 1.7719856109084943e-08,
+      "loss": 0.7484,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 1.7719856109084941e-06,
+      "step": 9401
+    },
+    {
+      "epoch": 0.96,
+      "gumbel_tau": 0.1088,
+      "image_masks_loss": 0.0025,
+      "llm_learning_rate": 1.7621391710076764e-08,
+      "loss": 0.7133,
+      "output_text_mask_loss": 0.0618,
+      "predictor_learning_rate": 1.7621391710076762e-06,
+      "step": 9402
+    },
+    {
+      "epoch": 0.96,
+      "gumbel_tau": 0.1088,
+      "image_masks_loss": 0.0233,
+      "llm_learning_rate": 1.752320067613028e-08,
+      "loss": 0.6871,
+      "output_text_mask_loss": 0.0071,
+      "predictor_learning_rate": 1.752320067613028e-06,
+      "step": 9403
+    },
+    {
+      "epoch": 0.96,
+      "gumbel_tau": 0.1088,
+      "image_masks_loss": 0.0282,
+      "llm_learning_rate": 1.7425283018058236e-08,
+      "loss": 0.6992,
+      "output_text_mask_loss": 0.0081,
+      "predictor_learning_rate": 1.7425283018058236e-06,
+      "step": 9404
+    },
+    {
+      "epoch": 0.96,
+      "gumbel_tau": 0.1088,
+      "image_masks_loss": 0.0076,
+      "llm_learning_rate": 1.7327638746643947e-08,
+      "loss": 0.7169,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 1.7327638746643948e-06,
+      "step": 9405
+    },
+    {
+      "epoch": 0.96,
+      "gumbel_tau": 0.1087,
+      "image_masks_loss": 0.0023,
+      "llm_learning_rate": 1.7230267872639926e-08,
+      "loss": 0.6933,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 1.7230267872639926e-06,
+      "step": 9406
+    },
+    {
+      "epoch": 0.96,
+      "gumbel_tau": 0.1087,
+      "image_masks_loss": 0.0202,
+      "llm_learning_rate": 1.7133170406768706e-08,
+      "loss": 0.7403,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 1.7133170406768705e-06,
+      "step": 9407
+    },
+    {
+      "epoch": 0.96,
+      "gumbel_tau": 0.1087,
+      "image_masks_loss": 0.0105,
+      "llm_learning_rate": 1.70363463597234e-08,
+      "loss": 0.6804,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 1.7036346359723399e-06,
+      "step": 9408
+    },
+    {
+      "epoch": 0.96,
+      "gumbel_tau": 0.1087,
+      "image_masks_loss": 0.0186,
+      "llm_learning_rate": 1.693979574216631e-08,
+      "loss": 0.6895,
+      "output_text_mask_loss": 0.0119,
+      "predictor_learning_rate": 1.6939795742166308e-06,
+      "step": 9409
+    },
+    {
+      "epoch": 0.96,
+      "gumbel_tau": 0.1086,
+      "image_masks_loss": 0.005,
+      "llm_learning_rate": 1.684351856472949e-08,
+      "loss": 0.7342,
+      "output_text_mask_loss": 0.0149,
+      "predictor_learning_rate": 1.6843518564729487e-06,
+      "step": 9410
+    },
+    {
+      "epoch": 0.96,
+      "gumbel_tau": 0.1086,
+      "image_masks_loss": 0.0046,
+      "llm_learning_rate": 1.6747514838015566e-08,
+      "loss": 0.745,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 1.6747514838015565e-06,
+      "step": 9411
+    },
+    {
+      "epoch": 0.96,
+      "gumbel_tau": 0.1086,
+      "image_masks_loss": 0.0075,
+      "llm_learning_rate": 1.665178457259664e-08,
+      "loss": 0.7395,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 1.665178457259664e-06,
+      "step": 9412
+    },
+    {
+      "epoch": 0.96,
+      "gumbel_tau": 0.1086,
+      "image_masks_loss": 0.0058,
+      "llm_learning_rate": 1.655632777901428e-08,
+      "loss": 0.7718,
+      "output_text_mask_loss": 0.0078,
+      "predictor_learning_rate": 1.655632777901428e-06,
+      "step": 9413
+    },
+    {
+      "epoch": 0.96,
+      "gumbel_tau": 0.1085,
+      "image_masks_loss": 0.0046,
+      "llm_learning_rate": 1.646114446778091e-08,
+      "loss": 0.6292,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 1.646114446778091e-06,
+      "step": 9414
+    },
+    {
+      "epoch": 0.96,
+      "gumbel_tau": 0.1085,
+      "image_masks_loss": 0.0075,
+      "llm_learning_rate": 1.6366234649378144e-08,
+      "loss": 0.7803,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 1.6366234649378143e-06,
+      "step": 9415
+    },
+    {
+      "epoch": 0.96,
+      "gumbel_tau": 0.1085,
+      "image_masks_loss": 0.0146,
+      "llm_learning_rate": 1.6271598334257898e-08,
+      "loss": 0.692,
+      "output_text_mask_loss": 0.0029,
+      "predictor_learning_rate": 1.6271598334257899e-06,
+      "step": 9416
+    },
+    {
+      "epoch": 0.96,
+      "gumbel_tau": 0.1085,
+      "image_masks_loss": 0.0129,
+      "llm_learning_rate": 1.6177235532841284e-08,
+      "loss": 0.761,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 1.6177235532841283e-06,
+      "step": 9417
+    },
+    {
+      "epoch": 0.96,
+      "gumbel_tau": 0.1084,
+      "image_masks_loss": 0.0044,
+      "llm_learning_rate": 1.6083146255520266e-08,
+      "loss": 0.719,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 1.6083146255520265e-06,
+      "step": 9418
+    },
+    {
+      "epoch": 0.97,
+      "gumbel_tau": 0.1084,
+      "image_masks_loss": 0.0086,
+      "llm_learning_rate": 1.598933051265572e-08,
+      "loss": 0.725,
+      "output_text_mask_loss": 0.0107,
+      "predictor_learning_rate": 1.598933051265572e-06,
+      "step": 9419
+    },
+    {
+      "epoch": 0.97,
+      "gumbel_tau": 0.1084,
+      "image_masks_loss": 0.0011,
+      "llm_learning_rate": 1.5895788314579107e-08,
+      "loss": 0.7376,
+      "output_text_mask_loss": 0.0063,
+      "predictor_learning_rate": 1.5895788314579106e-06,
+      "step": 9420
+    },
+    {
+      "epoch": 0.97,
+      "gumbel_tau": 0.1084,
+      "image_masks_loss": 0.0087,
+      "llm_learning_rate": 1.5802519671591355e-08,
+      "loss": 0.7207,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 1.5802519671591353e-06,
+      "step": 9421
+    },
+    {
+      "epoch": 0.97,
+      "gumbel_tau": 0.1083,
+      "image_masks_loss": 0.0216,
+      "llm_learning_rate": 1.570952459396341e-08,
+      "loss": 0.7571,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 1.570952459396341e-06,
+      "step": 9422
+    },
+    {
+      "epoch": 0.97,
+      "gumbel_tau": 0.1083,
+      "image_masks_loss": 0.006,
+      "llm_learning_rate": 1.561680309193625e-08,
+      "loss": 0.7758,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 1.561680309193625e-06,
+      "step": 9423
+    },
+    {
+      "epoch": 0.97,
+      "gumbel_tau": 0.1083,
+      "image_masks_loss": 0.0103,
+      "llm_learning_rate": 1.5524355175720318e-08,
+      "loss": 0.752,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 1.5524355175720318e-06,
+      "step": 9424
+    },
+    {
+      "epoch": 0.97,
+      "gumbel_tau": 0.1082,
+      "image_masks_loss": 0.0112,
+      "llm_learning_rate": 1.5432180855496636e-08,
+      "loss": 0.6154,
+      "output_text_mask_loss": 0.0231,
+      "predictor_learning_rate": 1.5432180855496636e-06,
+      "step": 9425
+    },
+    {
+      "epoch": 0.97,
+      "gumbel_tau": 0.1082,
+      "image_masks_loss": 0.0082,
+      "llm_learning_rate": 1.5340280141415142e-08,
+      "loss": 0.7076,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 1.5340280141415142e-06,
+      "step": 9426
+    },
+    {
+      "epoch": 0.97,
+      "gumbel_tau": 0.1082,
+      "image_masks_loss": 0.0093,
+      "llm_learning_rate": 1.524865304359635e-08,
+      "loss": 0.7854,
+      "output_text_mask_loss": 0.0156,
+      "predictor_learning_rate": 1.524865304359635e-06,
+      "step": 9427
+    },
+    {
+      "epoch": 0.97,
+      "gumbel_tau": 0.1082,
+      "image_masks_loss": 0.0044,
+      "llm_learning_rate": 1.5157299572130522e-08,
+      "loss": 0.7593,
+      "output_text_mask_loss": 0.0137,
+      "predictor_learning_rate": 1.5157299572130524e-06,
+      "step": 9428
+    },
+    {
+      "epoch": 0.97,
+      "gumbel_tau": 0.1081,
+      "image_masks_loss": 0.0107,
+      "llm_learning_rate": 1.506621973707767e-08,
+      "loss": 0.6983,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 1.5066219737077668e-06,
+      "step": 9429
+    },
+    {
+      "epoch": 0.97,
+      "gumbel_tau": 0.1081,
+      "image_masks_loss": 0.0036,
+      "llm_learning_rate": 1.4975413548467544e-08,
+      "loss": 0.701,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 1.4975413548467542e-06,
+      "step": 9430
+    },
+    {
+      "epoch": 0.97,
+      "gumbel_tau": 0.1081,
+      "image_masks_loss": 0.0073,
+      "llm_learning_rate": 1.4884881016300202e-08,
+      "loss": 0.6713,
+      "output_text_mask_loss": 0.021,
+      "predictor_learning_rate": 1.48848810163002e-06,
+      "step": 9431
+    },
+    {
+      "epoch": 0.97,
+      "gumbel_tau": 0.1081,
+      "image_masks_loss": 0.0062,
+      "llm_learning_rate": 1.4794622150544891e-08,
+      "loss": 0.6484,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 1.479462215054489e-06,
+      "step": 9432
+    },
+    {
+      "epoch": 0.97,
+      "gumbel_tau": 0.108,
+      "image_masks_loss": 0.0111,
+      "llm_learning_rate": 1.4704636961141439e-08,
+      "loss": 0.6806,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 1.4704636961141438e-06,
+      "step": 9433
+    },
+    {
+      "epoch": 0.97,
+      "gumbel_tau": 0.108,
+      "image_masks_loss": 0.0035,
+      "llm_learning_rate": 1.4614925457999418e-08,
+      "loss": 0.7017,
+      "output_text_mask_loss": 0.0073,
+      "predictor_learning_rate": 1.4614925457999418e-06,
+      "step": 9434
+    },
+    {
+      "epoch": 0.97,
+      "gumbel_tau": 0.108,
+      "image_masks_loss": 0.0056,
+      "llm_learning_rate": 1.4525487650997593e-08,
+      "loss": 0.7486,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 1.4525487650997593e-06,
+      "step": 9435
+    },
+    {
+      "epoch": 0.97,
+      "gumbel_tau": 0.108,
+      "image_masks_loss": 0.0239,
+      "llm_learning_rate": 1.443632354998531e-08,
+      "loss": 0.7689,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 1.4436323549985308e-06,
+      "step": 9436
+    },
+    {
+      "epoch": 0.97,
+      "gumbel_tau": 0.1079,
+      "image_masks_loss": 0.0192,
+      "llm_learning_rate": 1.4347433164781378e-08,
+      "loss": 0.7009,
+      "output_text_mask_loss": 0.0303,
+      "predictor_learning_rate": 1.4347433164781376e-06,
+      "step": 9437
+    },
+    {
+      "epoch": 0.97,
+      "gumbel_tau": 0.1079,
+      "image_masks_loss": 0.0054,
+      "llm_learning_rate": 1.4258816505174911e-08,
+      "loss": 0.6675,
+      "output_text_mask_loss": 0.0326,
+      "predictor_learning_rate": 1.4258816505174909e-06,
+      "step": 9438
+    },
+    {
+      "epoch": 0.97,
+      "gumbel_tau": 0.1079,
+      "image_masks_loss": 0.0143,
+      "llm_learning_rate": 1.4170473580923937e-08,
+      "loss": 0.6566,
+      "output_text_mask_loss": 0.0877,
+      "predictor_learning_rate": 1.4170473580923937e-06,
+      "step": 9439
+    },
+    {
+      "epoch": 0.97,
+      "gumbel_tau": 0.1079,
+      "image_masks_loss": 0.0052,
+      "llm_learning_rate": 1.4082404401757621e-08,
+      "loss": 0.7161,
+      "output_text_mask_loss": 0.0195,
+      "predictor_learning_rate": 1.408240440175762e-06,
+      "step": 9440
+    },
+    {
+      "epoch": 0.97,
+      "gumbel_tau": 0.1078,
+      "image_masks_loss": 0.008,
+      "llm_learning_rate": 1.3994608977374036e-08,
+      "loss": 0.774,
+      "output_text_mask_loss": 0.0088,
+      "predictor_learning_rate": 1.3994608977374034e-06,
+      "step": 9441
+    },
+    {
+      "epoch": 0.97,
+      "gumbel_tau": 0.1078,
+      "image_masks_loss": 0.0223,
+      "llm_learning_rate": 1.390708731744156e-08,
+      "loss": 0.6611,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 1.3907087317441558e-06,
+      "step": 9442
+    },
+    {
+      "epoch": 0.97,
+      "gumbel_tau": 0.1078,
+      "image_masks_loss": 0.0117,
+      "llm_learning_rate": 1.3819839431598037e-08,
+      "loss": 0.7271,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 1.3819839431598035e-06,
+      "step": 9443
+    },
+    {
+      "epoch": 0.97,
+      "gumbel_tau": 0.1078,
+      "image_masks_loss": 0.01,
+      "llm_learning_rate": 1.3732865329451617e-08,
+      "loss": 0.6753,
+      "output_text_mask_loss": NaN,
+      "predictor_learning_rate": 1.3732865329451616e-06,
+      "step": 9444
+    },
+    {
+      "epoch": 0.97,
+      "gumbel_tau": 0.1077,
+      "image_masks_loss": 0.0067,
+      "llm_learning_rate": 1.3646165020580194e-08,
+      "loss": 0.7971,
+      "output_text_mask_loss": 0.0427,
+      "predictor_learning_rate": 1.3646165020580192e-06,
+      "step": 9445
+    },
+    {
+      "epoch": 0.97,
+      "gumbel_tau": 0.1077,
+      "image_masks_loss": 0.0071,
+      "llm_learning_rate": 1.3559738514530852e-08,
+      "loss": 0.7181,
+      "output_text_mask_loss": 0.0004,
+      "predictor_learning_rate": 1.355973851453085e-06,
+      "step": 9446
+    },
+    {
+      "epoch": 0.97,
+      "gumbel_tau": 0.1077,
+      "image_masks_loss": 0.0121,
+      "llm_learning_rate": 1.3473585820821533e-08,
+      "loss": 0.7348,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 1.3473585820821532e-06,
+      "step": 9447
+    },
+    {
+      "epoch": 0.97,
+      "gumbel_tau": 0.1077,
+      "image_masks_loss": 0.0073,
+      "llm_learning_rate": 1.338770694893965e-08,
+      "loss": 0.7176,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 1.338770694893965e-06,
+      "step": 9448
+    },
+    {
+      "epoch": 0.97,
+      "gumbel_tau": 0.1076,
+      "image_masks_loss": 0.0062,
+      "llm_learning_rate": 1.330210190834208e-08,
+      "loss": 0.8041,
+      "output_text_mask_loss": 0.0217,
+      "predictor_learning_rate": 1.330210190834208e-06,
+      "step": 9449
+    },
+    {
+      "epoch": 0.97,
+      "gumbel_tau": 0.1076,
+      "image_masks_loss": 0.0096,
+      "llm_learning_rate": 1.3216770708456006e-08,
+      "loss": 0.7007,
+      "output_text_mask_loss": 0.0039,
+      "predictor_learning_rate": 1.3216770708456004e-06,
+      "step": 9450
+    },
+    {
+      "epoch": 0.97,
+      "gumbel_tau": 0.1076,
+      "image_masks_loss": 0.0141,
+      "llm_learning_rate": 1.3131713358678356e-08,
+      "loss": 0.7401,
+      "output_text_mask_loss": 0.0187,
+      "predictor_learning_rate": 1.3131713358678355e-06,
+      "step": 9451
+    },
+    {
+      "epoch": 0.97,
+      "gumbel_tau": 0.1076,
+      "image_masks_loss": 0.0058,
+      "llm_learning_rate": 1.3046929868375524e-08,
+      "loss": 0.7128,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 1.3046929868375523e-06,
+      "step": 9452
+    },
+    {
+      "epoch": 0.97,
+      "gumbel_tau": 0.1075,
+      "image_masks_loss": 0.0038,
+      "llm_learning_rate": 1.296242024688421e-08,
+      "loss": 0.7038,
+      "output_text_mask_loss": 0.0202,
+      "predictor_learning_rate": 1.2962420246884211e-06,
+      "step": 9453
+    },
+    {
+      "epoch": 0.97,
+      "gumbel_tau": 0.1075,
+      "image_masks_loss": 0.0121,
+      "llm_learning_rate": 1.287818450351086e-08,
+      "loss": 0.744,
+      "output_text_mask_loss": 0.0141,
+      "predictor_learning_rate": 1.2878184503510858e-06,
+      "step": 9454
+    },
+    {
+      "epoch": 0.97,
+      "gumbel_tau": 0.1075,
+      "image_masks_loss": 0.02,
+      "llm_learning_rate": 1.2794222647531663e-08,
+      "loss": 0.6751,
+      "output_text_mask_loss": 0.0172,
+      "predictor_learning_rate": 1.279422264753166e-06,
+      "step": 9455
+    },
+    {
+      "epoch": 0.97,
+      "gumbel_tau": 0.1075,
+      "image_masks_loss": 0.0044,
+      "llm_learning_rate": 1.2710534688192833e-08,
+      "loss": 0.8432,
+      "output_text_mask_loss": 0.0778,
+      "predictor_learning_rate": 1.2710534688192833e-06,
+      "step": 9456
+    },
+    {
+      "epoch": 0.97,
+      "gumbel_tau": 0.1074,
+      "image_masks_loss": 0.0078,
+      "llm_learning_rate": 1.2627120634710055e-08,
+      "loss": 0.6252,
+      "output_text_mask_loss": 0.0078,
+      "predictor_learning_rate": 1.2627120634710054e-06,
+      "step": 9457
+    },
+    {
+      "epoch": 0.97,
+      "gumbel_tau": 0.1074,
+      "image_masks_loss": 0.012,
+      "llm_learning_rate": 1.2543980496269037e-08,
+      "loss": 0.6782,
+      "output_text_mask_loss": 0.0034,
+      "predictor_learning_rate": 1.2543980496269037e-06,
+      "step": 9458
+    },
+    {
+      "epoch": 0.97,
+      "gumbel_tau": 0.1074,
+      "image_masks_loss": 0.0087,
+      "llm_learning_rate": 1.2461114282025789e-08,
+      "loss": 0.6872,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 1.2461114282025787e-06,
+      "step": 9459
+    },
+    {
+      "epoch": 0.97,
+      "gumbel_tau": 0.1074,
+      "image_masks_loss": 0.0153,
+      "llm_learning_rate": 1.237852200110523e-08,
+      "loss": 0.711,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 1.237852200110523e-06,
+      "step": 9460
+    },
+    {
+      "epoch": 0.97,
+      "gumbel_tau": 0.1073,
+      "image_masks_loss": 0.0211,
+      "llm_learning_rate": 1.2296203662602869e-08,
+      "loss": 0.6727,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 1.2296203662602867e-06,
+      "step": 9461
+    },
+    {
+      "epoch": 0.97,
+      "gumbel_tau": 0.1073,
+      "image_masks_loss": 0.0073,
+      "llm_learning_rate": 1.2214159275583392e-08,
+      "loss": 0.6871,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 1.2214159275583393e-06,
+      "step": 9462
+    },
+    {
+      "epoch": 0.97,
+      "gumbel_tau": 0.1073,
+      "image_masks_loss": 0.0082,
+      "llm_learning_rate": 1.2132388849082355e-08,
+      "loss": 0.7519,
+      "output_text_mask_loss": 0.0027,
+      "predictor_learning_rate": 1.2132388849082354e-06,
+      "step": 9463
+    },
+    {
+      "epoch": 0.97,
+      "gumbel_tau": 0.1073,
+      "image_masks_loss": 0.0078,
+      "llm_learning_rate": 1.2050892392104219e-08,
+      "loss": 0.7578,
+      "output_text_mask_loss": 0.0115,
+      "predictor_learning_rate": 1.2050892392104218e-06,
+      "step": 9464
+    },
+    {
+      "epoch": 0.97,
+      "gumbel_tau": 0.1072,
+      "image_masks_loss": 0.0134,
+      "llm_learning_rate": 1.1969669913623194e-08,
+      "loss": 0.8035,
+      "output_text_mask_loss": 0.0132,
+      "predictor_learning_rate": 1.1969669913623193e-06,
+      "step": 9465
+    },
+    {
+      "epoch": 0.97,
+      "gumbel_tau": 0.1072,
+      "image_masks_loss": 0.0074,
+      "llm_learning_rate": 1.1888721422584071e-08,
+      "loss": 0.7567,
+      "output_text_mask_loss": 0.0496,
+      "predictor_learning_rate": 1.188872142258407e-06,
+      "step": 9466
+    },
+    {
+      "epoch": 0.97,
+      "gumbel_tau": 0.1072,
+      "image_masks_loss": 0.0044,
+      "llm_learning_rate": 1.1808046927901107e-08,
+      "loss": 0.7605,
+      "output_text_mask_loss": 0.0231,
+      "predictor_learning_rate": 1.1808046927901107e-06,
+      "step": 9467
+    },
+    {
+      "epoch": 0.97,
+      "gumbel_tau": 0.1072,
+      "image_masks_loss": 0.0042,
+      "llm_learning_rate": 1.1727646438458307e-08,
+      "loss": 0.6406,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 1.1727646438458306e-06,
+      "step": 9468
+    },
+    {
+      "epoch": 0.97,
+      "gumbel_tau": 0.1071,
+      "image_masks_loss": 0.0206,
+      "llm_learning_rate": 1.1647519963109421e-08,
+      "loss": 0.7933,
+      "output_text_mask_loss": 0.0004,
+      "predictor_learning_rate": 1.1647519963109422e-06,
+      "step": 9469
+    },
+    {
+      "epoch": 0.97,
+      "gumbel_tau": 0.1071,
+      "image_masks_loss": 0.0046,
+      "llm_learning_rate": 1.1567667510678226e-08,
+      "loss": 0.8251,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 1.1567667510678226e-06,
+      "step": 9470
+    },
+    {
+      "epoch": 0.97,
+      "gumbel_tau": 0.1071,
+      "image_masks_loss": 0.0117,
+      "llm_learning_rate": 1.1488089089958243e-08,
+      "loss": 0.6264,
+      "output_text_mask_loss": 0.0054,
+      "predictor_learning_rate": 1.1488089089958243e-06,
+      "step": 9471
+    },
+    {
+      "epoch": 0.97,
+      "gumbel_tau": 0.1071,
+      "image_masks_loss": 0.0041,
+      "llm_learning_rate": 1.1408784709712739e-08,
+      "loss": 0.6907,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 1.1408784709712738e-06,
+      "step": 9472
+    },
+    {
+      "epoch": 0.97,
+      "gumbel_tau": 0.107,
+      "image_masks_loss": 0.0163,
+      "llm_learning_rate": 1.1329754378675007e-08,
+      "loss": 0.6789,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 1.1329754378675006e-06,
+      "step": 9473
+    },
+    {
+      "epoch": 0.97,
+      "gumbel_tau": 0.107,
+      "image_masks_loss": 0.0016,
+      "llm_learning_rate": 1.1250998105548083e-08,
+      "loss": 0.7523,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 1.1250998105548083e-06,
+      "step": 9474
+    },
+    {
+      "epoch": 0.97,
+      "gumbel_tau": 0.107,
+      "image_masks_loss": 0.0063,
+      "llm_learning_rate": 1.1172515899004754e-08,
+      "loss": 0.743,
+      "output_text_mask_loss": 0.0115,
+      "predictor_learning_rate": 1.1172515899004753e-06,
+      "step": 9475
+    },
+    {
+      "epoch": 0.97,
+      "gumbel_tau": 0.107,
+      "image_masks_loss": 0.0038,
+      "llm_learning_rate": 1.1094307767687551e-08,
+      "loss": 0.717,
+      "output_text_mask_loss": 0.037,
+      "predictor_learning_rate": 1.109430776768755e-06,
+      "step": 9476
+    },
+    {
+      "epoch": 0.97,
+      "gumbel_tau": 0.1069,
+      "image_masks_loss": 0.0091,
+      "llm_learning_rate": 1.1016373720209027e-08,
+      "loss": 0.7372,
+      "output_text_mask_loss": 0.0023,
+      "predictor_learning_rate": 1.1016373720209027e-06,
+      "step": 9477
+    },
+    {
+      "epoch": 0.97,
+      "gumbel_tau": 0.1069,
+      "image_masks_loss": 0.0164,
+      "llm_learning_rate": 1.093871376515121e-08,
+      "loss": 0.7109,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 1.093871376515121e-06,
+      "step": 9478
+    },
+    {
+      "epoch": 0.97,
+      "gumbel_tau": 0.1069,
+      "image_masks_loss": 0.0092,
+      "llm_learning_rate": 1.08613279110667e-08,
+      "loss": 0.7232,
+      "output_text_mask_loss": 0.0025,
+      "predictor_learning_rate": 1.08613279110667e-06,
+      "step": 9479
+    },
+    {
+      "epoch": 0.97,
+      "gumbel_tau": 0.1069,
+      "image_masks_loss": 0.0073,
+      "llm_learning_rate": 1.0784216166477012e-08,
+      "loss": 0.7273,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 1.0784216166477013e-06,
+      "step": 9480
+    },
+    {
+      "epoch": 0.97,
+      "gumbel_tau": 0.1068,
+      "image_masks_loss": 0.0029,
+      "llm_learning_rate": 1.070737853987397e-08,
+      "loss": 0.8027,
+      "output_text_mask_loss": 0.0139,
+      "predictor_learning_rate": 1.0707378539873969e-06,
+      "step": 9481
+    },
+    {
+      "epoch": 0.97,
+      "gumbel_tau": 0.1068,
+      "image_masks_loss": 0.0108,
+      "llm_learning_rate": 1.0630815039719133e-08,
+      "loss": 0.7165,
+      "output_text_mask_loss": 0.0124,
+      "predictor_learning_rate": 1.0630815039719133e-06,
+      "step": 9482
+    },
+    {
+      "epoch": 0.97,
+      "gumbel_tau": 0.1068,
+      "image_masks_loss": 0.0043,
+      "llm_learning_rate": 1.0554525674443539e-08,
+      "loss": 0.7227,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 1.0554525674443538e-06,
+      "step": 9483
+    },
+    {
+      "epoch": 0.97,
+      "gumbel_tau": 0.1068,
+      "image_masks_loss": 0.014,
+      "llm_learning_rate": 1.0478510452448798e-08,
+      "loss": 0.7112,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 1.0478510452448797e-06,
+      "step": 9484
+    },
+    {
+      "epoch": 0.97,
+      "gumbel_tau": 0.1067,
+      "image_masks_loss": 0.0334,
+      "llm_learning_rate": 1.0402769382105716e-08,
+      "loss": 0.7993,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 1.0402769382105715e-06,
+      "step": 9485
+    },
+    {
+      "epoch": 0.97,
+      "gumbel_tau": 0.1067,
+      "image_masks_loss": 0.014,
+      "llm_learning_rate": 1.032730247175484e-08,
+      "loss": 0.7241,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 1.032730247175484e-06,
+      "step": 9486
+    },
+    {
+      "epoch": 0.97,
+      "gumbel_tau": 0.1067,
+      "image_masks_loss": 0.0039,
+      "llm_learning_rate": 1.0252109729707305e-08,
+      "loss": 0.7724,
+      "output_text_mask_loss": 0.0158,
+      "predictor_learning_rate": 1.0252109729707304e-06,
+      "step": 9487
+    },
+    {
+      "epoch": 0.97,
+      "gumbel_tau": 0.1067,
+      "image_masks_loss": 0.0271,
+      "llm_learning_rate": 1.0177191164242872e-08,
+      "loss": 0.7076,
+      "output_text_mask_loss": 0.005,
+      "predictor_learning_rate": 1.0177191164242872e-06,
+      "step": 9488
+    },
+    {
+      "epoch": 0.97,
+      "gumbel_tau": 0.1066,
+      "image_masks_loss": 0.0099,
+      "llm_learning_rate": 1.0102546783611888e-08,
+      "loss": 0.7093,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 1.0102546783611888e-06,
+      "step": 9489
+    },
+    {
+      "epoch": 0.97,
+      "gumbel_tau": 0.1066,
+      "image_masks_loss": 0.0054,
+      "llm_learning_rate": 1.0028176596034722e-08,
+      "loss": 0.8161,
+      "output_text_mask_loss": 0.02,
+      "predictor_learning_rate": 1.0028176596034722e-06,
+      "step": 9490
+    },
+    {
+      "epoch": 0.97,
+      "gumbel_tau": 0.1066,
+      "image_masks_loss": 0.0112,
+      "llm_learning_rate": 9.954080609700933e-09,
+      "loss": 0.7995,
+      "output_text_mask_loss": 0.0124,
+      "predictor_learning_rate": 9.954080609700933e-07,
+      "step": 9491
+    },
+    {
+      "epoch": 0.97,
+      "gumbel_tau": 0.1066,
+      "image_masks_loss": 0.0082,
+      "llm_learning_rate": 9.880258832770106e-09,
+      "loss": 0.7981,
+      "output_text_mask_loss": 0.0565,
+      "predictor_learning_rate": 9.880258832770106e-07,
+      "step": 9492
+    },
+    {
+      "epoch": 0.97,
+      "gumbel_tau": 0.1065,
+      "image_masks_loss": 0.0215,
+      "llm_learning_rate": 9.806711273371572e-09,
+      "loss": 0.7432,
+      "output_text_mask_loss": 0.004,
+      "predictor_learning_rate": 9.806711273371572e-07,
+      "step": 9493
+    },
+    {
+      "epoch": 0.97,
+      "gumbel_tau": 0.1065,
+      "image_masks_loss": 0.0052,
+      "llm_learning_rate": 9.733437939604962e-09,
+      "loss": 0.5389,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 9.733437939604962e-07,
+      "step": 9494
+    },
+    {
+      "epoch": 0.97,
+      "gumbel_tau": 0.1065,
+      "image_masks_loss": 0.0036,
+      "llm_learning_rate": 9.66043883953882e-09,
+      "loss": 0.7599,
+      "output_text_mask_loss": 0.0004,
+      "predictor_learning_rate": 9.660438839538821e-07,
+      "step": 9495
+    },
+    {
+      "epoch": 0.97,
+      "gumbel_tau": 0.1065,
+      "image_masks_loss": 0.0093,
+      "llm_learning_rate": 9.587713981212276e-09,
+      "loss": 0.6785,
+      "output_text_mask_loss": 0.0181,
+      "predictor_learning_rate": 9.587713981212276e-07,
+      "step": 9496
+    },
+    {
+      "epoch": 0.97,
+      "gumbel_tau": 0.1064,
+      "image_masks_loss": 0.0101,
+      "llm_learning_rate": 9.515263372633643e-09,
+      "loss": 0.6633,
+      "output_text_mask_loss": 0.0256,
+      "predictor_learning_rate": 9.515263372633642e-07,
+      "step": 9497
+    },
+    {
+      "epoch": 0.97,
+      "gumbel_tau": 0.1064,
+      "image_masks_loss": 0.0152,
+      "llm_learning_rate": 9.443087021781538e-09,
+      "loss": 0.6733,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 9.443087021781538e-07,
+      "step": 9498
+    },
+    {
+      "epoch": 0.97,
+      "gumbel_tau": 0.1064,
+      "image_masks_loss": 0.0075,
+      "llm_learning_rate": 9.371184936604328e-09,
+      "loss": 0.6806,
+      "output_text_mask_loss": 0.0389,
+      "predictor_learning_rate": 9.371184936604327e-07,
+      "step": 9499
+    },
+    {
+      "epoch": 0.97,
+      "gumbel_tau": 0.1064,
+      "image_masks_loss": 0.0121,
+      "llm_learning_rate": 9.299557125019843e-09,
+      "loss": 0.7538,
+      "output_text_mask_loss": 0.0275,
+      "predictor_learning_rate": 9.299557125019842e-07,
+      "step": 9500
+    },
+    {
+      "epoch": 0.97,
+      "gumbel_tau": 0.1063,
+      "image_masks_loss": 0.0115,
+      "llm_learning_rate": 9.228203594915942e-09,
+      "loss": 0.7859,
+      "output_text_mask_loss": 0.0041,
+      "predictor_learning_rate": 9.22820359491594e-07,
+      "step": 9501
+    },
+    {
+      "epoch": 0.97,
+      "gumbel_tau": 0.1063,
+      "image_masks_loss": 0.0082,
+      "llm_learning_rate": 9.157124354150504e-09,
+      "loss": 0.7657,
+      "output_text_mask_loss": 0.0019,
+      "predictor_learning_rate": 9.157124354150504e-07,
+      "step": 9502
+    },
+    {
+      "epoch": 0.97,
+      "gumbel_tau": 0.1063,
+      "image_masks_loss": 0.0031,
+      "llm_learning_rate": 9.086319410550327e-09,
+      "loss": 0.7258,
+      "output_text_mask_loss": 0.0099,
+      "predictor_learning_rate": 9.086319410550326e-07,
+      "step": 9503
+    },
+    {
+      "epoch": 0.97,
+      "gumbel_tau": 0.1063,
+      "image_masks_loss": 0.0074,
+      "llm_learning_rate": 9.015788771913059e-09,
+      "loss": 0.6558,
+      "output_text_mask_loss": 0.0084,
+      "predictor_learning_rate": 9.015788771913058e-07,
+      "step": 9504
+    },
+    {
+      "epoch": 0.97,
+      "gumbel_tau": 0.1062,
+      "image_masks_loss": 0.0051,
+      "llm_learning_rate": 8.945532446005545e-09,
+      "loss": 0.709,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 8.945532446005544e-07,
+      "step": 9505
+    },
+    {
+      "epoch": 0.97,
+      "gumbel_tau": 0.1062,
+      "image_masks_loss": 0.0041,
+      "llm_learning_rate": 8.875550440564651e-09,
+      "loss": 0.7156,
+      "output_text_mask_loss": 0.0065,
+      "predictor_learning_rate": 8.875550440564651e-07,
+      "step": 9506
+    },
+    {
+      "epoch": 0.97,
+      "gumbel_tau": 0.1062,
+      "image_masks_loss": 0.0019,
+      "llm_learning_rate": 8.80584276329699e-09,
+      "loss": 0.7065,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 8.805842763296989e-07,
+      "step": 9507
+    },
+    {
+      "epoch": 0.97,
+      "gumbel_tau": 0.1062,
+      "image_masks_loss": 0.0083,
+      "llm_learning_rate": 8.736409421878644e-09,
+      "loss": 0.6996,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 8.736409421878643e-07,
+      "step": 9508
+    },
+    {
+      "epoch": 0.97,
+      "gumbel_tau": 0.1061,
+      "image_masks_loss": 0.0147,
+      "llm_learning_rate": 8.667250423955997e-09,
+      "loss": 0.6954,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 8.667250423955996e-07,
+      "step": 9509
+    },
+    {
+      "epoch": 0.97,
+      "gumbel_tau": 0.1061,
+      "image_masks_loss": 0.0128,
+      "llm_learning_rate": 8.59836577714518e-09,
+      "loss": 0.7199,
+      "output_text_mask_loss": 0.0156,
+      "predictor_learning_rate": 8.598365777145178e-07,
+      "step": 9510
+    },
+    {
+      "epoch": 0.97,
+      "gumbel_tau": 0.1061,
+      "image_masks_loss": 0.0074,
+      "llm_learning_rate": 8.52975548903151e-09,
+      "loss": 0.6872,
+      "output_text_mask_loss": 0.0097,
+      "predictor_learning_rate": 8.529755489031509e-07,
+      "step": 9511
+    },
+    {
+      "epoch": 0.97,
+      "gumbel_tau": 0.1061,
+      "image_masks_loss": 0.0068,
+      "llm_learning_rate": 8.461419567170615e-09,
+      "loss": 0.6899,
+      "output_text_mask_loss": 0.034,
+      "predictor_learning_rate": 8.461419567170614e-07,
+      "step": 9512
+    },
+    {
+      "epoch": 0.97,
+      "gumbel_tau": 0.106,
+      "image_masks_loss": 0.0095,
+      "llm_learning_rate": 8.393358019088138e-09,
+      "loss": 0.7349,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 8.393358019088137e-07,
+      "step": 9513
+    },
+    {
+      "epoch": 0.97,
+      "gumbel_tau": 0.106,
+      "image_masks_loss": 0.0025,
+      "llm_learning_rate": 8.32557085227892e-09,
+      "loss": 0.7422,
+      "output_text_mask_loss": 0.0277,
+      "predictor_learning_rate": 8.325570852278918e-07,
+      "step": 9514
+    },
+    {
+      "epoch": 0.97,
+      "gumbel_tau": 0.106,
+      "image_masks_loss": 0.0113,
+      "llm_learning_rate": 8.258058074207543e-09,
+      "loss": 0.5734,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 8.258058074207542e-07,
+      "step": 9515
+    },
+    {
+      "epoch": 0.97,
+      "gumbel_tau": 0.106,
+      "image_masks_loss": 0.007,
+      "llm_learning_rate": 8.190819692309172e-09,
+      "loss": 0.7306,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 8.19081969230917e-07,
+      "step": 9516
+    },
+    {
+      "epoch": 0.98,
+      "gumbel_tau": 0.1059,
+      "image_masks_loss": 0.0028,
+      "llm_learning_rate": 8.123855713987883e-09,
+      "loss": 0.7439,
+      "output_text_mask_loss": 0.0081,
+      "predictor_learning_rate": 8.123855713987882e-07,
+      "step": 9517
+    },
+    {
+      "epoch": 0.98,
+      "gumbel_tau": 0.1059,
+      "image_masks_loss": 0.0087,
+      "llm_learning_rate": 8.057166146618057e-09,
+      "loss": 0.7185,
+      "output_text_mask_loss": 0.0641,
+      "predictor_learning_rate": 8.057166146618056e-07,
+      "step": 9518
+    },
+    {
+      "epoch": 0.98,
+      "gumbel_tau": 0.1059,
+      "image_masks_loss": 0.015,
+      "llm_learning_rate": 7.990750997543816e-09,
+      "loss": 0.7379,
+      "output_text_mask_loss": 0.0025,
+      "predictor_learning_rate": 7.990750997543816e-07,
+      "step": 9519
+    },
+    {
+      "epoch": 0.98,
+      "gumbel_tau": 0.1059,
+      "image_masks_loss": 0.013,
+      "llm_learning_rate": 7.924610274078758e-09,
+      "loss": 0.7712,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 7.924610274078758e-07,
+      "step": 9520
+    },
+    {
+      "epoch": 0.98,
+      "gumbel_tau": 0.1058,
+      "image_masks_loss": 0.016,
+      "llm_learning_rate": 7.858743983506223e-09,
+      "loss": 0.6799,
+      "output_text_mask_loss": 0.0021,
+      "predictor_learning_rate": 7.858743983506222e-07,
+      "step": 9521
+    },
+    {
+      "epoch": 0.98,
+      "gumbel_tau": 0.1058,
+      "image_masks_loss": 0.0031,
+      "llm_learning_rate": 7.793152133080128e-09,
+      "loss": 0.6732,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 7.793152133080129e-07,
+      "step": 9522
+    },
+    {
+      "epoch": 0.98,
+      "gumbel_tau": 0.1058,
+      "image_masks_loss": 0.009,
+      "llm_learning_rate": 7.72783473002331e-09,
+      "loss": 0.7643,
+      "output_text_mask_loss": 0.0039,
+      "predictor_learning_rate": 7.727834730023308e-07,
+      "step": 9523
+    },
+    {
+      "epoch": 0.98,
+      "gumbel_tau": 0.1058,
+      "image_masks_loss": 0.0123,
+      "llm_learning_rate": 7.662791781528344e-09,
+      "loss": 0.6868,
+      "output_text_mask_loss": 0.0223,
+      "predictor_learning_rate": 7.662791781528344e-07,
+      "step": 9524
+    },
+    {
+      "epoch": 0.98,
+      "gumbel_tau": 0.1057,
+      "image_masks_loss": 0.0093,
+      "llm_learning_rate": 7.598023294758671e-09,
+      "loss": 0.7788,
+      "output_text_mask_loss": 0.0025,
+      "predictor_learning_rate": 7.598023294758671e-07,
+      "step": 9525
+    },
+    {
+      "epoch": 0.98,
+      "gumbel_tau": 0.1057,
+      "image_masks_loss": 0.0136,
+      "llm_learning_rate": 7.533529276846085e-09,
+      "loss": 0.8171,
+      "output_text_mask_loss": 0.0006,
+      "predictor_learning_rate": 7.533529276846085e-07,
+      "step": 9526
+    },
+    {
+      "epoch": 0.98,
+      "gumbel_tau": 0.1057,
+      "image_masks_loss": 0.0078,
+      "llm_learning_rate": 7.46930973489296e-09,
+      "loss": 0.706,
+      "output_text_mask_loss": 0.0259,
+      "predictor_learning_rate": 7.469309734892959e-07,
+      "step": 9527
+    },
+    {
+      "epoch": 0.98,
+      "gumbel_tau": 0.1057,
+      "image_masks_loss": 0.0046,
+      "llm_learning_rate": 7.405364675971138e-09,
+      "loss": 0.7613,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 7.405364675971138e-07,
+      "step": 9528
+    },
+    {
+      "epoch": 0.98,
+      "gumbel_tau": 0.1056,
+      "image_masks_loss": 0.012,
+      "llm_learning_rate": 7.341694107123043e-09,
+      "loss": 0.6777,
+      "output_text_mask_loss": 0.0179,
+      "predictor_learning_rate": 7.341694107123043e-07,
+      "step": 9529
+    },
+    {
+      "epoch": 0.98,
+      "gumbel_tau": 0.1056,
+      "image_masks_loss": 0.016,
+      "llm_learning_rate": 7.278298035359455e-09,
+      "loss": 0.7314,
+      "output_text_mask_loss": 0.0343,
+      "predictor_learning_rate": 7.278298035359454e-07,
+      "step": 9530
+    },
+    {
+      "epoch": 0.98,
+      "gumbel_tau": 0.1056,
+      "image_masks_loss": 0.0217,
+      "llm_learning_rate": 7.2151764676620105e-09,
+      "loss": 0.7885,
+      "output_text_mask_loss": 0.0271,
+      "predictor_learning_rate": 7.215176467662011e-07,
+      "step": 9531
+    },
+    {
+      "epoch": 0.98,
+      "gumbel_tau": 0.1056,
+      "image_masks_loss": 0.0284,
+      "llm_learning_rate": 7.152329410981818e-09,
+      "loss": 0.6933,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 7.152329410981817e-07,
+      "step": 9532
+    },
+    {
+      "epoch": 0.98,
+      "gumbel_tau": 0.1055,
+      "image_masks_loss": 0.0163,
+      "llm_learning_rate": 7.089756872240006e-09,
+      "loss": 0.722,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 7.089756872240005e-07,
+      "step": 9533
+    },
+    {
+      "epoch": 0.98,
+      "gumbel_tau": 0.1055,
+      "image_masks_loss": 0.0085,
+      "llm_learning_rate": 7.027458858326619e-09,
+      "loss": 0.6964,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 7.027458858326618e-07,
+      "step": 9534
+    },
+    {
+      "epoch": 0.98,
+      "gumbel_tau": 0.1055,
+      "image_masks_loss": 0.0237,
+      "llm_learning_rate": 6.965435376102281e-09,
+      "loss": 0.7707,
+      "output_text_mask_loss": 0.0175,
+      "predictor_learning_rate": 6.96543537610228e-07,
+      "step": 9535
+    },
+    {
+      "epoch": 0.98,
+      "gumbel_tau": 0.1055,
+      "image_masks_loss": 0.0061,
+      "llm_learning_rate": 6.903686432397361e-09,
+      "loss": 0.752,
+      "output_text_mask_loss": 0.0416,
+      "predictor_learning_rate": 6.90368643239736e-07,
+      "step": 9536
+    },
+    {
+      "epoch": 0.98,
+      "gumbel_tau": 0.1054,
+      "image_masks_loss": 0.0065,
+      "llm_learning_rate": 6.842212034011419e-09,
+      "loss": 0.7303,
+      "output_text_mask_loss": 0.0185,
+      "predictor_learning_rate": 6.842212034011419e-07,
+      "step": 9537
+    },
+    {
+      "epoch": 0.98,
+      "gumbel_tau": 0.1054,
+      "image_masks_loss": 0.002,
+      "llm_learning_rate": 6.781012187714597e-09,
+      "loss": 0.7437,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 6.781012187714597e-07,
+      "step": 9538
+    },
+    {
+      "epoch": 0.98,
+      "gumbel_tau": 0.1054,
+      "image_masks_loss": 0.0068,
+      "llm_learning_rate": 6.7200869002462256e-09,
+      "loss": 0.723,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 6.720086900246225e-07,
+      "step": 9539
+    },
+    {
+      "epoch": 0.98,
+      "gumbel_tau": 0.1054,
+      "image_masks_loss": 0.0054,
+      "llm_learning_rate": 6.659436178315381e-09,
+      "loss": 0.7806,
+      "output_text_mask_loss": 0.0006,
+      "predictor_learning_rate": 6.659436178315382e-07,
+      "step": 9540
+    },
+    {
+      "epoch": 0.98,
+      "gumbel_tau": 0.1053,
+      "image_masks_loss": 0.0086,
+      "llm_learning_rate": 6.59906002860089e-09,
+      "loss": 0.7147,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 6.59906002860089e-07,
+      "step": 9541
+    },
+    {
+      "epoch": 0.98,
+      "gumbel_tau": 0.1053,
+      "image_masks_loss": 0.0093,
+      "llm_learning_rate": 6.538958457752154e-09,
+      "loss": 0.6926,
+      "output_text_mask_loss": 0.0039,
+      "predictor_learning_rate": 6.538958457752153e-07,
+      "step": 9542
+    },
+    {
+      "epoch": 0.98,
+      "gumbel_tau": 0.1053,
+      "image_masks_loss": 0.0098,
+      "llm_learning_rate": 6.4791314723869346e-09,
+      "loss": 0.691,
+      "output_text_mask_loss": 0.0387,
+      "predictor_learning_rate": 6.479131472386935e-07,
+      "step": 9543
+    },
+    {
+      "epoch": 0.98,
+      "gumbel_tau": 0.1053,
+      "image_masks_loss": 0.006,
+      "llm_learning_rate": 6.419579079093852e-09,
+      "loss": 0.7193,
+      "output_text_mask_loss": 0.0084,
+      "predictor_learning_rate": 6.419579079093851e-07,
+      "step": 9544
+    },
+    {
+      "epoch": 0.98,
+      "gumbel_tau": 0.1052,
+      "image_masks_loss": 0.0078,
+      "llm_learning_rate": 6.360301284431269e-09,
+      "loss": 0.7542,
+      "output_text_mask_loss": 0.005,
+      "predictor_learning_rate": 6.360301284431269e-07,
+      "step": 9545
+    },
+    {
+      "epoch": 0.98,
+      "gumbel_tau": 0.1052,
+      "image_masks_loss": 0.0095,
+      "llm_learning_rate": 6.3012980949264666e-09,
+      "loss": 0.7429,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 6.301298094926467e-07,
+      "step": 9546
+    },
+    {
+      "epoch": 0.98,
+      "gumbel_tau": 0.1052,
+      "image_masks_loss": 0.0116,
+      "llm_learning_rate": 6.242569517077302e-09,
+      "loss": 0.7427,
+      "output_text_mask_loss": 0.0277,
+      "predictor_learning_rate": 6.242569517077301e-07,
+      "step": 9547
+    },
+    {
+      "epoch": 0.98,
+      "gumbel_tau": 0.1052,
+      "image_masks_loss": 0.0036,
+      "llm_learning_rate": 6.184115557350823e-09,
+      "loss": 0.7102,
+      "output_text_mask_loss": 0.0004,
+      "predictor_learning_rate": 6.184115557350823e-07,
+      "step": 9548
+    },
+    {
+      "epoch": 0.98,
+      "gumbel_tau": 0.1051,
+      "image_masks_loss": 0.0068,
+      "llm_learning_rate": 6.125936222184381e-09,
+      "loss": 0.7207,
+      "output_text_mask_loss": 0.0275,
+      "predictor_learning_rate": 6.125936222184381e-07,
+      "step": 9549
+    },
+    {
+      "epoch": 0.98,
+      "gumbel_tau": 0.1051,
+      "image_masks_loss": 0.0173,
+      "llm_learning_rate": 6.068031517984796e-09,
+      "loss": 0.7004,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 6.068031517984796e-07,
+      "step": 9550
+    },
+    {
+      "epoch": 0.98,
+      "gumbel_tau": 0.1051,
+      "image_masks_loss": 0.007,
+      "llm_learning_rate": 6.010401451128634e-09,
+      "loss": 0.7839,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 6.010401451128633e-07,
+      "step": 9551
+    },
+    {
+      "epoch": 0.98,
+      "gumbel_tau": 0.1051,
+      "image_masks_loss": 0.0111,
+      "llm_learning_rate": 5.9530460279619284e-09,
+      "loss": 0.7132,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 5.953046027961928e-07,
+      "step": 9552
+    },
+    {
+      "epoch": 0.98,
+      "gumbel_tau": 0.105,
+      "image_masks_loss": 0.0103,
+      "llm_learning_rate": 5.895965254801295e-09,
+      "loss": 0.6932,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 5.895965254801294e-07,
+      "step": 9553
+    },
+    {
+      "epoch": 0.98,
+      "gumbel_tau": 0.105,
+      "image_masks_loss": 0.0089,
+      "llm_learning_rate": 5.839159137932537e-09,
+      "loss": 0.7177,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 5.839159137932537e-07,
+      "step": 9554
+    },
+    {
+      "epoch": 0.98,
+      "gumbel_tau": 0.105,
+      "image_masks_loss": 0.0136,
+      "llm_learning_rate": 5.782627683610931e-09,
+      "loss": 0.6676,
+      "output_text_mask_loss": 0.0019,
+      "predictor_learning_rate": 5.782627683610931e-07,
+      "step": 9555
+    },
+    {
+      "epoch": 0.98,
+      "gumbel_tau": 0.105,
+      "image_masks_loss": 0.0233,
+      "llm_learning_rate": 5.726370898061773e-09,
+      "loss": 0.746,
+      "output_text_mask_loss": 0.004,
+      "predictor_learning_rate": 5.726370898061773e-07,
+      "step": 9556
+    },
+    {
+      "epoch": 0.98,
+      "gumbel_tau": 0.1049,
+      "image_masks_loss": 0.0147,
+      "llm_learning_rate": 5.670388787480663e-09,
+      "loss": 0.7514,
+      "output_text_mask_loss": 0.0078,
+      "predictor_learning_rate": 5.670388787480663e-07,
+      "step": 9557
+    },
+    {
+      "epoch": 0.98,
+      "gumbel_tau": 0.1049,
+      "image_masks_loss": 0.0225,
+      "llm_learning_rate": 5.614681358032114e-09,
+      "loss": 0.7489,
+      "output_text_mask_loss": 0.0099,
+      "predictor_learning_rate": 5.614681358032115e-07,
+      "step": 9558
+    },
+    {
+      "epoch": 0.98,
+      "gumbel_tau": 0.1049,
+      "image_masks_loss": 0.0051,
+      "llm_learning_rate": 5.55924861585122e-09,
+      "loss": 0.7195,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 5.55924861585122e-07,
+      "step": 9559
+    },
+    {
+      "epoch": 0.98,
+      "gumbel_tau": 0.1049,
+      "image_masks_loss": 0.0057,
+      "llm_learning_rate": 5.504090567041709e-09,
+      "loss": 0.6691,
+      "output_text_mask_loss": 0.0172,
+      "predictor_learning_rate": 5.504090567041709e-07,
+      "step": 9560
+    },
+    {
+      "epoch": 0.98,
+      "gumbel_tau": 0.1048,
+      "image_masks_loss": 0.0059,
+      "llm_learning_rate": 5.449207217678165e-09,
+      "loss": 0.8099,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 5.449207217678165e-07,
+      "step": 9561
+    },
+    {
+      "epoch": 0.98,
+      "gumbel_tau": 0.1048,
+      "image_masks_loss": 0.0088,
+      "llm_learning_rate": 5.394598573804088e-09,
+      "loss": 0.6878,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 5.394598573804088e-07,
+      "step": 9562
+    },
+    {
+      "epoch": 0.98,
+      "gumbel_tau": 0.1048,
+      "image_masks_loss": 0.0134,
+      "llm_learning_rate": 5.3402646414335564e-09,
+      "loss": 0.7082,
+      "output_text_mask_loss": 0.0019,
+      "predictor_learning_rate": 5.340264641433556e-07,
+      "step": 9563
+    },
+    {
+      "epoch": 0.98,
+      "gumbel_tau": 0.1048,
+      "image_masks_loss": 0.011,
+      "llm_learning_rate": 5.286205426549562e-09,
+      "loss": 0.7942,
+      "output_text_mask_loss": 0.0063,
+      "predictor_learning_rate": 5.286205426549561e-07,
+      "step": 9564
+    },
+    {
+      "epoch": 0.98,
+      "gumbel_tau": 0.1047,
+      "image_masks_loss": 0.0105,
+      "llm_learning_rate": 5.232420935105675e-09,
+      "loss": 0.671,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 5.232420935105675e-07,
+      "step": 9565
+    },
+    {
+      "epoch": 0.98,
+      "gumbel_tau": 0.1047,
+      "image_masks_loss": 0.0066,
+      "llm_learning_rate": 5.1789111730241036e-09,
+      "loss": 0.7919,
+      "output_text_mask_loss": 0.0046,
+      "predictor_learning_rate": 5.178911173024104e-07,
+      "step": 9566
+    },
+    {
+      "epoch": 0.98,
+      "gumbel_tau": 0.1047,
+      "image_masks_loss": 0.0085,
+      "llm_learning_rate": 5.125676146198189e-09,
+      "loss": 0.7283,
+      "output_text_mask_loss": 0.0221,
+      "predictor_learning_rate": 5.125676146198188e-07,
+      "step": 9567
+    },
+    {
+      "epoch": 0.98,
+      "gumbel_tau": 0.1047,
+      "image_masks_loss": 0.0056,
+      "llm_learning_rate": 5.07271586048963e-09,
+      "loss": 0.6855,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 5.07271586048963e-07,
+      "step": 9568
+    },
+    {
+      "epoch": 0.98,
+      "gumbel_tau": 0.1046,
+      "image_masks_loss": 0.0151,
+      "llm_learning_rate": 5.0200303217309845e-09,
+      "loss": 0.8597,
+      "output_text_mask_loss": 0.0153,
+      "predictor_learning_rate": 5.020030321730984e-07,
+      "step": 9569
+    },
+    {
+      "epoch": 0.98,
+      "gumbel_tau": 0.1046,
+      "image_masks_loss": 0.0087,
+      "llm_learning_rate": 4.967619535723723e-09,
+      "loss": 0.6853,
+      "output_text_mask_loss": 0.0336,
+      "predictor_learning_rate": 4.967619535723722e-07,
+      "step": 9570
+    },
+    {
+      "epoch": 0.98,
+      "gumbel_tau": 0.1046,
+      "image_masks_loss": 0.0078,
+      "llm_learning_rate": 4.9154835082398934e-09,
+      "loss": 0.7219,
+      "output_text_mask_loss": 0.0055,
+      "predictor_learning_rate": 4.915483508239893e-07,
+      "step": 9571
+    },
+    {
+      "epoch": 0.98,
+      "gumbel_tau": 0.1046,
+      "image_masks_loss": 0.0129,
+      "llm_learning_rate": 4.8636222450207384e-09,
+      "loss": 0.7875,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 4.863622245020738e-07,
+      "step": 9572
+    },
+    {
+      "epoch": 0.98,
+      "gumbel_tau": 0.1045,
+      "image_masks_loss": 0.0133,
+      "llm_learning_rate": 4.812035751776967e-09,
+      "loss": 0.6822,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 4.812035751776966e-07,
+      "step": 9573
+    },
+    {
+      "epoch": 0.98,
+      "gumbel_tau": 0.1045,
+      "image_masks_loss": 0.0099,
+      "llm_learning_rate": 4.76072403418959e-09,
+      "loss": 0.6069,
+      "output_text_mask_loss": 0.0027,
+      "predictor_learning_rate": 4.7607240341895894e-07,
+      "step": 9574
+    },
+    {
+      "epoch": 0.98,
+      "gumbel_tau": 0.1045,
+      "image_masks_loss": 0.0049,
+      "llm_learning_rate": 4.7096870979093655e-09,
+      "loss": 0.6701,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 4.709687097909365e-07,
+      "step": 9575
+    },
+    {
+      "epoch": 0.98,
+      "gumbel_tau": 0.1045,
+      "image_masks_loss": 0.015,
+      "llm_learning_rate": 4.658924948556243e-09,
+      "loss": 0.6731,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 4.6589249485562424e-07,
+      "step": 9576
+    },
+    {
+      "epoch": 0.98,
+      "gumbel_tau": 0.1044,
+      "image_masks_loss": 0.0071,
+      "llm_learning_rate": 4.608437591720749e-09,
+      "loss": 0.7673,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 4.6084375917207487e-07,
+      "step": 9577
+    },
+    {
+      "epoch": 0.98,
+      "gumbel_tau": 0.1044,
+      "image_masks_loss": 0.0153,
+      "llm_learning_rate": 4.558225032962049e-09,
+      "loss": 0.7163,
+      "output_text_mask_loss": 0.017,
+      "predictor_learning_rate": 4.5582250329620487e-07,
+      "step": 9578
+    },
+    {
+      "epoch": 0.98,
+      "gumbel_tau": 0.1044,
+      "image_masks_loss": 0.0087,
+      "llm_learning_rate": 4.508287277810441e-09,
+      "loss": 0.7608,
+      "output_text_mask_loss": 0.0437,
+      "predictor_learning_rate": 4.508287277810441e-07,
+      "step": 9579
+    },
+    {
+      "epoch": 0.98,
+      "gumbel_tau": 0.1044,
+      "image_masks_loss": 0.0045,
+      "llm_learning_rate": 4.458624331764305e-09,
+      "loss": 0.7649,
+      "output_text_mask_loss": 0.0101,
+      "predictor_learning_rate": 4.458624331764305e-07,
+      "step": 9580
+    },
+    {
+      "epoch": 0.98,
+      "gumbel_tau": 0.1043,
+      "image_masks_loss": 0.0026,
+      "llm_learning_rate": 4.4092362002931545e-09,
+      "loss": 0.7371,
+      "output_text_mask_loss": 0.0044,
+      "predictor_learning_rate": 4.409236200293154e-07,
+      "step": 9581
+    },
+    {
+      "epoch": 0.98,
+      "gumbel_tau": 0.1043,
+      "image_masks_loss": 0.0071,
+      "llm_learning_rate": 4.360122888835694e-09,
+      "loss": 0.7227,
+      "output_text_mask_loss": 0.0107,
+      "predictor_learning_rate": 4.360122888835694e-07,
+      "step": 9582
+    },
+    {
+      "epoch": 0.98,
+      "gumbel_tau": 0.1043,
+      "image_masks_loss": 0.0123,
+      "llm_learning_rate": 4.3112844028003755e-09,
+      "loss": 0.6831,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 4.311284402800375e-07,
+      "step": 9583
+    },
+    {
+      "epoch": 0.98,
+      "gumbel_tau": 0.1043,
+      "image_masks_loss": 0.0048,
+      "llm_learning_rate": 4.262720747565119e-09,
+      "loss": 0.7082,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 4.262720747565119e-07,
+      "step": 9584
+    },
+    {
+      "epoch": 0.98,
+      "gumbel_tau": 0.1042,
+      "image_masks_loss": 0.0163,
+      "llm_learning_rate": 4.214431928478424e-09,
+      "loss": 0.7911,
+      "output_text_mask_loss": 0.0039,
+      "predictor_learning_rate": 4.2144319284784237e-07,
+      "step": 9585
+    },
+    {
+      "epoch": 0.98,
+      "gumbel_tau": 0.1042,
+      "image_masks_loss": 0.0052,
+      "llm_learning_rate": 4.166417950857426e-09,
+      "loss": 0.6688,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 4.166417950857426e-07,
+      "step": 9586
+    },
+    {
+      "epoch": 0.98,
+      "gumbel_tau": 0.1042,
+      "image_masks_loss": 0.009,
+      "llm_learning_rate": 4.118678819989841e-09,
+      "loss": 0.7317,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 4.118678819989841e-07,
+      "step": 9587
+    },
+    {
+      "epoch": 0.98,
+      "gumbel_tau": 0.1042,
+      "image_masks_loss": 0.0148,
+      "llm_learning_rate": 4.071214541132573e-09,
+      "loss": 0.6789,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 4.071214541132573e-07,
+      "step": 9588
+    },
+    {
+      "epoch": 0.98,
+      "gumbel_tau": 0.1041,
+      "image_masks_loss": 0.0134,
+      "llm_learning_rate": 4.024025119512554e-09,
+      "loss": 0.7773,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 4.0240251195125533e-07,
+      "step": 9589
+    },
+    {
+      "epoch": 0.98,
+      "gumbel_tau": 0.1041,
+      "image_masks_loss": 0.0112,
+      "llm_learning_rate": 3.9771105603264586e-09,
+      "loss": 0.7382,
+      "output_text_mask_loss": 0.0261,
+      "predictor_learning_rate": 3.977110560326458e-07,
+      "step": 9590
+    },
+    {
+      "epoch": 0.98,
+      "gumbel_tau": 0.1041,
+      "image_masks_loss": 0.0162,
+      "llm_learning_rate": 3.930470868740988e-09,
+      "loss": 0.6728,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 3.9304708687409873e-07,
+      "step": 9591
+    },
+    {
+      "epoch": 0.98,
+      "gumbel_tau": 0.1041,
+      "image_masks_loss": 0.0122,
+      "llm_learning_rate": 3.884106049891478e-09,
+      "loss": 0.7442,
+      "output_text_mask_loss": 0.0056,
+      "predictor_learning_rate": 3.884106049891478e-07,
+      "step": 9592
+    },
+    {
+      "epoch": 0.98,
+      "gumbel_tau": 0.104,
+      "image_masks_loss": 0.016,
+      "llm_learning_rate": 3.838016108884124e-09,
+      "loss": 0.6716,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 3.8380161088841236e-07,
+      "step": 9593
+    },
+    {
+      "epoch": 0.98,
+      "gumbel_tau": 0.104,
+      "image_masks_loss": 0.0033,
+      "llm_learning_rate": 3.792201050794586e-09,
+      "loss": 0.7743,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 3.792201050794586e-07,
+      "step": 9594
+    },
+    {
+      "epoch": 0.98,
+      "gumbel_tau": 0.104,
+      "image_masks_loss": 0.0039,
+      "llm_learning_rate": 3.7466608806677186e-09,
+      "loss": 0.7623,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 3.746660880667718e-07,
+      "step": 9595
+    },
+    {
+      "epoch": 0.98,
+      "gumbel_tau": 0.104,
+      "image_masks_loss": 0.0025,
+      "llm_learning_rate": 3.7013956035189536e-09,
+      "loss": 0.6973,
+      "output_text_mask_loss": 0.0368,
+      "predictor_learning_rate": 3.7013956035189533e-07,
+      "step": 9596
+    },
+    {
+      "epoch": 0.98,
+      "gumbel_tau": 0.1039,
+      "image_masks_loss": 0.0121,
+      "llm_learning_rate": 3.6564052243326376e-09,
+      "loss": 0.7614,
+      "output_text_mask_loss": 0.0111,
+      "predictor_learning_rate": 3.6564052243326373e-07,
+      "step": 9597
+    },
+    {
+      "epoch": 0.98,
+      "gumbel_tau": 0.1039,
+      "image_masks_loss": 0.014,
+      "llm_learning_rate": 3.611689748063696e-09,
+      "loss": 0.7671,
+      "output_text_mask_loss": 0.0292,
+      "predictor_learning_rate": 3.611689748063696e-07,
+      "step": 9598
+    },
+    {
+      "epoch": 0.98,
+      "gumbel_tau": 0.1039,
+      "image_masks_loss": 0.0104,
+      "llm_learning_rate": 3.56724917963569e-09,
+      "loss": 0.7531,
+      "output_text_mask_loss": 0.0534,
+      "predictor_learning_rate": 3.56724917963569e-07,
+      "step": 9599
+    },
+    {
+      "epoch": 0.98,
+      "gumbel_tau": 0.1039,
+      "image_masks_loss": 0.0141,
+      "llm_learning_rate": 3.5230835239427607e-09,
+      "loss": 0.7188,
+      "output_text_mask_loss": 0.0041,
+      "predictor_learning_rate": 3.5230835239427606e-07,
+      "step": 9600
+    },
+    {
+      "epoch": 0.98,
+      "gumbel_tau": 0.1038,
+      "image_masks_loss": 0.0065,
+      "llm_learning_rate": 3.4791927858485175e-09,
+      "loss": 0.7208,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 3.479192785848517e-07,
+      "step": 9601
+    },
+    {
+      "epoch": 0.98,
+      "gumbel_tau": 0.1038,
+      "image_masks_loss": 0.0088,
+      "llm_learning_rate": 3.435576970186594e-09,
+      "loss": 0.6954,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 3.4355769701865934e-07,
+      "step": 9602
+    },
+    {
+      "epoch": 0.98,
+      "gumbel_tau": 0.1038,
+      "image_masks_loss": 0.0053,
+      "llm_learning_rate": 3.3922360817595366e-09,
+      "loss": 0.8373,
+      "output_text_mask_loss": 0.0198,
+      "predictor_learning_rate": 3.392236081759537e-07,
+      "step": 9603
+    },
+    {
+      "epoch": 0.98,
+      "gumbel_tau": 0.1038,
+      "image_masks_loss": 0.0077,
+      "llm_learning_rate": 3.3491701253407504e-09,
+      "loss": 0.833,
+      "output_text_mask_loss": 0.0101,
+      "predictor_learning_rate": 3.34917012534075e-07,
+      "step": 9604
+    },
+    {
+      "epoch": 0.98,
+      "gumbel_tau": 0.1037,
+      "image_masks_loss": 0.0131,
+      "llm_learning_rate": 3.3063791056719973e-09,
+      "loss": 0.7593,
+      "output_text_mask_loss": 0.0173,
+      "predictor_learning_rate": 3.306379105671997e-07,
+      "step": 9605
+    },
+    {
+      "epoch": 0.98,
+      "gumbel_tau": 0.1037,
+      "image_masks_loss": 0.0099,
+      "llm_learning_rate": 3.263863027466174e-09,
+      "loss": 0.6951,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 3.263863027466174e-07,
+      "step": 9606
+    },
+    {
+      "epoch": 0.98,
+      "gumbel_tau": 0.1037,
+      "image_masks_loss": 0.0109,
+      "llm_learning_rate": 3.2216218954050916e-09,
+      "loss": 0.7534,
+      "output_text_mask_loss": 0.0248,
+      "predictor_learning_rate": 3.221621895405091e-07,
+      "step": 9607
+    },
+    {
+      "epoch": 0.98,
+      "gumbel_tau": 0.1037,
+      "image_masks_loss": 0.0064,
+      "llm_learning_rate": 3.179655714140306e-09,
+      "loss": 0.6679,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 3.179655714140306e-07,
+      "step": 9608
+    },
+    {
+      "epoch": 0.98,
+      "gumbel_tau": 0.1037,
+      "image_masks_loss": 0.0113,
+      "llm_learning_rate": 3.1379644882931216e-09,
+      "loss": 0.7262,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 3.137964488293121e-07,
+      "step": 9609
+    },
+    {
+      "epoch": 0.98,
+      "gumbel_tau": 0.1036,
+      "image_masks_loss": 0.0036,
+      "llm_learning_rate": 3.0965482224551423e-09,
+      "loss": 0.6993,
+      "output_text_mask_loss": 0.071,
+      "predictor_learning_rate": 3.0965482224551424e-07,
+      "step": 9610
+    },
+    {
+      "epoch": 0.98,
+      "gumbel_tau": 0.1036,
+      "image_masks_loss": 0.0026,
+      "llm_learning_rate": 3.05540692118661e-09,
+      "loss": 0.6734,
+      "output_text_mask_loss": 0.0004,
+      "predictor_learning_rate": 3.05540692118661e-07,
+      "step": 9611
+    },
+    {
+      "epoch": 0.98,
+      "gumbel_tau": 0.1036,
+      "image_masks_loss": 0.0104,
+      "llm_learning_rate": 3.014540589018622e-09,
+      "loss": 0.7218,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 3.014540589018622e-07,
+      "step": 9612
+    },
+    {
+      "epoch": 0.98,
+      "gumbel_tau": 0.1036,
+      "image_masks_loss": 0.0208,
+      "llm_learning_rate": 2.9739492304511895e-09,
+      "loss": 0.7609,
+      "output_text_mask_loss": 0.001,
+      "predictor_learning_rate": 2.9739492304511897e-07,
+      "step": 9613
+    },
+    {
+      "epoch": 0.99,
+      "gumbel_tau": 0.1035,
+      "image_masks_loss": 0.0074,
+      "llm_learning_rate": 2.9336328499543487e-09,
+      "loss": 0.7746,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 2.9336328499543484e-07,
+      "step": 9614
+    },
+    {
+      "epoch": 0.99,
+      "gumbel_tau": 0.1035,
+      "image_masks_loss": 0.0077,
+      "llm_learning_rate": 2.8935914519676035e-09,
+      "loss": 0.726,
+      "output_text_mask_loss": 0.0025,
+      "predictor_learning_rate": 2.893591451967603e-07,
+      "step": 9615
+    },
+    {
+      "epoch": 0.99,
+      "gumbel_tau": 0.1035,
+      "image_masks_loss": 0.0067,
+      "llm_learning_rate": 2.853825040901037e-09,
+      "loss": 0.7299,
+      "output_text_mask_loss": 0.0263,
+      "predictor_learning_rate": 2.8538250409010367e-07,
+      "step": 9616
+    },
+    {
+      "epoch": 0.99,
+      "gumbel_tau": 0.1035,
+      "image_masks_loss": 0.0128,
+      "llm_learning_rate": 2.814333621133092e-09,
+      "loss": 0.7588,
+      "output_text_mask_loss": 0.0088,
+      "predictor_learning_rate": 2.8143336211330915e-07,
+      "step": 9617
+    },
+    {
+      "epoch": 0.99,
+      "gumbel_tau": 0.1034,
+      "image_masks_loss": 0.0181,
+      "llm_learning_rate": 2.775117197013344e-09,
+      "loss": 0.6073,
+      "output_text_mask_loss": 0.0235,
+      "predictor_learning_rate": 2.775117197013344e-07,
+      "step": 9618
+    },
+    {
+      "epoch": 0.99,
+      "gumbel_tau": 0.1034,
+      "image_masks_loss": 0.0122,
+      "llm_learning_rate": 2.7361757728597283e-09,
+      "loss": 0.6724,
+      "output_text_mask_loss": 0.0162,
+      "predictor_learning_rate": 2.736175772859728e-07,
+      "step": 9619
+    },
+    {
+      "epoch": 0.99,
+      "gumbel_tau": 0.1034,
+      "image_masks_loss": 0.0189,
+      "llm_learning_rate": 2.6975093529610364e-09,
+      "loss": 0.7688,
+      "output_text_mask_loss": 0.0208,
+      "predictor_learning_rate": 2.697509352961036e-07,
+      "step": 9620
+    },
+    {
+      "epoch": 0.99,
+      "gumbel_tau": 0.1034,
+      "image_masks_loss": 0.0069,
+      "llm_learning_rate": 2.659117941575251e-09,
+      "loss": 0.7298,
+      "output_text_mask_loss": 0.0553,
+      "predictor_learning_rate": 2.6591179415752507e-07,
+      "step": 9621
+    },
+    {
+      "epoch": 0.99,
+      "gumbel_tau": 0.1033,
+      "image_masks_loss": 0.0046,
+      "llm_learning_rate": 2.6210015429298243e-09,
+      "loss": 0.7126,
+      "output_text_mask_loss": 0.0443,
+      "predictor_learning_rate": 2.6210015429298237e-07,
+      "step": 9622
+    },
+    {
+      "epoch": 0.99,
+      "gumbel_tau": 0.1033,
+      "image_masks_loss": 0.0027,
+      "llm_learning_rate": 2.5831601612225087e-09,
+      "loss": 0.6828,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 2.583160161222509e-07,
+      "step": 9623
+    },
+    {
+      "epoch": 0.99,
+      "gumbel_tau": 0.1033,
+      "image_masks_loss": 0.0151,
+      "llm_learning_rate": 2.5455938006202495e-09,
+      "loss": 0.7071,
+      "output_text_mask_loss": 0.0034,
+      "predictor_learning_rate": 2.5455938006202494e-07,
+      "step": 9624
+    },
+    {
+      "epoch": 0.99,
+      "gumbel_tau": 0.1033,
+      "image_masks_loss": 0.0036,
+      "llm_learning_rate": 2.508302465260293e-09,
+      "loss": 0.7213,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 2.5083024652602924e-07,
+      "step": 9625
+    },
+    {
+      "epoch": 0.99,
+      "gumbel_tau": 0.1032,
+      "image_masks_loss": 0.006,
+      "llm_learning_rate": 2.471286159248798e-09,
+      "loss": 0.7468,
+      "output_text_mask_loss": 0.0147,
+      "predictor_learning_rate": 2.471286159248798e-07,
+      "step": 9626
+    },
+    {
+      "epoch": 0.99,
+      "gumbel_tau": 0.1032,
+      "image_masks_loss": 0.008,
+      "llm_learning_rate": 2.4345448866622268e-09,
+      "loss": 0.6657,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 2.4345448866622265e-07,
+      "step": 9627
+    },
+    {
+      "epoch": 0.99,
+      "gumbel_tau": 0.1032,
+      "image_masks_loss": 0.0038,
+      "llm_learning_rate": 2.398078651546787e-09,
+      "loss": 0.7376,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 2.3980786515467866e-07,
+      "step": 9628
+    },
+    {
+      "epoch": 0.99,
+      "gumbel_tau": 0.1032,
+      "image_masks_loss": 0.0048,
+      "llm_learning_rate": 2.361887457918155e-09,
+      "loss": 0.7201,
+      "output_text_mask_loss": 0.0256,
+      "predictor_learning_rate": 2.3618874579181548e-07,
+      "step": 9629
+    },
+    {
+      "epoch": 0.99,
+      "gumbel_tau": 0.1031,
+      "image_masks_loss": 0.011,
+      "llm_learning_rate": 2.325971309761754e-09,
+      "loss": 0.7561,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 2.325971309761754e-07,
+      "step": 9630
+    },
+    {
+      "epoch": 0.99,
+      "gumbel_tau": 0.1031,
+      "image_masks_loss": 0.0078,
+      "llm_learning_rate": 2.2903302110327542e-09,
+      "loss": 0.698,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 2.290330211032754e-07,
+      "step": 9631
+    },
+    {
+      "epoch": 0.99,
+      "gumbel_tau": 0.1031,
+      "image_masks_loss": 0.0051,
+      "llm_learning_rate": 2.2549641656560705e-09,
+      "loss": 0.6771,
+      "output_text_mask_loss": 0.0265,
+      "predictor_learning_rate": 2.2549641656560704e-07,
+      "step": 9632
+    },
+    {
+      "epoch": 0.99,
+      "gumbel_tau": 0.1031,
+      "image_masks_loss": 0.0111,
+      "llm_learning_rate": 2.2198731775260883e-09,
+      "loss": 0.7454,
+      "output_text_mask_loss": 0.0237,
+      "predictor_learning_rate": 2.2198731775260884e-07,
+      "step": 9633
+    },
+    {
+      "epoch": 0.99,
+      "gumbel_tau": 0.103,
+      "image_masks_loss": 0.0194,
+      "llm_learning_rate": 2.1850572505072167e-09,
+      "loss": 0.7646,
+      "output_text_mask_loss": 0.0275,
+      "predictor_learning_rate": 2.1850572505072165e-07,
+      "step": 9634
+    },
+    {
+      "epoch": 0.99,
+      "gumbel_tau": 0.103,
+      "image_masks_loss": 0.0018,
+      "llm_learning_rate": 2.15051638843361e-09,
+      "loss": 0.7056,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 2.1505163884336098e-07,
+      "step": 9635
+    },
+    {
+      "epoch": 0.99,
+      "gumbel_tau": 0.103,
+      "image_masks_loss": 0.0077,
+      "llm_learning_rate": 2.116250595108893e-09,
+      "loss": 0.6895,
+      "output_text_mask_loss": 0.0107,
+      "predictor_learning_rate": 2.116250595108893e-07,
+      "step": 9636
+    },
+    {
+      "epoch": 0.99,
+      "gumbel_tau": 0.103,
+      "image_masks_loss": 0.0098,
+      "llm_learning_rate": 2.082259874306436e-09,
+      "loss": 0.7141,
+      "output_text_mask_loss": 0.0025,
+      "predictor_learning_rate": 2.0822598743064357e-07,
+      "step": 9637
+    },
+    {
+      "epoch": 0.99,
+      "gumbel_tau": 0.1029,
+      "image_masks_loss": 0.0108,
+      "llm_learning_rate": 2.0485442297696336e-09,
+      "loss": 0.7013,
+      "output_text_mask_loss": 0.0668,
+      "predictor_learning_rate": 2.0485442297696334e-07,
+      "step": 9638
+    },
+    {
+      "epoch": 0.99,
+      "gumbel_tau": 0.1029,
+      "image_masks_loss": 0.01,
+      "llm_learning_rate": 2.0151036652107936e-09,
+      "loss": 0.6746,
+      "output_text_mask_loss": 0.0515,
+      "predictor_learning_rate": 2.0151036652107936e-07,
+      "step": 9639
+    },
+    {
+      "epoch": 0.99,
+      "gumbel_tau": 0.1029,
+      "image_masks_loss": 0.0113,
+      "llm_learning_rate": 1.9819381843130815e-09,
+      "loss": 0.7021,
+      "output_text_mask_loss": 0.0238,
+      "predictor_learning_rate": 1.9819381843130813e-07,
+      "step": 9640
+    },
+    {
+      "epoch": 0.99,
+      "gumbel_tau": 0.1029,
+      "image_masks_loss": 0.0076,
+      "llm_learning_rate": 1.9490477907282977e-09,
+      "loss": 0.6997,
+      "output_text_mask_loss": 0.034,
+      "predictor_learning_rate": 1.9490477907282978e-07,
+      "step": 9641
+    },
+    {
+      "epoch": 0.99,
+      "gumbel_tau": 0.1028,
+      "image_masks_loss": 0.0046,
+      "llm_learning_rate": 1.9164324880785458e-09,
+      "loss": 0.686,
+      "output_text_mask_loss": 0.0164,
+      "predictor_learning_rate": 1.9164324880785454e-07,
+      "step": 9642
+    },
+    {
+      "epoch": 0.99,
+      "gumbel_tau": 0.1028,
+      "image_masks_loss": 0.002,
+      "llm_learning_rate": 1.8840922799553964e-09,
+      "loss": 0.6648,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 1.8840922799553961e-07,
+      "step": 9643
+    },
+    {
+      "epoch": 0.99,
+      "gumbel_tau": 0.1028,
+      "image_masks_loss": 0.0086,
+      "llm_learning_rate": 1.8520271699204452e-09,
+      "loss": 0.7362,
+      "output_text_mask_loss": 0.0238,
+      "predictor_learning_rate": 1.852027169920445e-07,
+      "step": 9644
+    },
+    {
+      "epoch": 0.99,
+      "gumbel_tau": 0.1028,
+      "image_masks_loss": 0.0145,
+      "llm_learning_rate": 1.8202371615047566e-09,
+      "loss": 0.7153,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 1.8202371615047563e-07,
+      "step": 9645
+    },
+    {
+      "epoch": 0.99,
+      "gumbel_tau": 0.1028,
+      "image_masks_loss": 0.0138,
+      "llm_learning_rate": 1.7887222582088638e-09,
+      "loss": 0.6743,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 1.7887222582088635e-07,
+      "step": 9646
+    },
+    {
+      "epoch": 0.99,
+      "gumbel_tau": 0.1027,
+      "image_masks_loss": 0.0191,
+      "llm_learning_rate": 1.7574824635036015e-09,
+      "loss": 0.7306,
+      "output_text_mask_loss": 0.0839,
+      "predictor_learning_rate": 1.7574824635036013e-07,
+      "step": 9647
+    },
+    {
+      "epoch": 0.99,
+      "gumbel_tau": 0.1027,
+      "image_masks_loss": 0.0149,
+      "llm_learning_rate": 1.7265177808287181e-09,
+      "loss": 0.7034,
+      "output_text_mask_loss": 0.0023,
+      "predictor_learning_rate": 1.726517780828718e-07,
+      "step": 9648
+    },
+    {
+      "epoch": 0.99,
+      "gumbel_tau": 0.1027,
+      "image_masks_loss": 0.0071,
+      "llm_learning_rate": 1.695828213594819e-09,
+      "loss": 0.7793,
+      "output_text_mask_loss": 0.001,
+      "predictor_learning_rate": 1.695828213594819e-07,
+      "step": 9649
+    },
+    {
+      "epoch": 0.99,
+      "gumbel_tau": 0.1027,
+      "image_masks_loss": 0.0132,
+      "llm_learning_rate": 1.6654137651805901e-09,
+      "loss": 0.7122,
+      "output_text_mask_loss": 0.0175,
+      "predictor_learning_rate": 1.66541376518059e-07,
+      "step": 9650
+    },
+    {
+      "epoch": 0.99,
+      "gumbel_tau": 0.1026,
+      "image_masks_loss": 0.0073,
+      "llm_learning_rate": 1.6352744389361297e-09,
+      "loss": 0.6965,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 1.6352744389361295e-07,
+      "step": 9651
+    },
+    {
+      "epoch": 0.99,
+      "gumbel_tau": 0.1026,
+      "image_masks_loss": 0.0074,
+      "llm_learning_rate": 1.6054102381801718e-09,
+      "loss": 0.6937,
+      "output_text_mask_loss": 0.0208,
+      "predictor_learning_rate": 1.6054102381801716e-07,
+      "step": 9652
+    },
+    {
+      "epoch": 0.99,
+      "gumbel_tau": 0.1026,
+      "image_masks_loss": 0.0282,
+      "llm_learning_rate": 1.575821166201197e-09,
+      "loss": 0.7278,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 1.575821166201197e-07,
+      "step": 9653
+    },
+    {
+      "epoch": 0.99,
+      "gumbel_tau": 0.1026,
+      "image_masks_loss": 0.0096,
+      "llm_learning_rate": 1.5465072262579872e-09,
+      "loss": 0.6164,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 1.546507226257987e-07,
+      "step": 9654
+    },
+    {
+      "epoch": 0.99,
+      "gumbel_tau": 0.1025,
+      "image_masks_loss": 0.006,
+      "llm_learning_rate": 1.5174684215782386e-09,
+      "loss": 0.7107,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 1.5174684215782386e-07,
+      "step": 9655
+    },
+    {
+      "epoch": 0.99,
+      "gumbel_tau": 0.1025,
+      "image_masks_loss": 0.0084,
+      "llm_learning_rate": 1.4887047553602262e-09,
+      "loss": 0.7082,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 1.4887047553602263e-07,
+      "step": 9656
+    },
+    {
+      "epoch": 0.99,
+      "gumbel_tau": 0.1025,
+      "image_masks_loss": 0.0121,
+      "llm_learning_rate": 1.4602162307711388e-09,
+      "loss": 0.7083,
+      "output_text_mask_loss": 0.0193,
+      "predictor_learning_rate": 1.4602162307711387e-07,
+      "step": 9657
+    },
+    {
+      "epoch": 0.99,
+      "gumbel_tau": 0.1025,
+      "image_masks_loss": 0.0055,
+      "llm_learning_rate": 1.4320028509484663e-09,
+      "loss": 0.7718,
+      "output_text_mask_loss": 0.0269,
+      "predictor_learning_rate": 1.432002850948466e-07,
+      "step": 9658
+    },
+    {
+      "epoch": 0.99,
+      "gumbel_tau": 0.1024,
+      "image_masks_loss": 0.0072,
+      "llm_learning_rate": 1.4040646189988905e-09,
+      "loss": 0.7497,
+      "output_text_mask_loss": 0.028,
+      "predictor_learning_rate": 1.4040646189988905e-07,
+      "step": 9659
+    },
+    {
+      "epoch": 0.99,
+      "gumbel_tau": 0.1024,
+      "image_masks_loss": 0.0031,
+      "llm_learning_rate": 1.3764015379991168e-09,
+      "loss": 0.7468,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 1.376401537999117e-07,
+      "step": 9660
+    },
+    {
+      "epoch": 0.99,
+      "gumbel_tau": 0.1024,
+      "image_masks_loss": 0.0048,
+      "llm_learning_rate": 1.3490136109955975e-09,
+      "loss": 0.7537,
+      "output_text_mask_loss": 0.0034,
+      "predictor_learning_rate": 1.3490136109955976e-07,
+      "step": 9661
+    },
+    {
+      "epoch": 0.99,
+      "gumbel_tau": 0.1024,
+      "image_masks_loss": 0.0172,
+      "llm_learning_rate": 1.3219008410042533e-09,
+      "loss": 0.6985,
+      "output_text_mask_loss": 0.0542,
+      "predictor_learning_rate": 1.3219008410042532e-07,
+      "step": 9662
+    },
+    {
+      "epoch": 0.99,
+      "gumbel_tau": 0.1023,
+      "image_masks_loss": 0.0088,
+      "llm_learning_rate": 1.295063231010474e-09,
+      "loss": 0.7202,
+      "output_text_mask_loss": 0.0124,
+      "predictor_learning_rate": 1.2950632310104738e-07,
+      "step": 9663
+    },
+    {
+      "epoch": 0.99,
+      "gumbel_tau": 0.1023,
+      "image_masks_loss": 0.0106,
+      "llm_learning_rate": 1.268500783970228e-09,
+      "loss": 0.7463,
+      "output_text_mask_loss": 0.0599,
+      "predictor_learning_rate": 1.2685007839702278e-07,
+      "step": 9664
+    },
+    {
+      "epoch": 0.99,
+      "gumbel_tau": 0.1023,
+      "image_masks_loss": 0.008,
+      "llm_learning_rate": 1.2422135028083982e-09,
+      "loss": 0.7448,
+      "output_text_mask_loss": 0.0093,
+      "predictor_learning_rate": 1.2422135028083982e-07,
+      "step": 9665
+    },
+    {
+      "epoch": 0.99,
+      "gumbel_tau": 0.1023,
+      "image_masks_loss": 0.0039,
+      "llm_learning_rate": 1.2162013904196134e-09,
+      "loss": 0.7411,
+      "output_text_mask_loss": 0.0162,
+      "predictor_learning_rate": 1.2162013904196135e-07,
+      "step": 9666
+    },
+    {
+      "epoch": 0.99,
+      "gumbel_tau": 0.1022,
+      "image_masks_loss": 0.0142,
+      "llm_learning_rate": 1.1904644496685268e-09,
+      "loss": 0.7211,
+      "output_text_mask_loss": 0.0397,
+      "predictor_learning_rate": 1.1904644496685268e-07,
+      "step": 9667
+    },
+    {
+      "epoch": 0.99,
+      "gumbel_tau": 0.1022,
+      "image_masks_loss": 0.0135,
+      "llm_learning_rate": 1.1650026833895377e-09,
+      "loss": 0.727,
+      "output_text_mask_loss": 0.0307,
+      "predictor_learning_rate": 1.1650026833895377e-07,
+      "step": 9668
+    },
+    {
+      "epoch": 0.99,
+      "gumbel_tau": 0.1022,
+      "image_masks_loss": 0.0111,
+      "llm_learning_rate": 1.139816094386237e-09,
+      "loss": 0.7152,
+      "output_text_mask_loss": 0.0429,
+      "predictor_learning_rate": 1.1398160943862368e-07,
+      "step": 9669
+    },
+    {
+      "epoch": 0.99,
+      "gumbel_tau": 0.1022,
+      "image_masks_loss": 0.0049,
+      "llm_learning_rate": 1.1149046854325162e-09,
+      "loss": 0.7573,
+      "output_text_mask_loss": 0.001,
+      "predictor_learning_rate": 1.1149046854325163e-07,
+      "step": 9670
+    },
+    {
+      "epoch": 0.99,
+      "gumbel_tau": 0.1021,
+      "image_masks_loss": 0.0081,
+      "llm_learning_rate": 1.0902684592714596e-09,
+      "loss": 0.7173,
+      "output_text_mask_loss": 0.0542,
+      "predictor_learning_rate": 1.0902684592714595e-07,
+      "step": 9671
+    },
+    {
+      "epoch": 0.99,
+      "gumbel_tau": 0.1021,
+      "image_masks_loss": 0.0085,
+      "llm_learning_rate": 1.0659074186164519e-09,
+      "loss": 0.7513,
+      "output_text_mask_loss": 0.0004,
+      "predictor_learning_rate": 1.065907418616452e-07,
+      "step": 9672
+    },
+    {
+      "epoch": 0.99,
+      "gumbel_tau": 0.1021,
+      "image_masks_loss": 0.0058,
+      "llm_learning_rate": 1.0418215661495146e-09,
+      "loss": 0.6956,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 1.0418215661495145e-07,
+      "step": 9673
+    },
+    {
+      "epoch": 0.99,
+      "gumbel_tau": 0.1021,
+      "image_masks_loss": 0.0153,
+      "llm_learning_rate": 1.0180109045235254e-09,
+      "loss": 0.7323,
+      "output_text_mask_loss": 0.016,
+      "predictor_learning_rate": 1.0180109045235253e-07,
+      "step": 9674
+    },
+    {
+      "epoch": 0.99,
+      "gumbel_tau": 0.102,
+      "image_masks_loss": 0.0117,
+      "llm_learning_rate": 9.94475436360276e-10,
+      "loss": 0.707,
+      "output_text_mask_loss": 0.0572,
+      "predictor_learning_rate": 9.944754363602759e-08,
+      "step": 9675
+    },
+    {
+      "epoch": 0.99,
+      "gumbel_tau": 0.102,
+      "image_masks_loss": 0.0107,
+      "llm_learning_rate": 9.712151642518598e-10,
+      "loss": 0.678,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 9.712151642518597e-08,
+      "step": 9676
+    },
+    {
+      "epoch": 0.99,
+      "gumbel_tau": 0.102,
+      "image_masks_loss": 0.0052,
+      "llm_learning_rate": 9.48230090759561e-10,
+      "loss": 0.7983,
+      "output_text_mask_loss": 0.0095,
+      "predictor_learning_rate": 9.482300907595609e-08,
+      "step": 9677
+    },
+    {
+      "epoch": 0.99,
+      "gumbel_tau": 0.102,
+      "image_masks_loss": 0.0116,
+      "llm_learning_rate": 9.255202184146883e-10,
+      "loss": 0.6854,
+      "output_text_mask_loss": 0.005,
+      "predictor_learning_rate": 9.255202184146883e-08,
+      "step": 9678
+    },
+    {
+      "epoch": 0.99,
+      "gumbel_tau": 0.102,
+      "image_masks_loss": 0.0092,
+      "llm_learning_rate": 9.03085549718019e-10,
+      "loss": 0.7729,
+      "output_text_mask_loss": 0.1205,
+      "predictor_learning_rate": 9.030855497180191e-08,
+      "step": 9679
+    },
+    {
+      "epoch": 0.99,
+      "gumbel_tau": 0.1019,
+      "image_masks_loss": 0.0097,
+      "llm_learning_rate": 8.809260871397995e-10,
+      "loss": 0.7736,
+      "output_text_mask_loss": 0.0231,
+      "predictor_learning_rate": 8.809260871397995e-08,
+      "step": 9680
+    },
+    {
+      "epoch": 0.99,
+      "gumbel_tau": 0.1019,
+      "image_masks_loss": 0.0067,
+      "llm_learning_rate": 8.59041833120855e-10,
+      "loss": 0.6878,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 8.590418331208549e-08,
+      "step": 9681
+    },
+    {
+      "epoch": 0.99,
+      "gumbel_tau": 0.1019,
+      "image_masks_loss": 0.0023,
+      "llm_learning_rate": 8.374327900706469e-10,
+      "loss": 0.7229,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 8.374327900706468e-08,
+      "step": 9682
+    },
+    {
+      "epoch": 0.99,
+      "gumbel_tau": 0.1019,
+      "image_masks_loss": 0.0062,
+      "llm_learning_rate": 8.16098960369216e-10,
+      "loss": 0.8922,
+      "output_text_mask_loss": 0.0467,
+      "predictor_learning_rate": 8.160989603692159e-08,
+      "step": 9683
+    },
+    {
+      "epoch": 0.99,
+      "gumbel_tau": 0.1018,
+      "image_masks_loss": 0.0048,
+      "llm_learning_rate": 7.950403463655165e-10,
+      "loss": 0.7287,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 7.950403463655165e-08,
+      "step": 9684
+    },
+    {
+      "epoch": 0.99,
+      "gumbel_tau": 0.1018,
+      "image_masks_loss": 0.0068,
+      "llm_learning_rate": 7.742569503788044e-10,
+      "loss": 0.7709,
+      "output_text_mask_loss": 0.0006,
+      "predictor_learning_rate": 7.742569503788044e-08,
+      "step": 9685
+    },
+    {
+      "epoch": 0.99,
+      "gumbel_tau": 0.1018,
+      "image_masks_loss": 0.0049,
+      "llm_learning_rate": 7.537487746978045e-10,
+      "loss": 0.6817,
+      "output_text_mask_loss": 0.0231,
+      "predictor_learning_rate": 7.537487746978045e-08,
+      "step": 9686
+    },
+    {
+      "epoch": 0.99,
+      "gumbel_tau": 0.1018,
+      "image_masks_loss": 0.006,
+      "llm_learning_rate": 7.335158215807104e-10,
+      "loss": 0.7865,
+      "output_text_mask_loss": 0.0065,
+      "predictor_learning_rate": 7.335158215807103e-08,
+      "step": 9687
+    },
+    {
+      "epoch": 0.99,
+      "gumbel_tau": 0.1017,
+      "image_masks_loss": 0.0156,
+      "llm_learning_rate": 7.135580932557395e-10,
+      "loss": 0.7434,
+      "output_text_mask_loss": 0.034,
+      "predictor_learning_rate": 7.135580932557395e-08,
+      "step": 9688
+    },
+    {
+      "epoch": 0.99,
+      "gumbel_tau": 0.1017,
+      "image_masks_loss": 0.0075,
+      "llm_learning_rate": 6.938755919208562e-10,
+      "loss": 0.698,
+      "output_text_mask_loss": 0.0025,
+      "predictor_learning_rate": 6.938755919208562e-08,
+      "step": 9689
+    },
+    {
+      "epoch": 0.99,
+      "gumbel_tau": 0.1017,
+      "image_masks_loss": 0.0058,
+      "llm_learning_rate": 6.744683197432155e-10,
+      "loss": 0.6906,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 6.744683197432156e-08,
+      "step": 9690
+    },
+    {
+      "epoch": 0.99,
+      "gumbel_tau": 0.1017,
+      "image_masks_loss": 0.0028,
+      "llm_learning_rate": 6.553362788602746e-10,
+      "loss": 0.6923,
+      "output_text_mask_loss": 0.0099,
+      "predictor_learning_rate": 6.553362788602745e-08,
+      "step": 9691
+    },
+    {
+      "epoch": 0.99,
+      "gumbel_tau": 0.1016,
+      "image_masks_loss": 0.0105,
+      "llm_learning_rate": 6.364794713786814e-10,
+      "loss": 0.7966,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 6.364794713786815e-08,
+      "step": 9692
+    },
+    {
+      "epoch": 0.99,
+      "gumbel_tau": 0.1016,
+      "image_masks_loss": 0.0259,
+      "llm_learning_rate": 6.178978993751083e-10,
+      "loss": 0.6711,
+      "output_text_mask_loss": 0.0097,
+      "predictor_learning_rate": 6.178978993751084e-08,
+      "step": 9693
+    },
+    {
+      "epoch": 0.99,
+      "gumbel_tau": 0.1016,
+      "image_masks_loss": 0.0188,
+      "llm_learning_rate": 5.995915648956963e-10,
+      "loss": 0.7857,
+      "output_text_mask_loss": 0.0121,
+      "predictor_learning_rate": 5.995915648956962e-08,
+      "step": 9694
+    },
+    {
+      "epoch": 0.99,
+      "gumbel_tau": 0.1016,
+      "image_masks_loss": 0.0141,
+      "llm_learning_rate": 5.815604699566102e-10,
+      "loss": 0.7162,
+      "output_text_mask_loss": 0.0221,
+      "predictor_learning_rate": 5.8156046995661014e-08,
+      "step": 9695
+    },
+    {
+      "epoch": 0.99,
+      "gumbel_tau": 0.1015,
+      "image_masks_loss": 0.0024,
+      "llm_learning_rate": 5.638046165432065e-10,
+      "loss": 0.8127,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 5.638046165432065e-08,
+      "step": 9696
+    },
+    {
+      "epoch": 0.99,
+      "gumbel_tau": 0.1015,
+      "image_masks_loss": 0.0162,
+      "llm_learning_rate": 5.463240066108655e-10,
+      "loss": 0.6937,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 5.463240066108655e-08,
+      "step": 9697
+    },
+    {
+      "epoch": 0.99,
+      "gumbel_tau": 0.1015,
+      "image_masks_loss": 0.0189,
+      "llm_learning_rate": 5.291186420847139e-10,
+      "loss": 0.7804,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 5.2911864208471386e-08,
+      "step": 9698
+    },
+    {
+      "epoch": 0.99,
+      "gumbel_tau": 0.1015,
+      "image_masks_loss": 0.0032,
+      "llm_learning_rate": 5.121885248593473e-10,
+      "loss": 0.7266,
+      "output_text_mask_loss": 0.0236,
+      "predictor_learning_rate": 5.121885248593472e-08,
+      "step": 9699
+    },
+    {
+      "epoch": 0.99,
+      "gumbel_tau": 0.1014,
+      "image_masks_loss": 0.0069,
+      "llm_learning_rate": 4.955336567991076e-10,
+      "loss": 0.679,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 4.9553365679910754e-08,
+      "step": 9700
+    },
+    {
+      "epoch": 0.99,
+      "gumbel_tau": 0.1014,
+      "image_masks_loss": 0.0046,
+      "llm_learning_rate": 4.79154039738361e-10,
+      "loss": 0.7152,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 4.791540397383609e-08,
+      "step": 9701
+    },
+    {
+      "epoch": 0.99,
+      "gumbel_tau": 0.1014,
+      "image_masks_loss": 0.0213,
+      "llm_learning_rate": 4.63049675480387e-10,
+      "loss": 0.7413,
+      "output_text_mask_loss": 0.0069,
+      "predictor_learning_rate": 4.6304967548038704e-08,
+      "step": 9702
+    },
+    {
+      "epoch": 0.99,
+      "gumbel_tau": 0.1014,
+      "image_masks_loss": 0.0057,
+      "llm_learning_rate": 4.472205657990447e-10,
+      "loss": 0.7238,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 4.472205657990447e-08,
+      "step": 9703
+    },
+    {
+      "epoch": 0.99,
+      "gumbel_tau": 0.1014,
+      "image_masks_loss": 0.0035,
+      "llm_learning_rate": 4.3166671243738414e-10,
+      "loss": 0.8686,
+      "output_text_mask_loss": 0.0322,
+      "predictor_learning_rate": 4.316667124373841e-08,
+      "step": 9704
+    },
+    {
+      "epoch": 0.99,
+      "gumbel_tau": 0.1013,
+      "image_masks_loss": 0.0206,
+      "llm_learning_rate": 4.163881171079243e-10,
+      "loss": 0.6276,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 4.1638811710792425e-08,
+      "step": 9705
+    },
+    {
+      "epoch": 0.99,
+      "gumbel_tau": 0.1013,
+      "image_masks_loss": 0.0121,
+      "llm_learning_rate": 4.013847814934857e-10,
+      "loss": 0.7388,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 4.013847814934857e-08,
+      "step": 9706
+    },
+    {
+      "epoch": 0.99,
+      "gumbel_tau": 0.1013,
+      "image_masks_loss": 0.0106,
+      "llm_learning_rate": 3.8665670724635784e-10,
+      "loss": 0.7962,
+      "output_text_mask_loss": 0.0103,
+      "predictor_learning_rate": 3.8665670724635784e-08,
+      "step": 9707
+    },
+    {
+      "epoch": 0.99,
+      "gumbel_tau": 0.1013,
+      "image_masks_loss": 0.0121,
+      "llm_learning_rate": 3.7220389598802143e-10,
+      "loss": 0.7799,
+      "output_text_mask_loss": 0.0154,
+      "predictor_learning_rate": 3.722038959880214e-08,
+      "step": 9708
+    },
+    {
+      "epoch": 0.99,
+      "gumbel_tau": 0.1012,
+      "image_masks_loss": 0.0146,
+      "llm_learning_rate": 3.5802634931053625e-10,
+      "loss": 0.8393,
+      "output_text_mask_loss": 0.0353,
+      "predictor_learning_rate": 3.580263493105362e-08,
+      "step": 9709
+    },
+    {
+      "epoch": 0.99,
+      "gumbel_tau": 0.1012,
+      "image_masks_loss": 0.0127,
+      "llm_learning_rate": 3.441240687748759e-10,
+      "loss": 0.8106,
+      "output_text_mask_loss": 0.0343,
+      "predictor_learning_rate": 3.441240687748759e-08,
+      "step": 9710
+    },
+    {
+      "epoch": 0.99,
+      "gumbel_tau": 0.1012,
+      "image_masks_loss": 0.0061,
+      "llm_learning_rate": 3.304970559117604e-10,
+      "loss": 0.8019,
+      "output_text_mask_loss": 0.0429,
+      "predictor_learning_rate": 3.304970559117604e-08,
+      "step": 9711
+    },
+    {
+      "epoch": 1.0,
+      "gumbel_tau": 0.1012,
+      "image_masks_loss": 0.0088,
+      "llm_learning_rate": 3.1714531222221123e-10,
+      "loss": 0.7062,
+      "output_text_mask_loss": 0.0078,
+      "predictor_learning_rate": 3.171453122222112e-08,
+      "step": 9712
+    },
+    {
+      "epoch": 1.0,
+      "gumbel_tau": 0.1011,
+      "image_masks_loss": 0.014,
+      "llm_learning_rate": 3.040688391767188e-10,
+      "loss": 0.7514,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 3.0406883917671877e-08,
+      "step": 9713
+    },
+    {
+      "epoch": 1.0,
+      "gumbel_tau": 0.1011,
+      "image_masks_loss": 0.0164,
+      "llm_learning_rate": 2.912676382146873e-10,
+      "loss": 0.6928,
+      "output_text_mask_loss": 0.001,
+      "predictor_learning_rate": 2.9126763821468726e-08,
+      "step": 9714
+    },
+    {
+      "epoch": 1.0,
+      "gumbel_tau": 0.1011,
+      "image_masks_loss": 0.0122,
+      "llm_learning_rate": 2.7874171074637745e-10,
+      "loss": 0.744,
+      "output_text_mask_loss": 0.0137,
+      "predictor_learning_rate": 2.7874171074637744e-08,
+      "step": 9715
+    },
+    {
+      "epoch": 1.0,
+      "gumbel_tau": 0.1011,
+      "image_masks_loss": 0.0023,
+      "llm_learning_rate": 2.664910581509639e-10,
+      "loss": 0.7469,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 2.6649105815096387e-08,
+      "step": 9716
+    },
+    {
+      "epoch": 1.0,
+      "gumbel_tau": 0.101,
+      "image_masks_loss": 0.0037,
+      "llm_learning_rate": 2.5451568177736754e-10,
+      "loss": 0.7414,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 2.545156817773675e-08,
+      "step": 9717
+    },
+    {
+      "epoch": 1.0,
+      "gumbel_tau": 0.101,
+      "image_masks_loss": 0.0047,
+      "llm_learning_rate": 2.4281558294453334e-10,
+      "loss": 0.788,
+      "output_text_mask_loss": 0.046,
+      "predictor_learning_rate": 2.4281558294453332e-08,
+      "step": 9718
+    },
+    {
+      "epoch": 1.0,
+      "gumbel_tau": 0.101,
+      "image_masks_loss": 0.0068,
+      "llm_learning_rate": 2.313907629405976e-10,
+      "loss": 0.7961,
+      "output_text_mask_loss": 0.0058,
+      "predictor_learning_rate": 2.3139076294059757e-08,
+      "step": 9719
+    },
+    {
+      "epoch": 1.0,
+      "gumbel_tau": 0.101,
+      "image_masks_loss": 0.01,
+      "llm_learning_rate": 2.2024122302427565e-10,
+      "loss": 0.6596,
+      "output_text_mask_loss": 0.0557,
+      "predictor_learning_rate": 2.2024122302427563e-08,
+      "step": 9720
+    },
+    {
+      "epoch": 1.0,
+      "gumbel_tau": 0.1009,
+      "image_masks_loss": 0.004,
+      "llm_learning_rate": 2.0936696442291904e-10,
+      "loss": 0.6441,
+      "output_text_mask_loss": 0.0233,
+      "predictor_learning_rate": 2.0936696442291902e-08,
+      "step": 9721
+    },
+    {
+      "epoch": 1.0,
+      "gumbel_tau": 0.1009,
+      "image_masks_loss": 0.0129,
+      "llm_learning_rate": 1.9876798833390332e-10,
+      "loss": 0.79,
+      "output_text_mask_loss": 0.0332,
+      "predictor_learning_rate": 1.9876798833390332e-08,
+      "step": 9722
+    },
+    {
+      "epoch": 1.0,
+      "gumbel_tau": 0.1009,
+      "image_masks_loss": 0.0176,
+      "llm_learning_rate": 1.8844429592490554e-10,
+      "loss": 0.8387,
+      "output_text_mask_loss": 0.0065,
+      "predictor_learning_rate": 1.8844429592490554e-08,
+      "step": 9723
+    },
+    {
+      "epoch": 1.0,
+      "gumbel_tau": 0.1009,
+      "image_masks_loss": 0.0064,
+      "llm_learning_rate": 1.7839588833251653e-10,
+      "loss": 0.6821,
+      "output_text_mask_loss": 0.0048,
+      "predictor_learning_rate": 1.7839588833251652e-08,
+      "step": 9724
+    },
+    {
+      "epoch": 1.0,
+      "gumbel_tau": 0.1009,
+      "image_masks_loss": 0.0162,
+      "llm_learning_rate": 1.686227666633511e-10,
+      "loss": 0.708,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 1.686227666633511e-08,
+      "step": 9725
+    },
+    {
+      "epoch": 1.0,
+      "gumbel_tau": 0.1008,
+      "image_masks_loss": 0.0176,
+      "llm_learning_rate": 1.5912493199349287e-10,
+      "loss": 0.7068,
+      "output_text_mask_loss": 0.0376,
+      "predictor_learning_rate": 1.5912493199349286e-08,
+      "step": 9726
+    },
+    {
+      "epoch": 1.0,
+      "gumbel_tau": 0.1008,
+      "image_masks_loss": 0.0149,
+      "llm_learning_rate": 1.4990238536904955e-10,
+      "loss": 0.7205,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 1.4990238536904953e-08,
+      "step": 9727
+    },
+    {
+      "epoch": 1.0,
+      "gumbel_tau": 0.1008,
+      "image_masks_loss": 0.0062,
+      "llm_learning_rate": 1.4095512780532005e-10,
+      "loss": 0.809,
+      "output_text_mask_loss": 0.0208,
+      "predictor_learning_rate": 1.4095512780532004e-08,
+      "step": 9728
+    },
+    {
+      "epoch": 1.0,
+      "gumbel_tau": 0.1008,
+      "image_masks_loss": 0.0101,
+      "llm_learning_rate": 1.3228316028818243e-10,
+      "loss": 0.7183,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 1.3228316028818244e-08,
+      "step": 9729
+    },
+    {
+      "epoch": 1.0,
+      "gumbel_tau": 0.1007,
+      "image_masks_loss": 0.0116,
+      "llm_learning_rate": 1.238864837718734e-10,
+      "loss": 0.7368,
+      "output_text_mask_loss": 0.0031,
+      "predictor_learning_rate": 1.238864837718734e-08,
+      "step": 9730
+    },
+    {
+      "epoch": 1.0,
+      "gumbel_tau": 0.1007,
+      "image_masks_loss": 0.0135,
+      "llm_learning_rate": 1.1576509918176381e-10,
+      "loss": 0.6587,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 1.1576509918176382e-08,
+      "step": 9731
+    },
+    {
+      "epoch": 1.0,
+      "gumbel_tau": 0.1007,
+      "image_masks_loss": 0.0165,
+      "llm_learning_rate": 1.0791900741186079e-10,
+      "loss": 0.7329,
+      "output_text_mask_loss": 0.0191,
+      "predictor_learning_rate": 1.0791900741186078e-08,
+      "step": 9732
+    },
+    {
+      "epoch": 1.0,
+      "gumbel_tau": 0.1007,
+      "image_masks_loss": 0.0057,
+      "llm_learning_rate": 1.0034820932591782e-10,
+      "loss": 0.6922,
+      "output_text_mask_loss": 0.0019,
+      "predictor_learning_rate": 1.0034820932591782e-08,
+      "step": 9733
+    },
+    {
+      "epoch": 1.0,
+      "gumbel_tau": 0.1006,
+      "image_masks_loss": 0.0103,
+      "llm_learning_rate": 9.305270575826752e-11,
+      "loss": 0.752,
+      "output_text_mask_loss": 0.0122,
+      "predictor_learning_rate": 9.305270575826752e-09,
+      "step": 9734
+    },
+    {
+      "epoch": 1.0,
+      "gumbel_tau": 0.1006,
+      "image_masks_loss": 0.0099,
+      "llm_learning_rate": 8.60324975118787e-11,
+      "loss": 0.6419,
+      "output_text_mask_loss": 0.0298,
+      "predictor_learning_rate": 8.60324975118787e-09,
+      "step": 9735
+    },
+    {
+      "epoch": 1.0,
+      "gumbel_tau": 0.1006,
+      "image_masks_loss": 0.0141,
+      "llm_learning_rate": 7.928758535974412e-11,
+      "loss": 0.7124,
+      "output_text_mask_loss": 0.0004,
+      "predictor_learning_rate": 7.928758535974413e-09,
+      "step": 9736
+    },
+    {
+      "epoch": 1.0,
+      "gumbel_tau": 0.1006,
+      "image_masks_loss": 0.0079,
+      "llm_learning_rate": 7.281797004515812e-11,
+      "loss": 0.759,
+      "output_text_mask_loss": 0.0376,
+      "predictor_learning_rate": 7.281797004515811e-09,
+      "step": 9737
+    },
+    {
+      "epoch": 1.0,
+      "gumbel_tau": 0.1005,
+      "image_masks_loss": 0.0123,
+      "llm_learning_rate": 6.662365228005119e-11,
+      "loss": 0.6841,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 6.662365228005118e-09,
+      "step": 9738
+    },
+    {
+      "epoch": 1.0,
+      "gumbel_tau": 0.1005,
+      "image_masks_loss": 0.0047,
+      "llm_learning_rate": 6.070463274665539e-11,
+      "loss": 0.7152,
+      "output_text_mask_loss": 0.004,
+      "predictor_learning_rate": 6.0704632746655386e-09,
+      "step": 9739
+    },
+    {
+      "epoch": 1.0,
+      "gumbel_tau": 0.1005,
+      "image_masks_loss": 0.0044,
+      "llm_learning_rate": 5.506091209694919e-11,
+      "loss": 0.7196,
+      "output_text_mask_loss": 0.005,
+      "predictor_learning_rate": 5.506091209694919e-09,
+      "step": 9740
+    },
+    {
+      "epoch": 1.0,
+      "gumbel_tau": 0.1005,
+      "image_masks_loss": 0.0104,
+      "llm_learning_rate": 4.969249095210238e-11,
+      "loss": 0.6895,
+      "output_text_mask_loss": 0.0175,
+      "predictor_learning_rate": 4.969249095210238e-09,
+      "step": 9741
+    },
+    {
+      "epoch": 1.0,
+      "gumbel_tau": 0.1004,
+      "image_masks_loss": 0.0062,
+      "llm_learning_rate": 4.459936990386382e-11,
+      "loss": 0.6866,
+      "output_text_mask_loss": 0.0008,
+      "predictor_learning_rate": 4.459936990386382e-09,
+      "step": 9742
+    },
+    {
+      "epoch": 1.0,
+      "gumbel_tau": 0.1004,
+      "image_masks_loss": 0.0153,
+      "llm_learning_rate": 3.978154951261859e-11,
+      "loss": 0.7409,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 3.978154951261859e-09,
+      "step": 9743
+    },
+    {
+      "epoch": 1.0,
+      "gumbel_tau": 0.1004,
+      "image_masks_loss": 0.0057,
+      "llm_learning_rate": 3.523903030933085e-11,
+      "loss": 0.6704,
+      "output_text_mask_loss": 0.0576,
+      "predictor_learning_rate": 3.5239030309330844e-09,
+      "step": 9744
+    },
+    {
+      "epoch": 1.0,
+      "gumbel_tau": 0.1004,
+      "image_masks_loss": 0.0028,
+      "llm_learning_rate": 3.0971812793600954e-11,
+      "loss": 0.7797,
+      "output_text_mask_loss": 0.0429,
+      "predictor_learning_rate": 3.097181279360095e-09,
+      "step": 9745
+    },
+    {
+      "epoch": 1.0,
+      "gumbel_tau": 0.1004,
+      "image_masks_loss": 0.004,
+      "llm_learning_rate": 2.6979897435885915e-11,
+      "loss": 0.7027,
+      "output_text_mask_loss": 0.0065,
+      "predictor_learning_rate": 2.6979897435885913e-09,
+      "step": 9746
+    },
+    {
+      "epoch": 1.0,
+      "gumbel_tau": 0.1003,
+      "image_masks_loss": 0.0096,
+      "llm_learning_rate": 2.3263284675556498e-11,
+      "loss": 0.7351,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 2.3263284675556496e-09,
+      "step": 9747
+    },
+    {
+      "epoch": 1.0,
+      "gumbel_tau": 0.1003,
+      "image_masks_loss": 0.0101,
+      "llm_learning_rate": 1.9821974922007436e-11,
+      "loss": 0.8101,
+      "output_text_mask_loss": 0.0017,
+      "predictor_learning_rate": 1.9821974922007436e-09,
+      "step": 9748
+    },
+    {
+      "epoch": 1.0,
+      "gumbel_tau": 0.1003,
+      "image_masks_loss": 0.011,
+      "llm_learning_rate": 1.6655968554379897e-11,
+      "loss": 0.7385,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 1.6655968554379897e-09,
+      "step": 9749
+    },
+    {
+      "epoch": 1.0,
+      "gumbel_tau": 0.1003,
+      "image_masks_loss": 0.008,
+      "llm_learning_rate": 1.37652659207288e-11,
+      "loss": 0.7242,
+      "output_text_mask_loss": 0.0341,
+      "predictor_learning_rate": 1.3765265920728798e-09,
+      "step": 9750
+    },
+    {
+      "epoch": 1.0,
+      "gumbel_tau": 0.1002,
+      "image_masks_loss": 0.0065,
+      "llm_learning_rate": 1.1149867339965704e-11,
+      "loss": 0.6849,
+      "output_text_mask_loss": 0.0381,
+      "predictor_learning_rate": 1.1149867339965703e-09,
+      "step": 9751
+    },
+    {
+      "epoch": 1.0,
+      "gumbel_tau": 0.1002,
+      "image_masks_loss": 0.0055,
+      "llm_learning_rate": 8.809773099638375e-12,
+      "loss": 0.7118,
+      "output_text_mask_loss": 0.05,
+      "predictor_learning_rate": 8.809773099638374e-10,
+      "step": 9752
+    },
+    {
+      "epoch": 1.0,
+      "gumbel_tau": 0.1002,
+      "image_masks_loss": 0.0288,
+      "llm_learning_rate": 6.744983457873666e-12,
+      "loss": 0.7047,
+      "output_text_mask_loss": 0.013,
+      "predictor_learning_rate": 6.744983457873666e-10,
+      "step": 9753
+    },
+    {
+      "epoch": 1.0,
+      "gumbel_tau": 0.1002,
+      "image_masks_loss": 0.0064,
+      "llm_learning_rate": 4.955498641989742e-12,
+      "loss": 0.6901,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 4.955498641989741e-10,
+      "step": 9754
+    },
+    {
+      "epoch": 1.0,
+      "gumbel_tau": 0.1001,
+      "image_masks_loss": 0.0204,
+      "llm_learning_rate": 3.441318848773634e-12,
+      "loss": 0.6843,
+      "output_text_mask_loss": 0.0,
+      "predictor_learning_rate": 3.4413188487736334e-10,
+      "step": 9755
+    },
+    {
+      "epoch": 1.0,
+      "gumbel_tau": 0.1001,
+      "image_masks_loss": 0.0117,
+      "llm_learning_rate": 2.2024442450363505e-12,
+      "loss": 0.7701,
+      "output_text_mask_loss": 0.0118,
+      "predictor_learning_rate": 2.2024442450363502e-10,
+      "step": 9756
+    },
+    {
+      "epoch": 1.0,
+      "gumbel_tau": 0.1001,
+      "image_masks_loss": 0.0126,
+      "llm_learning_rate": 1.238874967335324e-12,
+      "loss": 0.6589,
+      "output_text_mask_loss": 0.0002,
+      "predictor_learning_rate": 1.238874967335324e-10,
+      "step": 9757
+    },
+    {
+      "epoch": 1.0,
+      "gumbel_tau": 0.1001,
+      "image_masks_loss": 0.0082,
+      "llm_learning_rate": 5.506111219744093e-13,
+      "loss": 0.723,
+      "output_text_mask_loss": 0.0116,
+      "predictor_learning_rate": 5.506111219744092e-11,
+      "step": 9758
+    },
+    {
+      "epoch": 1.0,
+      "gumbel_tau": 0.1,
+      "image_masks_loss": 0.0109,
+      "llm_learning_rate": 1.376527841712161e-13,
+      "loss": 0.7294,
+      "output_text_mask_loss": 0.0399,
+      "predictor_learning_rate": 1.3765278417121608e-11,
+      "step": 9759
+    },
+    {
+      "epoch": 1.0,
+      "gumbel_tau": 0.1,
+      "image_masks_loss": 0.0053,
+      "llm_learning_rate": 0.0,
+      "loss": 0.6993,
+      "output_text_mask_loss": 0.0038,
+      "predictor_learning_rate": 0.0,
+      "step": 9760
+    },
+    {
+      "epoch": 1.0,
+      "step": 9760,
+      "total_flos": 1446408346820608.0,
+      "train_loss": 0.7873950216980254,
+      "train_runtime": 63786.8455,
+      "train_samples_per_second": 9.792,
+      "train_steps_per_second": 0.153
+    }
+  ],
+  "logging_steps": 1.0,
+  "max_steps": 9760,
+  "num_input_tokens_seen": 0,
+  "num_train_epochs": 1,
+  "save_steps": 40000,
+  "total_flos": 1446408346820608.0,
+  "train_batch_size": 8,
+  "trial_name": null,
+  "trial_params": null
+}