Repository: CitrineInformatics/lolo (language: python)

Issue:
Expose multitask random forest in Lolopy
PLA-8731

Makes it so that the `RandomForestRegressor` model in Lolopy can be trained on 2-dimensional label data, `y`. The model then returns predictions as a 2-dimensional array. This brings its behavior more in line with scikit-learn's random forest regressor, which can also operate on multiple labels. The following changes were made to enable this feature:
* The `RandomForest` learner in Lolo now checks to see if the output data is a `Seq` (previously it only checked whether or not it was a `Double`). If it's a `Seq`, then it produces a `MultiTaskBagger` around a standardized `MultiTaskTreeLearner`.
* The top-level `fit` method of `BaseLoloLearner` inspects the shape of the training labels `y`, and if it's a 2-d array then it gets sent to Java as an array instead of as a vector
* The `predict` method of `BaseLoloRegressor` changes the way it expects to receive predictions and uncertainties from Java, depending on the number of labels (for >1 labels, it reshapes the stream into a 2-dimensional array).
* The `predict` method now has a `return_cov_matrix` option in addition to the `return_std` option. Only one can be requested. `return_cov_matrix` returns an array of covariance matrices for each prediction.
* `MultiTaskTree`'s interface is expanded to include arguments for `minLeafInstances`, `maxDepth`, and `numFeatures`. 
* Although the random forest interface also includes a `leafLearner` argument, that one is not piped through to `MultiTaskTree`. If it is specified, then we print a message warning the user that it is being ignored. At this time, nothing except a guess-the-mean learner is sensible.

There are a few other small changes:
* I decided to upgrade scikit-learn to the recently-released version 1, on the belief that stable releases are better than unstable releases.
* This requires bumping the minimum python version from 3.6 to 3.7. Version 3.7 has been available since mid-2018.
* Scikit-learn v1 deprecates the boston housing dataset, so I replaced its use in tests with the diabetes dataset.
* The uncertainty for a multitask bagged model has been changed from `Option[Seq[Seq[Option[Any]]]]` to `Option[Seq[Seq[Any]]]`. The former was too clunky and it also made it more difficult to send the uncertainties over the wire. So what do we do if the uncertainty for a given label isn't defined? I decided to use `NaN`. This does have precedent in Lolo, which uses `NaN` and `null` to indicate missing labels in training data.

Despite many of the changes to Lolopy being at an abstract level, this PR does not enable any multitask models in Lolopy besides `RandomForestRegressor`.
* The linear regressor model in Lolo only accepts a single label
* The Extra Random forest model in Lolo makes use of customized extra-random regression/classification splitters, and those are not generalized to multitask models
* Lolopy makes a distinction between regression and classification learners, and so even though Lolo can train a mixed regression-classification multitask tree, that doesn't map onto Lolopy
* `RandomForestClassifier` would be the easiest to generalized to a multitask model. It would require modifying the `BaseLoloClassifier`'s `predict` and `proba` methods to handle the more complex data shape. This is a good follow-up candidate, especially since it would bring Lolpy more in line with sklearn, but it didn't seem high priority enough to delay the release.


Discussion / hints:

Fix patch:
diff --git a/.travis.yml b/.travis.yml
index 7e66da4f..681c81ac 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -32,7 +32,7 @@ jobs:
         - ./miniconda.sh -b
         - export PATH=/home/travis/miniconda3/bin:$PATH
         - conda update --yes conda
-        - conda create --yes -n condaenv python=3.6
+        - conda create --yes -n condaenv python=3.7
         - conda install --yes -n condaenv pip
         - source activate condaenv
     install:
@@ -57,7 +57,7 @@ jobs:
       - ./miniconda.sh -b
       - export PATH=/home/travis/miniconda3/bin:$PATH
       - conda update --yes conda
-      - conda create --yes -n condaenv python=3.6
+      - conda create --yes -n condaenv python=3.7
       - conda install --yes -n condaenv pip
       - source activate condaenv
     install:
diff --git a/README.md b/README.md
index ec02f017..4d0876af 100644
--- a/README.md
+++ b/README.md
@@ -34,7 +34,7 @@ Lolo is on the central repository, and can be used by simply adding the followin
 <dependency>
     <groupId>io.citrine</groupId>
     <artifactId>lolo</artifactId>
-    <version>3.1.1</version>
+    <version>4.0.0</version>
 </dependency>
 ```
 Lolo provides higher level wrappers for common learner combinations.
diff --git a/RELEASE.md b/RELEASE.md
deleted file mode 100644
index 53f70304..00000000
--- a/RELEASE.md
+++ /dev/null
@@ -1,20 +0,0 @@
-# v3.0 is released!
-
-Lolo version 3.0 is released! Notably, this release introduces a shiny new interface for uncertainty quantification that provides new and improved granular uncertainty quantification (UQ). This release also introduces two new random forest features, and improves the accuracy of predictions and UQ on certain types of machine learning problems.
-
-### What’s New
-
-* Now you can distinguish estimated uncertainty in observations from uncertainty in the mean! You can query each of these with a new interface that exposes granular estimates of the [bias-variance-noise decomposition](https://en.wikipedia.org/wiki/Bias%E2%80%93variance_tradeoff). (#198)
-* A new `Splitter` interface allowing custom tree splitting protocols such as the new `BoltzmannSplitter` used to train [Boltzmann trees](https://www.youtube.com/watch?v=wWChMOkNlWk). (#95)
-* Support for [random rotation ensembles](http://www.jmlr.org/papers/v17/blaser16a.html), which is the technique of applying a different random rotation of feature space on each tree in a random forest. This can be helpful when a function to be learned aligns poorly to the coordinates of feature space. (#199 & #200)
-
-### Improvements
-
-* Bias-corrected jackknife variances are now rectified after summation instead of being rectified individually, improving the quality of importance scores. (#194)
-* The implementation of BaggedResult has been simplified. (#195)
-
-### Fixes
-
-* When training data contained identical inputs associated with more than one categorical label, or when trees were not grown to full depth, GuessTheMeanLearner, ClassificationTreeLearner, and RandomForest were biased toward predicting labels that appeared earlier in the training dataset.
-GuessTheMeanLearner now randomizes its tie-breaking choice of class label to predict. Consequently, GuessTheMeanLearner, ClassificationTreeLearner, and RandomForest are unbiased with respect to the order of training data and RandomForest performs substantially better on datasets with duplicated inputs. (#202)
-* Release 3.0.0 lacked a rescaling factor and a square root in getStdDevObs. Version 3.0.1 fixes this, improves test coverage, and introduces a Bessel correction.
diff --git a/python/README.md b/python/README.md
index 2bf31e3f..b82f30b8 100644
--- a/python/README.md
+++ b/python/README.md
@@ -19,9 +19,9 @@ Lolo is a Scala library that contains a variety of machine learning algorithms,
 
 ### Development 
 
-Lolopy requires Python >= 3.6, Java JDK >= 1.8, and sbt to be installed on your system when developing lolopy.  
+Lolopy requires Python >= 3.7, Java JDK >= 1.8, and sbt to be installed on your system when developing lolopy.
 
-Before developing `lolopy`, compile `lolo` on your system using Maven.
+Before developing `lolopy`, compile `lolo` on your system using sbt.
 We have provided a `Makefile` that contains the needed operations.
 To build and install `lolopy` call `make` in this directory.
 
diff --git a/python/lolopy/learners.py b/python/lolopy/learners.py
index 1882c68b..e55d41cf 100644
--- a/python/lolopy/learners.py
+++ b/python/lolopy/learners.py
@@ -3,7 +3,6 @@
 import numpy as np
 from lolopy.loloserver import get_java_gateway
 from lolopy.utils import send_feature_array, send_1D_array
-import random
 from sklearn.base import BaseEstimator, RegressorMixin, ClassifierMixin, is_regressor
 from sklearn.exceptions import NotFittedError
 
@@ -22,12 +21,12 @@ class BaseLoloLearner(BaseEstimator, metaclass=ABCMeta):
     In lolo, learners are not specific to a regression or classification problem and the type of problem is determined
     when fitting data is provided to the algorithm.
     In contrast, Scikit-learn learners for regression or classification problems are different classes.
-    We have implemented `BaseLoloRegressor` and `BaseLoloClassifier` abstract classes to make it easier for creating
+    We have implemented `BaseLoloRegressor` and `BaseLoloClassifier` abstract classes to make it easier to create
     a classification or regression version of a Lolo base class.
     The pattern for creating a scikit-learn compatible learner is to first implement the `_make_learner` and `__init__`
     operations in a special "Mixin" class that inherits from `BaseLoloLearner`, and then create a regression- or
     classification-specific class that inherits from both `BaseClassifier` or `BaseRegressor` and your new "Mixin".
-    See the RandomForest models as an example for this approach.
+    See the RandomForest models as an example of this approach.
     """
 
     def __init__(self):
@@ -35,6 +34,7 @@ def __init__(self):
 
         # Create a placeholder for the model
         self.model_ = None
+        self._num_outputs = None
         self._compress_level = 9
         self.feature_importances_ = None
         
@@ -70,6 +70,15 @@ def fit(self, X, y, weights=None):
         # Instantiate the JVM object
         learner = self._make_learner()
 
+        # Determine the number of outputs
+        y_shape = np.asarray(y).shape
+        if len(y_shape) == 1:
+            self._num_outputs = 1
+        elif len(y_shape) == 2:
+            self._num_outputs = y.shape[1]
+        else:
+            raise ValueError("Output array must be either 1- or 2-dimensional")
+
         # Convert all of the training data to Java arrays
         train_data, weights_java = self._convert_train_data(X, y, weights)
         assert train_data.length() == len(X), "Array copy failed"
@@ -123,15 +132,17 @@ def _convert_train_data(self, X, y, weights=None):
         if weights is None:
             weights = np.ones(len(y))
 
-        # Convert x, y, and w to float64 and int8 with native ordering
-
+        # Convert y and w to float64 or int32 with native ordering
         y = np.array(y, dtype=np.float64 if is_regressor(self) else np.int32)
         weights = np.array(weights, dtype=np.float64)
 
-
-        # Convert X and y to Java Objects
+        # Convert X, y, and w to Java Objects
         X_java = send_feature_array(self.gateway, X)
-        y_java = send_1D_array(self.gateway, y, is_regressor(self))
+        if self._num_outputs == 1:
+            y_java = send_1D_array(self.gateway, y, is_regressor(self))
+        else:
+            y_java = send_feature_array(self.gateway, y)
+
         assert y_java.length() == len(y) == len(X)
         w_java = send_1D_array(self.gateway, weights, True)
         assert w_java.length() == len(weights)
@@ -192,27 +203,75 @@ def _get_prediction_result(self, X):
 class BaseLoloRegressor(BaseLoloLearner, RegressorMixin):
     """Abstract class for models that produce regression models.
 
-    Implements the predict operation"""
+    As written, this allows for both single-task and multi-task models.
+    Implements the predict operation."""
+
+    def predict(self, X, return_std = False, return_cov_matrix = False):
+        """
+        Apply the model to a matrix of inputs, producing predictions and optionally some measure of uncertainty
 
-    def predict(self, X, return_std=False):
+        Args:
+            X (ndarray): Input array
+            return_std (bool): if True, return the standard deviations along with the predictions
+            return_cov_matrix (bool): If True, return the covariance matrix along with the predictions
+        Returns
+            Sequence of predictions OR
+            (Sequence of predictions, Sequence of standard deviations) OR
+            (Sequence of predictions, Sequence of covariance matrices).
+            Each prediction and standard deviation is a float (for single-output learners) or an array (for multi-output learners).
+            Each covariance matrix entry is a (# outputs x # outputs) matrix.
+        """
+        if return_std and return_cov_matrix:
+            raise ValueError("Only one of return_std or return_cov_matrix can be True")
         # Start the prediction process
         pred_result = self._get_prediction_result(X)
 
         # Pull out the expected values
-        y_pred_byte = self.gateway.jvm.io.citrine.lolo.util.LoloPyDataLoader.getRegressionExpected(pred_result)
-        y_pred = np.frombuffer(y_pred_byte, dtype='float')  # Lolo gives a byte array back
+        if self._num_outputs == 1:
+            y_pred_byte = self.gateway.jvm.io.citrine.lolo.util.LoloPyDataLoader.getRegressionExpected(pred_result)
+            y_pred = np.frombuffer(y_pred_byte, dtype='float')  # Lolo gives a byte array back
+        else:
+            y_pred_byte = self.gateway.jvm.io.citrine.lolo.util.LoloPyDataLoader.getMultiRegressionExpected(pred_result)
+            y_pred = np.frombuffer(y_pred_byte, dtype='float').reshape(-1, self._num_outputs)
 
-        # If desired, return the uncertainty too
         if return_std:
-            # TODO: This part fails on Windows because the NativeSystemBLAS is not found. Fix that
-            # TODO: This is only valid for regression models. Perhaps make a "LoloRegressor" class
-            y_std_bytes = self.gateway.jvm.io.citrine.lolo.util.LoloPyDataLoader.getRegressionUncertainty(pred_result)
-            y_std = np.frombuffer(y_std_bytes, 'float')
+            y_std = self._get_std(X, pred_result)
             return y_pred, y_std
 
+        if return_cov_matrix:
+            corr_matrix = self._get_corr_matrix(X, pred_result)
+            y_std = self._get_std(X, pred_result).reshape(-1, self._num_outputs)
+            sigma_sq_matrix = np.array([np.outer(y_std[i, :], y_std[i, :]) for i in range(len(X))])
+            # both sigma_squared and correlation matrices have size (# predictions, # outputs, # outputs).
+            # They are multiplied term-by-term to produce the covariance matrix.
+            cov_matrix = sigma_sq_matrix * corr_matrix
+            return y_pred, cov_matrix
+
         # Get the expected values
         return y_pred
 
+    def _get_std(self, X, pred_result):
+        # TODO: This part fails on Windows because the NativeSystemBLAS is not found. Fix that
+        if self._num_outputs == 1:
+            y_std_bytes = self.gateway.jvm.io.citrine.lolo.util.LoloPyDataLoader.getRegressionUncertainty(pred_result)
+            return np.frombuffer(y_std_bytes, 'float')
+        else:
+            y_std_bytes = self.gateway.jvm.io.citrine.lolo.util.LoloPyDataLoader.getMultiRegressionUncertainty(pred_result)
+            return np.frombuffer(y_std_bytes, 'float').reshape(-1, self._num_outputs)
+
+    def _get_corr_matrix(self, X, pred_result):
+        num_predictions = len(X)
+        corr_matrix = np.zeros((num_predictions, self._num_outputs, self._num_outputs))
+        idx = np.arange(self._num_outputs)
+        corr_matrix[:, idx, idx] = 1.0
+        for i in range(self._num_outputs - 1):
+            for j in range(i + 1, self._num_outputs):
+                rho_bytes = self.gateway.jvm.io.citrine.lolo.util.LoloPyDataLoader.getRegressionCorrelation(pred_result, i, j)
+                rho = np.frombuffer(rho_bytes, 'float')
+                corr_matrix[:, i, j] = rho
+                corr_matrix[:, j, i] = rho
+        return corr_matrix
+
 
 class BaseLoloClassifier(BaseLoloLearner, ClassifierMixin):
     """Base class for classification models
diff --git a/python/lolopy/version.py b/python/lolopy/version.py
index b736fa8e..854524dc 100644
--- a/python/lolopy/version.py
+++ b/python/lolopy/version.py
@@ -1,3 +1,3 @@
 # single source of truth for package version,
 # see https://packaging.python.org/en/latest/single_source_version/
-__version__ = "1.2.0"
+__version__ = "1.3.0"
diff --git a/python/requirements.txt b/python/requirements.txt
index cef58eab..b035d131 100644
--- a/python/requirements.txt
+++ b/python/requirements.txt
@@ -1,2 +1,2 @@
-scikit-learn==0.20.0
+scikit-learn==1.0.1
 py4j==0.10.8.1
diff --git a/src/main/scala/io/citrine/lolo/Learner.scala b/src/main/scala/io/citrine/lolo/Learner.scala
index b0bb7a63..3d9f68a2 100644
--- a/src/main/scala/io/citrine/lolo/Learner.scala
+++ b/src/main/scala/io/citrine/lolo/Learner.scala
@@ -25,3 +25,18 @@ trait Learner extends Serializable {
   }
 
 }
+
+/**
+  * A learner that trains on multiple labels, outputting a single model that makes predictions for all labels.
+  */
+trait MultiTaskLearner extends Serializable {
+
+  /**
+    * Train a model
+    *
+    * @param trainingData  to train on. Each entry is a tuple (vector of inputs, vector of labels)
+    * @param weights for the training rows, if applicable
+    * @return A training result that encompasses model(s) for all labels.
+    */
+  def train(trainingData: Seq[(Vector[Any], Vector[Any])], weights: Option[Seq[Double]] = None): MultiTaskTrainingResult
+}
diff --git a/src/main/scala/io/citrine/lolo/Model.scala b/src/main/scala/io/citrine/lolo/Model.scala
index 794e3915..5e95e288 100644
--- a/src/main/scala/io/citrine/lolo/Model.scala
+++ b/src/main/scala/io/citrine/lolo/Model.scala
@@ -30,3 +30,35 @@ trait Model[+T <: PredictionResult[Any]] extends Serializable {
              ): Option[DenseMatrix[Double]] = None
 
 }
+
+/** A model that predicts a sequence of values, corresponding to multiple labels. */
+trait MultiTaskModel extends Model[MultiTaskModelPredictionResult] {
+  /** The number of labels. Every prediction must have this length. */
+  val numLabels: Int
+
+  /** A boolean sequence indicating which labels are real-valued. Its length must be equal to `numLabels`. */
+  def getRealLabels: Seq[Boolean]
+
+  /** Individual models corresponding to each label */
+  def getModels: Seq[Model[PredictionResult[Any]]]
+}
+
+/**
+  * A set of parallel models, one for each label.
+  *
+  * @param models     sequence of models, one for each label
+  * @param realLabels boolean sequence indicating which labels are real-valued
+  */
+class ParallelModels(models: Seq[Model[PredictionResult[Any]]], realLabels: Seq[Boolean]) extends MultiTaskModel {
+  override val numLabels: Int = models.length
+
+  override def getRealLabels: Seq[Boolean] = realLabels
+
+  override def getModels: Seq[Model[PredictionResult[Any]]] = models
+
+  override def transform(inputs: Seq[Vector[Any]]) = new ParallelModelsPredictionResult(
+    models.map { model =>
+      model.transform(inputs).getExpected()
+    }.transpose
+  )
+}
diff --git a/src/main/scala/io/citrine/lolo/MultiTaskLearner.scala b/src/main/scala/io/citrine/lolo/MultiTaskLearner.scala
deleted file mode 100644
index 8a573faf..00000000
--- a/src/main/scala/io/citrine/lolo/MultiTaskLearner.scala
+++ /dev/null
@@ -1,17 +0,0 @@
-package io.citrine.lolo
-
-/**
-  * Created by maxhutch on 11/14/16.
-  */
-trait MultiTaskLearner extends Serializable {
-
-  /**
-    * Train a model
-    *
-    * @param inputs  to train on
-    * @param labels  sequence of sequences of labels
-    * @param weights for the training rows, if applicable
-    * @return a sequence of training results, one for each label
-    */
-  def train(inputs: Seq[Vector[Any]], labels: Seq[Seq[Any]], weights: Option[Seq[Double]] = None): Seq[TrainingResult]
-}
diff --git a/src/main/scala/io/citrine/lolo/PredictionResult.scala b/src/main/scala/io/citrine/lolo/PredictionResult.scala
index 2414f1c7..a554287f 100644
--- a/src/main/scala/io/citrine/lolo/PredictionResult.scala
+++ b/src/main/scala/io/citrine/lolo/PredictionResult.scala
@@ -162,13 +162,25 @@ trait RegressionResult extends PredictionResult[Double] {
   }
 }
 
-case class MultiResult[T](values: Seq[T]) extends PredictionResult[T] {
+/** Container for predictions made on multiple labels simultaneously. */
+trait MultiTaskModelPredictionResult extends PredictionResult[Seq[Any]] {
+
+  override def getUncertainty(observational: Boolean = true): Option[Seq[Seq[Any]]] = None
+
   /**
-    * Get the expected values for this prediction
+    * Get the correlation coefficients between the predictions made on two labels.
+    * Correlation coefficient is bounded between -1 and 1.
+    * If either index is out of bounds or does not correspond to a real-valued label, then this method must return None.
     *
-    * @return expected value of each prediction
+    * @param  i index of the first label
+    * @param  j index of the second label
+    * @param observational whether the uncertainty correlation should take observational noise into account
+    * @return optional sequence of correlation coefficients between specified labels for each prediction
     */
-  override def getExpected(): Seq[T] = values
+  def getUncertaintyCorrelation(i: Int, j: Int, observational: Boolean = true): Option[Seq[Double]] = None
+}
 
-  def append(other: this.type): MultiResult[T] = MultiResult(values ++ other.values)
+/** A container that holds the predictions of several parallel models for multiple labels. */
+class ParallelModelsPredictionResult(predictions: Seq[Seq[Any]]) extends MultiTaskModelPredictionResult {
+  override def getExpected(): Seq[Seq[Any]] = predictions
 }
diff --git a/src/main/scala/io/citrine/lolo/TrainingResult.scala b/src/main/scala/io/citrine/lolo/TrainingResult.scala
index 07b525c0..b0a50c90 100644
--- a/src/main/scala/io/citrine/lolo/TrainingResult.scala
+++ b/src/main/scala/io/citrine/lolo/TrainingResult.scala
@@ -1,7 +1,5 @@
 package io.citrine.lolo
 
-import breeze.linalg.DenseVector
-
 /**
   * Created by maxhutch on 12/4/16.
   */
@@ -35,3 +33,11 @@ trait TrainingResult extends Serializable {
     */
   def getPredictedVsActual(): Option[Seq[(Vector[Any], Any, Any)]] = None
 }
+
+trait MultiTaskTrainingResult extends TrainingResult {
+  override def getModel(): MultiTaskModel
+
+  def getModels(): Seq[Model[PredictionResult[Any]]]
+
+  override def getPredictedVsActual(): Option[Seq[(Vector[Any], Seq[Option[Any]], Seq[Option[Any]])]] = None
+}
\ No newline at end of file
diff --git a/src/main/scala/io/citrine/lolo/bags/BaggedResult.scala b/src/main/scala/io/citrine/lolo/bags/BaggedResult.scala
index 4821565b..616fbf41 100644
--- a/src/main/scala/io/citrine/lolo/bags/BaggedResult.scala
+++ b/src/main/scala/io/citrine/lolo/bags/BaggedResult.scala
@@ -1,7 +1,8 @@
 package io.citrine.lolo.bags
 
 import breeze.linalg.{DenseMatrix, DenseVector, norm}
-import io.citrine.lolo.{PredictionResult, RegressionResult}
+import io.citrine.lolo.stats.StatsUtils
+import io.citrine.lolo.{MultiTaskModelPredictionResult, ParallelModelsPredictionResult, PredictionResult, RegressionResult}
 import io.citrine.lolo.util.Async
 import org.slf4j.{Logger, LoggerFactory}
 
@@ -10,10 +11,15 @@ import org.slf4j.{Logger, LoggerFactory}
   *
   * This allows the implementation to depend on the number of simultaneous predictions, which has performance
   * implications.
+  * For background on the uncertainty calculation, see Wager, S.; Hastie, T and Efron, B. Confidence Intervals for
+  * Random Forests: The Jackknife and Infinitesimal Jackknife. Journal of Machine Learning Research 15 (2014).
   */
 trait BaggedResult[+T] extends PredictionResult[T] {
   def predictions: Seq[PredictionResult[T]]
 
+  /** The number of inputs that have been predicted on (NOT the number of bagged models). */
+  def numPredictions: Int
+
   /**
     * Average the gradients from the models in the ensemble
     *
@@ -35,27 +41,31 @@ trait BaggedResult[+T] extends PredictionResult[T] {
 
 
 /**
-  * Container with model-wise predictions and logic to compute variances and training row scores
+  * Container with model-wise predictions at a single input point.
+  * Assuming a single input allows for performance optimizations and more readable code.
+  * See [[MultiPredictionBaggedResult]] for a generic implementation.
   *
   * @param predictions for each constituent model
   * @param NibIn       the sample matrix as (N_models x N_training)
   * @param bias        model to use for estimating bias
   */
-case class BaggedSingleResult(
-                               predictions: Seq[PredictionResult[Double]],
-                               NibIn: Vector[Vector[Int]],
-                               bias: Option[Double] = None,
-                               rescale: Double = 1.0,
-                               disableBootstrap: Boolean = false
-                             ) extends BaggedResult[Double] with RegressionResult {
+case class SinglePredictionBaggedResult(
+                                         predictions: Seq[PredictionResult[Double]],
+                                         NibIn: Vector[Vector[Int]],
+                                         bias: Option[Double] = None,
+                                         rescale: Double = 1.0,
+                                         disableBootstrap: Boolean = false
+                                       ) extends BaggedResult[Double] with RegressionResult {
   private lazy val treePredictions: Array[Double] = predictions.map(_.getExpected().head).toArray
 
+  override def numPredictions: Int = 1
+
   /**
     * Return the ensemble average or maximum vote
     *
     * @return expected value of each prediction
     */
-  override def getExpected(): Seq[Double] = Seq(expected)
+  override def getExpected(): Seq[Double] = Seq(expected + bias.getOrElse(0.0))
 
   private lazy val expected = treePredictions.sum / treePredictions.length
   private lazy val treeVariance: Double = {
@@ -108,18 +118,17 @@ case class BaggedSingleResult(
   override def getImportanceScores(): Option[Seq[Seq[Double]]] = Some(Seq(singleScores))
 
   private lazy val singleScores: Vector[Double] = {
-    // Compute the Bessel-uncorrected variance of the ensemble of predicted values,
-    // and then divide by the size of the ensemble an extra time
-    val varT = treeVariance * (treePredictions.length - 1.0) / (treePredictions.length * treePredictions.length)
-
     // This will be more convenient later
     val nMat = NibIn.transpose
 
+    // Compute the Bessel-uncorrected variance of the ensemble of predicted values, then multiply by (n - 1) / (n * B)
+    // We later sum over the training data, introducing a factor of n and leaving us with the expected correction term
+    val correction = treeVariance * (treePredictions.length - 1.0) * (nMat.size - 1) / (treePredictions.length * treePredictions.length * nMat.size)
+
     // Loop over each of the training instances, computing its contribution to the uncertainty
     val trainingContributions = nMat.indices.toVector.map { idx =>
       // Pull the vector of the number of times this instance was used to train each tree
       val vecN = nMat(idx).toArray
-      val nTot = vecN.sum
 
       // Loop over the trees, computing the covariance for the IJ estimate and the predicted value of
       // the out-of-bag trees for the J(ackknife) estimate
@@ -128,7 +137,7 @@ case class BaggedSingleResult(
       var tNot: Double = 0.0
       var tNotCount: Int = 0
       vecN.indices.foreach { jdx =>
-        cov = cov + (vecN(jdx) - nTot) * (treePredictions(jdx) - expected)
+        cov = cov + (vecN(jdx) - 1) * (treePredictions(jdx) - expected)
 
         if (vecN(jdx) == 0) {
           tNot = tNot + treePredictions(jdx)
@@ -141,13 +150,10 @@ case class BaggedSingleResult(
       if (tNotCount > 0) {
         // Compute the Jackknife after bootstrap estimate
         val varJ = Math.pow(tNot / tNotCount - expected, 2.0) * (nMat.size - 1) / nMat.size
-        // Compute the sum of the corrections to the IJ and J estimates
-        val correction = Math.E * varT
-        // Averaged the correct IJ and J estimates
-        0.5 * (varJ + varIJ - correction)
+        // Averaged the corrected IJ and J estimates
+        0.5 * (varJ + varIJ - Math.E * correction)
       } else {
         // We can't compute the Jackknife after bootstrap estimate, so just correct the IJ estimate
-        val correction = varT
         varIJ - correction
       }
     }
@@ -166,6 +172,8 @@ case class BaggedClassificationResult(
   lazy val expected: Seq[Any] = expectedMatrix.map(ps => ps.groupBy(identity).maxBy(_._2.size)._1).seq
   lazy val uncertainty: Seq[Map[Any, Double]] = expectedMatrix.map(ps => ps.groupBy(identity).mapValues(_.size.toDouble / ps.size).toMap)
 
+  override def numPredictions: Int = expectedMatrix.length
+
   /**
    * Return the majority vote vote
    *
@@ -181,26 +189,26 @@ case class BaggedClassificationResult(
   *
   * These calculations are implemented using matrix arithmetic to make them more performant when the number
   * of predictions is large.  This obfuscates the algorithm significantly, however.  To see what is being computed,
-  * look at [[BaggedSingleResult]], which is more clear.  These two implementations are tested for consistency.
+  * look at [[SinglePredictionBaggedResult]], which is more clear.  These two implementations are tested for consistency.
   *
   * @param predictions for each constituent model
   * @param NibIn       the sample matrix as (N_models x N_training)
   * @param bias        model to use for estimating bias
   */
-case class BaggedMultiResult(
-                         predictions: Seq[PredictionResult[Double]],
-                         NibIn: Vector[Vector[Int]],
-                         bias: Option[Seq[Double]] = None,
-                         rescale: Double = 1.0,
-                         disableBootstrap: Boolean = false
-                       ) extends BaggedResult[Double] with RegressionResult {
+case class MultiPredictionBaggedResult(
+                                        predictions: Seq[PredictionResult[Double]],
+                                        NibIn: Vector[Vector[Int]],
+                                        bias: Option[Seq[Double]] = None,
+                                        rescale: Double = 1.0,
+                                        disableBootstrap: Boolean = false
+                                      ) extends BaggedResult[Double] with RegressionResult {
 
   /**
     * Return the ensemble average
     *
     * @return expected value of each prediction
     */
-  override def getExpected(): Seq[Double] = expected
+  override def getExpected(): Seq[Double] = expected.zip(biasCorrection).map(x => x._1 + x._2)
 
   override def getStdDevObs(): Option[Seq[Double]] = {
     if (disableBootstrap) {
@@ -249,32 +257,21 @@ case class BaggedMultiResult(
 
   override def getImportanceScores(): Option[Seq[Seq[Double]]] = Some(scores)
 
-  /* Subtract off 1 to make correlations easier; transpose to be prediction-wise */
-  lazy val Nib: Vector[Vector[Int]] = NibIn.transpose.map(_.map(_ - 1))
+  override def numPredictions: Int = expectedMatrix.length
+
+  /* transpose to be (# training) x (# models) */
+  lazy val Nib: Vector[Vector[Int]] = NibIn.transpose
 
   /* Make a matrix of the tree-wise predictions */
   lazy val expectedMatrix: Seq[Seq[Double]] = predictions.map(p => p.getExpected()).transpose
 
-  /* Extract the prediction by averaging for regression, taking the most popular response for classification */
+  /* Extract the prediction by averaging over trees and adding the bias correction. */
+  lazy val biasCorrection: Seq[Double] = bias.getOrElse(Seq.fill(expectedMatrix.length)(0))
   lazy val expected: Seq[Double] = expectedMatrix.map(ps => ps.sum / ps.size)
 
-  /* This matrix is used to compute the jackknife variance */
-  lazy val NibJMat = new DenseMatrix[Double](Nib.head.size, Nib.size,
-    Nib.flatMap { v =>
-      val itot = 1.0 / v.size
-      val icount = 1.0 / v.count(_ == -1.0)
-      v.map(n => if (n == -1) icount - itot else -itot)
-    }.toArray
-  )
-
-  /* This matrix is used to compute the IJ variance */
-  lazy val NibIJMat = new DenseMatrix[Double](Nib.head.size, Nib.size,
-    Nib.flatMap { v =>
-      val itot = 1.0 / v.size
-      val vtot = v.sum.toDouble / (v.size * v.size)
-      v.map(n => n * itot - vtot)
-    }.toArray
-  )
+  lazy val NibJMat = BaggedResult.getJackknifeAfterBootstrapMatrix(Nib)
+
+  lazy val NibIJMat = BaggedResult.getInfinitesimalJackknifeMatrix(Nib)
 
   /* This represents the variance of the estimate of the mean. */
   lazy val stdDevMean: Seq[Double] = variance(expected.asInstanceOf[Seq[Double]].toVector, expectedMatrix, NibJMat, NibIJMat).map{Math.sqrt}
@@ -339,13 +336,13 @@ case class BaggedMultiResult(
     val arg = IJMat2 + JMat2
     Async.canStop()
 
-    /* Avoid division in the loop */
-    val inverseSize = 1.0 / modelPredictions.head.size
+    /* Avoid division in the loop, calculate (n - 1) / (n * B^2) */
+    val prefactor = 1.0 / Math.pow(modelPredictions.head.size, 2.0) * (Nib.size - 1) / Nib.size
 
     modelPredictions.indices.map { i =>
       Async.canStop()
       /* Compute the first order bias correction for the variance estimators */
-      val correction = Math.pow(inverseSize * norm(predMat(::, i) - meanPrediction(i)), 2)
+      val correction = prefactor * Math.pow(norm(predMat(::, i) - meanPrediction(i)), 2)
 
       /* The correction is prediction dependent, so we need to operate on vectors */
       0.5 * (arg(::, i) - Math.E * correction)
@@ -395,10 +392,115 @@ case class BaggedMultiResult(
   }
 }
 
+/**
+  * Container with model-wise predictions for each label and the machinery to compute (co)variance.
+  *
+  * @param baggedPredictions  bagged prediction results for each label
+  * @param realLabels         a boolean sequence indicating which labels are real-valued
+  * @param NibIn              the sampling matrix as (# bags) x (# training)
+  */
+case class MultiTaskBaggedResult(
+                                  baggedPredictions: Seq[BaggedResult[Any]],
+                                  realLabels: Seq[Boolean],
+                                  NibIn: Vector[Vector[Int]]
+                                ) extends BaggedResult[Seq[Any]] with MultiTaskModelPredictionResult {
+
+  /* transpose to be (# training) x (# models) */
+  lazy val Nib: Vector[Vector[Int]] = NibIn.transpose
+
+  override def numPredictions: Int = baggedPredictions.head.numPredictions
+
+  override def getExpected(): Seq[Seq[Any]] = baggedPredictions.map(_.getExpected()).transpose
+
+  override def predictions: Seq[PredictionResult[Seq[Any]]] = baggedPredictions
+    .map(_.predictions.map(_.getExpected()))
+    .transpose
+    .map(x => new ParallelModelsPredictionResult(x.transpose))
+
+  // For each prediction, the uncertainty is a sequence of entries for each label. Missing uncertainty values are reported as NaN
+  override def getUncertainty(observational: Boolean = true): Option[Seq[Seq[Any]]] = {
+    Some(baggedPredictions.map { predictionResult =>
+      predictionResult.getUncertainty(observational) match {
+        case Some(value) => value
+        case None => Seq.fill(numPredictions)(Double.NaN)
+      }
+    }.transpose)
+  }
+
+  override def getUncertaintyCorrelation(i: Int, j: Int, observational: Boolean = true): Option[Seq[Double]] = {
+    (realLabels(i), realLabels(j)) match {
+      case (true, true) if i == j => Some(Seq.fill(numPredictions)(1.0))
+      case (true, true) =>
+        if (observational) {
+          Some(uncertaintyCorrelationObservational(i, j))
+        } else {
+          Some(uncertaintyCorrelationMean)
+        }
+      case _: Any => None
+    }
+  }
+
+  /** The uncertainty correlation of the observational distribution is the correlation coefficient calculated over the bootstrap ensemble predictions. */
+  private def uncertaintyCorrelationObservational(i: Int, j: Int): Seq[Double] = {
+    // make (# predictions) x (# bags) prediction matrices for each label
+    val baggedPredictionsI = baggedPredictions(i).predictions.map(_.getExpected()).transpose.asInstanceOf[Seq[Seq[Double]]]
+    val baggedPredictionsJ = baggedPredictions(j).predictions.map(_.getExpected()).transpose.asInstanceOf[Seq[Seq[Double]]]
+    baggedPredictionsI.zip(baggedPredictionsJ).map { case (bagsI, bagsJ) =>
+      StatsUtils.correlation(bagsI, bagsJ)
+    }
+  }
+
+  /**
+    * The uncertainty correlation of the mean distribution is 0.0. In theory it should be estimated using the jackknife,
+    * but in practice the jackknife performs poorly when estimating covariance, so we default to the trivial implementation for now.
+    */
+  private def uncertaintyCorrelationMean: Seq[Double] = Seq.fill(numPredictions)(0.0)
+
+}
+
 object BaggedResult {
 
   val logger: Logger = LoggerFactory.getLogger(getClass)
 
+  /**
+    * Generate a matrix that is useful for computing (co)variance via jackknife after bootstrap (JaB).
+    * The central term of the JaB calculation (Wager et. al. 2014, equation 6) is the difference between the
+    * out-of-bag prediction on a point and the mean prediction on that point. If this is written as a single sum over bags,
+    * then each point has a weight -1/B when it is in-bag and weight 1/|{N_{bi}=0}| - 1/B when it is out-of-bag (B is the number of bags).
+    * This matrix encodes those weights, so when it is multiplied by the (# bags) x (# predictions) prediction matrix
+    * we have a matrix of the \Delta terms from equation 6.
+    *
+    * @param Nib The (# training) x (# bags) matrix indicating how many times each training point is used in each bag
+    */
+  def getJackknifeAfterBootstrapMatrix(Nib: Vector[Vector[Int]]): DenseMatrix[Double] = {
+    new DenseMatrix[Double](Nib.head.size, Nib.size,
+      Nib.flatMap { v =>
+        val itot = 1.0 / v.size  // 1/B
+        val icount = 1.0 / v.count(_ == 0)  // 1/|{N_{bi}=0}|
+        v.map(n => if (n == 0) icount - itot else -itot)
+      }.toArray
+    )
+  }
+
+  /**
+    * Generate a matrix that is useful for computing (co)variance via infinitesimal jackknife (IJ).
+    * The central term of the IJ calculation (Wager et. al. 2014, equation 5) is the covariance between the number of
+    * times a training point appears in a bag and the prediction made by that bag.
+    * This matrix encodes (N - \bar{N})/B (B is the number of bags), so that when it is multiplied by the
+    *  (# bags) x (# predictions) prediction matrix, we have a matrix of the covariance terms from equation 5.
+    *
+    * @param Nib The (# training) x (# bags) matrix indicating how many times each training point is used in each bag
+    */
+  def getInfinitesimalJackknifeMatrix(Nib: Vector[Vector[Int]]): DenseMatrix[Double] = {
+    new DenseMatrix[Double](Nib.head.size, Nib.size,
+      Nib.flatMap { v =>
+        val itot = 1.0 / v.size  // 1/B
+        val vtot = v.sum.toDouble / (v.size * v.size)  // \bar{N} / B
+        v.map(n => n * itot - vtot)
+      }.toArray
+    )
+  }
+
   /**
    * Make sure the variance is non-negative
    *
diff --git a/src/main/scala/io/citrine/lolo/bags/Bagger.scala b/src/main/scala/io/citrine/lolo/bags/Bagger.scala
index 717152a0..9743dc1c 100644
--- a/src/main/scala/io/citrine/lolo/bags/Bagger.scala
+++ b/src/main/scala/io/citrine/lolo/bags/Bagger.scala
@@ -2,7 +2,7 @@ package io.citrine.lolo.bags
 
 import breeze.linalg.DenseMatrix
 import breeze.stats.distributions.{Poisson, Rand, RandBasis}
-import io.citrine.lolo.stats.metrics.ClassificationMetrics
+import io.citrine.lolo.stats.metrics.{ClassificationMetrics, RegressionMetrics}
 import io.citrine.lolo.util.{Async, InterruptibleExecutionContext}
 import io.citrine.lolo.{Learner, Model, PredictionResult, RegressionResult, TrainingResult}
 
@@ -186,8 +186,7 @@ class BaggedTrainingResult[+T : ClassTag](
                             biasModel: Option[Model[PredictionResult[T]]] = None,
                             rescale: Double = 1.0,
                             disableBootstrap: Boolean = false
-                          )
-  extends TrainingResult {
+                          ) extends TrainingResult {
 
   lazy val NibT = Nib.transpose
   lazy val model = new BaggedModel[T](models, Nib, useJackknife, biasModel, rescale, disableBootstrap)
@@ -211,10 +210,8 @@ class BaggedTrainingResult[+T : ClassTag](
   }
 
   lazy val loss: Double = rep match {
-    case x: Double => Math.sqrt(predictedVsActual.map(d => Math.pow(d._2.asInstanceOf[Double] - d._3.asInstanceOf[Double], 2)).sum / predictedVsActual.size)
-    case x: Any =>
-      val f1 = ClassificationMetrics.f1scores(predictedVsActual)
-      if (f1 > 0.0) 1.0 / f1 - 1.0 else Double.MaxValue
+    case _: Double => RegressionMetrics.RMSE(predictedVsActual.asInstanceOf[Seq[(Vector[Any], Double, Double)]])
+    case _: Any => ClassificationMetrics.loss(predictedVsActual)
   }
 
   /**
@@ -271,9 +268,9 @@ class BaggedModel[+T: ClassTag](
 
     val res = if (inputs.size == 1 && isRegression) {
       // In the special case of a single prediction on a real value, emit an optimized BaggedSingleResult
-      BaggedSingleResult(ensemblePredictions.map(_.asInstanceOf[PredictionResult[Double]]), Nib, bias.map(_.head), rescale, disableBootstrap)
+      SinglePredictionBaggedResult(ensemblePredictions.map(_.asInstanceOf[PredictionResult[Double]]), Nib, bias.map(_.head), rescale, disableBootstrap)
     } else if (isRegression) {
-      BaggedMultiResult(ensemblePredictions.map(_.asInstanceOf[PredictionResult[Double]]), Nib, bias, rescale, disableBootstrap)
+      MultiPredictionBaggedResult(ensemblePredictions.map(_.asInstanceOf[PredictionResult[Double]]), Nib, bias, rescale, disableBootstrap)
     } else {
       BaggedClassificationResult(ensemblePredictions)
     }
diff --git a/src/main/scala/io/citrine/lolo/bags/MultiTaskBagger.scala b/src/main/scala/io/citrine/lolo/bags/MultiTaskBagger.scala
index a708a323..fc674e33 100644
--- a/src/main/scala/io/citrine/lolo/bags/MultiTaskBagger.scala
+++ b/src/main/scala/io/citrine/lolo/bags/MultiTaskBagger.scala
@@ -2,6 +2,7 @@ package io.citrine.lolo.bags
 
 import breeze.stats.distributions.{Poisson, Rand, RandBasis}
 import io.citrine.lolo._
+import io.citrine.lolo.stats.metrics.{ClassificationMetrics, RegressionMetrics}
 import io.citrine.lolo.util.{Async, InterruptibleExecutionContext}
 
 import scala.collection.parallel.ExecutionContextTaskSupport
@@ -10,8 +11,12 @@ import scala.collection.parallel.immutable.{ParRange, ParSeq}
 /**
   * Create an ensemble of multi-task models
   *
-  * @param method  learner to train each model in the ensemble
-  * @param numBags number of models in the ensemble
+  * @param method                 learner to train each model in the ensemble
+  * @param numBags                number of models in the ensemble
+  * @param useJackknife           whether to enable jackknife uncertainty estimate
+  * @param biasLearner            learner to use for estimating bias
+  * @param uncertaintyCalibration whether to empirically recalibrate the predicted uncertainties
+  * @param randBasis              breeze RandBasis to use for generating breeze random numbers
   */
 case class MultiTaskBagger(
                             method: MultiTaskLearner,
@@ -22,103 +27,192 @@ case class MultiTaskBagger(
                             randBasis: RandBasis = Rand
                           ) extends MultiTaskLearner {
 
-  private def combineImportance(v1: Option[Vector[Double]], v2: Option[Vector[Double]]): Option[Vector[Double]] = {
+  override def train(trainingData: Seq[(Vector[Any], Vector[Any])], weights: Option[Seq[Double]] = None): MultiTaskBaggedTrainingResult = {
+    val (inputs, labels) = trainingData.unzip
+    val repInput = inputs.head
+    val repOutput = labels.head
+    /* Make sure the training data are the same size */
+    assert(inputs.forall(repInput.size == _.size))
+    assert(labels.forall(repOutput.size == _.size))
+    assert(trainingData.size >= Bagger.minimumTrainingSize, s"We need to have at least ${Bagger.minimumTrainingSize} rows, only ${trainingData.size} given")
+
+    // if numBags is non-positive, set # bags = # inputs
+    val actualBags = if (numBags > 0) numBags else trainingData.size
+
+    // Compute the number of instances of each training row in each training sample
+    val dist = new Poisson(1.0)(randBasis)
+    val Nib: Vector[Vector[Int]] = Vector.tabulate(actualBags) { _ =>
+      Vector.tabulate(trainingData.size) { _ =>
+        dist.draw()
+      }
+    }
+    val weightsActual = weights.getOrElse(Seq.fill(trainingData.size)(1.0))
+
+    val parIterator = new ParRange(Nib.indices)
+    parIterator.tasksupport = new ExecutionContextTaskSupport(InterruptibleExecutionContext())
+    val (models: ParSeq[MultiTaskModel], importances: ParSeq[Option[Vector[Double]]]) = parIterator.map { i =>
+      val meta = method.train(trainingData, Some(Nib(i).zip(weightsActual).map(p => p._1.toDouble * p._2)))
+      (meta.getModel(), meta.getFeatureImportance())
+    }.unzip
+
+    val averageImportance: Option[Vector[Double]] = importances
+      .reduce(combineImportance)
+      .map(_.map(_ / importances.size))
+
+    // Get bias model and rescale ratio for each label
+    val (biasModels, ratios) = Seq.tabulate(repOutput.length) { i =>
+      val thisLabelModels: ParSeq[Model[PredictionResult[Any]]] = models.map(_.getModels(i))
+      val isRegression = models.head.getRealLabels(i)
+      val thisTrainingData = trainingData.map{ case (inputs, outputs) => (inputs, outputs(i)) }
+      val helper = BaggerHelper(thisLabelModels, thisTrainingData, Nib, useJackknife, uncertaintyCalibration)
+      val biasModel = if (biasLearner.isDefined && isRegression) {
+        Async.canStop()
+        Some(biasLearner.get.train(helper.biasTraining).getModel().asInstanceOf[Model[PredictionResult[Double]]])
+      } else None
+      (biasModel, helper.ratio)
+    }.unzip
+
+    new MultiTaskBaggedTrainingResult(
+      models = models,
+      featureImportance = averageImportance,
+      Nib = Nib,
+      trainingData = trainingData,
+      useJackknife = useJackknife,
+      biasModels = biasModels,
+      rescaleRatios = ratios
+    )
+  }
+
+  /** Combine two optional feature importance vectors. */
+  def combineImportance(v1: Option[Vector[Double]], v2: Option[Vector[Double]]): Option[Vector[Double]] = {
     (v1, v2) match {
-      case (None, None) => None
       case (Some(v1: Vector[Double]), Some(v2: Vector[Double])) => Some(v1.zip(v2).map(p => p._1 + p._2))
       case _ => None
     }
   }
+}
 
-  /**
-    * Draw with replacement from the training data for each model
-    *
-    * @param inputs  to train on
-    * @param weights for the training rows, if applicable
-    * @return a model
-    */
-  override def train(inputs: Seq[Vector[Any]], labels: Seq[Seq[Any]], weights: Option[Seq[Double]] = None): Seq[BaggedTrainingResult[Any]] = {
-    /* Make sure the training data is the same size */
-    assert(inputs.forall(inputs.head.size == _.size))
-    assert(inputs.size > 8, s"We need to have at least 8 rows, only ${inputs.size} given")
-
-    /* Use unit weights if none are specified */
-    val weightsActual = weights.getOrElse(Seq.fill(inputs.size)(1.0))
-
-    /* Set default number of bags */
-    val actualBags = if (numBags > 0) {
-      numBags
-    } else {
-      inputs.size
+/**
+  * The result of training a bagger on a multi-label combined model.
+  *
+  * @param models             sequence of multi-models, one for each bag
+  * @param featureImportance  importance of input features
+  * @param Nib                matrix representing number of times each training datum appears in each bag
+  * @param trainingData       multi-label training data
+  * @param useJackknife       whether to enable jackknife uncertainty estimate
+  * @param biasModels         sequence of optional bias-correction models, one for each label
+  * @param rescaleRatios      sequence of uncertainty calibration ratios for each label
+  */
+class MultiTaskBaggedTrainingResult(
+                                     models: ParSeq[MultiTaskModel],
+                                     featureImportance: Option[Vector[Double]],
+                                     Nib: Vector[Vector[Int]],
+                                     trainingData: Seq[(Vector[Any], Seq[Any])],
+                                     useJackknife: Boolean,
+                                     biasModels: Seq[Option[Model[PredictionResult[Double]]]],
+                                     rescaleRatios: Seq[Double]
+                                   ) extends MultiTaskTrainingResult {
+
+  lazy val model = new MultiTaskBaggedModel(models, Nib, useJackknife, biasModels)
+
+  // Each entry is a tuple, (feature vector, seq of predicted labels, seq of actual labels).
+  // The labels are of type Option[Any] because a given training datum might not have a value for every single label.
+  // If the actual value for a label is None, then the corresponding prediction is recorded as None. The model could generate
+  // a prediction, but that's not useful in this context, since the point is to compare predictions with ground-truth values.
+  lazy val predictedVsActual: Seq[(Vector[Any], Seq[Option[Any]], Seq[Option[Any]])] =
+    trainingData.zip(Nib.transpose).flatMap { case ((features, labels), nb) =>
+      // Bagged models that were not trained on this input
+      val oob = models.zip(nb).filter(_._2 == 0).map(_._1)
+      if (oob.isEmpty) {
+        Seq()
+      } else {
+        // "Average" the predictions on each label over the out-of-bag models
+        val oobPredictions = oob.map(_.transform(Seq(features)).getExpected().head)
+        val predicted = oobPredictions.toVector.transpose.zipWithIndex.map {
+          case (predictions, labelIndex) if models.head.getRealLabels(labelIndex) =>
+            predictions.asInstanceOf[Seq[Double]].sum / predictions.size
+          case (predictions, _) => predictions.groupBy(identity).maxBy(_._2.size)._1
+        }
+        // Remove predictions for which the label was not specified
+        val (optionLabels, optionPredicted) = labels.zip(predicted).map {
+          case (l, _) if l == null || (l.isInstanceOf[Double] && l.asInstanceOf[Double].isNaN) => (None, None)
+          case (l, p) => (Some(l), Some(p))
+        }.unzip
+        Seq((features, optionPredicted, optionLabels))
+      }
     }
 
-    /* Compute the number of instances of each training row in each training sample */
-    val dist = new Poisson(1.0)(randBasis)
-    val Nib = Vector.tabulate(actualBags) { _ =>
-      Vector.tabulate(inputs.size) { _ =>
-        dist.draw()
+  lazy val loss = {
+    val allInputs = predictedVsActual.map(_._1)
+    val allPredicted: Seq[Seq[Option[Any]]] = predictedVsActual.map(_._2).transpose
+    val allActual: Seq[Seq[Option[Any]]] = predictedVsActual.map(_._3).transpose
+    (allPredicted, allActual, models.head.getRealLabels).zipped.map { case (labelPredicted, labelActual, isReal) =>
+      // Construct predicted-vs-actual for just this label, only keeping entries for which both predicted and actual are defined
+      val pva = (allInputs, labelPredicted, labelActual).zipped.flatMap {
+        case (input, Some(p), Some(a)) => Some((input, p, a))
+        case _ => None
+      }
+      if (isReal) {
+        RegressionMetrics.RMSE(pva.asInstanceOf[Seq[(Vector[Any], Double, Double)]])
+      } else {
+        ClassificationMetrics.loss(pva)
       }
     }
+  }.sum
 
-    /* Learn the actual models in parallel */
-    val parIterator = new ParRange(0 until actualBags)
-    parIterator.tasksupport = new ExecutionContextTaskSupport(InterruptibleExecutionContext())
-    val (models: ParSeq[Seq[Model[PredictionResult[Any]]]], importances: ParSeq[Seq[Option[Vector[Double]]]]) = parIterator.map { i =>
-      // Train the model
-      val meta = method.train(inputs.toVector, labels, Some(Nib(i).zip(weightsActual).map(p => p._1.toDouble * p._2)))
-      // Extract the model and feature importance from the TrainingResult
-      (meta.map(_.getModel()), meta.map(_.getFeatureImportance()))
-    }.unzip
+  override def getFeatureImportance(): Option[Vector[Double]] = featureImportance
 
-    /* Wrap the models in a BaggedModel object
-     *
-     * Transpose the models and importances so the bags are the inner index and the labels are the outer index.
-     * Foreach label, emit a BaggedTrainingResult
-     */
-    type arg = ((ParSeq[Model[PredictionResult[Any]]], Seq[Option[Vector[Double]]]), Int)
-    models.transpose.zip(importances.seq.transpose).zipWithIndex.map { xx: arg =>
-      xx match {
-        case ((m: ParSeq[Model[PredictionResult[Any]]], i: Seq[Option[Vector[Double]]]), k: Int) =>
-          val averageImportance: Option[Vector[Double]] = i.reduce {
-            combineImportance
-          }.map(_.map(_ / importances.size))
-          val trainingData = inputs.zip(labels(k))
-          val helper = BaggerHelper(m, trainingData, Nib, useJackknife, uncertaintyCalibration)
-
-          Async.canStop()
-          if (!helper.isRegression) {
-            new BaggedTrainingResult[Any](m, averageImportance, Nib, inputs.zip(labels(k)), useJackknife)
-          } else {
-            if (biasLearner.isEmpty) {
-              new BaggedTrainingResult[Double](m.asInstanceOf[ParSeq[Model[PredictionResult[Double]]]], averageImportance, Nib, trainingData, useJackknife, None, helper.ratio)
-            } else {
-              Async.canStop()
-              val baggedModel = new BaggedModel[Double](m.asInstanceOf[ParSeq[Model[PredictionResult[Double]]]], Nib, useJackknife, None, helper.ratio)
-              Async.canStop()
-              val baggedRes = baggedModel.transform(trainingData.map(_._1))
-              Async.canStop()
-              val foo = baggedRes.getUncertainty()
-              val biasTraining = trainingData.zip(
-                baggedRes.getExpected().zip(baggedRes.getUncertainty().get)
-              ).flatMap { case ((f, a), (p, u)) =>
-                if (a == null || (a.isInstanceOf[Double] && a.asInstanceOf[Double].isNaN)) {
-                  None
-                } else {
-                  // Math.E is only statistically correct.  It should be actualBags / Nib.transpose(i).count(_ == 0)
-                  // Or, better yet, filter the bags that don't include the training example
-                  val bias = Math.E * Math.max(Math.abs(p.asInstanceOf[Double] - a.asInstanceOf[Double]) - u.asInstanceOf[Double], 0.0)
-                  Some((f, bias))
-                }
-              }
-              Async.canStop()
-              val biasModel = biasLearner.get.train(helper.biasTraining).getModel().asInstanceOf[Model[PredictionResult[Double]]]
-              Async.canStop()
-
-              new BaggedTrainingResult[Double](m.asInstanceOf[ParSeq[Model[PredictionResult[Double]]]], averageImportance, Nib, trainingData, useJackknife, Some(biasModel), helper.ratio)
-            }
-          }
-      }
+  override def getModel(): MultiTaskBaggedModel = model
+
+  override def getModels(): Seq[Model[PredictionResult[Any]]] = {
+    val realLabels: Seq[Boolean] = models.head.getRealLabels
+    realLabels.zipWithIndex.map { case (isReal: Boolean, i: Int) =>
+      val thisLabelModels = models.map(_.getModels(i))
+      if (isReal) {
+        new BaggedModel[Double](thisLabelModels.asInstanceOf[ParSeq[Model[PredictionResult[Double]]]], Nib, useJackknife, biasModels(i), rescaleRatios(i))
+      } else {
+        new BaggedModel[Any](thisLabelModels.asInstanceOf[ParSeq[Model[PredictionResult[Any]]]], Nib, useJackknife, biasModels(i), rescaleRatios(i))
       }
-      .seq
+    }
+  }
+
+  override def getPredictedVsActual(): Option[Seq[(Vector[Any], Seq[Option[Any]], Seq[Option[Any]])]] = Some(predictedVsActual)
+
+  override def getLoss(): Option[Double] = {
+    if (predictedVsActual.nonEmpty) Some(loss) else None
   }
 }
+
+/**
+  * Container holding a parallel sequence of models, each of which predicts on multiple labels.
+  *
+  * @param models       sequence of multi-models, one for each bag
+  * @param Nib          matrix representing number of times each training datum appears in each bag
+  * @param useJackknife whether to enable jackknife uncertainty estimate
+  * @param biasModels   sequence of optional bias-correction models, one for each label
+  */
+class MultiTaskBaggedModel(
+                            models: ParSeq[MultiTaskModel],
+                            Nib: Vector[Vector[Int]],
+                            useJackknife: Boolean,
+                            biasModels: Seq[Option[Model[PredictionResult[Double]]]]
+                          ) extends MultiTaskModel {
+
+  lazy val groupedModels: Seq[BaggedModel[Any]] = Seq.tabulate(numLabels) { i =>
+    val thisLabelsModels = models.map(_.getModels(i))
+    if (getRealLabels(i)) {
+      new BaggedModel[Double](thisLabelsModels.asInstanceOf[ParSeq[Model[PredictionResult[Double]]]], Nib, useJackknife, biasModels(i))
+    } else {
+      new BaggedModel(thisLabelsModels, Nib, useJackknife, biasModels(i))
+    }
+  }
+
+  override def transform(inputs: Seq[Vector[Any]]) =
+    MultiTaskBaggedResult(groupedModels.map(_.transform(inputs)), getRealLabels, Nib)
+
+  override val numLabels: Int = models.head.numLabels
+
+  override def getRealLabels: Seq[Boolean] = models.head.getRealLabels
+
+  override def getModels: Seq[BaggedModel[Any]] = groupedModels
+}
diff --git a/src/main/scala/io/citrine/lolo/learners/ExtraRandomTrees.scala b/src/main/scala/io/citrine/lolo/learners/ExtraRandomTrees.scala
index b7f3424b..82adce13 100644
--- a/src/main/scala/io/citrine/lolo/learners/ExtraRandomTrees.scala
+++ b/src/main/scala/io/citrine/lolo/learners/ExtraRandomTrees.scala
@@ -52,23 +52,12 @@ case class ExtraRandomTrees(
     */
   override def train(trainingData: Seq[(Vector[Any], Any)], weights: Option[Seq[Double]]): TrainingResult = {
     val breezeRandBasis: RandBasis = new RandBasis(new MersenneTwister(rng.nextLong()))
+    val repOutput = trainingData.head._2
+    val isRegression = repOutput.isInstanceOf[Double]
+    val numFeatures: Int = RandomForest.getNumFeatures(subsetStrategy, trainingData.head._1.size, isRegression)
 
-    trainingData.head._2 match {
+    repOutput match {
       case _: Double => {
-        val numFeatures: Int = subsetStrategy match {
-          case x: String => {
-            x match {
-              case "auto" => trainingData.head._1.size
-              case "sqrt" => Math.ceil(Math.sqrt(trainingData.head._1.size)).toInt
-              case "log2" => Math.ceil(Math.log(trainingData.head._1.size) / Math.log(2)).toInt
-              case x: String =>
-                println(s"Unrecognized subsetStrategy ${x}; using auto")
-                trainingData.head._1.size
-            }
-          }
-          case x: Int => x
-          case x: Double => (trainingData.head._1.size * x).toInt
-        }
         val DTLearner = RegressionTreeLearner(
           leafLearner = leafLearner,
           numFeatures = numFeatures,
@@ -88,20 +77,6 @@ case class ExtraRandomTrees(
         bagger.train(trainingData, weights)
       }
       case _: Any => {
-        val numFeatures: Int = subsetStrategy match {
-          case x: String => {
-            x match {
-              case "auto" => Math.ceil(Math.sqrt(trainingData.head._1.size)).toInt
-              case "sqrt" => Math.ceil(Math.sqrt(trainingData.head._1.size)).toInt
-              case "log2" => Math.ceil(Math.log(trainingData.head._1.size) / Math.log(2)).toInt
-              case x: String =>
-                println(s"Unrecognized subsetStrategy ${x}; using auto")
-                trainingData.head._1.size
-            }
-          }
-          case x: Int => x
-          case x: Double => (trainingData.head._1.size * x).toInt
-        }
         val DTLearner = ClassificationTreeLearner(
           leafLearner = leafLearner,
           numFeatures = numFeatures,
diff --git a/src/main/scala/io/citrine/lolo/learners/RandomForest.scala b/src/main/scala/io/citrine/lolo/learners/RandomForest.scala
index 68faeecb..24814531 100644
--- a/src/main/scala/io/citrine/lolo/learners/RandomForest.scala
+++ b/src/main/scala/io/citrine/lolo/learners/RandomForest.scala
@@ -1,9 +1,11 @@
 package io.citrine.lolo.learners
 
 import breeze.stats.distributions.{RandBasis, ThreadLocalRandomGenerator}
-import io.citrine.lolo.bags.Bagger
-import io.citrine.lolo.transformers.FeatureRotator
+import io.citrine.lolo.bags.{Bagger, MultiTaskBagger}
+import io.citrine.lolo.linear.GuessTheMeanLearner
+import io.citrine.lolo.transformers.{FeatureRotator, MultiTaskFeatureRotator, MultiTaskStandardizer}
 import io.citrine.lolo.trees.classification.ClassificationTreeLearner
+import io.citrine.lolo.trees.multitask.MultiTaskTreeLearner
 import io.citrine.lolo.trees.regression.RegressionTreeLearner
 import io.citrine.lolo.trees.splits.{ClassificationSplitter, RegressionSplitter}
 import io.citrine.lolo.{Learner, TrainingResult}
@@ -24,7 +26,7 @@ import scala.util.Random
   * @param minLeafInstances minimum number of instances per leave in each tree
   * @param maxDepth       maximum depth of each tree in the forest (default: unlimited)
   * @param uncertaintyCalibration whether to empirically recalibrate the predicted uncertainties (default: false)
-  * @param randomizePivotLocation whether generate splits randomly between the data points (default: false)
+  * @param randomizePivotLocation whether to generate splits randomly between the data points (default: false)
   * @param randomlyRotateFeatures whether to randomly rotate real features for each tree in the forest (default: false)
   * @param rng            random number generator to use for stochastic functionality
   */
@@ -54,24 +56,13 @@ case class RandomForest(
     */
   override def train(trainingData: Seq[(Vector[Any], Any)], weights: Option[Seq[Double]]): TrainingResult = {
     val breezeRandBasis: RandBasis = new RandBasis(new ThreadLocalRandomGenerator(new MersenneTwister(rng.nextLong())))
+    val repOutput = trainingData.head._2
+    val isRegression = repOutput.isInstanceOf[Double] ||
+      (repOutput.isInstanceOf[Seq[Any]] && repOutput.asInstanceOf[Seq[Any]].exists(_.isInstanceOf[Double]))
+    val numFeatures: Int = RandomForest.getNumFeatures(subsetStrategy, trainingData.head._1.size, isRegression)
 
-    trainingData.head._2 match {
+    repOutput match {
       case _: Double =>
-        val numFeatures: Int = subsetStrategy match {
-          case x: String =>
-            x match {
-              case "auto" => trainingData.head._1.size
-              case "sqrt" => Math.ceil(Math.sqrt(trainingData.head._1.size)).toInt
-              case "log2" => Math.ceil(Math.log(trainingData.head._1.size) / Math.log(2)).toInt
-              case x: String =>
-                println(s"Unrecognized subsetStrategy ${x}; using auto")
-                trainingData.head._1.size
-            }
-          case x: Int =>
-            x
-          case x: Double =>
-            (trainingData.head._1.size * x).toInt
-        }
         val DTLearner = RegressionTreeLearner(
           leafLearner = leafLearner,
           numFeatures = numFeatures,
@@ -81,7 +72,8 @@ case class RandomForest(
           rng = rng
         )
 
-        val bagger = Bagger(if (randomlyRotateFeatures) FeatureRotator(DTLearner) else DTLearner,
+        val bagger = Bagger(
+          if (randomlyRotateFeatures) FeatureRotator(DTLearner) else DTLearner,
           numBags = numTrees,
           useJackknife = useJackknife,
           biasLearner = biasLearner,
@@ -89,22 +81,27 @@ case class RandomForest(
           randBasis = breezeRandBasis
         )
         bagger.train(trainingData, weights)
-      case _: Any =>
-        val numFeatures: Int = subsetStrategy match {
-          case x: String =>
-            x match {
-              case "auto" => Math.ceil(Math.sqrt(trainingData.head._1.size)).toInt
-              case "sqrt" => Math.ceil(Math.sqrt(trainingData.head._1.size)).toInt
-              case "log2" => Math.ceil(Math.log(trainingData.head._1.size) / Math.log(2)).toInt
-              case x: String =>
-                println(s"Unrecognized subsetStrategy ${x}; using auto")
-                Math.sqrt(trainingData.head._1.size).toInt
-            }
-          case x: Int =>
-            x
-          case x: Double =>
-            (trainingData.head._1.size * x).toInt
+      case _: Seq[Any] =>
+        if (leafLearner.isDefined && !leafLearner.get.isInstanceOf[GuessTheMeanLearner]) {
+          throw new IllegalArgumentException("Multitask random forest does not support leaf learners other than guess-the-mean")
         }
+        val DTLearner = new MultiTaskStandardizer(MultiTaskTreeLearner(
+          numFeatures = numFeatures,
+          maxDepth = maxDepth,
+          minLeafInstances = minLeafInstances,
+          randomizePivotLocation = randomizePivotLocation,
+          rng = rng
+        ))
+        val bagger = MultiTaskBagger(
+           if (randomlyRotateFeatures) MultiTaskFeatureRotator(DTLearner) else DTLearner,
+          numBags = numTrees,
+          useJackknife = useJackknife,
+          biasLearner = biasLearner,
+          uncertaintyCalibration = uncertaintyCalibration,
+          randBasis = breezeRandBasis
+        )
+        bagger.train(trainingData.asInstanceOf[Seq[(Vector[Any], Vector[Any])]], weights)
+      case _: Any =>
         val DTLearner = ClassificationTreeLearner(
           leafLearner = leafLearner,
           numFeatures = numFeatures,
@@ -122,3 +119,26 @@ case class RandomForest(
     }
   }
 }
+
+object RandomForest {
+
+  def getNumFeatures(subsetStrategy: Any, numTrainingFeatures: Int, isRegression: Boolean): Int = {
+    subsetStrategy match {
+      case x: String =>
+        x match {
+          case "auto" =>
+            if (isRegression) numTrainingFeatures else Math.ceil(Math.sqrt(numTrainingFeatures)).toInt
+          case "sqrt" => Math.ceil(Math.sqrt(numTrainingFeatures)).toInt
+          case "log2" => Math.ceil(Math.log(numTrainingFeatures) / Math.log(2)).toInt
+          case x: String =>
+            println(s"Unrecognized subsetStrategy $x; using auto")
+            getNumFeatures("auto", numTrainingFeatures, isRegression)
+        }
+      case x: Int =>
+        x
+      case x: Double =>
+        (numTrainingFeatures * x).toInt
+    }
+  }
+
+}
\ No newline at end of file
diff --git a/src/main/scala/io/citrine/lolo/stats/StatsUtils.scala b/src/main/scala/io/citrine/lolo/stats/StatsUtils.scala
new file mode 100644
index 00000000..9b308737
--- /dev/null
+++ b/src/main/scala/io/citrine/lolo/stats/StatsUtils.scala
@@ -0,0 +1,41 @@
+package io.citrine.lolo.stats
+
+object StatsUtils {
+
+  /** Compute the mean of a (weighted) vector, X */
+  def mean(X: Seq[Double], weights: Option[Seq[Double]] = None): Double = {
+    val actualWeights = weights.getOrElse(Seq.fill(X.length)(1.0))
+    val totalWeight = actualWeights.sum
+    require(totalWeight > 0.0, s"total weight must be positive, instead got $totalWeight")
+    X.zip(actualWeights).map { case (x, w) => x * w }.sum / totalWeight
+  }
+
+  /** Compute the variance of a (weighted) vector, X, with dof degrees of freedom. */
+  def variance(X: Seq[Double], weights: Option[Seq[Double]] = None, dof: Int = 0): Double = {
+    val actualWeights = weights.getOrElse(Seq.fill(X.length)(1.0))
+    val totalWeight = actualWeights.sum
+    require(dof >= 0, s"degrees of freedom must be non-negative, instead got $dof")
+    require(totalWeight > dof, s"Cannot compute variance on a sequence of weight $totalWeight with $dof degrees of freedom.")
+    val mu = mean(X, Some(actualWeights))
+    X.zip(actualWeights).map { case (x, w) => math.pow(x - mu, 2.0) * w }.sum / (totalWeight - dof)
+  }
+
+  /** Compute the (weighted) covariance between two vectors, X and Y, of the same length. */
+  def covariance(X: Seq[Double], Y: Seq[Double], weights: Option[Seq[Double]] = None): Double = {
+    require(X.length == Y.length, s"Cannot compute covariance between sequences of different lengths (lengths are ${X.length} and ${Y.length}).")
+    val actualWeights = weights.getOrElse(Seq.fill(X.length)(1.0))
+    val totalWeight = actualWeights.sum
+    val muX = mean(X, Some(actualWeights))
+    val muY = mean(Y, Some(actualWeights))
+    (X, Y, actualWeights).zipped.map { case (x, y, w) => (x - muX) * (y - muY) * w }.sum / totalWeight
+  }
+
+  /** Compute the (weighted) correlation coefficient between two vectors, X and Y, of the same length. */
+  def correlation(X: Seq[Double], Y: Seq[Double], weights: Option[Seq[Double]] = None): Double = {
+    val actualWeights = weights.getOrElse(Seq.fill(X.length)(1.0))
+    val sigma2X = variance(X, Some(actualWeights), dof = 0)
+    val sigma2Y = variance(Y, Some(actualWeights), dof = 0)
+    if (sigma2X == 0 || sigma2Y == 0) return 0.0
+    covariance(X, Y, Some(actualWeights)) / math.sqrt(sigma2X * sigma2Y)
+  }
+}
diff --git a/src/main/scala/io/citrine/lolo/stats/metrics/ClassificationMetrics.scala b/src/main/scala/io/citrine/lolo/stats/metrics/ClassificationMetrics.scala
index 0d8c469b..516f3a12 100644
--- a/src/main/scala/io/citrine/lolo/stats/metrics/ClassificationMetrics.scala
+++ b/src/main/scala/io/citrine/lolo/stats/metrics/ClassificationMetrics.scala
@@ -38,4 +38,10 @@ object ClassificationMetrics {
   def f1scores(predicted: Seq[Any], actual: Seq[Any]): Double = {
     f1scores(predicted.zip(actual).map(x => (Vector(), x._1, x._2)))
   }
+
+  /** Compute the loss on a scale from 0 (perfect model, f1 = 1.0) to infinity (f1 = 0, all predictions are incorrect) */
+  def loss(predictedVsActual: Seq[(Vector[Any], Any, Any)]): Double = {
+    val f1 = f1scores(predictedVsActual)
+    if (f1 > 0.0) 1.0 / f1 - 1.0 else Double.MaxValue
+  }
 }
diff --git a/src/main/scala/io/citrine/lolo/stats/metrics/RegressionMetrics.scala b/src/main/scala/io/citrine/lolo/stats/metrics/RegressionMetrics.scala
new file mode 100644
index 00000000..cffcf85f
--- /dev/null
+++ b/src/main/scala/io/citrine/lolo/stats/metrics/RegressionMetrics.scala
@@ -0,0 +1,10 @@
+package io.citrine.lolo.stats.metrics
+
+object RegressionMetrics {
+
+  /** Compute the root mean squared error. */
+  def RMSE(predictedVsActual: Seq[(Vector[Any], Double, Double)]): Double = {
+    math.sqrt(predictedVsActual.map { case (_, p, a) => math.pow(p - a, 2.0) }.sum / predictedVsActual.length)
+  }
+
+}
diff --git a/src/main/scala/io/citrine/lolo/transformers/FeatureRotator.scala b/src/main/scala/io/citrine/lolo/transformers/FeatureRotator.scala
index 3baeb833..9f88f45b 100644
--- a/src/main/scala/io/citrine/lolo/transformers/FeatureRotator.scala
+++ b/src/main/scala/io/citrine/lolo/transformers/FeatureRotator.scala
@@ -43,24 +43,21 @@ case class MultiTaskFeatureRotator(baseLearner: MultiTaskLearner) extends MultiT
   /**
     * Create linear transformations for continuous features and labels & pass data through to learner
     *
-    * @param inputs  to train on
-    * @param labels  sequence of sequences of labels
+    * @param trainingData  to train on
     * @param weights for the training rows, if applicable
     * @return a sequence of training results, one for each label
     */
   override def train(
-                     inputs: Seq[Vector[Any]],
-                     labels: Seq[Seq[Any]],
-                     weights: Option[Seq[Double]]
-                    ): Seq[RotatedFeatureTrainingResult] = {
+                      trainingData: Seq[(Vector[Any], Vector[Any])],
+                      weights: Option[Seq[Double]]
+                    ): MultiTaskRotatedFeatureTrainingResult = {
+    val inputs = trainingData.map(_._1)
+    val labels = trainingData.map(_._2)
     val featuresToRotate = FeatureRotator.getDoubleFeatures(inputs.head)
     val trans = FeatureRotator.getRandomRotation(inputs.head.length)
     val rotatedFeatures = FeatureRotator.applyRotation(inputs, featuresToRotate, trans)
-    val baseTrainingResult = baseLearner.train(rotatedFeatures, labels, weights)
-
-    baseTrainingResult.map { case (base) =>
-      RotatedFeatureTrainingResult(base, featuresToRotate, trans)
-    }
+    val baseTrainingResult = baseLearner.train(rotatedFeatures.zip(labels), weights)
+    MultiTaskRotatedFeatureTrainingResult(baseTrainingResult, featuresToRotate, trans)
   }
 }
 
@@ -97,6 +94,34 @@ case class RotatedFeatureTrainingResult(
   }
 }
 
+/**
+  * Training result bundling the base learner's multitask training result with the list of rotated features and the transformation
+  *
+  * @param baseTrainingResult training result to which to delegate prediction on rotated features
+  * @param rotatedFeatures indices of features to rotate
+  * @param trans matrix to apply to features
+  */
+case class MultiTaskRotatedFeatureTrainingResult(
+                                                  baseTrainingResult: MultiTaskTrainingResult,
+                                                  rotatedFeatures: IndexedSeq[Int],
+                                                  trans: DenseMatrix[Double]
+                                                ) extends MultiTaskTrainingResult {
+  override def getModel(): MultiTaskModel = new ParallelModels(getModels(), baseTrainingResult.getModel().getRealLabels)
+
+  override def getModels(): Seq[Model[PredictionResult[Any]]] = baseTrainingResult.getModels().map { model =>
+    RotatedFeatureModel(model, rotatedFeatures, trans)
+  }
+
+  override def getPredictedVsActual(): Option[Seq[(Vector[Any], Seq[Option[Any]], Seq[Option[Any]])]] = {
+    baseTrainingResult.getPredictedVsActual() match {
+      case None => None
+      case Some(predictedVsActual) => Some(predictedVsActual.map { case (inputs: Vector[Any], predicted: Seq[Option[Any]], actual: Seq[Option[Any]]) =>
+        (FeatureRotator.applyOneRotation(inputs, rotatedFeatures, trans), predicted, actual)
+      })
+    }
+  }
+}
+
 /**
   * Model bundling the base learner's model with the list of rotated features and the transformation
   *
diff --git a/src/main/scala/io/citrine/lolo/transformers/Standardizer.scala b/src/main/scala/io/citrine/lolo/transformers/Standardizer.scala
index 8a6c887f..68fe5abd 100644
--- a/src/main/scala/io/citrine/lolo/transformers/Standardizer.scala
+++ b/src/main/scala/io/citrine/lolo/transformers/Standardizer.scala
@@ -2,6 +2,15 @@ package io.citrine.lolo.transformers
 
 import io.citrine.lolo._
 
+case class Standardization(shift: Double, scale: Double) {
+  require(scale > 0 && scale < Double.PositiveInfinity)
+
+  def apply(x: Double): Double = (x - shift) / scale
+
+  def invert(x: Double): Double = x * scale + shift
+
+}
+
 /**
   * Standardize the training data to zero mean and unit variance before feeding it into another learner
   *
@@ -20,18 +29,17 @@ case class Standardizer(baseLearner: Learner) extends Learner {
     * @return training result containing a model
     */
   override def train(trainingData: Seq[(Vector[Any], Any)], weights: Option[Seq[Double]]): StandardizerTrainingResult = {
-    val rep = trainingData.head._1
     val inputTrans = Standardizer.getMultiStandardization(trainingData.map(_._1))
     val outputTrans = trainingData.head._2 match {
-      case x: Double => Some(Standardizer.getStandardization(trainingData.map(_._2.asInstanceOf[Double])))
-      case x: Any => None
+      case _: Double => Some(Standardizer.getStandardization(trainingData.map(_._2.asInstanceOf[Double])))
+      case _: Any => None
     }
 
     val (inputs, labels) = trainingData.unzip
     val standardTrainingData = Standardizer.applyStandardization(inputs, inputTrans).zip(Standardizer.applyStandardization(labels, outputTrans))
     val baseTrainingResult = baseLearner.train(standardTrainingData, weights)
 
-    new StandardizerTrainingResult(baseTrainingResult, Seq(outputTrans) ++ inputTrans)
+    new StandardizerTrainingResult(baseTrainingResult, outputTrans, inputTrans)
   }
 }
 
@@ -40,61 +48,115 @@ class MultiTaskStandardizer(baseLearner: MultiTaskLearner) extends MultiTaskLear
   /**
     * Train a model
     *
-    * @param inputs  to train on
-    * @param labels  sequence of sequences of labels
+    * @param trainingData  to train on
     * @param weights for the training rows, if applicable
     * @return a sequence of training results, one for each label
     */
-  override def train(inputs: Seq[Vector[Any]], labels: Seq[Seq[Any]], weights: Option[Seq[Double]]): Seq[TrainingResult] = {
+  override def train(trainingData: Seq[(Vector[Any], Vector[Any])], weights: Option[Seq[Double]]): MultiTaskStandardizerTrainingResult = {
+    val (inputs, labels) = trainingData.unzip
+    val labelsTransposed = labels.transpose.toVector
+    val repOutput = labels.head
     val inputTrans = Standardizer.getMultiStandardization(inputs)
-    val outputTrans: Seq[Option[(Double, Double)]] = labels.map { labelSeq =>
-      if (labelSeq.head != null && labelSeq.head.isInstanceOf[Double]) {
+    val outputTrans: Seq[Option[Standardization]] = repOutput.indices.map { i =>
+      if (repOutput(i) != null && repOutput(i).isInstanceOf[Double]) {
+        val labelSeq = labelsTransposed(i)
         Some(Standardizer.getStandardization(labelSeq.asInstanceOf[Seq[Double]].filterNot(_.isNaN())))
       } else {
         None
       }
     }
     val standardInputs = Standardizer.applyStandardization(inputs, inputTrans)
-    val standardLabels = labels.zip(outputTrans).map { case (labelSeq, trans) =>
-      Standardizer.applyStandardization(labelSeq, trans)
-    }
+    val standardLabels = labelsTransposed.zip(outputTrans).map { case (labelSeq, trans) =>
+      Standardizer.applyStandardization(labelSeq, trans).toVector
+    }.transpose
 
-    val baseTrainingResult = baseLearner.train(standardInputs, standardLabels, weights)
-
-    baseTrainingResult.zip(outputTrans).map { case (base, trans) =>
-      new StandardizerTrainingResult(base, Seq(trans) ++ inputTrans)
-    }
+    val baseTrainingResult = baseLearner.train(standardInputs.zip(standardLabels), weights)
+    new MultiTaskStandardizerTrainingResult(baseTrainingResult, outputTrans, inputTrans)
   }
 }
 
 /**
   * Training result wrapping the base learner's training result next to the transformations
   *
-  * @param baseTrainingResult
-  * @param trans
+  * @param baseTrainingResult result of training on the standardized inputs and output
+  * @param outputTrans optional transformation (rescale, offset) of output label
+  * @param inputTrans sequence of optional transformations (rescale, offset) of inputs
   */
 class StandardizerTrainingResult(
                                   baseTrainingResult: TrainingResult,
-                                  trans: Seq[Option[(Double, Double)]]
+                                  outputTrans: Option[Standardization],
+                                  inputTrans: Seq[Option[Standardization]]
                                 ) extends TrainingResult {
   /**
     * Get the model contained in the training result
     *
     * @return the model
     */
-  override def getModel(): Model[PredictionResult[Any]] = new StandardizerModel(baseTrainingResult.getModel(), trans)
+  override def getModel(): Model[PredictionResult[Any]] = new StandardizerModel(baseTrainingResult.getModel(), outputTrans, inputTrans)
+
+  override def getFeatureImportance(): Option[Vector[Double]] = baseTrainingResult.getFeatureImportance()
+
+  override def getPredictedVsActual(): Option[Seq[(Vector[Any], Any, Any)]] = {
+    baseTrainingResult.getPredictedVsActual() match {
+      case None => None
+      case Some(predictedVsActual) => Some(
+        (
+          Standardizer.invertStandardization(predictedVsActual.map(_._1), inputTrans),
+          Standardizer.invertStandardization(predictedVsActual.map(_._2), outputTrans),
+          Standardizer.invertStandardization(predictedVsActual.map(_._3), outputTrans),
+        ).zipped.toSeq
+      )
+    }
+  }
+}
+
+/**
+  * Training result wrapping the base learner's training result next to the transformations
+  *
+  * @param baseTrainingResult result of training on the standardized inputs and output
+  * @param outputTrans optional sequence of transformation (rescale, offset) of output labels
+  * @param inputTrans sequence of optional transformations (rescale, offset) of inputs
+  */
+class MultiTaskStandardizerTrainingResult(
+                                           baseTrainingResult: MultiTaskTrainingResult,
+                                           outputTrans: Seq[Option[Standardization]],
+                                           inputTrans: Seq[Option[Standardization]]
+                                         ) extends MultiTaskTrainingResult {
+  override def getModel(): MultiTaskModel = new ParallelModels(getModels(), baseTrainingResult.getModel().getRealLabels)
+
+  override def getModels(): Seq[Model[PredictionResult[Any]]] = baseTrainingResult.getModels().zipWithIndex.map { case (model, i) =>
+    new StandardizerModel(model, outputTrans(i), inputTrans)
+  }
 
   override def getFeatureImportance(): Option[Vector[Double]] = baseTrainingResult.getFeatureImportance()
+
+  override def getPredictedVsActual(): Option[Seq[(Vector[Any], Seq[Option[Any]], Seq[Option[Any]])]] = {
+    baseTrainingResult.getPredictedVsActual() match {
+      case None => None
+      case Some(predictedVsActual) => Some(
+        (
+          Standardizer.invertStandardization(predictedVsActual.map(_._1), inputTrans),
+          Standardizer.invertStandardizationOption(predictedVsActual.map(_._2), outputTrans),
+          Standardizer.invertStandardizationOption(predictedVsActual.map(_._3), outputTrans)
+        ).zipped.toSeq
+      )
+    }
+  }
 }
 
 /**
-  * Model that wrapps the base model next to the transformations
+  * Model that wraps the base model next to the transformations
   *
-  * @param baseModel
-  * @param trans
-  * @tparam T
+  * @param baseModel model trained on the standardized inputs and outputs
+  * @param outputTrans optional transformation (rescale, offset) of output label
+  * @param inputTrans sequence of optional transformations (rescale, offset) of inputs
+  * @tparam T type of prediction
   */
-class StandardizerModel[T](baseModel: Model[PredictionResult[T]], trans: Seq[Option[(Double, Double)]]) extends Model[PredictionResult[T]] {
+class StandardizerModel[T](
+                            baseModel: Model[PredictionResult[T]],
+                            outputTrans: Option[Standardization],
+                            inputTrans: Seq[Option[Standardization]]
+                          ) extends Model[PredictionResult[T]] {
 
   /**
     * Standardize the inputs and then apply the base model
@@ -103,19 +165,24 @@ class StandardizerModel[T](baseModel: Model[PredictionResult[T]], trans: Seq[Opt
     * @return a predictionresult which includes, at least, the expected outputs
     */
   override def transform(inputs: Seq[Vector[Any]]): StandardizerPrediction[T] = {
-    val standardInputs = Standardizer.applyStandardization(inputs, trans.tail)
-    new StandardizerPrediction(baseModel.transform(standardInputs), trans)
+    val standardInputs = Standardizer.applyStandardization(inputs, inputTrans)
+    new StandardizerPrediction(baseModel.transform(standardInputs), outputTrans, inputTrans)
   }
 }
 
 /**
   * Prediction that wraps the base prediction next to the transformation
   *
-  * @param baseResult
-  * @param trans
-  * @tparam T
+  * @param baseResult result of applying model to standardized inputs
+  * @param outputTrans optional transformation (rescale, offset) of output prediction
+  * @param inputTrans sequence of optional transformations (rescale, offset) of inputs
+  * @tparam T type of prediction
   */
-class StandardizerPrediction[T](baseResult: PredictionResult[T], trans: Seq[Option[(Double, Double)]]) extends PredictionResult[T] {
+class StandardizerPrediction[T](
+                                 baseResult: PredictionResult[T],
+                                 outputTrans: Option[Standardization],
+                                 inputTrans: Seq[Option[Standardization]]
+                               ) extends PredictionResult[T] {
   /**
     * Get the expected values for this prediction
     *
@@ -125,7 +192,7 @@ class StandardizerPrediction[T](baseResult: PredictionResult[T], trans: Seq[Opti
     */
   override def getExpected(): Seq[T] = {
     baseResult.getExpected().map {
-      case x: Double => x * rescale + intercept
+      case x: Double if outputTrans.isDefined => outputTrans.get.invert(x)
       case x: Any => x
     }.asInstanceOf[Seq[T]]
   }
@@ -139,7 +206,7 @@ class StandardizerPrediction[T](baseResult: PredictionResult[T], trans: Seq[Opti
     */
   override def getUncertainty(includeNoise: Boolean = true): Option[Seq[Any]] = {
     baseResult.getUncertainty(includeNoise) match {
-      case Some(x) if trans.head.isDefined => Some(x.map(_.asInstanceOf[Double] * rescale))
+      case Some(x) if outputTrans.isDefined => Some(x.map(_.asInstanceOf[Double] * outputRescale))
       case x: Any => x
     }
   }
@@ -155,19 +222,19 @@ class StandardizerPrediction[T](baseResult: PredictionResult[T], trans: Seq[Opti
   override def getGradient(): Option[Seq[Vector[Double]]] = {
     baseResult.getGradient() match {
       case None => None
-      case Some(x) =>
-        Some(x.map(g => g.zip(trans.tail).map {
-          case (0.0, Some((_, Double.PositiveInfinity))) => 0
-          // If there was a (linear) transformer used on that input, take the slope "m" and rescale by it
-          case (y: Double, Some((_, m))) => y * rescale * m
-          // Otherwise, just rescale by the output transformer
-          case (y, None) => y * rescale
-        }))
+      case Some(gradients) =>
+        Some(
+          gradients.map(g => g.zip(inputTrans).map {
+            // If there was a (linear) transformer used on that input, take the slope "m" and rescale by it
+            case (y, Some(inputStandardization)) => y * outputRescale / inputStandardization.scale
+            // Otherwise, just rescale by the output transformer
+            case (y, None) => y * outputRescale
+          })
+        )
     }
   }
 
-  val rescale = 1.0 / trans.head.map(_._2).getOrElse(1.0)
-  val intercept = trans.head.map(_._1).getOrElse(1.0)
+  val outputRescale = outputTrans.map(_.scale).getOrElse(1.0)
 }
 
 /**
@@ -179,17 +246,17 @@ object Standardizer {
     * The standardizations are just shifts and rescale.  The shift is by the mean and the re-scale is by the variance
     *
     * @param values to get a standardizer for
-    * @return
+    * @return (shift, rescaling)
     */
-  def getStandardization(values: Seq[Double]): (Double, Double) = {
+  def getStandardization(values: Seq[Double]): Standardization = {
     val mean = values.sum / values.size
     val scale = Math.sqrt(values.map(v => Math.pow(v - mean, 2)).sum / values.size)
 
     // If there is zero variance, then the scaling doesn't matter; default to 1.0
-    if (scale > 0) {
-      (mean, 1.0 / scale)
+    if (scale > 0 && scale < Double.PositiveInfinity) {
+      Standardization(mean, scale)
     } else {
-      (mean, 1.0)
+      Standardization(mean, 1.0)
     }
   }
 
@@ -201,28 +268,59 @@ object Standardizer {
     * @param values sequence of vectors to be standardized
     * @return sequence of standardization, each as an option
     */
-  def getMultiStandardization(values: Seq[Vector[Any]]): Seq[Option[(Double, Double)]] = {
+  def getMultiStandardization(values: Seq[Vector[Any]]): Seq[Option[Standardization]] = {
     val rep = values.head
     rep.indices.map { i =>
       rep(i) match {
-        case x: Double =>
-          Some(getStandardization(values.map(r => r(i).asInstanceOf[Double])))
-        case x: Any => None
+        case _: Double => Some(getStandardization(values.map(r => r(i).asInstanceOf[Double])))
+        case _: Any => None
       }
     }
   }
 
   /**
-    * Apply the standardization to vectors, which should result in an output with zero mean and unit variance
+    * Apply the standardizations to vectors, which should result in an output with zero mean and unit variance
     *
     * @param input to standardize
     * @param trans transformtions to apply.  None means no transformation
     * @return sequence of standardized vectors
     */
-  def applyStandardization(input: Seq[Vector[Any]], trans: Seq[Option[(Double, Double)]]): Seq[Vector[Any]] = {
+  def applyStandardization(input: Seq[Vector[Any]], trans: Seq[Option[Standardization]]): Seq[Vector[Any]] = {
     input.map { r =>
       r.zip(trans).map {
-        case (x: Double, Some(t)) => (x - t._1) * t._2
+        case (x: Double, Some(t)) => t.apply(x)
+        case (x: Any, _) => x
+      }
+    }
+  }
+
+  /**
+    * Invert the standardizations on vectors.
+    *
+    * @param input  to invert the standardization
+    * @param trans  transformations to un-apply. None means no transformation
+    * @return       sequence of restored vectors
+    */
+  def invertStandardization(input: Seq[Vector[Any]], trans: Seq[Option[Standardization]]): Seq[Vector[Any]] = {
+    input.map { r =>
+      r.zip(trans).map {
+        case (x: Double, Some(t)) => t.invert(x)
+        case (x: Any, _) => x
+      }
+    }
+  }
+
+  /**
+    * Invert the standardizations on sequences of optional values.
+    *
+    * @param input  to invert the standardization, if the value is defined
+    * @param trans  transformations to un-apply. None means no transformation
+    * @return       sequence of restored vectors
+    */
+  def invertStandardizationOption(input: Seq[Seq[Option[Any]]], trans: Seq[Option[Standardization]]): Seq[Seq[Option[Any]]] = {
+    input.map { r =>
+      r.zip(trans).map {
+        case (Some(x: Double), Some(t)) => Some(t.invert(x))
         case (x: Any, _) => x
       }
     }
@@ -235,11 +333,21 @@ object Standardizer {
     * @param trans transformation to apply
     * @return sequence of standardized values
     */
-  def applyStandardization(input: Seq[Any], trans: Option[(Double, Double)]): Seq[Any] = {
+  def applyStandardization(input: Seq[Any], trans: Option[Standardization]): Seq[Any] = {
     if (trans.isEmpty) return input
+    input.asInstanceOf[Seq[Double]].map(trans.get.apply)
+  }
 
-    input.asInstanceOf[Seq[Double]].map { r =>
-      (r - trans.get._1) * trans.get._2
-    }
+  /**
+    * Invert the standardization on a sequence of values
+    *
+    * @param input  to invert the standardization
+    * @param trans  transformation to un-apply
+    * @return       sequence of restored values
+    */
+  def invertStandardization(input: Seq[Any], trans: Option[Standardization]): Seq[Any] = {
+    if (trans.isEmpty) return input
+    input.asInstanceOf[Seq[Double]].map(trans.get.invert)
   }
+
 }
diff --git a/src/main/scala/io/citrine/lolo/trees/Nodes.scala b/src/main/scala/io/citrine/lolo/trees/Nodes.scala
index a3dd5582..d9c7c633 100644
--- a/src/main/scala/io/citrine/lolo/trees/Nodes.scala
+++ b/src/main/scala/io/citrine/lolo/trees/Nodes.scala
@@ -34,7 +34,7 @@ abstract class TrainingNode[T <: AnyVal, S](
   def getFeatureImportance(): mutable.ArraySeq[Double]
 }
 
-trait ModelNode[T <: PredictionResult[Any]] extends Serializable {
+trait ModelNode[+T <: PredictionResult[Any]] extends Serializable {
   def transform(input: Vector[AnyVal]): (T, TreeMeta)
 
   /**
@@ -82,7 +82,6 @@ trait ModelNode[T <: PredictionResult[Any]] extends Serializable {
   * @param split           to decide which branch to take
   * @param left            branch node
   * @param right           branch node
-  * @param numFeatures     number of features, used for Shapley computation
   * @param outputDimension dimension of model output, used for Shapley computation
   *                        1 for single-task regression, or equal to the number of classification categories.
   * @param trainingWeight  weight of training data in subtree (i.e. size of unweighted training set)
@@ -92,7 +91,6 @@ class InternalModelNode[T <: PredictionResult[Any]](
                                                      split: Split,
                                                      left: ModelNode[T],
                                                      right: ModelNode[T],
-                                                     numFeatures: Int,
                                                      outputDimension: Int,
                                                      trainingWeight: Double
                                                    ) extends ModelNode[T] {
@@ -200,13 +198,13 @@ class TrainingLeaf[T](
     * @return lightweight prediction node
     */
   def getNode(): ModelNode[PredictionResult[T]] = {
-    new ModelLeaf(leafLearner.train(trainingData).getModel().asInstanceOf[Model[PredictionResult[T]]], depth, trainingData.head._1.length, trainingData.size.toDouble)
+    new ModelLeaf(leafLearner.train(trainingData).getModel().asInstanceOf[Model[PredictionResult[T]]], depth, trainingData.size.toDouble)
   }
 
   override def getFeatureImportance(): mutable.ArraySeq[Double] = mutable.ArraySeq.fill(trainingData.head._1.size)(0.0)
 }
 
-class ModelLeaf[T](model: Model[PredictionResult[T]], depth: Int, numFeatures: Int, trainingWeight: Double) extends ModelNode[PredictionResult[T]] {
+class ModelLeaf[T](model: Model[PredictionResult[T]], depth: Int, trainingWeight: Double) extends ModelNode[PredictionResult[T]] {
   override def transform(input: Vector[AnyVal]): (PredictionResult[T], TreeMeta) = {
     (model.transform(Seq(input)), TreeMeta(depth))
   }
diff --git a/src/main/scala/io/citrine/lolo/trees/classification/ClassificationTrainingNode.scala b/src/main/scala/io/citrine/lolo/trees/classification/ClassificationTrainingNode.scala
index 224c2010..05b513bc 100644
--- a/src/main/scala/io/citrine/lolo/trees/classification/ClassificationTrainingNode.scala
+++ b/src/main/scala/io/citrine/lolo/trees/classification/ClassificationTrainingNode.scala
@@ -1,7 +1,6 @@
 package io.citrine.lolo.trees.classification
 
-import breeze.linalg.DenseVector
-import io.citrine.lolo.trees.splits.{ClassificationSplitter, NoSplit, Split, Splitter}
+import io.citrine.lolo.trees.splits.{NoSplit, Split, Splitter}
 import io.citrine.lolo.trees.{InternalModelNode, ModelNode, TrainingLeaf, TrainingNode}
 import io.citrine.lolo.{Learner, PredictionResult}
 
@@ -49,13 +48,13 @@ class ClassificationTrainingNode(
     * @return lightweight prediction node
     */
   override def getNode(): ModelNode[PredictionResult[Char]] = new InternalModelNode(
-    split, leftChild.getNode(), rightChild.getNode(), numClasses, numFeatures, trainingData.size.toDouble
+    split, leftChild.getNode(), rightChild.getNode(), numClasses, trainingData.size.toDouble
   )
 
   override def getFeatureImportance(): scala.collection.mutable.ArraySeq[Double] = {
     val improvement = deltaImpurity
-    var ans = leftChild.getFeatureImportance().zip(rightChild.getFeatureImportance()).map(p => p._1 + p._2)
-    ans(split.getIndex) = ans(split.getIndex) + improvement
+    val ans = leftChild.getFeatureImportance().zip(rightChild.getFeatureImportance()).map(p => p._1 + p._2)
+    ans(split.getIndex()) = ans(split.getIndex()) + improvement
     ans
   }
 }
diff --git a/src/main/scala/io/citrine/lolo/trees/classification/ClassificationTree.scala b/src/main/scala/io/citrine/lolo/trees/classification/ClassificationTree.scala
index 9f196ab1..3e8c6397 100644
--- a/src/main/scala/io/citrine/lolo/trees/classification/ClassificationTree.scala
+++ b/src/main/scala/io/citrine/lolo/trees/classification/ClassificationTree.scala
@@ -102,7 +102,7 @@ class ClassificationTrainingResult(
 
   /* Grab the feature influences */
   lazy val importance = rootTrainingNode.getFeatureImportance()
-  lazy val importanceNormalized = {
+  private lazy val importanceNormalized = {
     if (Math.abs(importance.sum) > 0) {
       importance.map(_ / importance.sum)
     } else {
diff --git a/src/main/scala/io/citrine/lolo/trees/multitask/MultiTaskTrainingNode.scala b/src/main/scala/io/citrine/lolo/trees/multitask/MultiTaskTrainingNode.scala
index e080ea47..4a76c6f8 100644
--- a/src/main/scala/io/citrine/lolo/trees/multitask/MultiTaskTrainingNode.scala
+++ b/src/main/scala/io/citrine/lolo/trees/multitask/MultiTaskTrainingNode.scala
@@ -2,36 +2,84 @@ package io.citrine.lolo.trees.multitask
 
 import io.citrine.lolo.PredictionResult
 import io.citrine.lolo.linear.GuessTheMeanLearner
+import io.citrine.lolo.trees.regression.RegressionTrainingLeaf
 import io.citrine.lolo.trees.splits.{MultiTaskSplitter, NoSplit}
 import io.citrine.lolo.trees.{InternalModelNode, ModelNode, TrainingLeaf}
 
+import scala.collection.mutable
 import scala.util.Random
 
 /**
-  * Node in a multi-task training tree, which can produce nodes for its model trees
+  * Node in a multi-task training tree, which can produce nodes for its model trees.
+  * Splits are chosen using a MultiTaskSplitter, which considers the sum impurity decrease across all tasks.
   *
   * @param inputs data on which to select splits and form models
+  * @param numFeatures to randomly select from at each split (negative values indicate that all features should be considered)
+  * @param maxDepth to grow the tree to
+  * @param minInstances minimum number of training instances per leaf
+  * @param randomizePivotLocation whether to generate splits randomly between the data points
+  * @param rng random number generator, for reproducibility
   */
 class MultiTaskTrainingNode(
                              inputs: Seq[(Vector[AnyVal], Array[AnyVal], Double)],
+                             numFeatures: Int,
+                             maxDepth: Int,
+                             minInstances: Int,
                              randomizePivotLocation: Boolean = false,
                              rng: Random = Random
                            ) {
 
   // Compute a split
-  val (split, _) = MultiTaskSplitter(rng = rng).getBestSplit(inputs, inputs.head._1.size, 1, randomizePivotLocation)
+  val (split, deltaImpurity) = if (maxDepth <= 0) {
+    (new NoSplit, 0.0)
+  } else {
+    MultiTaskSplitter(rng = rng, randomizePivotLocation = randomizePivotLocation).getBestSplit(inputs, numFeatures, minInstances)
+  }
 
   // Try to construct left and right children
   val (leftChild: Option[MultiTaskTrainingNode], rightChild: Option[MultiTaskTrainingNode]) = split match {
     case _: NoSplit => (None, None)
     case _: Any =>
       val (leftData, rightData) = inputs.partition(row => split.turnLeft(row._1))
-      (Some(new MultiTaskTrainingNode(leftData, randomizePivotLocation, rng = rng)), Some(new MultiTaskTrainingNode(rightData, randomizePivotLocation, rng = rng)))
+      (
+        Some(new MultiTaskTrainingNode(leftData, numFeatures, maxDepth - 1, minInstances, randomizePivotLocation, rng = rng)),
+        Some(new MultiTaskTrainingNode(rightData, numFeatures, maxDepth - 1, minInstances, randomizePivotLocation, rng = rng))
+      )
+  }
+
+  // get feature importance for the i'th label
+  def getFeatureImportance(index: Int): mutable.ArraySeq[Double] = {
+    // Filter out "missing" values, which are NaN for regression and 0 for encoded categoricals
+    val label = inputs.head._2(index)
+    val reducedData = if (label.isInstanceOf[Double]) {
+      inputs.map(x => (x._1, x._2(index).asInstanceOf[Double], x._3)).filterNot(_._2.isNaN)
+    } else {
+      inputs.map(x => (x._1, x._2(index).asInstanceOf[Char], x._3)).filter(_._2 > 0)
+    }
+    // Compute the valid data for each child
+    val (left, right) = reducedData.partition(r => split.turnLeft(r._1))
+
+    // get feature importance from the children if they exist, or from this node if it is a leaf
+    (leftChild, rightChild) match {
+      case (Some(theLeftChild), Some(theRightChild)) if left.nonEmpty && right.nonEmpty =>
+        val ans = theLeftChild.getFeatureImportance(index).zip(theRightChild.getFeatureImportance(index)).map(p => p._1 + p._2)
+        ans(split.getIndex()) = ans(split.getIndex()) + deltaImpurity
+        ans
+      case (Some(theLeftChild), _) if left.nonEmpty =>
+        theLeftChild.getFeatureImportance(index)
+      case (_, Some(theRightChild)) if right.nonEmpty =>
+        theRightChild.getFeatureImportance(index)
+      case (_, _) =>
+        if (label.isInstanceOf[Double]) {
+          new RegressionTrainingLeaf(reducedData.asInstanceOf[Seq[(Vector[AnyVal], Double, Double)]], GuessTheMeanLearner(), 1).getFeatureImportance()
+        } else {
+          new TrainingLeaf[Char](reducedData.asInstanceOf[Seq[(Vector[AnyVal], Char, Double)]], GuessTheMeanLearner(), 1).getFeatureImportance()
+        }
+    }
   }
 
   // Construct the model node for the `index`th label
   def getNode(index: Int): ModelNode[PredictionResult[Any]] = {
-
     // Filter out "missing" values, which are NaN for regression and 0 for encoded categoricals
     val label = inputs.head._2(index)
     val reducedData = if (label.isInstanceOf[Double]) {
@@ -43,40 +91,36 @@ class MultiTaskTrainingNode(
     val (left, right) = reducedData.partition(r => split.turnLeft(r._1))
 
     // Construct an internal node if the children are defined and actually have valid data
-    val node = if (leftChild.isDefined && rightChild.isDefined && left.nonEmpty && right.nonEmpty) {
-      // Not a case because of erasure :-(
-      if (label.isInstanceOf[Double]) {
-        new InternalModelNode[PredictionResult[Double]](
-          split,
-          leftChild.get.getNode(index).asInstanceOf[ModelNode[PredictionResult[Double]]],
-          rightChild.get.getNode(index).asInstanceOf[ModelNode[PredictionResult[Double]]],
-          numFeatures=inputs.head._1.length,
-          outputDimension = 0,  // Don't support multitask SHAP at this time.
-          trainingWeight = reducedData.length.toDouble
-        )
-      } else {
-        if (!label.isInstanceOf[Char]) throw new IllegalArgumentException("Training data wasn't double or char")
-        new InternalModelNode[PredictionResult[Char]](
-          split,
-          leftChild.get.getNode(index).asInstanceOf[ModelNode[PredictionResult[Char]]],
-          rightChild.get.getNode(index).asInstanceOf[ModelNode[PredictionResult[Char]]],
-          numFeatures=inputs.head._1.length,
-          outputDimension = 0,  // Don't support multitask SHAP at this time.
-          trainingWeight = reducedData.length.toDouble
-        )
-      }
-    } else if (leftChild.isDefined && left.nonEmpty) {
-      leftChild.get.getNode(index)
-    } else if (rightChild.isDefined && right.nonEmpty) {
-      rightChild.get.getNode(index)
-    } else {
-      // If there are no children or the children don't have valid data for this label, emit a leaf (with GTM)
-      if (label.isInstanceOf[Double]) {
-        new TrainingLeaf[Double](reducedData.asInstanceOf[Seq[(Vector[AnyVal], Double, Double)]], new GuessTheMeanLearner(), 1).getNode()
-      } else {
-        new TrainingLeaf[Char](reducedData.asInstanceOf[Seq[(Vector[AnyVal], Char, Double)]], new GuessTheMeanLearner(), 1).getNode()
-      }
+    (leftChild, rightChild) match {
+      case (Some(theLeftChild), Some(theRightChild)) if left.nonEmpty && right.nonEmpty =>
+        if (label.isInstanceOf[Double]) {
+          new InternalModelNode[PredictionResult[Double]](
+            split,
+            theLeftChild.getNode(index).asInstanceOf[ModelNode[PredictionResult[Double]]],
+            theRightChild.getNode(index).asInstanceOf[ModelNode[PredictionResult[Double]]],
+            outputDimension = 0,  // Don't support multitask SHAP at this time.
+            trainingWeight = reducedData.length.toDouble
+          )
+        } else {
+          if (!label.isInstanceOf[Char]) throw new IllegalArgumentException("Training data wasn't double or char")
+          new InternalModelNode[PredictionResult[Char]](
+            split,
+            theLeftChild.getNode(index).asInstanceOf[ModelNode[PredictionResult[Char]]],
+            theRightChild.getNode(index).asInstanceOf[ModelNode[PredictionResult[Char]]],
+            outputDimension = 0,  // Don't support multitask SHAP at this time.
+            trainingWeight = reducedData.length.toDouble
+          )
+        }
+      case (Some(theLeftChild), _) if left.nonEmpty =>
+        theLeftChild.getNode(index)
+      case (_, Some(theRightChild)) if right.nonEmpty =>
+        theRightChild.getNode(index)
+      case (_, _) =>
+        if (label.isInstanceOf[Double]) {
+          new RegressionTrainingLeaf(reducedData.asInstanceOf[Seq[(Vector[AnyVal], Double, Double)]], GuessTheMeanLearner(), 1).getNode()
+        } else {
+          new TrainingLeaf[Char](reducedData.asInstanceOf[Seq[(Vector[AnyVal], Char, Double)]], GuessTheMeanLearner(), 1).getNode()
+        }
     }
-    node.asInstanceOf[ModelNode[PredictionResult[Any]]]
   }
 }
diff --git a/src/main/scala/io/citrine/lolo/trees/multitask/MultiTaskTree.scala b/src/main/scala/io/citrine/lolo/trees/multitask/MultiTaskTree.scala
index ef713f9f..f5850ef5 100644
--- a/src/main/scala/io/citrine/lolo/trees/multitask/MultiTaskTree.scala
+++ b/src/main/scala/io/citrine/lolo/trees/multitask/MultiTaskTree.scala
@@ -4,32 +4,41 @@ import io.citrine.lolo.encoders.CategoricalEncoder
 import io.citrine.lolo.trees.ModelNode
 import io.citrine.lolo.trees.classification.ClassificationTree
 import io.citrine.lolo.trees.regression.RegressionTree
-import io.citrine.lolo.{Model, MultiTaskLearner, PredictionResult, TrainingResult}
+import io.citrine.lolo.{Model, MultiTaskLearner, MultiTaskTrainingResult, ParallelModels, PredictionResult}
 
 import scala.util.Random
 
 /**
-  * Multi-task tree learner, which produces multiple decision trees with the same split structure
+  * A tree learner that operates on multiple labels.
   *
+  * @param numFeatures to random select from at each split (numbers less than 0 indicate that all features are used)
+  * @param maxDepth to grow the tree to
+  * @param minLeafInstances minimum number of training instances per leaf
+  * @param randomizePivotLocation whether to generate splits randomly between the data points
+  * @param rng random number generator, for reproducibility
   */
 case class MultiTaskTreeLearner(
+                                 numFeatures: Int = -1,
+                                 maxDepth: Int = 30,
+                                 minLeafInstances: Int = 1,
                                  randomizePivotLocation: Boolean = false,
                                  rng: Random = Random
                                ) extends MultiTaskLearner {
 
   /**
-    * Train a model
+    * Construct one regression or classification tree for each label.
     *
-    * @param inputs  to train on
-    * @param labels  sequence of sequences of labels
-    * @param weights for the training rows, if applicable
-    * @return training result containing a model
+    * @param trainingData   to train on
+    * @param weights  for the training rows, if applicable
+    * @return         sequence of models, one for each label
     */
-  override def train(inputs: Seq[Vector[Any]], labels: Seq[Seq[Any]], weights: Option[Seq[Double]]): Seq[TrainingResult] = {
-    val labelsTransposed = labels.toVector.transpose
+  override def train(trainingData: Seq[(Vector[Any], Vector[Any])], weights: Option[Seq[Double]]): MultiTaskTreeTrainingResult = {
+    val (inputs, labels) = trainingData.unzip
+    val repInput = inputs.head
+    val repOutput = labels.head
+    val labelIndices = repOutput.indices
 
     /* Create encoders for any categorical features */
-    val repInput = inputs.head
     val inputEncoders: Seq[Option[CategoricalEncoder[Any]]] = repInput.zipWithIndex.map { case (v, i) =>
       if (v.isInstanceOf[Double]) {
         None
@@ -40,30 +49,43 @@ case class MultiTaskTreeLearner(
     val encodedInputs = inputs.map(r => CategoricalEncoder.encodeInput(r, inputEncoders))
 
     /* Create encoders for any categorical labels */
-    val repOutput = labels.map(_.head)
     val outputEncoders: Seq[Option[CategoricalEncoder[Any]]] = repOutput.zipWithIndex.map { case (v, i) =>
       if (v.isInstanceOf[Double]) {
         None
       } else {
-        Some(CategoricalEncoder.buildEncoder(labels(i).filterNot(_ == null)))
+        Some(CategoricalEncoder.buildEncoder(labels.map(_ (i)).filterNot(_ == null)))
       }
     }
-    val encodedLabels = labelsTransposed.map(CategoricalEncoder.encodeInput(_, outputEncoders))
+    val encodedLabels = labels.map(CategoricalEncoder.encodeInput(_, outputEncoders))
 
     // Encode the inputs, outputs, and filter out zero weight rows
     val collectedData = inputs.indices.map { i =>
       (encodedInputs(i), encodedLabels(i).toArray, weights.map(_ (i)).getOrElse(1.0))
     }.filter(_._3 > 0.0)
 
+    /* If the number of features isn't specified, use all of them */
+    val numFeaturesActual = if (numFeatures > 0) {
+      numFeatures
+    } else {
+      collectedData.head._1.size
+    }
+
     // Construct the training tree
-    val root = new MultiTaskTrainingNode(collectedData, randomizePivotLocation)
+    val root = new MultiTaskTrainingNode(
+      inputs = collectedData,
+      numFeatures = numFeaturesActual,
+      maxDepth = maxDepth,
+      minInstances = minLeafInstances,
+      randomizePivotLocation = randomizePivotLocation,
+      rng = rng
+    )
 
     // Construct the model trees
-    val nodes = labels.indices.map(root.getNode)
+    val nodes = labelIndices.map(root.getNode)
 
     // Stick the model trees into RegressionTree and ClassificationTree objects
-    val models = labels.indices.map { i =>
-      if (labels(i).head.isInstanceOf[Double]) {
+    val models = labelIndices.map { i =>
+      if (repOutput(i).isInstanceOf[Double]) {
         new RegressionTree(
           nodes(i).asInstanceOf[ModelNode[PredictionResult[Double]]],
           inputEncoders
@@ -77,16 +99,34 @@ case class MultiTaskTreeLearner(
       }
     }
 
-    // Wrap the models in dead-simple training results and return
-    models.map(new MultiTaskTreeTrainingResult(_))
+    val sumFeatureImportance: Vector[Double] = {
+      val startingImportances = Vector.fill(repInput.length)(0.0)
+      labelIndices.foldLeft(startingImportances) { (importance, i) =>
+        root.getFeatureImportance(i).toVector.zip(importance).map(p => p._1 + p._2)
+      }
+    }
+
+    new MultiTaskTreeTrainingResult(models, sumFeatureImportance)
   }
 }
 
-class MultiTaskTreeTrainingResult(model: Model[PredictionResult[Any]]) extends TrainingResult {
-  /**
-    * Get the model contained in the training result
-    *
-    * @return the model
-    */
-  override def getModel(): Model[PredictionResult[Any]] = model
+class MultiTaskTreeTrainingResult(
+                                   models: Seq[Model[PredictionResult[Any]]],
+                                   featureImportance: Vector[Double]
+                                 ) extends MultiTaskTrainingResult {
+  val model = new ParallelModels(models, models.map(_.isInstanceOf[RegressionTree]))
+  private lazy val importanceNormalized = {
+    if (Math.abs(featureImportance.sum) > 0) {
+      featureImportance.map(_ / featureImportance.sum)
+    } else {
+      featureImportance.map(_ => 1.0 / featureImportance.size)
+    }
+  }
+
+  override def getModel(): ParallelModels = model
+
+  override def getModels(): Seq[Model[PredictionResult[Any]]] = models
+
+  override def getFeatureImportance(): Option[Vector[Double]] = Some(importanceNormalized)
+
 }
diff --git a/src/main/scala/io/citrine/lolo/trees/regression/RegressionTrainingLeaf.scala b/src/main/scala/io/citrine/lolo/trees/regression/RegressionTrainingLeaf.scala
index 9e7c6d40..8fc45830 100644
--- a/src/main/scala/io/citrine/lolo/trees/regression/RegressionTrainingLeaf.scala
+++ b/src/main/scala/io/citrine/lolo/trees/regression/RegressionTrainingLeaf.scala
@@ -21,7 +21,7 @@ class RegressionTrainingLeaf(
     * @return lightweight prediction node
     */
   def getNode(): ModelNode[PredictionResult[Double]] = {
-    new ModelLeaf(model.asInstanceOf[Model[PredictionResult[Double]]], depth, 1, trainingData.size.toDouble)
+    new ModelLeaf(model.asInstanceOf[Model[PredictionResult[Double]]], depth, trainingData.size.toDouble)
   }
 
   /**
@@ -33,7 +33,7 @@ class RegressionTrainingLeaf(
     importance match {
       case Some(x) =>
         // Compute the weighted sum of the label, the square label, and the weights
-        val expectations: (Double, Double, Double) = trainingData.map { case (v, l, w) =>
+        val expectations: (Double, Double, Double) = trainingData.map { case (_, l, w) =>
           (l * w, l * l * w, w)
         }.reduce((u: (Double, Double, Double), v: (Double, Double, Double)) => (u._1 + v._1, u._2 + v._2, u._3 + v._3))
         // Use those sums to compute the variance as E[x^2] - E[x]^2
diff --git a/src/main/scala/io/citrine/lolo/trees/regression/RegressionTrainingNode.scala b/src/main/scala/io/citrine/lolo/trees/regression/RegressionTrainingNode.scala
index e7c81dd8..213a3bda 100644
--- a/src/main/scala/io/citrine/lolo/trees/regression/RegressionTrainingNode.scala
+++ b/src/main/scala/io/citrine/lolo/trees/regression/RegressionTrainingNode.scala
@@ -39,7 +39,7 @@ class RegressionTrainingNode(
     * @return lightweight prediction node
     */
   override def getNode(): ModelNode[PredictionResult[Double]] = {
-    new InternalModelNode[PredictionResult[Double]](split, leftChild.getNode(), rightChild.getNode(), numFeatures, 1, trainingData.size.toDouble)
+    new InternalModelNode[PredictionResult[Double]](split, leftChild.getNode(), rightChild.getNode(), 1, trainingData.size.toDouble)
   }
 
   /**
diff --git a/src/main/scala/io/citrine/lolo/trees/regression/RegressionTree.scala b/src/main/scala/io/citrine/lolo/trees/regression/RegressionTree.scala
index 97b79be9..6649a5b4 100644
--- a/src/main/scala/io/citrine/lolo/trees/regression/RegressionTree.scala
+++ b/src/main/scala/io/citrine/lolo/trees/regression/RegressionTree.scala
@@ -16,8 +16,10 @@ import scala.util.Random
   *
   * @param numFeatures to randomly select from at each split (default: all)
   * @param maxDepth    to grow the tree to
+  * @param minLeafInstances minimum number of training instances per leaf
   * @param leafLearner learner to train the leaves with
-  * @param minLeafInstances minimum number of instances per leaf
+  * @param splitter to determine the best split of the node data
+  * @param rng random number generator, for reproducibility
   */
 case class RegressionTreeLearner(
                                   numFeatures: Int = -1,
@@ -103,7 +105,7 @@ class RegressionTreeTrainingResult(
                                   ) extends TrainingResult {
   lazy val model = new RegressionTree(rootTrainingNode.getNode(), encoders)
   lazy val importance = rootTrainingNode.getFeatureImportance()
-  lazy val importanceNormalized = {
+  private lazy val importanceNormalized = {
     if (Math.abs(importance.sum) > 0) {
       importance.map(_ / importance.sum)
     } else {
@@ -135,7 +137,7 @@ class RegressionTree(
     * Apply the model by calling predict and wrapping the results
     *
     * @param inputs to apply the model to
-    * @return a predictionresult which includes only the expected outputs
+    * @return a prediction result which includes only the expected outputs
     */
   override def transform(inputs: Seq[Vector[Any]]): RegressionTreeResult = {
     new RegressionTreeResult(
diff --git a/src/main/scala/io/citrine/lolo/trees/splits/ClassificationSplitter.scala b/src/main/scala/io/citrine/lolo/trees/splits/ClassificationSplitter.scala
index 095384a9..dd28ede2 100644
--- a/src/main/scala/io/citrine/lolo/trees/splits/ClassificationSplitter.scala
+++ b/src/main/scala/io/citrine/lolo/trees/splits/ClassificationSplitter.scala
@@ -37,7 +37,7 @@ case class ClassificationSplitter(randomizedPivotLocation: Boolean = false, rng:
 
       /* Use different spliters for each type */
       val (possibleSplit, possibleImpurity) = rep._1(index) match {
-        case _: Double => getBestRealSplit(data, calculator, index, minInstances)
+        case _: Double => Splitter.getBestRealSplit[Char](data, calculator, index, minInstances, randomizedPivotLocation, rng)
         case _: Char => getBestCategoricalSplit(data, calculator, index, minInstances)
         case _: Any => throw new IllegalArgumentException("Trying to split unknown feature type")
       }
@@ -56,57 +56,6 @@ case class ClassificationSplitter(randomizedPivotLocation: Boolean = false, rng:
     }
   }
 
-  /**
-    * Find the best split on a continuous variable
-    *
-    * If randomizePivotLocation is true, the split pivots are drawn from a uniform random distribution between the two
-    * data points.  Each such pivot results in the same data split, but randomization can improve generalizability,
-    * particularly as part of an ensemble (i.e. random forests).
-    *
-    * @param data                 to split
-    * @param index                of the feature to split on
-    * @param minCount minimum number of data points to allow in each of the resulting splits
-    * @param randomizePivotLocation whether generate splits randomly between the data points (default: false)
-    * @return the best split of this feature
-    */
-  def getBestRealSplit(
-                        data: Seq[(Vector[AnyVal], Char, Double)],
-                        calculator: GiniCalculator,
-                        index: Int,
-                        minCount: Int,
-                        randomizePivotLocation: Boolean = false
-                      ): (RealSplit, Double) = {
-    /* Pull out the feature that's considered here and sort by it */
-    val thinData = data.map(dat => (dat._1(index).asInstanceOf[Double], dat._2, dat._3)).sortBy(_._1)
-
-    var bestPurity = Double.MaxValue
-    var bestPivot = Double.MinValue
-
-    /* Move the data from the right to the left partition one value at a time */
-    calculator.reset()
-    (0 until data.size - minCount).foreach { j =>
-      val totalPurity = calculator.add(thinData(j)._2, thinData(j)._3)
-
-      /* Keep track of the best split, avoiding splits in the middle of constant sets of feature values
-         It is really important for performance to keep these checks together so
-         1) there is only one branch and
-         2) it is usually false
-       */
-      if (totalPurity < bestPurity && j + 1 >= minCount && Math.abs((thinData(j + 1)._1 - thinData(j)._1) / thinData(j)._1) > 1.0e-9) {
-        bestPurity = totalPurity
-        /* Try pivots at the midpoints between consecutive member values */
-        val left = thinData(j + 1)._1
-        val right = thinData(j)._1
-        bestPivot = if (randomizePivotLocation) {
-          (left - right) * rng.nextDouble() + right
-        } else {
-          (left + right) / 2.0
-        }
-      }
-    }
-    (new RealSplit(index, bestPivot), bestPurity)
-  }
-
   def getBestCategoricalSplit(
                                data: Seq[(Vector[AnyVal], Char, Double)],
                                calculator: GiniCalculator,
diff --git a/src/main/scala/io/citrine/lolo/trees/splits/MultiTaskSplitter.scala b/src/main/scala/io/citrine/lolo/trees/splits/MultiTaskSplitter.scala
index 11608d13..f1fc89b5 100644
--- a/src/main/scala/io/citrine/lolo/trees/splits/MultiTaskSplitter.scala
+++ b/src/main/scala/io/citrine/lolo/trees/splits/MultiTaskSplitter.scala
@@ -8,20 +8,20 @@ import scala.util.Random
   *
   * Created by maxhutch on 11/29/16.
   */
-case class MultiTaskSplitter(rng: Random = Random) {
+case class MultiTaskSplitter(randomizePivotLocation: Boolean = false, rng: Random = Random) extends Splitter[Array[AnyVal]] {
 
   /**
     * Get the best split, considering numFeature random features (w/o replacement)
     *
     * @param data        to split
     * @param numFeatures to consider, randomly
+    * @param minInstances the minimum number of data points on a split node
     * @return a split object that optimally divides data
     */
   def getBestSplit(
                     data: Seq[(Vector[AnyVal], Array[AnyVal], Double)],
                     numFeatures: Int,
-                    minInstances: Int,
-                    randomizePivotLocation: Boolean = false
+                    minInstances: Int
                   ): (Split, Double) = {
     var bestSplit: Split = new NoSplit()
     var bestImpurity = Double.MaxValue
@@ -37,7 +37,7 @@ case class MultiTaskSplitter(rng: Random = Random) {
 
       /* Use different spliters for each type */
       val (possibleSplit, possibleImpurity) = rep._1(index) match {
-        case _: Double => getBestRealSplit(data, calculator, index, minInstances, randomizePivotLocation)
+        case _: Double => Splitter.getBestRealSplit[Array[AnyVal]](data, calculator, index, minInstances, randomizePivotLocation, rng)
         case _: Char => getBestCategoricalSplit(data, calculator, index, minInstances)
         case _: Any => throw new IllegalArgumentException("Trying to split unknown feature type")
       }
@@ -56,49 +56,6 @@ case class MultiTaskSplitter(rng: Random = Random) {
     }
   }
 
-  /**
-    * Find the best split on a continuous variable
-    *
-    * @param data  to split
-    * @param index of the feature to split on
-    * @return the best split of this feature
-    */
-  def getBestRealSplit(
-                        data: Seq[(Vector[AnyVal], Array[AnyVal], Double)],
-                        calculator: MultiImpurityCalculator,
-                        index: Int,
-                        minCount: Int,
-                        randomizePivotLocation: Boolean = false
-                      ): (Split, Double) = {
-    /* Pull out the feature that's considered here and sort by it */
-    val thinData = data.map(dat => (dat._1(index).asInstanceOf[Double], dat._2, dat._3)).sortBy(_._1)
-    val features = thinData.map(x => x._1)
-
-    /* Move the data from the right to the left partition one value at a time */
-    calculator.reset()
-    val pivots = (0 until data.size - minCount).flatMap { j =>
-      val totalImpurity = calculator.add(thinData(j)._2, thinData(j)._3)
-      if (j + 1 >= minCount && Math.abs((features(j + 1) - features(j)) / features(j)) > 1.0e-9) {
-        val left = features(j + 1)
-        val right = features(j)
-        val pivot = if (randomizePivotLocation) {
-          (left - right) * rng.nextDouble() + right
-        } else {
-          (left + right) / 2.0
-        }
-        Some((pivot, totalImpurity))
-      } else {
-        None
-      }
-    }
-    if (pivots.isEmpty) {
-      (new NoSplit, Double.MaxValue)
-    } else {
-      val best = pivots.minBy(_._2)
-      (new RealSplit(index, best._1), best._2)
-    }
-  }
-
   /**
     * Get find the best categorical splitter.
     *
diff --git a/src/main/scala/io/citrine/lolo/trees/splits/RegressionSplitter.scala b/src/main/scala/io/citrine/lolo/trees/splits/RegressionSplitter.scala
index 5a6b7881..63ddc8e7 100644
--- a/src/main/scala/io/citrine/lolo/trees/splits/RegressionSplitter.scala
+++ b/src/main/scala/io/citrine/lolo/trees/splits/RegressionSplitter.scala
@@ -48,7 +48,7 @@ case class RegressionSplitter(randomizePivotLocation: Boolean = false, rng: Rand
 
       /* Use different spliters for each type */
       val (possibleSplit, possibleVariance) = rep._1(index) match {
-        case _: Double => getBestRealSplit(data, calculator.copy(), index, minInstances, randomizePivotLocation)
+        case _: Double => Splitter.getBestRealSplit[Double](data, calculator.copy(), index, minInstances, randomizePivotLocation, rng)
         case _: Char => getBestCategoricalSplit(data, calculator.copy(), index, minInstances)
         case _: Any => throw new IllegalArgumentException("Trying to split unknown feature type")
       }
@@ -67,59 +67,6 @@ case class RegressionSplitter(randomizePivotLocation: Boolean = false, rng: Rand
     }
   }
 
-  /**
-    * Find the best split on a continuous variable.
-    *
-    * If randomizePivotLocation is true, the split pivots are drawn from a uniform random distribution between the two
-    * data points.  Each such pivot results in the same data split, but randomization can improve generalizability,
-    * particularly as part of an ensemble (i.e. random forests).
-    *
-    * @param data  to split
-    * @param index of the feature to split on
-    * @param minCount minimum number of data points to allow in each of the resulting splits
-    * @param randomizePivotLocation whether generate splits randomly between the data points (default: false)
-    * @return the best split of this feature
-    */
-  def getBestRealSplit(
-                        data: Seq[(Vector[AnyVal], Double, Double)],
-                        calculator: VarianceCalculator,
-                        index: Int,
-                        minCount: Int,
-                        randomizePivotLocation: Boolean = false
-                      ): (RealSplit, Double) = {
-    /* Pull out the feature that's considered here and sort by it */
-    val thinData = data.map(dat => (dat._1(index).asInstanceOf[Double], dat._2, dat._3)).sortBy(_._1)
-
-
-    /* Base cases for iteration */
-    var bestVariance = Double.MaxValue
-    var bestPivot = Double.MinValue
-
-    /* Move the data from the right to the left partition one value at a time */
-    calculator.reset()
-    (0 until data.size - minCount).foreach { j =>
-      val totalVariance = calculator.add(thinData(j)._2, thinData(j)._3)
-
-      /* Keep track of the best split, avoiding splits in the middle of constant sets of feature values
-         It is really important for performance to keep these checks together so
-         1) there is only one branch and
-         2) it is usually false
-       */
-      val left = thinData(j + 1)._1
-      val right = thinData(j)._1
-      if (totalVariance < bestVariance && j + 1 >= minCount && Splitter.isDifferent(left, right)) {
-        bestVariance = totalVariance
-        /* Try pivots at the midpoints between consecutive member values */
-        bestPivot = if (randomizePivotLocation) {
-          (left - right) * rng.nextDouble() + right
-        } else {
-          (left + right) / 2.0
-        }
-      }
-    }
-    (new RealSplit(index, bestPivot), bestVariance)
-  }
-
   /**
     * Get find the best categorical splitter.
     *
diff --git a/src/main/scala/io/citrine/lolo/trees/splits/Splitter.scala b/src/main/scala/io/citrine/lolo/trees/splits/Splitter.scala
index a5c83f20..044caf0d 100644
--- a/src/main/scala/io/citrine/lolo/trees/splits/Splitter.scala
+++ b/src/main/scala/io/citrine/lolo/trees/splits/Splitter.scala
@@ -1,5 +1,9 @@
 package io.citrine.lolo.trees.splits
 
+import io.citrine.lolo.trees.impurity.ImpurityCalculator
+
+import scala.util.Random
+
 /**
   * Created by maxhutch on 7/5/17.
   */
@@ -17,4 +21,50 @@ object Splitter {
       Math.abs((x1 - x2) / (x1 + x2)) > 1.0e-9
     }
   }
+
+  /**
+    * Find the best split on a continuous variable.
+    * This is a typical implementation for decision trees: it finds the pivot that maximizes the decrease in impurity.
+    * Some splitters might forego this method and implement different ways of choosing a pivot point.
+    *
+    * @param data  to split
+    * @param index of the feature to split on
+    * @param minCount minimum number of data points to allow in each of the resulting splits
+    * @param randomizePivotLocation whether generate splits by drawing a random value uniformly between the two split points.
+    *                               This can improve generalizability, particularly as part of an ensemble.
+    * @return the best split of this feature
+    */
+  def getBestRealSplit[T](
+                           data: Seq[(Vector[AnyVal], T, Double)],
+                           calculator: ImpurityCalculator[T],
+                           index: Int,
+                           minCount: Int,
+                           randomizePivotLocation: Boolean = false,
+                           rng: Random = Random
+                         ): (Split, Double) = {
+    /* Pull out the feature that's considered here and sort by it */
+    val thinData = data.map(dat => (dat._1(index).asInstanceOf[Double], dat._2, dat._3)).sortBy(_._1)
+    val features = thinData.map(x => x._1)
+
+    var bestImpurity = Double.MaxValue
+    var bestPivot = Double.MinValue
+
+    /* Move the data from the right to the left partition one value at a time */
+    calculator.reset()
+    (0 until data.size - minCount).foreach { j =>
+      val totalImpurity = calculator.add(thinData(j)._2, thinData(j)._3)
+      val left = features(j + 1)
+      val right = features(j)
+      if (totalImpurity < bestImpurity && j + 1 >= minCount && isDifferent(left, right)) {
+        bestImpurity = totalImpurity
+        /* Try pivots at the midpoints between consecutive member values */
+        bestPivot = if (randomizePivotLocation) {
+          (left - right) * rng.nextDouble() + right
+        } else {
+          (left + right) / 2.0
+        }
+      }
+    }
+    (new RealSplit(index, bestPivot), bestImpurity)
+  }
 }
diff --git a/src/main/scala/io/citrine/lolo/util/LoloPyDataLoader.scala b/src/main/scala/io/citrine/lolo/util/LoloPyDataLoader.scala
index 2f428c88..720b67d3 100644
--- a/src/main/scala/io/citrine/lolo/util/LoloPyDataLoader.scala
+++ b/src/main/scala/io/citrine/lolo/util/LoloPyDataLoader.scala
@@ -4,8 +4,7 @@ import java.io.{ByteArrayInputStream, ByteArrayOutputStream, ObjectInputStream,
 import java.nio.ByteBuffer
 import java.nio.ByteOrder
 import java.util.zip._
-
-import io.citrine.lolo.PredictionResult
+import io.citrine.lolo.{MultiTaskModelPredictionResult, PredictionResult}
 
 /**
   * Tool used to transfer data from LoloPy to the JVM
@@ -21,7 +20,7 @@ object LoloPyDataLoader {
     * @param bigEndian Whether the numbers are is big-endian or not
     * @return The array as a Scala array
     */
-  def getFeatureArray(input: Array[Byte], numAttributes : Integer, bigEndian: Boolean) : Seq[Vector[Double]] = {
+  def getFeatureArray(input: Array[Byte], numAttributes: Integer, bigEndian: Boolean) : Seq[Vector[Double]] = {
     // Get ordering
     val ordering = if (bigEndian) ByteOrder.BIG_ENDIAN else ByteOrder.LITTLE_ENDIAN
 
@@ -82,6 +81,16 @@ object LoloPyDataLoader {
     send1DArray(predResults)
   }
 
+  /**
+    * Generate the results of a multitask regression model, which are assumed to be all doubles
+    * @param predictionResult result of predicting on a multitask model
+    * @return Byte array of doubles in native system order (the caller must then reshape the result into a 2d array)
+    */
+  def getMultiRegressionExpected(predictionResult: MultiTaskModelPredictionResult): Array[Byte] = {
+    val predResults = predictionResult.getExpected().asInstanceOf[Seq[Seq[Double]]].flatten
+    send1DArray(predResults)
+  }
+
   /**
     * Send the training entry importance scores to the Python client
     * @param predictionResult Prediction result object
@@ -101,6 +110,31 @@ object LoloPyDataLoader {
     send1DArray(predResults)
   }
 
+  /**
+    * Get the uncertainties of a multitask regression model, which are assumed to be all doubles
+    * @param predictionResult result of predicting on a multitask model
+    * @return Byte array of doubles in native system order (the caller must then reshape the result into a 2d array)
+    */
+  def getMultiRegressionUncertainty(predictionResult: MultiTaskModelPredictionResult): Array[Byte] = {
+    val uncertaintyResults = predictionResult.getUncertainty().get.asInstanceOf[Seq[Seq[Double]]].flatten
+    send1DArray(uncertaintyResults)
+  }
+
+  /**
+    * Get the correlation coefficients between the uncertainties of a multitask regression model.
+    * By calling this method for all (i, j) pairs, one can construct a correlation matrix.
+    * Combined with getMultiRegressionUncertainty, one can construct the covariance matrix.
+    *
+    * @param predictionResult result of predicting on a multitask model
+    * @param i index of the first output
+    * @param j index of the second output
+    * @return Byte array of doubles in native system order
+    */
+  def getRegressionCorrelation(predictionResult: MultiTaskModelPredictionResult, i: Int, j: Int): Array[Byte] = {
+    val correlationResults = predictionResult.getUncertaintyCorrelation(i, j).get
+    send1DArray(correlationResults)
+  }
+
   /**
     * Prepare to send a 1D array of Doubles by converting it to a byte array
     * @param data Data to be sent


Test patch:
diff --git a/python/lolopy/tests/test_learners.py b/python/lolopy/tests/test_learners.py
index c82c9aae..9503aeea 100644
--- a/python/lolopy/tests/test_learners.py
+++ b/python/lolopy/tests/test_learners.py
@@ -8,7 +8,7 @@
 )
 from sklearn.exceptions import NotFittedError
 from sklearn.metrics import r2_score, accuracy_score, log_loss
-from sklearn.datasets import load_boston, load_iris
+from sklearn.datasets import load_iris, load_diabetes, load_linnerud
 from unittest import TestCase, main
 import pickle as pkl
 import numpy as np
@@ -35,7 +35,7 @@ def test_rf_regressor(self):
         rf = RandomForestRegressor(random_seed = 31247895)
 
         # Train the model
-        X, y = load_boston(True)
+        X, y = load_diabetes(return_X_y=True)
 
         # Make sure we get a NotFittedError
         with self.assertRaises(NotFittedError):
@@ -52,10 +52,10 @@ def test_rf_regressor(self):
         y_import = rf.get_importance_scores(X[:100, :])
         self.assertEqual((100, len(X)), y_import.shape)
 
-        # Basic test for functionality. R^2 above 0.98 was measured on 27Dec18
+        # Basic test for functionality. R^2 above 0.88 was measured on 2021-12-09
         score = r2_score(y_pred, y)
         print('R^2:', score)
-        self.assertGreater(score, 0.98)
+        self.assertGreater(score, 0.88)
 
         # Test with weights (make sure it doesn't crash)
         rf.fit(X, y, [2.0]*len(y))
@@ -70,15 +70,43 @@ def test_rf_regressor(self):
         self.assertTrue((y_std >= 0).all())  # They must be positive
         self.assertGreater(np.std(y_std), 0)  # Must have a variety of values
 
+        # For a single output, the covariance matrix is just the standard deviation squared
+        _, y_cov = rf.predict(X, return_cov_matrix=True)
+        assert np.all(y_cov.flatten() == y_std ** 2)
+
         # Make sure the detach operation functions
         rf.clear_model()
         self.assertIsNone(rf.model_)
 
+    def test_rf_multioutput_regressor(self):
+        rf = RandomForestRegressor(random_seed=810355)
+        # A regression dataset with 3 outputs
+        X, y = load_linnerud(return_X_y=True)
+        num_data = len(X)
+        num_outputs = y.shape[1]
+
+        rf.fit(X, y)
+        y_pred, y_std = rf.predict(X, return_std=True)
+        _, y_cov = rf.predict(X, return_cov_matrix=True)
+
+        # Assert that all returned values have the correct shape
+        assert y_pred.shape == (num_data, num_outputs)
+        assert y_std.shape == (num_data, num_outputs)
+        assert y_cov.shape == (num_data, num_outputs, num_outputs)
+
+        # The covariance matrices should be symmetric and the diagonals should be the squares of the standard deviations.
+        assert np.all(y_cov[:, 0, 1] == y_cov[:, 1, 0])
+        assert np.all(y_cov[:, 0, 0] == y_std[:, 0] ** 2)
+
+        # Make sure the user cannot call predict with both return_std and return_cov_matrix True
+        with self.assertRaises(ValueError):
+            rf.predict(X, return_std=True, return_cov_matrix=True)
+
     def test_classifier(self):
         rf = RandomForestClassifier(random_seed = 34789)
 
         # Load in the iris dataset
-        X, y = load_iris(True)
+        X, y = load_iris(return_X_y=True)
         rf.fit(X, y)
 
         # Predict the probability of membership in each class
@@ -105,7 +133,7 @@ def test_serialization(self):
         rf2 = pkl.loads(data)
 
         # Load in the iris dataset and train model
-        X, y = load_iris(True)
+        X, y = load_iris(return_X_y=True)
         rf.fit(X, y)
 
         # Try saving and loading the model
@@ -121,7 +149,7 @@ def test_regression_tree(self):
         tree = RegressionTreeLearner()
 
         # Make sure it trains and predicts properly
-        X, y = load_boston(True)
+        X, y = load_diabetes(return_X_y=True)
         tree.fit(X, y)
 
         # Make sure the prediction works
@@ -134,7 +162,7 @@ def test_regression_tree(self):
         tree.max_depth = 2
         tree.fit(X, y)
         y_pred = tree.predict(X)
-        self.assertAlmostEqual(0.695574477, r2_score(y, y_pred))  # Result is deterministic
+        self.assertAlmostEqual(0.433370098, r2_score(y, y_pred))  # Result is deterministic
 
         # Constrain the tree to a single node, using minimum count per split
         tree = RegressionTreeLearner(min_leaf_instances=1000)
@@ -194,7 +222,7 @@ def test_extra_random_trees_regressor(self):
         rf = ExtraRandomTreesRegressor(random_seed = 378456)
 
         # Train the model
-        X, y = load_boston(True)
+        X, y = load_diabetes(return_X_y=True)
 
         # Make sure we get a NotFittedError
         with self.assertRaises(NotFittedError):
@@ -211,10 +239,10 @@ def test_extra_random_trees_regressor(self):
         y_import = rf.get_importance_scores(X[:100, :])
         self.assertEqual((100, len(X)), y_import.shape)
 
-        # Basic test for functionality. R^2 above 0.98 was measured on 2020-04-06
+        # Basic test for functionality. R^2 above 0.88 was measured on 2021-12-09
         score = r2_score(y_pred, y)
         print("R2: ", score)
-        self.assertGreater(score, 0.98)
+        self.assertGreater(score, 0.88)
 
         # Test with weights (make sure it doesn't crash)
         rf.fit(X, y, [2.0]*len(y))
@@ -237,7 +265,7 @@ def test_extra_random_trees_classifier(self):
         rf = ExtraRandomTreesClassifier(random_seed = 378456)
 
         # Load in the iris dataset
-        X, y = load_iris(True)
+        X, y = load_iris(return_X_y=True)
         rf.fit(X, y)
 
         # Predict the probability of membership in each class
@@ -264,7 +292,7 @@ def test_serialization(self):
         rf2 = pkl.loads(data)
 
         # Load in the iris dataset and train model
-        X, y = load_iris(True)
+        X, y = load_iris(return_X_y=True)
         rf.fit(X, y)
 
         # Try saving and loading the model
diff --git a/src/test/scala/io/citrine/lolo/PerformanceTest.scala b/src/test/scala/io/citrine/lolo/PerformanceTest.scala
index eecfe1c1..700e4e7a 100644
--- a/src/test/scala/io/citrine/lolo/PerformanceTest.scala
+++ b/src/test/scala/io/citrine/lolo/PerformanceTest.scala
@@ -91,17 +91,17 @@ class PerformanceTest {
 
   def testMultitaskOverhead(N: Int): (Double, Double) = {
     val subset = trainingData.take(N)
-    val inputs = subset.map(_._1)
-    val realLabels = subset.map(_._2)
+    val (inputs, realLabels) = subset.unzip
     val catLabels = realLabels.map(_ > realLabels.sum / realLabels.size)
+    val labels = Vector(realLabels, catLabels).transpose
 
     val trainSingle: Double = Stopwatch.time({
-      new Bagger(new RegressionTreeLearner()).train(inputs.zip(realLabels)).getLoss()
-      new Bagger(new ClassificationTreeLearner()).train(inputs.zip(catLabels)).getLoss()
+      new Bagger(RegressionTreeLearner()).train(inputs.zip(realLabels)).getLoss()
+      new Bagger(ClassificationTreeLearner()).train(inputs.zip(catLabels)).getLoss()
     }, minRun = 1, maxRun = 1)
 
     val trainMulti: Double = Stopwatch.time({
-      new MultiTaskBagger(new MultiTaskTreeLearner()).train(inputs, Seq(realLabels, catLabels)).map(_.getLoss())
+      MultiTaskBagger(MultiTaskTreeLearner()).train(inputs.zip(labels)).getLoss()
     }, minRun = 1, maxRun = 1)
 
     (trainMulti, trainSingle)
diff --git a/src/test/scala/io/citrine/lolo/TestUtils.scala b/src/test/scala/io/citrine/lolo/TestUtils.scala
index d6e3b6a6..25299f47 100644
--- a/src/test/scala/io/citrine/lolo/TestUtils.scala
+++ b/src/test/scala/io/citrine/lolo/TestUtils.scala
@@ -1,11 +1,11 @@
 package io.citrine.lolo
 
 import breeze.stats.distributions.{RandBasis, ThreadLocalRandomGenerator}
+import io.citrine.lolo.linear.LinearRegressionLearner
+import io.citrine.lolo.stats.StatsUtils.variance
 import io.citrine.lolo.stats.functions.Friedman
 import org.apache.commons.math3.random.MersenneTwister
 
-import scala.collection.generic.CanBuildFrom
-import scala.collection.mutable
 import scala.util.{Random, Try}
 
 /**
@@ -46,6 +46,28 @@ object TestUtils {
     }
   }
 
+  /** Given a univariate data set, construct a corresponding data set with the desired Pearson correlation coefficient.
+    * The procedure is to generate new data randomly, calculate the residuals of a linear regression,
+    * and then create a suitable linear combination of X and the residuals.
+    * Source: https://stats.stackexchange.com/a/313138
+    *
+    * @param X    sequence of values
+    * @param rho  desired Pearson correlation coefficient
+    * @param rng  random number generator
+    * @return     sequence of values that have desired correlation with X
+    */
+  def makeLinearCorrelatedData(X: Seq[Double], rho: Double, rng: Random = new Random()): Seq[Double] = {
+    require(rho >= -1.0 && rho <= 1.0, "correlation coefficient must be between -1.0 and 1.0")
+    val Y = Seq.fill(X.length)(rng.nextGaussian())
+    val linearLearner = LinearRegressionLearner()
+    val linearModel = linearLearner.train(X.zip(Y).map { case (x, y) => (Vector(x), y) } ).getModel()
+    val yPred = linearModel.transform(X.map(Vector(_))).getExpected()
+    val residuals = Y.zip(yPred).map { case (actual, predicted) => actual - predicted }
+    val stdX = math.sqrt(variance(X))
+    val stdResiduals = math.sqrt(variance(residuals))
+    X.zip(residuals).map { case (x, residual) => rho * stdResiduals * x + math.sqrt(1 - rho * rho) * stdX * residual}
+  }
+
   def iterateTrainingData(
                            cols: Int,
                            function: (Seq[Double] => Double) = Friedman.friedmanGrosseSilverman,
diff --git a/src/test/scala/io/citrine/lolo/bags/MultiTaskBaggerTest.scala b/src/test/scala/io/citrine/lolo/bags/MultiTaskBaggerTest.scala
index 1f86d998..a0a12c64 100644
--- a/src/test/scala/io/citrine/lolo/bags/MultiTaskBaggerTest.scala
+++ b/src/test/scala/io/citrine/lolo/bags/MultiTaskBaggerTest.scala
@@ -29,12 +29,11 @@ class MultiTaskBaggerTest {
       TestUtils.generateTrainingData(512, 12, noise = 0.1, function = Friedman.friedmanSilverman),
       inputBins = Seq((0, 8))
     )
-    val inputs = trainingData.map(_._1)
-    val labels = trainingData.map(_._2)
+    val reshapedTrainingData = trainingData.map { case (input, label) => (input, Vector(label))}
     val DTLearner = MultiTaskTreeLearner()
     val baggedLearner = MultiTaskBagger(DTLearner, numBags = trainingData.size, randBasis = TestUtils.getBreezeRandBasis(10478L))
-    val RFMeta = baggedLearner.train(inputs, Seq(labels)).head
-    val RF = RFMeta.getModel()
+    val RFMeta = baggedLearner.train(reshapedTrainingData)
+    val RF = RFMeta.getModels().head
 
 
     val results = RF.transform(trainingData.map(_._1))
@@ -64,20 +63,19 @@ class MultiTaskBaggerTest {
             val sigmaObsAndSigmaMean: Seq[(Double, Double)] = (1 to 20).flatMap { _ =>
               val trainingDataTmp = TestUtils.generateTrainingData(nRows, nCols, noise = 0.0, function = _ => 0.0, seed = rng.nextLong())
               val trainingData = trainingDataTmp.map { x => (x._1, x._2 + noiseLevel * rng.nextDouble()) }
-              val inputs = trainingData.map(_._1)
-              val labels = trainingData.map(_._2)
+              val reshapedTrainingData = trainingData.map { case (input, label) => (input, Vector(label))}
               val baggedLearner = MultiTaskBagger(baseLearner, numBags = nBags, uncertaintyCalibration = true, randBasis = TestUtils.getBreezeRandBasis(7835178L))
-              val RFMeta = baggedLearner.train(inputs, Seq(labels)).head
-              val RF = RFMeta.getModel()
+              val RFMeta = baggedLearner.train(reshapedTrainingData)
+              val RF = RFMeta.getModels().head
               val results = RF.transform(trainingData.take(4).map(_._1))
 
               val sigmaMean: Seq[Double] = results.getUncertainty(observational = false).get.asInstanceOf[Seq[Double]]
-              sigmaMean.zip(results.asInstanceOf[BaggedMultiResult].getStdDevMean().get).foreach { case (a, b) =>
+              sigmaMean.zip(results.asInstanceOf[MultiPredictionBaggedResult].getStdDevMean().get).foreach { case (a, b) =>
                 assert(a == b, s"Expected getUncertainty(observational=false)=getStdDevMean() for $configDescription")
               }
 
               val sigmaObs: Seq[Double] = results.getUncertainty().get.asInstanceOf[Seq[Double]]
-              sigmaObs.zip(results.asInstanceOf[BaggedMultiResult].getStdDevObs().get).foreach { case (a, b) =>
+              sigmaObs.zip(results.asInstanceOf[MultiPredictionBaggedResult].getStdDevObs().get).foreach { case (a, b) =>
                 assert(a == b, s"Expected getUncertainty()=getStdDevObs() for $configDescription")
               }
 
@@ -141,12 +139,11 @@ class MultiTaskBaggerTest {
       TestUtils.generateTrainingData(128, 12, noise = 0.1, function = Friedman.friedmanSilverman),
       inputBins = Seq((0, 8)), responseBins = Some(8)
     )
-    val inputs = trainingData.map(_._1)
-    val labels = trainingData.map(_._2)
+    val reshapedTrainingData = trainingData.map { case (input, label) => (input, Vector(label))}
     val DTLearner = MultiTaskTreeLearner()
     val baggedLearner = MultiTaskBagger(DTLearner, numBags = trainingData.size, randBasis = TestUtils.getBreezeRandBasis(478L))
-    val RFMeta = baggedLearner.train(inputs, Seq(labels)).head
-    val RF = RFMeta.getModel()
+    val RFMeta = baggedLearner.train(reshapedTrainingData)
+    val RF = RFMeta.getModels().head
 
     /* Inspect the results */
     val results = RF.transform(trainingData.map(_._1))
@@ -165,23 +162,77 @@ class MultiTaskBaggerTest {
     assert(results.getGradient().isEmpty, "Returned a gradient when there shouldn't be one")
   }
 
-  /**
-    * Test that multi-task with dense labels works, and remembers all its inputs
-    */
+
+  /** Test that a multi-task bagged model properly stores and transposes individual trees, and remembers labels. */
   @Test
-  def testMixed(): Unit = {
-    /* Setup some data */
-    val raw: Seq[(Vector[Double], Double)] = TestUtils.generateTrainingData(256, 12, noise = 0.1, function = Friedman.friedmanSilverman)
-    val inputs: Seq[Vector[Double]] = raw.map(_._1)
-    val realLabel: Seq[Double] = raw.map(_._2)
+  def testCombinedMultiTaskModel(): Unit = {
+    val numTrain = 256
+    val numBags = 64
+    val numTest = 32
+    val raw: Seq[(Vector[Double], Double)] = TestUtils.generateTrainingData(numTrain, 12, noise = 0.1, function = Friedman.friedmanSilverman)
+    val (inputs: Seq[Vector[Double]], realLabel: Seq[Double]) = raw.unzip
     val catLabel: Seq[Boolean] = raw.map(_._2 > realLabel.max / 2.0)
-    val DTLearner = MultiTaskTreeLearner()
-    val baggedLearner = MultiTaskBagger(DTLearner, numBags = inputs.size, biasLearner = Some(new RegressionTreeLearner(maxDepth = 2)), randBasis = TestUtils.getBreezeRandBasis(78495L))
-    val RFMeta = baggedLearner.train(inputs, Seq(realLabel, catLabel)).last
-    val RF = RFMeta.getModel()
+    val labels = Vector(realLabel, catLabel).transpose
+
+    val learner = MultiTaskTreeLearner()
+    val baggedLearner = MultiTaskBagger(
+      learner,
+      numBags = numBags,
+      biasLearner = Some(RegressionTreeLearner(maxDepth = 2)),
+      randBasis = TestUtils.getBreezeRandBasis(78495L)
+    )
+    val RF = baggedLearner.train(inputs.zip(labels))
+
+    val testInputs = inputs.take(numTest)
+    val predictionResult = RF.getModel().transform(testInputs)
+    assert(predictionResult.predictions.length == numBags)
+
+    // The prediction made by the full model and the prediction made by just the categorical model should agree
+    // and both be equal to the training label.
+    val expected = predictionResult.getExpected()
+    val expectedCat = RF.getModels()(1).transform(testInputs).getExpected()
+    (0 until numTest).foreach { i =>
+      assert(expected(i)(1) == catLabel(i))
+      assert(catLabel(i) == expectedCat(i))
+    }
+  }
+
+  /** Test the correlation coefficient calculation. */
+  @Test
+  def testCorrelation(): Unit = {
+    val numTrain = 256
+    val numBags = 64
+    val numTest = 32
+    val trainingRho = 0.45 // desired correlation between two real-valued training labels
+    val raw: Seq[(Vector[Double], Double)] = TestUtils.generateTrainingData(numTrain, 12, noise = 0.1, function = Friedman.friedmanSilverman)
+    val (inputs: Seq[Vector[Double]], realLabel: Seq[Double]) = raw.unzip
+    val catLabel: Seq[Boolean] = raw.map(_._2 > realLabel.max / 2.0)
+    val correlatedLabel: Seq[Double] = TestUtils.makeLinearCorrelatedData(realLabel, trainingRho)
+    val labels = Vector(realLabel, catLabel, correlatedLabel).transpose
+
+    val learner = MultiTaskTreeLearner()
+    val baggedLearner = MultiTaskBagger(
+      learner,
+      numBags = numBags,
+      biasLearner = Some(RegressionTreeLearner(maxDepth = 2)),
+      randBasis = TestUtils.getBreezeRandBasis(78495L)
+    )
+    val RF = baggedLearner.train(inputs.zip(labels)).getModel()
+
+    val testInputs = TestUtils.generateTrainingData(numTest, 12, function = Friedman.friedmanSilverman).map(_._1)
+    val predictionResult = RF.transform(testInputs)
+
+    Seq(true, false).foreach { observational =>
+      // All real-valued predictions should be perfectly correlated with themselves
+      assert(predictionResult.getUncertaintyCorrelation(0, 0, observational).get == Seq.fill(numTest)(1.0))
+      // Correlation with a non-real-valued label should be empty
+      assert(predictionResult.getUncertaintyCorrelation(0, 1, observational).isEmpty)
+      // Otherwise, all we can assert is that -1.0 <= rho <= 1.0
+      predictionResult.getUncertaintyCorrelation(0, 2, observational).get.foreach { calcRho =>
+        assert(calcRho >= -1.0 && calcRho <= 1.0)
+      }
+    }
 
-    val catResults = RF.transform(inputs).getExpected().asInstanceOf[Seq[Boolean]]
-    assert(catResults.zip(catLabel).forall(p => p._1 == p._2))
   }
 
   /**
@@ -209,15 +260,15 @@ class MultiTaskBaggerTest {
         x
       }
     )
+    val labels = Vector(sparseReal, sparseCat).transpose
 
     val DTLearner = MultiTaskTreeLearner()
     val baggedLearner = MultiTaskBagger(DTLearner, numBags = inputs.size, biasLearner = Some(GuessTheMeanLearner(rng = rng)), randBasis = TestUtils.getBreezeRandBasis(7839L))
-    val trainingResult = baggedLearner.train(inputs, Seq(sparseReal, sparseCat))
-    val RFMeta = trainingResult.last
-    val RF = RFMeta.getModel()
+    val trainingResult = baggedLearner.train(inputs.zip(labels))
+    val RF = trainingResult.getModels().last
 
     val catResults = RF.transform(inputs).getExpected()
-    val realUncertainty = trainingResult.head.getModel().transform(inputs).getUncertainty().get
+    val realUncertainty = trainingResult.getModels().head.transform(inputs).getUncertainty().get
     assert(realUncertainty.forall(!_.asInstanceOf[Double].isNaN), s"Some uncertainty values were NaN")
 
     val referenceModel = Bagger(ClassificationTreeLearner(), numBags = inputs.size)
@@ -231,12 +282,8 @@ class MultiTaskBaggerTest {
     val multiF1 = ClassificationMetrics.f1scores(catResults, catLabel)
 
     // Make sure we can grab the loss without issue
-    val singleLoss = referenceModel.getLoss().get
-    val multiLoss = RFMeta.getLoss().get
-    val regressionLoss = trainingResult.head.getLoss().get
-    assert(!singleLoss.isNaN, "Single task classification loss was NaN")
-    assert(!multiLoss.isNaN, "Sparse multitask classification loss was NaN")
-    assert(!regressionLoss.isNaN, "Sparse regression loss was NaN")
+    assert(!referenceModel.getLoss().get.isNaN, "Single task classification loss was NaN")
+    assert(!trainingResult.getLoss().get.isNaN, "Sparse multitask loss was NaN")
 
     assert(multiF1 > singleF1, s"Multi-task is under-performing single-task")
     assert(multiF1 <= 1.0, "Multitask classification F1 score was greater than 1.0")
diff --git a/src/test/scala/io/citrine/lolo/learners/RandomForestTest.scala b/src/test/scala/io/citrine/lolo/learners/RandomForestTest.scala
index 86abcf25..f74827f8 100644
--- a/src/test/scala/io/citrine/lolo/learners/RandomForestTest.scala
+++ b/src/test/scala/io/citrine/lolo/learners/RandomForestTest.scala
@@ -3,6 +3,7 @@ package io.citrine.lolo.learners
 import breeze.linalg.DenseMatrix
 import breeze.stats.distributions.Beta
 import io.citrine.lolo.TestUtils
+import io.citrine.lolo.bags.MultiTaskBaggedResult
 import io.citrine.lolo.stats.functions.Friedman
 import org.junit.Test
 import org.scalatest.Assertions._
@@ -48,6 +49,25 @@ class RandomForestTest {
     }
   }
 
+  /** Test that a random forest with multiple outputs produces a multitask bagger. */
+  @Test
+  def testMultitaskForest(): Unit = {
+    val (inputs: Seq[Vector[Double]], realLabel: Seq[Double]) = TestUtils.binTrainingData(
+      TestUtils.generateTrainingData(256, 12, noise = 0.1, function = Friedman.friedmanSilverman, seed = rng.nextLong()),
+      inputBins = Seq((0, 8))
+    ).unzip
+    val catLabel: Seq[Boolean] = realLabel.map(_ > realLabel.max / 2.0)
+    val quadLabel: Seq[Double] = realLabel.map(x => x * x)
+    val labels = Vector(realLabel, catLabel, quadLabel).transpose
+
+    val RFMeta = RandomForest(rng = rng).train(inputs.zip(labels))
+    val model = RFMeta.getModel()
+
+    val results = model.transform(inputs).asInstanceOf[MultiTaskBaggedResult]
+    assert(results.getUncertainty().isDefined)
+    assert(results.getUncertaintyCorrelation(0, 2).isDefined)
+  }
+
   /**
     * Test that the regression forest does the same thing as the classification bagger
     */
diff --git a/src/test/scala/io/citrine/lolo/transformers/FeatureRotatorTest.scala b/src/test/scala/io/citrine/lolo/transformers/FeatureRotatorTest.scala
index c677d072..0747212e 100644
--- a/src/test/scala/io/citrine/lolo/transformers/FeatureRotatorTest.scala
+++ b/src/test/scala/io/citrine/lolo/transformers/FeatureRotatorTest.scala
@@ -1,6 +1,6 @@
 package io.citrine.lolo.transformers
 
-import breeze.linalg.{DenseMatrix, DenseVector, det}
+import breeze.linalg.{DenseMatrix, det}
 import io.citrine.lolo.TestUtils
 import io.citrine.lolo.linear.{GuessTheMeanLearner, LinearRegressionLearner}
 import io.citrine.lolo.stats.functions.Friedman
@@ -248,17 +248,18 @@ class FeatureRotatorTest {
 
     // Sparsify the categorical labels
     val sparseCatLabel = catLabel.map(x => if (rng.nextBoolean()) x else null)
+    val labels = Vector(doubleLabel, sparseCatLabel).transpose
 
     // Train and evaluate rotated models on original and rotated features
-    val baseLearner = new MultiTaskTreeLearner()
-    val rotatedLearner = new MultiTaskFeatureRotator(MultiTaskTreeLearner())
-
-    val baseTrainingResult = baseLearner.train(inputs, Seq(doubleLabel, sparseCatLabel))
-    val baseDoubleModel = baseTrainingResult.head.getModel()
-    val baseCatModel = baseTrainingResult.last.getModel()
-    val rotatedTrainingResult = rotatedLearner.train(inputs, Seq(doubleLabel, sparseCatLabel))
-    val rotatedDoubleModel = rotatedTrainingResult.head.getModel().asInstanceOf[RotatedFeatureModel[Double]]
-    val rotatedCatModel = rotatedTrainingResult.last.getModel().asInstanceOf[RotatedFeatureModel[Boolean]]
+    val baseLearner = MultiTaskTreeLearner()
+    val rotatedLearner = MultiTaskFeatureRotator(MultiTaskTreeLearner())
+
+    val baseTrainingResult = baseLearner.train(inputs.zip(labels))
+    val baseDoubleModel = baseTrainingResult.getModels().head
+    val baseCatModel = baseTrainingResult.getModels().last
+    val rotatedTrainingResult = rotatedLearner.train(inputs.zip(labels))
+    val rotatedDoubleModel = rotatedTrainingResult.getModels().head.asInstanceOf[RotatedFeatureModel[Double]]
+    val rotatedCatModel = rotatedTrainingResult.getModels().last.asInstanceOf[RotatedFeatureModel[Boolean]]
 
     // Check double labels are the same
     val baseDoubleResult = baseDoubleModel.transform(inputs).getExpected()
diff --git a/src/test/scala/io/citrine/lolo/transformers/StandardizerTest.scala b/src/test/scala/io/citrine/lolo/transformers/StandardizerTest.scala
index 045b4178..e125ba95 100644
--- a/src/test/scala/io/citrine/lolo/transformers/StandardizerTest.scala
+++ b/src/test/scala/io/citrine/lolo/transformers/StandardizerTest.scala
@@ -179,15 +179,17 @@ class StandardizerTest {
 
     // Sparsify the categorical labels
     val sparseCatLabel = catLabel.map(x => if (rng.nextBoolean()) x else null)
+    val labels = Vector(doubleLabel, sparseCatLabel).transpose
 
     // Screw up the scale of the double labels
     val scale = 10000.0
     val rescaledLabel = doubleLabel.map(_ * scale)
+    val rescaledLabels = rescaledLabel.zip(sparseCatLabel).map { case (r, c) => Vector(r, c) }
 
     // Train and evaluate standard models on original and rescaled labels
     val standardizer = new MultiTaskStandardizer(MultiTaskTreeLearner())
-    val baseRes = standardizer.train(inputs, Seq(doubleLabel, sparseCatLabel)).last.getModel().transform(inputs).getExpected()
-    val standardRes = standardizer.train(inputs, Seq(rescaledLabel, sparseCatLabel)).last.getModel().transform(inputs).getExpected()
+    val baseRes = standardizer.train(inputs.zip(labels)).getModels().last.transform(inputs).getExpected()
+    val standardRes = standardizer.train(inputs.zip(rescaledLabels)).getModels().last.transform(inputs).getExpected()
     // Train and evaluate unstandardized model on rescaled labels
 
     // Compute metrics for each of the models
diff --git a/src/test/scala/io/citrine/lolo/trees/multitask/MultiTaskTreeTest.scala b/src/test/scala/io/citrine/lolo/trees/multitask/MultiTaskTreeTest.scala
index cc07b733..1144028b 100644
--- a/src/test/scala/io/citrine/lolo/trees/multitask/MultiTaskTreeTest.scala
+++ b/src/test/scala/io/citrine/lolo/trees/multitask/MultiTaskTreeTest.scala
@@ -4,7 +4,7 @@ import io.citrine.lolo.TestUtils
 import io.citrine.lolo.stats.functions.Friedman
 import io.citrine.lolo.stats.metrics.ClassificationMetrics
 import io.citrine.lolo.trees.classification.{ClassificationTree, ClassificationTreeLearner}
-import io.citrine.lolo.trees.regression.RegressionTree
+import io.citrine.lolo.trees.regression.{RegressionTree, RegressionTreeLearner}
 import org.junit.Test
 import org.scalatest.Assertions._
 
@@ -19,9 +19,9 @@ class MultiTaskTreeTest {
 
   /* Setup some data */
   val raw: Seq[(Vector[Double], Double)] = TestUtils.generateTrainingData(1024, 12, noise = 0.1, function = Friedman.friedmanSilverman, seed = rng.nextLong())
-  val inputs: Seq[Vector[Double]] = raw.map(_._1)
-  val realLabel: Seq[Double] = raw.map(_._2)
+  val (inputs: Seq[Vector[Double]], realLabel: Seq[Double]) = raw.unzip
   val catLabel: Seq[Boolean] = raw.map(_._2 > realLabel.max / 2.0)
+  val labels = Vector(realLabel, catLabel).transpose
 
   /**
     * Test that the tree memorizes the data when it is complete
@@ -29,7 +29,7 @@ class MultiTaskTreeTest {
   @Test
   def testTwoLabels(): Unit = {
     val learner = MultiTaskTreeLearner()
-    val models = learner.train(inputs, Seq(realLabel, catLabel)).map(_.getModel())
+    val models = learner.train(inputs.zip(labels)).getModels()
     assert(models.size == 2)
     assert(models.head.isInstanceOf[RegressionTree])
     assert(models.last.isInstanceOf[ClassificationTree])
@@ -52,9 +52,10 @@ class MultiTaskTreeTest {
         x
       }
     )
+    val sparseLabels = Vector(realLabel, sparseCat).transpose
 
     val learner = MultiTaskTreeLearner()
-    val models = learner.train(inputs, Seq(realLabel, sparseCat)).map(_.getModel())
+    val models = learner.train(inputs.zip(sparseLabels)).getModels()
     val realResults = models.head.transform(inputs).getExpected().asInstanceOf[Seq[Double]]
     val catResults = models.last.transform(inputs).getExpected().asInstanceOf[Seq[Boolean]]
 
@@ -74,12 +75,13 @@ class MultiTaskTreeTest {
         x
       }
     )
+    val sparseLabels = Vector(realLabel, sparseCat).transpose
 
-    val learner = new MultiTaskTreeLearner()
-    val models = learner.train(inputs, Seq(realLabel, sparseCat)).map(_.getModel())
+    val learner = MultiTaskTreeLearner()
+    val models = learner.train(inputs.zip(sparseLabels)).getModels()
     val catResults = models.last.transform(inputs).getExpected().asInstanceOf[Seq[Boolean]]
 
-    val reference = new ClassificationTreeLearner()
+    val reference = ClassificationTreeLearner()
       .train(inputs.zip(sparseCat).filterNot(_._2 == null))
       .getModel()
       .transform(inputs)
@@ -91,6 +93,38 @@ class MultiTaskTreeTest {
     assert(multiF1 > singleF1)
   }
 
+  /** Test that the resulting predictions are independent of whether the trees are stored in one model or several models.*/
+  @Test
+  def testSingleModelEquality(): Unit = {
+    // Train twice with the same seed, first outputting two models and then outputting a combined model.
+    val seed = 817235L
+    val rng = new Random(seed)
+    val combinedModelRng = new Random(seed)
+    val learner = MultiTaskTreeLearner(rng = rng)
+    val combinedModelLearner = MultiTaskTreeLearner(rng = combinedModelRng)
+    val models = learner.train(inputs.zip(labels)).getModels()
+    val combinedModel = combinedModelLearner.train(inputs.zip(labels)).getModel()
+
+    // Generate new inputs to test equality on.
+    val testInputs = TestUtils
+      .generateTrainingData(32, 12, noise = 0.1, function = Friedman.friedmanSilverman, seed = rng.nextLong())
+      .map(_._1)
+    val realResults = models.head.transform(testInputs).getExpected().asInstanceOf[Seq[Double]]
+    val catResults = models.last.transform(testInputs).getExpected().asInstanceOf[Seq[Boolean]]
+    val allResults = combinedModel.transform(testInputs).getExpected().asInstanceOf[Seq[Seq[Any]]]
+    assert(Seq(realResults, catResults).transpose == allResults)
+  }
+
+  /** Test that feature importance correctly identifies the 1st feature as the most important for Friedman-Silverman. */
+  @Test
+  def testFeatureImportance(): Unit = {
+    val seed = 253768L
+    val multiTaskLearner = MultiTaskTreeLearner(rng = new Random(seed))
+    val multiTaskTrainingResult = multiTaskLearner.train(inputs.zip(labels))
+    val importances = multiTaskTrainingResult.getFeatureImportance().get
+    assert(importances(1) == importances.max)
+  }
+
 }
 
 
diff --git a/src/test/scala/io/citrine/lolo/trees/splits/MultiTaskSplitterTest.scala b/src/test/scala/io/citrine/lolo/trees/splits/MultiTaskSplitterTest.scala
index 4505cc20..0d6acf66 100644
--- a/src/test/scala/io/citrine/lolo/trees/splits/MultiTaskSplitterTest.scala
+++ b/src/test/scala/io/citrine/lolo/trees/splits/MultiTaskSplitterTest.scala
@@ -27,7 +27,7 @@ class MultiTaskSplitterTest {
     val calculator = MultiImpurityCalculator.build(data.map(_._2), data.map(_._3))
 
     val splitter = MultiTaskSplitter(rng = rng)
-    val (pivot, impurity) = splitter.getBestRealSplit(data, calculator, 0, 1)
+    val (pivot, impurity) = Splitter.getBestRealSplit(data, calculator, 0, 1)
     val turns = data.map(x => pivot.turnLeft(x._1))
     assert(turns == Seq(true, false, false, false))
     assert(impurity == 0.5)
@@ -69,7 +69,7 @@ class MultiTaskSplitterTest {
       (inputs(i), labels(i).asInstanceOf[Array[AnyVal]], weights(i))
     }
     val splitter = MultiTaskSplitter(rng = rng)
-    val (pivot, impurity) = splitter.getBestSplit(data, data.head._1.size, 1)
+    val (pivot, _) = splitter.getBestSplit(data, data.head._1.size, 1)
     assert(!pivot.isInstanceOf[NoSplit])
   }
 }
diff --git a/src/test/scala/io/citrine/lolo/trees/splits/SplitterTest.scala b/src/test/scala/io/citrine/lolo/trees/splits/SplitterTest.scala
index 54cc99aa..b16daabb 100644
--- a/src/test/scala/io/citrine/lolo/trees/splits/SplitterTest.scala
+++ b/src/test/scala/io/citrine/lolo/trees/splits/SplitterTest.scala
@@ -60,7 +60,7 @@ class SplitterTest {
 
     val calculator = new VarianceCalculator(0.0, 0.0, 8.0)
 
-    val (_, variance) = RegressionSplitter().getBestRealSplit(trainingData, calculator, 0, 1)
+    val (_, variance) = Splitter.getBestRealSplit(trainingData, calculator, 0, 1)
     assert(variance == Double.MaxValue, "didn't expect to find a valid split")
   }
 
diff --git a/src/test/scala/io/citrine/lolo/validation/MeritTest.scala b/src/test/scala/io/citrine/lolo/validation/MeritTest.scala
index f7fd07d6..e19b30f1 100644
--- a/src/test/scala/io/citrine/lolo/validation/MeritTest.scala
+++ b/src/test/scala/io/citrine/lolo/validation/MeritTest.scala
@@ -2,14 +2,13 @@ package io.citrine.lolo.validation
 
 import io.citrine.lolo.PredictionResult
 import org.apache.commons.math3.distribution.MultivariateNormalDistribution
+import org.apache.commons.math3.random.MersenneTwister
 import org.junit.Test
 import org.scalatest.Assertions._
 
 import scala.util.{Random, Try}
 
 class MeritTest {
-  val rng = new Random(34578L)
-
   /**
     * Generate test data by adding Gaussian noise to a uniformly distributed response
     *
@@ -26,13 +25,15 @@ class MeritTest {
                             noiseScale: Double = 1.0,
                             uncertaintyCorrelation: Double = 0.0,
                             batchSize: Int = 32,
-                            numBatch: Int = 1
+                            numBatch: Int = 1,
+                            rng: Random = Random
                           ): Iterable[(PredictionResult[Double], Seq[Double])] = {
     val maximumCorrelation = 0.999
 
     val noiseVariance = noiseScale * noiseScale
     val noiseUncertaintyCovariance = noiseVariance * Math.min(uncertaintyCorrelation, maximumCorrelation) // avoid singular matrices
     val errorDistribution = new MultivariateNormalDistribution(
+      new MersenneTwister(rng.nextLong()),
       Array(0.0, 0.0),
       Array(Array(noiseVariance, noiseUncertaintyCovariance), Array(noiseUncertaintyCovariance, noiseVariance))
     )
@@ -60,55 +61,51 @@ class MeritTest {
 
   /**
     * Test that the RMSE comes out correctly
-    *
-    * This has an false negative rate below 1/100
     */
   @Test
   def testRMSE(): Unit = {
-    val pva = getNormalPVA(batchSize = 256, numBatch = 32)
-    val (rmse, uncertainty) = RootMeanSquareError.estimate(pva)
+    val rng = new Random(34578L)
+    val pva = getNormalPVA(batchSize = 256, numBatch = 32, rng = rng)
+    val (rmse, uncertainty) = RootMeanSquareError.estimate(pva, rng = rng)
     assert(Math.abs(rmse - 1.0) < 3 * uncertainty, "RMSE estimate was not accurate enough")
     assert(uncertainty < 0.05, s"RMSE estimate was not precise enough")
   }
 
   /**
     * Test that the coefficient of determination comes out right
-    *
-    * This has an false negative rate below 1/100
     */
   @Test
   def testCoefficientOfDeterminization(): Unit = {
-    val pva = getNormalPVA(noiseScale = 0.1, batchSize = 256, numBatch = 32)
+    val rng = new Random(34578L)
+    val pva = getNormalPVA(noiseScale = 0.1, batchSize = 256, numBatch = 32, rng = rng)
     val expected = 1 - 12 * 0.1 * 0.1
-    val (r2, uncertainty) = CoefficientOfDetermination.estimate(pva)
+    val (r2, uncertainty) = CoefficientOfDetermination.estimate(pva, rng = rng)
     assert(Math.abs(r2 - expected) < 3 * uncertainty, "R^2 estimate was not accurate enough")
     assert(uncertainty < 0.05, s"R^2 estimate was not precise enough")
   }
 
   /**
     * Test that the standard confidence comes out right
-    *
-    * This has an false negative rate below 1/200
     */
   @Test
   def testStandardConfidence(): Unit = {
-    val pva = getNormalPVA(uncertaintyCorrelation = 1.0, batchSize = 256, numBatch = 32)
+    val rng = new Random(34578L)
+    val pva = getNormalPVA(uncertaintyCorrelation = 1.0, batchSize = 256, numBatch = 32, rng = rng)
     val expected = 0.68
-    val (confidence, uncertainty) = StandardConfidence.estimate(pva)
+    val (confidence, uncertainty) = StandardConfidence.estimate(pva, rng = rng)
     assert(Math.abs(confidence - expected) < 3 * uncertainty, "Confidence estimate was not accurate enough")
     assert(uncertainty < 0.05, s"Confidence estimate was not precise enough")
   }
 
   /**
     * Test that the standard error comes out right
-    *
-    * This has an false negative rate below 1/100
     */
   @Test
   def testStandardError(): Unit = {
-    val pva = getNormalPVA(noiseScale = 0.01, uncertaintyCorrelation = 1.0, batchSize = 256, numBatch = 32)
+    val rng = new Random(34578L)
+    val pva = getNormalPVA(noiseScale = 0.01, uncertaintyCorrelation = 1.0, batchSize = 256, numBatch = 32, rng = rng)
     val expected = 1.0
-    val (error, uncertainty) = StandardError().estimate(pva)
+    val (error, uncertainty) = StandardError().estimate(pva, rng = rng)
     assert(Math.abs(error - expected) < 3 * uncertainty, "Standard error estimate was not accurate enough")
     assert(uncertainty < 0.05, s"Standard error estimate was not precise enough")
   }
@@ -116,28 +113,26 @@ class MeritTest {
 
   /**
     * Test that the UncertaintyCorrelation is correct when perfectly correlated
-    *
-    * This has an false negative rate below 1/100
     */
   @Test
   def testPerfectUncertaintyCorrelation(): Unit = {
-    val pva = getNormalPVA(noiseScale = 0.01, uncertaintyCorrelation = 1.0, batchSize = 256, numBatch = 32)
+    val rng = new Random(34578L)
+    val pva = getNormalPVA(noiseScale = 0.01, uncertaintyCorrelation = 1.0, batchSize = 256, numBatch = 32, rng = rng)
     val expected = 1.0
-    val (corr, uncertainty) = UncertaintyCorrelation.estimate(pva)
+    val (corr, uncertainty) = UncertaintyCorrelation.estimate(pva, rng = rng)
     assert(Math.abs(corr - expected) < 3 * uncertainty, "Uncertainty correlation estimate was not accurate enough")
     assert(uncertainty < 0.05, s"Uncertainty correlation estimate was not precise enough")
   }
 
   /**
     * Test that the UncertaintyCorrelation is correct when there is no correlation
-    *
-    * This has an false negative rate below 1/100
     */
   @Test
   def testZeroUncertaintyCorrelation(): Unit = {
-    val pva = getNormalPVA(noiseScale = 0.01, uncertaintyCorrelation = 0.0, batchSize = 256, numBatch = 32)
+    val rng = new Random(34578L)
+    val pva = getNormalPVA(noiseScale = 0.01, uncertaintyCorrelation = 0.0, batchSize = 256, numBatch = 32, rng = rng)
     val expected = 0.0
-    val (corr, uncertainty) = UncertaintyCorrelation.estimate(pva)
+    val (corr, uncertainty) = UncertaintyCorrelation.estimate(pva, rng = rng)
     assert(Math.abs(corr - expected) < 3 * uncertainty, "Uncertainty correlation estimate was not accurate enough")
     assert(uncertainty < 0.05, s"Uncertainty correlation estimate was not precise enough")
   }