IterativeLoss may have a bug(?)

[naruya]

現在pixyz/DMMを参考にしてSSMを使った映像生成をしています。
(こちらを見ていただきたいです。bair pushing datasetを使っています)
https://github.com/naruya/ssm-pixyz/blob/issue1/SSM
(このコミットのコードで実験を回しました)
https://github.com/naruya/ssm-pixyz/commit/fb044a89c1581bd415def38caf9163cda458d33b

以下の4つのloss関数(は同じ値をとると思うんですが、)をIterativeLossを使って定義したところ、それぞれ異なる値をとっていました。

4つのlossの定義はこちらで
https://github.com/naruya/ssm-pixyz/blob/issue1/SSM/model.py
print_latex結果はこちらにも上げています。
https://github.com/naruya/ssm-pixyz/blob/issue1/SSM/_sample.ipynb

以下のmain.pyを実行して4つのlossの値を比較しました。
https://github.com/naruya/ssm-pixyz/blob/issue1/SSM/main.py

seq_len == 1 のときとseq_len == 3 のときをそれぞれ3回ずつ試しました。
loss1, loss2, loss3, loss4が上の画像に相当するloss関数の値で、
_ce, _klは参考までに出力してみています、具体的な内容はコードを見ていただきたいです。

seq_len == 1 のとき

loss1, loss2, loss3, loss4 を比較すると、loss2だけですが異常な値をとっていました。
loss_ce2, loss_kl2 は loss3, loss4 のときと同じような値をとっているので、自分の書き方が悪かったのかもしれませんが、IterativeLoss や LossOperator の挙動が期待と違うのかもしれないと思い載せました。

seq_len == 3 のとき

loss1, loss2, loss3, loss4 を比較すると、loss3とloss4だけ一致し、その他はバラバラな値をとっていました。(loss3とloss4の実装はもともとほとんど同じではあります。)
seq_len == 1 の比較しloss1も異なる値をとるようになっているので、IterativeLoss内のupdate周りにバグがある気がしています。

以上報告になります。