Nonlinear Regression in Apache Spark. Do it yourself

unlocker December 8, 2016 at 22:53

Nonlinear Regression in Apache Spark. Do it yourself

Tutorial

When solving signal processing problems, the method of approximating raw data by a regression model is often used. Based on the structure, models can be divided into three types - linear, reduced to linear and non-linear. In the Spark ML machine learning module Apache Spark, the functionality for the first two types is represented by the LinearRegression and GeneralizedLinearRegression classes, respectively. Training of nonlinear models in the standard library is not presented and requires independent development.

First, we briefly review the theoretical foundations of building non-linear models, and then move on to the practical development of the Spark ML extension.

A bit of math

Learning non-linear models as compared to linear ones is a more complicated task. This may be due to the presence of more than one extremum and / or the “ravine” nature of the response surface. The main stimulus for using nonlinear functions is the possibility of obtaining more compact models. It should also be noted that many analytical equations from the field of physics and engineering are initially non-linear, so the use of appropriate models may be forced.

For training nonlinear models, there is a fairly wide range of tools, the choice of which depends on the specific type of function, the presence and type of restrictions applied, etc. The article will use a combination of the quadratic error function and the Newton-Gauss method — a first-order quasi-algorithm Newtonian type. This algorithm has fairly good convergence in most cases.

The iteration step in the Newton-Gauss method is determined by the relation:,

$d = - (J ^ TJ) ^ {- 1} J ^ T r$ where J is the Jacobi matrix, r is the column vector of the residues $r_i = y_i-f (w; x_i)$ .

This formula logically consists of two parts: approximation of the Hessian matrix $H \ approx (J ^ TJ)$ and approximation of the gradient $\ nabla f \ approx J ^ T r$ .

The number of rows of the Jacobi matrix is determined by the number of training examples n , the number of columns by the size of the weight vector m . As shown in [1], the approximation of the Hessian matrix can be calculated by reading two rows of the matrix of the Jacobi matrix and multiplying them. The resulting $n ^ 2$ matrices in size $m \ times m$ can only be added up. The proposed permutation of operations does not change the overall computational complexity, but allows us not to load the entire Jacobi matrix into the memory and conduct the calculation in parallel. Similarly, the gradient approximation is calculated only n vectors of length m are subject to addition. The inversion of the obtained Hessian matrix is not very difficult in view of the relatively small size. To ensure the convergence of the algorithm, it is necessary to monitor the positive definiteness of the calculated matrix $(J ^ TJ)$ , which is realized by calculating the eigenvalues and vectors.

In articles [2, 3], a general scheme was proposed for applying the approach described above for Apache Spark. The only drawback of these works, in my opinion, is the lack of a clear link with the existing Spark ML API. We will try to fill this gap in the next section.

Spark ML API implementation

For the successful implementation of non-linear models, we need to understand the structure and purpose of the base classes. The machine learning API in the Spark system has two versions: 1.x is located inside the mllib package , 2.x in the ml package . The documentation for the Spark ML module [4] states that the transition from the API version 1.x to version 2.x is aimed at providing the possibility of embedding in chains (Pipelines) and working with a typed DataFrame structure. In our example, we will use a newer class structure, but if necessary, it can quite easily be implemented under the old structure.

Important Spark API Classes

class org.apache.spark.ml.feature.Instance describes an instance of a training example, including a real label, an example weight, and a vector of feature values;
org.apache.spark.ml.regression. {Regressor, RegressionModel} are the key classes that you want to extend. The first is a builder of the model, and the second is an already trained model.
Using the trait we introduced, org.apache.spark.ml.regression.NonLinearFunction , a contract for a nonlinear function is defined for which the goal is to select a vector of weights. There are only 3 methods in the contract: eval returns the value of the function at the point, grad - the gradient vector, dim - the dimension of the model (length of the weight vector).
The Breeze linear algebra library [5] takes over operations with matrices and has ready-made implementations of optimization functions. To use the Newton-Gauss algorithm or another first-order algorithm, we need to implement the breeze.optimize.DiffFunction trait in our loss function . The calculate method based on the transmitted coefficient vector should return a tuple of two values: the value of the penalty function and the vector of its gradient at the point.

Regressor class extension

The org.apache.spark.ml.regression.NonLinearRegression class extends the Regressor contract and, as a result of training, returns an instance of NonLinearRegressionModel . To create a nonlinear model, you need to specify a unique string identifier and the kernel function of the NonLinearFunction model. Among the optional parameters, one can list: the maximum number of training iterations, the initial approximation of the coefficient vector, and the required accuracy. Nonlinear functions often have many extrema and the choice of the initial approximation, based on a priori ideas about the behavior of a particular nuclear function, allows us to direct the search precisely to the global extremum region. It is worth noting the use of a ready-made implementation of the Broyden-Fletcher-Goldfarb-Shanno algorithm with limited memory consumption (LBFGS) from the Breeze library. The model training code is shown below.

Code org.apache.spark.ml.regression.NonLinearRegression # train

override protected def train(dataset: Dataset[_]): NonLinearRegressionModel = {
  // set instance weight to 1 if not defined the column
  val instanceWeightCol: Column = if (!isDefined(weightCol) || $(weightCol).isEmpty) lit(1.0) else col($(weightCol))
  val instances: RDD[Instance] = dataset.select(
    col($(labelCol)).cast(DoubleType), instanceWeightCol, col($(featuresCol))).rdd.map {
    case Row(label: Double, weight: Double, features: Vector) => Instance(label, weight, features)
  }
  // persists dataset if defined any storage level
  val handlePersistence = dataset.rdd.getStorageLevel == StorageLevel.NONE
  if (handlePersistence) instances.persist(StorageLevel.MEMORY_AND_DISK)
  val costFunc = new SquaresLossFunctionRdd(kernel, instances)
  val optimizer = new LBFGS[BDV[Double]]($(maxIter), 10, $(tol))
  // checks and assigns the initial coefficients
  val initial = {
    if (!isDefined(initCoeffs) || $(initCoeffs).length != kernel.dim) {
      if ($(initCoeffs).length != kernel.dim)
        logWarning(s"Provided initial coefficients vector (${$(initCoeffs)}) not corresponds with model dimensionality equals to ${kernel.dim}")
      BDV.zeros[Double](kernel.dim)
    } else
      BDV($(initCoeffs).clone())
  }
  val states = optimizer.iterations(new CachedDiffFunction[BDV[Double]](costFunc), initial)
  val (coefficients, objectiveHistory) = {
    val builder = mutable.ArrayBuilder.make[Double]
    var state: optimizer.State = null
    while (states.hasNext) {
      state = states.next
      builder += state.adjustedValue
    }
    // checks is method failed
    if (state == null) {
      val msg = s"${optimizer.getClass.getName} failed."
      logError(msg)
      throw new SparkException(msg)
    }
    (Vectors.dense(state.x.toArray.clone()).compressed, builder.result)
  }
  // unpersists the instances RDD
  if (handlePersistence) instances.unpersist()
  // produces the model and saves training summary
  val model = copyValues(new NonLinearRegressionModel(uid, kernel, coefficients))
  val (summaryModel: NonLinearRegressionModel, predictionColName: String) = model.findSummaryModelAndPredictionCol()
  val trainingSummary = new NonLinearRegressionSummary(
    summaryModel.transform(dataset), predictionColName, $(labelCol), $(featuresCol), objectiveHistory, model
  )
  model.setSummary(trainingSummary)
}

The given code of the train method can be divided into three parts: obtaining training examples from a data set; initiation of the penalty function and the search for the optimal solution; saving the vector of coefficients and learning outcomes in the model instance.

RegressionModel class extension

The implementation of the org.apache.spark.ml.regression.NonLinearRegressionModel class is pretty trivial. The predict method uses the kernel function to get the value at a point:

override protected def predict(features: Vector): Double = {
  kernel.eval(coefficients.asBreeze.toDenseVector, features.asBreeze.toDenseVector)
}

The only requirement of the Spark ML API that you should pay attention to is the requirement for serializability of the model. The desired behavior is ensured in the companion object by expanding the abstract classes org.apache.spark.ml.util. {MLReader, MLWriter} [6]. The state of the trained model consists of two parts: the coefficient vector and the kernel function. If everything has already been invented with the coefficient vector, then with the kernel function it’s a bit more complicated. It is not possible to save the kernel function directly to the DataFrame, but there are several alternative options.

For simplicity, the option was selected to binary serialize the function to a Base64 string. The disadvantages include the inaccessibility of a person reading the resulting string, as well as the need to support versioning of implementations.

A much more promising approach is to preserve the function in a symbolic form. This can be done in the image of objects of the class class formula of the stats package in the R language, for example, log (y) ~ a + log (x) . This method is more complicated than the first, but solves a number of problems: a human-readable representation of functions and the possibility of deserialization by different versions of parsers while maintaining backward compatibility. Significant complexity here is the development of a sufficiently flexible parser of symbolic expressions of functions.

Perhaps a useful improvement would be the ability to choose a step for numerically differentiating the core function. The complexity of saving the model is not significantly affected.

The implementation of the quadratic loss function

The last element necessary for training is the loss function. In our example, we use the quadratic loss function in the form of two realizations. In one, the training examples are specified in the form of a Breeze matrix, in the other, in the form of a RDD [Instance] Spark structure. The first implementation is easy to understand (it directly uses matrix expressions) and is suitable for small training sets. It serves as a test benchmark for us.

Code org.apache.spark.ml.regression.SquaresLossFunctionBreeze

package org.apache.spark.ml.regression
import breeze.linalg.{DenseMatrix => BDM, DenseVector => BDV}
/**
  * Breeze implementation for the squares loss function.
  *
  * @param fitmodel concrete model implementation
  * @param xydata   labeled data combined into the matrix:
  *                 the first n-th columns consist of feature values, (n+1)-th columns contains labels
  */
class SquaresLossFunctionBreeze(val fitmodel: NonLinearFunction, xydata: BDM[Double])
  extends SquaresLossFunction {
  /**
    * The number of instances.
    */
  val instanceCount: Int = xydata.rows
  /**
    * The number of features.
    */
  val featureCount: Int = xydata.cols - 1
  /**
    * Feature matrix.
    */
  val X: BDM[Double] = xydata(::, 0 until featureCount)
  /**
    * Label vector.
    */
  val y: BDV[Double] = xydata(::, featureCount)
  /**
    * The model dimensionality (the number of weights).
    *
    * @return dimensionality
    */
  override def dim: Int = fitmodel.dim
  /**
    * Evaluates loss function value and the gradient vector
    *
    * @param weights weights
    * @return (loss function value, gradient vector)
    */
  override def calculate(weights: BDV[Double]): (Double, BDV[Double]) = {
    val r: BDV[Double] = diff(weights)
    (0.5 * (r.t * r), gradient(weights))
  }
  /**
    * Calculates a positive definite approximation of the Hessian matrix.
    *
    * @param weights weights
    * @return Hessian matrix approximation
    */
  override def hessian(weights: BDV[Double]): BDM[Double] = {
    val J: BDM[Double] = jacobian(weights)
    posDef(J.t * J)
  }
  /**
    * Calculates the Jacobian matrix
    *
    * @param weights weights
    * @return the Jacobian
    */
  def jacobian(weights: BDV[Double]): BDM[Double] = {
    val gradData = (0 until instanceCount) map { i => fitmodel.grad(weights, X(i, ::).t).toArray }
    BDM(gradData: _*)
  }
  /**
    * Calculates the difference vector between the label and the approximated values.
    *
    * @param weights weights
    * @return difference vector
    */
  def diff(weights: BDV[Double]): BDV[Double] = {
    val diff = (0 until instanceCount) map (i => fitmodel.eval(weights, X(i, ::).t) - y(i))
    BDV(diff.toArray)
  }
  /**
    * Calculates the gradient vector
    *
    * @param weights weights
    * @return gradient vector
    */
  def gradient(weights: BDV[Double]): BDV[Double] = {
    val J: BDM[Double] = jacobian(weights)
    val r = diff(weights)
    2.0 * J.t * r
  }
}

The second option is designed to run in a distributed environment. For calculation, the RDD.treeAggregate function is used , which allows implementing the algorithm in the style of "Map-Reduce".

Code org.apache.spark.ml.regression.SquaresLossFunctionRdd

package org.apache.spark.ml.regression
import breeze.linalg.{DenseMatrix => BDM, DenseVector => BDV}
import org.apache.spark.broadcast.Broadcast
import org.apache.spark.ml.feature.Instance
import org.apache.spark.rdd.RDD
/**
  * Spark RDD implementation for the squares loss function.
  *
  * @param fitmodel concrete model implementation
  * @param xydata   RDD with instances
  */
class SquaresLossFunctionRdd(val fitmodel: NonLinearFunction, val xydata: RDD[Instance])
  extends SquaresLossFunction {
  /**
    * The model dimensionality (the number of weights).
    *
    * @return dimensionality
    */
  override def dim: Int = fitmodel.dim
  /**
    * Evaluates loss function value and the gradient vector
    *
    * @param weights weights
    * @return (loss function value, gradient vector)
    */
  override def calculate(weights: BDV[Double]): (Double, BDV[Double]) = {
    val bcW: Broadcast[BDV[Double]] = xydata.context.broadcast(weights)
    val (f: Double, grad: BDV[Double]) = xydata.treeAggregate((0.0, BDV.zeros[Double](dim)))(
      seqOp = (comb, item) => (comb, item) match {
        case ((loss, oldGrad), Instance(label, _, features)) =>
          val featuresBdv = features.asBreeze.toDenseVector
          val w: BDV[Double] = bcW.value
          val prediction = fitmodel.eval(w, featuresBdv)
          val addedLoss: Double = 0.5 * math.pow(label - prediction, 2)
          val addedGrad: BDV[Double] = 2.0 * (prediction - label) * fitmodel.grad(w, featuresBdv)
          (loss + addedLoss, oldGrad + addedGrad)
      },
      combOp = (comb1, comb2) => (comb1, comb2) match {
        case ((loss1, grad1: BDV[Double]), (loss2, grad2: BDV[Double])) => (loss1 + loss2, grad1 + grad2)
      })
    (f, grad)
  }
  /**
    * Calculates a positive definite approximation of the Hessian matrix.
    *
    * @param weights weights
    * @return Hessian matrix approximation
    */
  override def hessian(weights: BDV[Double]): BDM[Double] = {
    val bcW = xydata.context.broadcast(weights)
    val (hessian: BDM[Double]) = xydata.treeAggregate(new BDM[Double](dim, dim, Array.ofDim[Double](dim * dim)))(
      seqOp = (comb, item) => (comb, item) match {
        case ((oldHessian), Instance(_, _, features)) =>
          val grad = fitmodel.grad(bcW.value, features.asBreeze.toDenseVector)
          val subHessian: BDM[Double] = grad * grad.t
          oldHessian + subHessian
      },
      combOp = (comb1, comb2) => (comb1, comb2) match {
        case (hessian1, hessian2) => hessian1 + hessian2
      }
    )
    posDef(hessian)
  }
}

Project assembly

To simplify development and testing from the original Spark ML project, we borrowed pom.xml with a slight modification. We fix the version of Spark to one of the released ones, in our case 2.0.1 . You should pay attention to the inheritance of our POM file from org.apache.spark: spark-parent_2.11: 2.0.1 , which allows us not to re -arrange the configuration of Maven plugins.

To run tests that require SparkContext , we add org.apache.spark: spark-mllib_2.11: 2.0.1: test-jar : traits org.apache.spark.mllib.util.MLlibTestSparkContext , org.apache.spark in test dependencies .ml.util.TempDirectorywe implement in the corresponding test classes. Also useful for testing are extensions to the Suite classes from the org.apache.spark package that help you work with contexts, such as SparkFunSuite .

On the right of conclusion

There are several points that are not covered in this article, but their study seems extremely interesting:

use of a weighted training sample;
application of restrictions on the optimization domain, both soft (regularization) and hard (boundary conditions);
assessment of statistical indicators of the model (confidence intervals of coefficients, significance, etc.).

At the moment, I do not have enough information on the above aspects, but I will be grateful to all the shared sources.

The full code can be viewed on Github .

Full and comprehensive testing of this solution has yet to be carried out, so please treat the material as a concept and topic for discussion of improvements.

For comments and suggestions, it is preferable to use private messages, comments are better for discussions.

Thank you all for your attention.

Used materials

Tags:

Nonlinear Regression in Apache Spark. Do it yourself