RankMixer (TensorFlow / Keras 3, CUDA + MUSA)

A faithful TensorFlow / Keras 3 implementation of

RankMixer: Scaling Up Ranking Models in Industrial Recommenders Zhu et al., ByteDance, 2025. arXiv:2507.15551

with optimizations and bug fixes called out throughout the source (see also docs/rankmixer.pdf, the original paper).

The model:

• Multi-head Token Mixing (parameter-free, paper Eq. 3-5)
• Per-token FFN (PFFN, paper Eq. 6-9), implemented as a single fused einsum over the token bank.
• Optional Sparse-MoE PFFN with three routing strategies:
  • relu -- single ReLU router with L1 sparsity penalty (paper Eq. 10-11)
  • relu_dtsi -- two-router DS-MoE / DTSI (paper §3.4)
  • topk_softmax -- vanilla MoE baseline for ablations
• LayerNorm + residual block, mean-pool head, BCE-from-logits task loss.

The training pipeline targets both NVIDIA CUDA (stock TF) and Moore Threads MUSA (TF + plugin), with deterministic alignment flags held identical so loss / AUC curves can be diffed run-to-run.

File layout

rankmixer/
├── train.py                   # main training entry point
├── train_cuda.sh              # paper-spec CUDA launcher
├── train_musa.sh              # paper-spec MUSA launcher (mirror)
├── model/
│   ├── rankmixer.py           # tokenization + token mixer + PFFN + MoE
│   ├── embedding.py           # per-feature sparse + dense embedding
│   ├── mlp.py                 # prediction head MLP
│   ├── lr_schedule.py         # linear warmup
│   └── sparse_optim.py        # SparseAwareAdagrad / SparseAwareRMSProp
│                              # (workaround for a Keras 3 sparse bug)
├── data/
│   ├── criteo_kaggle_dataset.py
│   └── amazon_beauty_dataset.py
├── docs/
│   └── rankmixer.pdf
├── requirements.txt
└── README.md

Setup

cd /root/code-space/rankmixer
python -m venv .venv && source .venv/bin/activate
pip install -r requirements.txt

(The HuggingFace datasets and xxhash packages are only needed for the Amazon Beauty Reviews loader; skip them if you only train on Criteo.)

Datasets

Two datasets ship with the loaders:

Name	Loader	Default path
Criteo Kaggle	data/criteo_kaggle_dataset.py	${RANKMIXER_DATASET_DIR}/criteo_data.npz
Amazon Beauty Reviews	data/amazon_beauty_dataset.py	${RANKMIXER_DATASET_DIR}/jhan21___amazon-beauty-reviews-dataset/...

RANKMIXER_DATASET_DIR defaults to /root/code-space/tensorflow_musa_playground/dataset (where the data originally lived). You can also point at a per-run path with --data_path /path/to/criteo_data.npz.

The Criteo loader expects the standard DLRM .npz layout (y, X_int (N,13), X_cat (N,26), optional counts (26,)).

Training

CUDA

./train_cuda.sh                                   # dense PFFN
./train_cuda.sh --use_moe --num_experts 4         # Sparse-MoE
./train_cuda.sh --vocab_cap 50000 \
                --max_steps_per_epoch 200 \
                --epochs 1                        # smoke test

MUSA

# Plugin .so default: ../tensorflow_musa_extension_2.15/build/libmusa_plugin.so
./train_musa.sh
./train_musa.sh --musa_lib_path /custom/path/libmusa_plugin.so
MUSA_LIB_PATH=/some/path ./train_musa.sh

Both scripts hold every numerics-affecting flag identical (--seed, --dropout, --batch_size, --peak_lr, --init_lr, optimizers, --deterministic), so the runs are designed to be diffed for backend alignment. See the comments inside the scripts for the CUDA-only determinism knobs (TF_CUDNN_DETERMINISTIC, CUBLAS_WORKSPACE_CONFIG, PYTHONHASHSEED).

Optimizers

The paper's §4.1.3 recipe is RMSProp for the dense parameters and Adagrad for the sparse embeddings. Both are the defaults here.

Keras 3 ships a sparse-tensor wrapper that mis-handles the IndexedSlices / sqrt(accumulator) step shared by Adagrad and RMSProp, so out of the box those two optimizers crash on the first train_step with::

indices[k] = N is not in [0, B)

model/sparse_optim.py works around the bug by going straight to tf.raw_ops.Resource{,Sparse}ApplyAdagradV2 / RMSProp. Any of {adam, sgd, adagrad, rmsprop} is a valid value for --sparse_optimizer / --dense_optimizer and is plumbed through _make_optimizer in train.py.

Notes on paper deviations

• DTSI router: implemented as gate = h_train + h_infer during training, gate = h_infer (with optional Top-K) at inference. L1 only on h_infer. The paper's wording ("both updated during training") leaves the exact training forward under-specified; the sum-routers form is the most common DS-MoE formulation and gives h_infer direct task gradient. The previous stop-grad distillation variant is still selectable via --moe_distill_coef > 0.

• Tokenization is a positional split of the concatenated feature embeddings (paper Eq. 2 with a single Proj per token). The paper envisions a semantic split via domain knowledge -- supply your features in the desired groupings upstream if you want this.

• Head: plain ReLU MLP, no BatchNorm by default. The paper doesn't specify a head; --head_use_bn is reserved for future work.

• No LR decay after the linear warmup. The paper doesn't specify decay either, but a cosine schedule is the obvious next step.

Status / known limitations

• tf.data.from_tensor_slices materializes the Criteo arrays in RAM; expect ~14 GB for the full 45M-row Kaggle dump. For larger datasets, switch to a sharded TFRecord pipeline.
• --moe_routing relu_dtsi does best with lots of training steps; on small data you may see the AUC curve wobble inside the validation noise band. See the inline notes in model/rankmixer.py.
• MoE from scratch is structurally less stable than dense PFFN -- a warm-start path (initialize each expert from a trained dense PFFN) is a planned addition.

License

TBD.