Skip to content

refactor: migration rbm tensorflow to pytorch - #2359

Open
ds-wook wants to merge 3 commits into
recommenders-team:stagingfrom
ds-wook:refactor/migration-rbm-pytorch
Open

refactor: migration rbm tensorflow to pytorch#2359
ds-wook wants to merge 3 commits into
recommenders-team:stagingfrom
ds-wook:refactor/migration-rbm-pytorch

Conversation

@ds-wook

@ds-wook ds-wook commented Jul 16, 2026

Copy link
Copy Markdown
Contributor

Description

Migrate the RBM collaborative-filtering model from TensorFlow 1.x to PyTorch.

recommenders/models/rbm/rbm.py is rewritten as a standalone torch.nn.Module, following the same single-file migration pattern established by the LightGCN port (#2315) — no porting of the TF1 graph/sess.run base classes. The public API is kept unchanged so downstream code and the example notebook require no interface changes: RBM(...) constructor signature, fit(), recommend_k_items(), predict(), save()/load(), and the rmse_train attribute all behave as before.

Key changes:

Parameters (w, bv, bh) become nn.Parameter; the tf.lookup.StaticHashTable rating lookup becomes a registered tensor buffer.
The tf.data input pipeline + sess.run training loop is replaced by torch.randperm minibatching with an Adam optimizer.
Contrastive Divergence: the negative phase (gibbs_sampling) runs under torch.no_grad(), and the free-energy difference F(v) - F(v_k) is backpropagated.
Automatic GPU/CPU device selection; checkpoints are saved/loaded via state_dict (.pt).
Behavior fix: the original TF graph applied dropout at inference time (injecting noise into recommendations). The PyTorch version disables dropout in eval() mode, which improves ranking quality (MovieLens-100k MAP@10 ≈ 0.49 / nDCG@10 ≈ 0.64 vs. ≈ 0.14 / 0.41 previously).
examples/00_quick_start/rbm_movielens.ipynb is updated to import torch instead of tensorflow, refresh the TF-specific narrative to PyTorch, switch checkpoint paths to .pt, and use f-strings. The notebook was executed end-to-end on GPU with all cells passing.

Related Issues

References

Checklist:

  • I have followed the contribution guidelines and code style for this project.
  • I have added tests covering my contributions.
  • I have updated the documentation accordingly.
  • I have signed the commits, e.g. git commit -s -m "your commit message".
  • This PR is being made to staging branch AND NOT TO main branch.

@review-notebook-app

Copy link
Copy Markdown

Check out this pull request on  ReviewNB

See visual diffs & provide feedback on Jupyter Notebooks.


Powered by ReviewNB

@ds-wook
ds-wook force-pushed the refactor/migration-rbm-pytorch branch from 685f404 to 6904422 Compare July 16, 2026 08:50

@miguelgfierro miguelgfierro left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

@ds-wook really good stuff again! congrats!

I added some comments, please take a look

Comment on lines -339 to -346
"model = RBM(\n",
" possible_ratings=np.setdiff1d(np.unique(Xtr), np.array([0])),\n",
" visible_units=Xtr.shape[1],\n",
" hidden_units=600,\n",
" training_epoch=30,\n",
" minibatch_size=60,\n",
" keep_prob=0.9,\n",
" with_metrics=True\n",

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

The old implementation was not very well written. I think it would be good to improve the new one now that we are making the change.
The constructor should have only properties of the model and the trainer fit properties of the training process.
So for example, visible_units=Xtr.shape[1], hidden_units=600, is part of the model and training_epoch=30, minibatch_size=60, keep_prob=0.9, with_metrics=True, would be part of the training

Comment thread recommenders/models/rbm/rbm.py
Comment thread examples/00_quick_start/rbm_movielens.ipynb Outdated
" k=K,\n",
" )\n",
"\n",
" eval_ndcg = ndcg_at_k(\n",

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

the difference in the metrics could be because of this:
Silent behavior change in Gibbs annealing (not a regression — arguably a fix). In the original TF code,
gibbs_protocol re-called self.gibbs_sampling() to rebuild self.v_k, but the optimizer op had already been built
against the initial k=1 graph, so increasing k had no effect on training. The PyTorch version uses k directly in the loop, so CD-k annealing now genuinely takes effect. This changes training dynamics vs. the TF baseline — worth a sentence in the PR body, since it (not just the dropout fix) contributes to the metric differences.

  • gibbs_protocol bounds fix is correct and good. The added guard self.l + 1 < len(self.sampling_protocol) fixes a latent IndexError in the original (which indexed sampling_protocol[self.l + 1] unguarded). Good catch — but there's no test exercising it.
  • num_minibatches = int(n_users / self.minibatch) drops the remainder and can be 0. If minibatch_size > n_users the training loop never executes and the model is returned untrained (rmse_train all zeros) with no warning. This matches the original, so it's not a regression, but a fail-fast check (or at least max(1, ...)) would be a cheap robustness win.

@ds-wook

ds-wook commented Jul 23, 2026

Copy link
Copy Markdown
Contributor Author

@miguelgfierro I've incorporated the requested changes. Could you please review them when you have a chance? Thanks!

@@ -1,870 +1,906 @@
{
"cells": [

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

There is a big difference in the metrics:

Image

There is something that is not working.

Might be that we are leaking test data into the train process?

Also, please notice that the deep dive notebook also needs to be updated: https://github.com/recommenders-team/recommenders/blob/main/examples/02_model_collaborative_filtering/rbm_deep_dive.ipynb

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

Thanks for addressing the previous comments — the constructor/fit() split, the ported tests, and the new fail-fast checks all look good.

Before merging, I investigated the metric numbers in the notebook (MAP going from 0.14 to 0.53). I ran both the old TF code and this PR on the same data split. Here is what I found, in simple terms.

1. The old implementation was already wrong: it leaks the test set

The notebook evaluates with model.recommend_k_items(Xtst). This gives the model the test set as input, and then we check if it recommends those same test items. The model just has to echo its input to score well. This has been the case since the original TF notebook.

Proof: if we give the model the correct input (Xtr, the training ratings), the metrics drop to random level for both implementations:

Input Old TF (precision@10) This PR (precision@10)
Xtst (what the notebook does) 0.32 0.57
Xtr (correct) 0.003 0.014

So neither 0.14 nor 0.53 is a real number. The PyTorch model scores higher only because it echoes its input better, not because it recommends better. On top of that, map_at_k changed definition in 2023 (#2004), which inflates the new MAP even more. The dropout explanation in the PR description is not the cause — I tested it, and forcing dropout at inference barely changes the results. Please remove the improvement claim from the PR description and the notebook.

2. The implementation does not match the paper

I also evaluated the model the way the paper does (predict the held-out ratings, measure RMSE). The result: the model is worse than just predicting the average rating for everything.

Predictor Test RMSE (ML100k)
Always predict the global mean 1.13
Item mean 1.03
RBM (this PR, notebook config) 1.67
Paper's RBM (Netflix) 0.91, beating the 0.95 baseline

The reason is that our implementation simplifies the paper too much:

  • The paper uses one-hot softmax visible units: 5 binary units per item, with separate weights and biases per rating value. Our code uses a single scalar unit per item with p(v=l) ∝ exp(l·φ). This forces the rating distribution to be monotone — the model cannot even represent "most people rate this movie a 3".
  • multinomial_sampling is not real categorical sampling. It compares all probabilities against one shared random threshold. predict() uses it, which adds a lot of noise (RMSE 2.49 sampled vs. 1.67 deterministic).
  • The default learning rate is effectively 0.004/60 ≈ 0.00007. At that rate the training RMSE does not decrease at all.

3. What we would need to change to get the paper's numbers

  1. Replace the visible layer with the paper's one-hot softmax units: weights of shape (n_items, n_ratings, n_hidden) and biases of shape (n_items, n_ratings).
  2. Use real categorical sampling (torch.multinomial) during training, and a deterministic expected rating (Σ rating · probability) for predict().
  3. Fix the training defaults: remove the learning_rate / minibatch_size division (or raise the default), and train until the training RMSE converges.
  4. Evaluate the way the paper does: input Xtr, measure RMSE on the held-out ratings. Target on ML100k: ≈ 0.95, and at minimum beat the item-mean baseline (1.03). If we also report ranking metrics, the input must be Xtr, never Xtst.

Reference implementations to compare against: the paper (https://www.cs.toronto.edu/~rsalakhu/papers/rbmcf.pdf), https://github.com/felipecruz/CFRBM (Theano, ML100k, one-hot units), and https://github.com/erwtokritos/collaborativefiltering-rbm (Java, follows the paper exactly).

One minor leftover from the earlier round: in the metrics cell, please pass the arguments literally in each call instead of the shared args / kwargs variables — same style as https://github.com/recommenders-team/recommenders/blob/main/examples/02_model_collaborative_filtering/sar_deep_dive.ipynb.

@ds-wook

ds-wook commented Jul 28, 2026

Copy link
Copy Markdown
Contributor Author

@miguelgfierro
Thanks for the detailed investigation — you were right on both counts, and the numbers reproduced exactly on my side. I've reworked the implementation to follow the paper. Summary of what changed and the resulting numbers.

1. The leak is closed

Confirmed your finding first. On the same split, with the old TF-ported code:

Input precision@10
Xtst (what the notebook did) 0.592
Xtr (correct) 0.014

I've removed the improvement claim from the PR description. Both notebooks now pass Xtr to recommend_k_items() / predict() and use Xtst only as ground truth, with a comment at the call site explaining why.

Worth noting: remove_seen used to mask against the train matrix stored in fit(), which is what let the leaky call score so well. It now masks against the input matrix — the same thing standard_vae.py and multinomial_vae.py already do — so it also works for users that were not in the train set. As a side effect, the old leaky call now returns exactly 0.0 on every ranking metric, which is a nice sanity check that the ground truth is no longer being fed in.

2. Implementation now matches the paper

  • One-hot softmax visible units. w is now (n_items, n_ratings, n_hidden) and bv is (n_items, n_ratings), so every rating value of every item has its own weights and bias. The old p(v=l) ∝ exp(l·φ) parametrization was monotone in l by construction and could not represent "most users rate this movie a 3".
  • Real categorical sampling. multinomial_sampling() uses torch.multinomial, one independent draw per (user, item) pair, instead of comparing the whole distribution against a single shared threshold.
  • Deterministic inference. predict() returns the expected rating Σ l · P(v=l|h) with mean-field hidden units, so it is the MMSE estimator and two consecutive calls agree.
  • Training defaults fixed. Removed the learning_rate / minibatch_size division (the free energy is now averaged over the minibatch, so the learning rate is batch-size independent). New defaults: learning_rate=0.001, l2=0.01, training_epoch=100. The training RMSE used to increase (1.65 → 1.72); it now decreases and flattens (0.97 → 0.69).
  • bv is initialized to the empirical log rating frequencies, guarded by a flag in the state dict so that loading a checkpoint and resuming training does not overwrite the learned bias.

3. Results

Evaluated the way the paper does: input Xtr, RMSE on the held-out ratings.

ML100k (600 hidden → 200, 150 epochs, l2=0.02, ~6 s on GPU)

Predictor Test RMSE
Global mean 1.128
Movie mean 1.025
RBM (this PR) 0.950
RBM (old code) 2.532

ML1m (400 hidden, 150 epochs, ~45 s)

Predictor Test RMSE
Global mean 1.117
Movie mean 0.979
RBM (this PR) 0.889

Ranking metrics, now with the correct input:

Dataset MAP nDCG@10 Precision@10 Recall@10
ML100k 0.0622 0.1241 0.1107 0.0398
ML1m 0.0915 0.1630 0.1421 0.0350

For reference, precision@10 on ML100k went from 0.014 to 0.111 against an honest input.

4. Notebooks and tests

  • rbm_deep_dive.ipynb is updated too: PyTorch instead of TensorFlow, the theory section (1.2 / 1.3) rewritten for one-hot softmax units with the corrected P(v_i^l = 1|h) softmax expression and the expected-rating inference formula, and a new "Rating prediction" section reporting RMSE against the two trivial baselines for both datasets.
  • The metrics cells now pass the arguments literally in each call instead of the shared args / kwargs variables.
  • Tests updated for the new parameter shapes, plus five new ones: one-hot encoding round-trip, that multinomial_sampling converges to a deliberately non-monotone target distribution, predict() determinism, remove_seen behaviour, and that resuming training preserves the visible bias. 11 tests pass; test_groups.yml updated with the new entries and timings.

@SimonYansenZhao SimonYansenZhao mentioned this pull request Aug 4, 2026
5 tasks
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants