Experiment Tracking: You Need It, Don't Overthink It
Every ML team reaches a point where someone asks "which hyperparameters produced that model from three weeks ago?" and nobody can answer. That's when you need experiment tracking. The question isn't whether to track — it's which tool fits your workflow without becoming overhead.
I've used MLflow, Weights and Biases (W&B), and Neptune across different teams. They all solve the core problem. The differences are in ergonomics, collaboration features, and hosting model.
MLflow: The Self-Hosted Default
MLflow is the safe choice for teams that want to own their infrastructure. Open-source, self-hosted (or managed via Databricks), and integrates with almost everything:
import mlflow
mlflow.set_tracking_uri("http://mlflow.internal:5000")
mlflow.set_experiment("fraud-detection-v3")
with mlflow.start_run(run_name="xgboost-tuned"):
mlflow.log_params({"n_estimators": 500, "max_depth": 8,
"learning_rate": 0.05, "subsample": 0.8})
model = train_model(params)
metrics = evaluate(model, test_data)
mlflow.log_metrics({"auc_roc": metrics["auc"],
"precision_at_95_recall": metrics["p95"]})
mlflow.sklearn.log_model(model, "model")MLflow's strength is its model registry. Promoting a run's model to the registry, tagging it as staging or production, and having serving systems pull from the registry — this flow works well. For Databricks teams the integration is especially smooth.
The weakness is the UI. Comparing runs involves manual selection and scrolling. Visualizing training curves requires logging metrics at each step, and chart rendering is slow for runs with thousands of steps.
Self-Hosting Costs
Running MLflow means managing PostgreSQL for metadata, S3/GCS for artifacts, and the tracking server. It's not hard to set up but needs maintenance — backups, artifact cleanup, access control. We run MLflow in Kubernetes with S3 for artifacts. Total infrastructure cost is about $200/month for 15 ML engineers.
Weights and Biases: The Collaboration Tool
W&B is the tool I reach for when the project involves multiple people actively iterating. The UI is leagues ahead of MLflow. Side-by-side run comparison, automatic hyperparameter importance plots, and the collaborative reports feature make it easy to share findings.
import wandb
wandb.init(project="fraud-detection", name="xgboost-tuned-v3",
config={"n_estimators": 500, "max_depth": 8,
"learning_rate": 0.05, "subsample": 0.8})
for epoch in range(100):
train_loss, val_loss = train_one_epoch(model, epoch)
wandb.log({"train_loss": train_loss, "val_loss": val_loss, "epoch": epoch})
wandb.finish()The Sweeps feature for hyperparameter search is excellent. Define a search space, choose Bayesian optimization or grid search, and W&B distributes sweep runs across your compute. Not as powerful as Optuna, but the integration with experiment tracking is seamless.
See also: Data Pipeline Orchestration for ML: Airflow, Prefect, and Da.
The downside: it's a SaaS product. Your data lives on their servers unless you pay for the on-premises version. For teams with data residency requirements, this can be a deal-breaker.
Neptune: The Data Science Workbench
Neptune positions itself between MLflow and W&B with better visualization than MLflow and a more flexible metadata system. The key differentiator is nested run structure — organizing experiments hierarchically maps well to research workflows.
In practice, Neptune's adoption is lower, which means smaller community, fewer integrations, and less institutional knowledge. Technically solid but harder to get team buy-in compared to W&B.
What Actually Matters
If you need self-hosting and use Databricks: MLflow. If your team is 5+ people collaborating and SaaS is acceptable: W&B. If you're an individual or small team wanting quick setup: W&B free tier (5 minutes, zero infrastructure) or MLflow local mode.
Whichever tool you pick, log everything from day one. The cost of logging an extra parameter or metric is near zero. The cost of not having logged it when you need it — rerunning a week of experiments — is enormous. Log aggressively. Filter later.
Artifact Management and Model Lineage
Experiment tracking tools store more than just metrics. Artifacts — trained model files, training data snapshots, configuration files, evaluation reports — are equally important for reproducibility. When a production model starts failing six months after training, you need to reproduce the exact training environment: the same data, code, hyperparameters, and dependencies.
MLflow handles artifact storage natively through its artifact URI system. Models, plots, data files — anything can be logged as an artifact and stored in the configured backend (local filesystem, S3, GCS, Azure Blob). The lineage from experiment run to model registry to production deployment creates a traceable chain.
For a related perspective, see Feature Store Design Patterns for Real-Time and Batch Servin.
W&B takes this further with its Artifacts feature (capitalized Artifact, their product name, not generic artifacts). Artifacts have versions, aliases, and dependency graphs. You can trace a production model back through the training run that produced it, the dataset artifact that fed it, and the preprocessing code that created the dataset. This dependency tracking is useful for compliance — regulated industries need to prove exactly which data trained which model.
Team Workflows and Governance
In larger organizations, experiment tracking tools serve a governance function beyond individual productivity. Model review workflows — where a senior ML engineer reviews training runs before promoting to the model registry — prevent poorly validated models from reaching production.
# MLflow model stage transitions with approval
import mlflow
from mlflow.tracking import MlflowClient
client = MlflowClient()
# Promote model to staging after review
client.transition_model_version_stage(
name="fraud-detection",
version=12,
stage="Staging",
archive_existing_versions=False
)
# After validation in staging environment
client.transition_model_version_stage(
name="fraud-detection",
version=12,
stage="Production",
archive_existing_versions=True
)W&B's Reports feature lets engineers write narrative analyses of experiment results — comparing approaches, documenting why one architecture was chosen over another, and recording lessons learned. These reports become institutional knowledge that survives team turnover. In my experience, the knowledge capture aspect of experiment tracking is undervalued. The metrics and models are artifacts; the reasoning behind decisions is the real intellectual property.
Cost Considerations
For self-hosted MLflow, the main costs are storage (model artifacts accumulate fast — a large neural network is 500MB-2GB per checkpoint, and teams easily generate hundreds of checkpoints per week) and compute for the tracking server. Budget for aggressive artifact cleanup policies: delete failed runs after 7 days, archive old experiment runs after 90 days, keep only the top 3 checkpoints per successful training run.
For W&B, the Team plan costs $50/user/month. For a 15-person ML team, that's $750/month — more than the $200 self-hosted MLflow setup but less than the engineering time to maintain the self-hosted infrastructure. The break-even calculation depends on your team's DevOps capacity and the value you place on W&B's superior collaboration features.
Neptune's pricing is usage-based, which can be unpredictable. Heavy experiment logging (thousands of metrics per run, large artifacts) can generate unexpected bills. Monitor usage carefully during the trial period before committing.
Integration with Training Infrastructure
Experiment tracking tools are most valuable when integrated into your training infrastructure. Auto-logging — where the tracking library automatically captures parameters, metrics, and artifacts without explicit logging calls — reduces adoption friction dramatically. MLflow supports auto-logging for scikit-learn, TensorFlow, PyTorch Lightning, and XGBoost. W&B integrates with PyTorch Lightning, Keras, and Hugging Face Transformers via callbacks.
We covered a related topic in Streaming Feature Computation with Apache Flink for Real-Tim.
# PyTorch Lightning + W&B integration
from pytorch_lightning.loggers import WandbLogger
from pytorch_lightning import Trainer
logger = WandbLogger(project="fraud-detection", log_model=True)
trainer = Trainer(
max_epochs=100,
logger=logger,
callbacks=[
ModelCheckpoint(monitor="val_auc", mode="max"),
EarlyStopping(monitor="val_auc", patience=10),
]
)
trainer.fit(model, train_loader, val_loader)The one-line logger setup captures training loss, validation metrics, learning rate schedule, gradient norms, and model checkpoints automatically. No manual logging calls needed. This level of integration means new team members start logging experiments correctly from their first training run.
Hyperparameter Optimization Integration
Combining experiment tracking with hyperparameter optimization creates a powerful feedback loop. Every HPO trial is automatically a logged experiment, making it easy to analyze which hyperparameters matter and which don't.
With W&B Sweeps, you define the search space declaratively and the agent distributes trials across available GPUs. The parallel coordinate plots and importance analysis show you which hyperparameters affect performance and which are noise. In one project, this analysis revealed that our learning rate search range was too narrow — expanding it from [1e-4, 1e-3] to [1e-5, 1e-2] found a significantly better optimum at 5e-5 that we'd been missing.
For Optuna users (a more powerful HPO library), both MLflow and W&B have callback integrations that log each trial as a tracked experiment. Optuna's pruning feature — early-stopping unpromising trials based on intermediate results — pairs especially well with experiment tracking because you can see exactly when and why each trial was stopped.
Long-Term Data Management
After a year of active experimentation, teams generate tens of thousands of logged runs with terabytes of artifacts. Without cleanup policies, storage costs grow linearly while the value of old experiments diminishes. We implement tiered retention: active experiments (last 30 days) keep all artifacts at full resolution. Older experiments keep only metadata, metrics, and the best model checkpoint. Experiments older than 6 months keep only metadata unless explicitly preserved by tagging them as "landmark" runs.
This policy requires discipline. In practice, teams resist deleting old artifacts because "we might need them." Setting up the automated cleanup before the storage cost becomes painful is important. Both MLflow and W&B support programmatic deletion through their APIs, making it straightforward to implement a cleanup cron job that enforces the retention policy.