CQ-FiLM

Category-Quantized Feature-wise Linear Modulation Regressor

A novel deep learning architecture for tabular data that treats categorical features in a discrete quantized space, inspired by Vector Quantization from VQ-VAE.

🎯 Core Philosophy

Categorical features should be processed in a discrete space, not as continuous embeddings.

Traditional tabular deep learning models treat categorical features as continuous embeddings, which may not capture their inherently discrete nature. CQFiLM addresses this by:

1. Vector Quantization: Inspired by VQ-VAE's discrete representation learning
2. Discrete Latent Space: Categorical features are mapped to a finite codebook
3. FiLM Conditioning: Quantized categorical information modulates numerical feature processing

---

🏗️ Architecture Overview

Input: Categorical Features + Numerical Features
  │
  ├─► Categorical Branch:
  │    ├─ Embedding Layer (per category)
  │    ├─ Shared Encoder
  │    └─ Category Quantizer (VQ-inspired) ──► z_q (discrete codes)
  │
  └─► Numerical Branch:
       ├─ Linear Layer
       ├─ FiLM Layer (modulated by z_q)
       └─ Linear Layer ──► Conditioned Numerical Features
  
  ┌───────────────────────────────┐
  │ Concatenate [z_q, num_latent] │
  └───────────────────────────────┘
              │
         Regression Head
              │
           Output

---

📐 Mathematical Framework

1. Category Quantization (VQ Mechanism)

Given categorical embeddings z_e, find the nearest codebook vector:

z_q = argmin_{e_k ∈ Codebook} ||z_e - e_k||²

Loss Components:

L_codebook = ||z_e - sg(z_q)||²        # Move codebook towards encoder
L_commit = β ||sg(z_e) - z_q||²        # Commitment loss
L_vq = L_codebook + L_commit

where sg() denotes stop-gradient operation.

2. Feature-wise Linear Modulation (FiLM)

Categorical information modulates numerical features:

FiLM(x, γ, β) = γ ⊙ x + β

where:
  γ = W_γ(z_q)  # Scale parameter
  β = W_β(z_q)  # Shift parameter

3. Final Prediction

h_num = FiLM(ReLU(W₁(x_num)), γ, β)
h_final = Concat([z_q, ReLU(W₂(h_num))])
ŷ = MLP_head(h_final)

---

🔑 Key Components

1. CategoryQuantizer

class CategoryQuantizer(nn.Module):
    """
    Vector Quantization layer inspired by VQ-VAE.
    Maps continuous embeddings to discrete codebook entries.
    """
    def __init__(self, num_embeddings, embedding_dim, commitment_cost):
        # num_embeddings: Size of the codebook (K)
        # embedding_dim: Dimension of each code vector
        # commitment_cost: Weight for commitment loss (β)

Key Features:

• Discrete latent space representation
• Straight-through estimator for backpropagation
• Learnable codebook with commitment loss

2. FiLMLayer

class FiLMLayer(nn.Module):
    """
    Feature-wise Linear Modulation.
    Applies affine transformation conditioned on categorical codes.
    """
    def forward(self, x, gamma, beta):
        return gamma * x + beta

Purpose: Allows categorical information to dynamically adjust numerical feature processing.

3. CQFiLM_RegressionModel

class CQFiLM_RegressionModel(nn.Module):
    """
    Main model architecture combining:
    - Categorical embedding + quantization
    - Numerical feature processing with FiLM
    - Regression head
    """

Architecture Details:

• Categorical Path: Embedding → Encoder → Quantizer → z_q
• Numerical Path: Linear → FiLM(conditioned by z_q) → Linear
• Output: Concat → BatchNorm → Dropout → Regression

---

💡 Why Vector Quantization for Categorical Features?

Traditional Approach (Continuous Embeddings)

Category "A" → [0.23, 0.45, 0.12, ...]  (any real values)
Category "B" → [0.67, 0.21, 0.89, ...]

❌ Categories can drift anywhere in continuous space

CQFiLM Approach (Discrete Quantization)

Category "A" → Codebook Entry #5 → [0.2, 0.5, 0.1, ...]
Category "B" → Codebook Entry #12 → [0.7, 0.2, 0.9, ...]

✅ Categories are forced to choose from a finite set of "prototypes"

Benefits:

1. Structured Representation: Categories share a common discrete codebook
2. Regularization: Prevents overfitting by limiting representation space
3. Interpretability: Each code in the codebook represents a categorical "prototype"
4. Generalization: Similar categories map to nearby codebook entries

---

🚀 Usage

Basic Example (Regression)

from CQFiLMRegressor import CQFiLM_Wrapper

# Define features
cat_features = ['gender', 'department', 'diagnosis']
num_features = ['age', 'blood_pressure', 'heart_rate', 'temperature']
target_col = 'length_of_stay'

# Initialize model
model = CQFiLM_Wrapper(
    cat_features=cat_features,
    num_features=num_features,
    target_col=target_col,
    num_cq_embeddings=64,      # Codebook size (K)
    embedding_dim=16,          # Code dimension
    commitment_cost=0.25,      # β in VQ loss
    cat_embedding_dim=8,       # Initial categorical embedding
    num_hidden_dim=32,         # Hidden dimension for numerical path
    learning_rate=0.001,
    batch_size=32,
    epochs=100,
    seed=42
)

# Train
model.fit(X_train, y_train)

# Predict
predictions = model.predict(X_test)

# Evaluate
r2 = model.score(X_test, y_test)

Scikit-learn Pipeline Integration

from sklearn.model_selection import GridSearchCV

param_grid = {
    'num_cq_embeddings': [32, 64, 128],
    'embedding_dim': [8, 16, 32],
    'commitment_cost': [0.1, 0.25, 0.5],
    'learning_rate': [0.001, 0.0001]
}

grid_search = GridSearchCV(
    CQFiLM_Wrapper(cat_features=cat_features, num_features=num_features),
    param_grid,
    cv=5,
    scoring='r2'
)

grid_search.fit(X_train, y_train)

---

🎨 Model Variants

Current: Regression

The current implementation (CQFiLM_RegressionModel) is designed for regression tasks with:

• L1 Loss (MAE)
• MinMaxScaler for target normalization
• Single continuous output

Future: Classification

CQFiLM can be easily adapted for classification tasks:

# Classification variant (conceptual)
class CQFiLM_ClassificationModel(nn.Module):
    def __init__(self, ..., num_classes):
        # Same architecture as regression
        # Replace regression head with:
        self.classification_head = nn.Sequential(
            nn.Linear(final_input_dim, 64),
            nn.ReLU(),
            nn.BatchNorm1d(64),
            nn.Dropout(0.1),
            nn.Linear(64, num_classes)  # Multi-class output
        )

Loss Function:

• Cross-Entropy Loss for classification
• Combined with VQ loss: L_total = L_ce + L_vq

---

⚙️ Hyperparameters

Vector Quantization Parameters

Parameter	Description	Default	Range
num_cq_embeddings	Codebook size (K)	64	16-256
embedding_dim	Dimension of each code	16	8-64
commitment_cost	β in VQ loss	0.25	0.1-1.0

Guidelines:

• Larger codebook: More expressive but may lead to codebook collapse
• Higher commitment cost: Forces encoder to commit to codes
• Embedding dim: Should match the complexity of categorical relationships

Architecture Parameters

Parameter	Description	Default
cat_embedding_dim	Initial categorical embedding	8
num_hidden_dim	Numerical path hidden dim	32
learning_rate	Adam learning rate	0.001
batch_size	Training batch size	32
epochs	Training epochs	100

---

🔬 Training Process

Loss Function

L_total = L_regression + L_vq
        = L_regression + (L_codebook + β * L_commit)

where:

• L_regression: L1 Loss (MAE) between prediction and target
• L_codebook: Moves codebook towards encoder output
• L_commit: Encourages encoder to commit to chosen codes

Training Loop

1. Forward Pass:

   • Encode categorical features → z_e
   • Quantize → z_q (with VQ loss)
   • Modulate numerical features via FiLM
   • Predict output

2. Loss Calculation:

   • Compute regression loss
   • Add VQ loss from quantizer
   • Total loss = regression + VQ

3. Backward Pass:

   • Straight-through estimator for quantization
   • Update encoder, quantizer, and predictor

---

📊 Data Preprocessing

Categorical Features

• Missing values: Handled with placeholder '__MISSING__'
• Encoding: Automatic factorization (0, 1, 2, ...)
• Unseen categories: Mapped to missing value code

Numerical Features

• Missing values: Imputed with column means
• Scaling: StandardScaler (z-score normalization)
• NaN handling: Automatic imputation pipeline

Target Variable

• Scaling: MinMaxScaler [0, 1]
• Inverse transform: Automatic in .predict()

---

⚠️ Known Issues & Future Work

1. Quantization Collapse

Problem: Some codebook entries may never be used, leading to an effectively smaller codebook.

Symptoms:

• Only a few codes are selected during training
• Reduced model expressiveness
• Suboptimal performance

Potential Solutions:

# 1. Exponential Moving Average (EMA) updates
# 2. Codebook reset: Replace unused codes
# 3. Entropy regularization: Encourage uniform code usage
# 4. Gumbel-Softmax relaxation
# 5. Commitment cost scheduling

Future Implementation:

class CategoryQuantizer(nn.Module):
    def __init__(self, ..., use_ema=True, decay=0.99):
        # EMA updates for codebook
        self.use_ema = use_ema
        self.decay = decay
        self.cluster_size = nn.Parameter(
            torch.zeros(num_embeddings), requires_grad=False
        )

2. Codebook Initialization

Current: Uniform initialization U(-1/K, 1/K)

Alternatives:

• K-means clustering on training data
• Random selection from encoder outputs
• Xavier/Kaiming initialization

3. Scalability

Limitation: Memory scales with O(K * d) for codebook

Solutions:

• Product quantization (split into subspaces)
• Hierarchical codebooks
• Adaptive codebook size

---

🎯 When to Use CQFiLM

✅ Ideal Scenarios

1. High-cardinality categorical features

   • Many categories benefit from shared codebook structure

2. Mixed data types

   • Both categorical and numerical features present
   • Complex interactions between feature types

3. Structured categorical relationships

   • Categories have underlying similarity structure
   • VQ can learn meaningful category prototypes

4. Small to medium datasets

   • Discrete representation provides regularization
   • Prevents overfitting compared to free embeddings

❌ Not Recommended

1. Pure numerical data

   • Use standard MLP/TabNet instead

2. Very few categorical features

   • Overhead of VQ not justified

3. Extremely large datasets

   • Training time may be prohibitive
   • Consider simpler models first

---

📈 Performance Tips

1. Codebook Size Selection

# Rule of thumb: K ≈ sqrt(num_categories) to 2*sqrt(num_categories)
num_categories = sum(n_unique for cat in categorical_features)
num_cq_embeddings = int(np.sqrt(num_categories) * 1.5)

2. Commitment Cost Tuning

# Start with β = 0.25
# Increase if codes aren't being used (collapse)
# Decrease if encoder output deviates too much from codes

3. Learning Rate Scheduling

# Cosine annealing or ReduceLROnPlateau
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
    optimizer, T_max=epochs
)

4. Early Stopping

# Monitor validation loss
# Stop if no improvement for N epochs

---

🔍 Model Interpretation

1. Codebook Analysis

# Extract learned codebook
codebook = model.model.quantizer.embedding.weight.data.cpu().numpy()

# Visualize with t-SNE/UMAP
from sklearn.manifold import TSNE
tsne = TSNE(n_components=2)
codebook_2d = tsne.fit_transform(codebook)
plt.scatter(codebook_2d[:, 0], codebook_2d[:, 1])

2. Code Usage Statistics

# Track which codes are used during inference
def get_code_usage(model, X):
    model.eval()
    with torch.no_grad():
        z_e = model.encode(cat_embeddings)
        _, _, codes = model.quantizer(z_e)
    return codes.cpu().numpy()

3. Feature Importance (FiLM Attention)

# Analyze γ and β values
# High magnitude indicates strong modulation
gamma_values = model.model.film_gamma_generator1.weight.data
beta_values = model.model.film_beta_generator1.weight.data

---

📚 References

Vector Quantization

• VQ-VAE: Neural Discrete Representation Learning (van den Oord et al., 2017)
• VQ-VAE-2: Generating Diverse High-Fidelity Images (Razavi et al., 2019)

Feature-wise Linear Modulation

• FiLM: Visual Reasoning with a General Conditioning Layer (Perez et al., 2017)

🛠️ Implementation Details

Dependencies

torch>=2.0.0
numpy>=1.20.0
pandas>=1.3.0
scikit-learn>=1.0.0

Device Support

# Automatic GPU/CPU selection
device = 'cuda' if torch.cuda.is_available() else 'cpu'

# Manual override
model = CQFiLM_Wrapper(..., device='cuda:0')

Reproducibility

# Fixed seed for all random operations
model = CQFiLM_Wrapper(..., seed=42)
# Sets: random, numpy, torch, torch.cuda