RAG retrieval degradation after MLflow model validation

Problem – Retrieval‑Augmented Generation (RAG) Hybrid Search Degradation During MLflow Validation

During automated pre‑deployment validation runs orchestrated by an MLflow Project in the CI/CD pipeline, the hybrid retriever’s relevance metrics collapse:

  • Recall@10 drops by 40‑50 % compared with manual runs.
  • Logs contain errors such as ScoreNormalizationError: Sparse scores exceed dense scores range and warnings like HybridSearchWeightOverflow.
  • Qualitative inspection shows BM25 lexical matches (score range 0‑12) dominating cosine similarity scores (range –1‑1), injecting irrelevant passages into the generation context.
  • Metrics logged by MLflow Tracking API flag a steep increase in MetricComputationFailed: recall@k undefined due to zero valid retrievals.

These symptoms appear only in the staging validation step; the same model version performs well in local notebooks.

Root Cause – Unnormalized Scoring Scales in Hybrid Retrieval

The hybrid retriever combines a sparse lexical scorer (e.g., BM25) and a dense embedding scorer (cosine similarity). In the production pipeline the two scores are summed directly:

final_score = lexical_score + dense_score

Because the lexical score range (0‑12) is orders of magnitude larger than the dense score range (‑1‑1), the dense contribution becomes negligible. The issue is triggered during the MLflow validation run when:

  • The HybridRetriever is instantiated without an explicit score_normalizer (see LangChain issue #4123).
  • MLflow reloads the embedding model with a different dimension, breaking any hard‑coded scaling factor (real incident “InvalidScoringScaleException”).
  • The CI environment uses a default hybrid weight of 0.9 for the lexical component (legal‑document incident), further amplifying the imbalance.

Consequently, the final ranking is driven almost entirely by lexical relevance, causing “irrelevant context injection” during automated evaluation.

Debug – Investigation Steps

1. Examine MLflow Validation Logs

2024-07-31 02:14:07,321 - mlflow.tracking - INFO - Starting validation run for model version 3
2024-07-31 02:14:08,112 - retriever.hybrid - WARNING - HybridSearchWeightOverflow: lexical weight=0.9 exceeds 0.5 without normalization
2024-07-31 02:14:08,115 - retriever.hybrid - ERROR - ScoreNormalizationError: Sparse scores exceed dense scores range
2024-07-31 02:14:09,543 - mlflow.tracking - ERROR - MetricComputationFailed: recall@10 undefined due to zero valid retrievals

2. Reproduce the Scoring Path Locally

Run the same validation script outside MLflow to print raw scores:

python validate_rag.py --print-scores
Query: "What are the compliance requirements for GDPR?"
Lexical (BM25) top‑3 scores: [9.2, 8.7, 7.5]
Dense (cosine) top‑3 scores: [0.42, 0.38, 0.35]
Final combined scores (without normalization): [9.62, 9.08, 7.85]

3. Verify Embedding Model Load

mlflow models serve -m runs://model
# In a separate terminal
curl -X POST http://127.0.0.1:5000/invocations -d '{"inputs": ["sample query"]}'
# Response includes embedding dimension
"embedding_dim": 768

If the dimension differs from the one used when the hybrid weight was calibrated (e.g., 1024), the scaling factor is invalid.

4. Check Configuration in MLflow Project

# mlflow_project.yaml
name: rag_validation
conda_env: env.yml
entry_points:
  validate:
    command: "python -m validation.run"
    parameters:
      hybrid_weight: {type: float, default: 0.5}

Confirm that hybrid_weight is not overridden by CI environment variables.

Solution – Normalizing Scores and Calibrating Hybrid Weights

1. Introduce Min‑Max Scaling for Sparse Scores

Apply a MinMaxScaler (range 0‑1) to the lexical scores before combination.

# validation/run.py
from sklearn.preprocessing import MinMaxScaler
import numpy as np

def normalize_lexical(scores):
    scaler = MinMaxScaler(feature_range=(0, 1))
    return scaler.fit_transform(np.array(scores).reshape(-1, 1)).flatten()

def hybrid_score(lexical, dense, weight_lexical=0.5):
    lexical_norm = normalize_lexical(lexical)
    # dense scores are already in [-1, 1]; shift to [0, 1] for consistency
    dense_norm = (np.array(dense) + 1) / 2
    return weight_lexical * lexical_norm + (1 - weight_lexical) * dense_norm

2. Parameterize the Hybrid Weight and Enforce Bounds

# mlflow_project.yaml (updated)
parameters:
  hybrid_weight:
    type: float
    default: 0.5
    min: 0.0
    max: 1.0

In the CI pipeline, pass the calibrated weight explicitly:

mlflow run . -e validate -P hybrid_weight=0.4

3. Align Embedding Dimensions and Scaling Factors

When the embedding model version changes, recompute the dense score range:

# utility to compute dense score statistics
def dense_score_stats(embeddings):
    # cosine similarity yields [-1, 1]; after shift becomes [0, 2]
    shifted = (embeddings + 1) / 2
    return shifted.min(), shifted.max()

Update the normalization step accordingly.

4. Before / After Comparison

Before (no normalization)

Lexical: [9.2, 8.7, 7.5]
Dense:   [0.42, 0.38, 0.35]
Combined: [9.62, 9.08, 7.85]   # lexical dominates

After (MinMax + calibrated weight 0.4)

Lexical normalized: [1.0, 0.95, 0.68]
Dense normalized:   [0.71, 0.69, 0.68]
Combined: [0.4*1.0 + 0.6*0.71 = 0.826,
           0.4*0.95 + 0.6*0.69 = 0.782,
           0.4*0.68 + 0.6*0.68 = 0.68]

The dense component now contributes meaningfully, restoring balanced rankings.

Verify – Confirming the Fix

  • Metric Re‑run: Execute the MLflow validation run and observe Recall@10 returning to baseline (e.g., 0.78 vs. 0.31 pre‑fix).
  • Log Inspection: No longer see ScoreNormalizationError or HybridSearchWeightOverflow warnings.
  • Sample Query Check:
    Lexical norm: [0.92, 0.88, 0.75]
    Dense norm:   [0.65, 0.62, 0.60]
    Final scores: [0.77, 0.73, 0.68]
    Top‑k documents now include dense‑similar passages.
  • Health Endpoint (if exposed):
    curl http://staging-rag/api/health
    {"retriever_status":"ok","hybrid_weight":0.4,"normalization":"minmax"}
    

Prevent – Operational Guardrails

  • Include a score_normalizer in every hybrid retriever configuration; make it a required field in the model registry metadata.
  • Store the calibrated hybrid_weight as a model artifact in MLflow Model Registry; retrieve it programmatically during deployment.
  • Add a CI test that asserts the ratio between max lexical and max dense scores stays within a configurable bound (e.g., max_lexical / max_dense < 5).
  • Version‑lock the embedding model and record its dimension in the mlflow.log_params call; fail the validation if a mismatch is detected.
  • Instrument Prometheus metrics:
    # HELP rag_hybrid_score_ratio Ratio of lexical to dense scores
    # TYPE rag_hybrid_score_ratio gauge
    rag_hybrid_score_ratio 1.2
    

    Alert when the ratio exceeds a threshold.

Related Topic Hub: Model Serving Troubleshooting Hub

FAQ

  1. Why does the hybrid score collapse only during CI runs?

    The CI environment reloads the embedding model with a different dimension, breaking any hard‑coded scaling factor. Additionally, default CI parameters often set hybrid_weight to a high lexical value without normalization.

  2. Can I use a different normalization technique?

    Yes. Alternatives such as Z‑score standardization or a custom sigmoid scaling work, provided the resulting lexical and dense ranges are comparable. The key is to apply the same transformation consistently across runs.

  3. How do I detect score imbalance before it impacts metrics?

    Instrument a pre‑validation check that logs the min/max of each scorer and raises ScoreNormalizationError if max_lexical / max_dense > 10. Integrate this check into the MLflow run entry point.

  4. Is it safe to set hybrid_weight=0.0 to rely solely on dense vectors?

    Setting the weight to 0 disables lexical scoring, which may improve precision but can hurt recall for queries with rare terms. Prefer a balanced weight (e.g., 0.4‑0.6) and monitor both Recall@k and Precision@k.

  5. What if my dense scorer outputs negative cosine similarities?

    Shift the dense scores to a non‑negative range before combination, e.g., (score + 1) / 2. This aligns with the typical 0‑1 range after MinMax scaling of lexical scores.