PyTorch Frame: A Modular Framework for Multi-Modal Tabular Learning

31 Mar 2024  ยท  Weihua Hu, Yiwen Yuan, Zecheng Zhang, Akihiro Nitta, Kaidi Cao, Vid Kocijan, Jure Leskovec, Matthias Fey ยท

We present PyTorch Frame, a PyTorch-based framework for deep learning over multi-modal tabular data. PyTorch Frame makes tabular deep learning easy by providing a PyTorch-based data structure to handle complex tabular data, introducing a model abstraction to enable modular implementation of tabular models, and allowing external foundation models to be incorporated to handle complex columns (e.g., LLMs for text columns). We demonstrate the usefulness of PyTorch Frame by implementing diverse tabular models in a modular way, successfully applying these models to complex multi-modal tabular data, and integrating our framework with PyTorch Geometric, a PyTorch library for Graph Neural Networks (GNNs), to perform end-to-end learning over relational databases.

PDF Abstract

Results from the Paper


Task Dataset Model Metric Name Metric Value Global Rank Result Benchmark
Toxic Comment Classification Civil Comments LightGBM + RoBERTa embedding AUROC 0.865 # 15
Toxic Comment Classification Civil Comments ResNet + RoBERTa embedding AUROC 0.882 # 14
Toxic Comment Classification Civil Comments Trompt + RoBERTa embedding AUROC 0.885 # 13
Toxic Comment Classification Civil Comments ResNet + RoBERTa finetune AUROC 0.97 # 7
Toxic Comment Classification Civil Comments Trompt + OpenAI embedding AUROC 0.947 # 11
Toxic Comment Classification Civil Comments ResNet + OpenAI embedding AUROC 0.945 # 12
Binary Classification fake ResNet + RoBERTa embedding AUROC 0.934 # 6
Binary Classification fake FTTransformer + OpenAI embedding AUROC 0.911 # 8
Binary Classification fake FTTransformer + RoBERTa fintune AUROC 0.96 # 3
Binary Classification fake LightGBM + OpenAI embedding AUROC 0.966 # 2
Binary Classification fake LightGBM + RoBERTa embedding AUROC 0.954 # 4
Binary Classification fake Trompt + OpenAI embedding AUROC 0.979 # 1
Binary Classification fake FTTransformer + RoBERTa embedding AUROC 0.936 # 5
Binary Classification fake ResNet + OpenAI embedding AUROC 0.923 # 7
Binary Classification kickstarter LightGBM + RoBERTa embedding AUROC 0.767 # 4
Binary Classification kickstarter ResNet + RoBERTa finetune AUROC 0.786 # 3
Binary Classification kickstarter Trompt + OpenAI embedding AUROC 0.81 # 1

Methods