Aux loss decreases at the beginning but increases again

[audreyeternal]

Hi, I want to fine tune the compression model on our own dataset. I use the bmshj2018-factorized model(quality:8, criterion: ms-ssim)set both the aux optimizer and the normal optimizer to 1e-5. Also the lambda is set according to the recommendation in the paper. But during the training process, the aux loss decreases at first but rapidly goes up, like the following:

Train epoch 0: [0/3420 (0%)]	Loss: 1.046 |	ms_ssim loss: 0.998 |	Bpp loss: 0.67 |	Aux loss: 143.82
Train epoch 0: [40/3420 (1%)]	Loss: 1.240 |	ms_ssim loss: 0.998 |	Bpp loss: 0.81 |	Aux loss: 122.47
Train epoch 0: [80/3420 (2%)]	Loss: 1.117 |	ms_ssim loss: 0.998 |	Bpp loss: 0.73 |	Aux loss: 110.71
Train epoch 0: [120/3420 (4%)]	Loss: 1.085 |	ms_ssim loss: 0.998 |	Bpp loss: 0.72 |	Aux loss: 104.04
Train epoch 0: [160/3420 (5%)]	Loss: 1.204 |	ms_ssim loss: 0.998 |	Bpp loss: 0.80 |	Aux loss: 99.57
Train epoch 0: [200/3420 (6%)]	Loss: 1.138 |	ms_ssim loss: 0.998 |	Bpp loss: 0.76 |	Aux loss: 96.21
Train epoch 0: [240/3420 (7%)]	Loss: 1.086 |	ms_ssim loss: 0.998 |	Bpp loss: 0.73 |	Aux loss: 93.07
Train epoch 0: [280/3420 (8%)]	Loss: 1.098 |	ms_ssim loss: 0.998 |	Bpp loss: 0.73 |	Aux loss: 90.06
Train epoch 0: [320/3420 (9%)]	Loss: 1.141 |	ms_ssim loss: 0.998 |	Bpp loss: 0.76 |	Aux loss: 87.25
Train epoch 0: [360/3420 (11%)]	Loss: 1.189 |	ms_ssim loss: 0.998 |	Bpp loss: 0.81 |	Aux loss: 84.61
Train epoch 0: [400/3420 (12%)]	Loss: 1.118 |	ms_ssim loss: 0.998 |	Bpp loss: 0.75 |	Aux loss: 81.90
Train epoch 0: [440/3420 (13%)]	Loss: 1.121 |	ms_ssim loss: 0.998 |	Bpp loss: 0.75 |	Aux loss: 79.35
Train epoch 0: [480/3420 (14%)]	Loss: 1.104 |	ms_ssim loss: 0.998 |	Bpp loss: 0.75 |	Aux loss: 76.93
Train epoch 0: [520/3420 (15%)]	Loss: 1.169 |	ms_ssim loss: 0.998 |	Bpp loss: 0.79 |	Aux loss: 74.64
Train epoch 0: [560/3420 (16%)]	Loss: 1.160 |	ms_ssim loss: 0.998 |	Bpp loss: 0.79 |	Aux loss: 72.47
Train epoch 0: [600/3420 (18%)]	Loss: 1.082 |	ms_ssim loss: 0.998 |	Bpp loss: 0.74 |	Aux loss: 70.18
Train epoch 0: [640/3420 (19%)]	Loss: 1.070 |	ms_ssim loss: 0.998 |	Bpp loss: 0.72 |	Aux loss: 67.87
Train epoch 0: [680/3420 (20%)]	Loss: 1.091 |	ms_ssim loss: 0.998 |	Bpp loss: 0.75 |	Aux loss: 65.38
Train epoch 0: [720/3420 (21%)]	Loss: 1.083 |	ms_ssim loss: 0.998 |	Bpp loss: 0.74 |	Aux loss: 62.85
Train epoch 0: [760/3420 (22%)]	Loss: 1.140 |	ms_ssim loss: 0.998 |	Bpp loss: 0.78 |	Aux loss: 60.37
Train epoch 0: [800/3420 (23%)]	Loss: 1.147 |	ms_ssim loss: 0.998 |	Bpp loss: 0.79 |	Aux loss: 58.04
Train epoch 0: [840/3420 (25%)]	Loss: 0.999 |	ms_ssim loss: 0.999 |	Bpp loss: 0.67 |	Aux loss: 55.86
Train epoch 0: [880/3420 (26%)]	Loss: 1.102 |	ms_ssim loss: 0.998 |	Bpp loss: 0.75 |	Aux loss: 53.68
Train epoch 0: [920/3420 (27%)]	Loss: 1.117 |	ms_ssim loss: 0.998 |	Bpp loss: 0.77 |	Aux loss: 51.57
Train epoch 0: [960/3420 (28%)]	Loss: 1.084 |	ms_ssim loss: 0.998 |	Bpp loss: 0.74 |	Aux loss: 49.33
Train epoch 0: [1000/3420 (29%)]	Loss: 1.022 |	ms_ssim loss: 0.998 |	Bpp loss: 0.69 |	Aux loss: 47.28
Train epoch 0: [1040/3420 (30%)]	Loss: 1.096 |	ms_ssim loss: 0.998 |	Bpp loss: 0.75 |	Aux loss: 45.25
Train epoch 0: [1080/3420 (32%)]	Loss: 1.058 |	ms_ssim loss: 0.998 |	Bpp loss: 0.73 |	Aux loss: 43.33
Train epoch 0: [1120/3420 (33%)]	Loss: 1.156 |	ms_ssim loss: 0.998 |	Bpp loss: 0.80 |	Aux loss: 41.65
Train epoch 0: [1160/3420 (34%)]	Loss: 1.075 |	ms_ssim loss: 0.998 |	Bpp loss: 0.74 |	Aux loss: 40.04
Train epoch 0: [1200/3420 (35%)]	Loss: 1.080 |	ms_ssim loss: 0.999 |	Bpp loss: 0.75 |	Aux loss: 38.63
Train epoch 0: [1240/3420 (36%)]	Loss: 1.191 |	ms_ssim loss: 0.998 |	Bpp loss: 0.84 |	Aux loss: 37.36
Train epoch 0: [1280/3420 (37%)]	Loss: 1.091 |	ms_ssim loss: 0.999 |	Bpp loss: 0.77 |	Aux loss: 36.08
Train epoch 0: [1320/3420 (39%)]	Loss: 1.139 |	ms_ssim loss: 0.998 |	Bpp loss: 0.79 |	Aux loss: 34.90
Train epoch 0: [1360/3420 (40%)]	Loss: 1.069 |	ms_ssim loss: 0.998 |	Bpp loss: 0.73 |	Aux loss: 33.77
Train epoch 0: [1400/3420 (41%)]	Loss: 1.112 |	ms_ssim loss: 0.998 |	Bpp loss: 0.77 |	Aux loss: 32.84
Train epoch 0: [1440/3420 (42%)]	Loss: 1.098 |	ms_ssim loss: 0.998 |	Bpp loss: 0.76 |	Aux loss: 32.02
Train epoch 0: [1480/3420 (43%)]	Loss: 1.122 |	ms_ssim loss: 0.998 |	Bpp loss: 0.78 |	Aux loss: 31.42
Train epoch 0: [1520/3420 (44%)]	Loss: 1.016 |	ms_ssim loss: 0.999 |	Bpp loss: 0.69 |	Aux loss: 30.93
Train epoch 0: [1560/3420 (46%)]	Loss: 1.159 |	ms_ssim loss: 0.998 |	Bpp loss: 0.81 |	Aux loss: 30.52
Train epoch 0: [1600/3420 (47%)]	Loss: 1.078 |	ms_ssim loss: 0.999 |	Bpp loss: 0.75 |	Aux loss: 30.32
Train epoch 0: [1640/3420 (48%)]	Loss: 1.051 |	ms_ssim loss: 0.999 |	Bpp loss: 0.73 |	Aux loss: 30.38
Train epoch 0: [1680/3420 (49%)]	Loss: 1.066 |	ms_ssim loss: 0.998 |	Bpp loss: 0.73 |	Aux loss: 30.64
Train epoch 0: [1720/3420 (50%)]	Loss: 1.035 |	ms_ssim loss: 0.999 |	Bpp loss: 0.71 |	Aux loss: 31.15
Train epoch 0: [1760/3420 (51%)]	Loss: 1.110 |	ms_ssim loss: 0.998 |	Bpp loss: 0.77 |	Aux loss: 31.84
Train epoch 0: [1800/3420 (53%)]	Loss: 1.087 |	ms_ssim loss: 0.998 |	Bpp loss: 0.75 |	Aux loss: 32.69
Train epoch 0: [1840/3420 (54%)]	Loss: 1.029 |	ms_ssim loss: 0.999 |	Bpp loss: 0.71 |	Aux loss: 33.65
Train epoch 0: [1880/3420 (55%)]	Loss: 1.079 |	ms_ssim loss: 0.998 |	Bpp loss: 0.74 |	Aux loss: 34.72
Train epoch 0: [1920/3420 (56%)]	Loss: 1.099 |	ms_ssim loss: 0.998 |	Bpp loss: 0.76 |	Aux loss: 35.87
Train epoch 0: [1960/3420 (57%)]	Loss: 1.050 |	ms_ssim loss: 0.999 |	Bpp loss: 0.73 |	Aux loss: 37.08
Train epoch 0: [2000/3420 (58%)]	Loss: 1.097 |	ms_ssim loss: 0.998 |	Bpp loss: 0.76 |	Aux loss: 38.35
Train epoch 0: [2040/3420 (60%)]	Loss: 1.039 |	ms_ssim loss: 0.999 |	Bpp loss: 0.71 |	Aux loss: 39.81
Train epoch 0: [2080/3420 (61%)]	Loss: 1.122 |	ms_ssim loss: 0.998 |	Bpp loss: 0.79 |	Aux loss: 41.32
Train epoch 0: [2120/3420 (62%)]	Loss: 1.104 |	ms_ssim loss: 0.998 |	Bpp loss: 0.77 |	Aux loss: 42.79
Train epoch 0: [2160/3420 (63%)]	Loss: 1.055 |	ms_ssim loss: 0.999 |	Bpp loss: 0.73 |	Aux loss: 44.47
Train epoch 0: [2200/3420 (64%)]	Loss: 1.101 |	ms_ssim loss: 0.999 |	Bpp loss: 0.77 |	Aux loss: 46.08
Train epoch 0: [2240/3420 (65%)]	Loss: 1.100 |	ms_ssim loss: 0.998 |	Bpp loss: 0.77 |	Aux loss: 47.84
Train epoch 0: [2280/3420 (67%)]	Loss: 1.116 |	ms_ssim loss: 0.998 |	Bpp loss: 0.78 |	Aux loss: 49.56
Train epoch 0: [2320/3420 (68%)]	Loss: 1.082 |	ms_ssim loss: 0.999 |	Bpp loss: 0.76 |	Aux loss: 51.33
Train epoch 0: [2360/3420 (69%)]	Loss: 1.073 |	ms_ssim loss: 0.999 |	Bpp loss: 0.75 |	Aux loss: 53.11
Train epoch 0: [2400/3420 (70%)]	Loss: 1.004 |	ms_ssim loss: 0.999 |	Bpp loss: 0.69 |	Aux loss: 55.06
Train epoch 0: [2440/3420 (71%)]	Loss: 1.051 |	ms_ssim loss: 0.999 |	Bpp loss: 0.73 |	Aux loss: 57.04
Train epoch 0: [2480/3420 (73%)]	Loss: 1.031 |	ms_ssim loss: 0.999 |	Bpp loss: 0.70 |	Aux loss: 58.93
Train epoch 0: [2520/3420 (74%)]	Loss: 1.076 |	ms_ssim loss: 0.999 |	Bpp loss: 0.75 |	Aux loss: 60.76
Train epoch 0: [2560/3420 (75%)]	Loss: 1.036 |	ms_ssim loss: 0.999 |	Bpp loss: 0.73 |	Aux loss: 62.62
Train epoch 0: [2600/3420 (76%)]	Loss: 1.099 |	ms_ssim loss: 0.999 |	Bpp loss: 0.77 |	Aux loss: 64.57
Train epoch 0: [2640/3420 (77%)]	Loss: 1.062 |	ms_ssim loss: 0.999 |	Bpp loss: 0.74 |	Aux loss: 66.61
Train epoch 0: [2680/3420 (78%)]	Loss: 1.011 |	ms_ssim loss: 0.999 |	Bpp loss: 0.70 |	Aux loss: 68.64
Train epoch 0: [2720/3420 (80%)]	Loss: 1.034 |	ms_ssim loss: 0.999 |	Bpp loss: 0.72 |	Aux loss: 70.75
Train epoch 0: [2760/3420 (81%)]	Loss: 1.084 |	ms_ssim loss: 0.998 |	Bpp loss: 0.75 |	Aux loss: 72.82
Train epoch 0: [2800/3420 (82%)]	Loss: 1.075 |	ms_ssim loss: 0.999 |	Bpp loss: 0.75 |	Aux loss: 74.92
Train epoch 0: [2840/3420 (83%)]	Loss: 1.082 |	ms_ssim loss: 0.998 |	Bpp loss: 0.75 |	Aux loss: 77.03
Train epoch 0: [2880/3420 (84%)]	Loss: 1.089 |	ms_ssim loss: 0.999 |	Bpp loss: 0.77 |	Aux loss: 79.07
Train epoch 0: [2920/3420 (85%)]	Loss: 1.139 |	ms_ssim loss: 0.999 |	Bpp loss: 0.82 |	Aux loss: 81.23
Train epoch 0: [2960/3420 (87%)]	Loss: 1.137 |	ms_ssim loss: 0.998 |	Bpp loss: 0.80 |	Aux loss: 83.36
Train epoch 0: [3000/3420 (88%)]	Loss: 1.125 |	ms_ssim loss: 0.999 |	Bpp loss: 0.81 |	Aux loss: 85.50
Train epoch 0: [3040/3420 (89%)]	Loss: 1.058 |	ms_ssim loss: 0.999 |	Bpp loss: 0.74 |	Aux loss: 87.58
Train epoch 0: [3080/3420 (90%)]	Loss: 1.137 |	ms_ssim loss: 0.998 |	Bpp loss: 0.80 |	Aux loss: 89.53
Train epoch 0: [3120/3420 (91%)]	Loss: 1.084 |	ms_ssim loss: 0.999 |	Bpp loss: 0.76 |	Aux loss: 91.55
Train epoch 0: [3160/3420 (92%)]	Loss: 1.051 |	ms_ssim loss: 0.999 |	Bpp loss: 0.73 |	Aux loss: 93.47
Train epoch 0: [3200/3420 (94%)]	Loss: 1.060 |	ms_ssim loss: 0.999 |	Bpp loss: 0.74 |	Aux loss: 95.48
Train epoch 0: [3240/3420 (95%)]	Loss: 1.084 |	ms_ssim loss: 0.999 |	Bpp loss: 0.76 |	Aux loss: 97.50
Train epoch 0: [3280/3420 (96%)]	Loss: 1.054 |	ms_ssim loss: 0.999 |	Bpp loss: 0.73 |	Aux loss: 99.55
Train epoch 0: [3320/3420 (97%)]	Loss: 1.045 |	ms_ssim loss: 0.999 |	Bpp loss: 0.73 |	Aux loss: 101.71
Train epoch 0: [3360/3420 (98%)]	Loss: 1.020 |	ms_ssim loss: 0.999 |	Bpp loss: 0.70 |	Aux loss: 103.92
Train epoch 0: [3400/3420 (99%)]	Loss: 0.978 |	ms_ssim loss: 0.999 |	Bpp loss: 0.67 |	Aux loss: 106.15
Test epoch 0: Average losses:	Loss: 1.023 |	ms_ssim loss: 0.998 |	Bpp loss: 0.68 |	Aux loss: 107.07

Should I stop the training when the aux loss is the lowest? Or what should I do? Thank you!

[YodaEmbedding]

Focus on the loss only. aux_loss can mostly be ignored since it only plays a small role, as described here: https://interdigitalinc.github.io/CompressAI/models.html#compressai.models.CompressionModel.aux_loss

If you really want to bring it down to 0, try: #231.

Related links:

• Training using MSE than finetuning using MS-SSIM but aux loss goes bigger #232
• Aux loss increases after a couple thousand iterations #148
• About Auxiliary loss optimization #167
• i have two gaussian_conditional ,when update model, why one update,another one dont update #217 (comment)