AUC-ROC and AUC-PR Analysis for Machine Learning Systems

AUC-ROC and AUC-PR metrics are essential for evaluating binary classification models, with AUC-ROC measuring overall discriminative ability and AUC-PR focusing on precision-recall trade-offs in imbalanced datasets.

1. Problem Statement & ML Task Definition

Clarifying questions

1. What is the expected scale of training data and model inference?

   • Are we dealing with millions of users/items with real-time serving requirements or batch processing at scale?

2. What is the primary business use case?

   • Is this for content recommendation, fraud detection, ad targeting, or another specific application?

3. What is the latency budget for real-time inference?

   • Is sub-100ms required, or can we tolerate higher latencies for better accuracy?

4. Is this a batch or online serving system?

   • Will predictions be made in real-time for individual user requests or in batches?

5. What is the expected label distribution?

   • Is there significant class imbalance, and how should we handle it?

6. What is the primary business metric for success?

   • Is it revenue, engagement, or another KPI that drives business value?

7. What is the frequency of model retraining and data freshness requirements?

   • Is daily retraining acceptable or do we need hourly updates?

Business use case and success criteria

The system is designed to predict user engagement with content items (e.g., click probability) in a real-time recommendation engine. The primary goal is to maximize user engagement and revenue by delivering relevant content.

Success criteria:

• Offline: AUC-ROC > 0.90, AUC-PR > 0.75 on holdout set
• Online: CTR improvement of 15% over baseline, 20% increase in session duration
• Business impact: Revenue per user increase by 10%

ML problem formulation

Task type: Binary classification (click vs no-click)

Justification:

• The problem is framed as a binary decision: will the user engage with the content?
• User-item interaction data (clicks, views) provides clear positive/negative labels
• The system needs to rank items for each user in real-time, making it suitable for ranking-based approaches
• Binary classification is the most direct way to model engagement probability

ML evaluation metrics

Offline metrics:

• AUC-ROC: Measures overall discriminative power of the model across all thresholds
• AUC-PR: More informative for imbalanced datasets, measures precision-recall trade-off
• Precision@K: Top-K precision for recommendation lists (e.g., 5, 10, 20)
• NDCG@K: Normalized Discounted Cumulative Gain for ranked list evaluation

Online metrics:

• CTR (Click-Through Rate): Direct business KPI for engagement
• Engagement Duration: Average time spent on content after recommendation
• Revenue per User: Monetization impact of recommendations

Scale and constraints

Back-of-envelope estimates:

• Training data size: 50M user-item interactions (100GB raw data)
• Number of features: 50K sparse features + 500 dense embeddings (1M total features)
• Embedding matrix memory footprint: 1M × 128 dimensions × 4 bytes = ~500MB
• Serving QPS math: 100ms latency budget → 10 QPS per model instance, 500 instances for 5K QPS

Constraints:

• QPS: 5,000 requests per second (peak)
• Latency budget: <100ms end-to-end (50ms for model inference, 20ms for feature lookup, 30ms for scoring)
• Data freshness: 1-hour update window for features, daily retraining of model
• Model size: <50MB for efficient serving

Non-ML baseline strategy

Rule-based fallback:

• For new users or items with insufficient data, fall back to popularity-based ranking (most viewed content)
• Use simple heuristics like recency, category popularity, or user demographics

Cold start handling:

• For new users: Use demographic features and content similarity to bootstrap recommendations
• For new items: Use item content embeddings and collaborative filtering with similar items

Fallback logic:

def get_baseline_recommendations(user_id, item_id):
    if user_id in new_users or item_id in new_items:
        return get_popular_items()
    else:
        return get_user_based_recommendations(user_id)

2. Data Strategy

Training data sources

• Explicit feedback: User ratings, clicks, purchases, likes, and explicit relevance labels from user interactions with items (e.g., "This article was helpful" or "I don't want to see this again")
• Implicit signals: Page views, time spent on page, scroll depth, bounce rate, search queries, and item recommendation clicks to infer user interest or disinterest
• Third-party data: Demographic information, content metadata from publishers, topic categories, and contextual signals like time of day or device type

Labeling pipeline design

• Human annotation: For high-value, sensitive use cases (e.g., content moderation) where precision is critical. Annotation pipeline includes quality control, inter-annotator agreement checks, and labeling guidelines.
• Weak supervision: For large-scale data labeling using heuristics or pre-trained models. For example, use keyword matching for relevance or content classifiers to label user-item interactions.
• Programmatic labeling: Automatically label based on business logic or event tracking (e.g., if a user clicks an item, it's labeled as positive). This is used for rapid data generation and can be refined with human review.

Class imbalance handling strategy

• Oversampling: Apply SMOTE or random oversampling to minority classes (e.g., rare positive labels) to balance class distribution.
• Weighted loss: Incorporate class weights in the model's loss function (e.g., weighted cross-entropy) to penalize misclassification of rare classes more heavily.
• Threshold tuning: Adjust decision thresholds for inference to optimize precision/recall trade-offs rather than using a default 0.5 threshold.
• Justification: In recommendation systems, most items are not clicked or relevant, leading to highly imbalanced datasets. Without handling imbalance, models may become biased toward majority class and underperform on rare but important items.

Negative sampling approach

• Unigram distribution: Sample negative items based on item popularity (more popular items are sampled more frequently). This mimics real-world item distribution and improves model generalization.
• In-batch negative sampling: For training efficiency, sample negatives from within the same batch to reduce memory overhead and increase training speed.
• Hard negative mining: Periodically mine hard negatives (items that are likely to be relevant but were not selected) to improve model robustness.

Data freshness and update frequency requirements

• Training data: Updated daily to reflect recent user behavior and item changes. Weekly for content metadata or long-term trends.
• Real-time signals: Updated every few minutes to ensure model responsiveness to recent user activity.
• Back-of-envelope estimate:
  • Training data size: ~50M user-item interactions (50GB raw, ~200GB with features)
  • Feature dimensions: ~10K sparse + 500 dense (e.g., embeddings)
  • Model size: ~1GB for a DNN model with 500-dim embeddings and 2 hidden layers

Privacy compliance (GDPR, CCPA)

• Anonymization: Remove or hash PII fields like user IDs, email addresses, and device identifiers before training.
• Data retention: Retain only the minimum required data for model training and remove after 90 days or per legal requirement.
• Consent tracking: Ensure that user consent is properly logged and can be revoked.

Data quality validation and monitoring

• Validation checks:
  • Missing or malformed features (e.g., NaNs in embeddings)
  • Inconsistent timestamps
  • Duplicate entries or user-item pairs
• Monitoring:
  • Feature drift detection using statistical tests (e.g., Kolmogorov-Smirnov)
  • Label drift detection to identify shifts in user preferences or content changes
  • Data freshness monitoring to ensure training data is up-to-date

import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.utils import resample

# Example: Handling class imbalance in training data
def balance_dataset(df, target_col='is_relevant'):
    # Separate majority and minority classes
    df_majority = df[df[target_col] == 0]
    df_minority = df[df[target_col] == 1]

    # Oversample minority class
    df_minority_upsampled = resample(df_minority,
                                     replace=True,
                                     n_samples=len(df_majority),
                                     random_state=42)

    # Combine datasets
    df_balanced = pd.concat([df_majority, df_minority_upsampled])
    return df_balanced

# Example: Data validation function
def validate_training_data(df):
    # Check for missing values in key features
    assert not df['user_id'].isnull().any(), "Missing user IDs"
    assert not df['item_id'].isnull().any(), "Missing item IDs"
    assert not df['timestamp'].isnull().any(), "Missing timestamps"

    # Check for duplicates
    assert df.duplicated(subset=['user_id', 'item_id']).sum() == 0, "Duplicate user-item pairs"

    return True

3. Feature Engineering

Feature Taxonomy

Features are categorized into four main types:

1. User Features: User demographics, behavioral history, engagement patterns, and preferences
2. Item Features: Content metadata, category, popularity metrics, and quality indicators
3. Context Features: Temporal information, device type, geographic location, and session context
4. Interaction Features: User-item interaction signals, temporal decay factors, and sequence patterns

Embedding Design

Entity Embeddings:

• User ID embeddings: 128-dim, trained via click-through rate prediction
• Item ID embeddings: 256-dim, learned from item popularity and engagement signals

Text Embeddings:

• Item title/content embeddings: 384-dim using BERT-base, fine-tuned on item descriptions
• User preference embeddings: 128-dim from user history sequences

Graph Embeddings:

• User-item interaction graph: 64-dim Node2Vec embeddings for capturing user preferences and item similarities

Real-time vs Batch Feature Computation

Feature Type	Computation Mode	Reasoning
User session features	Real-time	Immediate feedback on user behavior
Item popularity metrics	Batch	Updated daily, no real-time need
User preference embeddings	Real-time	Dynamic user interests require frequent updates
Temporal context features	Real-time	Time-sensitive signals like hour-of-day

Feature Store Design

Online Store:

• Low-latency key-value store (e.g., Redis or DynamoDB)
• Optimized for sub-10ms lookup times
• TTL of 5 minutes for real-time features

Offline Store:

• Data warehouse (e.g., Snowflake or BigQuery) for training and batch processing
• Partitioned by date for efficient querying

Feature Freshness and TTL

Feature	Source	Update Frequency	Store Type
User session features	Real-time stream	Every 5 seconds	Online
Item popularity	Batch job	Daily	Offline
User preference embeddings	Real-time stream	Every 10 minutes	Online
Temporal context	System clock	Every minute	Online
User demographics	User profile API	On update	Online

Cross Features and Feature Interactions

Cross features are created for high-value combinations:

• User-item interaction time windows (e.g., last 1h, last 24h)
• User-category engagement patterns (e.g., user's preference for tech vs fashion)
• Temporal-spatial features (e.g., time-of-day × location)

Cross-feature engineering uses:

• Feature crosses: user_age × item_category
• Polynomial features: user_engagement_score²
• Bucketed features: user_age_group × item_price_bucket

Feature Importance and Selection

Methodology:

• SHAP values for global feature importance analysis
• Permutation importance from trained model (XGBoost or LightGBM baseline)
• Recursive feature elimination with cross-validation

Selection Strategy:

• Top 500 features by SHAP importance retained
• Features with SHAP mean absolute value < 0.001 excluded
• Regular re-evaluation every quarter to adapt to changing user behavior

Core Feature Engineering Pipeline Code

import numpy as np
from sklearn.preprocessing import StandardScaler
from sklearn.feature_extraction.text import TfidfVectorizer

class FeatureExtractor:
    def __init__(self):
        self.tfidf_vectorizer = TfidfVectorizer(max_features=1000, stop_words='english')
        self.scaler = StandardScaler()
        
    def extract_user_features(self, user_data):
        """Extract user behavioral features"""
        features = {
            'user_engagement_score': np.log1p(user_data['total_clicks']),
            'user_session_count': user_data['session_count'],
            'user_avg_session_duration': user_data['avg_session_duration'],
            'user_category_preference': self._calculate_category_preference(user_data),
        }
        return features
    
    def extract_item_features(self, item_data):
        """Extract item content and metadata features"""
        features = {
            'item_popularity': np.log1p(item_data['click_count']),
            'item_category_encoded': item_data['category_id'],
            'item_title_embedding': self._get_text_embedding(item_data['title']),
        }
        return features
    
    def extract_interaction_features(self, interaction_data):
        """Extract temporal and behavioral interaction signals"""
        features = {
            'time_since_last_interaction': self._calculate_time_diff(interaction_data),
            'user_item_engagement_score': interaction_data['engagement_score'],
            'temporal_decay_factor': self._calculate_temporal_decay(interaction_data),
        }
        return features
    
    def _calculate_category_preference(self, user_data):
        """Calculate user preference for each category"""
        return np.array([
            user_data.get(f'category_{i}_engagement', 0) 
            for i in range(10)
        ])
    
    def _get_text_embedding(self, text):
        """Generate TF-IDF embedding for item title"""
        return self.tfidf_vectorizer.fit_transform([text]).toarray()[0]
    
    def _calculate_time_diff(self, interaction_data):
        """Calculate time difference in seconds"""
        return (pd.Timestamp.now() - interaction_data['timestamp']).total_seconds()
    
    def _calculate_temporal_decay(self, interaction_data):
        """Calculate exponential decay factor"""
        time_diff = self._calculate_time_diff(interaction_data)
        return np.exp(-0.001 * time_diff)  # Decay rate of 0.001 per second

# Example usage
extractor = FeatureExtractor()
user_features = extractor.extract_user_features(user_data)
item_features = extractor.extract_item_features(item_data)
interaction_features = extractor.extract_interaction_features(interaction_data)

4. Model Architecture

Model Selection with Justification

For a recommendation system, I'll select a Two-Tower Deep Neural Network architecture with the following rationale:

1. Scalability: Two-tower models can scale to large embedding spaces efficiently
2. Expressiveness: DNNs capture complex user-item interactions better than linear models
3. Flexibility: Can incorporate multiple modalities (text, images, user behavior)
4. Online Serving: Well-suited for real-time scoring with low latency requirements

Comparison Matrix:

Model	Pros	Cons	Fit for Recommendation
Logistic Regression	Fast, interpretable	Limited expressiveness	❌
GBM (XGBoost/LightGBM)	Good baseline, handles mixed data	Slower inference, less scalable	⚠️
DNN (Two-Tower)	High expressiveness, scalable	More complex to train	✅
Transformer	Captures sequential patterns	High latency, memory intensive	⚠️

Architecture Diagram

graph TD
    A[User Features] --> B[User Tower]
    C[Item Features] --> D[Item Tower]
    B --> E[Merged Embeddings]
    D --> E
    E --> F[MLP Layers]
    F --> G[Scoring Layer]
    G --> H[Ranking Score]

Two-Tower Design

The model consists of:

• User Tower: Processes user features (demographics, behavior history)
• Item Tower: Processes item features (content, category, metadata)
• Interaction Layer: Computes dot product or MLP-based interaction
• Scoring Layer: Outputs relevance score for ranking

Core Model Implementation

import torch
import torch.nn as nn
import torch.nn.functional as F

class TwoTowerModel(nn.Module):
    def __init__(self, user_feature_dim, item_feature_dim, 
                 user_embedding_dim=128, item_embedding_dim=128,
                 hidden_dims=[256, 128, 64]):
        super().__init__()
        
        # User tower
        self.user_tower = nn.Sequential(
            nn.Linear(user_feature_dim, user_embedding_dim),
            nn.ReLU(),
            nn.Linear(user_embedding_dim, user_embedding_dim)
        )
        
        # Item tower
        self.item_tower = nn.Sequential(
            nn.Linear(item_feature_dim, item_embedding_dim),
            nn.ReLU(),
            nn.Linear(item_embedding_dim, item_embedding_dim)
        )
        
        # MLP for interaction
        self.mlp = nn.Sequential(
            nn.Linear(user_embedding_dim + item_embedding_dim, hidden_dims[0]),
            nn.ReLU(),
            nn.Linear(hidden_dims[0], hidden_dims[1]),
            nn.ReLU(),
            nn.Linear(hidden_dims[1], 1)
        )
        
    def forward(self, user_features, item_features):
        # Get embeddings from towers
        user_emb = self.user_tower(user_features)
        item_emb = self.item_tower(item_features)
        
        # Concatenate embeddings
        combined = torch.cat([user_emb, item_emb], dim=1)
        
        # Score prediction
        score = self.mlp(combined).squeeze()
        return torch.sigmoid(score)
    
    def compute_loss(self, user_features, item_features, labels):
        predictions = self.forward(user_features, item_features)
        loss = F.binary_cross_entropy(predictions, labels.float())
        return loss

# Example usage for inference
def rank_items(user_features, item_features_list, model):
    """Rank items for a given user"""
    scores = []
    for item_features in item_features_list:
        score = model(user_features, item_features)
        scores.append(score.item())
    return sorted(zip(item_features_list, scores), key=lambda x: x[1], reverse=True)

Model Complexity vs Latency Trade-off

Model	Parameters	Inference Time (ms)	Memory Usage	Accuracy (AUC-ROC)
Logistic Regression	~10K	0.1	1MB	0.72
Two-Tower DNN (Small)	~500K	2.5	10MB	0.85
Two-Tower DNN (Medium)	~2M	5.0	50MB	0.89
Two-Tower DNN (Large)	~10M	15.0	200MB	0.92

Offline Evaluation Baseline

Baseline model: Logistic Regression with TF-IDF features Expected AUC-ROC: 0.72 Proposed model: Two-Tower DNN with embeddings Expected AUC-ROC: 0.89

Key Design Decisions

1. Two-Tower Architecture: Separates user and item processing for better scalability
2. Embedding Dimensions: 128-dimensional embeddings provide good balance between expressiveness and efficiency
3. MLP Layers: Three-layer MLP for capturing complex interactions
4. Activation Functions: ReLU for non-linearity, sigmoid for final probability output
5. Loss Function: Binary cross-entropy for binary classification task

The model is designed to be production-ready with efficient inference characteristics while maintaining strong performance for recommendation tasks.

5. Serving & Inference Architecture

System Diagram

graph TD
    A[User Request] --> B[API Gateway]
    B --> C[Feature Store Lookup]
    C --> D[Embedding Cache]
    D --> E[Model Serving Layer]
    E --> F[Scoring/Re-ranking]
    F --> G[Response Formatter]
    G --> H[User Response]

    subgraph "Online Serving"
        C
        D
        E
        F
        G
    end

    subgraph "Model Store"
        I[Model Registry]
        J[Versioned Models]
    end

    I --> E
    J --> E

Online vs Batch Serving Decision

Decision: Online serving for real-time personalization use case.

Justification:

• Latency requirement: < 100ms for user-facing personalization (e.g., content ranking)
• User interaction: Real-time relevance scoring for each user request
• Business need: Immediate feedback loop for content recommendations
• Batch serving would introduce 10x+ latency overhead (500ms+), unsuitable for real-time engagement

Two-Stage Pipeline

Pipeline Design:

1. Candidate Retrieval:
   • Fast approximate nearest neighbor search (e.g., FAISS, Annoy)
   • Initial candidate pool of 1000 items from user/item embeddings
2. Scoring & Re-ranking:
   • Full model scoring on candidate set (e.g., two-tower DNN)
   • Final ranking using learned relevance scores

Justification:

• Reduces inference cost from O(N) to O(log N) for candidate selection
• Maintains high-quality ranking while preserving low latency

Serving Infrastructure

Selected Framework: TorchServe with custom model wrapper

Justification:

• Flexibility: Custom inference logic for multi-modal features
• Performance: Optimized for PyTorch models with CUDA support
• Scalability: Built-in model versioning and multi-model support
• Integration: Native Kubernetes deployment with Prometheus metrics

Latency Budget Breakdown

Stage	p50 Latency	p99 Latency	Budget
Feature Lookup	5ms	15ms	20ms
Model Inference	15ms	35ms	40ms
Scoring/Re-ranking	5ms	15ms	20ms
Response Formatting	2ms	5ms	10ms
Total	27ms	75ms	90ms

Model Compression Strategy

Approach: Quantization + Pruning

Strategy	Accuracy Impact	Latency Improvement	Memory Reduction
FP32 → INT8 Quantization	-1.5% AUC	30% inference speedup	75% memory
Pruning (20%)	-0.8% AUC	25% inference speedup	20% model size
Distillation (Student Model)	-3.2% AUC	50% speedup	60% size

Implementation:

• Use PyTorch's torch.quantization for INT8 quantization
• Apply pruning with torch.nn.utils.prune

Caching Strategy

Cache Type	Purpose	TTL	Size
Prediction Cache	Reuse recent predictions for same users	5min	10M entries
Embedding Cache	User/item embeddings for frequent lookups	30min	1M entries
KV Cache	For LLM-based content generation	2min	500K entries

Auto-scaling and Capacity Planning

Scaling Strategy:

• Horizontal Scaling: Kubernetes pods with HPA (Horizontal Pod Autoscaler)
• Vertical Scaling: GPU instance types (P3, P4) with 16GB VRAM
• Warm-up: Preloaded model instances to handle burst traffic

Capacity Calculation:

• QPS Target: 10,000 requests/sec
• Model Inference Time: ~25ms (p99)
• Required GPU Resources: 100x P4 instances (25ms × 10,000 = 250,000ms = 250s per second)
• Peak Load: 10x QPS during traffic spikes (100k/sec)

Monitoring:

• CPU/GPU Utilization: Target <85% for stable performance
• Queue Depth: Alert if >100 pending requests
• Error Rate: <0.1% for production stability

Core Inference Code Sample

import torch
from torch import nn
import numpy as np

class TwoTowerModel(nn.Module):
    def __init__(self, user_emb_dim=128, item_emb_dim=128, hidden_dim=256):
        super().__init__()
        self.user_embedding = nn.Embedding(100000, user_emb_dim)
        self.item_embedding = nn.Embedding(500000, item_emb_dim)
        self.user_mlp = nn.Sequential(
            nn.Linear(user_emb_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, 64)
        )
        self.item_mlp = nn.Sequential(
            nn.Linear(item_emb_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, 64)
        )
        self.final_layer = nn.Linear(64, 1)

    def forward(self, user_ids, item_ids):
        u_emb = self.user_embedding(user_ids)
        i_emb = self.item_embedding(item_ids)
        
        u_repr = self.user_mlp(u_emb)
        i_repr = self.item_mlp(i_emb)
        
        # Dot product similarity
        scores = torch.sum(u_repr * i_repr, dim=1, keepdim=True)
        return torch.sigmoid(scores)

    def predict(self, user_ids, item_ids):
        self.eval()
        with torch.no_grad():
            scores = self.forward(user_ids, item_ids)
            return scores.squeeze().cpu().numpy()

# Example usage for online serving
model = TwoTowerModel()
user_ids = torch.tensor([1, 2, 3])
item_ids = torch.tensor([100, 200, 300])
scores = model.predict(user_ids, item_ids)
print(f"Relevance scores: {scores}")

This implementation provides:

• Real-time scoring capability with minimal latency
• Production-ready inference interface
• Support for batched requests (vectorized operations)
• Memory-efficient embedding lookups
• Clear separation between model logic and serving infrastructure

6. Evaluation Framework

Offline Evaluation Methodology

For offline evaluation, we use temporal cross-validation to ensure our model generalizes across time and avoids data leakage. This approach splits the dataset chronologically, training on older data and evaluating on newer data to simulate real-world deployment conditions. For imbalanced datasets, we apply stratified sampling during cross-validation folds to maintain class distribution in both training and test sets.

Temporal splits are made at 3-month intervals, with a 6-month holdout for final validation. We also implement time-based shuffling to prevent temporal autocorrelation in training data, ensuring that the model learns from patterns that will persist in future data.

Slice-Based Evaluation

We define critical slices for evaluation to detect performance degradation across user segments and environmental conditions:

Slice	Definition	Rationale
User Demographics	Age groups (18-25, 26-35, 36-50, 50+)	Different usage patterns and preferences
Geographic Region	US, EU, APAC, LATAM	Varying content consumption and cultural preferences
Device Type	Mobile, Desktop, Tablet	Different engagement patterns and feature availability
Content Category	News, Entertainment, Education	Varying user expectations and engagement dynamics

Each slice is evaluated for AUC-ROC, Precision@K, and Recall@K to detect potential model drift or bias in specific user segments.

A/B Test Design

We design a robust A/B test with the following specifications:

• Randomization Unit: User ID (to prevent spill-over between control and treatment groups)
• Sample Size Calculation: Using a power analysis with 80% power, 5% significance level, and expected lift of 2% in CTR. With a baseline CTR of 3%, we require approximately 15,000 users per group.
• Duration: 30-day test period to capture seasonal trends and ensure statistical significance
• Primary Metric: Click-through rate (CTR) with secondary metrics including conversion rate and revenue per user

We implement a multi-armed bandit approach for the first 10 days to optimize for performance while maintaining statistical validity.

Guardrail Metrics and Blocking Conditions

To ensure safe launch, we define the following guardrail metrics:

Metric	Threshold	Action
AUC-ROC Drop	< 0.05 from baseline	Block launch
Precision@K Drop	< 0.1 from baseline	Block launch
Recall@K Drop	< 0.1 from baseline	Block launch
Model Confidence Variance	> 25% increase	Investigate before launch

If any guardrail metric crosses its threshold, the system triggers an automatic rollback to the previous stable model version and alerts the ML team.

Shadow Mode Rollout Plan

Before full A/B testing, we implement a shadow mode rollout:

1. Phase 1 (Week 1): 5% of traffic routed to new model with shadow logging
2. Phase 2 (Week 2): 15% of traffic with shadow logging and real-time performance monitoring
3. Phase 3 (Week 3): 50% of traffic with shadow logging and model performance comparison
4. Phase 4 (Week 4): Full A/B test launch with real traffic routing

During shadow mode, all model predictions are logged for offline comparison against the baseline model to validate performance before full rollout.

Long-Term Metric Measurement

We track long-term user engagement metrics to ensure model impact extends beyond immediate performance:

Metric	Measurement Window	Purpose
Delayed Conversion Rate	30-day post-click	Measure long-term user value
Engagement Decay	7, 14, 30-day engagement	Track retention impact
Revenue Per User	30-day period	Validate monetization impact
User Lifetime Value	90-day period	Assess long-term business value

These metrics are tracked using cohort analysis to understand the model's impact on user retention and monetization over time.

# Example evaluation code for offline AUC-ROC calculation
import numpy as np
from sklearn.metrics import roc_auc_score

def calculate_offline_metrics(y_true, y_pred_proba):
    """
    Calculate offline evaluation metrics for binary classification.
    
    Args:
        y_true: Ground truth labels (binary)
        y_pred_proba: Predicted probabilities for positive class
    
    Returns:
        dict: Dictionary containing AUC-ROC and other key metrics
    """
    # Calculate AUC-ROC
    auc_roc = roc_auc_score(y_true, y_pred_proba)
    
    # Calculate AUC-PR (Precision-Recall AUC)
    from sklearn.metrics import average_precision_score
    auc_pr = average_precision_score(y_true, y_pred_proba)
    
    # Calculate Precision@K for top K recommendations
    k_values = [1, 5, 10, 20]
    precision_at_k = {}
    
    # Sort predictions by probability descending
    sorted_indices = np.argsort(y_pred_proba)[::-1]
    
    for k in k_values:
        top_k_indices = sorted_indices[:k]
        relevant_items = np.sum(y_true[top_k_indices])
        precision_at_k[k] = relevant_items / k if k > 0 else 0
    
    return {
        "auc_roc": auc_roc,
        "auc_pr": auc_pr,
        "precision_at_k": precision_at_k
    }

# Example usage:
# y_true = [0, 1, 1, 0, 1, 0, 1]
# y_pred_proba = [0.1, 0.9, 0.8, 0.2, 0.7, 0.3, 0.6]
# metrics = calculate_offline_metrics(y_true, y_pred_proba)

7. Training Pipeline

Training Data Pipeline

The training data pipeline follows a multi-stage ETL process:

1. Data Extraction: Raw data from user interactions, item metadata, and contextual signals is extracted from various sources (click logs, user profiles, item catalogs) every 2 hours
2. Data Transformation: Features are computed and normalized using the feature store, with missing values handled through imputation strategies
3. Data Validation: Data quality checks include:
   • Duplicate detection and removal
   • Feature value range validation (e.g., age must be between 0-150)
   • Label consistency checks (positive/negative ratio within expected bounds)
4. Data Shuffling and Batching: Data is shuffled using reservoir sampling to ensure temporal distribution and batched for training with size 1024

def process_training_batch(raw_data):
    # Data validation and cleaning
    cleaned_data = raw_data.drop_duplicates().fillna(0)
    
    # Feature engineering (using pre-computed features from feature store)
    features = cleaned_data[['user_id', 'item_id', 'timestamp', 'category']].copy()
    
    # Label generation (e.g., click = 1, no click = 0)
    labels = cleaned_data['clicked'].astype(int)
    
    return features, labels

Distributed Training Strategy

For large-scale training, we use data parallelism with model parallelism for embedding layers:

• Data Parallelism: Training data is distributed across multiple GPUs, each processing a shard of the dataset
• Model Parallelism: Embedding tables are sharded across GPUs to handle large embedding matrices (e.g., 10M users × 128-dim embeddings)
• Pipeline Parallelism: For very large models, we pipeline forward and backward passes to overlap computation with data loading

Training Infrastructure

• Compute Cluster: 8x NVIDIA A100 GPUs (40GB VRAM each) for distributed training
• Storage: SSD-backed storage with 50TB capacity for training data and model checkpoints
• Training Time Estimate: With 10M samples, batch size 1024, and 8 GPUs, training takes ~2 hours per epoch

Hyperparameter Tuning Approach

We use Bayesian optimization for hyperparameter tuning:

• Key Hyperparameters:
  • Learning rate (0.001, 0.01, 0.1)
  • Batch size (256, 512, 1024)
  • Embedding dimension (32, 64, 128)
  • Dropout rate (0.1, 0.3, 0.5)
  • Number of layers (2, 3, 4)

Experiment Tracking and Reproducibility

• MLflow is used for experiment tracking:
  • Logging model parameters, metrics, and artifacts
  • Version control of models with tags (e.g., "v1.2.0", "production")
  • Model registry integration for model versioning
• Reproducibility:
  • Fixed random seeds (42) for all components
  • Docker containers with pinned dependencies

Model Versioning and Registry Strategy

• Versioning: Semantic versioning (v1.0.0, v1.1.0) with automatic increment based on performance improvements
• Registry: MLflow Model Registry for model lifecycle management:
  • Staging and production models are tagged
  • Rollback capability with version control
• Model Artifacts: Saved as PyTorch checkpoints with serialized model weights and configuration

Continual / Online Learning Design

• Online Learning Strategy:
  • For new user/item interactions, we use incremental updates with a small learning rate (0.0001)
  • For concept drift, we implement periodic retraining with recent data (last 30 days)
• Feedback Loop:
  • User feedback is collected every hour and used to update model weights via online gradient descent
  • A/B testing of new model versions with shadow mode for validation before full rollout

Training Pipeline Monitoring

• Metrics:
  • Loss curves and AUC-ROC for each epoch
  • Data drift detection on training features
• Alerts:
  • Training loss spikes > 20% trigger manual investigation
  • AUC-ROC drop below threshold (0.85) triggers retraining alert

8. Monitoring & Observability

Data Drift Detection

Data drift monitoring is critical for maintaining model performance over time. We implement drift detection using:

• Population Stability Index (PSI) for categorical features and discrete numerical features
• Kullback-Leibler divergence for continuous feature distributions
• Chi-square test for categorical distribution comparison

import numpy as np
from scipy import stats

def calculate_psi(expected, actual, bins=10):
    """Calculate Population Stability Index between two distributions"""
    # Ensure both distributions have same bins
    _, bin_edges = np.histogram(expected, bins=bins)
    
    # Count occurrences in each bin
    expected_counts, _ = np.histogram(expected, bins=bin_edges)
    actual_counts, _ = np.histogram(actual, bins=bin_edges)
    
    # Avoid zero division
    expected_counts = np.where(expected_counts == 0, 1e-6, expected_counts)
    actual_counts = np.where(actual_counts == 0, 1e-6, actual_counts)
    
    # Calculate PSI
    psi = np.sum((actual_counts - expected_counts) * np.log(actual_counts / expected_counts))
    return psi

def kl_divergence(p, q):
    """Calculate KL divergence between two distributions"""
    return np.sum(np.where(p != 0, p * np.log(p / q), 0))

Prediction Drift and Label Drift Monitoring

Prediction drift is monitored by tracking distribution shifts in model outputs:

from sklearn.metrics import roc_auc_score

class PredictionDriftDetector:
    def __init__(self, drift_threshold=0.1):
        self.drift_threshold = drift_threshold
        self.reference_predictions = []
        
    def detect_prediction_drift(self, current_predictions):
        """Detect drift in prediction distributions"""
        if not self.reference_predictions:
            self.reference_predictions = current_predictions
            return False
            
        # Compare distribution using KL divergence
        reference_dist = np.histogram(self.reference_predictions, bins=50)[0]
        current_dist = np.histogram(current_predictions, bins=50)[0]
        
        drift_score = kl_divergence(reference_dist / len(reference_dist), 
                                   current_dist / len(current_dist))
        
        return drift_score > self.drift_threshold

Model Performance Degradation Alerting

We implement alerting thresholds for key metrics:

Metric	Alert Threshold	Escalation
AUC-ROC	< 0.85	Alert
AUC-PR	< 0.75	Alert
Precision@K	< 0.60	Alert
Recall@K	< 0.50	Alert

Retraining Trigger Strategy

We use a hybrid approach for retraining triggers:

1. Scheduled retraining: Every 30 days or when model performance drops below threshold
2. Drift-triggered: When PSI > 0.1 or prediction drift detected
3. Performance-triggered: When AUC drops by 5% over 7-day window

Logging Design

Prediction logs include:

class PredictionLogger:
    def __init__(self):
        self.prediction_log = []
        
    def log_prediction(self, user_id, item_id, features, prediction_score, 
                       actual_label=None, timestamp=None):
        log_entry = {
            "user_id": user_id,
            "item_id": item_id,
            "features": features,
            "prediction_score": prediction_score,
            "actual_label": actual_label,
            "timestamp": timestamp or time.time(),
            "model_version": "v1.2"
        }
        self.prediction_log.append(log_entry)
        
        # Write to monitoring system (e.g., Kafka or CloudWatch)
        self.write_to_monitoring_system(log_entry)

Feedback Loop for Continuous Improvement

We implement a feedback collection pipeline:

class FeedbackCollector:
    def __init__(self):
        self.feedback_queue = []
        
    def collect_feedback(self, user_id, item_id, feedback_type, 
                        feedback_value, timestamp):
        """Collect user feedback for model improvement"""
        feedback = {
            "user_id": user_id,
            "item_id": item_id,
            "feedback_type": feedback_type,  # click, view, like, etc.
            "feedback_value": feedback_value,
            "timestamp": timestamp
        }
        
        # Store for retraining pipeline
        self.feedback_queue.append(feedback)
        
        # Update feature store with new signals
        self.update_feature_store(user_id, item_id, feedback)

Model Card Documentation Plan

We maintain model cards with:

• Performance metrics (AUC-ROC, AUC-PR, Precision@K)
• Data drift indicators
• Feature importance and drift status
• Model version history and retraining logs

Incident Response Plan for Model Failures

1. Immediate fallback: Route traffic to rule-based baseline
2. Alert escalation: Notify ML team within 15 minutes
3. Root cause analysis: Investigate drift, data quality, or model instability
4. Recovery plan: Re-train with updated data or rollback to previous version

Key Metrics for Monitoring

Metric	Frequency	Alert Threshold
AUC-ROC	Hourly	< 0.85
AUC-PR	Hourly	< 0.75
PSI (categorical)	Daily	> 0.1
KL divergence (continuous)	Daily	> 0.05
Prediction drift	Hourly	> 0.1
Model version	Real-time	Version mismatch

Back-of-Envelope Drift Detection Calculations

For a system with:

• 1M daily active users
• 100K items served per day
• 500M total training samples (50GB)
• 100 features per user/item

Expected drift detection frequency:

• PSI calculation: ~50ms per feature
• Total drift check: ~50ms × 100 = 5 seconds per batch
• Batch processing every 24 hours = ~86,400 seconds

This enables real-time drift detection with minimal overhead on production systems.

9. Responsible AI & Safety

Bias Analysis

Bias in ML systems can manifest through skewed training data, algorithmic bias in feature selection, or model decision-making. We conduct a comprehensive audit of training data to identify potential demographic imbalances in user representation, item categories, and geographic distribution. We evaluate model performance across slices defined by:

• Demographics (age, gender, region)
• User behavior segments (engagement level, recency)
• Content categories (topic, toxicity score)
• Platform characteristics (device type, OS version)

We implement slice-based fairness evaluation using metrics like demographic parity and equalized odds to detect and quantify bias in model predictions.

Fairness Metrics and Constraints

We prioritize demographic parity for user-level fairness, ensuring that positive prediction rates are similar across protected groups. Equalized odds is applied to maintain similar true positive and false positive rates across slices.

For binary classification tasks, we define fairness constraints as:

# Example constraint for demographic parity (simplified)
def demographic_parity_constraint(y_pred, sensitive_attr):
    # Group A and B predictions
    pred_A = y_pred[sensitive_attr == 'A']
    pred_B = y_pred[sensitive_attr == 'B']
    
    # Ratio of positive predictions should be similar
    ratio_A = np.mean(pred_A)
    ratio_B = np.mean(pred_B)
    
    return abs(ratio_A - ratio_B) <= 0.05  # Within 5% tolerance

Explainability Requirements and Approach

We implement SHAP (SHapley Additive exPlanations) for model interpretability, particularly for high-stakes decisions like content moderation or user targeting. For transformer-based models, we utilize attention visualization to understand which input tokens contribute most to predictions.

For LLM Systems: Safety Guardrails

For systems involving large language models, we implement:

• Content filtering using toxicity detection models (e.g., toxicity classifier from Perspective API)
• Output validation with regex patterns and semantic consistency checks
• Rate limiting to prevent abuse or spam generation
• Red teaming with adversarial examples to identify model weaknesses

Model Card Documentation Plan

We maintain comprehensive model cards that include:

• Training data sources and preprocessing steps
• Model architecture and hyperparameters
• Performance metrics across different slices
• Known limitations and bias considerations
• Usage guidelines and safety recommendations

Incident Response Plan for Model Failures or Harmful Outputs

We establish a multi-tiered incident response process:

1. Automated Detection: Alerting on prediction drift, label drift, or safety violations
2. Manual Review: Human-in-the-loop for flagged outputs or model degradation
3. Rollback Strategy: A/B testing rollback to previous stable model version
4. Post-mortem Analysis: Root cause analysis and mitigation planning for any safety incidents

The incident response team includes ML engineers, platform reliability engineers, and product stakeholders to ensure rapid and effective handling of any model-related safety or fairness issues.

10. Scalability & Reliability

Bottleneck Analysis per Pipeline Stage

Stage	Bottleneck Type	Back-of-Envelope Estimate	Mitigation Strategy
Data Ingestion	I/O throughput	10GB/s peak, 500MB/s average	Use distributed file systems (e.g., HDFS) and parallel readers
Feature Computation	CPU/memory	500GB feature matrix, 100ms per user	Implement distributed feature store with caching
Model Training	GPU memory	20GB per node, 1000 nodes	Use model parallelism and gradient checkpointing
Model Serving	Latency	100ms p99, 500 QPS	Implement two-stage pipeline with candidate retrieval
Model Updates	Freshness	24h window for retraining	Implement continuous learning with drift detection

Scale-up Strategy

Horizontal Scaling of Serving:

• Deploy model replicas across multiple regions to reduce latency and increase availability
• Use Kubernetes with auto-scaling based on QPS and latency metrics

Distributed Feature Computation:

• Implement feature store with distributed compute (e.g., Spark or Ray)
• Use partitioning strategies to scale feature computation linearly with data volume

Fault Tolerance

Model Serving Fallbacks:

• Implement graceful degradation to rule-based fallback when model fails
• Use circuit breaker pattern to prevent cascading failures

Stale Model Serving on Failure:

• Maintain last-known-good model version for serving
• Implement model version rollback capability with monitoring

Cold Start Handling Strategy

New Users:

• Use demographic-based cold start with default ranking
• Implement content-based filtering using user profile features

New Items:

• Use popularity-based ranking for new items
• Implement candidate generation from similar items using item embeddings

Multi-Region Serving Strategy

Global Distribution:

• Deploy model instances in multiple regions (US, EU, Asia) to reduce latency
• Use CDN for static model artifacts and feature embeddings

Cross-Region Coordination:

• Implement global model versioning with region-specific serving
• Use consistent hashing for feature lookups across regions

Traffic Management:

• Route traffic to nearest region based on latency
• Implement failover mechanisms with health checks for each region

graph TD
    A[User Request] --> B[Load Balancer]
    B --> C[Regional Model Server]
    C --> D[Feature Store]
    D --> E[Model Serving]
    E --> F[Response]
    
    subgraph RegionalDeployment
        C
        D
        E
    end
    
    subgraph GlobalInfrastructure
        A
        B
    end

11. Open Questions, Risks & Trade-offs

Decisions requiring product / engineering / policy alignment

• Label quality thresholds: What level of human annotation quality is acceptable for training? How do we balance labeling cost vs. model performance?
• Latency vs. accuracy trade-off: Should we prioritize sub-100ms response times or maximize AUC-ROC/AUC-PR at the cost of higher latency?
• Data privacy constraints: How do we balance model performance with strict data retention and anonymization policies?
• Bias mitigation strategy: Should we prioritize demographic parity or equalized odds across user segments?

Key ML risks

Risk	Impact	Likelihood	Mitigation
Data drift in training distribution	High - model performance degradation	Medium - continuous monitoring can detect drift	Implement feature drift detection, scheduled retraining
Label noise in training data	High - performance degradation	Medium - weak supervision and human review can help	Implement label quality checks, use robust loss functions
Cold start problem for new users/items	Medium - reduced recommendation quality	High - frequent model updates and fallback strategies	Implement hybrid rule-based + ML approach for new entities
Model overfitting to training data	Medium - poor generalization	Medium - cross-validation and regularization can help	Use proper train/validation/test splits, early stopping
Imbalanced class distribution	Medium - poor recall for rare events	High - common in recommendation systems	Implement stratified sampling, weighted loss functions

Trade-offs made with justification

Trade-off	Justification
Accuracy vs. latency	Prioritize sub-100ms serving latency over marginal AUC improvement to maintain user experience and system responsiveness
Freshness vs. cost	Accept 24-hour data freshness window to reduce infrastructure costs and training overhead
Model size vs. quality	Use model compression (quantization) to reduce serving cost while maintaining >95% of original AUC-ROC
Online vs. batch training	Use batch training for stability and reproducibility, with online learning for real-time feedback
Feature richness vs. latency	Select core features to reduce serving complexity, defer complex cross-features to batch training

Alternatives considered and rejected

Alternative	Reason for rejection
Transformer-based models	Too computationally expensive for real-time serving, not justified by marginal gains in AUC-ROC
Ensemble of multiple models	Increased complexity and latency without significant performance gains for this use case
Real-time feature computation	Too high latency cost; batch features are sufficient for this recommendation system
Multi-task learning	Not needed for a single ranking objective; adds complexity without clear benefit
Full online learning	Too unstable for production use, requires extensive monitoring and fallback handling

Technical debt acknowledged

• Limited A/B testing infrastructure in initial phases may lead to suboptimal model rollouts
• Manual feature engineering approach will require significant manual effort as the system scales
• Simplified model architecture may not capture complex user-item interactions present in the data
• Basic drift detection using simple statistical methods may miss subtle but impactful changes in user behavior or data distribution

12. Phased Implementation Roadmap

Phase 1: MVP (Weeks 1-4)

Objective: Establish baseline system with rule-based fallback and simple ML model.

Deliverables:

• Rule-based recommendation engine (cold start handling, popularity-based)
• Basic feature extraction pipeline
• Simple logistic regression model trained on historical data
• Online serving infrastructure with 100 QPS capacity
• Basic A/B testing framework (5% shadow mode)

Dependencies:

• Data pipeline ready with 1M training samples
• Feature store MVP (batch features only)
• Model registry and experiment tracking (MLflow)

Critical Path:

1. Data pipeline validation → 2. Model training and validation → 3. Basic serving infrastructure → 4. A/B test setup

Phase 2: Scale (Weeks 5-8)

Objective: Expand to real-time features, distributed training, and full A/B testing.

Deliverables:

• Real-time feature computation (user behavior, item metadata)
• Distributed training pipeline (data parallelism on 4 GPUs)
• Feature store with online and offline stores
• Multi-stage serving pipeline (candidate retrieval → scoring)
• Full A/B test infrastructure with 10% randomization

Dependencies:

• Feature store fully operational
• Distributed training cluster provisioned
• Real-time data ingestion pipeline (Kafka/Beam)
• Model versioning and deployment automation

Critical Path:

1. Feature store integration → 2. Distributed training pipeline → 3. Multi-stage serving → 4. A/B testing with proper metrics

Phase 3: Excellence (Weeks 9-12)

Objective: Advanced models, continuous learning, and performance optimization.

Deliverables:

• Transformer-based ranking model with attention mechanisms
• Online learning pipeline (incremental model updates)
• Real-time feature enrichment and contextual signals
• Advanced monitoring and drift detection system
• Multi-region serving with 1000 QPS capacity

Dependencies:

• Model compression and quantization pipeline
• Continuous learning framework (online model updates)
• Multi-region infrastructure (AWS/GCP)
• Production-grade observability stack

Critical Path:

1. Model compression and optimization → 2. Online learning implementation → 3. Multi-region deployment → 4. Advanced monitoring and alerting

Cross-Phase Dependencies:

• Feature store must be available before model training in Phase 2
• A/B testing infrastructure must be ready for production rollout in Phase 3
• Model versioning and experiment tracking must support all phases

Risk Mitigation:

• Shadow mode rollout before full A/B testing in Phase 2
• Fallback to rule-based system during model outages
• Gradual QPS scaling with monitoring in Phase 3

Success Metrics:

• 20% improvement in CTR over rule-based baseline (Phase 1)
• 30% improvement in engagement metrics (Phase 2)
• 50% reduction in model latency (Phase 3)