Skip to content

Add ListMLE Loss #130

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Open
wants to merge 10 commits into
base: main
Choose a base branch
from
Open

Add ListMLE Loss #130

wants to merge 10 commits into from

Conversation

@LakshmiKalaKadali
Copy link

@LakshmiKalaKadali LakshmiKalaKadali commented Jun 24, 2025 •
edited
Loading

Added ListMLELoss code to listwise ranking. This code does not consider Lambda weights.
Here is the gist for verified results with TFRS ListMLELoss

@LakshmiKalaKadali LakshmiKalaKadali marked this pull request as ready for review June 25, 2025 16:37
@abheesht17
Copy link
Collaborator

Not sure why tests did not run, commenting so that tests run

abheesht17
abheesht17 requested changes Aug 11, 2025
View reviewed changes
Copy link
Collaborator

@abheesht17 abheesht17 left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Sorry for the delay, did one pass! Some questions below, let me know what you think.

I'll take a closer look again!

keras_rs/api/layers/__init__.py Outdated
Comment on lines 8 to 16
from keras_rs.src.layers.embedding.distributed_embedding import DistributedEmbedding as DistributedEmbedding
from keras_rs.src.layers.embedding.distributed_embedding_config import FeatureConfig as FeatureConfig
from keras_rs.src.layers.embedding.distributed_embedding_config import TableConfig as TableConfig
from keras_rs.src.layers.embedding.embed_reduce import EmbedReduce as EmbedReduce
from keras_rs.src.layers.feature_interaction.dot_interaction import DotInteraction as DotInteraction
from keras_rs.src.layers.feature_interaction.feature_cross import FeatureCross as FeatureCross
from keras_rs.src.layers.retrieval.brute_force_retrieval import BruteForceRetrieval as BruteForceRetrieval
from keras_rs.src.layers.retrieval.hard_negative_mining import HardNegativeMining as HardNegativeMining
from keras_rs.src.layers.retrieval.remove_accidental_hits import RemoveAccidentalHits as RemoveAccidentalHits
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hmmm, these shouldn't be getting re-formatted, in my opinion. Are you on the correct Ruff version? Same for other files

@LakshmiKalaKadali
Copy link
Author

Hi @abheesht17, Done suggested changes. Could you please review the PR. Thank you

abheesht17
abheesht17 reviewed Oct 8, 2025
View reviewed changes
Copy link
Collaborator

@abheesht17 abheesht17 left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This LGTM! I'm going to drop an approval on this, but could you please format the code first?

@abheesht17
Copy link
Collaborator

Not sure why we don't have a way to run tests here, need to figure that out before merging

abheesht17
abheesht17 reviewed Oct 8, 2025
View reviewed changes
Copy link
Collaborator

@abheesht17 abheesht17 left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Noticed a few nits. Instead of using +, -, *, /, let's use Keras ops everywhere

@abheesht17
Copy link
Collaborator

@LakshmiKalaKadali - can you look into the test failures above?

@LakshmiKalaKadali
Copy link
Author

Hi @abheesht17,

The torch test fails are due to precision mismatch in between actual and expected values. That precision mismatch might be due to the mask AFTER exp(). Applied the mask to set invalid positions to (-1e9) before calling ops.exp(), and removed the masking that happens after exp(). I hope that this resolves the torch test failure. Could you please run the tests. Thank You

@hertschuh
Copy link
Collaborator

@LakshmiKalaKadali

The version of torch was updated for the tests. Please rebase to see if it helps.

@LakshmiKalaKadali LakshmiKalaKadali force-pushed the listmleloss branch 2 times, most recently from e986ef5 to 3da3fa2 Compare October 23, 2025 07:43
@hertschuh
Copy link
Collaborator

@LakshmiKalaKadali

I do not understand why the results are different. But the differences you see are significant enough that it has to be a bug. Can you try to debug by printing all intermediate values in compute_unreduced_loss and call and comparing the values for Torch and JAX or TF? Hopefully we can narrow it down to a single line of code.

@LakshmiKalaKadali
Copy link
Author

Sure @hertschuh, I will deep dive into the intermediate values. Thank You

LakshmiKalaKadali and others added 10 commits October 24, 2025 13:21
@LakshmiKalaKadali
Add STU layer
730bc21
@LakshmiKalaKadali @gemini-code-assist
Update keras_rs/src/layers/stu.py
1906f2d 
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
@LakshmiKalaKadali @gemini-code-assist
Update keras_rs/src/layers/common.py
d0add4f 
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
@LakshmiKalaKadali @gemini-code-assist
Update keras_rs/src/layers/jagged_tensors.py
eb98ae1 
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
@LakshmiKalaKadali @gemini-code-assist
Update keras_rs/src/layers/hstu_mha_attention.py
0347b51 
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
@LakshmiKalaKadali @gemini-code-assist
Update keras_rs/src/layers/stu.py
570df2b 
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
@LakshmiKalaKadali @gemini-code-assist
Update keras_rs/src/layers/stu.py
6f3dcef 
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
@LakshmiKalaKadali @gemini-code-assist
Update keras_rs/src/layers/jagged_tensors.py
c80e598 
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
@LakshmiKalaKadali @gemini-code-assist
Update keras_rs/src/layers/hstu_compute_output.py
8073dcb 
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
@LakshmiKalaKadali
Add list mle loss
8cb2d98
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@abheesht17 abheesht17 abheesht17 approved these changes

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

3 participants

@LakshmiKalaKadali @abheesht17 @hertschuh