Prior-Data Fitted Networks as Tabular Foundation Models for Ranking in Low-Data Settings

Learning to rank (LTR) traditionally requires large-scale training data to generalize effectively. In low-data domains where expert annotation is scarce, the performance of LTR methods degrades sharply. Foundation models have alleviated similar data dependencies in other domains via in-context learning, but a foundation model for ranking with tabular features has not been explored yet.

We propose prior-data fitted networks (PFNs) as a strong method for ranking in low-data settings. First, we demonstrate that PFNs, which are originally trained for classification, successfully outperform classification baselines on ranking data. Next, we evaluate PFNs as rankers, showing that they surpass state-of-the-art tuned baselines in low-data regimes. We introduce a novel sampling and inference scheme to obtain pairwise predictions from PFNs' native pointwise architecture, analogous to pairwise LTR. To address the limited context window of the transformers underlying PFNs, we propose a dynamic support set selection strategy for queries that scales PFNs beyond random subsampling. Our experimental results show that PFNs are an effective foundation model for ranking that provides significant gains when data is limited.

• src/models/: Implementation of baseline models and PFN-based rankers.
• src/tuning/: Hyperparameter optimization logic using Optuna.
• src/data/: Dataset loaders for MSLR30k, Yahoo LTRC, and Istella.
• src/utils/: Unified utilities for ranking evaluation and data processing.
• src/conf/: Hydra configuration files for models and datasets.
• scripts/: Data preprocessing scripts for raw LETOR files.
• train.py: Main entry point for training and evaluation.
• tune.py: Main entry point for hyperparameter optimization.

This project uses uv for dependency management.

# Install uv if not installed
curl -LsSf https://astral.sh/uv/install.sh | sh

# Sync dependencies
uv sync

Alternatively, you can use pip:

pip install -e .

This project supports three major LETOR datasets: MSLR-WEB30K, Yahoo! LTRC, and Istella.

Raw LETOR files must be converted into tensor formats (.pt) for efficient loading. Use the provided scripts in the scripts/ directory:

# MSLR-WEB30K
python scripts/preprocess_mslr30k.py --input-dir data/MSLR-WEB30K

# Yahoo! LTRC
python scripts/preprocess_yahoo_ltrc.py --input-dir data/ltrc_yahoo

# Istella
python scripts/preprocess_istella.py --input-dir data/istella

The processed data will be saved in data/processed/, which is where the training scripts expect it.

Use train.py to run experiments. Configuration is managed via Hydra.

XGBoost Baseline:

uv run python train.py model=xgboost_ranking dataset=mslr30k fold=1

PFN Pointwise:

uv run python train.py model=pfn_pointwise dataset=mslr30k fold=1

PFN Pairwise:

uv run python train.py model=pfn_pairwise dataset=mslr30k fold=1

Use tune.py to optimize hyperparameters using Optuna.

uv run python tune.py model=xgboost_ranking dataset=istella fold=1 tuning.n_trials=50

submit_jobs.py generates and optionally submits SLURM job scripts.

# Print a job script without submitting
python submit_jobs.py --model pfn_pointwise --dataset mslr30k --size 500

# Submit a single job
python submit_jobs.py --model pfn_pointwise --dataset mslr30k --size 500 --submit

# Submit all sizes for a model/dataset combo
python submit_jobs.py --model pfn_pairwise --dataset yahoo_ltrc --size all_sizes --submit

# Submit across multiple datasets
python submit_jobs.py --model xgboost_ranking --dataset mslr30k yahoo_ltrc --size 1000 5000 --submit

This project uses Hydra for flexible configuration. You can override any setting from the command line.

Predefined presets for different training scales are available in src/conf/sizes/.

# PFN Pointwise on MSLR30k with 1,000 samples
uv run python train.py model=pfn_pointwise dataset=mslr30k sizes=pfn_pointwise/1000

# PFN Dynamic (KNN strategy) 
uv run python train.py model=pfn_dynamic model.dynamic.n_neighbors=16

# XGBoost on Istella with 10,000 samples
uv run python train.py model=xgboost_ranking dataset=istella sizes=xgboost/10000

# Change fold, seed, or verbosity
uv run python train.py fold=2 seed=42 verbose=false

# Override specific model hyperparameters
uv run python train.py model=xgboost_ranking model.booster.max_depth=6 model.booster.eta=0.1

# Adjust sampling parameters
uv run python train.py model.max_train_samples=500 model.pooling_k=5

To see the final configuration that will be used for a job without running it, append --cfg job:

uv run python train.py model=pfn_pointwise sizes=pfn_pointwise/1000 --cfg job