UniDot

A Unified Network for Sequence Modeling and Feature Interaction in Large-scale Recommendation

UniDot is a post-click conversion-rate (pCVR) model that unifies the two model families industrial recommenders usually keep separate — feature-interaction models over multi-field user/item features, and sequential models over user-behavior histories — into a single stackable block. It was the runner-up on the Industrial track of the TAAC × KDD Cup 2026 (Tencent Uni-Rec Challenge), scoring 0.83217 test AUC — 0.037% behind the winning entry.

The design starts from the factorization-machine point of view: the embedding inner product that powers collaborative filtering is the same primitive as attention's query·key scoring, so one dot-product of tokens can underlie both feature interaction and sequence modeling. UniDot follows the SlimPer framework, which frames ranking as iterative refinement of a personalization state.

The paper, slides, and poster are in paper/.

Non-sequential fields and multi-domain behavioral sequences are tokenized into one shared token space. The model stacks a single macro-block (UniDotBlock) in which two buses run in parallel and exchange state every layer:

• Token-mixing bus — a swappable cross-token slot (Wukong: parallel LCB + FMB by default) over concatenated user / item / pooled-sequence tokens.
• Sequence-retrieval bus — item tokens cross-attend the behavioral histories (per-domain + a merged cross-domain stream).
• MLP-Mixer fusion (FuseFFN) — exchanges residual deltas between the two buses each layer.
• FM Highway — routes explicit per-layer dot-product interactions (per-sequence dots, aggregated grams, cross-bus user×item dots) around the residual stack straight to the classifier, keeping second-order signal alive at negligible cost.

The sequence side is embedded once per forward pass and shared by all consumers, which bounds inference latency. Training uses a dual sparse/dense optimizer (Adagrad for embedding tables, Muon for matrix weights with an AdamW auxiliary group for 1-D params), an auxiliary conversion-delay head, and multi-path mutual learning (DML) — used by the submitted configuration.

See .claude/CLAUDE.md for a full architectural walkthrough and paper/ for the formal description.

Official metric is ROC AUC on a 12M-example held-out test set.

★ Serving a single path of the multi-path model — at 1× inference cost — still scores 0.83184, which would itself have placed second. Multi-path training pulls even a single served model to a better minimum, so DML is a training boost you don't have to pay for at serving time (config.DML_INFER_MODE='single').

From the competition baseline to the final submission, +1.818 AUC points:

The architecture itself is the single largest jump (+1.10), from the released schema alone — no hand-crafted cross features.

Component ablations (tiny mode, single path, d_model=64, in-distribution held-out — so absolute numbers differ from the leaderboard):

The FM Highway is the costliest component to remove (−0.127%), which is the paper's central claim: explicit second-order signal routed around the residual stack is worth keeping.

Honest caveat: dropping FuseFFN slightly raises AUC (+0.022%) while worsening LogLoss, so that module's contribution is not conclusive. An input-conditioned (second-order) fuser is the clearest open direction — hence the --fuse_type {mixer,cross_attn,bilinear,asym} variants in the code.

Scaling: held-out AUC grows log-linearly at ≈+0.0025 AUC per doubling of training data (4M → 32M) and had not saturated. Dense capacity was the productive axis; doubling embedding width alone gave no A/B win.

bf16 autocast costs <0.0001 AUC. The submitted configuration is d_model=128, 6 macro-layers × 2 mix-blocks, sequence windows a:256 b:256 c:512 d:512 plus a merged abcd:512 stream. Full hyperparameters are in the paper's appendix.

Requires Python 3.10+ and a CUDA-capable GPU.

pip install torch pyarrow numpy scikit-learn tqdm tensorboard python-dotenv

Attention uses PyTorch's built-in scaled-dot-product attention (SDPA); no FlashAttention install is required.

Optional:

• muon-optimize — needed only for --dense_optimizer muon (the default in run.sh; use --dense_optimizer adamw otherwise).

The model streams Parquet files containing per-event user/item features and behavioral sequences (see schema.json). The task is binary: predict conversion (label_type == 2), with roughly 12% positives.

The competition data is 35M training / 12M test events over 142 columns: 54 user-int and 17 user-dense fids, 17 item-int and 4 item-dense fids, 11 pre-trained embedding fids, and 4 behavioral domains (9 / 14 / 12 / 10 fields each). The full dataset is not redistributed here, but a 1,000-row sample with the same schema is available for smoke tests: TAAC2026/second_round_sample_1000 (CC-BY-NC-4.0 — note this is more restrictive than this repo's MIT code license).

Paths are supplied via environment variables (loaded from .env in local mode). Copy the template and fill in your paths:

cp .env.example .env
# then edit .env:
#   TRAIN_DATA_PATH       a parquet file or a directory of shards
#   TRAIN_CKPT_PATH       where checkpoints are written
#   TRAIN_LOG_PATH        text logs
#   TRAIN_TF_EVENTS_PATH  TensorBoard events
#   SCHEMA_SRC_PARQUET    a parquet whose schema matches the training data

config.MODE == 'local' (the default) auto-loads .env and pins one GPU. The launcher picks torchrun DDP when more than one GPU is visible, so pin a single device for a smoke test. Point --train_data_path at any small parquet with the expected schema (a single-row-group file needs --split_by_time so the row-level train/val split has rows on both sides).

Grab the 1,000-row sample:

huggingface-cli download TAAC2026/second_round_sample_1000 --repo-type dataset --local-dir ./sample_data

CUDA_VISIBLE_DEVICES=0 bash run.sh \
    --train_data_path ./sample_data \
    --split_by_time \
    --num_epochs 2 \
    --reinit_cardinality_threshold 0 \
    --skip_num_slots 0

run.sh holds the reference configuration (6 macro-layers, 4 behavioral domains plus a merged stream, per-position fid compression, DIN conditioning, Muon dense optimizer). Override any flag on the command line:

bash run.sh                                     # reference config
bash run.sh --num_epochs 20 --dropout_rate 0.1  # override flags

bash run_dml.sh

run_dml.sh mirrors run.sh's training recipe and trains N mutually-distilling peers; at inference config.DML_INFER_MODE selects 'mean' (average all peers, best accuracy) or 'single' (serve one peer at 1× cost).

infer.py reads paths from environment variables, rebuilds the model from the checkpoint's train_config.json, and writes predictions:

MODEL_OUTPUT_PATH=/path/to/checkpoint \
EVAL_DATA_PATH=/path/to/test.parquet \
EVAL_RESULT_PATH=/path/to/predictions \
python3 infer.py

Inference-time behavior (batch size, torch.compile, autocast, matmul precision) is configured in config.py.

config.MODE selects the run profile:

• local — auto-pins one GPU and loads .env; for demo / single-box development.
• tiny — real-data path, single GPU, capped sample count for rapid A/B.
• remote — full multi-GPU DDP run.

run.sh reads MODE and appends the appropriate flags (sample caps, learning rates, DDP launcher).

If you use this code, please cite the paper:

@inproceedings{unidot2026,
  title     = {UniDot: A Unified Network for Sequence Modeling and
               Feature Interaction in Large-scale Recommendation},
  author    = {Lin, Rongcheng and Sun, Yan and Zhang, Jamey and
               Xiong, Guanglei and Ji, Ivan and Chen, Xianjie and Bu, Shujian},
  booktitle = {KDD Cup 2026 Tencent UniRec Challenge Workshop},
  year      = {2026}
}

Released under the MIT License — free to use, modify, and distribute, including for commercial purposes.