TabTransformer for MovieLens Binary Classification

A PyTorch implementation of TabTransformer for binary classification on the MovieLens dataset. This project predicts whether a user will rate a movie highly (≥4.0 stars) based on user, movie, and contextual features using transformer architecture designed for tabular data.

🚀 Features

• TabTransformer Architecture: Attention-based model specifically designed for tabular data
• Multi-Modal Features: Handles categorical (users, movies, genres, tags) and continuous features
• Adaptive Thresholding: Automatically finds optimal classification thresholds
• Flexible Dataset Support: Works with different MovieLens dataset sizes
• GPU Acceleration: CUDA support with automatic fallback to CPU/MPS
• Comprehensive Evaluation: Multiple metrics with threshold optimization

📋 Requirements

• Python 3.11+
• CUDA-compatible GPU (optional, but recommended)
• uv package manager

🛠️ Setup

1. Create Virtual Environment

uv venv --python 3.11

2. Activate Virtual Environment

Windows (PowerShell):

.venv\Scripts\Activate.ps1

macOS/Linux:

source .venv/bin/activate

3. Install Dependencies

uv sync

This will install all required dependencies including PyTorch with CUDA support on Windows/Linux, or CPU/MPS support on other platforms.

🎯 Quick Start

Download Data and Train Model

# Train model without tags (recommended for faster training)
python main.py --mode train

# Train model with tag features (slower but potentially more accurate)
python main.py --mode train --use-tags

Evaluate Trained Model

# Basic evaluation with standard 0.5 threshold
python main.py --mode evaluate

# Evaluation with adaptive threshold optimization
python main.py --mode evaluate --adaptive-threshold

# Optimize threshold for specific metric
python main.py --mode evaluate --adaptive-threshold --threshold-metric f1

Full Pipeline Example

# 1. Train a model (downloads data automatically)
python main.py --mode train --use-tags

# 2. Evaluate with adaptive thresholding
python main.py --mode evaluate --use-tags --adaptive-threshold --threshold-metric accuracy

📊 Usage Examples

Training Options

# Different dataset sizes
python main.py --mode train --dataset ml-latest-small
python main.py --mode train --dataset ml-latest

# Custom rating threshold
python main.py --mode train --rating-threshold 3.5

# Enable tag features
python main.py --mode train --use-tags

Evaluation Options

# Evaluate specific model
python main.py --mode evaluate --model-path path/to/model.pth

# Different threshold optimization metrics
python main.py --mode evaluate --adaptive-threshold --threshold-metric precision
python main.py --mode evaluate --adaptive-threshold --threshold-metric recall

🏗️ Project Structure

tab-transformer/
├── src/                          # Core implementation
│   ├── __init__.py
│   ├── config.py                 # Configuration settings
│   ├── dataset.py                # PyTorch Dataset implementation
│   ├── download_and_preprocess.py # Data downloading and preprocessing
│   ├── tabtransformer.py         # TabTransformer model architecture
│   ├── train_utils.py            # Training utilities and metrics
│   └── trainer.py                # Training loop implementation
├── data/                         # Dataset storage (created automatically)
├── artifacts/                    # Training artifacts and logs
├── notebook/                     # Jupyter notebooks for experimentation
│   └── bertclassifier.ipynb
├── main.py                       # Main pipeline script
├── pyproject.toml               # Project dependencies and configuration
├── python-version.txt           # Python version specification
└── uv.lock                      # Dependency lock file

Key Components

• main.py: Complete pipeline orchestration with CLI interface
• src/tabtransformer.py: Core transformer architecture for tabular data
• src/dataset.py: Custom PyTorch dataset with feature engineering
• src/trainer.py: Training loop with validation and checkpointing
• src/config.py: Centralized configuration management
• src/train_utils.py: Metrics computation and adaptive thresholding

🧠 Model Architecture

The TabTransformer combines the power of transformer attention with tabular data processing:

graph TD
    A[User ID] --> E[User Embedding]
    B[Movie ID] --> F[Movie Embedding]
    C[Genre IDs] --> G[Genre Embeddings]
    D[Tag IDs] --> H[Tag Embeddings]
    I[Continuous Features] --> J[Linear Projection]
    
    E --> K[Column-wise Contextual Embeddings]
    F --> K
    G --> K
    H --> K
    J --> K
    
    K --> L[Multi-Head Self-Attention]
    L --> M[Feed Forward Network]
    M --> N[LayerNorm + Residual]
    N --> O{More Layers?}
    O -->|Yes| L
    O -->|No| P[Global Pooling]
    
    P --> Q[Classification Head]
    Q --> R[Binary Prediction]
    
    style E fill:#e1f5fe
    style F fill:#e1f5fe
    style G fill:#e1f5fe
    style H fill:#e1f5fe
    style J fill:#fff3e0
    style L fill:#f3e5f5
    style Q fill:#e8f5e8

Loading

Architecture Details

• Categorical Embeddings: Users, movies, genres, and tags are embedded into dense vectors
• Column-aware Processing: Each feature type gets distinct column ID embeddings
• Multi-Head Attention: Captures complex feature interactions
• Continuous Feature Integration: Numerical features are projected and combined
• Adaptive Classification: Optional threshold optimization for deployment

⚙️ Configuration

Key configuration options in src/config.py:

MODEL_CONFIG = {
    'd_model': 128,           # Model dimension
    'n_heads': 8,             # Attention heads
    'n_layers': 3,            # Transformer layers
    'd_ff': 256,              # Feed-forward dimension
    'dropout': 0.1            # Dropout rate
}

TRAINING_CONFIG = {
    'num_epochs': 50,         # Training epochs
    'batch_size': 512,        # Batch size
    'learning_rate': 1e-3,    # Learning rate
    'weight_decay': 1e-4      # L2 regularization
}

📈 Performance Features

Adaptive Thresholding

The model includes automatic threshold optimization:

• Metrics: Optimize for accuracy, F1-score, precision, or recall
• Validation-based: Uses validation set to find optimal threshold
• Multiple Thresholds: Tests 100 different threshold values
• Comprehensive Reporting: Shows improvement over standard 0.5 threshold

Efficient Training

• GPU Acceleration: Automatic CUDA detection and usage
• Data Loading: Multi-worker data loading with memory pinning
• Mixed Precision: Optional for faster training on modern GPUs
• Checkpointing: Automatic model saving and configuration persistence

🔍 Model Evaluation

The evaluation pipeline provides comprehensive metrics:

EVALUATION RESULTS
==================================================
Test Samples: 10,000

🎯 ADAPTIVE THRESHOLD RESULTS (threshold = 0.6234):
Accuracy:     0.8756
Precision:    0.8934
Recall:       0.8123
F1 Score:     0.8512
AUC-ROC:      0.9234
Test Loss:    0.3421

📊 STANDARD THRESHOLD RESULTS (threshold = 0.5):
Accuracy:     0.8543
Precision:    0.8721
Recall:       0.7891
F1 Score:     0.8289

📈 IMPROVEMENT:
Accuracy improvement: +0.0213
F1 improvement:       +0.0223

🎛️ Advanced Usage

Custom Dataset Configuration

# Use larger dataset
python main.py --mode train --dataset ml-latest

# Custom rating threshold
python main.py --mode train --rating-threshold 3.5

Model Analysis

# Enable tag features for richer representations
python main.py --mode train --use-tags

# Evaluate with different optimization targets
python main.py --mode evaluate --adaptive-threshold --threshold-metric precision

🐛 Troubleshooting

Common Issues

1. CUDA Out of Memory: Reduce batch size in src/config.py
2. Download Failures: Check internet connection and retry
3. Missing Data: Run with --mode train first to download datasets
4. Model Not Found: Ensure training completed successfully

Performance Tips

• Use --use-tags only if you need maximum accuracy (slower training)
• Enable GPU acceleration for faster training
• Use adaptive thresholding for production deployments
• Monitor training logs in the artifacts directory

📝 License

This project is available under the MIT License.

🤝 Contributing

1. Fork the repository
2. Create a feature branch
3. Make your changes
4. Add tests if applicable
5. Submit a pull request

---

Happy modeling! 🚀