Spark MLlib
DataFrame
rawPrediction
probability
machine learning

What do columns ‘rawPrediction’ and ‘probability’ of DataFrame mean in Spark MLlib?

Master System Design with Codemia

Enhance your system design skills with over 120 practice problems, detailed solutions, and hands-on exercises.

Spark MLlib is a powerful library within the Apache Spark ecosystem, which facilitates scalable machine learning. When using Spark MLlib's models, you often encounter columns such as rawPrediction and probability in the resulting DataFrame from transformations made using models such as logistic regression or decision trees. Understanding these columns is crucial for interpreting model outputs accurately.

Understanding rawPrediction

and probability Columns in Spark MLlib

When you train a machine learning model using Spark MLlib, the output typically includes predictions in the form of a DataFrame. For classification problems, this DataFrame can have up to three primary columns: prediction , rawPrediction , and probability . Here, we focus on rawPrediction and probability .

rawPrediction

The rawPrediction column contains the raw output scores from the classification model. These scores often represent the decision function—for instance, in SVM, these would be the margins. In the context of logistic regression, these are the log-odds or the values before the logistic function is applied. As such, they don’t constitute probabilities directly but are instead unbounded values which need further processing to become probabilities.

For example, in a binary logistic regression model:

• The model might output a rawPrediction of 2.0 for class 0 and -2.0 for class 1. • These results reflect the distance of the sample point from the decision boundary in feature space.

probability

The probability column, on the other hand, displays the predicted probabilities of each class. This is generated by applying a link function—most commonly the softmax function—to the rawPrediction . For binary classification, this is equivalent to transforming the log-odds from logistic regression to a probability using the logistic function.

For example:

• Given a rawPrediction of 2.0, the probability for class 0 might be calculated as 1/(1+e2.0)1 / (1 + e^{-2.0}), resulting in approximately 0.88. • For class 1, this would instead be e2.0/(1+e2.0)e^{-2.0} / (1 + e^{-2.0}), resulting in approximately 0.12.

Technical Example: Logistic Regression

Consider the use of Spark MLlib's logistic regression for binary classification:

  1. Training the Model:
    • Here, rawPrediction might look like [-0.5, 0.5] indicating the log-odds for both classes. • probability might convert this into [0.38, 0.62] using the logistic function, indicating a 38% chance of class 0 and 62% chance of class 1. • Multi-Class Classification: For multi-class problems, rawPrediction extends to vector outputs corresponding to each class, and probability provides probability vectors summing up to one. • Pipeline Component: Both columns are typically part of the DataFrame produced by the transform method in an MLlib-based pipeline, ideal for further evaluation or debugging. • Threshold Adjustments: While probabilities are often used directly to assign predictions, custom thresholds can be applied to adjust sensitivity/specificity tradeoffs.

Course illustration
Course illustration

All Rights Reserved.