Extend ContrastiveOutput to support sequential encoders

[sararb]

Goals ⚽

This PR adds support for negative sampling to the ContrastiveOutput class for session-based models where the query encoder returns a 3-D ragged tensor.

Implementation Details 🚧

• Flatten the values of the query embeddings to match with the sampled negative embeddings.
• Reconstruct the ragged representation using the mask information of the input query.
• Apply the same transformation to positive candidates (in case they are sequential).

Testing Details 🔍

• Add a unit test for defining a transformer-based model as a retrieval model and training it with sampled softmax.

Benchmark 🔍

I used the session-based script (implemented here)) to perform a benchmark of sampled softmax in various configurations, similar to the study conducted in T4Rec (available here)).

Commandline
I use the first four days for training and evaluation is computed for the fifth day. Here is the base command line with the utilized hparam:

python3 session_based.py --metrics_log_frequency 20 --train_path /models/examples/session_based_script/ecomrees_five_days/train --eval_path /models/examples/session_based_script/ecomrees_five_days/valid --schema_path /models/examples/session_based_script/ecomrees_five_days --task multi_class_classification --embedding_dim 448  --d_model 192 --n_layer 3 --n_head 16 --label_smoothing 0.0  --model_type xlnet --eval_batch_size 128 --train_batch_size 128 --epochs 5 --weight_tying --xlnet_attn_type bi --training_task masked --evaluation_task last --masking_probability 0.30000000000000004 --lr 0.0006667377132554976 --transformer_dropout 0.0  --log_to_wandb --transformer_activation gelu --feature_normalization --input_dropout 0.1 --optimizer adamw --weight_decay 3.910060265627374e-05 --save_topk_predictions --emb_init_std 0.11 --sampled_softmax --num_negatives 1000 --logq_correction

The hparams that are changed for the experiments are --sampled_softmax (enables sampled softmax if provided), --logq_correction, and --num_negatives (number of negative samples).

Results
The results can be seen in the following table. Average examples/sec represents the throughtput and Recall and NDCG are accuracy top-k metrics.

[github-actions]

Documentation preview

https://nvidia-merlin.github.io/models/review/pr-1086

[gabrielspmoreira]

Sounds good to me

[gabrielspmoreira]

This with_flat_values is very useful! Seems faster than rebuilding the full ragged tensor.