Return scalar losses instead of per-sample means #18013
Conversation
|
The documentation is not available anymore as the PR was closed or merged. |
There was a problem hiding this comment.
LGTM, thank you @Rocketknight1!
- I didn't see the change in RAG model (which you modified in the last PR and we have some test failures)
- You want to work on this in a separate PR?
- Could you run (some of) the current (related) failed tests before merge?
|
@ydshieh will investigate RAG and make sure tests pass before merging! |
|
@ydshieh I reverted the RAG loss function to the pre-XLA version, so hopefully those tests pass now. All other tests are passing! Do you think there's anything else I'm missing before I merge? |
Nothing I can think of as the tests pass now 🎉 . Thank you! |
| reduced_masked_loss = tf.reduce_sum(masked_loss, axis=1) / loss_denominator | ||
| return reduced_masked_loss | ||
| reduced_masked_loss = tf.reduce_sum(masked_loss) / tf.reduce_sum(loss_mask) | ||
| return tf.reshape(reduced_masked_loss, (1,)) |
There was a problem hiding this comment.
@Rocketknight1 Sorry to bother, but I didn't notice this part yesterday, and now wondering why we need (1,) as returned shapre?
There was a problem hiding this comment.
@ydshieh Keras expects the loss to have at least one dimension, and so standard methods will often fail if the loss is a pure scalar (with shape None). Everything works fine if it has shape (1,) instead.
* Return scalar losses instead of per-sample means * Make loss shape (1,) instead of scalar * Allow scalar losses in test_loss_computation * Allow scalar losses in test_loss_computation * Allow scalar losses in test_loss_computation * Remove XLA loss function for RAG
This updates the TF XLA-compatible losses to return scalars instead of per-sample means. As @ydshieh pointed out, per-sample means give too much weight to samples with fewer masked positions. The new approach should match PyTorch losses exactly (up to floating-point error).
TODO: