On the Mechanisms of Collaborative Learning in VAE Recommenders

This repository contains the official implementation of the ICLR 2026 paper On the Mechanisms of Collaborative Learning in VAE Recommenders.

VAE-based collaborative filtering mostly exploits local collaboration among input-similar users while under-using global collaboration between distant-but-related users. This work analyzes the mechanisms behind this (latent proximity, β-KL regularization, input masking) and proposes Personalized Item Alignment (PIA) — an anchor regularizer that aligns user posteriors with item embeddings, enabling global consistency while preserving user identity. The method is validated on MovieLens-20M, Netflix, and Million Song Data, and was deployed on an Amazon streaming platform following a successful online experiment.

How it works

Binary input masking creates stochastic overlaps between users' interaction vectors, letting distant-but-related users collaborate — e.g. masking can make two users with disjoint histories look alike, so a valid recommendation can propagate between them:

Collaboration is governed by latent proximity: the expected Wasserstein distance between two users' masked posteriors is bounded by the KL term, so gradient sharing fades as users move apart in latent space:

PIA anchors each user's posterior to the embeddings of their interacted items, keeping user identity intact (structured by activity level, below) while enabling global consistency:

Quick start on Colab

The fastest way to try PIA is the demo notebook demo.ipynb — it downloads MovieLens-20M, trains Multi-VAE + PIA for a few epochs on a free Colab GPU, evaluates Recall/NDCG, and shows live recommendations for a held-out user. Click the Colab badge above to launch it.

Installation

pip install -r requirements.txt

The code was tested with Python 3.8, PyTorch 2.4, and CUDA 12.1.

Repository structure

File	Purpose
demo.ipynb	End-to-end Colab demo (download → train → evaluate → recommend)
train.py	Train a model
evaluation.py	Evaluate a trained checkpoint (Recall@K, NDCG@K)
plot_tsne.py	Visualize user representations with t-SNE
model.py	VAE architectures (shallow, RecVAE, hierarchical gated)
data.py	Data loading
utils/	Dataset preprocessing, losses, layers, optimizer utilities

Datasets and preprocessing

• MovieLens 20M
• Netflix Prize
• Million Song Data

After downloading the datasets:

1. Place them under a root directory of your choice, e.g. /your_root_directory/data.
2. Rename the unzipped folders as needed — e.g. the code expects the MovieLens data in /your_root_directory/data/ml-20m.

Then generate the train/validation/test splits:

# MovieLens 20M
python -m utils.movieLens -root /your_root_directory

# Netflix
python -m utils.netflix -root /your_root_directory

# Million Song Data
python -m utils.MSD -root /your_root_directory

This produces the following files per dataset:

• train.csv — training data.
• validation_tr.csv / validation_te.csv — model input and held-out targets for the validation stage.
• test_tr.csv / test_te.csv — model input and held-out targets for the testing stage.
• unique_sid.txt — item-ID mapping used when reading train.csv (optional).

Training

python -m train -root /your_root_directory -dataset ml-20m -lambda_alignment 8 -lambda_scale 1.2

The checkpoint with the best validation score is saved to /your_root_directory/checkpoints/<dataset>/ with an auto-generated name encoding the configuration (printed at the start of training).

Key options

Argument	Choices / default	Description
-dataset	ml-20m, netflix, MSD	Dataset to train on
-model_type	shallow (default), rec, hierachical_gated	Encoder/decoder architecture
-prior_type	normal (default), composition, vamp	Latent prior
-posterior_type	vae (default), iaf	Posterior family
-lambda_alignment	default 8	PIA alignment strength; set 0 to disable PIA
-lambda_scale	default 1.2	PIA scale weight
-rho	default 5	PIA temperature
-alternative_training	flag	Alternate encoder/decoder updates (RecVAE-style)
-n_epochs / -lr / -batch_size	400 / 5e-4 / 500	Optimization settings

Examples

Multi-VAE + PIA

python -m train -dataset ml-20m -lambda_alignment 8 -lambda_scale 1.2
python -m train -dataset MSD -lambda_alignment 8 -lambda_scale 1.2
python -m train -dataset netflix -lambda_alignment 8 -lambda_scale 1.2

RecVAE + PIA, full settings — alternating training, composition prior, and adaptive beta (see the RecVAE repository):

python -m train -dataset ml-20m -lambda_alignment 8 -lambda_scale 1.2 -model_type rec -prior_type composition -alternative_training

RecVAE architecture only — reduced training time:

python -m train -dataset MSD -lambda_alignment 8 -lambda_scale 1.2 -model_type rec
python -m train -dataset netflix -lambda_alignment 8 -lambda_scale 1.2 -model_type rec

H+Vamp + PIA (see the EVCF repository):

python -m train -dataset ml-20m -prior_type vamp -model_type hierachical_gated -lambda_alignment 2 -lambda_scale 1.0

Evaluation

To evaluate a trained checkpoint, pass the model file name (without the .pt extension) along with the same architecture options used for training:

python -m evaluation -root /your_root_directory -dataset ml-20m -model_type shallow -model <model_name>

Visualization

To visualize user representations with t-SNE:

python -m plot_tsne -root /your_root_directory -dataset ml-20m -model <model_name>

Citation

If you find this code useful, please cite:

@inproceedings{vuong2026mechanisms,
  title     = {On the Mechanisms of Collaborative Learning in {VAE} Recommenders},
  author    = {Vuong, Tung-Long and Monteil, Julien and Dang, Hien and Vaskovych, Volodymyr and Le, Trung and Nguyen, Vu},
  booktitle = {International Conference on Learning Representations (ICLR)},
  year      = {2026}
}

License

This project is licensed under the CC-BY-NC-4.0 License — see LICENSE.txt for details.